Add python_files_only to job setting

Author: trungleduc

src/sagemaker/remote_function/client.py

@@ -85,6 +85,7 @@ def remote(
     spark_config: SparkConfig = None,
     use_spot_instances=False,
     max_wait_time_in_seconds=None,
+    python_files_only=True,
 ):
     """Decorator for running the annotated function as a SageMaker training job.
 
@@ -265,6 +266,8 @@ 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``.
+
+        python_files_only (bool): Specifies whether non-python files are allowed in job dependencies. Defaults to ``True``.
     """
 
     def _remote(func):
@@ -296,6 +299,7 @@ def _remote(func):
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            python_files_only=python_files_only,
         )
 
         @functools.wraps(func)
@@ -506,6 +510,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        python_files_only=True,
     ):
         """Constructor for RemoteExecutor
 
@@ -692,6 +697,8 @@ 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``.
+
+            python_files_only (bool): Specifies whether non-python files are allowed in job dependencies. Defaults to ``True``.
         """
         self.max_parallel_jobs = max_parallel_jobs
 
@@ -731,6 +738,7 @@ def __init__(
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            python_files_only=python_files_only,
         )
 
         self._state_condition = threading.Condition()

src/sagemaker/remote_function/job.py

@@ -193,6 +193,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        python_files_only: bool = True,
     ):
         """Initialize a _JobSettings instance which configures the remote job.
 
@@ -363,6 +364,9 @@ 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``.
+
+            python_files_only (bool): Specifies whether non-python files are allowed in job
+              dependencies. Defaults to ``True``.
         """
         self.sagemaker_session = sagemaker_session or Session()
         self.environment_variables = resolve_value_from_config(
@@ -450,6 +454,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.python_files_only = python_files_only
         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 +654,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,
+            python_files_only=job_settings.python_files_only,
         )
 
         stored_function = StoredFunction(
@@ -890,6 +896,7 @@ def _prepare_and_upload_dependencies(
     s3_base_uri: str,
     s3_kms_key: str,
     sagemaker_session: Session,
+    python_files_only: bool,
 ) -> str:
     """Upload the job dependencies to S3 if present"""
 
@@ -906,12 +913,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 = _filter_non_python_files if python_files_only else None
         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)