fix: add short version to smdataparallel supported list (#2239)

* fix: add short version to smdataparallel supported list

* sort version string

Co-authored-by: Rui Wang Napieralski <wru@amazon.com>

Author: icywang86rui

src/sagemaker/fw_utils.py

@@ -59,8 +59,8 @@
     "local_gpu",
 )
 SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {
-    "tensorflow": ["2.3.1", "2.3.2", "2.4.1"],
-    "pytorch": ["1.6.0", "1.7.1", "1.8.0"],
+    "tensorflow": ["2.3", "2.3.1", "2.3.2", "2.4", "2.4.1"],
+    "pytorch": ["1.6", "1.6.0", "1.7", "1.7.1", "1.8", "1.8.0"],
 }
 SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
 

tests/unit/test_fw_utils.py

@@ -632,10 +632,15 @@ def test_validate_smdataparallel_args_not_raises():
         (None, None, None, None, smdataparallel_disabled),
         ("ml.p3.16xlarge", "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.3.2", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "tensorflow", "2.3", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "tensorflow", "2.4", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.6.0", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.6", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.7.1", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.7", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.8.0", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.8", "py3", smdataparallel_enabled),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(