From 9595ef09fd4c8509932b356a038d355ba6620d30 Mon Sep 17 00:00:00 2001 From: Greg Hendrickson Date: Tue, 3 Feb 2026 15:42:19 +0000 Subject: [PATCH] fix: add aws_session_token support for SSO credentials --- components/runpod_trainer/handler.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/components/runpod_trainer/handler.py b/components/runpod_trainer/handler.py index f74ffdd..2ad2f18 100644 --- a/components/runpod_trainer/handler.py +++ b/components/runpod_trainer/handler.py @@ -32,6 +32,7 @@ def upload_to_s3(local_path: str, s3_bucket: str, s3_prefix: str, aws_credential 's3', aws_access_key_id=aws_credentials.get('aws_access_key_id'), aws_secret_access_key=aws_credentials.get('aws_secret_access_key'), + aws_session_token=aws_credentials.get('aws_session_token'), region_name=aws_credentials.get('aws_region', 'us-east-1') ) @@ -277,6 +278,7 @@ def train_llm_lora(job_input: Dict[str, Any]) -> Dict[str, Any]: aws_creds = { 'aws_access_key_id': job_input.get('aws_access_key_id'), 'aws_secret_access_key': job_input.get('aws_secret_access_key'), + 'aws_session_token': job_input.get('aws_session_token'), 'aws_region': job_input.get('aws_region', 'us-east-1'), } model_short = model_name.split('/')[-1] @@ -426,6 +428,7 @@ def train_bert_classifier(job_input: Dict[str, Any]) -> Dict[str, Any]: aws_creds = { 'aws_access_key_id': job_input.get('aws_access_key_id'), 'aws_secret_access_key': job_input.get('aws_secret_access_key'), + 'aws_session_token': job_input.get('aws_session_token'), 'aws_region': job_input.get('aws_region', 'us-east-1'), } s3_prefix = job_input.get('s3_prefix', 'ddi-models/bert')