breaking: addressing comments from PR #1559

Author: metrizable

src/sagemaker/fw_utils.py

@@ -686,17 +686,19 @@ def _region_supports_debugger(region_name):
 def validate_version_or_image_args(framework_version, py_version, image_name):
     """Checks if version or image arguments are specified.
 
-    Used to validate framework and model arguments to enforce version or image specification.
-    Raises ValueError if version or image arguments are not specified.
+    Validates framework and model arguments to enforce version or image specification.
 
     Args:
-        framework_version (str): the version of the framework
-        py_version (str): the version of python
-        image_name (str): the uri of the image
+        framework_version (str): The version of the framework.
+        py_version (str): The version of python.
+        image_name (str): The URI of the image.
+
+    Raises:
+        ValueError: if `image_name` is None and either `framework_version` or `py_version` is
+                        None.
     """
     if (framework_version is None or py_version is None) and image_name is None:
         raise ValueError(
             "framework_version or py_version was None, yet image_name was also None. "
             "Either specify both framework_version and py_version, or specify image_name."
         )
-    return True

src/sagemaker/mxnet/estimator.py

@@ -74,10 +74,12 @@ def __init__(
                 If ``source_dir`` is specified, then ``entry_point``
                 must point to a file located at the root of ``source_dir``.
             framework_version (str): MXNet version you want to use for executing
-                your model training code. List of supported versions. Defaults to ``None``.
+                your model training code. Defaults to `None`. Required unless
+                ``image_name`` is provided. List of supported versions.
                 https://github.com/aws/sagemaker-python-sdk#mxnet-sagemaker-estimators.
             py_version (str): Python version you want to use for executing your
-                model training code. One of 'py2' or 'py3'. Defaults to ``None``.
+                model training code. One of 'py2' or 'py3'. Defaults to ``None``. Required
+                unless ``image_name`` is provided.
             source_dir (str): Path (absolute, relative or an S3 URI) to a directory
                 with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must

src/sagemaker/mxnet/model.py

@@ -87,9 +87,11 @@ def __init__(
                 hosting. If ``source_dir`` is specified, then ``entry_point``
                 must point to a file located at the root of ``source_dir``.
             framework_version (str): MXNet version you want to use for executing
-                your model training code. Defaults to ``None``.
+                your model training code. Defaults to ``None``. Required unless
+                ``image_name`` is provided.
             py_version (str): Python version you want to use for executing your
-                model training code. Defaults to ``None``.
+                model training code. Defaults to ``None``. Required unless
+                ``image_name`` is provided.
             image (str): A Docker image URI (default: None). If not specified, a
                 default image for MXNet will be used.
 
@@ -112,7 +114,6 @@ def __init__(
             :class:`~sagemaker.model.FrameworkModel` and
             :class:`~sagemaker.model.Model`.
         """
-        # TODO: rename/unify image attribute to match across code base
         validate_version_or_image_args(framework_version, py_version, image)
         if py_version and py_version == "py2":
             logger.warning(

tests/unit/test_fw_utils.py

@@ -1275,13 +1275,16 @@ def test_warn_if_parameter_server_with_multi_gpu(caplog):
 
 
 def test_validate_version_or_image_args():
-    for good_args in [("", "", None), (None, "", ""), ("", None, "")]:
-        kwargs = dict(zip(("framework_version", "py_version", "image_name"), good_args))
-        assert fw_utils.validate_version_or_image_args(**kwargs)
+    good_args = [("1.0", "py3", None), (None, "py3", "my:uri"), ("1.0", None, "my:uri")]
+    for framework_version, py_version, image_name in good_args:
+        assert (
+            fw_utils.validate_version_or_image_args(framework_version, py_version, image_name)
+            is None
+        )
 
 
 def test_validate_version_or_image_args_raises():
-    for bad_args in [(None, None, None), (None, "", None), ("", None, None)]:
-        kwargs = dict(zip(("framework_version", "py_version", "image_name"), bad_args))
+    bad_args = [(None, None, None), (None, "py3", None), ("1.0", None, None)]
+    for framework_version, py_version, image_name in bad_args:
         with pytest.raises(ValueError):
-            fw_utils.validate_version_or_image_args(**kwargs)
+            fw_utils.validate_version_or_image_args(framework_version, py_version, image_name)

tests/unit/test_mxnet.py

@@ -37,7 +37,7 @@
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = "ml.c4.4xlarge"
 ACCELERATOR_TYPE = "ml.eia.medium"
-IMAGE = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.6.0-cpu-py3"
+IMAGE = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.4.0-cpu-py3"
 COMPILATION_JOB_NAME = "{}-{}".format("compilation-sagemaker-mxnet", TIMESTAMP)
 FRAMEWORK = "mxnet"
 ROLE = "Dummy"
@@ -85,25 +85,25 @@ def sagemaker_session():
     return session
 
 
-def is_mms_version(mxnet_version):
+def _is_mms_version(mxnet_version):
     return parse_version(MXNetModel._LOWEST_MMS_VERSION) <= parse_version(mxnet_version)
 
 
 @pytest.fixture()
 def skip_if_mms_version(mxnet_version):
-    if is_mms_version(mxnet_version):
+    if _is_mms_version(mxnet_version):
         pytest.skip("Skipping because this version uses MMS")
 
 
 @pytest.fixture()
 def skip_if_not_mms_version(mxnet_version):
-    if not is_mms_version(mxnet_version):
+    if not _is_mms_version(mxnet_version):
         pytest.skip("Skipping because this version does not use MMS")
 
 
-def _get_train_args(job_name, image_name):
+def _get_train_args(job_name):
     return {
-        "image": image_name,
+        "image": IMAGE,
         "input_mode": "File",
         "input_config": [
             {
@@ -321,7 +321,7 @@ def test_mxnet(
 
     actual_train_args = sagemaker_session.method_calls[0][2]
     job_name = actual_train_args["job_name"]
-    expected_train_args = _get_train_args(job_name, IMAGE)
+    expected_train_args = _get_train_args(job_name)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["experiment_config"] = EXPERIMENT_CONFIG
 
@@ -338,7 +338,7 @@ def test_mxnet(
     assert "cpu" in model.prepare_container_def(CPU)["Image"]
     predictor = mx.deploy(1, GPU)
     assert isinstance(predictor, MXNetPredictor)
-    assert is_mms_version(mxnet_version) ^ (create_tar_file.called and not repack_model.called)
+    assert _is_mms_version(mxnet_version) ^ (create_tar_file.called and not repack_model.called)
 
 
 @patch("sagemaker.utils.create_tar_file", MagicMock())
@@ -464,7 +464,7 @@ def test_model_image_accelerator(
     )
     container_def = model.prepare_container_def(INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
     assert container_def["Image"] == IMAGE
-    assert is_mms_version(mxnet_version) ^ (tar_and_upload.called and not repack_model.called)
+    assert _is_mms_version(mxnet_version) ^ (tar_and_upload.called and not repack_model.called)
 
 
 def test_attach(sagemaker_session, mxnet_version, mxnet_py_version):