convert source_dir (str) argument to include_local_workdir (bool) (aws#883)

nmadan · Namrata Madan · Namrata Madan · commit 87b51d45589e · 2023-04-18T13:18:20.000-07:00
Co-authored-by: Namrata Madan &lt;nmmadan@amazon.com&gt;
diff --git a/src/sagemaker/config/config_schema.py b/src/sagemaker/config/config_schema.py
@@ -50,7 +50,6 @@
 INSTANCE_TYPE = "InstanceType"
 S3_KMS_KEY_ID = "S3KmsKeyId"
 S3_ROOT_URI = "S3RootUri"
-SOURCE_DIR = "SourceDir"
 TAGS = "Tags"
 JOB_CONDA_ENV = "JobCondaEnvironment"
 OFFLINE_STORE_CONFIG = "OfflineStoreConfig"
@@ -195,7 +194,6 @@
                         ),
                         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_VpcConfig.html#sagemaker-Type-VpcConfig-SecurityGroupIds
                         SchemaOptional(SECURITY_GROUP_IDS): [And(str, Regex(r"[-0-9a-zA-Z]+"))],
-                        SchemaOptional(SOURCE_DIR): str,
                         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_VpcConfig.html#sagemaker-Type-VpcConfig-Subnets
                         SchemaOptional(SUBNETS): [And(str, Regex(r"[-0-9a-zA-Z]+"))],
                         SchemaOptional(TAGS): [TAG],
diff --git a/src/sagemaker/remote_function/client.py b/src/sagemaker/remote_function/client.py
@@ -61,6 +61,7 @@ def remote(
     dependencies: str = None,
     environment_variables: Dict[str, str] = None,
     image_uri: str = None,
+    include_local_workdir: bool = False,
     instance_count: int = 1,
     instance_type: str = None,
     job_conda_env: str = None,
@@ -73,7 +74,6 @@ def remote(
     s3_root_uri: str = None,
     sagemaker_session: Session = None,
     security_group_ids: List[str] = None,
-    source_dir: str = None,
     subnets: List[str] = None,
     tags: List[Tuple[str, str]] = None,
     volume_kms_key: str = None,
@@ -87,6 +87,8 @@ def remote(
             ``from_active_conda_env``. Defaults to None.
         environment_variables (Dict): environment variables
         image_uri (str): Docker image URI on ECR.
+        include_local_workdir (bool): 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``.
         instance_count (int): Number of instance to use. Default is 1.
         instance_type (str): EC2 instance type.
         job_conda_env (str): Name of the conda environment to activate during execution of the job.
@@ -106,8 +108,6 @@ def remote(
             AWS service calls are delegated to (default: None). If not provided, one is created
             with default AWS configuration chain.
         security_group_ids (List[str]): List of security group IDs.
-        source_dir (str): Path to locally defined modules that are used in the remote function.
-            Default is None.
         subnets (List[str]): List of subnet IDs.
         tags (List[Tuple[str, str]]): List of tags attached to the job.
         volume_kms_key (str): KMS key used for encrypting EBS volume attached to the training
@@ -126,6 +126,7 @@ def wrapper(*args, **kwargs):
                 dependencies=dependencies,
                 environment_variables=environment_variables,
                 image_uri=image_uri,
+                include_local_workdir=include_local_workdir,
                 instance_count=instance_count,
                 instance_type=instance_type,
                 job_conda_env=job_conda_env,
@@ -138,7 +139,6 @@ def wrapper(*args, **kwargs):
                 s3_root_uri=s3_root_uri,
                 sagemaker_session=sagemaker_session,
                 security_group_ids=security_group_ids,
-                source_dir=source_dir,
                 subnets=subnets,
                 tags=tags,
                 volume_kms_key=volume_kms_key,
@@ -309,6 +309,7 @@ def __init__(
         dependencies: str = None,
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
+        include_local_workdir: bool = False,
         instance_count: int = 1,
         instance_type: str = None,
         job_conda_env: str = None,
@@ -322,7 +323,6 @@ def __init__(
         s3_root_uri: str = None,
         sagemaker_session: Session = None,
         security_group_ids: List[str] = None,
-        source_dir: str = None,
         subnets: List[str] = None,
         tags: List[Tuple[str, str]] = None,
         volume_kms_key: str = None,
@@ -336,6 +336,9 @@ def __init__(
             environment_variables (Dict): Environment variables passed to the underlying sagemaker
                 job. Defaults to None
             image_uri (str): Docker image URI on ECR. Defaults to base Python image.
+            include_local_workdir (bool): 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``.
             instance_count (int): Number of instance to use. Defaults to 1.
             instance_type (str): EC2 instance type.
             job_conda_env (str): Name of the conda environment to activate during execution
@@ -360,7 +363,6 @@ def __init__(
                 AWS service calls are delegated to (default: None). If not provided, one is created
                 with default AWS configuration chain.
             security_group_ids (List[str]): List of security group IDs. Defaults to None.
-            source_dir: path to local dependencies/modules.
             subnets (List[str]): List of subnet IDs. Defaults to None.
             tags (List[Tuple[str, str]]): List of tags attached to the job. Defaults to None.
             volume_kms_key (str): KMS key used for encrypting EBS volume attached to the training
@@ -377,6 +379,7 @@ def __init__(
             dependencies=dependencies,
             environment_variables=environment_variables,
             image_uri=image_uri,
+            include_local_workdir=include_local_workdir,
             instance_count=instance_count,
             instance_type=instance_type,
             job_conda_env=job_conda_env,
@@ -389,7 +392,6 @@ def __init__(
             s3_root_uri=s3_root_uri,
             sagemaker_session=sagemaker_session,
             security_group_ids=security_group_ids,
-            source_dir=source_dir,
             subnets=subnets,
             tags=tags,
             volume_kms_key=volume_kms_key,
diff --git a/src/sagemaker/remote_function/job.py b/src/sagemaker/remote_function/job.py
@@ -107,6 +107,7 @@ def __init__(
         dependencies: str = None,
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
+        include_local_workdir: bool = False,
         instance_count: int = 1,
         instance_type: str = None,
         job_conda_env: str = None,
@@ -119,7 +120,6 @@ def __init__(
         s3_root_uri: str = None,
         sagemaker_session: Session = None,
         security_group_ids: List[str] = None,
-        source_dir: str = None,
         subnets: List[str] = None,
         tags: List[Tuple[str, str]] = None,
         volume_kms_key: str = None,
@@ -142,7 +142,7 @@ def __init__(
             self.image_uri = self._get_default_image(self.sagemaker_session)
 
         self.dependencies = self._get_from_config(dependencies, config_schema.DEPENDENCIES)
-
+        self.include_local_workdir = include_local_workdir
         self.instance_type = self._get_from_config(
             instance_type, config_schema.INSTANCE_TYPE, required=True
         )
@@ -151,7 +151,6 @@ def __init__(
         self.max_runtime_in_seconds = max_runtime_in_seconds
         self.max_retry_attempts = max_retry_attempts
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
-        self.source_dir = self._get_from_config(source_dir, config_schema.SOURCE_DIR)
         self.job_conda_env = self._get_from_config(job_conda_env, config_schema.JOB_CONDA_ENV)
         self.job_name_prefix = job_name_prefix
 
@@ -280,7 +279,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
         dependencies_list_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)
         user_dependencies_s3uri = _prepare_and_upload_dependencies(
             local_dependencies_path=dependencies_list_path,
-            workdir_path=job_settings.source_dir,
+            include_local_workdir=job_settings.include_local_workdir,
             s3_base_uri=s3_base_uri,
             s3_kms_key=job_settings.s3_kms_key,
             sagemaker_session=job_settings.sagemaker_session,
@@ -481,23 +480,23 @@ def _prepare_and_upload_runtime_scripts(
 
 def _prepare_and_upload_dependencies(
     local_dependencies_path: str,
-    workdir_path: str,
+    include_local_workdir: bool,
     s3_base_uri: str,
     s3_kms_key: str,
     sagemaker_session: Session,
 ) -> str:
     """Upload the job dependencies to S3 if present"""
 
-    if not local_dependencies_path and not workdir_path:
+    if not local_dependencies_path and not include_local_workdir:
         return None
 
     with _tmpdir() as tmp_workspace:
         # TODO Remove the following hack to avoid dir_exists error in the copy_tree call below.
         tmp_workspace = os.path.join(tmp_workspace, "remote_function/")
 
-        if workdir_path:
+        if include_local_workdir:
             shutil.copytree(
-                workdir_path,
+                os.getcwd(),
                 tmp_workspace,
                 ignore=_filter_non_python_files,
             )
@@ -516,12 +515,14 @@ def _prepare_and_upload_dependencies(
         workspace_archive_path = shutil.make_archive(workspace_archive_path, "zip", tmp_workspace)
         logger.info("Successfully created workdir archive at '%s'", workspace_archive_path)
 
-        return S3Uploader.upload(
+        upload_path = S3Uploader.upload(
             workspace_archive_path,
             s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE),
             s3_kms_key,
             sagemaker_session,
         )
+        logger.info("Successfully uploaded workdir to '%s'", upload_path)
+        return upload_path
 
 
 def _convert_run_to_json(run: Run) -> str:
diff --git a/tests/data/config/config.yaml b/tests/data/config/config.yaml
@@ -12,7 +12,6 @@ SageMaker:
         S3KmsKeyId: "somekmskeyid"
         S3RootUri: "s3://bucket/key"
         SecurityGroupIds: ["sg123"]
-        SourceDir: "../mymodule"
         Subnets: ["subnet-1234"]
         Tags: [{"someTagKey": "someTagValue"}]
         VolumeKmsKeyId: "somekmskeyid"
diff --git a/tests/data/remote_function/config.yaml b/tests/data/remote_function/config.yaml
@@ -9,7 +9,6 @@ SageMaker:
         JobCondaEnvironment: "my_conda_env"
         S3KmsKeyId: "someS3KmsKey"
         SecurityGroupIds: ["sg123"]
-        SourceDir: "../mymodule"
         Subnets: ["subnet-1234"]
         Tags: [{"someTagKey": "someTagValue"}, {"someTagKey2": "someTagValue2"}]
         VolumeKmsKeyId: "someVolumeKmsKey"
diff --git a/tests/integ/sagemaker/remote_function/test_decorator.py b/tests/integ/sagemaker/remote_function/test_decorator.py
@@ -113,26 +113,28 @@ def divide(x, y):
 
 
 def test_with_local_dependencies(
-    sagemaker_session, dummy_container_without_error, cpu_instance_type
+    sagemaker_session, dummy_container_without_error, cpu_instance_type, monkeypatch
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
     source_dir_path = os.path.join(os.path.dirname(__file__))
+    with monkeypatch.context() as m:
+        m.chdir(source_dir_path)
+        dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
 
-    @remote(
-        role=ROLE,
-        image_uri=dummy_container_without_error,
-        dependencies=dependencies_path,
-        instance_type=cpu_instance_type,
-        sagemaker_session=sagemaker_session,
-        source_dir=source_dir_path,
-    )
-    def train(x):
-        from helpers import local_module
-        from helpers.nested_helper import local_module2
+        @remote(
+            role=ROLE,
+            image_uri=dummy_container_without_error,
+            dependencies=dependencies_path,
+            instance_type=cpu_instance_type,
+            sagemaker_session=sagemaker_session,
+            include_local_workdir=True,
+        )
+        def train(x):
+            from helpers import local_module
+            from helpers.nested_helper import local_module2
 
-        return local_module.square(x) + local_module2.cube(x)
+            return local_module.square(x) + local_module2.cube(x)
 
-    assert train(2) == 12
+        assert train(2) == 12
 
 
 def test_with_additional_dependencies(
diff --git a/tests/unit/sagemaker/config/conftest.py b/tests/unit/sagemaker/config/conftest.py
@@ -63,7 +63,6 @@ def valid_remote_function_module(valid_iam_role_arn, valid_vpc_security_group, v
             "S3KmsKeyId": "somekmskeyid",
             "S3RootUri": "s3://bucket/key",
             "SecurityGroupIds": [valid_vpc_security_group],
-            "SourceDir": "../mymodule",
             "Subnets": [valid_vpc_subnet],
             "Tags": [{"someTagKey": "someTagValue"}],
             "VolumeKmsKeyId": "somekmskeyid",
diff --git a/tests/unit/sagemaker/remote_function/test_job.py b/tests/unit/sagemaker/remote_function/test_job.py
