pathways: standalone job entrypoint scripts

Author: Rohan Gujarathi

setup.py

@@ -103,7 +103,6 @@ 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",

src/sagemaker/remote_function/bootstrap.py

@@ -38,11 +38,6 @@ class ServiceError(RemoteFunctionError):
     """Raised when errors encountered during interaction with SageMaker, S3 service APIs"""
 
 
-@pickling_support.install
-class RuntimeEnvironmentError(RemoteFunctionError):
-    """Raised when errors encountered during remote function runtime environment setup"""
-
-
 @pickling_support.install
 class SerializationError(RemoteFunctionError):
     """Raised when errors encountered during serialization of remote function objects"""

src/sagemaker/remote_function/errors.py

@@ -15,21 +15,68 @@
 
 import os
 import re
+import shutil
 from typing import Dict, List, Tuple
 
 from sagemaker.config import config_schema
 from sagemaker.config.config_factory import SageMakerConfigFactory
 from sagemaker.session import get_execution_role, _logs_for_job, Session
-from sagemaker.utils import name_from_base
-from sagemaker.s3 import s3_path_join
+from sagemaker.utils import name_from_base, _tmpdir
+from sagemaker.s3 import s3_path_join, S3Uploader
 from sagemaker import vpc_utils
 from sagemaker.remote_function.core.stored_function import StoredFunction
-from sagemaker.remote_function.core.runtime_environment import RuntimeEnvironmentManager
+from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
+    RuntimeEnvironmentManager,
+)
 from sagemaker.remote_function import logging_config
 
 
-JOBS_CONTAINER_ENTRYPOINT = ["/bin/bash", "job_driver.sh"]
+# runtime script names
+BOOTSTRAP_SCRIPT_NAME = "bootstrap_runtime_environment.py"
+ENTRYPOINT_SCRIPT_NAME = "job_driver.sh"
+RUNTIME_MANAGER_SCRIPT_NAME = "runtime_environment_manager.py"
 
+# training channel names
+RUNTIME_SCRIPTS_CHANNEL_NAME = "remote_function_scripts"
+USER_DEPENDENCIES_CHANNEL_NAME = "remote_function_dependencies"
+SAGEMAKER_WHL_CHANNEL_NAME = "sagemaker_whl_file"
+
+SAGEMAKER_SDK_WHL_FILE = (
+    "s3://sagemaker-pathways/test/pysdk/sagemaker-2.120.1.dev0-py2.py3-none-any.whl"
+)
+
+JOBS_CONTAINER_ENTRYPOINT = [
+    "/bin/bash",
+    f"/opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{ENTRYPOINT_SCRIPT_NAME}",
+]
+
+ENTRYPOINT_SCRIPT = r"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function
+
+set -eu
+
+printf "INFO: Bootstraping runtime environment.\n"
+python /opt/ml/input/data/remote_function_scripts/bootstrap_runtime_environment.py "$@"
+
+if [ -f "remote_function_conda_env.txt" ]
+then
+    conda_env=$(cat remote_function_conda_env.txt)
+
+    if which mamba >/dev/null; then
+        conda_exe="mamba"
+    else
+        conda_exe="conda"
+    fi
+
+    printf "INFO: Invoking remote function inside conda environment: $conda_env.\n"
+    $conda_exe run -n $conda_env python -m sagemaker.remote_function.invoke_function "$@"
+else
+    printf "INFO: No conda env provided. Invoking remote function\n"
+    python -m sagemaker.remote_function.invoke_function "$@"
+fi
+"""
 
 logger = logging_config.get_logger()
 
@@ -173,12 +220,20 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
 
         s3_base_uri = s3_path_join(job_settings.s3_root_uri, job_name)
 
-        runtime_environment_manager = RuntimeEnvironmentManager(
+        local_dependencies_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)
+
+        remote_function_scripts_s3uri = _prepare_and_upload_runtime_scripts(
+            s3_base_uri=s3_base_uri,
+            s3_kms_key=job_settings.s3_kms_key,
+            sagemaker_session=job_settings.sagemaker_session,
+        )
+
+        user_dependencies_s3uri = _prepare_and_upload_dependencies(
+            local_dependencies_path=local_dependencies_path,
             s3_base_uri=s3_base_uri,
             s3_kms_key=job_settings.s3_kms_key,
             sagemaker_session=job_settings.sagemaker_session,
         )
-        uploaded_dependencies_path = runtime_environment_manager.snapshot(job_settings.dependencies)
 
         stored_function = StoredFunction(
             sagemaker_session=job_settings.sagemaker_session,
@@ -198,6 +253,46 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
             RetryStrategy={"MaximumRetryAttempts": job_settings.max_retry_attempts},
         )
 
+        input_data_config = [
+            dict(
+                ChannelName=RUNTIME_SCRIPTS_CHANNEL_NAME,
+                DataSource={
+                    "S3DataSource": {
+                        "S3Uri": remote_function_scripts_s3uri,
+                        "S3DataType": "S3Prefix",
+                    }
+                },
+            )
+        ]
+
+        if user_dependencies_s3uri:
+            input_data_config.append(
+                dict(
+                    ChannelName=USER_DEPENDENCIES_CHANNEL_NAME,
+                    DataSource={
+                        "S3DataSource": {
+                            "S3Uri": user_dependencies_s3uri,
+                            "S3DataType": "S3Prefix",
+                        }
+                    },
+                )
+            )
+
+        # temporary solution for public beta to make sagemaker installer available
+        # in the images, this should be removed before pathways GA.
+        input_data_config.append(
+            dict(
+                ChannelName=SAGEMAKER_WHL_CHANNEL_NAME,
+                DataSource={
+                    "S3DataSource": {
+                        "S3Uri": SAGEMAKER_SDK_WHL_FILE,
+                        "S3DataType": "S3Prefix",
+                    }
+                },
+            )
+        )
+        request_dict["InputDataConfig"] = input_data_config
+
         output_config = {"S3OutputPath": s3_base_uri}
         if job_settings.s3_kms_key is not None:
             output_config["KmsKeyId"] = job_settings.s3_kms_key
@@ -207,8 +302,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
         container_args.extend(["--region", job_settings.sagemaker_session.boto_region_name])
         if job_settings.s3_kms_key:
             container_args.extend(["--s3_kms_key", job_settings.s3_kms_key])
-        if uploaded_dependencies_path:
-            container_args.extend(["--dependencies", uploaded_dependencies_path])
+
         if job_settings.job_conda_env:
             container_args.extend(["--job_conda_env", job_settings.job_conda_env])
 
@@ -282,3 +376,48 @@ def wait(self, timeout: int = None):
             wait=True,
             timeout=timeout,
         )
+
+
+def _prepare_and_upload_runtime_scripts(
+    s3_base_uri: str, s3_kms_key: str, sagemaker_session: Session
+):
+    """Copy runtime scripts to a folder and upload to S3"""
+
+    with _tmpdir() as remote_function_scripts:
+
+        # write entrypoint script to tmpdir
+        entrypoint_script_path = os.path.join(remote_function_scripts, ENTRYPOINT_SCRIPT_NAME)
+        with open(entrypoint_script_path, "w") as file:
+            file.writelines(ENTRYPOINT_SCRIPT)
+
+        bootstrap_script_path = os.path.join(
+            os.path.dirname(__file__), "runtime_environment", BOOTSTRAP_SCRIPT_NAME
+        )
+        runtime_manager_script_path = os.path.join(
+            os.path.dirname(__file__), "runtime_environment", RUNTIME_MANAGER_SCRIPT_NAME
+        )
+
+        # copy runtime scripts to tmpdir
+        shutil.copy2(bootstrap_script_path, remote_function_scripts)
+        shutil.copy2(runtime_manager_script_path, remote_function_scripts)
+
+        return S3Uploader.upload(
+            remote_function_scripts,
+            s3_path_join(s3_base_uri, RUNTIME_SCRIPTS_CHANNEL_NAME),
+            s3_kms_key,
+            sagemaker_session,
+        )
+
+
+def _prepare_and_upload_dependencies(
+    local_dependencies_path: str, s3_base_uri: str, s3_kms_key: str, sagemaker_session: Session
+):
+    """Upload dependency file to S3 if present"""
+    if local_dependencies_path:
+        return S3Uploader.upload(
+            local_dependencies_path,
+            s3_path_join(s3_base_uri, USER_DEPENDENCIES_CHANNEL_NAME),
+            s3_kms_key,
+            sagemaker_session,
+        )
+    return None