fix: use py3 when using pytorch_full_version

Author: metrizable

tests/integ/test_git.py

@@ -21,6 +21,7 @@
 
 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
@@ -51,15 +52,18 @@
 
 
 @pytest.mark.local_mode
-def test_github(sagemaker_local_session, pytorch_full_version):
+def test_github(sagemaker_local_session):
     script_path = "mnist.py"
     data_path = os.path.join(DATA_DIR, "pytorch_mnist")
     git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+
+    # TODO: fails for newer pytorch versions when using MNIST from torchvision due to missing dataset
+    # "algo-1-v767u_1  | RuntimeError: Dataset not found. You can use download=True to download it"
     pytorch = PyTorch(
         entry_point=script_path,
         role="SageMakerRole",
         source_dir="pytorch",
-        framework_version=pytorch_full_version,
+        framework_version=PYTORCH_VERSION,
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",

tests/integ/test_pytorch_train.py

@@ -191,7 +191,7 @@ def _get_pytorch_estimator(
         entry_point=entry_point,
         role="SageMakerRole",
         framework_version=pytorch_full_version,
-        py_version=PYTHON_VERSION,
+        py_version="py3",
         train_instance_count=1,
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,

tests/integ/test_source_dirs.py

@@ -19,23 +19,26 @@
 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
 
 
 @pytest.mark.local_mode
-def test_source_dirs(tmpdir, sagemaker_local_session, pytorch_full_version):
+def test_source_dirs(tmpdir, sagemaker_local_session):
     source_dir = os.path.join(DATA_DIR, "pytorch_source_dirs")
     lib = os.path.join(str(tmpdir), "alexa.py")
 
     with open(lib, "w") as f:
         f.write("def question(to_anything): return 42")
 
+    # TODO: fails on newer versions of pytorch in call to np.load(BytesIO(stream.read()))
+    # "ValueError: Cannot load file containing pickled data when allow_pickle=False"
     estimator = PyTorch(
         entry_point="train.py",
         role="SageMakerRole",
         source_dir=source_dir,
         dependencies=[lib],
-        framework_version=pytorch_full_version,
+        framework_version=PYTORCH_VERSION,
         py_version=PYTHON_VERSION,
         train_instance_count=1,
         train_instance_type="local",

tests/integ/test_transformer.py

@@ -178,7 +178,7 @@ def test_transform_pytorch_vpc_custom_model_bucket(
         entry_point=os.path.join(data_dir, "mnist.py"),
         role="SageMakerRole",
         framework_version=pytorch_full_version,
-        py_version=PYTHON_VERSION,
+        py_version="py3",
         sagemaker_session=sagemaker_session,
         vpc_config={"Subnets": subnet_ids, "SecurityGroupIds": [security_group_id]},
         code_location="s3://{}".format(custom_bucket_name),

tests/integ/test_tuner.py

@@ -828,7 +828,7 @@ def test_attach_tuning_pytorch(sagemaker_session, cpu_instance_type, pytorch_ful
         role="SageMakerRole",
         train_instance_count=1,
         framework_version=pytorch_full_version,
-        py_version=PYTHON_VERSION,
+        py_version="py3",
         train_instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
     )