Added torchrun execution for remote jobs

Author: brunopistone

src/sagemaker/remote_function/client.py

@@ -35,13 +35,12 @@
 
 from sagemaker.session import Session
 from sagemaker.s3 import s3_path_join
-from sagemaker.remote_function.job import _JobSettings, _Job, _RunInfo
+from sagemaker.remote_function.job import _Job, _JobSettings
+from sagemaker.remote_function.job import _RunInfo
 from sagemaker.remote_function import logging_config
 from sagemaker.utils import name_from_base, base_from_name
 from sagemaker.remote_function.spark_config import SparkConfig
 from sagemaker.remote_function.custom_file_filter import CustomFileFilter
-from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
-from sagemaker.telemetry.constants import Feature
 
 _API_CALL_LIMIT = {
     "SubmittingIntervalInSecs": 1,
@@ -59,7 +58,6 @@
 logger = logging_config.get_logger()
 
 
-@_telemetry_emitter(feature=Feature.REMOTE_FUNCTION, func_name="remote_function.remote")
 def remote(
     _func=None,
     *,
@@ -90,6 +88,8 @@ def remote(
     spark_config: SparkConfig = None,
     use_spot_instances=False,
     max_wait_time_in_seconds=None,
+    use_torchrun=False,
+    nproc_per_node=1,
 ):
     """Decorator for running the annotated function as a SageMaker training job.
 
@@ -310,6 +310,8 @@ def _remote(func):
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            use_torchrun=use_torchrun,
+            nproc_per_node=nproc_per_node,
         )
 
         @functools.wraps(func)

src/sagemaker/remote_function/job.py

@@ -162,6 +162,51 @@
 fi
 """
 
+ENTRYPOINT_TORCHRUN_SCRIPT = f"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function with torchrun
+
+set -eu
+
+PERSISTENT_CACHE_DIR=${{SAGEMAKER_MANAGED_WARMPOOL_CACHE_DIRECTORY:-/opt/ml/cache}}
+export CONDA_PKGS_DIRS=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/conda/pkgs
+printf "INFO: CONDA_PKGS_DIRS is set to '$CONDA_PKGS_DIRS'\\n"
+export PIP_CACHE_DIR=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/pip
+printf "INFO: PIP_CACHE_DIR is set to '$PIP_CACHE_DIR'\\n"
+
+
+printf "INFO: Bootstraping runtime environment.\\n"
+python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
+
+if [ -d {JOB_REMOTE_FUNCTION_WORKSPACE} ]
+then
+    if [ -f "remote_function_conda_env.txt" ]
+    then
+        cp remote_function_conda_env.txt {JOB_REMOTE_FUNCTION_WORKSPACE}/remote_function_conda_env.txt
+    fi
+    printf "INFO: Changing workspace to {JOB_REMOTE_FUNCTION_WORKSPACE}.\\n"
+    cd {JOB_REMOTE_FUNCTION_WORKSPACE}
+fi
+
+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 with torchrun inside conda environment: $conda_env.\\n"
+    $conda_exe run -n $conda_env torchrun --nproc_per_node $NPROC_PER_NODE -m sagemaker.remote_function.invoke_function "$@"
+else
+    printf "INFO: No conda env provided. Invoking remote function with torchrun\\n"
+    torchrun --nproc_per_node $NPROC_PER_NODE -m sagemaker.remote_function.invoke_function "$@"
+fi
+"""
+
 SPARK_ENTRYPOINT_SCRIPT = f"""
 #!/bin/bash
 
@@ -216,6 +261,8 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        use_torchrun=False,
+        nproc_per_node=1,
     ):
         """Initialize a _JobSettings instance which configures the remote job.
 
@@ -555,6 +602,9 @@ def __init__(
         tags = format_tags(tags)
         self.tags = self.sagemaker_session._append_sagemaker_config_tags(tags, REMOTE_FUNCTION_TAGS)
 
+        self.use_torchrun = use_torchrun
+        self.nproc_per_node = nproc_per_node
+
     @staticmethod
     def _get_default_image(session):
         """Return Studio notebook image, if in Studio env. Else, base python.
@@ -951,7 +1001,7 @@ def _get_job_name(job_settings, func):
 
 
 def _prepare_and_upload_runtime_scripts(
-    spark_config: SparkConfig, s3_base_uri: str, s3_kms_key: str, sagemaker_session: Session
+    spark_config: SparkConfig, s3_base_uri: str, s3_kms_key: str, sagemaker_session: Session, use_torchrun: bool = False, nproc_per_node: int = 1
 ):
     """Copy runtime scripts to a folder and upload to S3.
 
@@ -988,6 +1038,10 @@ def _prepare_and_upload_runtime_scripts(
             )
             shutil.copy2(spark_script_path, bootstrap_scripts)
 
+        if use_torchrun:
+            entry_point_script = ENTRYPOINT_TORCHRUN_SCRIPT
+            entry_point_script = entry_point_script.replace("$NPROC_PER_NODE", str(nproc_per_node))
+
         with open(entrypoint_script_path, "w", newline="\n") as file:
             file.writelines(entry_point_script)
 
@@ -1025,6 +1079,8 @@ def _generate_input_data_config(job_settings: _JobSettings, s3_base_uri: str):
         s3_base_uri=s3_base_uri,
         s3_kms_key=job_settings.s3_kms_key,
         sagemaker_session=job_settings.sagemaker_session,
+        use_torchrun=job_settings.use_torchrun,
+        nproc_per_node=job_settings.nproc_per_node,
     )
 
     input_data_config = [