feature: allow non-python files in job dependencies (#4138)

* feature: add custom_file_filter argument to job setting

* change: Add more tests

Author: trungleduc

src/sagemaker/feature_store/feature_processor/_config_uploader.py

@@ -12,7 +12,7 @@
 # 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
+from typing import Callable, Dict, Optional, Tuple, List
 
 import attr
 
@@ -70,6 +70,7 @@ def prepare_step_input_channel_for_spark_mode(
             s3_base_uri,
             self.remote_decorator_config.s3_kms_key,
             sagemaker_session,
+            self.remote_decorator_config.custom_file_filter,
         )
 
         (
@@ -134,6 +135,7 @@ def _prepare_and_upload_dependencies(
         s3_base_uri: str,
         s3_kms_key: str,
         sagemaker_session: Session,
+        custom_file_filter: Optional[Callable[[str, List], List]] = None,
     ) -> str:
         """Upload the training step dependencies to S3 if present"""
         return _prepare_and_upload_dependencies(
@@ -144,6 +146,7 @@ def _prepare_and_upload_dependencies(
             s3_base_uri=s3_base_uri,
             s3_kms_key=s3_kms_key,
             sagemaker_session=sagemaker_session,
+            custom_file_filter=custom_file_filter,
         )
 
     def _prepare_and_upload_runtime_scripts(

src/sagemaker/remote_function/client.py

@@ -17,7 +17,7 @@
 from collections import deque
 import time
 import threading
-from typing import Dict, List, Tuple, Any
+from typing import Callable, Dict, List, Optional, Tuple, Any
 import functools
 import itertools
 import inspect
@@ -85,6 +85,7 @@ def remote(
     spark_config: SparkConfig = None,
     use_spot_instances=False,
     max_wait_time_in_seconds=None,
+    custom_file_filter: Optional[Callable[[str, List], List]] = None,
 ):
     """Decorator for running the annotated function as a SageMaker training job.
 
@@ -265,6 +266,11 @@ def remote(
         max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
           After this amount of time Amazon SageMaker will stop waiting for managed spot training
           job to complete. Defaults to ``None``.
+
+        custom_file_filter (Callable[[str, List], List]): A function that filters job
+          dependencies to be uploaded to S3. This function is passed to the ``ignore``
+          argument of ``shutil.copytree``. Defaults to ``None``, which means only python
+          files are accepted.
     """
 
     def _remote(func):
@@ -296,6 +302,7 @@ def _remote(func):
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            custom_file_filter=custom_file_filter,
         )
 
         @functools.wraps(func)
@@ -506,6 +513,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        custom_file_filter: Optional[Callable[[str, List], List]] = None,
     ):
         """Constructor for RemoteExecutor
 
@@ -692,6 +700,11 @@ def __init__(
             max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
               After this amount of time Amazon SageMaker will stop waiting for managed spot training
               job to complete. Defaults to ``None``.
+
+            custom_file_filter (Callable[[str, List], List]): A function that filters job
+              dependencies to be uploaded to S3. This function is passed to the ``ignore``
+              argument of ``shutil.copytree``. Defaults to ``None``, which means only python
+              files are accepted.
         """
         self.max_parallel_jobs = max_parallel_jobs
 
@@ -731,6 +744,7 @@ def __init__(
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            custom_file_filter=custom_file_filter,
         )
 
         self._state_condition = threading.Condition()

src/sagemaker/remote_function/job.py

@@ -20,7 +20,7 @@
 import sys
 import json
 import secrets
-from typing import Dict, List, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 from urllib.parse import urlparse
 from io import BytesIO
 
@@ -193,6 +193,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        custom_file_filter: Optional[Callable[[str, List], List]] = None,
     ):
         """Initialize a _JobSettings instance which configures the remote job.
 
@@ -363,6 +364,11 @@ def __init__(
             max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
               After this amount of time Amazon SageMaker will stop waiting for managed spot
               training job to complete. Defaults to ``None``.
+
+            custom_file_filter (Callable[[str, List], List]): A function that filters job
+              dependencies to be uploaded to S3. This function is passed to the ``ignore``
+              argument of ``shutil.copytree``. Defaults to ``None``, which means only python
+              files are accepted.
         """
         self.sagemaker_session = sagemaker_session or Session()
         self.environment_variables = resolve_value_from_config(
@@ -450,6 +456,7 @@ def __init__(
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
         self.spark_config = spark_config
         self.use_spot_instances = use_spot_instances
+        self.custom_file_filter = custom_file_filter
         self.max_wait_time_in_seconds = max_wait_time_in_seconds
         self.job_conda_env = resolve_value_from_config(
             direct_input=job_conda_env,
@@ -649,6 +656,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
             s3_base_uri=s3_base_uri,
             s3_kms_key=job_settings.s3_kms_key,
             sagemaker_session=job_settings.sagemaker_session,
+            custom_file_filter=job_settings.custom_file_filter,
         )
 
         stored_function = StoredFunction(
@@ -890,6 +898,7 @@ def _prepare_and_upload_dependencies(
     s3_base_uri: str,
     s3_kms_key: str,
     sagemaker_session: Session,
+    custom_file_filter: Optional[Callable[[str, List], List]] = None,
 ) -> str:
     """Upload the job dependencies to S3 if present"""
 
@@ -906,12 +915,12 @@ def _prepare_and_upload_dependencies(
         os.mkdir(tmp_workspace_dir)
         # TODO Remove the following hack to avoid dir_exists error in the copy_tree call below.
         tmp_workspace = os.path.join(tmp_workspace_dir, JOB_REMOTE_FUNCTION_WORKSPACE)
-
+        ignore = custom_file_filter if custom_file_filter is not None else _filter_non_python_files
         if include_local_workdir:
             shutil.copytree(
                 os.getcwd(),
                 tmp_workspace,
-                ignore=_filter_non_python_files,
+                ignore=ignore,
             )
             logger.info("Copied user workspace python scripts to '%s'", tmp_workspace)
 

tests/unit/sagemaker/feature_store/feature_processor/test_config_uploader.py

@@ -50,6 +50,10 @@ def runtime_env_manager():
     return mocked_runtime_env_manager
 
 
+def custom_file_filter():
+    pass
+
+
 @pytest.fixture
 def remote_decorator_config(sagemaker_session):
     return Mock(
@@ -64,6 +68,7 @@ def remote_decorator_config(sagemaker_session):
         pre_execution_script="some_path",
         python_sdk_whl_s3_uri=SAGEMAKER_WHL_FILE_S3_PATH,
         environment_variables={"REMOTE_FUNCTION_SECRET_KEY": "some_secret_key"},
+        custom_file_filter=None,
     )
 
 
@@ -72,6 +77,24 @@ def config_uploader(remote_decorator_config, runtime_env_manager):
     return ConfigUploader(remote_decorator_config, runtime_env_manager)
 
 
+@pytest.fixture
+def remote_decorator_config_with_filter(sagemaker_session):
+    return Mock(
+        _JobSettings,
+        sagemaker_session=sagemaker_session,
+        s3_root_uri="some_s3_uri",
+        s3_kms_key="some_kms",
+        spark_config=SparkConfig(),
+        dependencies=None,
+        include_local_workdir=True,
+        pre_execution_commands="some_commands",
+        pre_execution_script="some_path",
+        python_sdk_whl_s3_uri=SAGEMAKER_WHL_FILE_S3_PATH,
+        environment_variables={"REMOTE_FUNCTION_SECRET_KEY": "some_secret_key"},
+        custom_file_filter=custom_file_filter,
+    )
+
+
 @patch("sagemaker.feature_store.feature_processor._config_uploader.StoredFunction")
 def test_prepare_and_upload_callable(mock_stored_function, config_uploader, wrapped_func):
     mock_stored_function.save(wrapped_func).return_value = None
@@ -108,6 +131,42 @@ def test_prepare_and_upload_dependencies(mock_upload, config_uploader):
         s3_base_uri=remote_decorator_config.s3_root_uri,
         s3_kms_key=remote_decorator_config.s3_kms_key,
         sagemaker_session=sagemaker_session,
+        custom_file_filter=None,
+    )
+
+
+@patch(
+    "sagemaker.feature_store.feature_processor._config_uploader._prepare_and_upload_dependencies",
+    return_value="some_s3_uri",
+)
+def test_prepare_and_upload_dependencies_with_filter(
+    mock_job_upload, remote_decorator_config_with_filter, runtime_env_manager
+):
+    config_uploader_with_filter = ConfigUploader(
+        remote_decorator_config=remote_decorator_config_with_filter,
+        runtime_env_manager=runtime_env_manager,
+    )
+    remote_decorator_config = config_uploader_with_filter.remote_decorator_config
+    config_uploader_with_filter._prepare_and_upload_dependencies(
+        local_dependencies_path="some/path/to/dependency",
+        include_local_workdir=True,
+        pre_execution_commands=remote_decorator_config.pre_execution_commands,
+        pre_execution_script_local_path=remote_decorator_config.pre_execution_script,
+        s3_base_uri=remote_decorator_config.s3_root_uri,
+        s3_kms_key=remote_decorator_config.s3_kms_key,
+        sagemaker_session=sagemaker_session,
+        custom_file_filter=remote_decorator_config_with_filter.custom_file_filter,
+    )
+
+    mock_job_upload.assert_called_once_with(
+        local_dependencies_path="some/path/to/dependency",
+        include_local_workdir=True,
+        pre_execution_commands=remote_decorator_config.pre_execution_commands,
+        pre_execution_script_local_path=remote_decorator_config.pre_execution_script,
+        s3_base_uri=remote_decorator_config.s3_root_uri,
+        s3_kms_key=remote_decorator_config.s3_kms_key,
+        sagemaker_session=sagemaker_session,
+        custom_file_filter=custom_file_filter,
     )
 
 
@@ -201,6 +260,7 @@ def test_prepare_step_input_channel(
         s3_base_uri=remote_decorator_config.s3_root_uri,
         s3_kms_key="some_kms",
         sagemaker_session=sagemaker_session,
+        custom_file_filter=None,
     )
 
     mock_spark_dependency_upload.assert_called_once_with(

tests/unit/sagemaker/feature_store/feature_processor/test_feature_scheduler.py

@@ -263,13 +263,7 @@ def test_to_pipeline(
     )
 
     mock_dependency_upload.assert_called_once_with(
-        local_dependencies_path,
-        True,
-        None,
-        None,
-        f"{S3_URI}/pipeline_name",
-        None,
-        session,
+        local_dependencies_path, True, None, None, f"{S3_URI}/pipeline_name", None, session, None
     )
 
     mock_spark_dependency_upload.assert_called_once_with(
@@ -875,6 +869,7 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):
         "tags",
         "use_spot_instances",
         "max_wait_time_in_seconds",
+        "custom_file_filter",
     }
 
     job_settings = _JobSettings(

tests/unit/sagemaker/remote_function/test_client.py

@@ -175,6 +175,32 @@ def square(x):
     assert mock_job_settings.call_args.kwargs["image_uri"] == IMAGE
 
 
+@patch(
+    "sagemaker.remote_function.client.serialization.deserialize_obj_from_s3",
+    return_value=EXPECTED_JOB_RESULT,
+)
+@patch("sagemaker.remote_function.client._JobSettings")
+@patch("sagemaker.remote_function.client._Job.start")
+def test_decorator_with_custom_file_filter(
+    mock_start, mock_job_settings, mock_deserialize_obj_from_s3
+):
+    mock_job = Mock(job_name=TRAINING_JOB_NAME)
+    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
+
+    mock_start.return_value = mock_job
+
+    def custom_file_filter():
+        pass
+
+    @remote(image_uri=IMAGE, s3_root_uri=S3_URI, custom_file_filter=custom_file_filter)
+    def square(x):
+        return x * x
+
+    result = square(5)
+    assert result == EXPECTED_JOB_RESULT
+    assert mock_job_settings.call_args.kwargs["custom_file_filter"] == custom_file_filter
+
+
 @patch(
     "sagemaker.remote_function.client.serialization.deserialize_exception_from_s3",
     return_value=ZeroDivisionError("division by zero"),

tests/unit/sagemaker/remote_function/test_job.py

@@ -358,6 +358,7 @@ def test_start(
         s3_base_uri=f"{S3_URI}/{job.job_name}",
         s3_kms_key=None,
         sagemaker_session=session(),
+        custom_file_filter=None,
     )
 
     session().sagemaker_client.create_training_job.assert_called_once_with(
@@ -480,6 +481,7 @@ def test_start_with_complete_job_settings(
         s3_base_uri=f"{S3_URI}/{job.job_name}",
         s3_kms_key=job_settings.s3_kms_key,
         sagemaker_session=session(),
+        custom_file_filter=None,
     )
 
     session().sagemaker_client.create_training_job.assert_called_once_with(
@@ -778,6 +780,32 @@ def test_prepare_and_upload_dependencies(session, mock_copytree, mock_copy, mock
     )
 
 
+@patch("sagemaker.s3.S3Uploader.upload", return_value="some_uri")
+@patch("shutil.copy2")
+@patch("shutil.copytree")
+@patch("sagemaker.remote_function.job.Session", return_value=mock_session())
+def test_prepare_and_upload_dependencies_with_custom_filter(
+    session, mock_copytree, mock_copy, mock_s3_upload
+):
+    def custom_file_filter():
+        pass
+
+    s3_path = _prepare_and_upload_dependencies(
+        local_dependencies_path="some/path/to/dependency",
+        include_local_workdir=True,
+        pre_execution_commands=["cmd_1", "cmd_2"],
+        pre_execution_script_local_path=None,
+        s3_base_uri=S3_URI,
+        s3_kms_key=KMS_KEY_ARN,
+        sagemaker_session=session,
+        custom_file_filter=custom_file_filter,
+    )
+
+    assert s3_path == mock_s3_upload.return_value
+
+    mock_copytree.assert_called_with(os.getcwd(), ANY, ignore=custom_file_filter)
+
+
 @patch("sagemaker.remote_function.job.Session", return_value=mock_session())
 def test_prepare_and_upload_spark_dependent_file_without_spark_config(session):
     assert _prepare_and_upload_spark_dependent_files(