breaking: addressing comments in PR feedback

Author: metrizable

src/sagemaker/fw_utils.py

@@ -686,17 +686,15 @@ 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.
     """
     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/pytorch/estimator.py

@@ -70,10 +70,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): PyTorch version you want to use for
-                executing your model training code. Defaults to None. List of supported versions.
+                executing your model training code. Defaults to ``None``. Required unless
+                ``image_name`` is provided. List of supported versions:
                 https://github.com/aws/sagemaker-python-sdk#pytorch-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/pytorch/model.py

@@ -88,14 +88,16 @@ 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): PyTorch version you want to use for
-                executing your model training code. Defaults to None.
+                executing your model training code. Defaults to None. Required
+                unless ``image`` 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`` is provided.
             image (str): A Docker image URI (default: None). If not specified, a
                 default image for PyTorch will be used.
 
                 If ``framework_version`` or ``py_version`` are ``None``, then
-                ``image_name`` is required. If also ``None``, then a ``ValueError``
+                ``image`` is required. If also ``None``, then a ``ValueError``
                 will be raised.
             predictor_cls (callable[str, sagemaker.session.Session]): A function
                 to call to create a predictor with an endpoint name and

tests/integ/test_pytorch_train.py

@@ -98,10 +98,6 @@ def test_fit_deploy(sagemaker_local_session, pytorch_full_version):
         predictor.delete_endpoint()
 
 
-@pytest.mark.skipif(
-    PYTHON_VERSION == "py2",
-    reason="Python 2 is supported by PyTorch {} and lower versions.".format(LATEST_PY2_VERSION),
-)
 def test_deploy_model(
     pytorch_training_job, sagemaker_session, cpu_instance_type, pytorch_full_version
 ):
@@ -129,10 +125,6 @@ def test_deploy_model(
         assert output.shape == (batch_size, 10)
 
 
-@pytest.mark.skipif(
-    PYTHON_VERSION == "py2",
-    reason="Python 2 is supported by PyTorch {} and lower versions.".format(LATEST_PY2_VERSION),
-)
 def test_deploy_packed_model_with_entry_point_name(sagemaker_session, cpu_instance_type):
     endpoint_name = "test-pytorch-deploy-model-{}".format(sagemaker_timestamp())
 

tests/integ/test_source_dirs.py

@@ -20,11 +20,10 @@
 from tests.integ import DATA_DIR, PYTHON_VERSION
 
 from sagemaker.pytorch.estimator import PyTorch
-from sagemaker.pytorch.defaults import PYTORCH_VERSION
 
 
 @pytest.mark.local_mode
-def test_source_dirs(tmpdir, sagemaker_local_session):
+def test_source_dirs(tmpdir, sagemaker_local_session, pytorch_full_version):
     source_dir = os.path.join(DATA_DIR, "pytorch_source_dirs")
     lib = os.path.join(str(tmpdir), "alexa.py")
 
@@ -36,7 +35,7 @@ def test_source_dirs(tmpdir, sagemaker_local_session):
         role="SageMakerRole",
         source_dir=source_dir,
         dependencies=[lib],
-        framework_version=PYTORCH_VERSION,
+        framework_version=pytorch_full_version,
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",

tests/unit/test_pytorch.py

@@ -80,7 +80,6 @@ def fixture_sagemaker_session():
     return session
 
 
-# TODO: push assertions regarding uri generation particulars to create_image_uri tests
 def _get_full_cpu_image_uri(version, py_version=PYTHON_VERSION):
     return IMAGE_URI_FORMAT_STRING.format(REGION, IMAGE_NAME, version, "cpu", py_version)