breaking: updates based on PR feedback

Author: metrizable

src/sagemaker/pytorch/model.py

@@ -94,11 +94,9 @@ def __init__(
                 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`` is required. If also ``None``, then a ``ValueError``
-                will be raised.
+                default image for PyTorch will be used. If ``framework_version``
+                or ``py_version`` are ``None``, then ``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
                 SageMaker ``Session``. If specified, ``deploy()`` returns the

tests/integ/test_airflow_config.py

@@ -608,13 +608,13 @@ def test_xgboost_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu
 
 @pytest.mark.canary_quick
 def test_pytorch_airflow_config_uploads_data_source_to_s3_when_inputs_not_provided(
-    sagemaker_session, cpu_instance_type
+    sagemaker_session, cpu_instance_type, pytorch_full_version
 ):
     with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
         estimator = PyTorch(
             entry_point=PYTORCH_MNIST_SCRIPT,
             role=ROLE,
-            framework_version="1.3.1",
+            framework_version=pytorch_full_version,
             py_version="py3",
             train_instance_count=2,
             train_instance_type=cpu_instance_type,

tests/integ/test_git.py

@@ -21,7 +21,6 @@
 
 from tests.integ import lock as lock
 from sagemaker.mxnet.estimator import MXNet
-from sagemaker.pytorch.defaults import PYTORCH_VERSION
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.sklearn.estimator import SKLearn
 from sagemaker.sklearn.model import SKLearnModel
@@ -63,7 +62,7 @@ def test_github(sagemaker_local_session):
         entry_point=script_path,
         role="SageMakerRole",
         source_dir="pytorch",
-        framework_version=PYTORCH_VERSION,
+        framework_version="0.4",  # hard-code to last known good pytorch for now (see TODO above)
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",

tests/integ/test_source_dirs.py

@@ -19,7 +19,6 @@
 import tests.integ.lock as lock
 from tests.integ import DATA_DIR, PYTHON_VERSION
 
-from sagemaker.pytorch.defaults import PYTORCH_VERSION
 from sagemaker.pytorch.estimator import PyTorch
 
 
@@ -38,7 +37,7 @@ def test_source_dirs(tmpdir, sagemaker_local_session):
         role="SageMakerRole",
         source_dir=source_dir,
         dependencies=[lib],
-        framework_version=PYTORCH_VERSION,
+        framework_version="0.4",  # hard-code to last known good pytorch for now (see TODO above)
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",