Skip to content

Commit 462e07a

Browse files
author
jkasiraj
committed
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.
1 parent 4f3cc56 commit 462e07a

16 files changed

Lines changed: 1740 additions & 111 deletions

File tree

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Encode a ``CreateJob`` ``JobConfigDocument`` that may carry pipeline variables.
14+
15+
Unlike the other create APIs, ``CreateJob`` flattens the job configuration into
16+
one JSON **string** member, so nothing inside it sits at a position the pipeline
17+
definition serializer can rewrite, and ``json.dumps`` on a ``PipelineVariable``
18+
raises ``TypeError``. The encoder therefore emits the JSON text as a ``Join`` of
19+
literal fragments and the variables, which the pipeline service resolves during
20+
execution like any other variable leaf.
21+
22+
Limitations, both inherent to splicing values into finished JSON text:
23+
24+
* A variable is spliced in as raw text, so a value resolving to text containing
25+
a double quote or a backslash produces a malformed document. Pipeline
26+
variables resolve to ARNs, S3 URIs and identifiers in practice.
27+
* A variable always lands in a JSON **string** position; one standing in for a
28+
number or boolean reaches the service quoted. Parameterise string-valued
29+
fields only.
30+
"""
31+
from __future__ import absolute_import
32+
33+
import json
34+
import re
35+
import uuid
36+
from typing import List, Union
37+
38+
from sagemaker.core.helper.pipeline_variable import PipelineVariable
39+
40+
41+
def convert_job_config_document_to_string(job_config) -> Union[str, PipelineVariable]:
42+
"""Convert a job config dict to the string-typed ``JobConfigDocument`` value.
43+
44+
Args:
45+
job_config (Dict[str, Any]): The job configuration. May hold
46+
``PipelineVariable`` values at any depth.
47+
48+
Returns:
49+
Union[str, PipelineVariable]: ``json.dumps(job_config)`` when the config
50+
holds no pipeline variable. Otherwise a ``Join`` over the JSON text and
51+
those variables -- not yet a string, but the expression the pipeline
52+
service resolves to the document string during execution.
53+
"""
54+
# Import locally: sagemaker.core.workflow.functions imports from this package's
55+
# entities, and a module-level import here would be a cycle.
56+
from sagemaker.core.workflow.functions import Join
57+
58+
variables: List[PipelineVariable] = []
59+
# A run-unique prefix, drawn from characters json.dumps never escapes, so the
60+
# placeholder survives serialisation verbatim and cannot collide with real content.
61+
token_prefix = "__sagemaker_pipeline_variable_%s_" % uuid.uuid4().hex
62+
63+
def _placeholder(obj):
64+
# json.dumps invokes this for exactly the objects it cannot serialise, so
65+
# it drives the traversal: each variable is recorded and replaced by a
66+
# numbered placeholder in the emitted text.
67+
if isinstance(obj, PipelineVariable):
68+
variables.append(obj)
69+
return "%s%d__" % (token_prefix, len(variables) - 1)
70+
raise TypeError(
71+
"Object of type %s is not JSON serializable" % obj.__class__.__name__
72+
)
73+
74+
document = json.dumps(job_config, default=_placeholder)
75+
if not variables:
76+
return document
77+
78+
# With the capturing group, re.split alternates literal text and placeholder
79+
# indices: [literal, index, literal, index, ..., literal].
80+
pieces = re.split(re.escape(token_prefix) + r"(\d+)__", document)
81+
if (len(pieces) - 1) // 2 != len(variables):
82+
# Unreachable: every placeholder is a plain string value, so it must survive
83+
# json.dumps. Raised rather than silently shipping a literal placeholder.
84+
raise ValueError(
85+
"[PySDK Error] Could not encode JobConfigDocument: %d of %d pipeline "
86+
"variables were not found in the serialised document."
87+
% (len(variables) - (len(pieces) - 1) // 2, len(variables))
88+
)
89+
values: List[Union[str, PipelineVariable]] = [
90+
variables[int(piece)] if index % 2 else piece
91+
for index, piece in enumerate(pieces)
92+
if index % 2 or piece
93+
]
94+
return Join(on="", values=values)
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Shared capture paths for producers running under a ``PipelineSession``."""
14+
from __future__ import absolute_import
15+
16+
from typing import Any, Dict, List, Optional
17+
18+
from sagemaker.core.apiutils._boto_functions import to_pascal_case
19+
from sagemaker.core.shapes import Tag
20+
from sagemaker.core.utils.utils import serialize
21+
22+
23+
def capture_create_job_request(
24+
pipeline_session,
25+
*,
26+
job_name: str,
27+
role_arn: str,
28+
job_category: str,
29+
schema_version: str,
30+
job_config: Dict[str, Any],
31+
tags: Optional[List[Any]] = None,
32+
customer_details: Optional[Dict[str, Any]] = None,
33+
) -> None:
34+
"""Capture a ``CreateJob`` request for ``JobStep`` composition instead of submitting it.
35+
36+
``JobName`` stays in the request: ``JobStep``'s custom-job-prefix handling
37+
requires the key to be present. ``job_config`` is the raw config dict, passed
38+
through untouched -- ``JobStep`` scopes ``OutputDataConfig.S3OutputPath`` per
39+
execution and encodes the document at definition time, and the direct path
40+
encodes the same dict with ``json.dumps``, which keeps ``None`` as ``null``.
41+
42+
Args:
43+
pipeline_session (PipelineSession): The capturing session.
44+
job_name (str): Client-minted name; replaced by the service at execution.
45+
role_arn (str): The execution role for the job.
46+
job_category (str): ``CreateJob`` job category.
47+
schema_version (str): ``JobConfigDocument`` schema version.
48+
job_config (Dict[str, Any]): The raw job configuration dict.
49+
tags (Optional[List[Any]]): Tags in the caller's wire form; ``serialize``
50+
emits ``Tag`` shapes as ``{"Key", "Value"}``.
51+
customer_details (Optional[Dict[str, Any]]): ``CustomerDetails`` envelope
52+
member for producers whose direct path sends it (data preparation).
53+
Session-derived, so it is stable across executions.
54+
"""
55+
request: Dict[str, Any] = {
56+
"JobName": job_name,
57+
"RoleArn": role_arn,
58+
"JobCategory": job_category,
59+
"JobConfigSchemaVersion": schema_version,
60+
}
61+
if customer_details is not None:
62+
request["CustomerDetails"] = customer_details
63+
if tags is not None:
64+
request["Tags"] = tags
65+
captured = serialize(request)
66+
captured["JobConfigDocument"] = job_config
67+
pipeline_session._intercept_create_request(captured, None, "create_job")
68+
69+
70+
def capture_training_request(pipeline_session, create_args: Dict[str, Any]) -> None:
71+
"""Capture a ``CreateTrainingJob`` request for ``TrainingStep`` composition.
72+
73+
``create_args`` are ``TrainingJob.create`` keyword arguments with data channels
74+
already resolved, so datasets must be concrete at authoring time; only the job
75+
name is deferred to execution. Client-resolution members (``session``,
76+
``region``) are dropped, and dict-form tags are coerced through the ``Tag``
77+
model so ``serialize`` emits the wire form ``TrainingJob.create`` produces.
78+
79+
Args:
80+
pipeline_session (PipelineSession): The capturing session.
81+
create_args (Dict[str, Any]): ``TrainingJob.create`` keyword arguments.
82+
"""
83+
pipeline_args = {k: v for k, v in create_args.items() if k not in ("session", "region")}
84+
pipeline_args.pop("training_job_name", None)
85+
request = {to_pascal_case(k): v for k, v in pipeline_args.items()}
86+
if request.get("Tags"):
87+
# Dict-form tags arrive in both shapes: lowercase keys from the SDK's own
88+
# tag builders, PascalCase from callers passing wire-form dicts through.
89+
request["Tags"] = [
90+
Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value")))
91+
if isinstance(tag, dict)
92+
else tag
93+
for tag in request["Tags"]
94+
]
95+
pipeline_session._intercept_create_request(serialize(request), None, "train")

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

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,12 @@ def wrapper(*args, **kwargs):
357357
return wrapper
358358

359359

360+
# Names a producer may declare via `_pipeline_caller_name`. Restricted so a producer
361+
# cannot declare a name the structural checks below already resolve and compose the
362+
# wrong step over a different create API.
363+
_DECLARABLE_CALLER_NAMES = frozenset({"create_job"})
364+
365+
360366
def retrieve_caller_name(job_instance):
361367
"""Convenience method for runnable_by_pipeline decorator
362368
@@ -367,26 +373,45 @@ def retrieve_caller_name(job_instance):
367373
job_instance: A job class instance, one of the following types:
368374
- Processor (from sagemaker.core.processing)
369375
- ModelTrainer (from sagemaker.train.model_trainer)
376+
- a V3 finetune trainer (SFT/DPO/RLVR/RLAIF, from sagemaker.train), which
377+
subclasses BaseTrainer and carries no training_image
370378
- Transformer (from sagemaker.core.transformer)
371379
- HyperparameterTuner (from sagemaker.train.tuner)
380+
- a CreateJob-backed producer that declares `_pipeline_caller_name`
381+
(multi-turn RL)
372382
373383
Note:
374384
This function uses duck typing to avoid importing from Train package,
375385
which would create architecture violations (Core should not depend on Train).
376386
Instead of isinstance checks, we check for characteristic attributes/methods.
387+
388+
The CreateJob family is the exception, and declares its name rather than
389+
being duck typed, because its producers share no common signature to key on.
390+
See `_DECLARABLE_CALLER_NAMES`.
377391
"""
378392

379393
from sagemaker.core.processing import Processor
380394
from sagemaker.core.transformer import Transformer
381395

382396
# from sagemaker.utils.automl.automl import AutoML
383397

398+
# Resolved FIRST, and the order is load-bearing: MultiTurnRLTrainer inherits
399+
# `input_data_config` from BaseTrainer, so the train branch below would claim it
400+
# and compose a TrainingStep over CreateJob arguments.
401+
declared_caller_name = getattr(job_instance, "_pipeline_caller_name", None)
402+
if declared_caller_name in _DECLARABLE_CALLER_NAMES:
403+
return declared_caller_name
404+
384405
if isinstance(job_instance, Processor):
385406
return "run"
386407

387-
# Duck typing for ModelTrainer: has 'train' method and 'training_image' attribute
388-
# This avoids importing from sagemaker.train which would violate architecture
389-
if hasattr(job_instance, "train") and hasattr(job_instance, "training_image"):
408+
# Duck typing for the CreateTrainingJob family, avoiding an import from
409+
# sagemaker.train. Either marker suffices: `training_image` is ModelTrainer's,
410+
# `input_data_config` is BaseTrainer's (the V3 finetune trainers resolve their
411+
# image from the recipe and carry no `training_image`).
412+
if hasattr(job_instance, "train") and (
413+
hasattr(job_instance, "training_image") or hasattr(job_instance, "input_data_config")
414+
):
390415
return "train"
391416

392417
if isinstance(job_instance, Transformer):
Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Tests for the ``JobConfigDocument`` encoder."""
14+
from __future__ import absolute_import
15+
16+
import json
17+
18+
import pytest
19+
20+
from sagemaker.core.workflow.functions import Join
21+
from sagemaker.core.workflow.job_config_document import convert_job_config_document_to_string
22+
from sagemaker.core.workflow.parameters import ParameterString
23+
24+
25+
def _resolve(join, values):
26+
"""Splice concrete values into a Join, standing in for the pipeline service."""
27+
return "".join(
28+
values[piece.name] if isinstance(piece, ParameterString) else piece for piece in join.values
29+
)
30+
31+
32+
def test_config_without_variables_is_plain_json():
33+
config = {"A": "x", "B": {"C": [1, 2]}, "D": None}
34+
document = convert_job_config_document_to_string(config)
35+
assert isinstance(document, str)
36+
assert json.loads(document) == config
37+
38+
39+
def test_none_survives_as_null():
40+
document = convert_job_config_document_to_string({"A": None})
41+
assert document == '{"A": null}'
42+
43+
44+
def test_single_variable_becomes_a_join_over_the_json_text():
45+
path = ParameterString(name="OutputPath")
46+
document = convert_job_config_document_to_string({"Out": {"S3OutputPath": path}, "N": 1})
47+
assert isinstance(document, Join)
48+
assert document.on == ""
49+
assert document.values == ['{"Out": {"S3OutputPath": "', path, '"}, "N": 1}']
50+
assert json.loads(_resolve(document, {"OutputPath": "s3://b/p"})) == {
51+
"Out": {"S3OutputPath": "s3://b/p"},
52+
"N": 1,
53+
}
54+
55+
56+
def test_multiple_variables_keep_document_order():
57+
role = ParameterString(name="Role")
58+
path = ParameterString(name="OutputPath")
59+
document = convert_job_config_document_to_string({"Role": role, "Out": path})
60+
assert [v for v in document.values if isinstance(v, ParameterString)] == [role, path]
61+
assert json.loads(_resolve(document, {"Role": "r", "OutputPath": "p"})) == {
62+
"Role": "r",
63+
"Out": "p",
64+
}
65+
66+
67+
def test_adjacent_variables_emit_no_empty_literal_between_them():
68+
first = ParameterString(name="First")
69+
second = ParameterString(name="Second")
70+
document = convert_job_config_document_to_string({"L": [first, second]})
71+
assert document.values == ['{"L": ["', first, '", "', second, '"]}']
72+
assert "" not in document.values
73+
74+
75+
def test_same_variable_used_twice_is_spliced_twice():
76+
path = ParameterString(name="OutputPath")
77+
document = convert_job_config_document_to_string({"A": path, "B": path})
78+
assert document.values.count(path) == 2
79+
assert json.loads(_resolve(document, {"OutputPath": "p"})) == {"A": "p", "B": "p"}
80+
81+
82+
def test_variable_lands_in_a_string_position():
83+
"""A variable standing in for a number reaches the service quoted."""
84+
count = ParameterString(name="Count")
85+
document = convert_job_config_document_to_string({"N": count})
86+
assert json.loads(_resolve(document, {"Count": "3"})) == {"N": "3"}
87+
88+
89+
def test_quote_in_a_resolved_value_malforms_the_document():
90+
"""The documented limitation: values are spliced as raw text."""
91+
name = ParameterString(name="Name")
92+
document = convert_job_config_document_to_string({"Name": name})
93+
with pytest.raises(json.JSONDecodeError):
94+
json.loads(_resolve(document, {"Name": 'a"b'}))
95+
96+
97+
def test_unserialisable_non_variable_still_raises_type_error():
98+
with pytest.raises(TypeError, match="not JSON serializable"):
99+
convert_job_config_document_to_string({"A": object()})

0 commit comments

Comments
 (0)