feature: python module support to torch_distributed

Author: azzhipa

src/sagemaker/fw_utils.py

@@ -200,11 +200,12 @@ def validate_source_dir(script, directory):
             not contain ``script``.
     """
     if directory:
-        if not os.path.isfile(os.path.join(directory, script)):
+        if script.endswith(".py") and not os.path.isfile(os.path.join(directory, script)):
             raise ValueError(
                 'No file named "{}" was found in directory "{}".'.format(script, directory)
             )
 
+
     return True
 
 
@@ -1143,10 +1144,10 @@ def validate_torch_distributed_distribution(
             )
 
     # Check entry point type
-    if not entry_point.endswith(".py"):
+    if not entry_point.startswith("-m") and not entry_point.endswith(".py"):
         err_msg += (
             "Unsupported entry point type for the distribution torch_distributed.\n"
-            "Only python programs (*.py) are supported."
+            "Only python programs (*.py) or modules (-m *) are supported."
         )
 
     if err_msg:

tests/unit/test_fw_utils.py

@@ -186,7 +186,7 @@ def test_validate_source_dir_is_not_directory(sagemaker_session):
 
 
 def test_validate_source_dir_file_not_in_dir():
-    script = " !@#$%^&*() .myscript. !@#$%^&*() "
+    script = " !@#$%^&*() .myscript. !@#$%^&*().py"
     directory = "."
     with pytest.raises(ValueError):
         fw_utils.validate_source_dir(script, directory)