fix(train): add PipelineSession support to V3 trainers (SFT/DPO/RLAIF…

[nayan3107]

…/RLVR)

When a PipelineSession is passed as sagemaker_session, the serverless training path in SFTTrainer, DPOTrainer, RLAIFTrainer, and RLVRTrainer now intercepts the CreateTrainingJob request and returns step arguments instead of immediately launching a training job.

This enables V3 trainers to be used with SageMaker Pipelines TrainingStep, matching the existing behavior of ModelTrainer, Processor, Transformer, and HyperparameterTuner.

The fix follows the established SDK pattern: isinstance check for PipelineSession, call _intercept_create_request with the request args, and return session.context (the captured step arguments).

Issue #, if available:
Fixes: #6163

Description of changes:

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

Testing

Unit Tests (167 pass, 0 regressions):

• All existing SFT/DPO/RLAIF/RLVR trainer tests pass unchanged
• 4 new TestXxxTrainerPipelineSession classes added verifying:
  • TrainingJob.create() is NOT called when PipelineSession is used
  • session._intercept_create_request() IS called with correct create_args and func_name="train"
  • Return value is session.context (_JobStepArguments) — usable with TrainingStep

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

from sagemaker.core.workflow.pipeline_context import PipelineSession
from sagemaker.train.sft_trainer import SFTTrainer
from sagemaker.train.common import TrainingType

session = PipelineSession()
trainer = SFTTrainer(
    model='meta-textgeneration-llama-3-2-1b-instruct',
    training_type=TrainingType.LORA,
    training_dataset='s3://my-bucket/train.jsonl',
    model_package_group='my-group',
    sagemaker_session=session,
    accept_eula=True,
)
result = trainer.train()
# type(result) = _JobStepArguments
# result.caller_name = 'train'
# No CreateTrainingJob API call made. No billing.

Verified scenarios:

Scenario	Result
SFTTrainer + PipelineSession	✅ Returns step_args, no job launched
SFTTrainer + regular Session	✅ Launches job normally (existing behavior)
DPO/RLAIF/RLVR + PipelineSession	✅ Returns step_args (unit tested)
All 137 pre-existing trainer unit tests	✅ Pass (no regressions)

[github-actions]

Claude finished @nayan3107's task in 3m 23s —— View job

---

Code Review: PipelineSession support for V3 trainers

• Read the PR diff (/tmp/pr.diff)
• Traced the existing ModelTrainer → TrainingStep pipeline path
• Checked PipelineSession._intercept_create_request / runnable_by_pipeline / execute_job_functions
• Reviewed the added tests
• Posted findings

Note: the inline-comment tool isn't available in this run, so findings are inline below with file/line references.

---

I appreciate the intent here — closing the gap so SFT/DPO/RLAIF/RLVR match ModelTrainer for Pipelines. Unfortunately, as implemented this does not actually make the trainers usable with TrainingStep, and the tests don't exercise the real consumption path so they pass anyway. Details below.

🔴 Blocking: the returned object cannot be consumed by TrainingStep

The four trainers do:

if isinstance(sagemaker_session, PipelineSession):
    sagemaker_session._intercept_create_request(create_args, None, "train")
    return sagemaker_session.context

_intercept_create_request(..., "train") sets context = _JobStepArguments("train", create_args) (pipeline_context.py:231). That object has .args set but .func = None and .func_args = () (_StepArguments.__init__, pipeline_context.py:30-44).

Now look at how TrainingStep consumes step_args (sagemaker-mlops/.../workflow/steps.py:439-467):

execute_job_functions(self.step_args)                 # -> step_args.func(*step_args.func_args, ...)
model_trainer = self.step_args.func_args[0]           # expects the trainer instance
request_dict = model_trainer.sagemaker_session.context.args

and execute_job_functions (utilities.py:459):

chained_args = step_args.func(*step_args.func_args, **step_args.func_kwargs)

With the object these trainers return, that is None(...) → TypeError: 'NoneType' object is not callable, and func_args[0] would be an IndexError. So TrainingStep(step_args=trainer.train()).arguments raises at pipeline-compile time.

Root cause: ModelTrainer.train is decorated with @runnable_by_pipeline (model_trainer.py:791). In a PipelineSession that decorator does not run the body — it returns _StepArguments(caller_name, run_func, self_instance, ...) (pipeline_context.py:344), capturing the function and the trainer so TrainingStep can re-execute it later. The actual _intercept_create_request call inside train (model_trainer.py:831-833) only fires during that re-execution. These new trainers replicate only the inner interception and skip the capture half, so there's nothing for TrainingStep to re-run.

The fix should follow ModelTrainer exactly — decorate each train with @runnable_by_pipeline — rather than hand-rolling the intercept in the body.

🔴 Blocking: create_args is the wrong shape for a pipeline request

Even if the capture mechanism were fixed, the payload is wrong. create_args is built as kwargs for TrainingJob.create(**create_args) and includes non-serializable/irrelevant entries:

"session": sagemaker_session.boto_session,          # sft_trainer.py:432
"region": sagemaker_session.boto_session.region_name,
"training_job_name": current_training_job_name,

Compare ModelTrainer, which for the pipeline path builds a distinct request — PascalCase-converted, serialize()-d, with training_job_name popped and no session/region (model_trainer.py:767-773). The Arguments of a TrainingStep must be a serialized CreateTrainingJob shape (TrainingJobName, RoleArn, …), not snake_case kwargs carrying a live boto3 session object. As written, context.args would carry a boto session into the pipeline definition.

🟠 Tests don't cover the behavior that matters

test_train_with_pipeline_session_does_not_launch_job (all four files) mocks _intercept_create_request and manually sets pipeline_session.context = step_args, then asserts train() returns that same object and that TrainingJob.create wasn't called. Because the interceptor is mocked and no TrainingStep / execute_job_functions is ever built, the test is effectively tautological — it would pass regardless of whether the produced step_args is actually consumable. It cannot catch either blocking issue above.

A meaningful test would construct a real PipelineSession (or a faithful fake), build a TrainingStep(name=..., step_args=trainer.train(...)), and assert step.arguments returns a valid PascalCase CreateTrainingJob request (with TrainingJobName trimmed and no boto session). That is the assertion the issue (#6163) actually needs.

Suggested direction

1. Remove the manual intercept block from all four train methods.
2. Decorate each train with @runnable_by_pipeline (mirroring ModelTrainer.train), and inside the body build/serialize the pipeline request the same way ModelTrainer does (PascalCase, drop session/region, pop training_job_name) before the _intercept_create_request call, keeping the direct TrainingJob.create path for the non-pipeline case.
3. Add an end-to-end test that actually builds a TrainingStep and inspects .arguments.

Happy to help sketch the runnable_by_pipeline wiring if useful. I did not find security or performance concerns — the issues here are correctness/coverage.
· branch fix/pipeline-session-v3-trainers