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

{sagemaker_core/main → sagemaker/core/utils}/exceptions.py

@@ -91,21 +91,21 @@ class TimeoutExceededError(WaiterError):
         super().__init__(resource_type=resource_type, status=status, reason=reason)
 
 
-### Default Configs Errors
-class DefaultConfigsError(SageMakerCoreError):
-    """Raised when an error occurs in the Default Configs"""
+### Intelligent Defaults Errors
+class IntelligentDefaultsError(SageMakerCoreError):
+    """Raised when an error occurs in the Intelligent Defaults"""
 
-    fmt = "An error occurred while loading Default Configs. {message}"
+    fmt = "An error occurred while loading Intelligent Default. {message}"
 
     def __init__(self, message="", **kwargs):
-        """Initialize an DefaultConfigsError exception.
+        """Initialize an IntelligentDefaultsError exception.
         Args:
             message (str): A message describing the error.
         """
         super().__init__(message=message, **kwargs)
 
 
-class LocalConfigNotFoundError(DefaultConfigsError):
+class LocalConfigNotFoundError(IntelligentDefaultsError):
     """Raised when a configuration file is not found in local file system"""
 
     fmt = "Failed to load configuration file from location: {file_path}. {message}"
@@ -119,7 +119,7 @@ class LocalConfigNotFoundError(DefaultConfigsError):
         super().__init__(file_path=file_path, message=message)
 
 
-class S3ConfigNotFoundError(DefaultConfigsError):
+class S3ConfigNotFoundError(IntelligentDefaultsError):
     """Raised when a configuration file is not found in S3"""
 
     fmt = "Failed to load configuration file from S3 location: {s3_uri}. {message}"
@@ -133,7 +133,7 @@ class S3ConfigNotFoundError(DefaultConfigsError):
         super().__init__(s3_uri=s3_uri, message=message)
 
 
-class ConfigSchemaValidationError(DefaultConfigsError, ValidationError):
+class ConfigSchemaValidationError(IntelligentDefaultsError, ValidationError):
     """Raised when a configuration file does not adhere to the schema"""
 
     fmt = "Failed to validate configuration file from location: {file_path}. {message}"

sagemaker_core/main/default_configs_helper.py → sagemaker/core/utils/intelligent_defaults_helper.py

@@ -11,7 +11,6 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-import logging
 
 import os
 import jsonschema
@@ -25,14 +24,14 @@ from platformdirs import site_config_dir, user_config_dir
 
 from botocore.utils import merge_dicts
 from six.moves.urllib.parse import urlparse
-from sagemaker_core.main.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA
-from sagemaker_core.main.exceptions import (
+from sagemaker.core.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA
+from sagemaker.core.utils.exceptions import (
     LocalConfigNotFoundError,
     S3ConfigNotFoundError,
-    DefaultConfigsError,
+    IntelligentDefaultsError,
     ConfigSchemaValidationError,
 )
-from sagemaker_core.main.utils import get_textual_rich_logger
+from sagemaker.core.utils.utils import get_textual_rich_logger
 
 logger = get_textual_rich_logger(__name__)
 
@@ -116,7 +115,7 @@ def _load_config_from_s3(s3_uri, s3_resource_for_config) -> dict:
         boto_session = boto3.DEFAULT_SESSION or boto3.Session()
         boto_region_name = boto_session.region_name
         if boto_region_name is None:
-            raise DefaultConfigsError(
+            raise IntelligentDefaultsError(
                 message=(
                     "Valid region is not provided in the Boto3 session."
                     + "Setup local AWS configuration with a valid region supported by SageMaker."

{sagemaker_core/main → sagemaker/core/utils}/logs.py

@@ -1,11 +1,10 @@
-import boto3
 import botocore
 
 from boto3.session import Session
 import botocore.client
 from botocore.config import Config
 from typing import Generator, Tuple, List
-from sagemaker_core.main.utils import SingletonMeta
+from sagemaker.core.utils.utils import SingletonMeta
 
 
 class CloudWatchLogsClient(metaclass=SingletonMeta):

{sagemaker_core/main → sagemaker/core/utils}/utils.py

@@ -26,9 +26,9 @@ from rich.style import Style
 from rich.theme import Theme
 from rich.traceback import install
 from typing import Any, Dict, List, TypeVar, Generic, Type
-from sagemaker_core.main.code_injection.codec import transform
-from sagemaker_core.main.code_injection.constants import Color
-from sagemaker_core.main.user_agent import get_user_agent_extra_suffix
+from sagemaker.core.utils.code_injection.codec import transform
+from sagemaker.core.utils.code_injection.constants import Color
+from sagemaker.core.utils.user_agent import get_user_agent_extra_suffix
 
 
 def add_indent(text, num_spaces=4):
@@ -182,11 +182,8 @@ def get_textual_rich_logger(name: str, log_level: str = "INFO") -> logging.Logge
     """
     enable_textual_rich_console_and_traceback()
     handler = get_rich_handler()
+    logging.basicConfig(level=getattr(logging, log_level), handlers=[handler])
     logger = logging.getLogger(name)
-    for handler in logger.handlers:
-        logger.removeHandler(handler)
-    logger.addHandler(handler)
-    logger.setLevel(getattr(logging, log_level))
 
     return logger
 
@@ -336,7 +333,7 @@ class SageMakerClient(metaclass=SingletonMeta):
         Creates a boto3 client using the provided session, region, and service.
         """
         if session is None:
-            logger.warning("No boto3 session provided. Creating a new session.")
+            logger.debug("No boto3 session provided. Creating a new session.")
             session = Session(region_name=region_name)
 
         if region_name is None:
@@ -344,13 +341,25 @@ class SageMakerClient(metaclass=SingletonMeta):
             region_name = session.region_name
 
         if config is None:
-            logger.warning("No config provided. Using default config.")
+            logger.debug("No config provided. Using default config.")
             config = Config(retries={"max_attempts": 10, "mode": "standard"})
 
         self.config = Config(user_agent_extra=get_user_agent_extra_suffix())
         self.session = session
         self.region_name = region_name
-        self.sagemaker_client = session.client("sagemaker", region_name, config=self.config)
+        # Read region from environment variable, default to us-west-2
+        import os
+        env_region = os.environ.get('SAGEMAKER_REGION', region_name)
+        env_stage = os.environ.get('SAGEMAKER_STAGE', 'prod')  # default to gamma
+        logger.info(f"Runs on sagemaker {env_stage}, region:{env_region}")
+
+
+        self.sagemaker_client = session.client(
+            "sagemaker",
+            region_name=env_region,
+            config=self.config,
+        )
+        
         self.sagemaker_runtime_client = session.client(
             "sagemaker-runtime", region_name, config=self.config
         )
@@ -387,7 +396,6 @@ class ResourceIterator(Generic[T]):
         list_method: str,
         list_method_kwargs: dict = {},
         custom_key_mapping: dict = None,
-        extract_name_mapping: dict = None,
     ):
         """Initialize a ResourceIterator object
 
@@ -399,7 +407,6 @@ class ResourceIterator(Generic[T]):
             list_method (str): The list method string used to make list calls to the client.
             list_method_kwargs (dict, optional): The kwargs used to make list method calls. Defaults to {}.
             custom_key_mapping (dict, optional): The custom key mapping used to map keys from summary object to those expected from resource object during initialization. Defaults to None.
-            extract_name_mapping (dict, optional): The extract name mapping used to extract names from arn in summary object and map to those expected from resource object during initialization. Defaults to None.
         """
         self.summaries_key = summaries_key
         self.summary_name = summary_name
@@ -407,7 +414,6 @@ class ResourceIterator(Generic[T]):
         self.list_method = list_method
         self.list_method_kwargs = list_method_kwargs
         self.custom_key_mapping = custom_key_mapping
-        self.extract_name_mapping = extract_name_mapping
 
         self.resource_cls = resource_cls
         self.index = 0
@@ -436,12 +442,6 @@ class ResourceIterator(Generic[T]):
                 if self.custom_key_mapping:
                     init_data = {self.custom_key_mapping.get(k, k): v for k, v in init_data.items()}
 
-                # Extract name from arn. Currently implemented for HubContent and ImageVersion
-                if self.extract_name_mapping:
-                    for arn, target in self.extract_name_mapping.items():
-                        name = init_data[arn].split(target[0])[1].split("/")[0]
-                        init_data.update({target[1]: name})
-
                 # Filter out the fields that are not in the resource class
                 fields = self.resource_cls.__annotations__
                 init_data = {k: v for k, v in init_data.items() if k in fields}
@@ -491,8 +491,13 @@ def serialize(value: Any) -> Any:
     Returns:
         Any: The serialized object
     """
-    if value is None or isinstance(value, Unassigned):
+    from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+    if value is None or isinstance(value, type(Unassigned())):
         return None
+    elif isinstance(value, PipelineVariable):
+        # Return PipelineVariables as-is (Join, ExecutionVariables, etc.)
+        return value
     elif isinstance(value, Dict):
         # if the value is a dict, use _serialize_dict() to serialize it recursively
         return _serialize_dict(value)

sagemaker/core/workflow/__init__.py

@@ -0,0 +1,152 @@
+# 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.
+"""Workflow primitives for SageMaker pipelines.
+
+This module contains the foundation types and primitives used by workflow orchestration.
+These primitives can be used by Train, Serve, and MLOps packages without creating
+circular dependencies.
+
+For pipeline and step orchestration classes (Pipeline, TrainingStep, etc.),
+import from sagemaker.mlops.workflow instead.
+"""
+from __future__ import absolute_import
+
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+
+def is_pipeline_variable(var: object) -> bool:
+    """Check if the variable is a pipeline variable
+
+    Args:
+        var (object): The variable to be verified.
+    Returns:
+         bool: True if it is, False otherwise.
+    """
+    return isinstance(var, PipelineVariable)
+
+
+def is_pipeline_parameter_string(var: object) -> bool:
+    """Check if the variable is a pipeline parameter string
+
+    Args:
+        var (object): The variable to be verified.
+    Returns:
+         bool: True if it is, False otherwise.
+    """
+    from sagemaker.core.workflow.parameters import ParameterString
+
+    return isinstance(var, ParameterString)
+
+
+# Entities
+from sagemaker.core.workflow.entities import (
+    DefaultEnumMeta,
+    Entity,
+)
+
+# Execution Variables
+from sagemaker.core.workflow.execution_variables import (
+    ExecutionVariable,
+    ExecutionVariables,
+)
+
+# Functions
+from sagemaker.core.workflow.functions import (
+    Join,
+    JsonGet,
+)
+
+# Parameters
+from sagemaker.core.workflow.parameters import (
+    Parameter,
+    ParameterBoolean,
+    ParameterFloat,
+    ParameterInteger,
+    ParameterString,
+    ParameterTypeEnum,
+)
+
+# Properties
+from sagemaker.core.workflow.properties import (
+    Properties,
+    PropertiesList,
+    PropertiesMap,
+    PropertyFile,
+)
+
+# Step Outputs (primitive - used by properties)
+from sagemaker.core.workflow.step_outputs import (
+    StepOutput,
+    get_step,
+)
+
+# Conditions (primitive)
+from sagemaker.core.workflow.conditions import (
+    Condition,
+    ConditionComparison,
+    ConditionEquals,
+    ConditionGreaterThan,
+    ConditionGreaterThanOrEqualTo,
+    ConditionIn,
+    ConditionLessThan,
+    ConditionLessThanOrEqualTo,
+    ConditionNot,
+    ConditionOr,
+    ConditionTypeEnum,
+)
+
+# NOTE: Orchestration classes (Pipeline, TrainingStep, ProcessingStep, etc.) have moved
+# to sagemaker.mlops.workflow. Import them from there instead:
+#   from sagemaker.mlops.workflow import Pipeline, TrainingStep, ProcessingStep
+
+__all__ = [
+    # Helper functions
+    "is_pipeline_variable",
+    "is_pipeline_parameter_string",
+    # Entities (primitives)
+    "DefaultEnumMeta",
+    "Entity",
+    # Execution Variables (primitives)
+    "ExecutionVariable",
+    "ExecutionVariables",
+    # Functions (primitives)
+    "Join",
+    "JsonGet",
+    # Parameters (primitives)
+    "Parameter",
+    "ParameterBoolean",
+    "ParameterFloat",
+    "ParameterInteger",
+    "ParameterString",
+    "ParameterTypeEnum",
+    # Properties (primitives)
+    "Properties",
+    "PropertiesList",
+    "PropertiesMap",
+    "PropertyFile",
+    # Step Outputs (primitives)
+    "StepOutput",
+    "get_step",
+    # Conditions (primitives)
+    "Condition",
+    "ConditionComparison",
+    "ConditionEquals",
+    "ConditionGreaterThan",
+    "ConditionGreaterThanOrEqualTo",
+    "ConditionIn",
+    "ConditionLessThan",
+    "ConditionLessThanOrEqualTo",
+    "ConditionNot",
+    "ConditionOr",
+    "ConditionTypeEnum",
+]

sagemaker/core/workflow/conditions.py

@@ -0,0 +1,313 @@
+# 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.
+"""Conditions for condition steps.
+
+Ideally, some of these comparison conditions would be implemented as "partial classes",
+but use of functools.partial doesn't set correct metadata/type information.
+"""
+from __future__ import absolute_import
+
+import abc
+
+from enum import Enum
+from typing import Dict, List, Union
+
+import attr
+
+from sagemaker.core.workflow import is_pipeline_variable
+from sagemaker.core.workflow.entities import (
+    DefaultEnumMeta,
+    Entity,
+)
+from sagemaker.core.helper.pipeline_variable import (
+    PrimitiveType,
+    RequestType,
+    PipelineVariable,
+)
+from sagemaker.core.workflow.execution_variables import ExecutionVariable
+from sagemaker.core.workflow.parameters import Parameter
+from sagemaker.core.workflow.properties import Properties
+
+# TODO: consider base class for those with an expr method, rather than defining a type here
+from sagemaker.core.workflow.step_outputs import StepOutput
+
+ConditionValueType = Union[ExecutionVariable, Parameter, Properties, StepOutput]
+
+
+class ConditionTypeEnum(Enum, metaclass=DefaultEnumMeta):
+    """Condition type enum."""
+
+    EQ = "Equals"
+    GT = "GreaterThan"
+    GTE = "GreaterThanOrEqualTo"
+    IN = "In"
+    LT = "LessThan"
+    LTE = "LessThanOrEqualTo"
+    NOT = "Not"
+    OR = "Or"
+
+
+@attr.s
+class Condition(Entity):
+    """Abstract Condition entity.
+
+    Attributes:
+        condition_type (ConditionTypeEnum): The type of condition.
+    """
+
+    condition_type: ConditionTypeEnum = attr.ib(factory=ConditionTypeEnum.factory)
+
+    @property
+    @abc.abstractmethod
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+
+
+@attr.s
+class ConditionComparison(Condition):
+    """Generic comparison condition that can be used to derive specific condition comparisons.
+
+    Attributes:
+        left (Union[ConditionValueType, PrimitiveType]): The execution variable, parameter,
+            property, step output or Python primitive value to use in the comparison.
+        right (Union[ConditionValueType, PrimitiveType]): The execution variable,
+            parameter, property, step output or Python primitive value to compare to.
+    """
+
+    left: Union[ConditionValueType, PrimitiveType] = attr.ib(default=None)
+    right: Union[ConditionValueType, PrimitiveType] = attr.ib(default=None)
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        return {
+            "Type": self.condition_type.value,
+            "LeftValue": self.left,
+            "RightValue": self.right,
+        }
+
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        if isinstance(self.left, PipelineVariable):
+            steps.extend(self.left._referenced_steps)
+        if isinstance(self.right, PipelineVariable):
+            steps.extend(self.right._referenced_steps)
+        return steps
+
+
+class ConditionEquals(ConditionComparison):
+    """A condition for equality comparisons."""
+
+    def __init__(
+        self,
+        left: Union[ConditionValueType, PrimitiveType],
+        right: Union[ConditionValueType, PrimitiveType],
+    ):
+        """Construct A condition for equality comparisons.
+
+        Args:
+            left (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property, or Python primitive value to use in the comparison.
+            right (Union[ConditionValueType, PrimitiveType]): The execution
+                variable, parameter, property, or Python primitive value to compare to.
+        """
+
+        super(ConditionEquals, self).__init__(ConditionTypeEnum.EQ, left, right)
+
+
+class ConditionGreaterThan(ConditionComparison):
+    """A condition for greater than comparisons."""
+
+    def __init__(
+        self,
+        left: Union[ConditionValueType, PrimitiveType],
+        right: Union[ConditionValueType, PrimitiveType],
+    ):
+        """Construct an instance of ConditionGreaterThan for greater than comparisons.
+
+        Args:
+            left (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property, or Python primitive value to use in the comparison.
+            right (Union[ConditionValueType, PrimitiveType]): The execution
+                variable, parameter, property, or Python primitive value to compare to.
+        """
+
+        super(ConditionGreaterThan, self).__init__(ConditionTypeEnum.GT, left, right)
+
+
+class ConditionGreaterThanOrEqualTo(ConditionComparison):
+    """A condition for greater than or equal to comparisons."""
+
+    def __init__(
+        self,
+        left: Union[ConditionValueType, PrimitiveType],
+        right: Union[ConditionValueType, PrimitiveType],
+    ):
+        """Construct of ConditionGreaterThanOrEqualTo for greater than or equal to comparisons.
+
+        Args:
+            left (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property, or Python primitive value to use in the comparison.
+            right (Union[ConditionValueType, PrimitiveType]): The execution
+                variable, parameter, property, or Python primitive value to compare to.
+        """
+
+        super(ConditionGreaterThanOrEqualTo, self).__init__(ConditionTypeEnum.GTE, left, right)
+
+
+class ConditionLessThan(ConditionComparison):
+    """A condition for less than comparisons."""
+
+    def __init__(
+        self,
+        left: Union[ConditionValueType, PrimitiveType],
+        right: Union[ConditionValueType, PrimitiveType],
+    ):
+        """Construct an instance of ConditionLessThan for less than comparisons.
+
+        Args:
+            left (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property, or Python primitive value to use in the comparison.
+            right (Union[ConditionValueType, PrimitiveType]): The execution
+                variable, parameter, property, or Python primitive value to compare to.
+        """
+
+        super(ConditionLessThan, self).__init__(ConditionTypeEnum.LT, left, right)
+
+
+class ConditionLessThanOrEqualTo(ConditionComparison):
+    """A condition for less than or equal to comparisons."""
+
+    def __init__(
+        self,
+        left: Union[ConditionValueType, PrimitiveType],
+        right: Union[ConditionValueType, PrimitiveType],
+    ):
+        """Construct ConditionLessThanOrEqualTo for less than or equal to comparisons.
+
+        Args:
+            left (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property, or Python primitive value to use in the comparison.
+            right (Union[ConditionValueType, PrimitiveType]): The execution
+                variable, parameter, property, or Python primitive value to compare to.
+        """
+
+        super(ConditionLessThanOrEqualTo, self).__init__(ConditionTypeEnum.LTE, left, right)
+
+
+class ConditionIn(Condition):
+    """A condition to check membership."""
+
+    def __init__(
+        self,
+        value: Union[ConditionValueType, PrimitiveType],
+        in_values: List[Union[ConditionValueType, PrimitiveType]],
+    ):
+        """Construct a `ConditionIn` condition to check membership.
+
+        Args:
+            value (Union[ConditionValueType, PrimitiveType]): The execution variable,
+                parameter, property or primitive value to check for membership.
+            in_values (List[Union[ConditionValueType, PrimitiveType]]): The list
+                of values to check for membership in.
+        """
+        super(ConditionIn, self).__init__(ConditionTypeEnum.IN)
+        self.value = value
+        self.in_values = in_values
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        return {
+            "Type": self.condition_type.value,
+            "QueryValue": self.value,
+            "Values": self.in_values,
+        }
+
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        if isinstance(self.value, PipelineVariable):
+            steps.extend(self.value._referenced_steps)
+        for in_value in self.in_values:
+            if isinstance(in_value, PipelineVariable):
+                steps.extend(in_value._referenced_steps)
+        return steps
+
+
+class ConditionNot(Condition):
+    """A condition for negating another `Condition`."""
+
+    def __init__(self, expression: Condition):
+        """Construct a `ConditionNot` condition for negating another `Condition`.
+
+        Attributes:
+            expression (Condition): A `Condition` to take the negation of.
+        """
+        super(ConditionNot, self).__init__(ConditionTypeEnum.NOT)
+        self.expression = expression
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        return {"Type": self.condition_type.value, "Condition": self.expression.to_request()}
+
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return self.expression._referenced_steps
+
+
+class ConditionOr(Condition):
+    """A condition for taking the logical OR of a list of `Condition` instances."""
+
+    def __init__(self, conditions: List[Condition] = None):
+        """Construct a `ConditionOr` condition.
+
+        Attributes:
+            conditions (List[Condition]): A list of `Condition` instances to logically OR.
+        """
+        super(ConditionOr, self).__init__(ConditionTypeEnum.OR)
+        self.conditions = conditions or []
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        return {
+            "Type": self.condition_type.value,
+            "Conditions": [condition.to_request() for condition in self.conditions],
+        }
+
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        for condition in self.conditions:
+            steps.extend(condition._referenced_steps)
+        return steps
+
+
+def primitive_or_expr(
+    value: Union[
+        ExecutionVariable, PipelineVariable, PrimitiveType, Parameter, Properties, StepOutput
+    ],
+) -> Union[Dict[str, str], PrimitiveType]:
+    """Provide the expression of the value or return value if it is a primitive.
+
+    Args:
+        value (Union[ConditionValueType, PrimitiveType]): The value to evaluate.
+    Returns:
+        Either the expression of the value or the primitive value.
+    """
+    if is_pipeline_variable(value):
+        return value.expr
+    return value