Skip to content

feat(train): add the interactive Training Sessions client surface #6452

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rsareddy0329 merged 1 commit into from
Oct 13, 2026
Merged

feat(train): add the interactive Training Sessions client surface #6452

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions docs/api/sagemaker_train.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,44 @@ Multi-Turn RL Training
:undoc-members:
:show-inheritance:

Interactive Training Sessions
-----------------------------

.. automodule:: sagemaker.train.training_session
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.service_client
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.training_client
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.sampling_client
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.types
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.futures
:members:
:undoc-members:
:show-inheritance:

.. automodule:: sagemaker.train.training_session.exceptions
:members:
:undoc-members:
:show-inheritance:

Model Evaluation
----------------

Expand Down
4 changes: 3 additions & 1 deletion sagemaker-train/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ Homepage = "https://github.com/aws/sagemaker-python-sdk"
[project.optional-dependencies]
rft-strands = ["strands-agents"]
rft-langchain = ["httpx>=0.24.0"]
tokenizer = ["transformers"]
test = [
"pytest",
"pytest-cov",
Expand All @@ -63,7 +64,8 @@ test = [
"scipy",
"omegaconf",
"graphene",
"IPython"
"IPython",

@github-actions github-actions Bot Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Low: transformers is added to the test extra (and as the tokenizer extra). transformers is a heavy dependency that often drags in large transitive packages and markedly increases install time and image size for every CI run of the sagemaker-train unit suite — even though only the get_tokenizer()/tokenizer tests need it. Consider gating the tokenizer tests behind an importorskip and keeping transformers out of the default test extra (or pinning a minimum version), so the common unit run doesn't pay for it.

"transformers",

@github-actions github-actions Bot Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Low: transformers is added to the test extra. transformers is a heavy dependency and, without torch/tensorflow installed, emits import-time warnings and can slow collection. The runtime guards the import (optional tokenizer extra), so consider whether the test suite truly needs the full package pulled into every sagemaker-train test run, or whether the tokenizer paths can be exercised with a lighter stub/mock.

]
notebook = [
"ipywidgets>=8.0.0",
Expand Down
12 changes: 12 additions & 0 deletions sagemaker-train/src/sagemaker/train/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,4 +158,16 @@ def __getattr__(name):
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters

return list_hyperparameters
elif name == "ServiceClient":

@github-actions github-actions Bot Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Low: The new __getattr__ branches expose only ServiceClient, TrainingClient, and SamplingClient from sagemaker.train, but not the supporting public types (AdamParams, Datum, ModelInput, TensorData, SamplingParams, etc.). That is consistent with the documented import path (from sagemaker.train.training_session import ...), so it is fine as-is — but worth confirming the intent, since a user who does from sagemaker.train import ServiceClient will reasonably expect from sagemaker.train import Datum to work too and it won't.

@github-actions github-actions Bot Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Low: Only the three clients (ServiceClient, TrainingClient, SamplingClient) are re-exported lazily at the sagemaker.train top level, but the public value types callers routinely need (AdamParams, Datum, ModelInput, TensorData, SamplingParams, VpcConfig, …) are only reachable via sagemaker.train.training_session. The quick-start examples all import those types, so a user who does from sagemaker.train import ServiceClient will still need a second, differently-rooted import for the types. Prior precedent on this repo favored exposing the customer-facing surface through sagemaker.train consistently (Source: #5832). Consider either documenting training_session as the single import root or re-exporting the common types too.

from sagemaker.train.training_session import ServiceClient

return ServiceClient
elif name == "TrainingClient":
from sagemaker.train.training_session import TrainingClient

return TrainingClient
elif name == "SamplingClient":
from sagemaker.train.training_session import SamplingClient

return SamplingClient
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
Loading
Loading