feature: add SageMaker FeatureStore feature processing (#3944)

Co-authored-by: Alex Tang <[email protected]>
Co-authored-by: cansun <[email protected]>
Co-authored-by: Alex Tang <[email protected]>
Co-authored-by: Can Sun <[email protected]>
Co-authored-by: jiapinw <[email protected]>
Co-authored-by: suryans-commit <[email protected]>

Author: 5 people

requirements/extras/feature-processor_requirements.txt

@@ -0,0 +1,2 @@
+pyspark==3.3.1
+sagemaker-feature-store-pyspark-3.3

requirements/extras/test_requirements.txt

@@ -27,3 +27,5 @@ scipy==1.7.3
 urllib3>=1.26.8,<1.26.15
 docker>=5.0.2,<7.0.0
 PyYAML==6.0
+pyspark==3.3.1
+sagemaker-feature-store-pyspark-3.3

setup.py

@@ -71,6 +71,9 @@ def read_requirements(filename):
 extras = {
     "local": read_requirements("requirements/extras/local_requirements.txt"),
     "scipy": read_requirements("requirements/extras/scipy_requirements.txt"),
+    "feature-processor": read_requirements(
+        "requirements/extras/feature-processor_requirements.txt"
+    ),
 }
 # Meta dependency groups
 extras["all"] = [item for group in extras.values() for item in group]

src/sagemaker/estimator.py

@@ -170,6 +170,8 @@ def __init__(
         instance_groups: Optional[List[InstanceGroup]] = None,
         training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
+        container_entry_point: Optional[List[str]] = None,
+        container_arguments: Optional[List[str]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -523,6 +525,11 @@ def __init__(
                 private Docker registry where your training image is hosted (default: None).
                 When it's set to None, SageMaker will not do authentication before pulling the image
                 in the private Docker registry.
+            container_entry_point (List[str]): Optional. The entrypoint script for a Docker
+                container used to run a training job. This script takes precedence over
+                the default train processing instructions.
+            container_arguments (List[str]): Optional. The arguments for a container used to run
+                a training job.
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -647,6 +654,10 @@ def __init__(
             training_repository_credentials_provider_arn
         )
 
+        # container entry point / arguments configs
+        self.container_entry_point = container_entry_point
+        self.container_arguments = container_arguments
+
         self.encrypt_inter_container_traffic = resolve_value_from_config(
             direct_input=encrypt_inter_container_traffic,
             config_path=TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
@@ -1786,6 +1797,16 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 "MetricsDefinition"
             ]
 
+        if "ContainerEntrypoint" in job_details["AlgorithmSpecification"]:
+            init_params["container_entry_point"] = job_details["AlgorithmSpecification"][
+                "ContainerEntrypoint"
+            ]
+
+        if "ContainerArguments" in job_details["AlgorithmSpecification"]:
+            init_params["container_arguments"] = job_details["AlgorithmSpecification"][
+                "ContainerArguments"
+            ]
+
         if "EnableInterContainerTrafficEncryption" in job_details:
             init_params["encrypt_inter_container_traffic"] = job_details[
                 "EnableInterContainerTrafficEncryption"
@@ -2266,6 +2287,12 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
                 ] = estimator.training_repository_credentials_provider_arn
             train_args["training_image_config"] = training_image_config
 
+        if estimator.container_entry_point is not None:
+            train_args["container_entry_point"] = estimator.container_entry_point
+
+        if estimator.container_arguments is not None:
+            train_args["container_arguments"] = estimator.container_arguments
+
         # encrypt_inter_container_traffic may be a pipeline variable place holder object
         # which is parsed in execution time
         # This does not check config because the EstimatorBase constuctor already did that check
@@ -2472,6 +2499,8 @@ def __init__(
         instance_groups: Optional[List[InstanceGroup]] = None,
         training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
+        container_entry_point: Optional[List[str]] = None,
+        container_arguments: Optional[List[str]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -2824,6 +2853,11 @@ def __init__(
                 private Docker registry where your training image is hosted (default: None).
                 When it's set to None, SageMaker will not do authentication before pulling the image
                 in the private Docker registry.
+            container_entry_point (List[str]): Optional. The entrypoint script for a Docker
+                container used to run a training job. This script takes precedence over
+                the default train processing instructions.
+            container_arguments (List[str]): Optional. The arguments for a container used to run
+                a training job.
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -2871,6 +2905,8 @@ def __init__(
             instance_groups=instance_groups,
             training_repository_access_mode=training_repository_access_mode,
             training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
+            container_entry_point=container_entry_point,
+            container_arguments=container_arguments,
             **kwargs,
         )
 

src/sagemaker/feature_store/feature_processor/__init__.py

@@ -0,0 +1,35 @@
+# 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.
+"""Exported classes for the sagemaker.feature_store.feature_processor module."""
+from __future__ import absolute_import
+
+from sagemaker.feature_store.feature_processor._data_source import (  # noqa: F401
+    CSVDataSource,
+    FeatureGroupDataSource,
+    ParquetDataSource,
+)
+from sagemaker.feature_store.feature_processor._exceptions import (  # noqa: F401
+    IngestionError,
+)
+from sagemaker.feature_store.feature_processor.feature_processor import (  # noqa: F401
+    feature_processor,
+)
+from sagemaker.feature_store.feature_processor.feature_scheduler import (  # noqa: F401
+    to_pipeline,
+    schedule,
+    describe,
+    delete_schedule,
+    list_pipelines,
+    execute,
+    TransformationCode,
+)

src/sagemaker/feature_store/feature_processor/_config_uploader.py

@@ -0,0 +1,180 @@
+# 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.
+"""Contains classes for preparing and uploading configs for a scheduled feature processor."""
+from __future__ import absolute_import
+from typing import Callable, Dict, Tuple, List
+
+import attr
+
+from sagemaker import Session
+from sagemaker.feature_store.feature_processor._constants import (
+    SPARK_JAR_FILES_PATH,
+    SPARK_PY_FILES_PATH,
+    SPARK_FILES_PATH,
+    S3_DATA_DISTRIBUTION_TYPE,
+)
+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_runtime_scripts,
+    _JobSettings,
+    RUNTIME_SCRIPTS_CHANNEL_NAME,
+    REMOTE_FUNCTION_WORKSPACE,
+    SPARK_CONF_WORKSPACE,
+    _prepare_and_upload_spark_dependent_files,
+)
+from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
+    RuntimeEnvironmentManager,
+)
+from sagemaker.remote_function.spark_config import SparkConfig
+from sagemaker.s3 import s3_path_join
+
+
+@attr.s
+class ConfigUploader:
+    """Prepares and uploads customer provided configs to S3"""
+
+    remote_decorator_config: _JobSettings = attr.ib()
+    runtime_env_manager: RuntimeEnvironmentManager = attr.ib()
+
+    def prepare_step_input_channel_for_spark_mode(
+        self, func: Callable, s3_base_uri: str, sagemaker_session: Session
+    ) -> Tuple[Dict, Dict]:
+        """Prepares input channels for SageMaker Pipeline Step."""
+        self._prepare_and_upload_callable(func, s3_base_uri, sagemaker_session)
+        bootstrap_scripts_s3uri = self._prepare_and_upload_runtime_scripts(
+            self.remote_decorator_config.spark_config,
+            s3_base_uri,
+            self.remote_decorator_config.s3_kms_key,
+            sagemaker_session,
+        )
+        dependencies_list_path = self.runtime_env_manager.snapshot(
+            self.remote_decorator_config.dependencies
+        )
+        user_dependencies_s3uri = self._prepare_and_upload_dependencies(
+            dependencies_list_path,
+            self.remote_decorator_config.include_local_workdir,
+            self.remote_decorator_config.pre_execution_commands,
+            self.remote_decorator_config.pre_execution_script,
+            s3_base_uri,
+            self.remote_decorator_config.s3_kms_key,
+            sagemaker_session,
+        )
+
+        (
+            submit_jars_s3_paths,
+            submit_py_files_s3_paths,
+            submit_files_s3_path,
+            config_file_s3_uri,
+        ) = self._prepare_and_upload_spark_dependent_files(
+            self.remote_decorator_config.spark_config,
+            s3_base_uri,
+            self.remote_decorator_config.s3_kms_key,
+            sagemaker_session,
+        )
+
+        input_data_config = {
+            RUNTIME_SCRIPTS_CHANNEL_NAME: TrainingInput(
+                s3_data=bootstrap_scripts_s3uri,
+                s3_data_type="S3Prefix",
+                distribution=S3_DATA_DISTRIBUTION_TYPE,
+            )
+        }
+        if user_dependencies_s3uri:
+            input_data_config[REMOTE_FUNCTION_WORKSPACE] = TrainingInput(
+                s3_data=s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE),
+                s3_data_type="S3Prefix",
+                distribution=S3_DATA_DISTRIBUTION_TYPE,
+            )
+
+        if config_file_s3_uri:
+            input_data_config[SPARK_CONF_WORKSPACE] = TrainingInput(
+                s3_data=config_file_s3_uri,
+                s3_data_type="S3Prefix",
+                distribution=S3_DATA_DISTRIBUTION_TYPE,
+            )
+
+        return input_data_config, {
+            SPARK_JAR_FILES_PATH: submit_jars_s3_paths,
+            SPARK_PY_FILES_PATH: submit_py_files_s3_paths,
+            SPARK_FILES_PATH: submit_files_s3_path,
+        }
+
+    def _prepare_and_upload_callable(
+        self, func: Callable, s3_base_uri: str, sagemaker_session: Session
+    ) -> None:
+        """Prepares and uploads callable to S3"""
+        stored_function = StoredFunction(
+            sagemaker_session=sagemaker_session,
+            s3_base_uri=s3_base_uri,
+            hmac_key=self.remote_decorator_config.environment_variables[
+                "REMOTE_FUNCTION_SECRET_KEY"
+            ],
+            s3_kms_key=self.remote_decorator_config.s3_kms_key,
+        )
+        stored_function.save(func)
+
+    def _prepare_and_upload_dependencies(
+        self,
+        local_dependencies_path: str,
+        include_local_workdir: bool,
+        pre_execution_commands: List[str],
+        pre_execution_script_local_path: str,
+        s3_base_uri: str,
+        s3_kms_key: str,
+        sagemaker_session: Session,
+    ) -> str:
+        """Upload the training step dependencies to S3 if present"""
+        return _prepare_and_upload_dependencies(
+            local_dependencies_path=local_dependencies_path,
+            include_local_workdir=include_local_workdir,
+            pre_execution_commands=pre_execution_commands,
+            pre_execution_script_local_path=pre_execution_script_local_path,
+            s3_base_uri=s3_base_uri,
+            s3_kms_key=s3_kms_key,
+            sagemaker_session=sagemaker_session,
+        )
+
+    def _prepare_and_upload_runtime_scripts(
+        self,
+        spark_config: SparkConfig,
+        s3_base_uri: str,
+        s3_kms_key: str,
+        sagemaker_session: Session,
+    ) -> str:
+        """Copy runtime scripts to a folder and upload to S3"""
+        return _prepare_and_upload_runtime_scripts(
+            spark_config=spark_config,
+            s3_base_uri=s3_base_uri,
+            s3_kms_key=s3_kms_key,
+            sagemaker_session=sagemaker_session,
+        )
+
+    def _prepare_and_upload_spark_dependent_files(
+        self,
+        spark_config: SparkConfig,
+        s3_base_uri: str,
+        s3_kms_key: str,
+        sagemaker_session: Session,
+    ) -> Tuple:
+        """Upload the spark dependencies to S3 if present"""
+        if not spark_config:
+            return None, None, None, None
+
+        return _prepare_and_upload_spark_dependent_files(
+            spark_config=spark_config,
+            s3_base_uri=s3_base_uri,
+            s3_kms_key=s3_kms_key,
+            sagemaker_session=sagemaker_session,
+        )

src/sagemaker/feature_store/feature_processor/_constants.py

@@ -0,0 +1,42 @@
+# 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.
+"""Module containing constants for feature_processor and feature_scheduler module."""
+from __future__ import absolute_import
+
+from sagemaker.workflow.parameters import Parameter, ParameterTypeEnum
+
+DEFAULT_INSTANCE_TYPE = "ml.m5.xlarge"
+DEFAULT_SCHEDULE_STATE = "ENABLED"
+UNDERSCORE = "_"
+RESOURCE_NOT_FOUND_EXCEPTION = "ResourceNotFoundException"
+RESOURCE_NOT_FOUND = "ResourceNotFound"
+EXECUTION_TIME_PIPELINE_PARAMETER = "scheduled_time"
+VALIDATION_EXCEPTION = "ValidationException"
+EVENT_BRIDGE_INVOCATION_TIME = "<aws.scheduler.scheduled-time>"
+SCHEDULED_TIME_PIPELINE_PARAMETER = Parameter(
+    name=EXECUTION_TIME_PIPELINE_PARAMETER, parameter_type=ParameterTypeEnum.STRING
+)
+EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT = "%Y-%m-%dT%H:%M:%SZ"  # 2023-01-01T07:00:00Z
+NO_FLEXIBLE_TIME_WINDOW = dict(Mode="OFF")
+PIPELINE_NAME_MAXIMUM_LENGTH = 80
+PIPELINE_CONTEXT_TYPE = "FeatureEngineeringPipeline"
+SPARK_JAR_FILES_PATH = "submit_jars_s3_paths"
+SPARK_PY_FILES_PATH = "submit_py_files_s3_paths"
+SPARK_FILES_PATH = "submit_files_s3_path"
+FEATURE_PROCESSOR_TAG_KEY = "sm-fs-fe:created-from"
+FEATURE_PROCESSOR_TAG_VALUE = "fp-to-pipeline"
+FEATURE_GROUP_ARN_REGEX_PATTERN = r"arn:(.*?):sagemaker:(.*?):(.*?):feature-group/(.*?)$"
+SAGEMAKER_WHL_FILE_S3_PATH = "s3://ada-private-beta/sagemaker-2.151.1.dev0-py2.py3-none-any.whl"
+S3_DATA_DISTRIBUTION_TYPE = "FullyReplicated"
+PIPELINE_CONTEXT_NAME_TAG_KEY = "sm-fs-fe:feature-engineering-pipeline-context-name"
+PIPELINE_VERSION_CONTEXT_NAME_TAG_KEY = "sm-fs-fe:feature-engineering-pipeline-version-context-name"