sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/enums.py

@@ -0,0 +1,51 @@
+# 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.
+"""Defines enum values."""
+
+from __future__ import absolute_import
+
+import logging
+from enum import Enum
+
+
+LOGGER = logging.getLogger("sagemaker")
+
+
+class EndpointType(Enum):
+    """Types of endpoint"""
+
+    MODEL_BASED = "ModelBased"  # Amazon SageMaker Model Based Endpoint
+    INFERENCE_COMPONENT_BASED = (
+        "InferenceComponentBased"  # Amazon SageMaker Inference Component Based Endpoint
+    )
+
+
+class RoutingStrategy(Enum):
+    """Strategy for routing https traffics."""
+
+    RANDOM = "RANDOM"
+    """The endpoint routes each request to a randomly chosen instance.
+    """
+    LEAST_OUTSTANDING_REQUESTS = "LEAST_OUTSTANDING_REQUESTS"
+    """The endpoint routes requests to the specific instances that have
+    more capacity to process them.
+    """
+
+
+class Tag(str, Enum):
+    """Enum class for tag keys to apply to models."""
+
+    OPTIMIZATION_JOB_NAME = "sagemaker-sdk:optimization-job-name"
+    SPECULATIVE_DRAFT_MODEL_PROVIDER = "sagemaker-sdk:speculative-draft-model-provider"
+    FINE_TUNING_MODEL_PATH = "sagemaker-sdk:fine-tuning-model-path"
+    FINE_TUNING_JOB_NAME = "sagemaker-sdk:fine-tuning-job-name"

sagemaker/core/environment_variables.py

@@ -0,0 +1,101 @@
+# 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.
+"""Accessors to retrieve environment variables for hosting containers."""
+
+from __future__ import absolute_import
+
+import logging
+from typing import Dict, Optional
+
+from sagemaker.core.jumpstart import utils as jumpstart_utils
+from sagemaker.core.jumpstart import artifacts
+from sagemaker.core.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.core.jumpstart.enums import JumpStartModelType, JumpStartScriptScope
+from sagemaker.core.helper.session_helper import Session
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_default(
+    region: Optional[str] = None,
+    model_id: Optional[str] = None,
+    model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    include_aws_sdk_env_vars: bool = True,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    instance_type: Optional[str] = None,
+    script: JumpStartScriptScope = JumpStartScriptScope.INFERENCE,
+    config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+) -> Dict[str, str]:
+    """Retrieves the default container environment variables for the model matching the arguments.
+
+    Args:
+        region (str): Optional. The AWS Region for which to retrieve the default environment
+             variables. (Default: None).
+        model_id (str): Optional. The model ID of the model for which to
+            retrieve the default environment variables. (Default: None).
+        model_version (str): Optional. The version of the model for which to retrieve the
+            default environment variables. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to
+            retrieve model details from. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
+        include_aws_sdk_env_vars (bool): True if environment variables for low-level AWS API call
+            should be included. The `Model` class of the SageMaker Python SDK inserts environment
+            variables that would be required when making the low-level AWS API call.
+            (Default: True).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        instance_type (str): An instance type to optionally supply in order to get environment
+            variables specific for the instance type.
+        script (JumpStartScriptScope): The JumpStart script for which to retrieve environment
+            variables.
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
+    Returns:
+        dict: The variables to use for the model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError(
+            "Must specify JumpStart `model_id` and `model_version` "
+            "when retrieving environment variables."
+        )
+
+    return artifacts._retrieve_default_environment_variables(
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        include_aws_sdk_env_vars=include_aws_sdk_env_vars,
+        sagemaker_session=sagemaker_session,
+        instance_type=instance_type,
+        script=script,
+        config_name=config_name,
+        model_type=model_type,
+    )

sagemaker/core/exceptions.py

@@ -0,0 +1,108 @@
+# 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.
+"""Custom exception classes for Sagemaker SDK"""
+from __future__ import absolute_import
+
+
+class UnexpectedStatusException(ValueError):
+    """Raised when resource status is not expected and thus not allowed for further execution"""
+
+    def __init__(self, message, allowed_statuses, actual_status):
+        self.allowed_statuses = allowed_statuses
+        self.actual_status = actual_status
+        super(UnexpectedStatusException, self).__init__(message)
+
+
+class CapacityError(UnexpectedStatusException):
+    """Raised when resource status is not expected and fails with a reason of CapacityError"""
+
+
+class AsyncInferenceError(Exception):
+    """The base exception class for Async Inference exceptions."""
+
+    fmt = "An unspecified error occurred"
+
+    def __init__(self, **kwargs):
+        msg = self.fmt.format(**kwargs)
+        Exception.__init__(self, msg)
+        self.kwargs = kwargs
+
+
+class ObjectNotExistedError(AsyncInferenceError):
+    """Raised when Amazon S3 object not exist in the given path"""
+
+    fmt = "Object not exist at {output_path}. {message}"
+
+    def __init__(self, message, output_path):
+        super().__init__(message=message, output_path=output_path)
+
+
+class PollingTimeoutError(AsyncInferenceError):
+    """Raised when wait longer than expected and no result object in Amazon S3 bucket yet"""
+
+    fmt = "No result at {output_path} after polling for {seconds} seconds. {message}"
+
+    def __init__(self, message, output_path, seconds):
+        super().__init__(message=message, output_path=output_path, seconds=seconds)
+
+
+class UnexpectedClientError(AsyncInferenceError):
+    """Raised when ClientError's error code is not expected"""
+
+    fmt = "Encountered unexpected client error: {message}"
+
+    def __init__(self, message):
+        super().__init__(message=message)
+
+
+class AutoMLStepInvalidModeError(Exception):
+    """Raised when the automl mode passed into AutoMLStep in invalid"""
+
+    fmt = (
+        "Mode in AutoMLJobConfig must be defined for AutoMLStep. "
+        "AutoMLStep currently only supports ENSEMBLING mode"
+    )
+
+    def __init__(self, **kwargs):
+        msg = self.fmt.format(**kwargs)
+        Exception.__init__(self, msg)
+        self.kwargs = kwargs
+
+
+class AsyncInferenceModelError(AsyncInferenceError):
+    """Raised when model returns errors for failed requests"""
+
+    fmt = "Model returned error: {message} "
+
+    def __init__(self, message):
+        super().__init__(message=message)
+
+
+class ModelStreamError(Exception):
+    """Raised when invoke_endpoint_with_response_stream Response returns ModelStreamError"""
+
+    def __init__(self, message="An error occurred", code=None):
+        self.message = message
+        self.code = code
+        if code is not None:
+            super().__init__(f"{message} (Code: {code})")
+        else:
+            super().__init__(message)
+
+
+class InternalStreamFailure(Exception):
+    """Raised when invoke_endpoint_with_response_stream Response returns InternalStreamFailure"""
+
+    def __init__(self, message="An error occurred"):
+        self.message = message
+        super().__init__(self.message)

sagemaker/core/experiments/__init__.py

@@ -0,0 +1,53 @@
+# 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 Experiments module for tracking experiments, trials, and runs."""
+from __future__ import absolute_import
+
+# Lazy imports to avoid circular dependencies during package initialization
+# Users should import directly from the specific modules:
+# from sagemaker.core.experiments.experiment import Experiment
+# from sagemaker.core.experiments.run import Run
+# etc.
+
+__all__ = [
+    "Experiment",
+    "Run",
+    "_RunContext",
+    "_Trial",
+    "_TrialComponent",
+]
+
+
+def __getattr__(name):
+    """Lazy import to avoid circular dependencies."""
+    if name == "Experiment":
+        from sagemaker.core.experiments.experiment import Experiment
+
+        return Experiment
+    elif name == "Run":
+        from sagemaker.core.experiments.run import Run
+
+        return Run
+    elif name == "_RunContext":
+        from sagemaker.core.experiments._run_context import _RunContext
+
+        return _RunContext
+    elif name == "_Trial":
+        from sagemaker.core.experiments.trial import _Trial
+
+        return _Trial
+    elif name == "_TrialComponent":
+        from sagemaker.core.experiments.trial_component import _TrialComponent
+
+        return _TrialComponent
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

sagemaker/core/experiments/_api_types.py

@@ -0,0 +1,251 @@
+# 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 API objects for SageMaker experiments."""
+from __future__ import absolute_import
+
+import enum
+import numbers
+
+from sagemaker.core.apiutils import _base_types
+
+
+class TrialComponentMetricSummary(_base_types.ApiObject):
+    """Summary model of a trial component.
+
+    Attributes:
+        metric_name (str): The name of the metric.
+        source_arn (str):  The ARN of the source.
+        time_stamp (datetime): Metric last updated value.
+        max (float): The max value of the metric.
+        min (float):  The min value of the metric.
+        last (float):  The last value of the metric.
+        count (float):  The number of samples used to generate the metric.
+        avg (float):  The average value of the metric.
+        std_dev (float):  The standard deviation of the metric.
+    """
+
+    metric_name = None
+    source_arn = None
+    time_stamp = None
+    max = None
+    min = None
+    last = None
+    count = None
+    avg = None
+    std_dev = None
+
+    def __init__(self, metric_name=None, source_arn=None, **kwargs):
+        super(TrialComponentMetricSummary, self).__init__(
+            metric_name=metric_name, source_arn=source_arn, **kwargs
+        )
+
+
+class TrialComponentParameters(_base_types.ApiObject):
+    """A dictionary of TrialComponentParameterValues"""
+
+    @classmethod
+    def from_boto(cls, boto_dict, **kwargs):
+        """Converts a boto dict to a dictionary of TrialComponentParameterValues
+
+        Args:
+            boto_dict (dict): boto response dictionary.
+            **kwargs:  Arbitrary keyword arguments.
+
+        Returns:
+            dict: Dictionary of parameter values.
+        """
+        return_map = {}
+        for key, value in boto_dict.items():
+            return_map[key] = value.get("NumberValue", value.get("StringValue", None))
+        return return_map
+
+    @classmethod
+    def to_boto(cls, parameters):
+        """Converts TrialComponentParameters to dict.
+
+        Args:
+            parameters (TrialComponentParameters): Dictionary to convert.
+
+        Returns:
+            dict: Dictionary of trial component parameters in boto format.
+        """
+        boto_map = {}
+        for key, value in parameters.items():
+            if isinstance(value, numbers.Number):
+                boto_map[key] = {"NumberValue": value}
+            else:
+                boto_map[key] = {"StringValue": str(value)}
+        return boto_map
+
+
+class TrialComponentArtifact(_base_types.ApiObject):
+    """Trial component artifact.
+
+    Attributes:
+        value (str): The artifact value.
+        media_type (str): The media type.
+    """
+
+    value = None
+    media_type = None
+
+    def __init__(self, value=None, media_type=None, **kwargs):
+        super(TrialComponentArtifact, self).__init__(value=value, media_type=media_type, **kwargs)
+
+
+class _TrialComponentStatusType(enum.Enum):
+    """The type of trial component status"""
+
+    InProgress = "InProgress"
+    Completed = "Completed"
+    Failed = "Failed"
+
+
+class TrialComponentStatus(_base_types.ApiObject):
+    """Status of the trial component.
+
+    Attributes:
+        primary_status (str): The status of a trial component.
+        message (str): Status message.
+    """
+
+    primary_status = None
+    message = None
+
+    def __init__(self, primary_status=None, message=None, **kwargs):
+        super(TrialComponentStatus, self).__init__(
+            primary_status=primary_status, message=message, **kwargs
+        )
+
+
+class TrialComponentSummary(_base_types.ApiObject):
+    """Summary model of a trial component.
+
+    Attributes:
+        trial_component_name (str): Name of trial component.
+        trial_component_arn (str): ARN of the trial component.
+        display_name (str): Friendly display name in UI.
+        source_arn (str): ARN of the trial component source.
+        status (str): Status.
+        start_time (datetime): Start time.
+        end_time (datetime): End time.
+        creation_time (datetime): Creation time.
+        created_by (str): Created by.
+        last_modified_time (datetime): Date last modified.
+        last_modified_by (datetime): User last modified.
+    """
+
+    _custom_boto_types = {
+        "status": (TrialComponentStatus, False),
+    }
+    trial_component_name = None
+    trial_component_arn = None
+    display_name = None
+    source_arn = None
+    status = None
+    start_time = None
+    end_time = None
+    creation_time = None
+    created_by = None
+    last_modified_time = None
+    last_modified_by = None
+
+
+class TrialComponentSource(_base_types.ApiObject):
+    """Trial Component Source
+
+    Attributes:
+        source_arn (str): The ARN of the source.
+    """
+
+    source_arn = None
+
+    def __init__(self, source_arn=None, **kwargs):
+        super(TrialComponentSource, self).__init__(source_arn=source_arn, **kwargs)
+
+
+class Parent(_base_types.ApiObject):
+    """The trial/experiment/run that a trial component is associated with.
+
+    Attributes:
+        trial_name (str): Name of the trial.
+        experiment_name (str): Name of the experiment.
+        run_name (str): Name of the run.
+    """
+
+    trial_name = None
+    experiment_name = None
+    run_name = None
+
+
+class TrialComponentSearchResult(_base_types.ApiObject):
+    """Summary model of an Trial Component search result.
+
+    Attributes:
+        trial_component_arn (str): ARN of the trial component.
+        trial_component_name (str): Name of the trial component.
+        display_name (str): Display name of the trial component for UI display.
+        source (dict): The source of the trial component.
+        status (dict): The status of the trial component.
+        start_time (datetime): Start time.
+        end_time (datetime): End time.
+        creation_time (datetime): Creation time.
+        created_by (str): Created by.
+        last_modified_time (datetime): Date last modified.
+        last_modified_by (datetime): User last modified.
+        parameters (dict): The hyperparameters of the component.
+        input_artifacts (dict): The input artifacts of the component.
+        output_artifacts (dict): The output artifacts of the component.
+        metrics (list): The metrics for the component.
+        source_detail (dict): The source of the trial component.
+        tags (list): The list of tags that are associated with the trial component.
+        parents (list[Parent]): The parent of trial component.
+    """
+
+    _custom_boto_types = {
+        "parents": (Parent, True),  # parents is a collection (list) of Parent objects
+    }
+    trial_component_arn = None
+    trial_component_name = None
+    display_name = None
+    source = None
+    status = None
+    start_time = None
+    end_time = None
+    creation_time = None
+    created_by = None
+    last_modified_time = None
+    last_modified_by = None
+    parameters = None
+    input_artifacts = None
+    output_artifacts = None
+    metrics = None
+    source_detail = None
+    tags = None
+    parents = None
+
+
+class TrialSummary(_base_types.ApiObject):
+    """Summary model of a trial.
+
+    Attributes:
+        trial_arn (str): The ARN of the trial.
+        trial_name (str): The name of the trial.
+        creation_time (datetime):  When the trial was created.
+        last_modified_time (datetime): When the trial was last modified.
+    """
+
+    trial_arn = None
+    trial_name = None
+    creation_time = None
+    last_modified_time = None

sagemaker/core/experiments/_environment.py

@@ -0,0 +1,124 @@
+# 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 _RunEnvironment class."""
+from __future__ import absolute_import
+
+import enum
+import json
+import logging
+import os
+
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.experiments import trial_component
+from sagemaker.core.common_utils import retry_with_backoff
+
+TRAINING_JOB_ARN_ENV = "TRAINING_JOB_ARN"
+PROCESSING_JOB_CONFIG_PATH = "/opt/ml/config/processingjobconfig.json"
+TRANSFORM_JOB_ARN_ENV = "TRANSFORM_JOB_ARN"
+MAX_RETRY_ATTEMPTS = 7
+
+logger = logging.getLogger(__name__)
+
+
+class _EnvironmentType(enum.Enum):
+    """SageMaker jobs which data can be pulled from the environment."""
+
+    SageMakerTrainingJob = 1
+    SageMakerProcessingJob = 2
+    SageMakerTransformJob = 3
+
+
+class _RunEnvironment(object):
+    """Retrieves job specific data from the environment."""
+
+    def __init__(self, environment_type: _EnvironmentType, source_arn: str):
+        """Init for _RunEnvironment.
+
+        Args:
+            environment_type (_EnvironmentType): The environment type.
+            source_arn (str): The ARN of the current job.
+        """
+        self.environment_type = environment_type
+        self.source_arn = source_arn
+
+    @classmethod
+    def load(
+        cls,
+        training_job_arn_env: str = TRAINING_JOB_ARN_ENV,
+        processing_job_config_path: str = PROCESSING_JOB_CONFIG_PATH,
+        transform_job_arn_env: str = TRANSFORM_JOB_ARN_ENV,
+    ):
+        """Loads source arn of current job from environment.
+
+        Args:
+            training_job_arn_env (str): The environment key for training job ARN
+                (default: `TRAINING_JOB_ARN`).
+            processing_job_config_path (str): The processing job config path
+                (default: `/opt/ml/config/processingjobconfig.json`).
+            transform_job_arn_env (str): The environment key for transform job ARN
+                (default: `TRANSFORM_JOB_ARN_ENV`).
+
+        Returns:
+            _RunEnvironment: Job data loaded from the environment. None if config does not exist.
+        """
+        if training_job_arn_env in os.environ:
+            environment_type = _EnvironmentType.SageMakerTrainingJob
+            source_arn = os.environ.get(training_job_arn_env)
+            return _RunEnvironment(environment_type, source_arn)
+        if os.path.exists(processing_job_config_path):
+            environment_type = _EnvironmentType.SageMakerProcessingJob
+            source_arn = json.loads(open(processing_job_config_path).read())["ProcessingJobArn"]
+            return _RunEnvironment(environment_type, source_arn)
+        if transform_job_arn_env in os.environ:
+            environment_type = _EnvironmentType.SageMakerTransformJob
+            # TODO: need to update to get source_arn from config file once Transform side ready
+            source_arn = os.environ.get(transform_job_arn_env)
+            return _RunEnvironment(environment_type, source_arn)
+
+        return None
+
+    def get_trial_component(self, sagemaker_session: Session):
+        """Retrieves the trial component from the job in the environment.
+
+        Args:
+            sagemaker_session (sagemaker.core.helper.session_helper.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:
+            _TrialComponent: The trial component created from the job. None if not found.
+        """
+
+        def _get_trial_component():
+            summaries = list(
+                trial_component._TrialComponent.list(
+                    source_arn=self.source_arn.lower(), sagemaker_session=sagemaker_session
+                )
+            )
+            if summaries:
+                summary = summaries[0]
+                return trial_component._TrialComponent.load(
+                    trial_component_name=summary.trial_component_name,
+                    sagemaker_session=sagemaker_session,
+                )
+            return None
+
+        job_tc = None
+        try:
+            job_tc = retry_with_backoff(_get_trial_component, MAX_RETRY_ATTEMPTS)
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.error(
+                "Failed to get trail component in the current environment due to %s", str(ex)
+            )
+        return job_tc