Pathways: update input channel names and create sylinks for sagemaker whl file

Author: Rohan Gujarathi

src/sagemaker/remote_function/job.py

@@ -37,8 +37,8 @@
 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"
+RUNTIME_SCRIPTS_CHANNEL_NAME = "sm_rf_bootstrap"
+REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
 SAGEMAKER_WHL_CHANNEL_NAME = "sagemaker_whl_file"
 
 SAGEMAKER_SDK_WHL_FILE = (
@@ -50,15 +50,15 @@
     f"/opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{ENTRYPOINT_SCRIPT_NAME}",
 ]
 
-ENTRYPOINT_SCRIPT = r"""
+ENTRYPOINT_SCRIPT = f"""
 #!/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 "$@"
+printf "INFO: Bootstraping runtime environment.\\n"
+python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
 
 if [ -f "remote_function_conda_env.txt" ]
 then
@@ -70,10 +70,10 @@
         conda_exe="conda"
     fi
 
-    printf "INFO: Invoking remote function inside conda environment: $conda_env.\n"
+    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"
+    printf "INFO: No conda env provided. Invoking remote function\\n"
     python -m sagemaker.remote_function.invoke_function "$@"
 fi
 """
@@ -222,7 +222,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
 
         local_dependencies_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)
 
-        remote_function_scripts_s3uri = _prepare_and_upload_runtime_scripts(
+        bootstrap_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,
@@ -258,7 +258,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
                 ChannelName=RUNTIME_SCRIPTS_CHANNEL_NAME,
                 DataSource={
                     "S3DataSource": {
-                        "S3Uri": remote_function_scripts_s3uri,
+                        "S3Uri": bootstrap_scripts_s3uri,
                         "S3DataType": "S3Prefix",
                     }
                 },
@@ -268,7 +268,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
         if user_dependencies_s3uri:
             input_data_config.append(
                 dict(
-                    ChannelName=USER_DEPENDENCIES_CHANNEL_NAME,
+                    ChannelName=REMOTE_FUNCTION_WORKSPACE,
                     DataSource={
                         "S3DataSource": {
                             "S3Uri": user_dependencies_s3uri,
@@ -383,10 +383,10 @@ def _prepare_and_upload_runtime_scripts(
 ):
     """Copy runtime scripts to a folder and upload to S3"""
 
-    with _tmpdir() as remote_function_scripts:
+    with _tmpdir() as bootstrap_scripts:
 
         # write entrypoint script to tmpdir
-        entrypoint_script_path = os.path.join(remote_function_scripts, ENTRYPOINT_SCRIPT_NAME)
+        entrypoint_script_path = os.path.join(bootstrap_scripts, ENTRYPOINT_SCRIPT_NAME)
         with open(entrypoint_script_path, "w") as file:
             file.writelines(ENTRYPOINT_SCRIPT)
 
@@ -398,11 +398,11 @@ def _prepare_and_upload_runtime_scripts(
         )
 
         # copy runtime scripts to tmpdir
-        shutil.copy2(bootstrap_script_path, remote_function_scripts)
-        shutil.copy2(runtime_manager_script_path, remote_function_scripts)
+        shutil.copy2(bootstrap_script_path, bootstrap_scripts)
+        shutil.copy2(runtime_manager_script_path, bootstrap_scripts)
 
         return S3Uploader.upload(
-            remote_function_scripts,
+            bootstrap_scripts,
             s3_path_join(s3_base_uri, RUNTIME_SCRIPTS_CHANNEL_NAME),
             s3_kms_key,
             sagemaker_session,
@@ -416,7 +416,7 @@ def _prepare_and_upload_dependencies(
     if local_dependencies_path:
         return S3Uploader.upload(
             local_dependencies_path,
-            s3_path_join(s3_base_uri, USER_DEPENDENCIES_CHANNEL_NAME),
+            s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE),
             s3_kms_key,
             sagemaker_session,
         )

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -28,9 +28,11 @@
 SUCCESS_EXIT_CODE = 0
 DEFAULT_FAILURE_CODE = 1
 
-DEPENDENCY_FILE_CHANNEL_NAME = "remote_function_dependencies"
+REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
 BASE_CHANNEL_PATH = "/opt/ml/input/data"
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
+SAGEMAKER_WHL_FILE_NAME = "sagemaker-2.120.1.dev0-py2.py3-none-any.whl"
+SAGEMAKER_WHL_CHANNEL = "sagemaker_whl_file"
 
 
 logging.basicConfig(
@@ -73,7 +75,7 @@ def _bootstrap_runtime_environment(
     Args:
         job_conda_env (str): conda environment to be activated. Default is None.
     """
-    dependency_file_path = os.path.join(BASE_CHANNEL_PATH, DEPENDENCY_FILE_CHANNEL_NAME)
+    dependency_file_path = os.path.join(BASE_CHANNEL_PATH, REMOTE_FUNCTION_WORKSPACE)
     dependency_file = None
 
     if not os.path.exists(dependency_file_path):
@@ -86,13 +88,30 @@ def _bootstrap_runtime_environment(
             break
 
     if dependency_file:
+        _create_sm_installer_symlinks()
         RuntimeEnvironmentManager().bootstrap(
             local_dependencies_file=dependency_file, job_conda_env=job_conda_env
         )
     else:
         logger.info("No dependency file found")
 
 
+def _create_sm_installer_symlinks():
+    """Creates a symlink for sagemaker installer
+
+    This function create a symlink of sagemaker whl such that public beta users won't need
+    to provide whl file path in their dependencies file. This logic should be removed before GA
+    """
+
+    whl_file_path = os.path.join(BASE_CHANNEL_PATH, SAGEMAKER_WHL_CHANNEL, SAGEMAKER_WHL_FILE_NAME)
+    destination_path1 = os.path.join(
+        BASE_CHANNEL_PATH, REMOTE_FUNCTION_WORKSPACE, SAGEMAKER_WHL_FILE_NAME
+    )
+    destination_path2 = os.path.join(os.getcwd(), SAGEMAKER_WHL_FILE_NAME)
+    os.symlink(whl_file_path, destination_path1)
+    os.symlink(whl_file_path, destination_path2)
+
+
 def _write_failure_reason_file(failure_msg):
     """Create a file 'failure' with failure reason written if bootstrap runtime env failed.
 

tests/integ/sagemaker/remote_function/conftest.py

@@ -70,7 +70,7 @@
     "dependencies:\n"
     "  - scipy=1.10.0\n"
     "  - pip:\n"
-    "    - /opt/ml/input/data/sagemaker_whl_file/sagemaker-2.120.1.dev0-py2.py3-none-any.whl[remote_function]\n"
+    "    - sagemaker-2.120.1.dev0-py2.py3-none-any.whl[remote_function]\n"
     "prefix: /opt/conda/bin/conda\n"
 )
 

tests/unit/sagemaker/remote_function/runtime_environment/test_bootstrap_runtime_environment.py

@@ -20,7 +20,7 @@
 import sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment as bootstrap
 
 TEST_JOB_CONDA_ENV = "conda_env"
-TEST_DEPENDENCIES_PATH = "/opt/ml/input/data/remote_function_dependencies"
+TEST_DEPENDENCIES_PATH = "/opt/ml/input/data/sm_rf_user_ws"
 
 
 def mock_args():
@@ -33,13 +33,16 @@ def mock_args():
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._create_sm_installer_symlinks"
+)
 @patch("sys.exit")
 @patch("os.path.exists", return_value=True)
 @patch("os.listdir", return_value=["fileA.py", "fileB.sh", "requirements.txt"])
 @patch(
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager.RuntimeEnvironmentManager.bootstrap"
 )
-def test_main_success(bootstrap_runtime, list_dir, path_exists, _exit_process):
+def test_main_success(bootstrap_runtime, list_dir, path_exists, _exit_process, symlink):
     bootstrap.main()
     path_exists.assert_called_once_with(TEST_DEPENDENCIES_PATH)
     list_dir.assert_called_once_with(TEST_DEPENDENCIES_PATH)
@@ -51,14 +54,17 @@ def test_main_success(bootstrap_runtime, list_dir, path_exists, _exit_process):
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._create_sm_installer_symlinks"
+)
 @patch("sys.exit")
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._write_failure_reason_file"
 )
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._bootstrap_runtime_environment"
 )
-def test_main_failure(bootstrap_runtime, write_failure, _exit_process):
+def test_main_failure(bootstrap_runtime, write_failure, _exit_process, symlink):
     runtime_err = RuntimeEnvironmentError("some failure reason")
     bootstrap_runtime.side_effect = runtime_err
 

tests/unit/sagemaker/remote_function/test_job.py

@@ -34,8 +34,8 @@
 DEFAULT_ROLE_ARN = "default_execution_role_arn"
 PATH_TO_SRC_DIR = "/path/to/source_dir"
 TEST_REGION = "us-west-2"
-RUNTIME_SCRIPTS_CHANNEL_NAME = "remote_function_scripts"
-USER_DEPENDENCIES_CHANNEL_NAME = "remote_function_dependencies"
+RUNTIME_SCRIPTS_CHANNEL_NAME = "sm_rf_bootstrap"
+REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
 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"
@@ -187,7 +187,7 @@ def test_start(
                 },
             ),
             dict(
-                ChannelName=USER_DEPENDENCIES_CHANNEL_NAME,
+                ChannelName=REMOTE_FUNCTION_WORKSPACE,
                 DataSource={
                     "S3DataSource": {
                         "S3Uri": mock_dependency_upload.return_value,
@@ -211,7 +211,7 @@ def test_start(
             TrainingInputMode="File",
             ContainerEntrypoint=[
                 "/bin/bash",
-                "/opt/ml/input/data/remote_function_scripts/job_driver.sh",
+                "/opt/ml/input/data/sm_rf_bootstrap/job_driver.sh",
             ],
             ContainerArguments=[
                 "--s3_base_uri",
@@ -297,7 +297,7 @@ def test_start_with_complete_job_settings(
                 },
             ),
             dict(
-                ChannelName=USER_DEPENDENCIES_CHANNEL_NAME,
+                ChannelName=REMOTE_FUNCTION_WORKSPACE,
                 DataSource={
                     "S3DataSource": {
                         "S3Uri": mock_dependency_upload.return_value,
@@ -321,7 +321,7 @@ def test_start_with_complete_job_settings(
             TrainingInputMode="File",
             ContainerEntrypoint=[
                 "/bin/bash",
-                "/opt/ml/input/data/remote_function_scripts/job_driver.sh",
+                "/opt/ml/input/data/sm_rf_bootstrap/job_driver.sh",
             ],
             ContainerArguments=[
                 "--s3_base_uri",