Support aws_session_token for S3 backends

Author: bghira

simpletuner/helpers/data_backend/aws.py

@@ -63,6 +63,7 @@ def __init__(
         endpoint_url: str = None,
         aws_access_key_id: str = None,
         aws_secret_access_key: str = None,
+        aws_session_token: str = None,
         read_retry_limit: int = 5,
         write_retry_limit: int = 5,
         read_retry_interval: int = 5,
@@ -79,6 +80,7 @@ def __init__(
         self.write_retry_interval = write_retry_interval
         self.compress_cache = compress_cache
         self.max_pool_connections = max_pool_connections
+        self.aws_session_token = aws_session_token
         self.type = "aws"
 
         extra_args = {"region_name": region_name}
@@ -91,6 +93,7 @@ def __init__(
             "s3",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
             config=s3_config,
             **extra_args,
         )
@@ -354,6 +357,7 @@ def test_s3_connection(
     endpoint_url: Optional[str] = None,
     aws_access_key_id: Optional[str] = None,
     aws_secret_access_key: Optional[str] = None,
+    aws_session_token: Optional[str] = None,
     max_keys: int = 5,
 ) -> Dict[str, Any]:
     """Run a minimal connectivity check against an S3 bucket."""
@@ -372,6 +376,7 @@ def test_s3_connection(
             "s3",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
             config=s3_config,
             **extra_args,
         )

simpletuner/helpers/data_backend/builders/aws.py

@@ -30,6 +30,7 @@ def _create_backend(self, config: BaseBackendConfig) -> S3DataBackend:
         aws_endpoint_url = getattr(config, "aws_endpoint_url", None)
         aws_access_key_id = getattr(config, "aws_access_key_id", None)
         aws_secret_access_key = getattr(config, "aws_secret_access_key", None)
+        aws_session_token = getattr(config, "aws_session_token", None)
 
         compress_cache = self._get_compression_setting(config)
         max_pool_connections = self._resolve_max_pool_connections(config)
@@ -42,6 +43,7 @@ def _create_backend(self, config: BaseBackendConfig) -> S3DataBackend:
             endpoint_url=aws_endpoint_url,
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
             compress_cache=compress_cache,
             max_pool_connections=max_pool_connections,
         )

simpletuner/helpers/data_backend/config/image.py

@@ -34,6 +34,7 @@ class ImageBackendConfig(BaseBackendConfig):
     aws_endpoint_url: Optional[str] = None
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
+    aws_session_token: Optional[str] = None
     aws_max_pool_connections: Optional[int] = None
 
     dataset_name: Optional[str] = None
@@ -131,6 +132,7 @@ def _get_arg(key: str, default: Any = None) -> Any:
             config.aws_endpoint_url = backend_dict.get("aws_endpoint_url")
             config.aws_access_key_id = backend_dict.get("aws_access_key_id")
             config.aws_secret_access_key = backend_dict.get("aws_secret_access_key")
+            config.aws_session_token = backend_dict.get("aws_session_token")
             config.aws_max_pool_connections = backend_dict.get(
                 "aws_max_pool_connections", _get_arg("aws_max_pool_connections")
             )
@@ -422,6 +424,8 @@ def to_dict(self) -> Dict[str, Any]:
             config["aws_endpoint_url"] = self.aws_endpoint_url
             config["aws_access_key_id"] = self.aws_access_key_id
             config["aws_secret_access_key"] = self.aws_secret_access_key
+            if self.aws_session_token is not None:
+                config["aws_session_token"] = self.aws_session_token
             if self.aws_max_pool_connections is not None:
                 config["aws_max_pool_connections"] = self.aws_max_pool_connections
 

simpletuner/helpers/data_backend/config/text_embed.py

@@ -45,6 +45,7 @@ def from_dict(cls, backend_dict: Dict[str, Any], args: Dict[str, Any]) -> "TextE
             "aws_endpoint_url",
             "aws_access_key_id",
             "aws_secret_access_key",
+            "aws_session_token",
             "aws_region_name",
         }
         for key in passthrough_keys: