change: Update Run init and add Run load and _RunContext (aws#707)

* change: Update Run init and add Run load

Add exp name and run group name to load and address comments

* Address nit comments

Co-authored-by: Dewen Qi <[email protected]>

Author: qidewenwhen

src/sagemaker/amazon/amazon_estimator.py

@@ -27,7 +27,7 @@
 from sagemaker.deprecations import renamed_warning
 from sagemaker.estimator import EstimatorBase, _TrainingJob
 from sagemaker.inputs import FileSystemInput, TrainingInput
-from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import sagemaker_timestamp, check_and_get_run_experiment_config
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 from sagemaker.workflow import is_pipeline_variable
@@ -255,6 +255,7 @@ def fit(
         """
         self._prepare_for_training(records, job_name=job_name, mini_batch_size=mini_batch_size)
 
+        experiment_config = check_and_get_run_experiment_config(experiment_config)
         self.latest_training_job = _TrainingJob.start_new(
             self, records, experiment_config=experiment_config
         )

src/sagemaker/estimator.py

@@ -79,6 +79,7 @@
     get_config_value,
     name_from_base,
     to_string,
+    check_and_get_run_experiment_config,
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
@@ -1122,6 +1123,7 @@ def fit(
         """
         self._prepare_for_training(job_name=job_name)
 
+        experiment_config = check_and_get_run_experiment_config(experiment_config)
         self.latest_training_job = _TrainingJob.start_new(self, inputs, experiment_config)
         self.jobs.append(self.latest_training_job)
         if wait:

src/sagemaker/experiments/_run_context.py

@@ -0,0 +1,58 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Contains the SageMaker Experiment _RunContext class."""
+from __future__ import absolute_import
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sagemaker.experiments.run import Run
+
+
+class _RunContext:
+    """A static context variable to keep track of the current Run object"""
+
+    _context_run = None
+
+    @classmethod
+    def add_run_object(cls, run: "Run"):
+        """Keep track of the current executing Run object
+
+        by adding it to a class static variable.
+
+        Args:
+            run (Run): The current Run object to be tracked.
+        """
+        cls._context_run = run
+
+    @classmethod
+    def drop_current_run(cls) -> "Run":
+        """Drop the Run object tracked in the global static variable
+
+        as its execution finishes (its "with" block ends).
+
+        Return:
+            Run: the dropped Run object.
+        """
+        current_run = cls._context_run
+        cls._context_run = None
+        return current_run
+
+    @classmethod
+    def get_current_run(cls) -> "Run":
+        """Return the current Run object without dropping it.
+
+        Return:
+            Run: the current Run object to be returned.
+        """
+        return cls._context_run

src/sagemaker/experiments/_utils.py

@@ -17,6 +17,8 @@
 
 import mimetypes
 import urllib
+from functools import wraps
+
 from sagemaker.apiutils import _utils
 
 
@@ -68,3 +70,16 @@ def verify_length_of_true_and_predicted(true_labels, predicted_attrs, predicted_
             "Lengths mismatch between true labels and {}: "
             "({} vs {}).".format(predicted_attrs_name, len(true_labels), len(predicted_attrs))
         )
+
+
+def validate_invoked_inside_run_context(func):
+    """A Decorator to force the decorated method called under Run context."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        self_instance = args[0]
+        if not self_instance._inside_load_context and not self_instance._inside_init_context:
+            raise RuntimeError("This method should be called inside context of 'with' statement.")
+        return func(*args, **kwargs)
+
+    return wrapper