feat: add JobStep and shared pipeline capture for CreateJob producers

sagemaker-core/src/sagemaker/core/workflow/job_config_document.py

@@ -0,0 +1,94 @@
+# 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.
+"""Encode a ``CreateJob`` ``JobConfigDocument`` that may carry pipeline variables.
+
+Unlike the other create APIs, ``CreateJob`` flattens the job configuration into
+one JSON **string** member, so nothing inside it sits at a position the pipeline
+definition serializer can rewrite, and ``json.dumps`` on a ``PipelineVariable``
+raises ``TypeError``. The encoder therefore emits the JSON text as a ``Join`` of
+literal fragments and the variables, which the pipeline service resolves during
+execution like any other variable leaf.
+
+Limitations, both inherent to splicing values into finished JSON text:
+
+* A variable is spliced in as raw text, so a value resolving to text containing
+  a double quote or a backslash produces a malformed document. Pipeline
+  variables resolve to ARNs, S3 URIs and identifiers in practice.
+* A variable always lands in a JSON **string** position; one standing in for a
+  number or boolean reaches the service quoted. Parameterise string-valued
+  fields only.
+"""
+from __future__ import absolute_import
+
+import json
+import re
+import uuid
+from typing import List, Union
+
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+
+def convert_job_config_document_to_string(job_config) -> Union[str, PipelineVariable]:
+    """Convert a job config dict to the string-typed ``JobConfigDocument`` value.
+
+    Args:
+        job_config (Dict[str, Any]): The job configuration. May hold
+            ``PipelineVariable`` values at any depth.
+
+    Returns:
+        Union[str, PipelineVariable]: ``json.dumps(job_config)`` when the config
+        holds no pipeline variable. Otherwise a ``Join`` over the JSON text and
+        those variables -- not yet a string, but the expression the pipeline
+        service resolves to the document string during execution.
+    """
+    # Import locally: sagemaker.core.workflow.functions imports from this package's
+    # entities, and a module-level import here would be a cycle.
+    from sagemaker.core.workflow.functions import Join
+
+    variables: List[PipelineVariable] = []
+    # A run-unique prefix, drawn from characters json.dumps never escapes, so the
+    # placeholder survives serialisation verbatim and cannot collide with real content.
+    token_prefix = "__sagemaker_pipeline_variable_%s_" % uuid.uuid4().hex
+
+    def _placeholder(obj):
+        # json.dumps invokes this for exactly the objects it cannot serialise, so
+        # it drives the traversal: each variable is recorded and replaced by a
+        # numbered placeholder in the emitted text.
+        if isinstance(obj, PipelineVariable):
+            variables.append(obj)
+            return "%s%d__" % (token_prefix, len(variables) - 1)
+        raise TypeError(
+            "Object of type %s is not JSON serializable" % obj.__class__.__name__
+        )
+
+    document = json.dumps(job_config, default=_placeholder)
+    if not variables:
+        return document
+
+    # With the capturing group, re.split alternates literal text and placeholder
+    # indices: [literal, index, literal, index, ..., literal].
+    pieces = re.split(re.escape(token_prefix) + r"(\d+)__", document)
+    if (len(pieces) - 1) // 2 != len(variables):
+        # Unreachable: every placeholder is a plain string value, so it must survive
+        # json.dumps. Raised rather than silently shipping a literal placeholder.
+        raise ValueError(
+            "[PySDK Error] Could not encode JobConfigDocument: %d of %d pipeline "
+            "variables were not found in the serialised document."
+            % (len(variables) - (len(pieces) - 1) // 2, len(variables))
+        )
+    values: List[Union[str, PipelineVariable]] = [
+        variables[int(piece)] if index % 2 else piece
+        for index, piece in enumerate(pieces)
+        if index % 2 or piece
+    ]
+    return Join(on="", values=values)

sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py

@@ -0,0 +1,93 @@
+# 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.
+"""Shared capture paths for producers running under a ``PipelineSession``."""
+from __future__ import absolute_import
+
+from typing import Any, Dict, List, Optional
+
+from sagemaker.core.apiutils._boto_functions import to_pascal_case
+from sagemaker.core.shapes import Tag
+from sagemaker.core.utils.utils import serialize
+
+
+def capture_create_job_request(
+    pipeline_session,
+    *,
+    job_name: str,
+    role_arn: str,
+    job_category: str,
+    schema_version: str,
+    job_config: Dict[str, Any],
+    tags: Optional[List[Any]] = None,
+    customer_details: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Capture a ``CreateJob`` request for ``JobStep`` composition instead of submitting it.
+
+    ``JobName`` stays in the request: ``JobStep``'s custom-job-prefix handling
+    requires the key to be present. ``job_config`` is the raw config dict --
+    ``JobStep`` scopes ``OutputDataConfig.S3OutputPath`` per execution and encodes
+    the document at definition time.
+
+    Args:
+        pipeline_session (PipelineSession): The capturing session.
+        job_name (str): Client-minted name; replaced by the service at execution.
+        role_arn (str): The execution role for the job.
+        job_category (str): ``CreateJob`` job category.
+        schema_version (str): ``JobConfigDocument`` schema version.
+        job_config (Dict[str, Any]): The raw job configuration dict.
+        tags (Optional[List[Any]]): Tags in the caller's wire form; ``serialize``
+            emits ``Tag`` shapes as ``{"Key", "Value"}``.
+        customer_details (Optional[Dict[str, Any]]): ``CustomerDetails`` envelope
+            member for producers whose direct path sends it (data preparation).
+            Session-derived, so it is stable across executions.
+    """
+    request: Dict[str, Any] = {
+        "JobName": job_name,
+        "RoleArn": role_arn,
+        "JobCategory": job_category,
+        "JobConfigSchemaVersion": schema_version,
+        "JobConfigDocument": job_config,
+    }
+    if customer_details is not None:
+        request["CustomerDetails"] = customer_details
+    if tags is not None:
+        request["Tags"] = tags
+    pipeline_session._intercept_create_request(serialize(request), None, "create_job")
+
+
+def capture_training_request(pipeline_session, create_args: Dict[str, Any]) -> None:
+    """Capture a ``CreateTrainingJob`` request for ``TrainingStep`` composition.
+
+    ``create_args`` are ``TrainingJob.create`` keyword arguments with data channels
+    already resolved, so datasets must be concrete at authoring time; only the job
+    name is deferred to execution. Client-resolution members (``session``,
+    ``region``) are dropped, and dict-form tags are coerced through the ``Tag``
+    model so ``serialize`` emits the wire form ``TrainingJob.create`` produces.
+
+    Args:
+        pipeline_session (PipelineSession): The capturing session.
+        create_args (Dict[str, Any]): ``TrainingJob.create`` keyword arguments.
+    """
+    pipeline_args = {k: v for k, v in create_args.items() if k not in ("session", "region")}
+    pipeline_args.pop("training_job_name", None)
+    request = {to_pascal_case(k): v for k, v in pipeline_args.items()}
+    if request.get("Tags"):
+        # Dict-form tags arrive in both shapes: lowercase keys from the SDK's own
+        # tag builders, PascalCase from callers passing wire-form dicts through.
+        request["Tags"] = [
+            Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value")))
+            if isinstance(tag, dict)
+            else tag
+            for tag in request["Tags"]
+        ]
+    pipeline_session._intercept_create_request(serialize(request), None, "train")

sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py

@@ -356,6 +356,12 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+# Names a producer may declare via `_pipeline_caller_name`. Restricted so a producer
+# cannot declare a name the structural checks below already resolve and compose the
+# wrong step over a different create API.
+_DECLARABLE_CALLER_NAMES = frozenset({"create_job"})
+
+
 def retrieve_caller_name(job_instance):
     """Convenience method for runnable_by_pipeline decorator
 
@@ -366,26 +372,45 @@ def retrieve_caller_name(job_instance):
         job_instance: A job class instance, one of the following types:
             - Processor (from sagemaker.core.processing)
             - ModelTrainer (from sagemaker.train.model_trainer)
+            - a V3 finetune trainer (SFT/DPO/RLVR/RLAIF, from sagemaker.train), which
+              subclasses BaseTrainer and carries no training_image
             - Transformer (from sagemaker.core.transformer)
             - HyperparameterTuner (from sagemaker.train.tuner)
+            - a CreateJob-backed producer that declares `_pipeline_caller_name`
+              (multi-turn RL)
 
     Note:
         This function uses duck typing to avoid importing from Train package,
         which would create architecture violations (Core should not depend on Train).
         Instead of isinstance checks, we check for characteristic attributes/methods.
+
+        The CreateJob family is the exception, and declares its name rather than
+        being duck typed, because its producers share no common signature to key on.
+        See `_DECLARABLE_CALLER_NAMES`.
     """
 
     from sagemaker.core.processing import Processor
     from sagemaker.core.transformer import Transformer
 
     # from sagemaker.utils.automl.automl import AutoML
 
+    # Resolved FIRST, and the order is load-bearing: MultiTurnRLTrainer inherits
+    # `input_data_config` from BaseTrainer, so the train branch below would claim it
+    # and compose a TrainingStep over CreateJob arguments.
+    declared_caller_name = getattr(job_instance, "_pipeline_caller_name", None)
+    if declared_caller_name in _DECLARABLE_CALLER_NAMES:
+        return declared_caller_name
+
     if isinstance(job_instance, Processor):
         return "run"
 
-    # Duck typing for ModelTrainer: has 'train' method and 'training_image' attribute
-    # This avoids importing from sagemaker.train which would violate architecture
-    if hasattr(job_instance, "train") and hasattr(job_instance, "training_image"):
+    # Duck typing for the CreateTrainingJob family, avoiding an import from
+    # sagemaker.train. Either marker suffices: `training_image` is ModelTrainer's,
+    # `input_data_config` is BaseTrainer's (the V3 finetune trainers resolve their
+    # image from the recipe and carry no `training_image`).
+    if hasattr(job_instance, "train") and (
+        hasattr(job_instance, "training_image") or hasattr(job_instance, "input_data_config")
+    ):
         return "train"
 
     if isinstance(job_instance, Transformer):