Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Allow any AWS authentication method in studios (#19336)

This commit is contained in:
Andy☼ McSherry☼ 2024-01-24 08:20:53 -08:00 committed by GitHub
parent 71bfdc3c60
commit 577bd85654
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 1 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/lightning/data/streaming/client.py
View File

@ -7,8 +7,6 @@ from lightning.data.streaming.constants import _BOTO3_AVAILABLE
if _BOTO3_AVAILABLE:
import boto3
import botocore
from botocore.credentials import InstanceMetadataProvider
from botocore.utils import InstanceMetadataFetcher
class S3Client:
@ -31,14 +29,8 @@ class S3Client:
# Re-generate credentials for EC2
if self._last_time is None or (time() - self._last_time) > self._refetch_interval:
provider = InstanceMetadataProvider(iam_role_fetcher=InstanceMetadataFetcher(timeout=3600, num_attempts=5))
credentials = provider.load()
self._client = boto3.client(
"s3",
aws_access_key_id=credentials.access_key,
aws_secret_access_key=credentials.secret_key,
aws_session_token=credentials.token,
config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
"s3", config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"})
)
self._last_time = time()

6
tests/tests_data/streaming/test_client.py
View File

@ -31,12 +31,6 @@ def test_s3_client_with_cloud_space_id(monkeypatch):
botocore = mock.MagicMock()
monkeypatch.setattr(client, "botocore", botocore)
instance_metadata_provider = mock.MagicMock()
monkeypatch.setattr(client, "InstanceMetadataProvider", instance_metadata_provider)
instance_metadata_fetcher = mock.MagicMock()
monkeypatch.setattr(client, "InstanceMetadataFetcher", instance_metadata_fetcher)
monkeypatch.setenv("LIGHTNING_CLOUD_SPACE_ID", "dummy")
s3 = client.S3Client(1)