Merge pull request #122 from dreadnode/users/raja/add-refresh-s3-token

feat: Add automatic credential refresh for S3 storage operations

Author: rdheekonda

docs/sdk/api.mdx

@@ -728,26 +728,41 @@ def get_user(self) -> UserResponse:
 ### get\_user\_data\_credentials
 
 ```python
-get_user_data_credentials() -> UserDataCredentials
+get_user_data_credentials(
+    duration: int = DEFAULT_FS_CREDENTIAL_DURATION,
+) -> UserDataCredentials
 ```
 
 Retrieves user data credentials for secondary storage access.
 
+**Parameters:**
+
+* **`duration`**
+  (`int`, default:
+  `DEFAULT_FS_CREDENTIAL_DURATION`
+  )
+  –Credential lifetime in seconds (default: 4 hours)
+
 **Returns:**
 
 * `UserDataCredentials`
   –The user data credentials object.
 
 <Accordion title="Source code in dreadnode/api/client.py" icon="code">
 ```python
-def get_user_data_credentials(self) -> UserDataCredentials:
+def get_user_data_credentials(
+    self, duration: int = DEFAULT_FS_CREDENTIAL_DURATION
+) -> UserDataCredentials:
     """
     Retrieves user data credentials for secondary storage access.
 
+    Args:
+        duration: Credential lifetime in seconds (default: 4 hours)
+
     Returns:
         The user data credentials object.
     """
-    response = self.request("GET", "/user-data/credentials")
+    response = self._request("GET", "/user-data/credentials", params={"duration": duration})
     return UserDataCredentials(**response.json())
 ```
 

docs/sdk/artifact.mdx

@@ -244,7 +244,10 @@ ArtifactStorage
 ---------------
 
 ```python
-ArtifactStorage(file_system: AbstractFileSystem)
+ArtifactStorage(
+    file_system: AbstractFileSystem,
+    credential_refresher: Callable[[], bool] | None = None,
+)
 ```
 
 Storage for artifacts with efficient handling of large files and directories.
@@ -260,17 +263,28 @@ Initialize artifact storage with a file system and prefix path.
 * **`file_system`**
   (`AbstractFileSystem`)
   –FSSpec-compatible file system
+* **`credential_refresher`**
+  (`Callable[[], bool] | None`, default:
+  `None`
+  )
+  –Optional function to refresh credentials when it's about to expire
 
 <Accordion title="Source code in dreadnode/artifact/storage.py" icon="code">
 ```python
-def __init__(self, file_system: fsspec.AbstractFileSystem):
+def __init__(
+    self,
+    file_system: fsspec.AbstractFileSystem,
+    credential_refresher: t.Callable[[], bool] | None = None,
+):
     """
     Initialize artifact storage with a file system and prefix path.
 
     Args:
         file_system: FSSpec-compatible file system
+        credential_refresher: Optional function to refresh credentials when it's about to expire
     """
     self._file_system = file_system
+    self._credential_refresher = credential_refresher
 ```
 
 
@@ -464,6 +478,7 @@ Store a file in the storage system, using multipart upload for large files.
 
 <Accordion title="Source code in dreadnode/artifact/storage.py" icon="code">
 ```python
+@with_credential_refresh
 def store_file(self, file_path: Path, target_key: str) -> str:
     """
     Store a file in the storage system, using multipart upload for large files.

docs/sdk/main.mdx

@@ -65,6 +65,8 @@ def __init__(
     self._fs_prefix: str = ".dreadnode/storage/"
 
     self._initialized = False
+    self._credentials: UserDataCredentials | None = None
+    self._credentials_expiry: datetime | None = None
 ```
 
 
@@ -380,6 +382,7 @@ def continue_run(self, run_context: RunContext) -> RunSpan:
         tracer=self._get_tracer(),
         file_system=self._fs,
         prefix_path=self._fs_prefix,
+        credential_refresher=self._refresh_storage_credentials if self._credentials else None,
     )
 ```
 
@@ -523,19 +526,21 @@ def initialize(self) -> None:
         #         )
         #     )
         # )
-
-        credentials = self._api.get_user_data_credentials()
-        resolved_endpoint = resolve_endpoint(credentials.endpoint)
+        self._credentials = self._api.get_user_data_credentials(
+            duration=DEFAULT_FS_CREDENTIAL_DURATION
+        )
+        self._credentials_expiry = self._credentials.expiration
+        resolved_endpoint = resolve_endpoint(self._credentials.endpoint)
         self._fs = S3FileSystem(
-            key=credentials.access_key_id,
-            secret=credentials.secret_access_key,
-            token=credentials.session_token,
+            key=self._credentials.access_key_id,
+            secret=self._credentials.secret_access_key,
+            token=self._credentials.session_token,
             client_kwargs={
                 "endpoint_url": resolved_endpoint,
-                "region_name": credentials.region,
+                "region_name": self._credentials.region,
             },
         )
-        self._fs_prefix = f"{credentials.bucket}/{credentials.prefix}/"
+        self._fs_prefix = f"{self._credentials.bucket}/{self._credentials.prefix}/"
 
     self._logfire = logfire.configure(
         local=not self.is_default,
@@ -1723,6 +1728,7 @@ def run(
         file_system=self._fs,
         prefix_path=self._fs_prefix,
         autolog=autolog,
+        credential_refresher=self._refresh_storage_credentials if self._credentials else None,
     )
 ```
 

docs/sdk/metric.mdx

@@ -31,8 +31,8 @@ Metric
 Metric(
     value: float,
     step: int = 0,
-    timestamp: datetime = lambda: datetime.now(
-        timezone.utc
+    timestamp: datetime = (
+        lambda: datetime.now(timezone.utc)
     )(),
     attributes: JsonDict = dict(),
 )

dreadnode/api/client.py

@@ -36,7 +36,11 @@
     process_run,
     process_task,
 )
-from dreadnode.constants import DEFAULT_MAX_POLL_TIME, DEFAULT_POLL_INTERVAL
+from dreadnode.constants import (
+    DEFAULT_FS_CREDENTIAL_DURATION,
+    DEFAULT_MAX_POLL_TIME,
+    DEFAULT_POLL_INTERVAL,
+)
 from dreadnode.util import logger
 from dreadnode.version import VERSION
 
@@ -517,12 +521,17 @@ def export_timeseries(
 
     # User data access
 
-    def get_user_data_credentials(self) -> UserDataCredentials:
+    def get_user_data_credentials(
+        self, duration: int = DEFAULT_FS_CREDENTIAL_DURATION
+    ) -> UserDataCredentials:
         """
         Retrieves user data credentials for secondary storage access.
 
+        Args:
+            duration: Credential lifetime in seconds (default: 4 hours)
+
         Returns:
             The user data credentials object.
         """
-        response = self.request("GET", "/user-data/credentials")
+        response = self._request("GET", "/user-data/credentials", params={"duration": duration})
         return UserDataCredentials(**response.json())

dreadnode/artifact/storage.py

@@ -4,10 +4,12 @@
 """
 
 import hashlib
+import typing as t
 from pathlib import Path
 
 import fsspec  # type: ignore[import-untyped]
 
+from dreadnode.storage_utils import with_credential_refresh
 from dreadnode.util import logger
 
 CHUNK_SIZE = 8 * 1024 * 1024  # 8MB
@@ -22,15 +24,27 @@ class ArtifactStorage:
     - Batch uploads for directories handled by fsspec
     """
 
-    def __init__(self, file_system: fsspec.AbstractFileSystem):
+    def __init__(
+        self,
+        file_system: fsspec.AbstractFileSystem,
+        credential_refresher: t.Callable[[], bool] | None = None,
+    ):
         """
         Initialize artifact storage with a file system and prefix path.
 
         Args:
             file_system: FSSpec-compatible file system
+            credential_refresher: Optional function to refresh credentials when it's about to expire
         """
         self._file_system = file_system
+        self._credential_refresher = credential_refresher
 
+    def _refresh_credentials_if_needed(self) -> None:
+        """Refresh credentials if refresher is available."""
+        if self._credential_refresher:
+            self._credential_refresher()
+
+    @with_credential_refresh
     def store_file(self, file_path: Path, target_key: str) -> str:
         """
         Store a file in the storage system, using multipart upload for large files.

dreadnode/constants.py

@@ -56,3 +56,7 @@
     # allow overriding the user config file via env variable
     os.getenv("DREADNODE_USER_CONFIG_FILE") or pathlib.Path.home() / ".dreadnode" / "config"
 )
+
+# Default values for the file system credential management
+DEFAULT_FS_CREDENTIAL_DURATION = 14400  # 4 hours in seconds
+FS_CREDENTIAL_REFRESH_BUFFER = 300  # 5 minutes in seconds

dreadnode/main.py

@@ -26,6 +26,7 @@
 from dreadnode.api.client import ApiClient
 from dreadnode.config import UserConfig
 from dreadnode.constants import (
+    DEFAULT_FS_CREDENTIAL_DURATION,
     DEFAULT_SERVER_URL,
     ENV_API_KEY,
     ENV_API_TOKEN,
@@ -35,6 +36,7 @@
     ENV_PROJECT,
     ENV_SERVER,
     ENV_SERVER_URL,
+    FS_CREDENTIAL_REFRESH_BUFFER,
 )
 from dreadnode.metric import (
     Metric,
@@ -64,7 +66,7 @@
     Inherited,
     JsonValue,
 )
-from dreadnode.util import clean_str, handle_internal_errors, resolve_endpoint
+from dreadnode.util import clean_str, handle_internal_errors, logger, resolve_endpoint
 from dreadnode.version import VERSION
 
 if t.TYPE_CHECKING:
@@ -73,6 +75,8 @@
     from opentelemetry.sdk.trace import SpanProcessor
     from opentelemetry.trace import Tracer
 
+    from dreadnode.api.models import UserDataCredentials
+
 
 ToObject = t.Literal["task-or-run", "run"]
 
@@ -137,6 +141,8 @@ def __init__(
         self._fs_prefix: str = ".dreadnode/storage/"
 
         self._initialized = False
+        self._credentials: UserDataCredentials | None = None
+        self._credentials_expiry: datetime | None = None
 
     def _get_profile_server(self, profile: str | None = None) -> str | None:
         with contextlib.suppress(Exception):
@@ -347,19 +353,21 @@ def initialize(self) -> None:
             #         )
             #     )
             # )
-
-            credentials = self._api.get_user_data_credentials()
-            resolved_endpoint = resolve_endpoint(credentials.endpoint)
+            self._credentials = self._api.get_user_data_credentials(
+                duration=DEFAULT_FS_CREDENTIAL_DURATION
+            )
+            self._credentials_expiry = self._credentials.expiration
+            resolved_endpoint = resolve_endpoint(self._credentials.endpoint)
             self._fs = S3FileSystem(
-                key=credentials.access_key_id,
-                secret=credentials.secret_access_key,
-                token=credentials.session_token,
+                key=self._credentials.access_key_id,
+                secret=self._credentials.secret_access_key,
+                token=self._credentials.session_token,
                 client_kwargs={
                     "endpoint_url": resolved_endpoint,
-                    "region_name": credentials.region,
+                    "region_name": self._credentials.region,
                 },
             )
-            self._fs_prefix = f"{credentials.bucket}/{credentials.prefix}/"
+            self._fs_prefix = f"{self._credentials.bucket}/{self._credentials.prefix}/"
 
         self._logfire = logfire.configure(
             local=not self.is_default,
@@ -406,6 +414,45 @@ def api(self, *, server: str | None = None, token: str | None = None) -> ApiClie
 
         return self._api
 
+    def _refresh_storage_credentials(self) -> bool:
+        """Refresh storage credentials if they are about to expire."""
+        if not self._api or not self._credentials:
+            return False
+
+        now = datetime.now(timezone.utc)
+
+        if (
+            self._credentials_expiry is None
+            or (self._credentials_expiry - now).total_seconds() < FS_CREDENTIAL_REFRESH_BUFFER
+        ):
+            try:
+                logger.info("Refreshing storage credentials")
+                self._credentials = self._api.get_user_data_credentials(
+                    duration=DEFAULT_FS_CREDENTIAL_DURATION
+                )
+                self._credentials_expiry = self._credentials.expiration
+
+                resolved_endpoint = resolve_endpoint(self._credentials.endpoint)
+                self._fs = S3FileSystem(
+                    key=self._credentials.access_key_id,
+                    secret=self._credentials.secret_access_key,
+                    token=self._credentials.session_token,
+                    client_kwargs={
+                        "endpoint_url": resolved_endpoint,
+                        "region_name": self._credentials.region,
+                    },
+                )
+                logger.info(
+                    f"Storage credentials refreshed, valid until {self._credentials_expiry}"
+                )
+                return True  # noqa: TRY300
+
+            except Exception as e:  # noqa: BLE001
+                logger.error(f"Failed to refresh storage credentials: {e}")
+                return False
+
+        return True
+
     def _get_tracer(self, *, is_span_tracer: bool = True) -> "Tracer":
         return self._logfire._tracer_provider.get_tracer(  # noqa: SLF001
             self.otel_scope,
@@ -778,6 +825,7 @@ def run(
             file_system=self._fs,
             prefix_path=self._fs_prefix,
             autolog=autolog,
+            credential_refresher=self._refresh_storage_credentials if self._credentials else None,
         )
 
     def get_run_context(self) -> RunContext:
@@ -824,6 +872,7 @@ def continue_run(self, run_context: RunContext) -> RunSpan:
             tracer=self._get_tracer(),
             file_system=self._fs,
             prefix_path=self._fs_prefix,
+            credential_refresher=self._refresh_storage_credentials if self._credentials else None,
         )
 
     def tag(self, *tag: str, to: ToObject = "task-or-run") -> None: