Support for AWS session token in TensorizerConfig when using temporary credentials

[khalilbenayed]

I'm trying to use Tensorizer to stream model weights from an S3 bucket within an EKS pod. The pod is configured with IAM Roles for Service Accounts (IRSA), and I'm using the boto3.assume_role_with_web_identity API to obtain temporary credentials.

Since Tensorizer currently requires manually setting the S3 credentials, I pass in the values returned from assume_role_with_web_identity. However, I encounter 403 Forbidden errors when attempting to stream weights. After investigation, I confirmed that the temporary credentials are valid—I can manually download S3 objects using the same credentials.

The issue appears to be that while I can provide aws_access_key_id and aws_secret_access_key, there's no way to supply the required session_token via TensorizerConfig. This token is essential for using temporary credentials granted by STS.

Feature request: Add support for an optional session_token parameter in TensorizerConfig to support IRSA and other temporary credential setups.

[pingleiwandb]

I think the implementation is using boto to create a signed url for download, so you can workaround by creating a signed url and let curl do the download. Though for upload it is using s3 client directly...

tensorizer/tensorizer/stream_io.py

Lines 1072 to 1099 in 79c5063

	def s3_download(
	path_uri: str,
	s3_access_key_id: Optional[str],
	s3_secret_access_key: Optional[str],
	s3_endpoint: Optional[str] = default_s3_read_endpoint,
	s3_region_name: Optional[str] = None,
	s3_signature_version: Optional[str] = None,
	buffer_size: Optional[int] = None,
	begin: Optional[int] = None,
	end: Optional[int] = None,
	certificate_handling: Optional[CAInfo] = None,
	) -> CURLStreamFile:
	url = _s3_download_url(
	path_uri=path_uri,
	s3_access_key_id=s3_access_key_id,
	s3_secret_access_key=s3_secret_access_key,
	s3_endpoint=s3_endpoint,
	s3_region_name=s3_region_name,
	s3_signature_version=s3_signature_version,
	)
	return CURLStreamFile(
	url,
	buffer_size=buffer_size,
	certificate_handling=certificate_handling,
	begin=begin,
	end=end,
	)

[khalilbenayed]

Using presigned urls as a workaround won't work for sharded models, as the presigned url would be unique per shard and the TensorizerConfig expects a single URI patterns for all shards...
Supporting IRSA and other temporary creds setup is essential for production applications, it would be great if this feature could be implemented.