edit unit tests for pt2.0.1 and pt2.1

Author: Sirut Buasai

src/sagemaker/fw_utils.py

@@ -138,6 +138,7 @@
         "1.12.1",
         "1.13.1",
         "2.0.0",
+        "2.0.1"
     ],
 }
 
@@ -153,10 +154,11 @@
     "1.13.1",
     "2.0.0",
     "2.0.1",
+    "2.1.0"
 ]
 
 
-TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = ["1.13.1", "2.0.0", "2.0.1"]
+TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = ["1.13.1", "2.0.0", "2.0.1", "2.1.0"]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
 TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS = [

tests/unit/test_fw_utils.py

@@ -937,6 +937,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("ml.p3.16xlarge", "pytorch", "1.12", "py38", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.13.1", "py39", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "2.0.0", "py310", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "2.0.1", "py310", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py3", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py37", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "tensorflow", "2.4.3", "py3", smdataparallel_enabled_custom_mpi),
@@ -959,6 +960,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("ml.p3.16xlarge", "pytorch", "1.12.1", "py38", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "pytorch", "1.13.1", "py39", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "pytorch", "2.0.0", "py310", smdataparallel_enabled_custom_mpi),
+        ("ml.p3.16xlarge", "pytorch", "2.0.1", "py310", smdataparallel_enabled_custom_mpi),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(
@@ -995,6 +997,10 @@ def test_validate_pytorchddp_not_raises():
         "1.12",
         "1.12.0",
         "1.12.1",
+        "1.13.1",
+        "2.0.0",
+        "2.0.1",
+        "2.1.0"
     ]
     for framework_version in pytorchddp_supported_fw_versions:
         fw_utils.validate_pytorch_distribution(
@@ -1060,6 +1066,8 @@ def test_validate_torch_distributed_not_raises():
     torch_distributed_gpu_supported_fw_versions = [
         "1.13.1",
         "2.0.0",
+        "2.0.1",
+        "2.1.0"
     ]
     for framework_version in torch_distributed_gpu_supported_fw_versions:
         fw_utils.validate_torch_distributed_distribution(

tests/unit/test_utils.py

@@ -384,6 +384,9 @@ def test_set_nested_value():
 
 
 def test_get_short_version():
+    assert sagemaker.utils.get_short_version("2.1.0") == "2.1"
+    assert sagemaker.utils.get_short_version("2.1") == "2.1"
+    assert sagemaker.utils.get_short_version("2.0.1") == "2.0"
     assert sagemaker.utils.get_short_version("2.0.0") == "2.0"
     assert sagemaker.utils.get_short_version("2.0") == "2.0"