Master pathways mega pr (aws#882)

* Pathways - Experiments+ integration.

* fix comments

* RemoteExecutor Exp+

* change: add __init__ py in runtime_environment package

* change: integ test fixes

* Some improvement to logging

* Use UTC for the log entry timestamp
* Only configure the "sagemaker" logger instead of root logger

* change: update s3 bucket for sagemaker whl file

* change: update user workdir packing and unpacking

* change: Change remote function dependency cache dir structure

* build(deps): bump apache-airflow from 2.4.1 to 2.5.1 (aws#3722)

* py37 backward compatibility

* minor fix for flake8 due to rebasing

* Revert "Some improvement to logging"

This reverts commit a0a37de179b00e5c3e69fb1c9ba0d7b871a0fbd8.

* minor integ test fix

* fix: list_futures integ test

---------

Co-authored-by: Ao Guo <[email protected]>
Co-authored-by: Rohan Gujarathi <[email protected]>
Co-authored-by: Zhankui Lu <[email protected]>
Co-authored-by: Dipankar Patro <[email protected]>
Co-authored-by: Mourya Baddam <[email protected]>
Co-authored-by: Kalyani Nikure <[email protected]>
Co-authored-by: Dipankar Patro <[email protected]>
Co-authored-by: Namrata Madan <[email protected]>

Author: 9 people

src/sagemaker/experiments/run.py

@@ -715,6 +715,14 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
 
         self.close()
 
+    def __getstate__(self):
+        """Overriding this method to prevent instance of Run from being pickled.
+
+        Raise:
+            NotImplementedError: If attempting to pickle this instance.
+        """
+        raise NotImplementedError("Instance of Run type is not allowed to be pickled.")
+
 
 def load_run(
     run_name: Optional[str] = None,
@@ -792,31 +800,34 @@ def load_run(
 
     verify_load_input_names(run_name=run_name, experiment_name=experiment_name)
 
-    if run_name or environment:
-        if run_name:
-            logger.warning(
-                "run_name is explicitly supplied in load_run, "
-                "which will be prioritized to load the Run object. "
-                "In other words, the run name in the experiment config, fetched from the "
-                "job environment or the current run context, will be ignored."
-            )
-        else:
-            exp_config = get_tc_and_exp_config_from_job_env(
-                environment=environment, sagemaker_session=sagemaker_session
-            )
-            run_name = Run._extract_run_name_from_tc_name(
-                trial_component_name=exp_config[RUN_NAME],
-                experiment_name=exp_config[EXPERIMENT_NAME],
-            )
-            experiment_name = exp_config[EXPERIMENT_NAME]
-
+    if run_name:
+        logger.warning(
+            "run_name is explicitly supplied in load_run, "
+            "which will be prioritized to load the Run object. "
+            "In other words, the run name in the experiment config, fetched from the "
+            "job environment or the current run context, will be ignored."
+        )
         run_instance = Run(
             experiment_name=experiment_name,
             run_name=run_name,
             sagemaker_session=sagemaker_session,
         )
     elif _RunContext.get_current_run():
         run_instance = _RunContext.get_current_run()
+    elif environment:
+        exp_config = get_tc_and_exp_config_from_job_env(
+            environment=environment, sagemaker_session=sagemaker_session
+        )
+        run_name = Run._extract_run_name_from_tc_name(
+            trial_component_name=exp_config[RUN_NAME],
+            experiment_name=exp_config[EXPERIMENT_NAME],
+        )
+        experiment_name = exp_config[EXPERIMENT_NAME]
+        run_instance = Run(
+            experiment_name=experiment_name,
+            run_name=run_name,
+            sagemaker_session=sagemaker_session,
+        )
     else:
         raise RuntimeError(
             "Failed to load a Run object. "

src/sagemaker/remote_function/client.py

@@ -24,6 +24,7 @@
 
 from botocore.exceptions import ClientError
 from sagemaker.exceptions import UnexpectedStatusException
+from sagemaker.experiments._run_context import _RunContext
 
 import sagemaker.remote_function.core.serialization as serialization
 from sagemaker.remote_function.errors import RemoteFunctionError, ServiceError, DeserializationError
@@ -33,7 +34,7 @@
 
 from sagemaker.session import Session
 from sagemaker.s3 import s3_path_join
-from sagemaker.remote_function.job import _JobSettings, _Job
+from sagemaker.remote_function.job import _JobSettings, _Job, _RunInfo
 from sagemaker.remote_function import logging_config
 from sagemaker.utils import name_from_base, base_from_name
 
@@ -143,8 +144,7 @@ def wrapper(*args, **kwargs):
                 volume_kms_key=volume_kms_key,
                 volume_size=volume_size,
             )
-            future = Future()
-            job = future._start_and_notify(job_settings, func, args, kwargs)
+            job = _Job.start(job_settings, func, args, kwargs)
 
             try:
                 job.wait()
@@ -205,12 +205,15 @@ def wrapper(*args, **kwargs):
 class _SubmitRequest:
     """Class that holds parameters and data for creating a new job."""
 
-    def __init__(self, future, job_settings: _JobSettings, func, func_args, func_kwargs):
+    def __init__(
+        self, future, job_settings: _JobSettings, func, func_args, func_kwargs, run_info=None
+    ):
         self.future = future
         self.job_settings = job_settings
         self.func = func
         self.args = func_args
         self.kwargs = func_kwargs
+        self.run_info = run_info
 
 
 def _submit_worker(executor):
@@ -237,7 +240,7 @@ def has_work_to_do():
             time.sleep(_API_CALL_LIMIT["SubmittingIntervalInSecs"])
             # submit a new job
             job = request.future._start_and_notify(
-                request.job_settings, request.func, request.args, request.kwargs
+                request.job_settings, request.func, request.args, request.kwargs, request.run_info
             )
 
             with executor._state_condition:
@@ -417,8 +420,14 @@ def submit(self, func, *args, **kwargs):
 
         with self._state_condition:
             future = Future()
+
+            run_info = None
+            if _RunContext.get_current_run() is not None:
+                run = _RunContext.get_current_run()
+                run_info = _RunInfo(run.experiment_name, run.run_name)
+
             self._pending_request_queue.append(
-                _SubmitRequest(future, self.job_settings, func, args, kwargs)
+                _SubmitRequest(future, self.job_settings, func, args, kwargs, run_info)
             )
 
             if self._workers is None:
@@ -605,7 +614,9 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
         future._return = job_return
         return future
 
-    def _start_and_notify(self, job_settings: _JobSettings, func, func_args, func_kwargs):
+    def _start_and_notify(
+        self, job_settings: _JobSettings, func, func_args, func_kwargs, run_info=None
+    ):
         """Start and record the newly created job in the future object.
 
         The job is recorded if one is successfully started. Otherwise, the exception is
@@ -615,7 +626,7 @@ def _start_and_notify(self, job_settings: _JobSettings, func, func_args, func_kw
             if self._state in [_PENDING]:
 
                 try:
-                    self._job = _Job.start(job_settings, func, func_args, func_kwargs)
+                    self._job = _Job.start(job_settings, func, func_args, func_kwargs, run_info)
                 except (Exception,) as e:  # pylint: disable=broad-except
                     self._exception = e
                     self._state = _FINISHED
@@ -675,11 +686,14 @@ def result(self, timeout: float = None) -> Any:
                                 "FailureReason" in self._job.describe()
                                 and self._job.describe()["FailureReason"]
                             ):
-                                raise RuntimeEnvironmentError(self._job.describe()["FailureReason"])
-                            self._exception = RemoteFunctionError(
-                                "Failed to execute remote function. "
-                                + "Check corresponding job for details."
-                            )
+                                self._exception = RuntimeEnvironmentError(
+                                    self._job.describe()["FailureReason"]
+                                )
+                            else:
+                                self._exception = RemoteFunctionError(
+                                    "Failed to execute remote function. "
+                                    + "Check corresponding job for details."
+                                )
                         else:
                             self._exception = serr
                     self._state = _FINISHED

src/sagemaker/remote_function/core/serialization.py

@@ -90,6 +90,13 @@ def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: st
     try:
         bytes_to_upload = cloudpickle.dumps(obj)
     except Exception as e:
+        if isinstance(
+            e, NotImplementedError
+        ) and "Instance of Run type is not allowed to be pickled." in str(e):
+            raise SerializationError(
+                "Remote function does not allow parameters of Run type."
+            ) from e
+
         raise SerializationError(
             "Error when serializing object of type [{}]: {}".format(type(obj).__name__, repr(e))
         ) from e

src/sagemaker/remote_function/core/stored_function.py

@@ -13,13 +13,7 @@
 """SageMaker job function serializer/deserializer."""
 from __future__ import absolute_import
 
-import os
-import pathlib
-import shutil
-import sys
-
-from sagemaker.utils import _tmpdir
-from sagemaker.s3 import s3_path_join, S3Uploader, S3Downloader
+from sagemaker.s3 import s3_path_join
 from sagemaker.remote_function import logging_config
 
 import sagemaker.remote_function.core.serialization as serialization
@@ -44,20 +38,16 @@ def __init__(self, sagemaker_session, s3_base_uri, s3_kms_key=None):
         self.s3_base_uri = s3_base_uri
         self.s3_kms_key = s3_kms_key
 
-    def save(self, func, source_dir=None, *args, **kwargs):
-        """Serialize and persist the function and it's dependencies and arguments.
+    def save(self, func, *args, **kwargs):
+        """Serialize and persist the function and arguments.
 
         Args:
             func: the python function.
-            source_dir: path to local dependencies/modules. Defaults to ``None``
             args: the positional arguments to func.
             kwargs: the keyword arguments to func.
         Returns:
             None
         """
-        if source_dir:
-            self._zip_and_upload_source_dir(source_dir)
-
         logger.info(
             f"Serializing function code to {s3_path_join(self.s3_base_uri, 'function.pkl')}"
         )
@@ -78,22 +68,6 @@ def save(self, func, source_dir=None, *args, **kwargs):
             self.s3_kms_key,
         )
 
-    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.")
-
-        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_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."""
 
@@ -111,8 +85,6 @@ 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)
 
@@ -125,27 +97,3 @@ 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/remote_function/invoke_function.py

@@ -16,8 +16,14 @@
 
 import argparse
 import sys
+import json
 
 import boto3
+from sagemaker.experiments.run import Run
+from sagemaker.remote_function.job import (
+    KEY_EXPERIMENT_NAME,
+    KEY_RUN_NAME,
+)
 
 from sagemaker.session import Session
 from sagemaker.remote_function.errors import handle_error
@@ -33,6 +39,7 @@ def _parse_agrs():
     parser.add_argument("--region", type=str, required=True)
     parser.add_argument("--s3_base_uri", type=str, required=True)
     parser.add_argument("--s3_kms_key", type=str)
+    parser.add_argument("--run_in_context", type=str)
 
     args, _ = parser.parse_known_args()
     return args
@@ -44,12 +51,28 @@ def _get_sagemaker_session(region):
     return Session(boto_session=boto_session)
 
 
-def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key):
+def _load_run_object(run_in_context: str, sagemaker_session: Session) -> Run:
+    """Load current run in json string into run object"""
+    run_dict = json.loads(run_in_context)
+    return Run(
+        experiment_name=run_dict.get(KEY_EXPERIMENT_NAME),
+        run_name=run_dict.get(KEY_RUN_NAME),
+        sagemaker_session=sagemaker_session,
+    )
+
+
+def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key, run_in_context):
     """Execute stored remote function"""
     from sagemaker.remote_function.core.stored_function import StoredFunction
 
     stored_function = StoredFunction(sagemaker_session, s3_base_uri, s3_kms_key)
-    stored_function.load_and_invoke()
+
+    if run_in_context:
+        run_obj = _load_run_object(run_in_context, sagemaker_session)
+        with run_obj:
+            stored_function.load_and_invoke()
+    else:
+        stored_function.load_and_invoke()
 
 
 def main():
@@ -65,9 +88,10 @@ def main():
         region = args.region
         s3_base_uri = args.s3_base_uri
         s3_kms_key = args.s3_kms_key
+        run_in_context = args.run_in_context
 
         sagemaker_session = _get_sagemaker_session(region)
-        _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key)
+        _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key, run_in_context)
 
     except Exception as e:  # pylint: disable=broad-except
         logger.exception("Error encountered while invoking the remote function.")