update data parallel to support pytorch 1.7.1

ChoiByungWook · ChoiByungWook · commit caf6580cf882 · 2021-03-03T12:29:38.000-08:00
diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
@@ -60,7 +60,7 @@
 )
 SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {
     "tensorflow": ["2.3.0", "2.3.1"],
-    "pytorch": ["1.6.0"],
+    "pytorch": ["1.6.0", "1.7.1"],
 }
 SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
 
diff --git a/tests/unit/test_fw_utils.py b/tests/unit/test_fw_utils.py
@@ -632,6 +632,7 @@ 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", "pytorch", "1.6.0", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.7.1", "py3", smdataparallel_enabled),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(

Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@`
`60`	`60`	`)`
`61`	`61`	`SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {`
`62`	`62`	`"tensorflow": ["2.3.0", "2.3.1"],`
`63`		`- "pytorch": ["1.6.0"],`
	`63`	`+ "pytorch": ["1.6.0", "1.7.1"],`
`64`	`64`	`}`
`65`	`65`	`SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]`
`66`	`66`
Original file line number	Diff line number	Diff line change
`@@ -632,6 +632,7 @@ def test_validate_smdataparallel_args_not_raises():`
`632`	`632`	`(None, None, None, None, smdataparallel_disabled),`
`633`	`633`	`("ml.p3.16xlarge", "tensorflow", "2.3.1", "py3", smdataparallel_enabled),`
`634`	`634`	`("ml.p3.16xlarge", "pytorch", "1.6.0", "py3", smdataparallel_enabled),`
	`635`	`+ ("ml.p3.16xlarge", "pytorch", "1.7.1", "py3", smdataparallel_enabled),`
`635`	`636`	`]`
`636`	`637`	`for instance_type, framework_name, framework_version, py_version, distribution in good_args:`
`637`	`638`	`fw_utils._validate_smdataparallel_args(`