dreadnode 1.13.0__tar.gz → 1.13.2__tar.gz

{dreadnode-1.13.0 → dreadnode-1.13.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dreadnode
-Version: 1.13.0
+Version: 1.13.2
 Summary: Dreadnode SDK
 Author: Nick Landers
 Author-email: monoxgas@gmail.com

{dreadnode-1.13.0 → dreadnode-1.13.2}/dreadnode/api/client.py

@@ -36,7 +36,10 @@ from dreadnode.api.util import (
     process_run,
     process_task,
 )
-from dreadnode.constants import DEFAULT_MAX_POLL_TIME, DEFAULT_POLL_INTERVAL
+from dreadnode.constants import (
+    DEFAULT_MAX_POLL_TIME,
+    DEFAULT_POLL_INTERVAL,
+)
 from dreadnode.util import logger
 from dreadnode.version import VERSION
 
@@ -524,5 +527,5 @@ class ApiClient:
         Returns:
             The user data credentials object.
         """
-        response = self.request("GET", "/user-data/credentials")
+        response = self._request("GET", "/user-data/credentials")
         return UserDataCredentials(**response.json())

{dreadnode-1.13.0 → dreadnode-1.13.2}/dreadnode/artifact/storage.py

@@ -4,10 +4,12 @@ Provides efficient uploading of files and directories with deduplication.
 """
 
 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-1.13.0 → dreadnode-1.13.2}/dreadnode/constants.py

@@ -39,6 +39,7 @@ ENV_API_KEY = "DREADNODE_API_KEY"  # pragma: allowlist secret (alternative to AP
 ENV_LOCAL_DIR = "DREADNODE_LOCAL_DIR"
 ENV_PROJECT = "DREADNODE_PROJECT"
 ENV_PROFILE = "DREADNODE_PROFILE"
+ENV_CONSOLE = "DREADNODE_CONSOLE"
 
 #
 # Environment
@@ -55,3 +56,6 @@ USER_CONFIG_PATH = pathlib.Path(
     # 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
+FS_CREDENTIAL_REFRESH_BUFFER = 300  # 5 minutes in seconds

{dreadnode-1.13.0 → dreadnode-1.13.2}/dreadnode/main.py

@@ -29,11 +29,13 @@ from dreadnode.constants import (
     DEFAULT_SERVER_URL,
     ENV_API_KEY,
     ENV_API_TOKEN,
+    ENV_CONSOLE,
     ENV_LOCAL_DIR,
     ENV_PROFILE,
     ENV_PROJECT,
     ENV_SERVER,
     ENV_SERVER_URL,
+    FS_CREDENTIAL_REFRESH_BUFFER,
 )
 from dreadnode.metric import (
     Metric,
@@ -63,7 +65,7 @@ from dreadnode.types import (
     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:
@@ -72,6 +74,8 @@ if t.TYPE_CHECKING:
     from opentelemetry.sdk.trace import SpanProcessor
     from opentelemetry.trace import Tracer
 
+    from dreadnode.api.models import UserDataCredentials
+
 
 ToObject = t.Literal["task-or-run", "run"]
 
@@ -100,7 +104,7 @@ class Dreadnode:
     project: str | None
     service_name: str | None
     service_version: str | None
-    console: logfire.ConsoleOptions | t.Literal[False, True]
+    console: logfire.ConsoleOptions | bool
     send_to_logfire: bool | t.Literal["if-token-present"]
     otel_scope: str
 
@@ -113,7 +117,7 @@ class Dreadnode:
         project: str | None = None,
         service_name: str | None = None,
         service_version: str | None = None,
-        console: logfire.ConsoleOptions | t.Literal[False, True] = True,
+        console: logfire.ConsoleOptions | bool = True,
         send_to_logfire: bool | t.Literal["if-token-present"] = False,
         otel_scope: str = "dreadnode",
     ) -> None:
@@ -136,6 +140,8 @@ class Dreadnode:
         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):
@@ -167,7 +173,7 @@ class Dreadnode:
         project: str | None = None,
         service_name: str | None = None,
         service_version: str | None = None,
-        console: logfire.ConsoleOptions | t.Literal[False, True] = True,
+        console: logfire.ConsoleOptions | bool | None = None,
         send_to_logfire: bool | t.Literal["if-token-present"] = False,
         otel_scope: str = "dreadnode",
     ) -> None:
@@ -195,7 +201,7 @@ class Dreadnode:
             project: The default project name to associate all runs with.
             service_name: The service name to use for OpenTelemetry.
             service_version: The service version to use for OpenTelemetry.
-            console: Whether to log span information to the console.
+            console: Whether to log span information to the console (`DREADNODE_CONSOLE` or the default is True).
             send_to_logfire: Whether to send data to Logfire.
             otel_scope: The OpenTelemetry scope name.
         """
@@ -252,7 +258,11 @@ class Dreadnode:
         self.project = project or os.environ.get(ENV_PROJECT)
         self.service_name = service_name
         self.service_version = service_version
-        self.console = console
+        self.console = console or os.environ.get(ENV_CONSOLE, "true").lower() in [
+            "true",
+            "1",
+            "yes",
+        ]
         self.send_to_logfire = send_to_logfire
         self.otel_scope = otel_scope
 
@@ -342,19 +352,19 @@ class Dreadnode:
             #         )
             #     )
             # )
-
-            credentials = self._api.get_user_data_credentials()
-            resolved_endpoint = resolve_endpoint(credentials.endpoint)
+            self._credentials = self._api.get_user_data_credentials()
+            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,
@@ -401,6 +411,43 @@ class Dreadnode:
 
         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()
+                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,
@@ -773,6 +820,7 @@ class Dreadnode:
             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:
@@ -819,6 +867,7 @@ class Dreadnode:
             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:

{dreadnode-1.13.0 → dreadnode-1.13.2}/dreadnode/serialization.py

@@ -519,7 +519,7 @@ def _serialize(obj: t.Any, seen: set[int] | None = None) -> tuple[JsonValue, Jso
     # Primitives early
 
     if isinstance(obj, str | int | float | bool) or obj is None:
-        return obj, EMPTY_SCHEMA
+        return obj, {}
 
     # Cycle tracking
 

dreadnode-1.13.2/dreadnode/storage_utils.py

@@ -0,0 +1,37 @@
+import functools
+import typing as t
+
+from dreadnode.util import logger
+
+
+def with_credential_refresh(func: t.Callable[..., t.Any]) -> t.Callable[..., t.Any]:
+    """Decorator that automatically handles credential refresh on storage errors."""
+
+    @functools.wraps(func)
+    def wrapper(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        # Try to refresh credentials before operation
+        if hasattr(self, "_refresh_credentials_if_needed"):
+            self._refresh_credentials_if_needed()
+
+        try:
+            return func(self, *args, **kwargs)
+        except Exception as e:
+            error_str = str(e)
+            if any(
+                error in error_str
+                for error in [
+                    "ExpiredToken",
+                    "TokenRefreshRequired",
+                    "InvalidAccessKeyId",
+                    "The Access Key Id you provided does not exist",
+                ]
+            ):
+                logger.info("Storage credential error, forcing refresh and retrying")
+
+                if hasattr(self, "_refresh_credentials_if_needed"):
+                    self._refresh_credentials_if_needed()
+
+                return func(self, *args, **kwargs)
+            raise
+
+    return wrapper

{dreadnode-1.13.0 → dreadnode-1.13.2}/dreadnode/tracing/span.py

@@ -36,6 +36,7 @@ from dreadnode.convert import run_span_to_graph
 from dreadnode.metric import Metric, MetricAggMode, MetricsDict
 from dreadnode.object import Object, ObjectRef, ObjectUri, ObjectVal
 from dreadnode.serialization import Serialized, serialize
+from dreadnode.storage_utils import with_credential_refresh
 from dreadnode.tracing.constants import (
     EVENT_ATTRIBUTE_LINK_HASH,
     EVENT_ATTRIBUTE_OBJECT_HASH,
@@ -365,6 +366,7 @@ class RunSpan(Span):
         update_frequency: int = 5,
         run_id: str | ULID | None = None,
         type: SpanType = "run",
+        credential_refresher: t.Callable[[], bool] | None = None,
     ) -> None:
         self.autolog = autolog
         self.project = project
@@ -375,7 +377,9 @@ class RunSpan(Span):
         self._object_schemas: dict[str, JsonDict] = {}
         self._inputs: list[ObjectRef] = []
         self._outputs: list[ObjectRef] = []
-        self._artifact_storage = ArtifactStorage(file_system=file_system)
+        self._artifact_storage = ArtifactStorage(
+            file_system=file_system, credential_refresher=credential_refresher
+        )
         self._artifacts: list[DirectoryNode] = []
         self._artifact_merger = ArtifactMerger()
         self._artifact_tree_builder = ArtifactTreeBuilder(
@@ -406,6 +410,7 @@ class RunSpan(Span):
             SPAN_ATTRIBUTE_PROJECT: project,
             **(attributes or {}),
         }
+        self._credential_refresher = credential_refresher
         super().__init__(name, tracer, attributes=attributes, type=type, tags=tags)
 
     @classmethod
@@ -415,6 +420,7 @@ class RunSpan(Span):
         tracer: Tracer,
         file_system: AbstractFileSystem,
         prefix_path: str,
+        credential_refresher: t.Callable[[], bool] | None = None,
     ) -> "RunSpan":
         self = RunSpan(
             name=f"run.{context['run_id']}.fragment",
@@ -425,6 +431,7 @@ class RunSpan(Span):
             prefix_path=prefix_path,
             type="run_fragment",
             run_id=context["run_id"],
+            credential_refresher=credential_refresher,
         )
 
         self._remote_context = context["trace_context"]
@@ -500,6 +507,10 @@ class RunSpan(Span):
         if self._context_token is not None:
             current_run_span.reset(self._context_token)
 
+    def _refresh_credentials_if_needed(self) -> None:
+        if self._credential_refresher:
+            self._credential_refresher()
+
     def push_update(self, *, force: bool = False) -> None:
         if self._span is None:
             return
@@ -604,6 +615,7 @@ class RunSpan(Span):
 
         return composite_hash
 
+    @with_credential_refresh
     def _store_file_by_hash(self, data: bytes, full_path: str) -> str:
         """
         Writes data to the given full_path in the object store if it doesn't already exist.

{dreadnode-1.13.0 → dreadnode-1.13.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreadnode"
-version = "1.13.0"
+version = "1.13.2"
 description = "Dreadnode SDK"
 requires-python = ">=3.10,<3.14"