feature: Add Pipeline step decorator, NotebookJobStep, and scheduler (#1271)

Co-authored-by: Rohan Gujarathi <[email protected]>
Co-authored-by: svia3 <[email protected]>
Co-authored-by: Zhankui Lu <[email protected]>
Co-authored-by: Dewen Qi <[email protected]>
Co-authored-by: Edward Sun <[email protected]>
Co-authored-by: Stephen Via <[email protected]>
Co-authored-by: Namrata Madan <[email protected]>
Co-authored-by: Stacia Choe <[email protected]>
Co-authored-by: Edward Sun <[email protected]>
Co-authored-by: Edward Sun <[email protected]>
Co-authored-by: Rohan Gujarathi <[email protected]>
fix: Multiple bug fixes including removing unsupported feature. (#1105)
Fix some problems with pipeline compilation (#1125)
fix: Refactor JsonGet s3 URI and add serialize_output_to_json flag (#1164)
fix: invoke_function circular import (#1262)
fix: pylint (#1264)
fix: Add logging for docker build failures (#1267)

Author: 12 people

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -96,6 +96,10 @@ Pipeline Context
 .. autoclass:: sagemaker.workflow.pipeline_context.LocalPipelineSession
     :members:
 
+Pipeline Schedule
+-----------------
+
+.. autoclass:: sagemaker.workflow.triggers.PipelineSchedule
 
 Parallelism Configuration
 -------------------------
@@ -120,7 +124,6 @@ Selective Execution Config
 
 .. autoclass:: sagemaker.workflow.selective_execution_config.SelectiveExecutionConfig
 
-
 Properties
 ----------
 
@@ -162,6 +165,8 @@ Steps
 
 .. autoclass:: sagemaker.workflow.steps.ProcessingStep
 
+.. autoclass:: sagemaker.workflow.notebook_job_step.NotebookJobStep
+
 .. autoclass:: sagemaker.workflow.steps.CreateModelStep
 
 .. autoclass:: sagemaker.workflow.callback_step.CallbackStep
@@ -185,3 +190,14 @@ Steps
 .. autoclass:: sagemaker.workflow.emr_step.EMRStep
 
 .. autoclass:: sagemaker.workflow.automl_step.AutoMLStep
+
+@step decorator
+---------------
+
+.. automethod:: sagemaker.workflow.function_step.step
+
+.. autoclass:: sagemaker.workflow.function_step.DelayedReturn
+
+.. autoclass:: sagemaker.workflow.step_outputs.StepOutput
+
+.. autofunction:: sagemaker.workflow.step_outputs.get_step

requirements/extras/test_requirements.txt

@@ -32,9 +32,8 @@ sagemaker-feature-store-pyspark-3.3
 # TODO find workaround
 xgboost>=1.6.2,<=1.7.6
 pillow>=9.5.0,<=10.0.0
-torch@https://download.pytorch.org/whl/cpu/torch-2.0.0%2Bcpu-cp310-cp310-linux_x86_64.whl
-torchvision@https://download.pytorch.org/whl/cpu/torchvision-0.15.1%2Bcpu-cp310-cp310-linux_x86_64.whl
 transformers==4.32.0
 sentencepiece==0.1.99
 # https://github.com/triton-inference-server/server/issues/6246
 tritonclient[http]<2.37.0
+nbformat>=5.9,<6

src/sagemaker/config/config_schema.py

@@ -112,6 +112,7 @@
 CONTAINER_ROOT = "container_root"
 REGION_NAME = "region_name"
 TELEMETRY_OPT_OUT = "TelemetryOptOut"
+NOTEBOOK_JOB = "NotebookJob"
 
 
 def _simple_path(*args: str):
@@ -276,6 +277,7 @@ def _simple_path(*args: str):
 MODEL_PACKAGE_VALIDATION_PROFILES_PATH = _simple_path(
     SAGEMAKER, MODEL_PACKAGE, VALIDATION_SPECIFICATION, VALIDATION_PROFILES
 )
+REMOTE_FUNCTION_PATH = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION)
 REMOTE_FUNCTION_DEPENDENCIES = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, DEPENDENCIES
 )
@@ -322,6 +324,20 @@ def _simple_path(*args: str):
     REMOTE_FUNCTION,
     ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
 )
+NOTEBOOK_JOB_ROLE_ARN = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, ROLE_ARN)
+NOTEBOOK_JOB_S3_ROOT_URI = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, S3_ROOT_URI)
+NOTEBOOK_JOB_S3_KMS_KEY_ID = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, S3_KMS_KEY_ID
+)
+NOTEBOOK_JOB_VOLUME_KMS_KEY_ID = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, VOLUME_KMS_KEY_ID
+)
+NOTEBOOK_JOB_VPC_CONFIG_SUBNETS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, VPC_CONFIG, SUBNETS
+)
+NOTEBOOK_JOB_VPC_CONFIG_SECURITY_GROUP_IDS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, NOTEBOOK_JOB, VPC_CONFIG, SECURITY_GROUP_IDS
+)
 MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
     SAGEMAKER,
     MONITORING_SCHEDULE,
@@ -709,6 +725,16 @@ def _simple_path(*args: str):
                                         },
                                         IMAGE_URI: {TYPE: "string"},
                                         INCLUDE_LOCAL_WORKDIR: {TYPE: "boolean"},
+                                        "WorkdirConfig": {
+                                            TYPE: OBJECT,
+                                            ADDITIONAL_PROPERTIES: False,
+                                            PROPERTIES: {
+                                                "IgnoreNamePatterns": {
+                                                    TYPE: "array",
+                                                    "items": {"type": "string"},
+                                                }
+                                            },
+                                        },
                                         INSTANCE_TYPE: {TYPE: "string"},
                                         JOB_CONDA_ENV: {TYPE: "string"},
                                         ROLE_ARN: {"$ref": "#/definitions/roleArn"},
@@ -719,6 +745,17 @@ def _simple_path(*args: str):
                                         VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
                                     },
                                 },
+                                NOTEBOOK_JOB: {
+                                    TYPE: OBJECT,
+                                    ADDITIONAL_PROPERTIES: False,
+                                    PROPERTIES: {
+                                        ROLE_ARN: {"$ref": "#/definitions/roleArn"},
+                                        S3_ROOT_URI: {"$ref": "#/definitions/s3Uri"},
+                                        S3_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"},
+                                        VOLUME_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"},
+                                        VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
+                                    },
+                                },
                                 SERVE: {
                                     TYPE: OBJECT,
                                     ADDITIONAL_PROPERTIES: False,

src/sagemaker/feature_store/feature_processor/_config_uploader.py

@@ -26,7 +26,7 @@
 from sagemaker.inputs import TrainingInput
 from sagemaker.remote_function.core.stored_function import StoredFunction
 from sagemaker.remote_function.job import (
-    _prepare_and_upload_dependencies,
+    _prepare_and_upload_workspace,
     _prepare_and_upload_runtime_scripts,
     _JobSettings,
     RUNTIME_SCRIPTS_CHANNEL_NAME,
@@ -38,6 +38,7 @@
     RuntimeEnvironmentManager,
 )
 from sagemaker.remote_function.spark_config import SparkConfig
+from sagemaker.remote_function.workdir_config import WorkdirConfig
 from sagemaker.s3 import s3_path_join
 
 
@@ -62,9 +63,10 @@ def prepare_step_input_channel_for_spark_mode(
         dependencies_list_path = self.runtime_env_manager.snapshot(
             self.remote_decorator_config.dependencies
         )
-        user_dependencies_s3uri = self._prepare_and_upload_dependencies(
+        user_workspace_s3uri = self._prepare_and_upload_workspace(
             dependencies_list_path,
             self.remote_decorator_config.include_local_workdir,
+            self.remote_decorator_config.workdir_config,
             self.remote_decorator_config.pre_execution_commands,
             self.remote_decorator_config.pre_execution_script,
             s3_base_uri,
@@ -92,7 +94,7 @@ def prepare_step_input_channel_for_spark_mode(
                 distribution=S3_DATA_DISTRIBUTION_TYPE,
             )
         }
-        if user_dependencies_s3uri:
+        if user_workspace_s3uri:
             input_data_config[REMOTE_FUNCTION_WORKSPACE] = TrainingInput(
                 s3_data=s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE),
                 s3_data_type="S3Prefix",
@@ -126,10 +128,11 @@ def _prepare_and_upload_callable(
         )
         stored_function.save(func)
 
-    def _prepare_and_upload_dependencies(
+    def _prepare_and_upload_workspace(
         self,
         local_dependencies_path: str,
         include_local_workdir: bool,
+        workdir_config: WorkdirConfig,
         pre_execution_commands: List[str],
         pre_execution_script_local_path: str,
         s3_base_uri: str,
@@ -138,9 +141,10 @@ def _prepare_and_upload_dependencies(
         custom_file_filter: Optional[Callable[[str, List], List]] = None,
     ) -> str:
         """Upload the training step dependencies to S3 if present"""
-        return _prepare_and_upload_dependencies(
+        return _prepare_and_upload_workspace(
             local_dependencies_path=local_dependencies_path,
             include_local_workdir=include_local_workdir,
+            workdir_config=workdir_config,
             pre_execution_commands=pre_execution_commands,
             pre_execution_script_local_path=pre_execution_script_local_path,
             s3_base_uri=s3_base_uri,

src/sagemaker/local/entities.py

@@ -215,7 +215,9 @@ def start(self, input_data_config, output_data_config, hyperparameters, environm
         """
         for channel in input_data_config:
             if channel["DataSource"] and "S3DataSource" in channel["DataSource"]:
-                data_distribution = channel["DataSource"]["S3DataSource"]["S3DataDistributionType"]
+                data_distribution = channel["DataSource"]["S3DataSource"].get(
+                    "S3DataDistributionType", None
+                )
                 data_uri = channel["DataSource"]["S3DataSource"]["S3Uri"]
             elif channel["DataSource"] and "FileDataSource" in channel["DataSource"]:
                 data_distribution = channel["DataSource"]["FileDataSource"][
@@ -230,7 +232,7 @@ def start(self, input_data_config, output_data_config, hyperparameters, environm
             # use a single Data URI - this makes handling S3 and File Data easier down the stack
             channel["DataUri"] = data_uri
 
-            if data_distribution != "FullyReplicated":
+            if data_distribution and data_distribution != "FullyReplicated":
                 raise RuntimeError(
                     "DataDistribution: %s is not currently supported in Local Mode"
                     % data_distribution

src/sagemaker/local/image.py

@@ -812,7 +812,13 @@ def _create_docker_host(
             "networks": {"sagemaker-local": {"aliases": [host]}},
         }
 
-        if command != "process":
+        is_train_with_entrypoint = False
+        if command == "train" and self.container_entrypoint:
+            # Remote function or Pipeline function step is translated into a training job
+            # with container_entrypoint configured
+            is_train_with_entrypoint = True
+
+        if command != "process" and not is_train_with_entrypoint:
             host_config["command"] = command
         else:
             if self.container_entrypoint:

src/sagemaker/local/local_session.py

@@ -197,6 +197,10 @@ def create_training_job(
             AlgorithmSpecification["TrainingImage"],
             sagemaker_session=self.sagemaker_session,
         )
+        if AlgorithmSpecification.get("ContainerEntrypoint", None):
+            container.container_entrypoint = AlgorithmSpecification["ContainerEntrypoint"]
+        if AlgorithmSpecification.get("ContainerArguments", None):
+            container.container_arguments = AlgorithmSpecification["ContainerArguments"]
         training_job = _LocalTrainingJob(container)
         hyperparameters = kwargs["HyperParameters"] if "HyperParameters" in kwargs else {}
         logger.info("Starting training job")
@@ -718,6 +722,8 @@ def _initialize(
                 else load_sagemaker_config(s3_resource=self.s3_resource)
             )
         else:
+            self.s3_resource = self.boto_session.resource("s3", region_name=self._region_name)
+            self.s3_client = self.boto_session.client("s3", region_name=self._region_name)
             self.sagemaker_config = (
                 sagemaker_config if sagemaker_config else load_sagemaker_config()
             )

src/sagemaker/local/pipeline.py

@@ -15,12 +15,12 @@
 from abc import ABC, abstractmethod
 
 import json
-from copy import deepcopy
 from datetime import datetime
 from typing import Dict, List, Union
 from botocore.exceptions import ClientError
 
 from sagemaker.workflow.conditions import ConditionTypeEnum
+from sagemaker.workflow.function_step import DelayedReturn
 from sagemaker.workflow.steps import StepTypeEnum, Step
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.entities import PipelineVariable
@@ -87,7 +87,7 @@ def evaluate_step_arguments(self, step):
     def _parse_arguments(self, obj, step_name):
         """Parse and evaluate arguments field"""
         if isinstance(obj, dict):
-            obj_copy = deepcopy(obj)
+            obj_copy = {}
             for k, v in obj.items():
                 obj_copy[k] = self._parse_arguments(v, step_name)
             return obj_copy
@@ -108,16 +108,17 @@ def evaluate_pipeline_variable(self, pipeline_variable, step_name):
         elif isinstance(pipeline_variable, Parameter):
             value = self.execution.pipeline_parameters.get(pipeline_variable.name)
         elif isinstance(pipeline_variable, Join):
-            evaluated = [
-                str(self.evaluate_pipeline_variable(v, step_name)) for v in pipeline_variable.values
-            ]
-            value = pipeline_variable.on.join(evaluated)
+            value = self._evaluate_join_function(pipeline_variable, step_name)
         elif isinstance(pipeline_variable, Properties):
             value = self._evaluate_property_reference(pipeline_variable, step_name)
         elif isinstance(pipeline_variable, ExecutionVariable):
             value = self._evaluate_execution_variable(pipeline_variable)
         elif isinstance(pipeline_variable, JsonGet):
             value = self._evaluate_json_get_function(pipeline_variable, step_name)
+        elif isinstance(pipeline_variable, DelayedReturn):
+            # DelayedReturn showing up in arguments, meaning that it's data referenced
+            # We should convert it to JsonGet and evaluate the JsonGet object
+            value = self._evaluate_json_get_function(pipeline_variable._to_json_get(), step_name)
         else:
             self.execution.update_step_failure(
                 step_name, f"Unrecognized pipeline variable {pipeline_variable.expr}."
@@ -127,6 +128,13 @@ def evaluate_pipeline_variable(self, pipeline_variable, step_name):
             self.execution.update_step_failure(step_name, f"{pipeline_variable.expr} is undefined.")
         return value
 
+    def _evaluate_join_function(self, pipeline_variable, step_name):
+        """Evaluate join function runtime value"""
+        evaluated = [
+            str(self.evaluate_pipeline_variable(v, step_name)) for v in pipeline_variable.values
+        ]
+        return pipeline_variable.on.join(evaluated)
+
     def _evaluate_property_reference(self, pipeline_variable, step_name):
         """Evaluate property reference runtime value."""
         try:
@@ -156,6 +164,43 @@ def _evaluate_execution_variable(self, pipeline_variable):
 
     def _evaluate_json_get_function(self, pipeline_variable, step_name):
         """Evaluate join function runtime value."""
+        s3_bucket = None
+        s3_key = None
+        try:
+            if pipeline_variable.property_file:
+                s3_bucket, s3_key = self._evaluate_json_get_property_file_reference(
+                    pipeline_variable=pipeline_variable, step_name=step_name
+                )
+            else:
+                # JsonGet's s3_uri can only be a Join function
+                # This has been validated in _validate_json_get_function
+                s3_uri = self._evaluate_join_function(pipeline_variable.s3_uri, step_name)
+                s3_bucket, s3_key = parse_s3_url(s3_uri)
+
+            file_content = self.sagemaker_session.read_s3_file(s3_bucket, s3_key)
+            file_json = json.loads(file_content)
+            return get_using_dot_notation(file_json, pipeline_variable.json_path)
+        except ClientError as e:
+            self.execution.update_step_failure(
+                step_name,
+                f"Received an error while reading file {s3_path_join('s3://', s3_bucket, s3_key)} "
+                f"from S3: {e.response.get('Code')}: {e.response.get('Message')}",
+            )
+        except json.JSONDecodeError:
+            self.execution.update_step_failure(
+                step_name,
+                f"Contents of file {s3_path_join('s3://', s3_bucket, s3_key)} are not "
+                f"in valid JSON format.",
+            )
+        except ValueError:
+            self.execution.update_step_failure(
+                step_name, f"Invalid json path '{pipeline_variable.json_path}'"
+            )
+
+    def _evaluate_json_get_property_file_reference(
+        self, pipeline_variable: JsonGet, step_name: str
+    ):
+        """Evaluate JsonGet's property file reference to get s3 bucket and key"""
         property_file_reference = pipeline_variable.property_file
         property_file = None
         if isinstance(property_file_reference, str):
@@ -180,28 +225,9 @@ def _evaluate_json_get_function(self, pipeline_variable, step_name):
         processing_output_s3_bucket = processing_step_response["ProcessingOutputConfig"]["Outputs"][
             property_file.output_name
         ]["S3Output"]["S3Uri"]
-        try:
-            s3_bucket, s3_key_prefix = parse_s3_url(processing_output_s3_bucket)
-            file_content = self.sagemaker_session.read_s3_file(
-                s3_bucket, s3_path_join(s3_key_prefix, property_file.path)
-            )
-            file_json = json.loads(file_content)
-            return get_using_dot_notation(file_json, pipeline_variable.json_path)
-        except ClientError as e:
-            self.execution.update_step_failure(
-                step_name,
-                f"Received an error while file reading file '{property_file.path}' from S3: "
-                f"{e.response.get('Code')}: {e.response.get('Message')}",
-            )
-        except json.JSONDecodeError:
-            self.execution.update_step_failure(
-                step_name,
-                f"Contents of property file '{property_file.name}' are not in valid JSON format.",
-            )
-        except ValueError:
-            self.execution.update_step_failure(
-                step_name, f"Invalid json path '{pipeline_variable.json_path}'"
-            )
+        s3_bucket, s3_key_prefix = parse_s3_url(processing_output_s3_bucket)
+        s3_key = s3_path_join(s3_key_prefix, property_file.path)
+        return s3_bucket, s3_key
 
 
 class _StepExecutor(ABC):

src/sagemaker/remote_function/__init__.py

@@ -15,3 +15,5 @@
 
 from sagemaker.remote_function.client import remote, RemoteExecutor  # noqa: F401
 from sagemaker.remote_function.checkpoint_location import CheckpointLocation  # noqa: F401
+from sagemaker.remote_function.workdir_config import WorkdirConfig  # noqa: F401
+from sagemaker.remote_function.spark_config import SparkConfig  # noqa: F401