fix: merge WorkdirConfig and custom_filter_filter parameters

Author: Namrata Madan

src/sagemaker/config/config_schema.py

@@ -52,6 +52,7 @@
 ENVIRONMENT_VARIABLES = "EnvironmentVariables"
 IMAGE_URI = "ImageUri"
 INCLUDE_LOCAL_WORKDIR = "IncludeLocalWorkDir"
+CUSTOM_FILE_FILTER = "CustomFileFilter"
 INSTANCE_TYPE = "InstanceType"
 S3_KMS_KEY_ID = "S3KmsKeyId"
 S3_ROOT_URI = "S3RootUri"
@@ -733,7 +734,7 @@ def _simple_path(*args: str):
                                         },
                                         IMAGE_URI: {TYPE: "string"},
                                         INCLUDE_LOCAL_WORKDIR: {TYPE: "boolean"},
-                                        "WorkdirConfig": {
+                                        CUSTOM_FILE_FILTER: {
                                             TYPE: OBJECT,
                                             ADDITIONAL_PROPERTIES: False,
                                             PROPERTIES: {

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, Optional, Tuple, List
+from typing import Callable, Dict, Optional, Tuple, List, Union
 
 import attr
 
@@ -38,7 +38,7 @@
     RuntimeEnvironmentManager,
 )
 from sagemaker.remote_function.spark_config import SparkConfig
-from sagemaker.remote_function.workdir_config import WorkdirConfig
+from sagemaker.remote_function.custom_file_filter import CustomFileFilter
 from sagemaker.s3 import s3_path_join
 
 
@@ -66,7 +66,6 @@ def prepare_step_input_channel_for_spark_mode(
         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,
@@ -132,19 +131,17 @@ 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,
         s3_kms_key: str,
         sagemaker_session: Session,
-        custom_file_filter: Optional[Callable[[str, List], List]] = None,
+        custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
     ) -> str:
         """Upload the training step dependencies to S3 if present"""
         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/remote_function/__init__.py

@@ -15,5 +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.custom_file_filter import CustomFileFilter  # noqa: F401
 from sagemaker.remote_function.spark_config import SparkConfig  # noqa: F401

src/sagemaker/remote_function/client.py

@@ -17,7 +17,7 @@
 from collections import deque
 import time
 import threading
-from typing import Callable, Dict, List, Optional, Tuple, Any
+from typing import Callable, Dict, List, Optional, Tuple, Any, Union
 import functools
 import itertools
 import inspect
@@ -39,7 +39,7 @@
 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.workdir_config import WorkdirConfig
+from sagemaker.remote_function.custom_file_filter import CustomFileFilter
 
 _API_CALL_LIMIT = {
     "SubmittingIntervalInSecs": 1,
@@ -66,7 +66,7 @@ def remote(
     environment_variables: Dict[str, str] = None,
     image_uri: str = None,
     include_local_workdir: bool = False,
-    workdir_config: WorkdirConfig = None,
+    custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
     instance_count: int = 1,
     instance_type: str = None,
     job_conda_env: str = None,
@@ -87,7 +87,6 @@ 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.
 
@@ -195,10 +194,12 @@ def remote(
           methods that are not available via PyPI or conda. Only python files are included.
           Default value is ``False``.
 
-        workdir_config (WorkdirConfig): A ``WorkdirConfig`` object that specifies the
-          local directories and files to be included in the remote function.
-          workdir_config takes precedence over include_local_workdir.
-          Default value is ``None``.
+        custom_file_filter (Callable[[str, List], List], CustomFileFilter): Either a function
+          that filters job dependencies to be uploaded to S3 or a ``CustomFileFilter`` object
+          that specifies the local directories and files to be included in the remote function.
+          If a callable is passed in, that function is passed to the ``ignore``  argument of
+          ``shutil.copytree``. Defaults to ``None``, which means only python
+          files are accepted and uploaded to S3.
 
         instance_count (int): The number of instances to use. Defaults to 1.
           NOTE: Remote function does not support instance_count > 1 for non Spark jobs.
@@ -274,11 +275,6 @@ 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):
@@ -290,7 +286,7 @@ def _remote(func):
             environment_variables=environment_variables,
             image_uri=image_uri,
             include_local_workdir=include_local_workdir,
-            workdir_config=workdir_config,
+            custom_file_filter=custom_file_filter,
             instance_count=instance_count,
             instance_type=instance_type,
             job_conda_env=job_conda_env,
@@ -311,7 +307,6 @@ 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)
@@ -501,7 +496,7 @@ def __init__(
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
         include_local_workdir: bool = False,
-        workdir_config: WorkdirConfig = None,
+        custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
         instance_count: int = 1,
         instance_type: str = None,
         job_conda_env: str = None,
@@ -523,7 +518,6 @@ 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
 
@@ -628,10 +622,12 @@ def __init__(
               local directories. Set to ``True`` if the remote function code imports local modules
               and methods that are not available via PyPI or conda. Default value is ``False``.
 
-            workdir_config (WorkdirConfig): A ``WorkdirConfig`` object that specifies the
-              local directories and files to be included in the remote function.
-              workdir_config takes precedence over include_local_workdir.
-              Default value is ``None``.
+            custom_file_filter (Callable[[str, List], List], CustomFileFilter): Either a function
+              that filters job dependencies to be uploaded to S3 or a ``CustomFileFilter`` object
+              that specifies the local directories and files to be included in the remote function.
+              If a callable is passed in, that function is passed to the ``ignore``  argument of
+              ``shutil.copytree``. Defaults to ``None``, which means only python
+              files are accepted and uploaded to S3.
 
             instance_count (int): The number of instances to use. Defaults to 1.
               NOTE: Remote function does not support instance_count > 1 for non Spark jobs.
@@ -715,11 +711,6 @@ 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
 
@@ -739,7 +730,7 @@ def __init__(
             environment_variables=environment_variables,
             image_uri=image_uri,
             include_local_workdir=include_local_workdir,
-            workdir_config=workdir_config,
+            custom_file_filter=custom_file_filter,
             instance_count=instance_count,
             instance_type=instance_type,
             job_conda_env=job_conda_env,
@@ -760,7 +751,6 @@ 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/workdir_config.py renamed to src/sagemaker/remote_function/custom_file_filter.py

@@ -16,17 +16,17 @@
 import fnmatch
 import os
 import shutil
-from typing import List, Optional, Callable
+from typing import List, Optional, Callable, Union
 
 from sagemaker.utils import resolve_value_from_config
-from sagemaker.config.config_schema import REMOTE_FUNCTION_PATH
+from sagemaker.config.config_schema import REMOTE_FUNCTION_PATH, CUSTOM_FILE_FILTER
 
 
-class WorkdirConfig:
+class CustomFileFilter:
     """Configuration that specifies how the local working directory should be packaged."""
 
     def __init__(self, *, ignore_name_patterns: List[str] = None):
-        """Initialize a WorkdirConfig.
+        """Initialize a CustomFileFilter.
 
         Args:
             ignore_name_patterns (List[str]): ignore files or directories with names
@@ -50,59 +50,51 @@ def workdir(self):
         return self._workdir
 
 
-def resolve_workdir_config_from_config_file(
-    direct_input: WorkdirConfig = None, sagemaker_session=None
-) -> WorkdirConfig:
-    """Resolve the workdir configuration from the config file.
+def resolve_custom_file_filter_from_config_file(
+    direct_input: Union[Callable[[str, List], List], CustomFileFilter] = None,
+    sagemaker_session=None,
+) -> Union[Callable[[str, List], List], CustomFileFilter, None]:
+    """Resolve the CustomFileFilter configuration from the config file.
 
     Args:
-        direct_input (WorkdirConfig): direct input from the user.
+        direct_input (Callable[[str, List], List], CustomFileFilter): direct input from the user.
         sagemaker_session (sagemaker.session.Session): sagemaker session.
     Returns:
-        WorkdirConfig: configuration that specifies how the local
+        CustomFileFilter: configuration that specifies how the local
             working directory should be packaged.
     """
     if direct_input is not None:
         return direct_input
     ignore_name_patterns = resolve_value_from_config(
         direct_input=None,
-        config_path=".".join([REMOTE_FUNCTION_PATH, "WorkdirConfig", "IgnoreNamePatterns"]),
+        config_path=".".join([REMOTE_FUNCTION_PATH, CUSTOM_FILE_FILTER, "IgnoreNamePatterns"]),
         default_value=None,
         sagemaker_session=sagemaker_session,
     )
     if ignore_name_patterns is not None:
-        return WorkdirConfig(ignore_name_patterns=ignore_name_patterns)
+        return CustomFileFilter(ignore_name_patterns=ignore_name_patterns)
     return None
 
 
-def copy_workdir(workdir_config: WorkdirConfig, dst: str):
+def copy_workdir(
+    dst: str,
+    custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
+):
     """Copy the local working directory to the destination.
 
     Args:
-        workdir_config (WorkdirConfig): configuration that specifies how the local
-            working directory should be packaged.
         dst (str): destination path.
+        custom_file_filter (Union[Callable[[str, List], List], CustomFileFilter): configuration that
+            specifies how the local working directory should be packaged.
     """
 
     def _ignore_patterns(path: str, names: List):  # pylint: disable=unused-argument
         ignored_names = set()
-        if workdir_config.ignore_name_patterns is not None:
-            for pattern in workdir_config.ignore_name_patterns:
+        if custom_file_filter.ignore_name_patterns is not None:
+            for pattern in custom_file_filter.ignore_name_patterns:
                 ignored_names.update(fnmatch.filter(names, pattern))
         return ignored_names
 
-    shutil.copytree(
-        workdir_config.workdir,
-        dst,
-        ignore=_ignore_patterns,
-    )
-
-
-def copy_local_files(
-    custom_file_filter: Optional[Callable[[str, List], List]], workdir: str, dst: str
-):
-    """Copy files from the local working directory to the destination."""
-
     def _filter_non_python_files(path: str, names: List) -> List:
         """Ignore function for filtering out non python files."""
         to_ignore = []
@@ -119,9 +111,18 @@ def _filter_non_python_files(path: str, names: List) -> List:
 
         return to_ignore
 
-    ignore = custom_file_filter if custom_file_filter is not None else _filter_non_python_files
+    _ignore = None
+    _src = os.getcwd()
+    if not custom_file_filter:
+        _ignore = _filter_non_python_files
+    elif callable(custom_file_filter):
+        _ignore = custom_file_filter
+    elif isinstance(custom_file_filter, CustomFileFilter):
+        _ignore = _ignore_patterns
+        _src = custom_file_filter.workdir
+
     shutil.copytree(
-        workdir,
+        _src,
         dst,
-        ignore=ignore,
+        ignore=_ignore,
     )