Skip to content

Truss: support new auth methods for training (OIDC, AWS_ASSUME_ROLE) #2633

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

Draft
MarkGotesman wants to merge 2 commits into
base: main
Choose a base branch
from
+26 −7
Draft

Truss: support new auth methods for training (OIDC, AWS_ASSUME_ROLE) #2633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
34df714
Support AssumeRole for weight auth
MarkGotesman Aug 31, 2026
3393295
Enable OIDC for weight auth
MarkGotesman Aug 31, 2026
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
17 changes: 12 additions & 5 deletions truss-train/truss_train/definitions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,14 +249,21 @@ class TrainingJob(custom_types.SafeModelNoExtra):
enable_baseten_workdir: bool = True

@model_validator(mode="after")
def _validate_weights_auth_only_custom_secret(self) -> "TrainingJob":
"""Training jobs only support CUSTOM_SECRET with auth_secret_name for weights; OIDC is not supported."""
def _validate_weights_auth_method(self) -> "TrainingJob":
"""Validate that weight authentication is supported for training jobs."""
supported_auth_methods = {
truss_config.WeightsAuthMethod.CUSTOM_SECRET,
truss_config.WeightsAuthMethod.AWS_ASSUME_ROLE,
truss_config.WeightsAuthMethod.AWS_OIDC,
truss_config.WeightsAuthMethod.GCP_OIDC,
}
for w in self.weights:
if w.auth is not None:
if w.auth.auth_method != truss_config.WeightsAuthMethod.CUSTOM_SECRET:
if w.auth.auth_method not in supported_auth_methods:
raise ValueError(
f"weight {w.source}: only auth_method CUSTOM_SECRET with auth_secret_name is supported for training jobs. "
"OIDC and assume-role methods (AWS_OIDC, GCP_OIDC, AWS_ASSUME_ROLE) are not supported."
f"weight {w.source}: auth_method {w.auth.auth_method.value} is not "
"supported for training jobs. Supported auth methods: AWS_ASSUME_ROLE, "
"AWS_OIDC, CUSTOM_SECRET, GCP_OIDC."
)
return self

Expand Down
12 changes: 12 additions & 0 deletions truss/base/truss_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@
DOCKER_AUTH_SECRET_NAME_PARAM = "secret_name"
AWS_OIDC_ROLE_ARN_PARAM = "aws_oidc_role_arn"
AWS_OIDC_REGION_PARAM = "aws_oidc_region"
AWS_ASSUME_ROLE_ARN_PARAM = "aws_assume_role_arn"
AWS_ASSUME_ROLE_REGION_PARAM = "aws_assume_role_region"
GCP_OIDC_SERVICE_ACCOUNT_PARAM = "gcp_oidc_service_account"
GCP_OIDC_WORKLOAD_ID_PROVIDER_PARAM = "gcp_oidc_workload_id_provider"
AWS_ASSUME_ROLE_ARN_PARAM = "aws_assume_role_arn"
Expand Down Expand Up @@ -349,6 +351,12 @@ class WeightsAuth(AuthFieldsMixin):
default=None,
description="Baseten secret name containing credentials for accessing the source.",
)
aws_assume_role_arn: Optional[str] = pydantic.Field(
default=None, description="AWS IAM role ARN for AssumeRole authentication."
)
aws_assume_role_region: Optional[str] = pydantic.Field(
default=None, description="AWS region for AssumeRole authentication."
)

@pydantic.field_validator("auth_method", mode="before")
@classmethod
Expand All @@ -364,6 +372,8 @@ def _validate_auth_fields(self) -> "WeightsAuth":
forbidden=[
AWS_OIDC_ROLE_ARN_PARAM,
AWS_OIDC_REGION_PARAM,
AWS_ASSUME_ROLE_ARN_PARAM,
AWS_ASSUME_ROLE_REGION_PARAM,
GCP_OIDC_SERVICE_ACCOUNT_PARAM,
GCP_OIDC_WORKLOAD_ID_PROVIDER_PARAM,
AWS_ASSUME_ROLE_ARN_PARAM,
Expand All @@ -376,6 +386,8 @@ def _validate_auth_fields(self) -> "WeightsAuth":
required=[AWS_OIDC_ROLE_ARN_PARAM, AWS_OIDC_REGION_PARAM],
forbidden=[
WEIGHTS_AUTH_SECRET_NAME_PARAM,
AWS_ASSUME_ROLE_ARN_PARAM,
AWS_ASSUME_ROLE_REGION_PARAM,
GCP_OIDC_SERVICE_ACCOUNT_PARAM,
GCP_OIDC_WORKLOAD_ID_PROVIDER_PARAM,
AWS_ASSUME_ROLE_ARN_PARAM,
Expand Down
4 changes: 2 additions & 2 deletions truss/config.schema.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1968,7 +1968,7 @@
}
],
"default": null,
"description": "AWS IAM role ARN that Baseten assumes with its own AWS principal, scoped by the sts:ExternalId Baseten assigns to your organization.",
"description": "AWS IAM role ARN for AssumeRole authentication.",
"title": "Aws Assume Role Arn"
},
"aws_assume_role_region": {
Expand All @@ -1981,7 +1981,7 @@
}
],
"default": null,
"description": "AWS region for AWS AssumeRole authentication.",
"description": "AWS region for AssumeRole authentication.",
"title": "Aws Assume Role Region"
},
"auth_method": {
Expand Down
Loading