From bb13f00e0eadcdce95d544cb6d6230bcff42a9c8 Mon Sep 17 00:00:00 2001 From: jkasiraj Date: Fri, 18 Sep 2026 18:11:17 +0000 Subject: [PATCH] feat: add JobStep and shared pipeline capture for CreateJob producers Port from the staging branch: JobStep (StepTypeEnum.JOB) composes a CreateJob request captured from a producer running under a PipelineSession, with the JobConfigDocument encoder deferring PipelineVariables via Join. MultiTurnRLTrainer opts in via _pipeline_caller_name. The four finetune trainers' inline capture blocks are replaced by a shared capture_training_request helper with identical wire behavior, including dual-shape dict tag tolerance. --- .../core/workflow/job_config_document.py | 94 ++++ .../core/workflow/pipeline_capture.py | 93 ++++ .../core/workflow/pipeline_context.py | 31 +- .../workflow/test_retrieve_caller_name.py | 243 +++++++++ .../src/sagemaker/mlops/workflow/__init__.py | 2 + .../src/sagemaker/mlops/workflow/steps.py | 128 +++++ .../integ/workflow/test_job_step_mtrl.py | 114 ++++ .../tests/unit/workflow/test_job_step.py | 255 +++++++++ .../src/sagemaker/train/dpo_trainer.py | 26 +- .../sagemaker/train/multi_turn_rl_trainer.py | 66 ++- .../src/sagemaker/train/rlaif_trainer.py | 26 +- .../src/sagemaker/train/rlvr_trainer.py | 26 +- .../src/sagemaker/train/sft_trainer.py | 27 +- .../unit/train/test_multi_turn_rl_trainer.py | 514 +++++++++++++++++- 14 files changed, 1539 insertions(+), 106 deletions(-) create mode 100644 sagemaker-core/src/sagemaker/core/workflow/job_config_document.py create mode 100644 sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py create mode 100644 sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_job_step.py diff --git a/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py b/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py new file mode 100644 index 0000000000..63ed77eee5 --- /dev/null +++ b/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) diff --git a/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py b/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py new file mode 100644 index 0000000000..f919994b75 --- /dev/null +++ b/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") diff --git a/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py b/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py index a6f3ffe171..0f20256067 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py +++ b/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py @@ -348,6 +348,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 @@ -358,13 +364,21 @@ 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 @@ -372,12 +386,23 @@ def retrieve_caller_name(job_instance): # 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): diff --git a/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py b/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py new file mode 100644 index 0000000000..ff08fa1d50 --- /dev/null +++ b/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py @@ -0,0 +1,243 @@ +# 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. +"""Tests for `retrieve_caller_name`, including the CreateJob declared-name path. + +`retrieve_caller_name` decides the `caller_name` carried on the `_StepArguments` a +`@runnable_by_pipeline` producer returns, and that value is what every step class +checks through `validate_step_args_input(expected_caller=...)`. It is a distinct +namespace from the `func_name` passed to `PipelineSession._intercept_create_request` +-- `ProcessingStep` expects `"run"` from here while interception is called with +`"process"` -- so a step guard cannot be satisfied by the interception name alone. +""" +from __future__ import absolute_import + +import pytest + +from sagemaker.core.workflow.pipeline_context import ( + _DECLARABLE_CALLER_NAMES, + _StepArguments, + retrieve_caller_name, +) +from sagemaker.core.workflow.utilities import validate_step_args_input + + +class _Declared: + """A CreateJob-backed producer, declaring its caller name.""" + + _pipeline_caller_name = "create_job" + + +class _DeclaredAndTrainerShaped(_Declared): + """Declares create_job AND carries the ModelTrainer structural signature. + + A synthetic shape, kept because it pins the ordering against the `training_image` + marker specifically -- no real producer has this combination. For MultiTurnRLTrainer's + ACTUAL shape, and the fixture that would catch a live regression in it, see + `_DeclaredAndFinetuneShaped` below. MTRL has no training_image, so this fixture does + not cover it. + """ + + training_image = "an-image" + + def train(self): # pragma: no cover - never invoked + raise AssertionError("train() must not be called by name resolution") + + +class _UndeclaredCreateJobShaped: + """A CreateJob producer that has NOT opted in -- e.g. GenAIEvaluator today.""" + + def evaluate(self): # pragma: no cover - never invoked + raise AssertionError("evaluate() must not be called by name resolution") + + def build_config(self): # pragma: no cover - never invoked + raise AssertionError("build_config() must not be called by name resolution") + + +class _TrainerShaped: + training_image = "an-image" + + def train(self): # pragma: no cover - never invoked + raise AssertionError("train() must not be called by name resolution") + + +class _TunerShaped: + model_trainer = object() + + def tune(self): # pragma: no cover - never invoked + raise AssertionError("tune() must not be called by name resolution") + + +def test_declared_create_job_producer_resolves_to_create_job(): + assert retrieve_caller_name(_Declared()) == "create_job" + + +def test_declared_name_is_resolved_before_the_structural_trainer_check(): + """The ordering guard. A declared CreateJob producer must not be read as a trainer.""" + assert retrieve_caller_name(_DeclaredAndTrainerShaped()) == "create_job" + + +def test_an_undeclared_create_job_producer_still_resolves_to_none(): + """Fail-closed: opting in is explicit, so today's undecorated producers are unchanged.""" + assert retrieve_caller_name(_UndeclaredCreateJobShaped()) is None + + +@pytest.mark.parametrize("claimed", ["train", "transform", "tune", "run", "process", ""]) +def test_a_producer_cannot_declare_a_name_outside_the_allowlist(claimed): + """A producer must not be able to claim a name resolved structurally elsewhere. + + Declaring "train" would let a CreateJob producer compose a TrainingStep over a + different create API -- the same class of confusion `JobStep`'s guard exists to + prevent. + """ + + class _Claiming: + _pipeline_caller_name = claimed + + assert retrieve_caller_name(_Claiming()) is None + + +def test_the_declarable_set_is_pinned(): + """Widening the set is a deliberate change, so it has to move this assertion too.""" + assert _DECLARABLE_CALLER_NAMES == frozenset({"create_job"}) + + +def test_structural_resolution_is_unchanged(): + assert retrieve_caller_name(_TrainerShaped()) == "train" + assert retrieve_caller_name(_TunerShaped()) == "tune" + assert retrieve_caller_name(object()) is None + + +def test_a_declared_producer_satisfies_the_job_step_guard(): + """The payoff: what `JobStep` pins is now reachable. + + `JobStep` calls validate_step_args_input(expected_caller={"create_job"}), which + reads `_StepArguments.caller_name`. Before this branch existed no producer could + produce that value, so the guard rejected every capturable producer. + """ + step_args = _StepArguments(retrieve_caller_name(_Declared())) + + validate_step_args_input( + step_args=step_args, + expected_caller={"create_job"}, + error_message="should not raise", + ) + + +def test_the_job_step_guard_still_rejects_a_trainer_capture(): + step_args = _StepArguments(retrieve_caller_name(_TrainerShaped())) + + with pytest.raises(ValueError): + validate_step_args_input( + step_args=step_args, + expected_caller={"create_job"}, + error_message="a TrainingStep capture must not compose a JobStep", + ) + + +# --------------------------------------------------------------------------- +# The BaseTrainer family. The train branch keys on `training_image` OR +# `input_data_config`, because the four V3 finetune trainers carry no +# training_image -- they resolve the image from the recipe. +# --------------------------------------------------------------------------- + + +class _FinetuneTrainerShaped: + """An SFT/DPO/RLVR/RLAIF trainer: train() and input_data_config, no training_image.""" + + input_data_config = None + + def train(self): # pragma: no cover - never invoked + raise AssertionError("train() must not be called by name resolution") + + +class _DeclaredAndFinetuneShaped(_Declared): + """MultiTurnRLTrainer's ACTUAL shape, which the training_image fixture does not cover. + + MTRL subclasses BaseTrainer, so it already carries `input_data_config` and already + satisfies the train branch structurally. It has no training_image, so + `_DeclaredAndTrainerShaped` above would not catch a regression here. + """ + + input_data_config = None + + def train(self): # pragma: no cover - never invoked + raise AssertionError("train() must not be called by name resolution") + + +class _DataPreparerShaped: + """A CreateJob producer with no train(): attach()/stop(), plus a config attribute.""" + + input_data_config = None + + def attach(self): # pragma: no cover - never invoked + raise AssertionError("attach() must not be called by name resolution") + + def stop(self): # pragma: no cover - never invoked + raise AssertionError("stop() must not be called by name resolution") + + +def test_a_finetune_trainer_resolves_to_train(): + """The fix. Keying on training_image alone left these four unrecognised. + + An unrecognised producer returns None, `_StepArguments` is built with it happily, + and TrainingStep then rejects it at construction -- before the producer's own + interception block is ever reached. + """ + assert retrieve_caller_name(_FinetuneTrainerShaped()) == "train" + + +def test_a_finetune_trainer_satisfies_the_training_step_guard(): + """What TrainingStep pins is now reachable for the BaseTrainer family.""" + step_args = _StepArguments(retrieve_caller_name(_FinetuneTrainerShaped())) + + validate_step_args_input( + step_args=step_args, + expected_caller={"train"}, + error_message="should not raise", + ) + + +def test_the_declaration_is_what_keeps_mtrl_out_of_the_train_branch(): + """The ordering guard, at MTRL's real shape rather than a training_image stand-in. + + This is now load-bearing rather than hypothetical: MTRL carries input_data_config + from BaseTrainer, so it satisfies the train branch structurally today. Only the + declared name, resolved first, keeps it composing a JobStep instead of a + TrainingStep over CreateJob arguments. + """ + assert retrieve_caller_name(_DeclaredAndFinetuneShaped()) == "create_job" + + +def test_both_train_markers_are_independently_sufficient(): + """A union, not a replacement -- neither marker may become required.""" + assert retrieve_caller_name(_TrainerShaped()) == "train" # training_image only + assert retrieve_caller_name(_FinetuneTrainerShaped()) == "train" # input_data_config only + + +def test_the_config_marker_alone_does_not_resolve_to_train(): + """`train` stays a required conjunct, so widening cannot reach a non-trainer.""" + + class _ConfigOnly: + input_data_config = None + + assert retrieve_caller_name(_ConfigOnly()) is None + + +def test_an_undeclared_create_job_producer_without_train_is_still_unrecognised(): + """DataPreparer's shape. It fails the `train` conjunct, so the widening cannot claim it. + + DataPreparer is the CreateJob producer that is not yet pipeline-capable. When it is + onboarded it declares `create_job` like the others; until then it must not be read + as a trainer merely for holding a config attribute. + """ + assert retrieve_caller_name(_DataPreparerShaped()) is None diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 129abb1c76..22405c0620 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -37,6 +37,7 @@ TrainingStep, ProcessingStep, TransformStep, + JobStep, TuningStep, ) @@ -89,6 +90,7 @@ "TrainingStep", "ProcessingStep", "TransformStep", + "JobStep", "TuningStep", # Step implementations "AutoMLStep", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 76e90a5309..fc158b2176 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -62,6 +62,7 @@ class StepTypeEnum(Enum): EMR_SERVERLESS = "EMRServerless" FAIL = "Fail" AUTOML = "AutoML" + JOB = "Job" class Step(Entity): @@ -789,3 +790,130 @@ def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = "") -> "output/model.tar.gz", ], ) + + +class JobStep(ConfigurableRetryStep): + """`JobStep` for SageMaker Pipelines Workflows. + + Wraps a `CreateJob` request captured from a producer running under a + `PipelineSession`. `CreateJob` is one control-plane API shared by every job + category: `JobCategory` names the category and all of its parameters travel + inside the `JobConfigDocument` envelope member, a serialized JSON string. + `properties.JobConfigDocument` is therefore a string, not a modeled structure, + but a property reference may descend into it: for example + ``step.properties.JobConfigDocument.OutputModelPackageArn`` resolves as a + String, with the backend parsing the JSON lazily at execution time. + + When the captured request carries the job configuration as a dict, `arguments` + scopes `OutputDataConfig.S3OutputPath` per execution and encodes the document + (see `sagemaker.core.workflow.job_config_document` for the encoding and its + limits). A pre-encoded string document is passed through unchanged. + + `SageMakerJobStepRetryPolicy` applies: `CreateJob` is a synchronous create + that can hit a resource limit, which is what that policy retries. + """ + + def __init__( + self, + name: str, + step_args: Optional[_JobStepArguments] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + cache_config: Optional[CacheConfig] = None, + depends_on: Optional[List[Union[str, Step]]] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + ): + """Construct a `JobStep` using step_args captured from a `CreateJob` producer. + + Args: + name (str): The name of the `JobStep`. + step_args (_JobStepArguments): The arguments for the `JobStep` definition. + display_name (str): The display name of the `JobStep`. + description (str): The description of the `JobStep`. + cache_config (CacheConfig): A `sagemaker.workflow.steps.CacheConfig` instance. + depends_on (List[Union[str, Step]]): A list of `Step` + names or `Step` instances that this `JobStep` + depends on. + retry_policies (List[RetryPolicy]): A list of retry policies. + """ + super(JobStep, self).__init__( + name, StepTypeEnum.JOB, display_name, description, depends_on, retry_policies + ) + + if step_args: + from sagemaker.core.workflow.utilities import validate_step_args_input + + validate_step_args_input( + step_args=step_args, + expected_caller={"create_job"}, + error_message="The step_args of JobStep must be obtained from a producer that " + "creates a SageMaker job via CreateJob.", + ) + + self.step_args = step_args + self._properties = Properties(step_name=name, step=self, shape_name="DescribeJobResponse") + self.cache_config = cache_config + + @property + def arguments(self) -> RequestType: + """The arguments dictionary that is used to call `create_job`. + + NOTE: `CreateJob` has no `ExperimentConfig` member, so unlike the training, + processing and transform steps there is no experiment config to trim. + """ + from sagemaker.core.workflow.execution_variables import ExecutionVariables + from sagemaker.core.workflow.functions import Join + from sagemaker.core.workflow.job_config_document import convert_job_config_document_to_string + from sagemaker.core.workflow.utilities import execute_job_functions + from sagemaker.core.workflow.utilities import _pipeline_config + + if self.step_args: + # execute the producer function with saved parameters, + # and store args in PipelineSession's _context + execute_job_functions(self.step_args) + + # populate request dict with args + producer = self.step_args.func_args[0] + request_dict = producer.sagemaker_session.context.args + else: + raise ValueError("step_args input is required.") + + document = request_dict.get("JobConfigDocument") + if isinstance(document, dict): + output_config = document.get("OutputDataConfig") + if isinstance(output_config, dict) and isinstance( + output_config.get("S3OutputPath"), str + ): + # Nothing scopes CreateJob output server-side (training appends the + # job name; CreateJob appends nothing), so executions sharing a + # prefix silently read each other's results. The execution id, not + # the job name: JobName is regenerated below, and a sibling step can + # rebuild this path from ExecutionVariables without a property + # reference. The trailing empty value yields the trailing slash. + output_config["S3OutputPath"] = Join( + on="/", + values=[ + output_config["S3OutputPath"].rstrip("/"), + ExecutionVariables.PIPELINE_EXECUTION_ID, + "", + ], + ) + request_dict["JobConfigDocument"] = convert_job_config_document_to_string(document) + + # Continue to pop job name if not explicitly opted-in via config + request_dict = trim_request_dict(request_dict, "JobName", _pipeline_config) + + return request_dict + + @property + def properties(self): + """A `Properties` object representing the `DescribeJobResponse` data model.""" + return self._properties + + def to_request(self) -> RequestType: + """Updates the request dictionary with cache configuration.""" + request_dict = super().to_request() + if self.cache_config: + request_dict.update(self.cache_config.config) + + return request_dict diff --git a/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py b/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py new file mode 100644 index 0000000000..1fd4ac646b --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py @@ -0,0 +1,114 @@ +"""Integration test for a `JobStep` captured from `MultiTurnRLTrainer`. + +Agent RFT needs real assets, so inputs come from environment variables and the +test skips when they are absent: + + SAGEMAKER_INTEG_MTRL_MODEL model id for multi-turn RL + SAGEMAKER_INTEG_MTRL_AGENT_ENV Bedrock AgentCore runtime ARN (or Lambda ARN) + SAGEMAKER_INTEG_MTRL_DATASET training dataset S3 URI + SAGEMAKER_INTEG_MODEL_PACKAGE_GROUP output model package group + SAGEMAKER_INTEG_MTRL_MLFLOW_APP MLflow app ARN + +Proves the second `CreateJob` producer end to end: the pipeline service accepts +a `JobStep` whose arguments were captured from `MultiTurnRLTrainer.train()`, +resolves the execution-scoped output path inside `JobConfigDocument`, and +creates the AgentRFT job from it. The execution is stopped once the job exists. +""" +from __future__ import absolute_import + +import json +import os +import time +import uuid + +import pytest + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.steps import JobStep +from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer + +_REQUIRED_ENV = ( + "SAGEMAKER_INTEG_MTRL_MODEL", + "SAGEMAKER_INTEG_MTRL_AGENT_ENV", + "SAGEMAKER_INTEG_MTRL_DATASET", + "SAGEMAKER_INTEG_MODEL_PACKAGE_GROUP", + "SAGEMAKER_INTEG_MTRL_MLFLOW_APP", +) + +pytestmark = pytest.mark.skipif( + any(not os.environ.get(name) for name in _REQUIRED_ENV), + reason="requires AgentRFT assets: %s" % ", ".join(_REQUIRED_ENV), +) + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def _wait_for_created_job(execution, timeout_seconds=900): + """Wait until the step reports its created job, or terminal failure.""" + deadline = time.time() + timeout_seconds + while time.time() < deadline: + steps = execution.list_steps() + if steps: + metadata = steps[0].get("Metadata", {}) + if "Job" in metadata: + return steps[0] + if steps[0].get("StepStatus") in ("Failed", "Stopped"): + raise AssertionError("step ended %s: %s" % (steps[0]["StepStatus"], steps[0])) + time.sleep(30) + raise TimeoutError("step did not create a job in time") + + +def test_job_step_from_mtrl_capture(sagemaker_session, pipeline_session, role): + bucket = sagemaker_session.default_bucket() + pipeline_name = "integ-mtrl-job-step-%s" % uuid.uuid4().hex[:8] + + trainer = MultiTurnRLTrainer( + model=os.environ["SAGEMAKER_INTEG_MTRL_MODEL"], + agent_env=os.environ["SAGEMAKER_INTEG_MTRL_AGENT_ENV"], + training_dataset=os.environ["SAGEMAKER_INTEG_MTRL_DATASET"], + output_model_package_group=os.environ["SAGEMAKER_INTEG_MODEL_PACKAGE_GROUP"], + mlflow_app_arn=os.environ["SAGEMAKER_INTEG_MTRL_MLFLOW_APP"], + s3_output_path="s3://%s/%s/output" % (bucket, pipeline_name), + accept_eula=True, + sagemaker_session=pipeline_session, + ) + + step = JobStep(name="mtrl", step_args=trainer.train(wait=False)) + pipeline = Pipeline(name=pipeline_name, steps=[step], sagemaker_session=pipeline_session) + + definition = json.loads(pipeline.definition()) + assert definition["Steps"][0]["Type"] == "Job" + document = definition["Steps"][0]["Arguments"]["JobConfigDocument"] + # JobStep scopes the output path per execution inside the document. + assert "Std:Join" in document + assert {"Get": "Execution.PipelineExecutionId"} in document["Std:Join"]["Values"] + + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + step_state = _wait_for_created_job(execution) + assert step_state["Metadata"]["Job"]["Arn"] + finally: + try: + execution.stop() + except Exception: # noqa: BLE001 -- may already be terminal + pass + try: + pipeline.delete() + except Exception: # noqa: BLE001 -- best-effort cleanup + pass diff --git a/sagemaker-mlops/tests/unit/workflow/test_job_step.py b/sagemaker-mlops/tests/unit/workflow/test_job_step.py new file mode 100644 index 0000000000..3a7bcc6da9 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_job_step.py @@ -0,0 +1,255 @@ +# 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. +"""Unit tests for `StepTypeEnum.JOB` and `JobStep`.""" +from __future__ import absolute_import + +import pytest +from unittest.mock import Mock + +from sagemaker.core.workflow.pipeline_context import _JobStepArguments, _StepArguments +from sagemaker.mlops.workflow.retry import ( + SageMakerJobExceptionTypeEnum, + SageMakerJobStepRetryPolicy, +) +from sagemaker.mlops.workflow.steps import CacheConfig, ConfigurableRetryStep, JobStep, StepTypeEnum + +# --- StepTypeEnum membership ------------------------------------------------- +# +# The member set is pinned as a whole. A new step type must be added here in the +# same change that adds it to the enum, so an accidental addition or removal +# fails rather than passing silently. This tree does not carry PR #6224, which +# would add ENDPOINT_CONFIG, ENDPOINT, INFERENCE_COMPONENT and LINEAGE. +EXPECTED_STEP_TYPES = { + "CONDITION": "Condition", + "CREATE_MODEL": "Model", + "PROCESSING": "Processing", + "REGISTER_MODEL": "RegisterModel", + "TRAINING": "Training", + "TRANSFORM": "Transform", + "CALLBACK": "Callback", + "TUNING": "Tuning", + "LAMBDA": "Lambda", + "QUALITY_CHECK": "QualityCheck", + "CLARIFY_CHECK": "ClarifyCheck", + "EMR": "EMR", + "EMR_SERVERLESS": "EMRServerless", + "FAIL": "Fail", + "AUTOML": "AutoML", + "JOB": "Job", +} + + +def test_step_type_enum_member_set_is_exact(): + assert {member.name: member.value for member in StepTypeEnum} == EXPECTED_STEP_TYPES + + +def test_step_type_enum_has_job_member(): + assert StepTypeEnum.JOB.value == "Job" + # The value must match the service's own name for the step, which is how a + # definition round-trips through StepTypeEnum(request_dict["Type"]). + assert StepTypeEnum("Job") is StepTypeEnum.JOB + + +# --- JobStep construction --------------------------------------------------- + + +def test_job_step_is_configurable_retry_step(): + # Retryable like TrainingStep, not a plain Step: CreateJob is a synchronous + # create that can hit a resource limit. + step = JobStep(name="my-job") + assert isinstance(step, ConfigurableRetryStep) + assert step.step_type is StepTypeEnum.JOB + assert step.name == "my-job" + + +def test_job_step_accepts_sagemaker_job_step_retry_policy(): + policy = SageMakerJobStepRetryPolicy( + exception_types=[SageMakerJobExceptionTypeEnum.RESOURCE_LIMIT], + max_attempts=2, + ) + step = JobStep(name="my-job", retry_policies=[policy]) + assert step.retry_policies == [policy] + step.add_retry_policy(policy) + assert len(step.retry_policies) == 2 + + +def test_job_step_optional_metadata(): + step = JobStep( + name="my-job", + display_name="My Job", + description="a job step", + depends_on=["upstream"], + ) + assert step.display_name == "My Job" + assert step.description == "a job step" + assert step.depends_on == ["upstream"] + + +# --- properties ------------------------------------------------------------- + + +def test_job_step_properties_expose_describe_job_response_members(): + step = JobStep(name="my-job") + properties = step.properties + + # Walked from the botocore DescribeJobResponse shape, not hand-listed. + for member in ("JobName", "JobArn", "JobCategory", "JobStatus", "JobConfigDocument"): + assert hasattr(properties, member), member + + assert properties.JobName.expr == {"Get": "Steps.my-job.JobName"} + + +def test_job_config_document_is_a_string_property(): + # JobConfigDocument is a JSON string in the describe response, so it has no + # modeled sub-members. A reference descending into it is resolved by the + # pipeline service at execution time, never type-checked here. + from sagemaker.core.workflow.properties import Properties + + step = JobStep(name="my-job") + document = step.properties.JobConfigDocument + + assert document.expr == {"Get": "Steps.my-job.JobConfigDocument"} + assert [key for key, value in document.__dict__.items() if isinstance(value, Properties)] == [] + + # Contrast: a modeled structure member does get walked. + assert isinstance(step.properties.SecondaryStatusTransitions, Properties) + + +def test_job_step_properties_reference_the_step_instance(): + step = JobStep(name="my-job") + assert step.properties._referenced_steps == [step] + + +# --- step_args validation --------------------------------------------------- + + +def test_job_step_rejects_non_step_args(): + with pytest.raises(TypeError, match="must be obtained from a producer"): + JobStep(name="my-job", step_args={"JobName": "not-step-args"}) + + +@pytest.mark.parametrize("caller_name", ["train", "transform", "run", "tune", "create_model"]) +def test_job_step_rejects_other_producers(caller_name): + # Every currently capturable producer builds a different create request. + step_args = _JobStepArguments(caller_name, {"JobName": "j"}) + with pytest.raises(ValueError, match="must be obtained from a producer"): + JobStep(name="my-job", step_args=step_args) + + +def test_job_step_accepts_create_job_caller(): + step_args = _JobStepArguments("create_job", {"JobName": "j"}) + step = JobStep(name="my-job", step_args=step_args) + assert step.step_args is step_args + + +# --- arguments -------------------------------------------------------------- + + +def test_job_step_arguments_requires_step_args(): + step = JobStep(name="my-job", step_args=None) + with pytest.raises(ValueError, match="step_args input is required"): + _ = step.arguments + + +def _capturing_step_args(request): + """Build step_args whose func writes `request` into the session context. + + This stands in for `@runnable_by_pipeline` + `_intercept_create_request`, + which no CreateJob producer implements yet. + """ + producer = Mock() + producer.sagemaker_session.context.args = request + + def capture(_producer): + return None + + return _StepArguments("create_job", capture, producer) + + +def test_job_step_arguments_trims_job_name_by_default(): + request = { + "JobName": "my-job-2026-09-10-22-31-57-000", + "RoleArn": "arn:aws:iam::123456789012:role/JobRole", + "JobCategory": "SyntheticDataGeneration", + "JobConfigSchemaVersion": "1.0", + "JobConfigDocument": '{"Foo": "bar"}', + } + step = JobStep(name="my-job", step_args=_capturing_step_args(request)) + + arguments = step.arguments + + # No custom job prefix opted in, so the generated name is dropped from the + # persisted definition. + assert "JobName" not in arguments + assert arguments["JobCategory"] == "SyntheticDataGeneration" + assert arguments["JobConfigDocument"] == '{"Foo": "bar"}' + + +def test_job_step_to_request_shape(): + request = {"JobName": "my-job", "JobCategory": "DataQualityEvaluation"} + step = JobStep( + name="my-job", + step_args=_capturing_step_args(request), + display_name="My Job", + description="a job step", + depends_on=["upstream"], + ) + + step_request = step.to_request() + + assert step_request["Name"] == "my-job" + assert step_request["Type"] == "Job" + assert step_request["DependsOn"] == ["upstream"] + assert step_request["DisplayName"] == "My Job" + assert step_request["Description"] == "a job step" + assert "CacheConfig" not in step_request + + +def test_job_step_to_request_includes_cache_config(): + request = {"JobName": "my-job"} + step = JobStep( + name="my-job", + step_args=_capturing_step_args(request), + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), + ) + + step_request = step.to_request() + + assert step_request["CacheConfig"] == {"Enabled": True, "ExpireAfter": "P30D"} + + +def test_job_step_retry_policies_land_in_request(): + request = {"JobName": "my-job"} + policy = SageMakerJobStepRetryPolicy( + exception_types=[SageMakerJobExceptionTypeEnum.RESOURCE_LIMIT], + max_attempts=3, + ) + step = JobStep( + name="my-job", + step_args=_capturing_step_args(request), + retry_policies=[policy], + ) + + step_request = step.to_request() + + assert step_request["RetryPolicies"] == [policy.to_request()] + + +# --- exports ---------------------------------------------------------------- + + +def test_job_step_is_exported_from_workflow_package(): + import sagemaker.mlops.workflow as workflow + + assert workflow.JobStep is JobStep + assert "JobStep" in workflow.__all__ diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index 159e7b4230..4d3fa77dc4 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py @@ -5,9 +5,8 @@ from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE from sagemaker.core.resources import TrainingJob, ModelPackageGroup, ModelPackage from sagemaker.core.shapes import VpcConfig +from sagemaker.core.workflow.pipeline_capture import capture_training_request from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline -from sagemaker.core.utils.utils import serialize -from sagemaker.core.apiutils._boto_functions import to_pascal_case from sagemaker.train.defaults import TrainDefaults from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.train.configs import StoppingCondition @@ -373,27 +372,10 @@ def train(self, if self.stopping_condition is not None: create_args["stopping_condition"] = self.stopping_condition - # If running within a PipelineSession, intercept the request and store - # step arguments instead of launching a training job. - # This must come before data path validation since in pipeline mode - # the data path may be a pipeline parameter that doesn't exist yet. + # Capture must come before data path validation: in pipeline mode the + # data path may be a pipeline parameter that doesn't exist yet. if isinstance(sagemaker_session, PipelineSession): - pipeline_args = {k: v for k, v in create_args.items() - if k not in ("session", "region")} - pipeline_args.pop("training_job_name", None) - pipeline_request = {to_pascal_case(k): v for k, v in pipeline_args.items()} - # Normalize Tags to PascalCase dicts. JumpStart tags come as lowercase - # dicts; user-provided tags come as Tag pydantic objects (typed - # Optional[List[Tag]]). Handle both. - if "Tags" in pipeline_request and pipeline_request["Tags"]: - pipeline_request["Tags"] = [ - {"Key": t.get("key", t.get("Key")), "Value": t.get("value", t.get("Value"))} - if isinstance(t, dict) - else {"Key": t.key, "Value": t.value} - for t in pipeline_request["Tags"] - ] - serialized_request = serialize(pipeline_request) - sagemaker_session._intercept_create_request(serialized_request, None, "train") + capture_training_request(sagemaker_session, create_args) return # Validate data paths exist before submission diff --git a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py index 2ab58b1cbf..8b8a9f6b2a 100644 --- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py +++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py @@ -17,15 +17,17 @@ import json import logging import re -from typing import Any, Dict, Optional, Union +from typing import Any, ClassVar, Dict, Optional, Union import boto3 from sagemaker.ai_registry.dataset import DataSet from sagemaker.core.resources import Job, ModelPackageGroup, ModelPackage, MlflowApp -from sagemaker.core.shapes import VpcConfig +from sagemaker.core.shapes import Tag, VpcConfig from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter, TelemetryParamType from sagemaker.core.telemetry.constants import Feature +from sagemaker.core.workflow.pipeline_capture import capture_create_job_request +from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline from sagemaker.train.custom_agent_lambda import CustomAgentLambda from sagemaker.train.agent_rft_job import AgentRFTJob from sagemaker.train.base_trainer import BaseTrainer @@ -182,6 +184,8 @@ class MultiTurnRLTrainer(BaseTrainer): _customization_technique = "MTRL" + _pipeline_caller_name: ClassVar[str] = "create_job" + def __init__( self, model: Union[str, ModelPackage], @@ -277,22 +281,28 @@ def __init__( ("wait", TelemetryParamType.KWARG_EXISTS), ], ) + @runnable_by_pipeline def train( self, training_dataset: Optional[Union[str, DataSet]] = None, wait: bool = True, dry_run: bool = False, - ) -> AgentRFTJob: + ) -> Optional[AgentRFTJob]: """Launch an Agentic RFT job. Args: training_dataset: Training dataset override. wait: If True (default), block until job reaches terminal status. + Ignored under a ``PipelineSession``, where there is no job to wait + on -- ``runnable_by_pipeline`` forces it to ``False`` and warns. dry_run: If True, runs validation without submitting a job. - Returns None on success. + Returns None on success. Ignored (with a warning) under a + ``PipelineSession``, where nothing is submitted either way. Returns: - AgentRFTJob instance for tracking the job, or None if dry_run=True. + AgentRFTJob instance for tracking the job. ``None`` if dry_run=True + or under a ``PipelineSession``, where the assembled request is + captured for ``JobStep`` composition instead of submitted. """ sagemaker_session = TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session @@ -313,17 +323,39 @@ def train( if training_dataset is not None: self.training_dataset = training_dataset - job_config_doc = self._build_job_config_document(dry_run=dry_run) - - if dry_run: - logger.info("Dry-run validation passed. No job submitted.") - return None tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) # Merge user-provided tags with the JumpStart tags tags.extend(self.tags or []) + # Capture must come before the document build: a `PipelineVariable` in the + # config is only encodable by `JobStep`'s deferred encoding, and `json.dumps` + # rejects it. The decorator gates on the session too, but only the body runs + # when `execute_job_functions` re-invokes the captured function, so the test + # here is the one that populates `session.context`. + if isinstance(sagemaker_session, PipelineSession): + if dry_run: + logger.warning( + "dry_run is ignored under a PipelineSession: the assembled " + "CreateJob request is captured for pipeline composition instead." + ) + return capture_create_job_request( + sagemaker_session, + job_name=current_job_name, + role_arn=role, + job_category=JOB_CATEGORY, + schema_version=JOB_CONFIG_SCHEMA_VERSION, + job_config=self._build_job_config(), + tags=[Tag(**tag) if isinstance(tag, dict) else tag for tag in tags], + ) + + job_config_doc = self._build_job_config_document(dry_run=dry_run) + + if dry_run: + logger.info("Dry-run validation passed. No job submitted.") + return None + try: job = Job.create( job_name=current_job_name, @@ -421,8 +453,12 @@ def attach(cls, job_name: str, session=None) -> AgentRFTJob: # ---- Private: JobConfigDocument construction ---- - def _build_job_config_document(self, dry_run: bool = False) -> str: - """Build the JobConfigDocument JSON string conforming to v1_0_0 schema.""" + def _build_job_config(self, dry_run: bool = False) -> dict: + """Build the job configuration conforming to the v1_0_0 schema, as a dict. + + The capture path hands this dict to ``JobStep``, which scopes the output + path per execution and encodes the document at definition time. + """ config = { "AgentConfig": self._build_agent_config(), "InputDataConfig": self._build_input_data_config(), @@ -435,7 +471,11 @@ def _build_job_config_document(self, dry_run: bool = False) -> str: "SecurityGroupIds": self.networking.security_group_ids, "Subnets": self.networking.subnets, } - doc = json.dumps(config, indent=2) + return config + + def _build_job_config_document(self, dry_run: bool = False) -> str: + """Build the JobConfigDocument JSON string conforming to v1_0_0 schema.""" + doc = json.dumps(self._build_job_config(dry_run=dry_run), indent=2) logger.info(f"JobConfigDocument:\n{doc}") return doc diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index 3a8197d74b..52c3ccc7e1 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py @@ -4,9 +4,8 @@ from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE from sagemaker.core.resources import TrainingJob, ModelPackageGroup, MlflowTrackingServer, ModelPackage from sagemaker.core.shapes import VpcConfig +from sagemaker.core.workflow.pipeline_capture import capture_training_request from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline -from sagemaker.core.utils.utils import serialize -from sagemaker.core.apiutils._boto_functions import to_pascal_case from sagemaker.train.defaults import TrainDefaults from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.train.common_utils.recipe_utils import _get_hub_content_metadata @@ -339,27 +338,10 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati if self.stopping_condition is not None: create_args["stopping_condition"] = self.stopping_condition - # If running within a PipelineSession, intercept the request and store - # step arguments instead of launching a training job. - # This must come before data path validation since in pipeline mode - # the data path may be a pipeline parameter that doesn't exist yet. + # Capture must come before data path validation: in pipeline mode the + # data path may be a pipeline parameter that doesn't exist yet. if isinstance(sagemaker_session, PipelineSession): - pipeline_args = {k: v for k, v in create_args.items() - if k not in ("session", "region")} - pipeline_args.pop("training_job_name", None) - pipeline_request = {to_pascal_case(k): v for k, v in pipeline_args.items()} - # Normalize Tags to PascalCase dicts. JumpStart tags come as lowercase - # dicts; user-provided tags come as Tag pydantic objects (typed - # Optional[List[Tag]]). Handle both. - if "Tags" in pipeline_request and pipeline_request["Tags"]: - pipeline_request["Tags"] = [ - {"Key": t.get("key", t.get("Key")), "Value": t.get("value", t.get("Value"))} - if isinstance(t, dict) - else {"Key": t.key, "Value": t.value} - for t in pipeline_request["Tags"] - ] - serialized_request = serialize(pipeline_request) - sagemaker_session._intercept_create_request(serialized_request, None, "train") + capture_training_request(sagemaker_session, create_args) return # Validate data paths exist before submission diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index a182b1a581..d3c3537f74 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py @@ -6,9 +6,8 @@ from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE from sagemaker.core.resources import TrainingJob, ModelPackageGroup, MlflowTrackingServer, ModelPackage from sagemaker.core.shapes import VpcConfig +from sagemaker.core.workflow.pipeline_capture import capture_training_request from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline -from sagemaker.core.utils.utils import serialize -from sagemaker.core.apiutils._boto_functions import to_pascal_case from sagemaker.train.defaults import TrainDefaults from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.ai_registry.dataset import DataSet @@ -559,27 +558,10 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, if self.stopping_condition is not None: create_args["stopping_condition"] = self.stopping_condition - # If running within a PipelineSession, intercept the request and store - # step arguments instead of launching a training job. - # This must come before data path validation since in pipeline mode - # the data path may be a pipeline parameter that doesn't exist yet. + # Capture must come before data path validation: in pipeline mode the + # data path may be a pipeline parameter that doesn't exist yet. if isinstance(sagemaker_session, PipelineSession): - pipeline_args = {k: v for k, v in create_args.items() - if k not in ("session", "region")} - pipeline_args.pop("training_job_name", None) - pipeline_request = {to_pascal_case(k): v for k, v in pipeline_args.items()} - # Normalize Tags to PascalCase dicts. JumpStart tags come as lowercase - # dicts; user-provided tags come as Tag pydantic objects (typed - # Optional[List[Tag]]). Handle both. - if "Tags" in pipeline_request and pipeline_request["Tags"]: - pipeline_request["Tags"] = [ - {"Key": t.get("key", t.get("Key")), "Value": t.get("value", t.get("Value"))} - if isinstance(t, dict) - else {"Key": t.key, "Value": t.value} - for t in pipeline_request["Tags"] - ] - serialized_request = serialize(pipeline_request) - sagemaker_session._intercept_create_request(serialized_request, None, "train") + capture_training_request(sagemaker_session, create_args) return # Validate data paths exist before submission diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index e810bdec52..17c69b6c2c 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py @@ -4,9 +4,8 @@ from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE from sagemaker.core.resources import TrainingJob, ModelPackageGroup, ModelPackage from sagemaker.core.shapes import VpcConfig +from sagemaker.core.workflow.pipeline_capture import capture_training_request from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline -from sagemaker.core.utils.utils import serialize -from sagemaker.core.apiutils._boto_functions import to_pascal_case from sagemaker.train.defaults import TrainDefaults from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.ai_registry.dataset import DataSet @@ -441,28 +440,10 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati if self.stopping_condition is not None: create_args["stopping_condition"] = self.stopping_condition - # If running within a PipelineSession, intercept the request and store - # step arguments instead of launching a training job. - # This must come before data path validation since in pipeline mode - # the data path may be a pipeline parameter that doesn't exist yet. + # Capture must come before data path validation: in pipeline mode the + # data path may be a pipeline parameter that doesn't exist yet. if isinstance(sagemaker_session, PipelineSession): - # Build pipeline-compatible request: PascalCase, serialized, no session/region - pipeline_args = {k: v for k, v in create_args.items() - if k not in ("session", "region")} - pipeline_args.pop("training_job_name", None) - pipeline_request = {to_pascal_case(k): v for k, v in pipeline_args.items()} - # Normalize Tags to PascalCase dicts. JumpStart tags come as lowercase - # dicts; user-provided tags come as Tag pydantic objects (typed - # Optional[List[Tag]]). Handle both. - if "Tags" in pipeline_request and pipeline_request["Tags"]: - pipeline_request["Tags"] = [ - {"Key": t.get("key", t.get("Key")), "Value": t.get("value", t.get("Value"))} - if isinstance(t, dict) - else {"Key": t.key, "Value": t.value} - for t in pipeline_request["Tags"] - ] - serialized_request = serialize(pipeline_request) - sagemaker_session._intercept_create_request(serialized_request, None, "train") + capture_training_request(sagemaker_session, create_args) return # Validate data paths exist before submission diff --git a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py index 360deebcb2..7d3ec65cce 100644 --- a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py +++ b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py @@ -1,11 +1,23 @@ """Unit tests for MultiTurnRLTrainer.""" import json +import warnings from unittest.mock import MagicMock, patch, PropertyMock import pytest from sagemaker.ai_registry.dataset import DataSet -from sagemaker.core.resources import ModelPackage, MlflowApp +from sagemaker.core.resources import Base, Job, ModelPackage, MlflowApp +from sagemaker.core.shapes import Tag +from sagemaker.core.workflow.execution_variables import ExecutionVariable +from sagemaker.core.workflow.functions import Join +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.core.workflow.pipeline_context import ( + PipelineSession, + _JobStepArguments, + _StepArguments, + retrieve_caller_name, +) +from sagemaker.core.workflow.utilities import execute_job_functions from sagemaker.train.custom_agent_lambda import CustomAgentLambda from sagemaker.train.multi_turn_rl_trainer import ( MultiTurnRLTrainer, @@ -832,3 +844,503 @@ def test_dry_run_passes_flag_to_mlflow_resolver(self, mock_get_role, mock_job_cl # Verify dry_run=True was passed call_kwargs = mock_resolve_mlflow.call_args[1] assert call_kwargs["dry_run"] is True + + +PINNED_JOB_NAME = "test-model-mtrl-1757000000-abc123" +ROLE_ARN = "arn:aws:iam::123456789012:role/SageMakerRole" + + +class _Hyperparameters: + """Minimal stand-in for the hyperparameters object `train()` snapshots.""" + + def __init__(self, values=None): + self._values = values or {} + + def to_dict(self): + return dict(self._values) + + +class TestCaptureTagCoercion: + """The capture path must emit the same wire tags as `Job.create`. + + The capture path never reaches `Job.create`, so `train()` coerces dict-form + tags through the `Tag` model before `serialize`, which is the same parse + `Job.create` applies. Population: `_get_jumpstart_tags` lowercase dicts and + `BaseTrainer.tags` `Tag` objects. + """ + + @staticmethod + def _capture_wire_tags(tags): + """What the capture path sends: coerce to `Tag`, then `serialize`.""" + from sagemaker.core.utils.utils import serialize + + return serialize([Tag(**tag) if isinstance(tag, dict) else tag for tag in tags]) + + @staticmethod + def _job_create_wire_tags(tags): + """What `Job.create` really sends for `tags`, measured not derived.""" + mock_client = MagicMock() + with patch.object(Base, "get_sagemaker_client", return_value=mock_client): + try: + Job.create( + job_name=PINNED_JOB_NAME, + role_arn=ROLE_ARN, + job_category=JOB_CATEGORY, + job_config_schema_version=JOB_CONFIG_SCHEMA_VERSION, + job_config_document="{}", + tags=tags, + ) + except Exception: + # Constructing the `Job` resource from a MagicMock response fails after + # the call. The call itself is what is under test. + pass + assert mock_client.create_job.call_args is not None, "create_job was not reached" + return mock_client.create_job.call_args.kwargs["Tags"] + + def test_matches_job_create_for_lowercase_dicts(self): + tags = [{"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}] + assert self._capture_wire_tags(tags) == self._job_create_wire_tags(tags) + + def test_matches_job_create_for_tag_objects(self): + tags = [Tag(key="Project", value="beta")] + assert self._capture_wire_tags(tags) == self._job_create_wire_tags(tags) + + def test_matches_job_create_for_both_forms_in_one_list(self): + tags = [ + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + Tag(key="Project", value="beta"), + ] + assert self._capture_wire_tags(tags) == self._job_create_wire_tags(tags) + + def test_empty_list_is_preserved_not_dropped(self): + """`Job.create` sends `Tags: []` rather than omitting the key.""" + assert self._capture_wire_tags([]) == self._job_create_wire_tags([]) == [] + + +class TestPipelineCapture: + """Under a `PipelineSession`, `train()` must capture rather than submit. + + The producer half of composing a `JobStep` over `CreateJob`: declare a caller name + the resolver recognises, assemble the request `Job.create` would have assembled, + and hand it to the session instead of the service. + """ + + @pytest.fixture(autouse=True) + def _skip_role_validation(self, monkeypatch): + """`TrainDefaults.get_role` validates even an explicit role against live IAM.""" + monkeypatch.setattr( + "sagemaker.train.defaults.resolve_and_validate_role", + lambda provided_role=None, **kwargs: provided_role or ROLE_ARN, + ) + + @staticmethod + def _pipeline_session(): + """A real `PipelineSession` with only its outbound edges stubbed. + + `_intercept_create_request` is deliberately NOT mocked: these tests assert on + the `_JobStepArguments` it really builds. + """ + session = PipelineSession() + session.sagemaker_client = MagicMock() + session.sagemaker_config = {} + return session + + @staticmethod + def _direct_session(): + """A non-pipeline session stub. + + `sagemaker_config` is a real dict because `_telemetry_emitter` resolves the + opt-out flag through it and jsonschema rejects a MagicMock there. + """ + session = MagicMock() + session.sagemaker_config = {} + return session + + @staticmethod + def _make_trainer(sagemaker_session=None, **overrides): + """A trainer with `__init__`'s resolution already done, as the other suites do.""" + trainer = object.__new__(MultiTurnRLTrainer) + trainer.agent_env = BEDROCK_AGENT_ARN + trainer.bedrock_agentcore_qualifier = "DEFAULT" + trainer.s3_output_path = overrides.get("s3_output_path", S3_OUTPUT) + trainer.output_model_package_group = MPG_ARN + trainer.intermediate_checkpoint_model_package_group = ( + "arn:aws:sagemaker:us-west-2:123456789012:model-package-group/default-ckpt-mpg" + ) + trainer.mlflow_app_arn = MLFLOW_ARN + trainer.mlflow_experiment_name = None + trainer.mlflow_run_name = None + trainer.accept_eula = True + trainer.kms_key_arn = overrides.get("kms_key_arn") + trainer.networking = None + trainer.model = "test-model-id" + trainer.validation_dataset = None + trainer._model_arn = MODEL_ARN + trainer.training_dataset = overrides.get("training_dataset", S3_DATA) + trainer._hp_defaults = {} + trainer.hyperparameters = _Hyperparameters(overrides.get("hyperparameters")) + trainer._model_name = "test-model" + trainer.base_job_name = "test-model-mtrl" + trainer.role = ROLE_ARN + trainer.tags = overrides.get("tags") + trainer.sagemaker_session = sagemaker_session + trainer._latest_job = overrides.get("latest_job") + trainer._recipe_path = None + trainer._overrides = None + trainer._resolved_recipe_cache = None + return trainer + + @staticmethod + def _resolve(value): + """Resolve an encoded document to text, the way the service would.""" + if isinstance(value, str): + return value + if isinstance(value, Join): + return value.on.join(TestPipelineCapture._resolve(item) for item in value.values) + if isinstance(value, ExecutionVariable): + return "EXEC-ID" + if isinstance(value, ParameterString): + return "RESOLVED-PARAM" + raise AssertionError("unexpected value in encoded document: %r" % (value,)) + + @staticmethod + def _pinned_name(): + return patch( + "sagemaker.train.multi_turn_rl_trainer._get_unique_name", + return_value=PINNED_JOB_NAME, + ) + + def _capture(self, trainer, session): + """Drive the decorator, then replay the captured call as a step would.""" + with self._pinned_name(): + step_args = trainer.train() + assert isinstance(step_args, _StepArguments) + execute_job_functions(step_args) + return session.context + + # --- the declared caller name --------------------------------------------- + + def test_declares_the_create_job_caller_name(self): + """`JobStep`'s `expected_caller={"create_job"}` guard accepts only this value.""" + assert MultiTurnRLTrainer._pipeline_caller_name == "create_job" + + def test_caller_name_is_a_plain_string_on_the_class(self): + assert isinstance(MultiTurnRLTrainer.__dict__["_pipeline_caller_name"], str) + + def test_resolver_returns_the_declared_name(self): + """Nothing else in the resolver produces "create_job". + + The duck-typed branches return run/train/transform/tune, so without this + declaration `JobStep`'s guard is unreachable from any producer. + """ + assert retrieve_caller_name(self._make_trainer()) == "create_job" + + def test_resolver_does_not_mistake_the_trainer_for_a_model_trainer(self): + trainer = self._make_trainer() + # The trainer carries BaseTrainer's attributes the train duck-typing keys on, + # so the declared name winning here is the resolution ORDER under test. + assert retrieve_caller_name(trainer) != "train" + + def test_declaration_wins_over_the_model_trainer_branch(self): + """Why the resolver checks the declaration first, for this class specifically. + + Load-bearing now, not hypothetically. This class has a `train()` method and, by + subclassing `BaseTrainer`, already carries `input_data_config`, which is one of + the two markers the CreateTrainingJob branch accepts, so it already matches that + branch structurally. The declaration is the only thing keeping it out. Removing + it does not yield `None`, it resolves "train" and composes a `TrainingStep` over + `CreateJob` arguments. + + Setting `training_image` here adds the branch's other marker, so the instance + carries both of them rather than only the inherited one. The declaration still + wins, which is what pins the ordering. + """ + trainer = self._make_trainer() + trainer.training_image = "123456789012.dkr.ecr.us-west-2.amazonaws.com/img:latest" + assert retrieve_caller_name(trainer) == "create_job" + + # --- capture instead of submission ---------------------------------------- + + def test_train_returns_step_arguments_and_submits_nothing(self): + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + with self._pinned_name(), patch( + "sagemaker.train.multi_turn_rl_trainer.Job.create" + ) as mock_create: + step_args = trainer.train() + assert isinstance(step_args, _StepArguments) + assert step_args.caller_name == "create_job" + mock_create.assert_not_called() + + def test_executing_the_captured_call_populates_the_session_context(self): + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + with patch("sagemaker.train.multi_turn_rl_trainer.Job.create") as mock_create: + context = self._capture(trainer, session) + mock_create.assert_not_called() + assert isinstance(context, _JobStepArguments) + assert context.caller_name == "create_job" + + def test_captured_request_carries_the_create_job_envelope(self): + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + args = self._capture(trainer, session).args + assert set(args) == { + "JobName", + "RoleArn", + "JobCategory", + "JobConfigSchemaVersion", + "JobConfigDocument", + "Tags", + } + assert args["JobCategory"] == JOB_CATEGORY + assert args["JobConfigSchemaVersion"] == JOB_CONFIG_SCHEMA_VERSION + assert args["RoleArn"] == ROLE_ARN + + def test_session_and_region_are_not_request_members(self): + """They are `Job.create`'s client-resolution arguments, not `CreateJob` keys.""" + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + args = self._capture(trainer, session).args + assert "session" not in args + assert "region" not in args + + def test_job_name_is_left_in_the_captured_request(self): + """Popping it would make `trim_request_dict`'s custom-prefix branch unreachable. + + That branch is `if job_key in request_dict:`, so a producer that pops the key + silently gives a user who opted into `use_custom_job_prefix` nothing at all. + The upstream finetune trainers pop it; this does not. + """ + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + assert self._capture(trainer, session).args["JobName"] == PINNED_JOB_NAME + + def test_capture_leaves_latest_job_untouched(self): + """The early return skips the assignment, so a prior handle is not clobbered.""" + previous = MagicMock() + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session, latest_job=previous) + self._capture(trainer, session) + assert trainer._latest_job is previous + + def test_output_model_package_arn_reports_none_after_capture(self): + """The only property reading `_latest_job` already guards on None.""" + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + self._capture(trainer, session) + assert trainer._latest_job is None + assert trainer.output_model_package_arn is None + + def test_wait_is_overridden_and_announced(self): + """`wait` is meaningless under a pipeline session, and is not silently dropped. + + `runnable_by_pipeline` forces it to False and warns before the body runs, so + the override is announced by the framework rather than absorbed here. + """ + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + with self._pinned_name(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + step_args = trainer.train(wait=True) + assert step_args.func_kwargs.get("wait") is False + assert any( + "No Wait" in str(w.message) for w in caught + ), [str(w.message) for w in caught] + + # --- parity with the direct submission path ------------------------------- + + def _direct_create_kwargs(self, trainer): + """Run the direct path and record what it hands to `Job.create`.""" + recorded = {} + + def _record(**kwargs): + recorded.update(kwargs) + return MagicMock() + + with self._pinned_name(), patch( + "sagemaker.train.multi_turn_rl_trainer.Job.create", side_effect=_record + ), patch( + "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job", + return_value=MagicMock(), + ): + trainer.train(wait=False) + return recorded + + @staticmethod + def _job_create_wire_request(create_kwargs, job_config_document): + """Push `create_kwargs` through the real `Job.create` and capture the wire dict. + + `session` and `region` are dropped: they select the client rather than form part + of the request, which is exactly why the captured request omits them. The + document is substituted so the comparison isolates the envelope from the one + deliberate difference between the routes (see the document tests below). + """ + replay = {k: v for k, v in create_kwargs.items() if k not in ("session", "region")} + replay["job_config_document"] = job_config_document + mock_client = MagicMock() + with patch.object(Base, "get_sagemaker_client", return_value=mock_client): + try: + Job.create(**replay) + except Exception: + pass + assert mock_client.create_job.call_args is not None, "create_job was not reached" + return mock_client.create_job.call_args.kwargs + + def test_captured_request_matches_what_job_create_would_send(self): + """Criterion: the hand-built dict equals the resource layer's own output. + + Demonstrated by running both routes and pushing the direct route's arguments + through the real `Job.create`, `populate_chained_attributes` and `serialize`, + rather than by reasoning about what they do. + """ + tags = [Tag(key="Project", value="beta")] + direct = self._direct_create_kwargs( + self._make_trainer(sagemaker_session=self._direct_session(), tags=tags) + ) + assert set(direct) == { + "job_name", + "job_category", + "role_arn", + "job_config_schema_version", + "job_config_document", + "tags", + "session", + "region", + } + + session = self._pipeline_session() + captured = dict(self._capture( + self._make_trainer(sagemaker_session=session, tags=tags), session + ).args) + # The document is captured as the raw config dict for JobStep to scope and + # encode; compare it to the direct route's config, and the rest to the wire. + document = captured.pop("JobConfigDocument") + assert document == json.loads(direct["job_config_document"]) + wire = self._job_create_wire_request(direct, direct["job_config_document"]) + wire.pop("JobConfigDocument") + assert captured == wire + + def test_captured_request_matches_job_create_with_both_tag_forms(self): + """The studio dicts and a user `Tag` object converge on the same wire form.""" + tags = [Tag(key="Project", value="beta"), {"key": "Team", "value": "verse"}] + direct = self._direct_create_kwargs( + self._make_trainer(sagemaker_session=self._direct_session(), tags=tags) + ) + session = self._pipeline_session() + captured = dict(self._capture( + self._make_trainer(sagemaker_session=session, tags=tags), session + ).args) + document = captured.pop("JobConfigDocument") + assert document == json.loads(direct["job_config_document"]) + wire = self._job_create_wire_request(direct, direct["job_config_document"]) + wire.pop("JobConfigDocument") + assert captured == wire + assert captured["Tags"] == [ + {"Key": "sagemaker-sdk:jumpstart-model-id", "Value": "test-model"}, + {"Key": "sagemaker-sdk:jumpstart-hub-name", "Value": "SageMakerPublicHub"}, + {"Key": "Project", "Value": "beta"}, + {"Key": "Team", "Value": "verse"}, + ] + + def test_tags_are_always_present_in_the_captured_request(self): + """`Job.create` passes `tags` unconditionally, and `serialize` keeps `[]`.""" + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session, tags=None) + args = self._capture(trainer, session).args + # The two JumpStart tags are always present, so this is never the empty case, + # but the key is unconditional either way. + assert "Tags" in args + assert {t["Key"] for t in args["Tags"]} == { + "sagemaker-sdk:jumpstart-model-id", + "sagemaker-sdk:jumpstart-hub-name", + } + + # --- the JobConfigDocument mechanism -------------------------------------- + + def test_document_describes_the_same_config_on_both_routes(self): + """The capture route hands JobStep the raw dict the direct route serializes.""" + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + captured = self._capture(trainer, session).args["JobConfigDocument"] + direct = self._make_trainer(sagemaker_session=self._direct_session())._build_job_config_document() + assert isinstance(captured, dict) + assert captured == json.loads(direct) + + def test_direct_document_is_byte_identical_to_the_previous_behaviour(self): + """Criterion: the direct path is unchanged, indentation included.""" + trainer = self._make_trainer(sagemaker_session=self._direct_session()) + assert trainer._build_job_config_document() == json.dumps( + trainer._build_job_config(), indent=2 + ) + + def test_pipeline_variable_in_the_job_config_survives_capture(self): + """`serialize` keeps a `PipelineVariable` intact inside the captured dict, + so JobStep can encode it into the definition.""" + session = self._pipeline_session() + trainer = self._make_trainer( + sagemaker_session=session, s3_output_path=ParameterString(name="OutputPath") + ) + document = self._capture(trainer, session).args["JobConfigDocument"] + assert isinstance(document, dict) + assert isinstance(document["OutputDataConfig"]["S3OutputPath"], ParameterString) + + def test_pipeline_variable_document_would_be_a_type_error_unencoded(self): + """Locks the reason the encoder is needed rather than assuming it.""" + trainer = self._make_trainer( + sagemaker_session=self._direct_session(), s3_output_path=ParameterString(name="OutputPath") + ) + with pytest.raises(TypeError, match="not JSON serializable"): + trainer._build_job_config_document() + + # --- the direct path is unchanged ----------------------------------------- + + def test_direct_path_still_submits_and_returns_a_job(self): + """Criterion: a normal session behaves exactly as before.""" + job_handle = MagicMock() + trainer = self._make_trainer(sagemaker_session=self._direct_session()) + with self._pinned_name(), patch( + "sagemaker.train.multi_turn_rl_trainer.Job.create" + ) as mock_create, patch( + "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job", + return_value=job_handle, + ): + returned = trainer.train(wait=False) + mock_create.assert_called_once() + assert returned is job_handle + assert trainer._latest_job is job_handle + job_handle.wait.assert_not_called() + + def test_direct_path_honours_wait(self): + job_handle = MagicMock() + trainer = self._make_trainer(sagemaker_session=self._direct_session()) + with self._pinned_name(), patch( + "sagemaker.train.multi_turn_rl_trainer.Job.create" + ), patch( + "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job", + return_value=job_handle, + ): + trainer.train(wait=True) + job_handle.wait.assert_called_once() + + def test_direct_path_forwards_tags_without_normalisation(self): + """Unchanged: `Job.create` performs the coercion on this route.""" + tags = [Tag(key="Project", value="beta")] + direct = self._direct_create_kwargs( + self._make_trainer(sagemaker_session=self._direct_session(), tags=tags) + ) + assert direct["tags"][-1] is tags[0] + + def test_dry_run_is_overridden_and_announced_under_a_pipeline_session(self, caplog): + """`dry_run` cannot validate here -- its path encodes the document eagerly. + + The capture still happens (nothing is submitted either way) and the + override is announced rather than silent, like the `wait` override. + """ + session = self._pipeline_session() + trainer = self._make_trainer(sagemaker_session=session) + with self._pinned_name(), caplog.at_level("WARNING"): + step_args = trainer.train(dry_run=True) + execute_job_functions(step_args) + assert session.context is not None, "dry_run suppressed the capture" + assert any("dry_run is ignored" in r.message for r in caplog.records)