feature: add custom_file_filter argument to job setting

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

@@ -108,6 +108,7 @@ 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,
     )