refactor: Remove Bedrock step types and clean up step test modules (#6223)

* refactor: Remove Bedrock step types and clean up step test modules

Removes the four Bedrock pipeline step classes (BedrockCustomModelStep,
BedrockCustomModelDeploymentStep, BedrockModelImportStep,
BedrockProvisionedModelThroughputStep) and their StepTypeEnum values
pending re-evaluation of SDK support for these step types. They can be
reintroduced later if needed. The remaining step types
(EndpointConfigStep, EndpointStep, InferenceComponentStep, LineageStep)
are unchanged.

Also renames the step test modules to descriptive names
(test_inference_lineage_steps.py, test_lineage_step.py) and removes the
now-unused camelCase-to-PascalCase member-name conversion from the
argument validation helper (it existed only for the Bedrock APIs).
---
X-AI-Prompt: Follow-up on the merged v2 backport: remove Bedrock step
types and internal-codename file names
X-AI-Tool: kiro-cli

* refactor(pipeline): Align inference and lineage steps with v3

Brings the four step types from #6153 onto the same design as v3, so both
majors present one API. These steps are not in use yet, so the change is
made directly rather than with a deprecation path.

Steps now take step_args captured under a PipelineSession instead of a raw
arguments dict, following the convention used by TrainingStep and
ModelStep. Capture goes through one seam:
PipelineSession._intercept_create_request returns the captured context, so
a producer returns its result directly.

Producers, all now routed through that seam:
- Session.create_endpoint_config, which creates only the endpoint
  configuration. It gained an optional production_variants argument,
  appended last so existing callers are unaffected, because an
  inference-component endpoint needs a variant with no model name that the
  single-model form cannot express.
- Session.create_endpoint and Session.create_inference_component.
- Action.create(), Artifact.create() and Context.create() for lineage,
  whose requests Record._invoke_api now routes through the same seam.

LineageStep takes actions, artifacts, contexts and associations. A step
creates a batch of entities and its associations reference them by name and
type, or by ARN for an entity created elsewhere. This replaces the
cross-step Steps.<Name>.ActionArns['<name>'] form, which could not
resolve: the service matches an association endpoint against the entities
created by the same step, so a reference to another step's entity fails
the execution.

Two fixes fall out of the seam change:

PipelineModel.register() tested the intercept result with
`"ModelPackageArn" in model_package`, which raises TypeError once that
result is a step arguments object rather than None. It now checks for a
dict first.

The lineage unit tests used a bare Mock session, whose auto-mocked
_intercept_create_request swallowed the request so the client was never
called. The double now invokes the callable, as the real base Session
does.

_argument_validation.py is deleted. With the Bedrock steps removed by the
previous commit and the four remaining steps on step_args, nothing uses
it.

Testing: 39 unit tests for the four steps. The lineage, step, pipeline,
pipeline session and pipeline model suites were run against the previous
commit and this one with identical results. The integration tests are
ported but were not run against the service on this branch; the equivalent
v3 tests pass end to end.
---
X-AI-Prompt: Re-port v2 onto the new design
X-AI-Tool: kiro-cli

* fix(pipeline): Let create_endpoint_config build an IC endpoint config

Found by running the deployment integration test against SageMaker.

An inference-component endpoint config carries a production variant with no
model name, and therefore needs ExecutionRoleArn in place of a model.
create_endpoint_config had no way to supply it, so the test failed with
"unexpected keyword argument 'role'". A role argument is added, appended
last so existing callers are unaffected, applying the same rule
endpoint_from_production_variants already uses: when no variant names a
model, expand the role and set ExecutionRoleArn.

The KMS check then failed too, because instance_supports_kms() was called
with the instance_type that explicit variants leave as None. It now derives
KMS support from the variants in that case, as v3 does, and keeps using
instance_type when one is given.

Also moves the deployment test's resource creation inside its try, so a
failure between creating the model and starting the pipeline still reaches
cleanup. The two failures above leaked a model and an S3 object each,
because both were created before the try began.

Testing: both integration tests now pass against SageMaker -- lineage in
11.70s and the chained EndpointConfig -> Endpoint -> InferenceComponent
pipeline in 396.23s, with the endpoint and component observed InService and
no resources left behind. The 25 existing create_endpoint_config and
endpoint_from_production_variants unit tests still pass, confirming the
single-model path is unchanged.
---
X-AI-Prompt: Run the v2 integ tests against the service, also make sure you
are removing the old code and adding the new code
X-AI-Tool: kiro-cli

* fix(tests): Use the shared session fixtures in the new integ tests

Both new integ tests defined local sagemaker_session and pipeline_session
fixtures returning a bare Session() / PipelineSession(), shadowing the
shared ones. CI runs the integ suite with AWS_DEFAULT_REGION unset, so a
bare session has no region and both tests errored at setup with "Must setup
local AWS configuration with a region supported by SageMaker". Reported by
@jam-jee on the PR.

The local role fixture went too: the workflow conftest already provides one.
The shared fixtures build on a boto3.Session carrying an explicit region,
which is exactly what was missing.

My earlier runs passed only because the harness exported
AWS_DEFAULT_REGION, which hid this. Both tests have now been re-run with
that variable unset, matching CI: lineage in 11.83s and the chained
EndpointConfig -> Endpoint -> InferenceComponent pipeline in 411.04s, with
no resources left behind.
---
X-AI-Prompt: did you read the comment on the pr?
X-AI-Tool: kiro-cli

---------

Co-authored-by: Rishabh Devnani <rdevnani@amazon.com>

Author: Rishabh0255

src/sagemaker/apiutils/_base_types.py

@@ -219,10 +219,27 @@ def with_boto(self, boto_dict):
         )
         return self
 
+    # Lineage entity creates that a pipeline step can compose rather than execute.
+    # Requests for these are routed through the session's
+    # ``_intercept_create_request`` seam, so a ``PipelineSession`` captures them as
+    # step arguments while a plain ``Session`` calls the service. Used by
+    # ``sagemaker.workflow.lineage_step.LineageStep``.
+    _PIPELINE_CAPTURABLE_METHODS = frozenset(
+        {"create_action", "create_artifact", "create_context", "add_association"}
+    )
+
     def _invoke_api(self, boto_method, boto_method_members):
         """Invoke a SageMaker API."""
         api_values = {k: v for k, v in vars(self).items() if k in boto_method_members}
         api_kwargs = self.to_boto(api_values)
-        api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method)
-        api_boto_response = api_method(**api_kwargs)
-        return self.with_boto(api_boto_response)
+
+        def submit(request):
+            api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method)
+            return self.with_boto(api_method(**request))
+
+        if boto_method in self._PIPELINE_CAPTURABLE_METHODS:
+            return self.sagemaker_session._intercept_create_request(
+                api_kwargs, submit, boto_method
+            )
+
+        return submit(api_kwargs)

src/sagemaker/pipeline.py

@@ -485,7 +485,7 @@ def register(
             **model_pkg_args
         )
 
-        if model_package is not None and "ModelPackageArn" in model_package:
+        if isinstance(model_package, dict) and "ModelPackageArn" in model_package:
             return ModelPackage(
                 role=self.role,
                 model_package_arn=model_package.get("ModelPackageArn"),

src/sagemaker/session.py

@@ -4857,9 +4857,9 @@ def describe_model(self, name):
     def create_endpoint_config(
         self,
         name,
-        model_name,
-        initial_instance_count,
-        instance_type,
+        model_name=None,
+        initial_instance_count=None,
+        instance_type=None,
         accelerator_type=None,
         tags=None,
         kms_key=None,
@@ -4872,6 +4872,8 @@ def create_endpoint_config(
         serverless_inference_config_dict=None,
         routing_config: Optional[Dict[str, Any]] = None,
         inference_ami_version: Optional[str] = None,
+        production_variants: Optional[List[Dict[str, Any]]] = None,
+        role: Optional[str] = None,
     ):
         """Create an Amazon SageMaker endpoint configuration.
 
@@ -4933,6 +4935,16 @@ def create_endpoint_config(
              Specifies an option from a collection of preconfigured
              Amazon Machine Image (AMI) images. For a full list of options, see:
              https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ProductionVariant.html
+            production_variants (Optional[List[Dict[str, Any]]]): Explicit
+                production variants to use as-is. When given, ``model_name``,
+                ``initial_instance_count`` and ``instance_type`` are ignored and
+                no variant is built. Use this for a variant a single-model
+                configuration cannot express, such as an inference-component
+                endpoint whose variant carries no model name.
+            role (Optional[str]): An AWS IAM role, name or full ARN. Required
+                when no production variant names a model, as with an
+                inference-component endpoint: the endpoint config then carries
+                ``ExecutionRoleArn`` in place of a model.
 
         Example:
             >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}]
@@ -4942,24 +4954,25 @@ def create_endpoint_config(
             .Client.add_tags
 
         Returns:
-            str: Name of the endpoint point configuration created.
+            str: Name of the endpoint point configuration created. Under a
+            ``PipelineSession`` the captured step arguments are returned instead
+            and no service call is made.
         """
-        logger.info("Creating endpoint-config with name %s", name)
-
         tags = format_tags(tags) or []
-        provided_production_variant = production_variant(
-            model_name,
-            instance_type,
-            initial_instance_count,
-            accelerator_type=accelerator_type,
-            serverless_inference_config=serverless_inference_config_dict,
-            volume_size=volume_size,
-            model_data_download_timeout=model_data_download_timeout,
-            container_startup_health_check_timeout=container_startup_health_check_timeout,
-            routing_config=routing_config,
-            inference_ami_version=inference_ami_version,
-        )
-        production_variants = [provided_production_variant]
+        if production_variants is None:
+            provided_production_variant = production_variant(
+                model_name,
+                instance_type,
+                initial_instance_count,
+                accelerator_type=accelerator_type,
+                serverless_inference_config=serverless_inference_config_dict,
+                volume_size=volume_size,
+                model_data_download_timeout=model_data_download_timeout,
+                container_startup_health_check_timeout=container_startup_health_check_timeout,
+                routing_config=routing_config,
+                inference_ami_version=inference_ami_version,
+            )
+            production_variants = [provided_production_variant]
         # Currently we just inject CoreDumpConfig.KmsKeyId from the config for production variant.
         # But if that parameter is injected, then CoreDumpConfig.DestinationS3Uri needs to be
         # present.
@@ -4975,17 +4988,43 @@ def create_endpoint_config(
             "ProductionVariants": production_variants,
         }
 
+        role = resolve_value_from_config(
+            role,
+            ENDPOINT_CONFIG_EXECUTION_ROLE_ARN_PATH,
+            sagemaker_session=self,
+        )
+        # For an Amazon SageMaker inference-component based endpoint, no Model name is
+        # passed during endpoint creation. ExecutionRoleArn is needed in the endpoint
+        # config instead, so that the Endpoint can be created.
+        model_names = [pv["ModelName"] for pv in production_variants if "ModelName" in pv]
+        if len(model_names) == 0:
+            # The SDK allows deploying with a role name rather than a full ARN, so
+            # expand it here.
+            role = self.expand_role(role)
+            request["ExecutionRoleArn"] = role
+
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(
             tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS)
         )
         if tags is not None:
             request["Tags"] = tags
+        # When explicit variants are supplied there is no single instance_type to
+        # consult, so derive KMS support from the variants themselves, as
+        # endpoint_from_production_variants does.
+        if instance_type is not None:
+            supports_kms = instance_supports_kms(instance_type)
+        else:
+            supports_kms = any(
+                instance_supports_kms(pv["InstanceType"])
+                for pv in production_variants
+                if "InstanceType" in pv
+            )
         kms_key = (
             resolve_value_from_config(
                 kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
             )
-            if instance_supports_kms(instance_type)
+            if supports_kms
             else kms_key
         )
         if kms_key is not None:
@@ -5008,8 +5047,14 @@ def create_endpoint_config(
         if explainer_config_dict is not None:
             request["ExplainerConfig"] = explainer_config_dict
 
-        self.sagemaker_client.create_endpoint_config(**request)
-        return name
+        def submit(request):
+            logger.info("Creating endpoint-config with name %s", name)
+            self.sagemaker_client.create_endpoint_config(**request)
+            return name
+
+        return self._intercept_create_request(
+            request, submit, self.create_endpoint_config.__name__
+        )
 
     def create_endpoint_config_from_existing(
         self,
@@ -5169,29 +5214,39 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live
                 (default: None).
 
         Returns:
-            str: Name of the Amazon SageMaker ``Endpoint`` created.
+            str: Name of the Amazon SageMaker ``Endpoint`` created. Under a
+            ``PipelineSession`` the captured step arguments are returned instead
+            and no service call is made.
 
         Raises:
             botocore.exceptions.ClientError: If Sagemaker throws an exception while creating
             endpoint.
         """
-        logger.info("Creating endpoint with name %s", endpoint_name)
-
         tags = format_tags(tags) or []
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(
             tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT, TAGS)
         )
-        try:
-            res = self.sagemaker_client.create_endpoint(
-                EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags
-            )
+        create_endpoint_request = {
+            "EndpointName": endpoint_name,
+            "EndpointConfigName": config_name,
+            "Tags": tags,
+        }
+
+        def submit(request):
+            logger.info("Creating endpoint with name %s", endpoint_name)
+            res = self.sagemaker_client.create_endpoint(**request)
             if res:
                 self.endpoint_arn = res["EndpointArn"]
 
             if wait:
                 self.wait_for_endpoint(endpoint_name, live_logging=live_logging)
             return endpoint_name
+
+        try:
+            return self._intercept_create_request(
+                create_endpoint_request, submit, self.create_endpoint.__name__
+            )
         except Exception as e:
             troubleshooting = (
                 "https://docs.aws.amazon.com/sagemaker/latest/dg/"
@@ -5361,13 +5416,9 @@ def create_inference_component(
 
         Returns:
             str: Name of the Amazon SageMaker ``InferenceComponent`` if created.
+            Under a ``PipelineSession`` the captured step arguments are returned
+            instead and no service call is made.
         """
-        LOGGER.info(
-            "Creating inference component with name %s for endpoint %s",
-            inference_component_name,
-            endpoint_name,
-        )
-
         if runtime_config is None:
             runtime_config = {"CopyCount": 1}
 
@@ -5387,10 +5438,20 @@ def create_inference_component(
         if tags and len(tags) != 0:
             request["Tags"] = tags
 
-        self.sagemaker_client.create_inference_component(**request)
-        if wait:
-            self.wait_for_inference_component(inference_component_name)
-        return inference_component_name
+        def submit(request):
+            LOGGER.info(
+                "Creating inference component with name %s for endpoint %s",
+                inference_component_name,
+                endpoint_name,
+            )
+            self.sagemaker_client.create_inference_component(**request)
+            if wait:
+                self.wait_for_inference_component(inference_component_name)
+            return inference_component_name
+
+        return self._intercept_create_request(
+            request, submit, self.create_inference_component.__name__
+        )
 
     def wait_for_inference_component(self, inference_component_name, poll=20):
         """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete.