pathways: refactor job_driver entrypoint (aws#825)

Co-authored-by: Rohan Gujarathi <[email protected]>

Author: 2 people

setup.py

@@ -103,10 +103,12 @@ def read_requirements(filename):
     ],
     install_requires=required_packages,
     extras_require=extras,
+    scripts=["src/sagemaker/remote_function/job_driver.sh"],
     entry_points={
         "console_scripts": [
             "sagemaker-upgrade-v2=sagemaker.cli.compatibility.v2.sagemaker_upgrade_v2:main",
-            "invoke-remote-function=sagemaker.remote_function.job_driver:main",
+            "invoke-remote-function=sagemaker.remote_function.invoke_function:main",
+            "bootstrap-runtime-env=sagemaker.remote_function.bootstrap:main",
         ]
     },
 )

src/sagemaker/remote_function/job_driver.py renamed to src/sagemaker/remote_function/bootstrap.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""An entry point for invoking remote function inside a job."""
+"""An entry point for runtime environment."""
 from __future__ import absolute_import
 
 import argparse
@@ -20,9 +20,8 @@
 
 from sagemaker.session import Session
 from sagemaker.remote_function.errors import handle_error
-from sagemaker.remote_function.core.runtime_environment import RuntimeEnvironmentManager
 from sagemaker.remote_function import logging_config
-
+from sagemaker.remote_function.core.runtime_environment import RuntimeEnvironmentManager
 
 SUCCESS_EXIT_CODE = 0
 DEFAULT_FAILURE_CODE = 1
@@ -40,12 +39,6 @@ def _parse_agrs():
     return parser.parse_args()
 
 
-def _get_sagemaker_session(region):
-    """Get sagemaker session for interacting with AWS or Sagemaker services"""
-    boto_session = boto3.session.Session(region_name=region)
-    return Session(boto_session=boto_session)
-
-
 def _execute_pre_exec_cmds():
     """Execute pre-flight commands before invkoing remote function"""
     # TODO: complete me
@@ -58,12 +51,19 @@ def _uncompress_src_dir():
     pass  # pylint: disable=W0107
 
 
-def _bootstrap_runtime_environment(
+def _get_sagemaker_session(region):
+    """Get sagemaker session for interacting with AWS or Sagemaker services"""
+    boto_session = boto3.session.Session(region_name=region)
+    return Session(boto_session=boto_session)
+
+
+def bootstrap_runtime_environment(
     runtime_manager: RuntimeEnvironmentManager,
     dependencies: str,
     job_conda_env: str = None,
 ):
     """Bootstrap runtime environment for remote function invocation"""
+
     _execute_pre_exec_cmds()
 
     if dependencies:
@@ -72,21 +72,14 @@ def _bootstrap_runtime_environment(
     _uncompress_src_dir()
 
 
-def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key):
-    """Execute stored remote function"""
-    from sagemaker.remote_function.core.stored_function import StoredFunction
-
-    stored_function = StoredFunction(sagemaker_session, s3_base_uri, s3_kms_key)
-    stored_function.load_and_invoke()
-
-
 def main():
-    """Entry point for job driver script"""
+    """Entry point for bootstrap script"""
 
     logging_config.basic_config()
     logger = logging_config.get_logger()
 
     exit_code = DEFAULT_FAILURE_CODE
+
     try:
         args = _parse_agrs()
         region = args.region
@@ -103,16 +96,15 @@ def main():
             sagemaker_session=sagemaker_session,
         )
 
-        _bootstrap_runtime_environment(
+        bootstrap_runtime_environment(
             runtime_manager=runtime_environment_manager,
             dependencies=dependencies,
             job_conda_env=job_conda_env,
         )
 
-        _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key)
         exit_code = SUCCESS_EXIT_CODE
     except Exception as e:  # pylint: disable=broad-except
-        logger.exception("Error encountered when invoking the remote function.")
+        logger.exception("Error encountered while bootstrapping runtime environment.")
         exit_code = handle_error(e, sagemaker_session, s3_base_uri, s3_kms_key)
     finally:
         sys.exit(exit_code)

src/sagemaker/remote_function/invoke_function.py

@@ -0,0 +1,81 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""An entry point for invoking remote function inside a job."""
+
+from __future__ import absolute_import
+
+import argparse
+import sys
+
+import boto3
+
+from sagemaker.session import Session
+from sagemaker.remote_function.errors import handle_error
+from sagemaker.remote_function import logging_config
+
+
+SUCCESS_EXIT_CODE = 0
+
+
+def _parse_agrs():
+    """Parses CLI arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--region", type=str, required=True)
+    parser.add_argument("--s3_base_uri", type=str, required=True)
+    parser.add_argument("--s3_kms_key", type=str)
+    parser.add_argument("--dependencies", type=str)
+    parser.add_argument("--job_conda_env", type=str)
+
+    return parser.parse_args()
+
+
+def _get_sagemaker_session(region):
+    """Get sagemaker session for interacting with AWS or Sagemaker services"""
+    boto_session = boto3.session.Session(region_name=region)
+    return Session(boto_session=boto_session)
+
+
+def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key):
+    """Execute stored remote function"""
+    from sagemaker.remote_function.core.stored_function import StoredFunction
+
+    stored_function = StoredFunction(sagemaker_session, s3_base_uri, s3_kms_key)
+    stored_function.load_and_invoke()
+
+
+def main():
+    """Entry point for invoke function script"""
+
+    logging_config.basic_config()
+    logger = logging_config.get_logger()
+
+    exit_code = SUCCESS_EXIT_CODE
+
+    try:
+        args = _parse_agrs()
+        region = args.region
+        s3_base_uri = args.s3_base_uri
+        s3_kms_key = args.s3_kms_key
+
+        sagemaker_session = _get_sagemaker_session(region)
+        _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key)
+
+    except Exception as e:  # pylint: disable=broad-except
+        logger.exception("Error encountered while invoking the remote function.")
+        exit_code = handle_error(e, sagemaker_session, s3_base_uri, s3_kms_key)
+    finally:
+        sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

src/sagemaker/remote_function/job.py

@@ -26,7 +26,7 @@
 from sagemaker.remote_function import logging_config
 
 
-JOBS_CONTAINER_ENTRYPOINT = ["invoke-remote-function"]
+JOBS_CONTAINER_ENTRYPOINT = ["/bin/bash", "job_driver.sh"]
 
 
 logger = logging_config.get_logger()

src/sagemaker/remote_function/job_driver.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function
+
+set -eu
+
+printf "INFO: Bootstraping runtime environment\n"
+bootstrap-runtime-env "$@"
+
+printf "INFO: Invoking remote function\n"
+invoke-remote-function "$@"

tests/unit/sagemaker/remote_function/test_bootstrap.py

@@ -0,0 +1,72 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from mock import patch, Mock
+from sagemaker.remote_function import bootstrap
+from sagemaker.remote_function.errors import RuntimeEnvironmentError
+
+TEST_REGION = "us-west-2"
+TEST_S3_BASE_URI = "s3://my-bucket/"
+TEST_S3_KMS_KEY = "my-kms-key"
+TEST_DEPENDENCIES = "s3://my-bucket/requirements.txt"
+
+
+def mock_args():
+    args = Mock()
+    args.region = TEST_REGION
+    args.s3_base_uri = TEST_S3_BASE_URI
+    args.s3_kms_key = TEST_S3_KMS_KEY
+    return args
+
+
+def mock_session():
+    return Mock()
+
+
+@patch("sagemaker.remote_function.bootstrap._parse_agrs", new=mock_args)
+@patch("sys.exit")
+@patch("sagemaker.remote_function.core.runtime_environment.RuntimeEnvironmentManager.bootstrap")
+@patch(
+    "sagemaker.remote_function.bootstrap._get_sagemaker_session",
+    return_value=mock_session(),
+)
+def test_main_success(_get_sagemaker_session, bootstrap_runtime, _exit_process):
+    bootstrap.main()
+
+    _get_sagemaker_session.assert_called_with(TEST_REGION)
+    bootstrap_runtime.assert_called()
+    _exit_process.assert_called_with(0)
+
+
+@patch("sagemaker.remote_function.bootstrap._parse_agrs", new=mock_args)
+@patch("sagemaker.remote_function.bootstrap.handle_error")
+@patch("sys.exit")
+@patch("sagemaker.remote_function.core.runtime_environment.RuntimeEnvironmentManager.bootstrap")
+@patch(
+    "sagemaker.remote_function.bootstrap._get_sagemaker_session",
+    return_value=mock_session(),
+)
+def test_main_failure(_get_sagemaker_session, bootstrap_runtime, _exit_process, handle_error):
+    runtime_err = RuntimeEnvironmentError("some failure reason")
+    bootstrap_runtime.side_effect = runtime_err
+    handle_error.return_value = 1
+
+    bootstrap.main()
+
+    _get_sagemaker_session.assert_called_with(TEST_REGION)
+    bootstrap_runtime.assert_called()
+    handle_error.assert_called_with(
+        runtime_err, _get_sagemaker_session(), TEST_S3_BASE_URI, TEST_S3_KMS_KEY
+    )
+    _exit_process.assert_called_with(1)

tests/unit/sagemaker/remote_function/test_job_driver.py renamed to tests/unit/sagemaker/remote_function/test_invoke_function.py

@@ -13,7 +13,7 @@
 from __future__ import absolute_import
 
 from mock import patch, Mock
-from sagemaker.remote_function import job_driver
+from sagemaker.remote_function import invoke_function
 from sagemaker.remote_function.errors import SerializationError
 
 TEST_REGION = "us-west-2"
@@ -33,35 +33,35 @@ def mock_session():
     return Mock()
 
 
-@patch("sagemaker.remote_function.job_driver._parse_agrs", new=mock_args)
+@patch("sagemaker.remote_function.invoke_function._parse_agrs", new=mock_args)
 @patch("sys.exit")
 @patch("sagemaker.remote_function.core.stored_function.StoredFunction.load_and_invoke")
 @patch(
-    "sagemaker.remote_function.job_driver._get_sagemaker_session",
+    "sagemaker.remote_function.invoke_function._get_sagemaker_session",
     return_value=mock_session(),
 )
 def test_main_success(_get_sagemaker_session, load_and_invoke, _exit_process):
-    job_driver.main()
+    invoke_function.main()
 
     _get_sagemaker_session.assert_called_with(TEST_REGION)
     load_and_invoke.assert_called()
     _exit_process.assert_called_with(0)
 
 
-@patch("sagemaker.remote_function.job_driver._parse_agrs", new=mock_args)
-@patch("sagemaker.remote_function.job_driver.handle_error")
+@patch("sagemaker.remote_function.invoke_function._parse_agrs", new=mock_args)
+@patch("sagemaker.remote_function.invoke_function.handle_error")
 @patch("sys.exit")
 @patch("sagemaker.remote_function.core.stored_function.StoredFunction.load_and_invoke")
 @patch(
-    "sagemaker.remote_function.job_driver._get_sagemaker_session",
+    "sagemaker.remote_function.invoke_function._get_sagemaker_session",
     return_value=mock_session(),
 )
 def test_main_failure(_get_sagemaker_session, load_and_invoke, _exit_process, handle_error):
     ser_err = SerializationError("some failure reason")
     load_and_invoke.side_effect = ser_err
     handle_error.return_value = 1
 
-    job_driver.main()
+    invoke_function.main()
 
     _get_sagemaker_session.assert_called_with(TEST_REGION)
     load_and_invoke.assert_called()

tests/unit/sagemaker/remote_function/test_job.py

@@ -128,7 +128,7 @@ def test_start(session, mock_stored_function, mock_runtime_manager):
         AlgorithmSpecification=dict(
             TrainingImage=IMAGE,
             TrainingInputMode="File",
-            ContainerEntrypoint=["invoke-remote-function"],
+            ContainerEntrypoint=["/bin/bash", "job_driver.sh"],
             ContainerArguments=[
                 "--s3_base_uri",
                 f"{S3_URI}/{job.job_name}",
@@ -190,7 +190,7 @@ def test_start_with_complete_job_settings(session, mock_stored_function, mock_ru
         AlgorithmSpecification=dict(
             TrainingImage=IMAGE,
             TrainingInputMode="File",
-            ContainerEntrypoint=["invoke-remote-function"],
+            ContainerEntrypoint=["/bin/bash", "job_driver.sh"],
             ContainerArguments=[
                 "--s3_base_uri",
                 f"{S3_URI}/{job.job_name}",