fix: run Python script entry point as script and install from requirements.txt (#64)

Author: ajaykarpur

src/sagemaker_training/entry_point.py

@@ -117,13 +117,10 @@ def install(name, path=environment.code_dir, capture_error=False):
 
     entry_point_type = _entry_point_type.get(path, name)
 
-    if (
-        entry_point_type is _entry_point_type.PYTHON_PACKAGE
-        or entry_point_type is _entry_point_type.PYTHON_PROGRAM
-        or modules.has_requirements(path)
-    ):
-        modules.prepare(path, name)
+    if entry_point_type is _entry_point_type.PYTHON_PACKAGE:
         modules.install(path, capture_error)
+    elif entry_point_type is _entry_point_type.PYTHON_PROGRAM and modules.has_requirements(path):
+        modules.install_requirements(path, capture_error)
 
     if entry_point_type is _entry_point_type.COMMAND:
         os.chmod(os.path.join(path, name), 511)

src/sagemaker_training/errors.py

@@ -52,6 +52,10 @@ class InstallModuleError(_CalledProcessError):
     """Error class indicating a module failed to install."""
 
 
+class InstallRequirementsError(_CalledProcessError):
+    """Error class indicating a module failed to install."""
+
+
 class ImportModuleError(ClientError):
     """Error class indicating a module failed to import."""
 

src/sagemaker_training/modules.py

@@ -129,6 +129,23 @@ def install(path, capture_error=False):  # type: (str, bool) -> None
     )
 
 
+def install_requirements(path, capture_error=False):  # type: (str, bool) -> None
+    """Install dependencies from requirements.txt in the executing Python environment.
+
+    Args:
+        path (str):  Real path location of the requirements.txt file.
+        capture_error (bool): Default false. If True, the running process captures the
+            stderr, and appends it to the returned Exception message in case of errors.
+    """
+    cmd = "{} -m pip install -r requirements.txt".format(process.python_executable())
+
+    logger.info("Installing dependencies from requirements.txt:\n{}".format(cmd))
+
+    process.check_error(
+        shlex.split(cmd), errors.InstallRequirementsError, cwd=path, capture_error=capture_error
+    )
+
+
 def import_module(uri, name=DEFAULT_MODULE_NAME):  # type: (str, str) -> module
     """Download, prepare and install a compressed tar file from S3 or provided directory as a
     module.

test/integration/local/test_dummy.py

@@ -47,6 +47,6 @@ def test_install_requirements(capsys):
 
     stdout = capsys.readouterr().out
 
-    assert "Installing collected packages: pyfiglet, train.py" in stdout
-    assert "Successfully installed pyfiglet-0.8.post1 train.py-1.0.0" in stdout
+    assert "Installing collected packages: pyfiglet" in stdout
+    assert "Successfully installed pyfiglet-0.8.post1" in stdout
     assert "Reporting training SUCCESS" in stdout

test/unit/test_modules.py

@@ -76,6 +76,20 @@ def test_install(check_error):
         )
 
 
+@patch("sagemaker_training.process.check_error", autospec=True)
+def test_install_requirements(check_error):
+    path = "c://sagemaker-pytorch-container"
+
+    cmd = [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"]
+
+    with patch("os.path.exists", return_value=True):
+        modules.install_requirements(path)
+
+        check_error.assert_called_with(
+            cmd, errors.InstallRequirementsError, cwd=path, capture_error=False
+        )
+
+
 @patch("sagemaker_training.process.check_error", autospec=True)
 def test_install_fails(check_error):
     check_error.side_effect = errors.ClientError()