fix(train): complete PipelineSession support for V3 trainers (SFT/DPO…

[nayan3107]

…/RLAIF/RLVR)

Add @runnable_by_pipeline decorator and PascalCase request serialization so TrainingStep can consume step_args from V3 fine-tuning trainers.

Changes:

• Add @runnable_by_pipeline decorator on train() for all 4 trainers
• Build PascalCase serialized request (remove session/region, pop job name)
• Fix Tags to PascalCase (JumpStart returns lowercase key/value)
• Add source_code=None to BaseTrainer (required by get_code_hash)
• Add consumer tests (TrainingStep.arguments validation)
• Add regular session regression tests

Fixes: #6163

Issue #, if available:

Description of changes:

fix(train): complete PipelineSession support for V3 trainers

Follow-up to #6213. Completes PipelineSession integration so V3 fine-tuning
trainers (SFTTrainer, DPOTrainer, RLAIFTrainer, RLVRTrainer) work with
TrainingStep in SageMaker Pipelines.

Fixes #6163

Changes

• @runnable_by_pipeline decorator on train() for all 4 trainers — matches
  ModelTrainer/Processor/Transformer pattern. Decorator captures function reference

  • trainer instance so TrainingStep can re-execute during pipeline compilation.

• PascalCase request serialization inside train() body — removes non-serializable
  session/region, pops training_job_name (engine generates at runtime), converts
  snake_case keys to PascalCase via to_pascal_case() + serialize().

• Tags PascalCase fix — _get_jumpstart_tags() returns lowercase key/value
  dicts (not Tag objects), so _serialize_dict() doesn't auto-convert them. Manually
  converts to Key/Value before serialization.

• source_code = None in BaseTrainer — get_code_hash() in pipeline compilation
  accesses model_trainer.source_code directly. ModelTrainer defines this field;
  BaseTrainer didn't. Fine-tuning trainers use built-in containers with no customer code.

Testing

Unit tests (12 new, all existing pass):

• Producer tests: train() returns _StepArguments with func + trainer captured
• Consumer tests: TrainingStep(step_args=...).arguments produces valid PascalCase dict
  with no session/region leaked, Tags in correct Key/Value format
• Regular session tests: TrainingJob.create() still called normally (no regression)

Manual E2E (real AWS, us-west-2, account 243179453435):

session = PipelineSession()
trainer = SFTTrainer(model='meta-textgeneration-llama-3-2-1b-instruct', ...)
result = trainer.train()                    # _StepArguments (no job)
step = TrainingStep(name="SFT", step_args=result)
step.arguments                              # PascalCase dict ✅
pipeline = Pipeline(steps=[step], sagemaker_session=session)
pipeline.upsert(role_arn="...")             # CreatePipeline ✅
execution = pipeline.start()               # Execution started ✅

Pipeline created and execution rg29qr88isq5 started successfully.

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

[rohangujarathi]

nit: we could reuse some existing dataset or move this to a file

[papriwal]

One small heads-up on the new integ test's teardown (test_v3_trainer_pipeline.py). Nothing blocking; the production PipelineSession change looks good.

[papriwal]

Minor: since the *_create_and_execute tests register a model package into model_package_group, delete_model_package_group here may fail on a non-empty group (and the except Exception: pass would quietly hide it), potentially leaving the group and its model packages behind. Might be worth deleting the member model packages before the group so nothing accumulates in the CI account over time. Non-blocking.

[github-actions]

Claude finished @nayan3107's task in 4m 5s —— View job

---

Review: complete PipelineSession support for V3 trainers

I reviewed the diff against origin/master, traced runnable_by_pipeline / execute_job_functions / TrainingStep.arguments, and confirmed the mechanism end-to-end. The overall approach is sound and mirrors the existing ModelTrainer pattern (model_trainer.py:767-773), and source_code = None in BaseTrainer correctly satisfies get_code_hash() (utilities.py:173). Tests are a solid improvement (producer + consumer + regression). A few issues below, one of which is a real bug.

---

🔴 Tags conversion breaks when a user passes Tag objects

sft_trainer.py:455-459 (and identically dpo_trainer.py, rlaif_trainer.py, rlvr_trainer.py):

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"))}
        for t in pipeline_request["Tags"]
    ]

The tag list is built as tags = _get_jumpstart_tags(...) (lowercase dicts) followed by tags.extend(self.tags or []). But the public field is typed tags: Optional[List[Tag]] (base_trainer.py:97), and Tag is a pydantic model (core/shapes/shapes.py:1109) with no .get method. When a user supplies tags=[Tag(key=..., value=...)], t.get(...) raises AttributeError and train() fails under a PipelineSession.

The non-pipeline path passes these same objects straight to TrainingJob.create(tags=tags), which serializes Tag shapes correctly — so the pipeline path is inconsistent with what the SDK otherwise accepts. Suggest normalizing both forms, e.g.:

def _tag_kv(t):
    if isinstance(t, dict):
        return {"Key": t.get("key", t.get("Key")), "Value": t.get("value", t.get("Value"))}
    return {"Key": t.key, "Value": t.value}

The new consumer tests only mock _get_jumpstart_tags to return lowercase dicts, so this exact path (Tag objects) is untested — worth adding a case with trainer.tags = [Tag(key="k", value="v")].

---

🟡 The serialization block is duplicated verbatim in 4 files

The ~11-line PascalCase/serialize/tag-fix block is copy-pasted identically into sft/dpo/rlaif/rlvr_trainer.py, and a near-variant already lives in model_trainer.py. Any fix (like the Tags one above) now has to be applied in five places. Consider extracting a small helper (e.g. in sagemaker/train/utils.py) — _build_pipeline_request(create_args) — and calling it from each trainer. This keeps the behavior in one place and shrinks each train() body.

---

🟡 Missing newline at EOF

sft_trainer.py ends with return training_job and no trailing newline (diff line 605, \ No newline at end of file). flake8/W292 will flag this in the CI lint gate.

---

🟢 Minor test notes (non-blocking)

• The PascalCase assertion assert any(k[0].isupper() for k in non_none_keys) only checks that at least one key is PascalCase, not all of them — a weaker guarantee than the intent. Consider assert all(...).
• _assert_pipeline_create_and_execute in the integ test starts a real ~30-min serverless fine-tuning job (with pytest.fail on timeout). That's fine for tests/integ/ (CI-gated, not run locally per AGENTS.md), just noting the real AWS cost. Only SFT/DPO have the execute variant; RLAIF/RLVR are definition-only, which seems intentional.

---

Verified correct

• retrieve_caller_name duck-typing resolves to "train" because all four trainers subclass BaseTrainer, which exposes training_image (base_trainer.py:102) — matching validate_step_args_input's expected_caller={ModelTrainer.train.__name__} (steps.py:429). The caller_name == "train" assertions hold.
• Changing return sagemaker_session.context → bare return is correct: execute_job_functions reads the request from sagemaker_session.context.args after re-executing train(), not from the return value.
• serialize() drops None values, so MlflowConfig/ModelPackageConfig = None do not leak into the request.

Nice work — the core wiring is correct and well-tested. The Tags-with-Tag-objects bug is the one I'd fix before merge; the duplication is worth addressing while you're in there.

[nayan3107]

Fixed -

• Tags conversion for Tag objects. The old code assumed all tag entries were dicts and called t.get(...), which crashes on the pydantic Tag objects the public type signature (tags: Optional[List[Tag]]) invites. Now normalizes both forms in all four trainers (sft/dpo/rlaif/rlvr_trainer.py):
  python

  {"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}

• Missing trailing newline in sft_trainer.py (flake8 W292).
• Unit test coverage for Tag objects. New test test_train_pipeline_session_normalizes_tag_objects in test_sft_trainer.py — mixes a JumpStart dict tag with user-supplied Tag(key="env", value="prod") and asserts all three end up in [{Key, Value}] form. This test fails on the old code with AttributeError: 'Tag' object has no attribute 'get'.
• Model package group cleanup (@papriwal). Fixture now paginates list_model_packages, deletes each member package, then the group. Cleanup errors are printed instead of silently swallowed, so real leaks show up in CI output.
  Sample training data moved to files (@rohangujarathi). Two 20-row files under sagemaker-mlops/tests/integ/data/v3_trainer/ (sft_train.jsonl, preference_train.jsonl); loaded and repeated 7× in-memory at fixture time to exceed the 128 batch size used by DPO/RLAIF/RLVR.