feature: support source directory on job side

Author: Namrata Madan

src/sagemaker/remote_function/core/stored_function.py

@@ -16,9 +16,10 @@
 import os
 import pathlib
 import shutil
+import sys
 
 from sagemaker.utils import _tmpdir
-from sagemaker.s3 import s3_path_join, S3Uploader
+from sagemaker.s3 import s3_path_join, S3Uploader, S3Downloader
 from sagemaker.remote_function import logging_config
 
 import sagemaker.remote_function.core.serialization as serialization
@@ -80,17 +81,18 @@ def save(self, func, source_dir=None, *args, **kwargs):
     def _zip_and_upload_source_dir(self, source_dir):
         source_dir_path = pathlib.Path(source_dir)
         if not source_dir_path.is_dir():
-            raise AttributeError(source_dir + "is not a valid directory.")
+            raise AttributeError(source_dir + " is not a valid directory.")
+
+        s3_path = s3_path_join(self.s3_base_uri, "source_dir")
+        logger.info(f"Uploading function source directory to {s3_path}")
         with _tmpdir() as tmp:
             archived_filepath = shutil.make_archive(
-                os.path.join(tmp, "source"), "zip", source_dir_path.parent, source_dir_path.name
-            )
-            S3Uploader.upload(
-                archived_filepath,
-                s3_path_join(self.s3_base_uri, "source_dir"),
-                self.s3_kms_key,
-                self.sagemaker_session,
+                os.path.join(tmp, source_dir_path.name),
+                "zip",
+                source_dir_path.parent,
+                source_dir_path.name,
             )
+            S3Uploader.upload(archived_filepath, s3_path, self.s3_kms_key, self.sagemaker_session)
 
     def load_and_invoke(self) -> None:
         """Load and deserialize the function and the arguments and then execute it."""
@@ -109,6 +111,8 @@ def load_and_invoke(self) -> None:
             self.sagemaker_session, s3_path_join(self.s3_base_uri, "arguments.pkl")
         )
 
+        self._download_and_unzip_source_dir()
+
         logger.info("Invoking the function")
         result = func(*args, **kwargs)
 
@@ -121,3 +125,27 @@ def load_and_invoke(self) -> None:
             s3_path_join(self.s3_base_uri, "results.pkl"),
             self.s3_kms_key,
         )
+
+    def _download_and_unzip_source_dir(self):
+        source_dir_s3_path = s3_path_join(self.s3_base_uri, "source_dir")
+        local_source_dir_path = os.path.join(os.getcwd(), "source_dir")
+
+        logger.info(
+            f"Downloading source modules from {source_dir_s3_path} to {local_source_dir_path}"
+        )
+
+        downloaded_paths = S3Downloader.download(
+            source_dir_s3_path,
+            local_source_dir_path,
+            kms_key=self.s3_kms_key,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        if len(downloaded_paths) < 1:
+            return
+
+        source_dir_archive_path = downloaded_paths[0]
+        shutil.unpack_archive(
+            source_dir_archive_path, pathlib.Path(source_dir_archive_path).parent.absolute()
+        )
+        sys.path.append(local_source_dir_path)

src/sagemaker/s3.py

@@ -171,6 +171,9 @@ def download(s3_uri, local_path, kms_key=None, sagemaker_session=None):
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created
                 using the default AWS configuration chain.
+
+        Returns:
+            list[str]: List of local paths of downloaded files
         """
         sagemaker_session = sagemaker_session or Session()
         bucket, key_prefix = parse_s3_url(url=s3_uri)
@@ -179,7 +182,7 @@ def download(s3_uri, local_path, kms_key=None, sagemaker_session=None):
         else:
             extra_args = None
 
-        sagemaker_session.download_data(
+        return sagemaker_session.download_data(
             path=local_path, bucket=bucket, key_prefix=key_prefix, extra_args=extra_args
         )
 

src/sagemaker/session.py

@@ -389,6 +389,9 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
                 download operation. Please refer to the ExtraArgs parameter in the boto3
                 documentation here:
                 https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-example-download-file.html
+
+        Returns:
+            list[str]: List of local paths of downloaded files
         """
         # Initialize the S3 client.
         if self.s3_client is None:
@@ -408,7 +411,12 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
             if next_token != "":
                 request_parameters.update({"ContinuationToken": next_token})
             response = s3.list_objects_v2(**request_parameters)
-            contents = response.get("Contents")
+            contents = response.get("Contents", None)
+            if not contents:
+                LOGGER.info(
+                    "Nothing to download from bucket: %s, key_prefix: %s.", bucket, key_prefix
+                )
+                return []
             # For each object, save its key or directory.
             for s3_object in contents:
                 key = s3_object.get("Key")
@@ -417,6 +425,7 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
 
         # For each object key, create the directory on the local machine if needed, and then
         # download the file.
+        downloaded_paths = []
         for key in keys:
             tail_s3_uri_path = os.path.basename(key)
             if not os.path.splitext(key_prefix)[1]:
@@ -427,6 +436,8 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
             s3.download_file(
                 Bucket=bucket, Key=key, Filename=destination_path, ExtraArgs=extra_args
             )
+            downloaded_paths.append(destination_path)
+        return downloaded_paths
 
     def read_s3_file(self, bucket, key_prefix):
         """Read a single file from S3.

tests/integ/sagemaker/remote_function/helpers/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import absolute_import
+
+
+def square(x):
+    return x * x

tests/integ/sagemaker/remote_function/helpers/local_module.py

@@ -98,10 +98,32 @@ def divide(x, y):
     assert divide(10, 2) == 5
 
 
+def test_with_local_dependencies(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
+    source_dir_path = os.path.join(os.path.dirname(__file__), "helpers")
+
+    @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 square(x):
+        from helpers import local_module
+
+        return local_module.square(x)
+
+    assert square(9) == 81
+
+
 def test_with_additional_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
 
     @remote(
         role=ROLE,
@@ -121,7 +143,7 @@ def cuberoot(x):
 def test_additional_dependencies_with_job_conda_env(
     sagemaker_session, dummy_container_with_conda, cpu_instance_type
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
 
     @remote(
         role=ROLE,
@@ -142,7 +164,7 @@ def cuberoot(x):
 def test_additional_dependencies_with_default_conda_env(
     sagemaker_session, dummy_container_with_conda, cpu_instance_type
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
 
     @remote(
         role=ROLE,
@@ -162,7 +184,7 @@ def cuberoot(x):
 def test_additional_dependencies_with_non_existent_conda_env(
     sagemaker_session, dummy_container_with_conda, cpu_instance_type
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
 
     @remote(
         role=ROLE,
@@ -184,7 +206,7 @@ def cuberoot(x):
 def test_with_non_existent_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/non_existent_requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "non_existent_requirements.txt")
 
     @remote(
         role=ROLE,
@@ -204,7 +226,7 @@ def test_with_incompatible_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
 
-    dependencies_path = os.path.join(DATA_DIR, "remote_function/old_deps_requirements.txt")
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "old_deps_requirements.txt")
 
     @remote(
         role=ROLE,

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
 import tempfile
 
 import pytest
@@ -31,12 +32,12 @@ def random_s3_uri():
     return "".join(random.choices(string.ascii_uppercase + string.digits, k=10))
 
 
-def upload(b, s3_uri, kms_key=None, sagemaker_session=None):
+def upload_bytes(b, s3_uri, kms_key=None, sagemaker_session=None):
     assert kms_key == KMS_KEY
     mock_s3[s3_uri] = b
 
 
-def read(s3_uri, sagemaker_session=None):
+def read_bytes(s3_uri, sagemaker_session=None):
     return mock_s3[s3_uri]
 
 
@@ -48,10 +49,11 @@ def quadratic(x=2, *, a=1, b=0, c=0):
     "args, kwargs",
     [([], {}), ([3], {}), ([], {"a": 2, "b": 1, "c": 1})],
 )
-@patch("sagemaker.s3.S3Uploader.upload_bytes", new=upload)
-@patch("sagemaker.s3.S3Downloader.read_bytes", new=read)
+@patch("sagemaker.s3.S3Uploader.upload_bytes", new=upload_bytes)
+@patch("sagemaker.s3.S3Downloader.read_bytes", new=read_bytes)
 @patch("sagemaker.s3.S3Uploader.upload")
-def test_save_and_load(mock_s3_upload, args, kwargs):
+@patch("sagemaker.s3.S3Downloader.download")
+def test_save_and_load(s3_source_dir_download, s3_source_dir_upload, args, kwargs):
     session = Mock()
     s3_base_uri = random_s3_uri()
     source_dir = tempfile.TemporaryDirectory()
@@ -61,12 +63,31 @@ def test_save_and_load(mock_s3_upload, args, kwargs):
     )
     stored_function.save(quadratic, source_dir.name, *args, **kwargs)
 
-    mock_s3_upload.assert_called_once_with(ANY, f"{s3_base_uri}/source_dir", KMS_KEY, session)
+    s3_source_dir_upload.assert_called_once_with(ANY, f"{s3_base_uri}/source_dir", KMS_KEY, session)
 
     stored_function.load_and_invoke()
 
+    s3_source_dir_download.assert_called_once_with(
+        f"{s3_base_uri}/source_dir",
+        os.path.join(os.getcwd(), "source_dir"),
+        kms_key=KMS_KEY,
+        sagemaker_session=session,
+    )
+
     assert deserialize_obj_from_s3(session, s3_uri=f"{s3_base_uri}/results.pkl") == quadratic(
         *args, **kwargs
     )
 
     source_dir.cleanup()
+
+
+def test_save_invalid_source_directory():
+    session = Mock()
+    s3_base_uri = random_s3_uri()
+
+    stored_function = StoredFunction(
+        sagemaker_session=session, s3_base_uri=s3_base_uri, s3_kms_key=KMS_KEY
+    )
+
+    with pytest.raises(AttributeError, match="not a valid directory."):
+        stored_function.save(quadratic, "invalid_source_dir")