change: Change Run.init and Run.load to constructor and module method respectively (aws#752)

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

Author: qidewenwhen

src/sagemaker/experiments/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+"""Sagemaker Experiment Module"""
+from __future__ import absolute_import
+
+from sagemaker.experiments.run import Run, load_run, list_runs, SortOrderType  # noqa: F401

src/sagemaker/experiments/_environment.py

@@ -99,6 +99,14 @@ def get_trial_component(self, sagemaker_session):
         Returns:
             _TrialComponent: The trial component created from the job. None if not found.
         """
+        # TODO: Remove this condition check once we have a way to retrieve source ARN
+        # from transform job env
+        if self.environment_type == _EnvironmentType.SageMakerTransformJob:
+            logger.error(
+                "Currently getting the job trial component from the transform job environment "
+                "is not supported. Returning None."
+            )
+            return None
 
         def _get_trial_component():
             summaries = list(

src/sagemaker/experiments/_run_context.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from sagemaker.experiments.run import Run
+    from sagemaker.experiments import Run
 
 
 class _RunContext:

src/sagemaker/experiments/_utils.py

@@ -13,13 +13,20 @@
 """Contains the SageMaker Experiment utility methods."""
 from __future__ import absolute_import
 
+import logging
 import os
 
 import mimetypes
 import urllib
 from functools import wraps
+from typing import Optional
 
+from sagemaker import Session
 from sagemaker.apiutils import _utils
+from sagemaker.experiments._environment import _RunEnvironment, _EnvironmentType
+from sagemaker.experiments.trial_component import _TrialComponent
+from sagemaker.utilities.search_expression import Filter, Operator, SearchExpression
+from sagemaker.utils import retry_with_backoff
 
 
 def resolve_artifact_name(file_path):
@@ -93,3 +100,119 @@ def is_already_exist_error(error):
             `botocore.exceptions.ClientError`
     """
     return error["Code"] == "ValidationException" and "already exists" in error["Message"]
+
+
+def get_tc_and_exp_config_from_job_env(
+    environment: _RunEnvironment,
+    sagemaker_session: Session,
+) -> dict:
+    """Retrieve an experiment config from the job environment.
+
+    Args:
+        environment (_RunEnvironment): The run environment object with job specific data.
+        sagemaker_session (sagemaker.session.Session): Session object which
+            manages interactions with Amazon SageMaker APIs and any other
+            AWS services needed. If not specified, one is created using the
+            default AWS configuration chain.
+    """
+    job_name = environment.source_arn.split("/")[-1]
+    if environment.environment_type == _EnvironmentType.SageMakerTrainingJob:
+        job_response = retry_with_backoff(
+            callable_func=lambda: sagemaker_session.describe_training_job(job_name),
+            num_attempts=4,
+        )
+    elif environment.environment_type == _EnvironmentType.SageMakerProcessingJob:
+        job_response = retry_with_backoff(
+            callable_func=lambda: sagemaker_session.describe_processing_job(job_name),
+            num_attempts=4,
+        )
+    else:  # environment.environment_type == _EnvironmentType.SageMakerTransformJob
+        raise RuntimeError(
+            "Failed to load the Run as loading experiment config "
+            "from transform job environment is not currently supported. "
+            "As a workaround, please explicitly pass in "
+            "the experiment_name and run_name in load_run."
+        )
+
+    job_exp_config = job_response.get("ExperimentConfig", dict())
+    from sagemaker.experiments.run import RUN_NAME
+
+    if job_exp_config.get(RUN_NAME, None):
+        return job_exp_config
+    raise RuntimeError(
+        "Not able to fetch RunName in ExperimentConfig of the sagemaker job. "
+        "Please make sure the ExperimentConfig is correctly set."
+    )
+
+
+def verify_load_input_names(
+    run_name: Optional[str] = None,
+    experiment_name: Optional[str] = None,
+):
+    """Verify the run_name and the experiment_name inputs in load_run.
+
+    Args:
+        run_name (str): The run_name supplied by the user (default: None).
+        experiment_name (str): The experiment_name supplied by the user
+            (default: None).
+
+    Raises:
+        ValueError: If run_name is supplied while experiment_name is not.
+    """
+    if not run_name and experiment_name:
+        logging.warning(
+            "No run_name is supplied. Ignoring the provided experiment_name "
+            "since it only takes effect along with run_name. "
+            "Will load the Run object from the job environment or current Run context."
+        )
+    if run_name and not experiment_name:
+        raise ValueError(
+            "Invalid input: experiment_name is missing when run_name is supplied. "
+            "Please supply a valid experiment_name when the run_name is not None."
+        )
+
+
+def is_run_trial_component(trial_component_name: str, sagemaker_session: Session) -> bool:
+    """Check if a trial component is generated by `sagemaker.experiments.Run`
+
+    Args:
+        trial_component_name (str): The name of the trial component.
+        sagemaker_session (sagemaker.session.Session): Session object which
+            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:
+        bool: Indicate whether the trial component is created by
+            `sagemaker.experiments.Run` or not.
+    """
+    search_filter = Filter(
+        name="TrialComponentName",
+        operator=Operator.EQUALS,
+        value=trial_component_name,
+    )
+    search_expression = SearchExpression(filters=[search_filter])
+
+    def search():
+        return list(
+            _TrialComponent.search(
+                search_expression=search_expression,
+                max_results=1,  # TrialComponentName is unique in an account
+                sagemaker_session=sagemaker_session,
+            )
+        )[0]
+
+    try:
+        tc_search_res = retry_with_backoff(search, 4)
+        from sagemaker.experiments.run import RUN_TC_TAG
+
+        if not tc_search_res.tags or RUN_TC_TAG not in tc_search_res.tags:
+            return False
+        return True
+    except Exception as ex:  # pylint: disable=broad-except
+        logging.warning(
+            "Failed to inspect the type of the trial component (%s), due to (%s)",
+            trial_component_name,
+            str(ex),
+        )
+        return False