feat(pipeline): Add inference and lineage step types

sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py

@@ -14,6 +14,7 @@
 functions, conditions, properties) and can import from sagemaker.train and sagemaker.serve
 for orchestration purposes.
 """
+
 from __future__ import absolute_import
 
 __version__ = "0.1.0"
@@ -46,8 +47,11 @@
 from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep
 from sagemaker.mlops.workflow.condition_step import ConditionStep
 from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig
+from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep
 from sagemaker.mlops.workflow.fail_step import FailStep
+from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep
 from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput
+from sagemaker.mlops.workflow.lineage_step import LineageStep
 from sagemaker.mlops.workflow.model_step import ModelStep
 from sagemaker.mlops.workflow.monitor_batch_transform_step import MonitorBatchTransformStep
 from sagemaker.mlops.workflow.notebook_job_step import NotebookJobStep
@@ -98,9 +102,13 @@
     "ConditionStep",
     "EMRStep",
     "EMRStepConfig",
+    "EndpointConfigStep",
+    "EndpointStep",
     "FailStep",
+    "InferenceComponentStep",
     "LambdaStep",
     "LambdaOutput",
+    "LineageStep",
     "ModelStep",
     "MonitorBatchTransformStep",
     "NotebookJobStep",

sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.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.
+"""Client-side validation for pipeline step ``arguments`` blocks.
+
+Validates the **top-level keys** of a step's ``arguments`` dict against
+the corresponding public AWS API input shape from botocore, and rejects
+fields that SageMaker Pipelines is known not to support. This fails fast
+at step construction with a clear error, instead of a server-side parse
+failure at ``CreatePipeline`` time.
+
+Values are intentionally not validated: they may be pipeline variables
+(parameter references, step property references, ``Join``/``JsonGet``
+expressions) that only resolve at pipeline compile or execution time.
+
+If the installed botocore release does not know the target operation
+(for example, a very old botocore release), shape
+validation is skipped and the service remains the authority.
+"""
+
+from __future__ import absolute_import
+
+import logging
+from typing import Any, Dict, FrozenSet, Optional, Sequence, Tuple
+
+import botocore.session
+from botocore.exceptions import UnknownServiceError
+from botocore.model import OperationNotFoundError
+
+logger = logging.getLogger(__name__)
+
+# Cache of (service, operation) -> allowed top-level keys.
+# ``None`` means botocore does not know the operation; skip shape checks.
+_SHAPE_CACHE: Dict[Tuple[str, str], Optional[FrozenSet[str]]] = {}
+
+
+def _allowed_top_level_keys(service_name: str, operation_name: str) -> Optional[FrozenSet[str]]:
+    """Return the allowed top-level keys for an operation input shape.
+
+    Args:
+        service_name (str): botocore service name (e.g. ``sagemaker``).
+        operation_name (str): operation name (e.g. ``CreateEndpointConfig``).
+
+    Returns:
+        The allowed key set, or ``None`` if the installed botocore does
+        not know the operation (validation should then be skipped).
+    """
+    cache_key = (service_name, operation_name)
+    if cache_key not in _SHAPE_CACHE:
+        try:
+            session = botocore.session.get_session()
+            service_model = session.get_service_model(service_name)
+            operation_model = service_model.operation_model(operation_name)
+            members = operation_model.input_shape.members.keys()
+            _SHAPE_CACHE[cache_key] = frozenset(members)
+        except (UnknownServiceError, OperationNotFoundError):
+            logger.warning(
+                "Installed botocore does not know %s.%s; skipping "
+                "client-side argument shape validation for this step.",
+                service_name,
+                operation_name,
+            )
+            _SHAPE_CACHE[cache_key] = None
+    return _SHAPE_CACHE[cache_key]
+
+
+def validate_step_arguments(
+    step_class_name: str,
+    arguments: Dict[str, Any],
+    service_name: str,
+    operation_name: str,
+    unsupported_fields: Sequence[str] = (),
+) -> None:
+    """Validate the top-level keys of a step ``arguments`` dict.
+
+    Args:
+        step_class_name (str): Step class name, used in error messages.
+        arguments (Dict[str, Any]): The user-provided ``arguments`` dict.
+        service_name (str): botocore service name of the wrapped API.
+        operation_name (str): Operation whose input shape defines the
+            allowed top-level fields.
+        unsupported_fields (Sequence[str]): Fields that exist in the
+            public API shape but are rejected by SageMaker Pipelines.
+
+    Raises:
+        ValueError: If ``arguments`` is not a non-empty dict with string
+            keys, contains an unsupported field, or contains a key that
+            is not part of the operation's input shape.
+    """
+    if arguments is None:
+        raise ValueError(f"arguments is required for {step_class_name}.")
+    if not isinstance(arguments, dict) or not arguments:
+        raise ValueError(f"{step_class_name}: arguments must be a non-empty dict.")
+    non_string_keys = [key for key in arguments if not isinstance(key, str)]
+    if non_string_keys:
+        raise ValueError(
+            f"{step_class_name}: argument keys must be strings; got {non_string_keys!r}."
+        )
+    rejected = sorted(field for field in unsupported_fields if field in arguments)
+    if rejected:
+        raise ValueError(
+            f"{step_class_name}: field(s) {rejected} are not supported by "
+            "SageMaker Pipelines and would be rejected at pipeline creation "
+            "time. Remove them from arguments."
+        )
+    allowed = _allowed_top_level_keys(service_name, operation_name)
+    if allowed is None:
+        return
+    unknown = sorted(set(arguments) - allowed)
+    if unknown:
+        raise ValueError(
+            f"{step_class_name}: unknown argument field(s) {unknown}. "
+            f"Allowed top-level fields (from {service_name}.{operation_name}): "
+            f"{sorted(allowed)}."
+        )

sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py

@@ -0,0 +1,223 @@
+# 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.
+"""Step definitions for SageMaker Endpoint deployment in Pipelines.
+
+Design note: the pipeline service models each step's
+``Arguments`` block as an opaque structure validated against
+the underlying SageMaker request model (``CreateEndpointConfigInput``
+or ``CreateEndpointInput``) minus a small exclusion set. This SDK
+validates the **top-level keys** of the ``arguments`` dict against the
+public ``CreateEndpointConfig``/``CreateEndpoint`` API input shape at
+construction time (values are not validated -- they may be pipeline
+variables) and forwards the dict to the service, which remains the
+authority on full schema validation.
+
+Excluded fields (the pipeline service rejects the pipeline if present):
+
+* ``EndpointConfig``: ``DataCaptureConfig``, ``ExplainerConfig``
+* ``Endpoint``: ``DeploymentConfig``
+"""
+
+from __future__ import absolute_import
+
+from typing import Any, Dict, List, Optional, Union
+
+from sagemaker.core.helper.pipeline_variable import RequestType
+from sagemaker.core.workflow.properties import Properties
+
+from sagemaker.mlops.workflow._argument_validation import validate_step_arguments
+from sagemaker.mlops.workflow.retry import RetryPolicy
+from sagemaker.mlops.workflow.step_collections import StepCollection
+from sagemaker.mlops.workflow.steps import (
+    CacheConfig,
+    ConfigurableRetryStep,
+    Step,
+    StepTypeEnum,
+)
+
+
+class EndpointConfigStep(ConfigurableRetryStep):
+    """Creates a SageMaker EndpointConfig within a pipeline.
+
+    Wraps the SageMaker ``CreateEndpointConfig`` API. The ``arguments``
+    dict is passed through to the service; it accepts any field of
+    ``CreateEndpointConfigInput`` **except** ``DataCaptureConfig`` and
+    ``ExplainerConfig``, which are rejected by the pipeline service.
+
+    Per the pipeline service's step contract, ``EndpointConfig`` is structurally
+    cacheable (``cache_config``) and retryable (``retry_policies``).
+    """
+
+    def __init__(
+        self,
+        name: str,
+        arguments: Dict[str, Any],
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
+        cache_config: Optional[CacheConfig] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
+    ):
+        """Construct an ``EndpointConfigStep``.
+
+        Args:
+            name (str): The name of the step.
+            arguments (Dict[str, Any]): The ``Arguments`` block for the
+                ``CreateEndpointConfig`` call. Required fields:
+                ``EndpointConfigName``, ``ProductionVariants``. Optional
+                fields include ``KmsKeyId``, ``AsyncInferenceConfig``,
+                ``ShadowProductionVariants``, ``ExecutionRoleArn``,
+                ``VpcConfig``, ``EnableNetworkIsolation``,
+                ``MetricsConfig``. Values may be pipeline variables
+                (parameter references, step property references) — the
+                pipeline compiler resolves them at definition time.
+                Do not include ``DataCaptureConfig`` or ``ExplainerConfig``
+                (the pipeline service rejects them).
+            display_name (str): Optional display name.
+            description (str): Optional description.
+            depends_on (List[Union[str, Step, StepCollection]]): Optional
+                explicit step dependencies.
+            cache_config (CacheConfig): Optional cache configuration.
+            retry_policies (List[RetryPolicy]): Optional retry policies.
+        """
+        super().__init__(
+            name=name,
+            step_type=StepTypeEnum.ENDPOINT_CONFIG,
+            display_name=display_name,
+            description=description,
+            depends_on=depends_on,
+            retry_policies=retry_policies,
+        )
+        if arguments is None:
+            raise ValueError("arguments is required for EndpointConfigStep.")
+        validate_step_arguments(
+            "EndpointConfigStep",
+            arguments,
+            service_name="sagemaker",
+            operation_name="CreateEndpointConfig",
+            unsupported_fields=("DataCaptureConfig", "ExplainerConfig"),
+        )
+        self._arguments = arguments
+        self.cache_config = cache_config
+        self._properties = Properties(
+            step_name=name, step=self, shape_name="DescribeEndpointConfigOutput"
+        )
+
+    @property
+    def arguments(self) -> RequestType:
+        """The ``Arguments`` block for the ``CreateEndpointConfig`` call."""
+        validate_step_arguments(
+            "EndpointConfigStep",
+            self._arguments,
+            service_name="sagemaker",
+            operation_name="CreateEndpointConfig",
+            unsupported_fields=("DataCaptureConfig", "ExplainerConfig"),
+        )
+        return self._arguments
+
+    @property
+    def properties(self):
+        """A ``Properties`` object shaped like ``DescribeEndpointConfigOutput``."""
+        return self._properties
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        request_dict = super().to_request()
+        if self.cache_config:
+            request_dict.update(self.cache_config.config)
+        return request_dict
+
+
+class EndpointStep(Step):
+    """Creates or updates a SageMaker Endpoint within a pipeline.
+
+    Wraps the SageMaker ``CreateEndpoint``/``UpdateEndpoint`` API — the
+    pipeline chooses create-vs-update based on endpoint existence. The
+    ``arguments`` dict is passed through to the service; it accepts any
+    field of ``CreateEndpointInput`` **except** ``DeploymentConfig``,
+    which is rejected by the pipeline service.
+
+    Per the pipeline service's step contract, ``Endpoint`` is structurally cacheable
+    but not retryable at the pipeline level.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        arguments: Dict[str, Any],
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
+        cache_config: Optional[CacheConfig] = None,
+    ):
+        """Construct an ``EndpointStep``.
+
+        Args:
+            name (str): The name of the step.
+            arguments (Dict[str, Any]): The ``Arguments`` block for the
+                ``CreateEndpoint`` / ``UpdateEndpoint`` call. Required
+                fields: ``EndpointName``, ``EndpointConfigName``. Optional
+                fields: ``GraphConfigName``, ``DeletionCondition``.
+                Values may be pipeline variables. Do not include
+                ``DeploymentConfig`` (the pipeline service rejects it).
+            display_name (str): Optional display name.
+            description (str): Optional description.
+            depends_on (List[Union[str, Step, StepCollection]]): Optional
+                explicit step dependencies.
+            cache_config (CacheConfig): Optional cache configuration.
+        """
+        super().__init__(
+            name=name,
+            display_name=display_name,
+            description=description,
+            step_type=StepTypeEnum.ENDPOINT,
+            depends_on=depends_on,
+        )
+        if arguments is None:
+            raise ValueError("arguments is required for EndpointStep.")
+        validate_step_arguments(
+            "EndpointStep",
+            arguments,
+            service_name="sagemaker",
+            operation_name="CreateEndpoint",
+            unsupported_fields=("DeploymentConfig",),
+        )
+        self._arguments = arguments
+        self.cache_config = cache_config
+        self._properties = Properties(
+            step_name=name, step=self, shape_name="DescribeEndpointOutput"
+        )
+
+    @property
+    def arguments(self) -> RequestType:
+        """The ``Arguments`` block for the ``CreateEndpoint``/``UpdateEndpoint`` call."""
+        validate_step_arguments(
+            "EndpointStep",
+            self._arguments,
+            service_name="sagemaker",
+            operation_name="CreateEndpoint",
+            unsupported_fields=("DeploymentConfig",),
+        )
+        return self._arguments
+
+    @property
+    def properties(self):
+        """A ``Properties`` object shaped like ``DescribeEndpointOutput``."""
+        return self._properties
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        request_dict = super().to_request()
+        if self.cache_config:
+            request_dict.update(self.cache_config.config)
+        return request_dict