breaking: small updates based on PR feedback

Author: metrizable

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
@@ -52,15 +51,15 @@
 
 
 @pytest.mark.local_mode
-def test_github(sagemaker_local_session):
+def test_github(sagemaker_local_session, pytorch_full_version):
     script_path = "mnist.py"
     data_path = os.path.join(DATA_DIR, "pytorch_mnist")
     git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
     pytorch = PyTorch(
         entry_point=script_path,
         role="SageMakerRole",
         source_dir="pytorch",
-        framework_version=PYTORCH_VERSION,
+        framework_version=pytorch_full_version,
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",

tests/integ/test_tuner.py

@@ -36,7 +36,6 @@
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.predictor import json_deserializer
 from sagemaker.pytorch import PyTorch
-from sagemaker.pytorch.defaults import PYTORCH_VERSION
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.tensorflow.defaults import LATEST_VERSION
 from sagemaker.tuner import (
@@ -820,15 +819,15 @@ def test_tuning_chainer(sagemaker_session, cpu_instance_type):
     reason="This test has always failed, but the failure was masked by a bug. "
     "This test should be fixed. Details in https://github.com/aws/sagemaker-python-sdk/pull/968"
 )
-def test_attach_tuning_pytorch(sagemaker_session, cpu_instance_type):
+def test_attach_tuning_pytorch(sagemaker_session, cpu_instance_type, pytorch_full_version):
     mnist_dir = os.path.join(DATA_DIR, "pytorch_mnist")
     mnist_script = os.path.join(mnist_dir, "mnist.py")
 
     estimator = PyTorch(
         entry_point=mnist_script,
         role="SageMakerRole",
         train_instance_count=1,
-        framework_version=PYTORCH_VERSION,
+        framework_version=pytorch_full_version,
         py_version=PYTHON_VERSION,
         train_instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,

tests/unit/test_pytorch.py

@@ -93,11 +93,7 @@ def _get_full_cpu_image_uri_with_ei(version, py_version=PYTHON_VERSION):
 
 
 def _pytorch_estimator(
-    sagemaker_session,
-    framework_version=defaults.PYTORCH_VERSION,
-    train_instance_type=None,
-    base_job_name=None,
-    **kwargs
+    sagemaker_session, framework_version, train_instance_type=None, base_job_name=None, **kwargs
 ):
     return PyTorch(
         entry_point=SCRIPT_PATH,
@@ -572,13 +568,17 @@ def test_model_py2_warning(warning, sagemaker_session, pytorch_version):
     warning.assert_called_with(model.__framework_name__, defaults.LATEST_PY2_VERSION)
 
 
-def test_pt_enable_sm_metrics(sagemaker_session):
-    pytorch = _pytorch_estimator(sagemaker_session, enable_sagemaker_metrics=True)
+def test_pt_enable_sm_metrics(sagemaker_session, pytorch_full_version):
+    pytorch = _pytorch_estimator(
+        sagemaker_session, framework_version=pytorch_full_version, enable_sagemaker_metrics=True
+    )
     assert pytorch.enable_sagemaker_metrics
 
 
-def test_pt_disable_sm_metrics(sagemaker_session):
-    pytorch = _pytorch_estimator(sagemaker_session, enable_sagemaker_metrics=False)
+def test_pt_disable_sm_metrics(sagemaker_session, pytorch_full_version):
+    pytorch = _pytorch_estimator(
+        sagemaker_session, framework_version=pytorch_full_version, enable_sagemaker_metrics=False
+    )
     assert not pytorch.enable_sagemaker_metrics
 
 
@@ -594,9 +594,9 @@ def test_pt_enable_sm_metrics_if_fw_ver_is_at_least_1_15(sagemaker_session):
         assert pytorch.enable_sagemaker_metrics
 
 
-def test_custom_image_estimator_deploy(sagemaker_session):
+def test_custom_image_estimator_deploy(sagemaker_session, pytorch_full_version):
     custom_image = "mycustomimage:latest"
-    pytorch = _pytorch_estimator(sagemaker_session)
+    pytorch = _pytorch_estimator(sagemaker_session, framework_version=pytorch_full_version)
     pytorch.fit(inputs="s3://mybucket/train", job_name="new_name")
     model = pytorch.create_model(image=custom_image)
     assert model.image == custom_image