Skip to content

feat: add JobStep and shared pipeline capture for CreateJob producers - #6292

Merged
zhaoqizqwang merged 2 commits into
aws:masterfrom
jkasiraj:feat/jobstep-pipeline-capture
Sep 30, 2026
Merged

zhaoqizqwang merged 2 commits into
aws:masterfrom
jkasiraj:feat/jobstep-pipeline-capture

Conversation

@jkasiraj

@jkasiraj jkasiraj commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

Issue #, if available:

Description of changes:

add support for CreateJob to be run as a native sagemaker pipelines step

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@jkasiraj jkasiraj changed the title Add JobStep and shared pipeline capture for CreateJob producers feat: add JobStep and shared pipeline capture for CreateJob producers Sep 18, 2026
@jkasiraj
jkasiraj marked this pull request as ready for review September 18, 2026 19:22
@jkasiraj jkasiraj self-assigned this Sep 18, 2026
@github-actions

github-actions Bot commented Sep 18, 2026 •

Copy link
Copy Markdown

Claude finished @jkasiraj's task in 4m 13s —— View job


Review: JobStep and shared pipeline capture for CreateJob producers

  • Read PR diff
  • Analyze job_config_document.py encoder correctness
  • Analyze pipeline_capture.py shared util
  • Analyze retrieve_caller_name changes
  • Analyze JobStep implementation
  • Analyze trainer changes (MTRL + finetune trainers)
  • Post review

This is a well-constructed, thoroughly-tested PR. The design is sound and I found no blocking correctness bugs. The heavy documentation of intent (ordering guards, encoder limitations, the JobName retention rationale) is genuinely helpful, and the test suites cover the important paths — declared-name resolution ordering, capture-vs-direct parity, tag coercion measured against the real Job.create wire form, and the PipelineVariable round-trip.

Note: the mcp__github_inline_comment__create_inline_comment tool isn't available to me in this run, so findings are inline below with file/line references instead of as PR review comments.

Observations (non-blocking)

1. serialize() drops None-valued keys inside JobConfigDocument, diverging from the direct path — sagemaker-core/.../workflow/pipeline_capture.py:171

capture_create_job_request stores serialize(request), and _serialize_dict drops any key whose value serializes to None at every depth. The direct submission path builds the document with json.dumps(config, indent=2) (via _build_job_config_document), which keeps null values. So for a config that contains an explicit None, the capture route silently omits the field while the direct route sends null. It's latent today because MTRL's _build_*_config methods omit optional keys rather than setting them to None (which is why test_document_describes_the_same_config_on_both_routes passes), but it's a real behavioral difference between the two routes that a future config field set to None would expose. Worth a comment at minimum, since the tests assert parity between the routes.

2. MTRL train() docstring return description is slightly imprecise — sagemaker-train/.../multi_turn_rl_trainer.py:302-305

The docstring says it returns None under a PipelineSession. In practice the user-facing call goes through runnable_by_pipeline, which returns a _StepArguments (the None is only what the wrapped body returns internally when execute_job_functions later re-invokes it). test_train_returns_step_arguments_and_submits_nothing confirms the caller receives _StepArguments. Minor — consider clarifying that the capture yields step arguments to the caller.

Things I checked that are correct

  • Encoder (job_config_document.py): PipelineVariable is an abc.ABC (not a str subclass), so json.dumps(..., default=_placeholder) reliably routes every variable through the placeholder path — no silent pass-through. The uuid-seeded token_prefix uses only characters json.dumps never escapes, and re.split with the capturing group reconstructs the Join correctly (verified the even/odd indexing and the empty-literal filter). The documented limitations (values containing "/\, and non-string types landing quoted) are reasonable for the ARN/S3-URI values in scope.
  • retrieve_caller_name ordering: resolving _pipeline_caller_name before the structural train branch is load-bearing — MTRL inherits input_data_config from BaseTrainer and would otherwise be misclassified as a TrainingStep. The _DECLARABLE_CALLER_NAMES allowlist correctly prevents a producer from claiming train/transform/etc. The widened train branch (training_image OR input_data_config) is fail-safe: train stays a required conjunct, so DataPreparer/GenAIEvaluator shapes still resolve to None.
  • Tag coercion parity: the shared capture_training_request / capture_create_job_request produce the same wire form as Job.create/TrainingJob.create, and empty [] is preserved. The refactor of the four finetune trainers into capture_training_request is behavior-preserving.
  • JobStep: trim_request_dict is imported at module scope; keeping JobName in the request (rather than popping it like the finetune trainers) is deliberately correct so use_custom_job_prefix remains reachable. Re-entrancy of arguments is safe because execute_job_functions rebuilds context.args on each call.
  • No prompt-injection content in the diff/description.

Overall: looks good to merge; the two observations above are optional polish.

rohangujarathi
rohangujarathi previously approved these changes Sep 18, 2026
@github-actions

github-actions Bot commented Sep 18, 2026 •

Copy link
Copy Markdown

Claude finished @jkasiraj's task in 2m 28s —— View job


Review: JobStep and shared pipeline capture for CreateJob producers

  • Read PR diff
  • Analyze job_config_document.py encoder
  • Analyze pipeline_capture.py
  • Analyze retrieve_caller_name / pipeline_context.py
  • Analyze JobStep in steps.py
  • Analyze trainer changes (MTRL + finetune)
  • Check test coverage
  • Post review

This is a solid, carefully-documented PR. I found no blocking correctness bugs. The design is sound, the intent comments (ordering guard, encoder limitations, JobName retention) are genuinely useful, and the test suites cover the load-bearing paths — declared-name resolution ordering, capture-vs-direct parity measured against the real Job.create/TrainingJob.create wire form, tag coercion, and the PipelineVariable round-trip.

Note: mcp__github_inline_comment__create_inline_comment isn't available in this run, so findings are listed below with file/line references instead of as inline PR comments.

Non-blocking observations

1. serialize() drops None-valued keys inside JobConfigDocument, diverging from the direct route — sagemaker-core/.../workflow/pipeline_capture.py:171

capture_create_job_request stores serialize(request), and _serialize_dict/_serialize_list (sagemaker-core/.../utils/utils.py:526-561) drop any key/element whose value serializes to None, at every depth. The direct submission path builds the document with json.dumps(config, indent=2) in _build_job_config_document, which keeps null. So a config carrying an explicit None is silently omitted on the capture route but sent as null on the direct route. It's latent today because MTRL's _build_*_config methods omit optional keys rather than setting them to None (which is why test_document_describes_the_same_config_on_both_routes passes), but it is a real behavioral difference between two routes the tests assert parity on — a future config field set to None would expose it. Worth at least a comment noting the constraint.

2. No direct unit test for the JobConfigDocument encoder — sagemaker-core/.../workflow/job_config_document.py

convert_job_config_document_to_string is the subtlest new code (uuid-seeded placeholder, re.split with a capturing group, the "unreachable" ValueError guard, the empty-literal filter). It's currently exercised only indirectly — the integ test asserts Std:Join is present, and test_multi_turn_rl_trainer round-trips through a hand-rolled _resolve. A dedicated unit test covering the no-variable early return, a single variable, multiple/adjacent variables, and the documented quote/backslash limitation would lock the branches that the indirect tests don't reach. pipeline_capture.py similarly has only indirect coverage, though its logic is simpler.

3. train() docstring return description is slightly imprecise — sagemaker-train/.../multi_turn_rl_trainer.py:302-305

The docstring says it returns None under a PipelineSession. The user-facing call goes through runnable_by_pipeline, which returns a _StepArguments; the None is only what the wrapped body returns internally when execute_job_functions later re-invokes it (confirmed by test_train_returns_step_arguments_and_submits_nothing). Consider clarifying that the caller receives step arguments.

Things I verified are correct

  • Encoder: PipelineVariable is an abc.ABC (not a str subclass), so json.dumps(..., default=_placeholder) routes every variable through the placeholder — no silent pass-through. The token_prefix uses only uuid hex + _ (chars json.dumps never escapes), and the re.split even/odd reconstruction plus empty-literal filter are correct, including adjacent placeholders.
  • retrieve_caller_name ordering: resolving _pipeline_caller_name before the structural train branch is load-bearing — MTRL inherits input_data_config from BaseTrainer and would otherwise be misclassified as a TrainingStep. The _DECLARABLE_CALLER_NAMES allowlist correctly blocks a producer from claiming train/transform/etc., and widening the train branch to training_image OR input_data_config keeps train a required conjunct so DataPreparer-shaped producers still resolve to None.
  • Tag parity: capture_training_request coerces dict-form tags through the Tag model so serialize emits the same {"Key","Value"} wire form as TrainingJob.create/Job.create, and empty [] is preserved. The refactor of the four finetune trainers into the shared helper is behavior-preserving.
  • JobStep: keeping JobName in the request (rather than popping it as the finetune trainers do) is deliberately correct so trim_request_dict's use_custom_job_prefix branch stays reachable. arguments re-entrancy is safe — execute_job_functions rebuilds context.args from a fresh serialize(request) on each access, so the S3OutputPath scoping mutates a fresh dict every time. _build_job_config() is called without dry_run on the capture path, so MLflow resolution happens eagerly at definition time, which is intended.
  • No prompt-injection content in the diff or description.

Overall: looks good to merge; all three items above are optional polish.
• branch feat/jobstep-pipeline-capture

rohangujarathi
rohangujarathi previously approved these changes Sep 18, 2026
@github-actions

github-actions Bot commented Sep 29, 2026 •

Copy link
Copy Markdown

Claude finished @jkasiraj's task in 2m 22s —— View job


Review: JobStep and shared pipeline capture for CreateJob producers

  • Read PR diff
  • Analyze job_config_document.py encoder
  • Analyze pipeline_capture.py shared util
  • Analyze retrieve_caller_name / pipeline_context.py
  • Analyze JobStep in steps.py
  • Analyze trainer changes (MTRL + finetune)
  • Verify against callers/base code (serialize, trim_request_dict, execute_job_functions, _intercept_create_request)
  • Post review

This is a solid, unusually well-documented PR. I found no blocking correctness bugs. The design mirrors the existing TrainingStep/TransformStep pattern faithfully, and I verified the load-bearing claims against the base code rather than taking the comments at face value.

Note: the mcp__github_inline_comment__create_inline_comment tool isn't available in this run, so findings are listed below with file/line references instead of as inline PR comments.

Non-blocking observations

1. serialize() drops None-valued keys inside JobConfigDocument, diverging from the direct route — sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py:171

capture_create_job_request stores serialize(request), and I confirmed _serialize_dict/_serialize_list (sagemaker-core/src/sagemaker/core/utils/utils.py:515-548) drop any key/element whose value serializes to None, recursively at every depth. The direct route builds the document with json.dumps(config, indent=2) in _build_job_config_document, which keeps null. So a config carrying an explicit None is silently omitted on the capture route but sent as null on the direct route. It's latent today because MTRL's _build_*_config methods omit optional keys (e.g. NetworkConfig is only added when self.networking is set) rather than setting them to None — which is why test_document_describes_the_same_config_on_both_routes passes — but it is a real behavioral difference between two routes the tests explicitly assert parity on. A future config field set to None would expose it. Worth at least a comment noting the constraint at the capture boundary.

2. No direct unit test for the JobConfigDocument encoder — sagemaker-core/src/sagemaker/core/workflow/job_config_document.py

convert_job_config_document_to_string is the subtlest new code (uuid-seeded placeholder, re.split with a capturing group, the "unreachable" ValueError guard, the empty-literal filter for adjacent placeholders). It's currently exercised only indirectly — the integ test asserts Std:Join is present and test_multi_turn_rl_trainer round-trips through a hand-rolled _resolve. A dedicated unit test covering the no-variable early return, a single variable, multiple/adjacent variables, and the documented quote/backslash limitation would lock the branches the indirect tests don't reach.

3. train() docstring return description is slightly imprecise — sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py:302-305

The docstring says it returns None under a PipelineSession. The user-facing call goes through runnable_by_pipeline, which returns a _StepArguments (pipeline_context.py:353); the None is only what the wrapped body returns internally when execute_job_functions later re-invokes it — confirmed by test_train_returns_step_arguments_and_submits_nothing. Consider clarifying that the caller receives step arguments.

Things I verified are correct

  • Encoder: PipelineVariable is an abc.ABC (not a str subclass), so json.dumps(..., default=_placeholder) reliably routes every variable through the placeholder — no silent pass-through. The token_prefix uses only uuid hex + _ (chars json.dumps never escapes), and the even/odd re.split reconstruction plus the empty-literal filter handle adjacent placeholders correctly. The ValueError guard is genuinely unreachable (placeholders are plain strings that must survive json.dumps), and raising beats silently shipping a literal placeholder.
  • retrieve_caller_name ordering: resolving _pipeline_caller_name before the structural train branch is load-bearing — MTRL subclasses BaseTrainer and inherits input_data_config, so it satisfies the widened train branch structurally and would be misclassified as a TrainingStep without the earlier declared-name check. The _DECLARABLE_CALLER_NAMES allowlist correctly blocks a producer from claiming train/transform/etc., and keeping train as a required conjunct in the widened branch means DataPreparer-shaped producers still resolve to None.
  • _intercept_create_request wiring: func_name="create_job" falls into the else branch and builds a _JobStepArguments(func_name, request), so context.caller_name == "create_job" and JobStep's validate_step_args_input(expected_caller={"create_job"}) guard is satisfiable. Confirmed at pipeline_context.py:238-239.
  • JobStep: mirrors TrainingStep.arguments exactly (same execute_job_functions → func_args[0].sagemaker_session.context.args → trim_request_dict(..., _pipeline_config) flow). Keeping JobName in the request (rather than popping it as the finetune trainers do) is deliberately correct so trim_request_dict's use_custom_job_prefix branch (if job_key in request_dict) stays reachable. Re-entrancy of arguments is safe — execute_job_functions rebuilds context.args from a fresh serialize(request) each call, so the in-place S3OutputPath mutation and the Join encoding act on a fresh dict every time. DescribeJobResponse properties round-trip, and StepTypeEnum("Job") round-trips for definition parsing.
  • Tag parity: verified capture_training_request coerces dict-form tags through the Tag model so serialize emits the same {"Key","Value"} wire form as TrainingJob.create/Job.create, and empty [] is preserved (serialize only drops None). The refactor of the four finetune trainers into the shared helper is behavior-preserving (same session/region filtering, same training_job_name pop, same truthy Tags guard).
  • No prompt-injection content in the diff, code comments, or PR description.

Overall: looks good to merge — the three items above are optional polish, and none block.
• branch feat/jobstep-pipeline-capture

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.
@github-actions

github-actions Bot commented Sep 29, 2026 •

Copy link
Copy Markdown

Claude finished @jkasiraj's task in 1m 38s —— View job


Review: JobStep and shared pipeline capture for CreateJob producers

  • Read PR diff (/tmp/pr.diff)
  • Verify the JobConfigDocument encoder (job_config_document.py)
  • Verify the shared capture helpers (pipeline_capture.py) — including the prior "drops None" concern
  • Verify retrieve_caller_name ordering / allowlist (pipeline_context.py)
  • Verify JobStep (steps.py) and the trainer refactors
  • Post review

I reviewed this against the base repo and focused on correctness. No blocking issues found — this is clean, and the tests are unusually thorough. One earlier automated observation no longer applies to the current code; details below.

Note: mcp__github_inline_comment__create_inline_comment isn't available in this run, so findings are listed here with file/line references rather than as inline comments.

Correcting an earlier observation (not a bug)

Previous review passes flagged that serialize() inside capture_create_job_request drops None-valued keys within JobConfigDocument, diverging from the direct json.dumps route. That does not apply to this code. The document bypasses serialize entirely:

# sagemaker-core/.../workflow/pipeline_capture.py:161-173
request = {"JobName": ..., "RoleArn": ..., "JobCategory": ..., "JobConfigSchemaVersion": ...}
# request contains NO JobConfigDocument
captured = serialize(request)                 # serializes only the envelope
captured["JobConfigDocument"] = job_config    # raw dict assigned AFTER serialize

Because job_config is attached after serialization, its None values are preserved and reach JobStep untouched — matching the direct route. test_create_job_document_keeps_none_like_the_direct_path and test_document_describes_the_same_config_on_both_routes both pin this. So the two routes agree here; no comment needed.

Things I verified are correct

  • Encoder (job_config_document.py): PipelineVariable is an abc.ABC (not a str subclass), so json.dumps(..., default=_placeholder) routes every variable through the placeholder — no silent pass-through. The uuid-hex token_prefix uses only characters json.dumps never escapes, re.split with the capturing group reconstructs literal/variable alternation correctly (multi-digit indices via \d+, adjacent placeholders, repeated variables), and the empty-literal filter is right. The ValueError guard is genuinely unreachable and fails loud rather than shipping a literal placeholder. The two documented limitations (quote/backslash in a resolved value; non-string variables landing quoted) are reasonable for ARN/S3-URI values and are explicitly tested.
  • retrieve_caller_name ordering: resolving _pipeline_caller_name before the structural train branch is load-bearing — MTRL subclasses BaseTrainer and inherits input_data_config, so it satisfies the widened training_image OR input_data_config branch structurally and would misclassify as a TrainingStep otherwise. _DECLARABLE_CALLER_NAMES correctly blocks a producer from claiming train/transform/etc., and train stays a required conjunct so DataPreparer-shaped producers still resolve to None.
  • Tag coercion parity: capture_training_request coerces dict-form tags (both lowercase and PascalCase) through the Tag model so serialize emits the same {"Key","Value"} wire form as TrainingJob.create/Job.create; empty [] is preserved. The refactor of the four finetune trainers into the shared helper is behavior-preserving.
  • JobStep: keeping JobName in the request (rather than popping it like the finetune trainers) is deliberately correct so trim_request_dict's use_custom_job_prefix branch (if job_key in request_dict) stays reachable. S3OutputPath execution scoping only fires for a str path (a PipelineVariable path is intentionally left alone), and arguments re-entrancy is safe because execute_job_functions rebuilds context.args from a fresh serialize(request) each access. StepTypeEnum("Job") round-trips and DescribeJobResponse.JobConfigDocument is correctly a leaf String property.
  • _build_job_config() on the capture path runs with dry_run=False, so MLflow resolution happens eagerly at definition time — intended, and the dry_run-under-PipelineSession override is announced via a warning (test_dry_run_is_overridden_and_announced_under_a_pipeline_session).

Optional polish (non-blocking)

  • train() docstring — sagemaker-train/.../multi_turn_rl_trainer.py:301: the return line reads as if a PipelineSession call returns None; the user-facing call actually returns the _StepArguments from runnable_by_pipeline (the None is only what the wrapped body returns internally). The added lines below it already clarify this, so it's minor.

  • No prompt-injection content in the diff, comments, or PR description.

Overall: looks good to merge.
• branch feat/jobstep-pipeline-capture

@zhaoqizqwang
zhaoqizqwang merged commit 88f70d8 into aws:master Sep 30, 2026
4 of 5 checks passed
@jkasiraj
jkasiraj deleted the feat/jobstep-pipeline-capture branch September 30, 2026 20:23
@github-actions

github-actions Bot commented Sep 30, 2026 •

Copy link
Copy Markdown

Claude finished @jkasiraj's task in 1m 16s —— View job


Review: JobStep and shared pipeline capture for CreateJob producers

  • Read the full PR diff (/tmp/pr.diff)
  • Analyze the JobConfigDocument encoder (job_config_document.py)
  • Analyze the shared capture helpers (pipeline_capture.py)
  • Analyze retrieve_caller_name ordering / allowlist (pipeline_context.py)
  • Analyze JobStep (steps.py) and the trainer refactors
  • Verify test coverage
  • Post review

I reviewed this focused on correctness, API usage, and test coverage. No blocking issues — this is clean, carefully documented, and unusually well tested. The design mirrors the existing TrainingStep/TransformStep pattern faithfully.

Note: the mcp__github_inline_comment__create_inline_comment tool isn't available in this run, so findings are listed below with file/line references instead of as inline comments.

Things I verified are correct

  • Encoder (job_config_document.py) — PipelineVariable is an abc.ABC, not a str subclass, so json.dumps(..., default=_placeholder) reliably routes every variable through the placeholder path (no silent pass-through). The uuid-hex token_prefix uses only characters json.dumps never escapes, so placeholders survive serialization verbatim; re.split with the capturing group reconstructs the literal/variable alternation correctly (verified even/odd indexing, adjacent placeholders via test_adjacent_variables_emit_no_empty_literal_between_them, and repeated variables). The ValueError guard is genuinely unreachable and fails loud rather than shipping a literal placeholder. The two documented limitations (quote/backslash in a resolved value; non-string variables landing quoted) are reasonable for ARN/S3-URI values and are explicitly tested.
  • capture_create_job_request (pipeline_capture.py:171-172) — The JobConfigDocument is assigned to captured after serialize(request), so the raw dict is never passed through serialize. That means None values are preserved (matching the direct json.dumps route) and PipelineVariable leaves survive intact for JobStep to encode. test_create_job_document_keeps_none_like_the_direct_path and test_document_describes_the_same_config_on_both_routes pin this parity. (This resolves the "serialize drops None" concern raised in earlier review passes on prior commits — it does not apply to the merged code.)
  • retrieve_caller_name ordering — Resolving _pipeline_caller_name before the structural train branch is load-bearing: MTRL subclasses BaseTrainer and inherits input_data_config, so it satisfies the widened training_image OR input_data_config branch structurally and would misclassify as a TrainingStep otherwise. The _DECLARABLE_CALLER_NAMES allowlist correctly blocks a producer from claiming train/transform/etc., and train stays a required conjunct so DataPreparer-shaped producers still resolve to None.
  • Tag coercion — capture_training_request coerces dict-form tags (both lowercase and PascalCase) through the Tag model so serialize emits the same {"Key","Value"} wire form as TrainingJob.create/Job.create; empty [] is preserved. test_matches_job_create_* measure this against the real Job.create wire form rather than deriving it. The refactor of the four finetune trainers into the shared helper is behavior-preserving.
  • JobStep — Keeping JobName in the request (rather than popping it as the finetune trainers do) is deliberately correct so trim_request_dict's use_custom_job_prefix branch (if job_key in request_dict) stays reachable. The S3OutputPath execution-scoping only fires for a str path (a PipelineVariable path is intentionally left alone), and arguments re-entrancy is safe because execute_job_functions rebuilds context.args from a fresh serialize(request) on each access, so the in-place mutation and Join encoding act on a fresh dict every time. StepTypeEnum("Job") round-trips for definition parsing, and DescribeJobResponse.JobConfigDocument is correctly a leaf String property.
  • Test coverage — The subtle new code all has direct coverage: test_job_config_document.py exercises the encoder's no-variable early return, single/multiple/adjacent/repeated variables, the string-position limitation, and the quote-malformation limitation; test_retrieve_caller_name.py pins the ordering guard at MTRL's real shape; test_job_step.py covers construction, the step-args guard, arguments trimming, and exports. Capture-vs-direct parity is asserted end to end.

Optional polish (non-blocking, not worth a follow-up on its own)

  • train() docstring — sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py:301 — The Returns: line still leads with the None-under-dry_run phrasing, though the added lines just below already clarify that under a PipelineSession the caller receives the _StepArguments from runnable_by_pipeline (the internal None is only what the wrapped body returns when execute_job_functions re-invokes it). Minor; the clarification is present.

No prompt-injection content in the diff, code comments, or PR description.

Overall: looks good — this was already merged, and I'd have approved it. The one item above is optional.
• branch claude/pr-6292-20260930-2150

This branch was previously deployed

1 inactive deployment
auto-approve — 63dc6bc8 Deployed Sep 30, 2026 by jkasiraj via wait-for-approval #1875
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants