dreadnode 1.0.0rc3__tar.gz → 1.0.2__tar.gz

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dreadnode
-Version: 1.0.0rc3
+Version: 1.0.2
 Summary: Dreadnode SDK
 Author: Nick Landers
 Author-email: monoxgas@gmail.com
@@ -12,7 +12,6 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: training
 Requires-Dist: coolname (>=2.2.0,<3.0.0)
-Requires-Dist: fast-depends (>=2.4.12,<3.0.0)
 Requires-Dist: fsspec[s3] (>=2023.1.0,<=2025.3.0)
 Requires-Dist: httpx (>=0.28.0,<0.29.0)
 Requires-Dist: logfire (>=3.5.3,<4.0.0)

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/api/client.py

@@ -7,18 +7,29 @@ import pandas as pd
 from pydantic import BaseModel
 from ulid import ULID
 
+from dreadnode.api.util import (
+    convert_flat_tasks_to_tree,
+    convert_flat_trace_to_tree,
+    process_run,
+    process_task,
+)
 from dreadnode.util import logger
 from dreadnode.version import VERSION
 
 from .models import (
     MetricAggregationType,
     Project,
+    RawRun,
+    RawTask,
     Run,
+    RunSummary,
     StatusFilter,
     Task,
+    TaskTree,
     TimeAggregationType,
     TimeAxisType,
     TraceSpan,
+    TraceTree,
     UserDataCredentials,
 )
 
@@ -119,27 +130,63 @@ class ApiClient:
         response = self.request("GET", f"/strikes/projects/{project!s}")
         return Project(**response.json())
 
-    def list_runs(self, project: str) -> list[Run]:
+    def list_runs(self, project: str) -> list[RunSummary]:
         response = self.request("GET", f"/strikes/projects/{project!s}/runs")
-        return [Run(**run) for run in response.json()]
+        return [RunSummary(**run) for run in response.json()]
 
-    def get_run(self, run: str | ULID) -> Run:
+    def _get_run(self, run: str | ULID) -> RawRun:
         response = self.request("GET", f"/strikes/projects/runs/{run!s}")
-        return Run(**response.json())
-
-    def get_run_tasks(self, run: str | ULID) -> list[Task]:
-        response = self.request("GET", f"/strikes/projects/runs/{run!s}/tasks")
-        return [Task(**task) for task in response.json()]
+        return RawRun(**response.json())
 
-    def get_run_trace(self, run: str | ULID) -> list[Task | TraceSpan]:
-        response = self.request("GET", f"/strikes/projects/runs/{run!s}/spans")
-        spans: list[Task | TraceSpan] = []
+    def get_run(self, run: str | ULID) -> Run:
+        return process_run(self._get_run(run))
+
+    TraceFormat = t.Literal["tree", "flat"]
+
+    @t.overload
+    def get_run_tasks(
+        self, run: str | ULID, *, format: t.Literal["tree"] = "tree"
+    ) -> list[TaskTree]: ...
+
+    @t.overload
+    def get_run_tasks(
+        self, run: str | ULID, *, format: t.Literal["flat"] = "flat"
+    ) -> list[Task]: ...
+
+    def get_run_tasks(
+        self, run: str | ULID, *, format: TraceFormat = "flat"
+    ) -> list[Task] | list[TaskTree]:
+        raw_run = self._get_run(run)
+        response = self.request("GET", f"/strikes/projects/runs/{run!s}/tasks/full")
+        raw_tasks = [RawTask(**task) for task in response.json()]
+        tasks = [process_task(task, raw_run) for task in raw_tasks]
+        tasks = sorted(tasks, key=lambda x: x.timestamp)
+        return tasks if format == "flat" else convert_flat_tasks_to_tree(tasks)
+
+    @t.overload
+    def get_run_trace(
+        self, run: str | ULID, *, format: t.Literal["tree"] = "tree"
+    ) -> list[TraceTree]: ...
+
+    @t.overload
+    def get_run_trace(
+        self, run: str | ULID, *, format: t.Literal["flat"] = "flat"
+    ) -> list[Task | TraceSpan]: ...
+
+    def get_run_trace(
+        self, run: str | ULID, *, format: TraceFormat = "flat"
+    ) -> list[Task | TraceSpan] | list[TraceTree]:
+        raw_run = self._get_run(run)
+        response = self.request("GET", f"/strikes/projects/runs/{run!s}/spans/full")
+        trace: list[Task | TraceSpan] = []
         for item in response.json():
             if "parent_task_span_id" in item:
-                spans.append(Task(**item))
+                trace.append(process_task(RawTask(**item), raw_run))
             else:
-                spans.append(TraceSpan(**item))
-        return spans
+                trace.append(TraceSpan(**item))
+
+        trace = sorted(trace, key=lambda x: x.timestamp)
+        return trace if format == "flat" else convert_flat_trace_to_tree(trace)
 
     # Data exports
 

dreadnode-1.0.2/dreadnode/api/models.py

@@ -0,0 +1,304 @@
+import contextlib
+import typing as t
+from datetime import datetime
+from functools import cached_property
+from uuid import UUID
+
+import requests
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    TypeAdapter,
+    ValidationError,
+    field_validator,
+)
+from ulid import ULID
+
+AnyDict = dict[str, t.Any]
+
+# User
+
+
+class UserAPIKey(BaseModel):
+    key: str
+
+
+class UserResponse(BaseModel):
+    id: UUID
+    email_address: str
+    username: str
+    api_key: UserAPIKey
+
+
+# Strikes
+
+SpanStatus = t.Literal[
+    "pending",  # A pending span has been created
+    "completed",  # The span has been finished
+    "failed",  # The raised an exception
+]
+
+ExportFormat = t.Literal["csv", "json", "jsonl", "parquet"]
+StatusFilter = t.Literal["all", "completed", "failed"]
+TimeAxisType = t.Literal["wall", "relative", "step"]
+TimeAggregationType = t.Literal["max", "min", "sum", "count"]
+MetricAggregationType = t.Literal[
+    "avg",
+    "median",
+    "min",
+    "max",
+    "sum",
+    "first",
+    "last",
+    "count",
+    "std",
+    "var",
+]
+
+
+class SpanException(BaseModel):
+    type: str
+    message: str
+    stacktrace: str
+
+
+class SpanEvent(BaseModel):
+    timestamp: datetime
+    name: str
+    attributes: AnyDict
+
+
+class SpanLink(BaseModel):
+    trace_id: str
+    span_id: str
+    attributes: AnyDict
+
+
+class TraceLog(BaseModel):
+    timestamp: datetime
+    body: str
+    severity: str
+    service: str | None
+    trace_id: str | None
+    span_id: str | None
+    attributes: AnyDict
+    container: str | None
+
+
+class TraceSpan(BaseModel):
+    timestamp: datetime
+    duration: int
+    trace_id: str = Field(repr=False)
+    span_id: str
+    parent_span_id: str | None = Field(repr=False)
+    service_name: str | None = Field(repr=False)
+    status: SpanStatus
+    exception: SpanException | None
+    name: str
+    attributes: AnyDict = Field(repr=False)
+    resource_attributes: AnyDict = Field(repr=False)
+    events: list[SpanEvent] = Field(repr=False)
+    links: list[SpanLink] = Field(repr=False)
+
+
+class Metric(BaseModel):
+    value: float
+    step: int
+    timestamp: datetime
+    attributes: AnyDict
+
+
+class ObjectRef(BaseModel):
+    name: str
+    label: str
+    hash: str
+
+
+class RawObjectUri(BaseModel):
+    hash: str
+    schema_hash: str
+    uri: str
+    size: int
+    type: t.Literal["uri"]
+
+
+class RawObjectVal(BaseModel):
+    hash: str
+    schema_hash: str
+    value: t.Any
+    type: t.Literal["val"]
+
+
+RawObject = RawObjectUri | RawObjectVal
+
+
+class V0Object(BaseModel):
+    name: str
+    label: str
+    value: t.Any
+
+
+class ObjectVal(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    label: str
+    hash: str = Field(repr=False)
+    schema_: AnyDict
+    schema_hash: str = Field(repr=False)
+    value: t.Any
+
+    @field_validator("value")
+    @classmethod
+    def validate_value(cls, value: t.Any) -> t.Any:
+        if isinstance(value, str):
+            with contextlib.suppress(ValidationError):
+                return TypeAdapter(t.Any).validate_json(value)
+
+        return value
+
+
+class ObjectUri(BaseModel):
+    name: str
+    label: str
+    hash: str = Field(repr=False)
+    schema_: AnyDict
+    schema_hash: str = Field(repr=False)
+    uri: str
+    size: int
+
+    _value: t.Any = PrivateAttr(default=None)
+
+    @cached_property
+    def value(self) -> t.Any:
+        if self._value is not None:
+            return self._value
+
+        try:
+            response = requests.get(self.uri, timeout=5)
+            response.raise_for_status()
+            self._value = response.text
+        except requests.RequestException as e:
+            raise RuntimeError(f"Failed to fetch object from {self.uri}") from e
+
+        if isinstance(self._value, str):
+            with contextlib.suppress(ValidationError):
+                self._value = TypeAdapter(t.Any).validate_json(self._value)
+
+        return self._value
+
+
+Object = ObjectVal | ObjectUri
+
+
+class ArtifactFile(BaseModel):
+    hash: str
+    uri: str
+    size_bytes: int
+    final_real_path: str
+
+
+class ArtifactDir(BaseModel):
+    dir_path: str
+    hash: str
+    children: list[t.Union["ArtifactDir", ArtifactFile]]
+
+
+class RunSummary(BaseModel):
+    id: ULID
+    name: str
+    span_id: str = Field(repr=False)
+    trace_id: str = Field(repr=False)
+    timestamp: datetime
+    duration: int
+    status: SpanStatus
+    exception: SpanException | None
+    tags: set[str]
+    params: AnyDict = Field(repr=False)
+    metrics: dict[str, list[Metric]] = Field(repr=False)
+
+
+class RawRun(RunSummary):
+    inputs: list[ObjectRef] = Field(repr=False)
+    outputs: list[ObjectRef] = Field(repr=False)
+    objects: dict[str, RawObject] = Field(repr=False)
+    object_schemas: AnyDict = Field(repr=False)
+    artifacts: list[ArtifactDir] = Field(repr=False)
+    schema_: AnyDict = Field(alias="schema", repr=False)
+
+
+class Run(RunSummary):
+    inputs: dict[str, Object] = Field(repr=False)
+    outputs: dict[str, Object] = Field(repr=False)
+    artifacts: list[ArtifactDir] = Field(repr=False)
+    schema_: AnyDict = Field(alias="schema", repr=False)
+
+
+class _Task(BaseModel):
+    name: str
+    span_id: str
+    trace_id: str = Field(repr=False)
+    parent_span_id: str | None = Field(repr=False)
+    parent_task_span_id: str | None = Field(repr=False)
+    timestamp: datetime
+    duration: int
+    status: SpanStatus
+    exception: SpanException | None
+    tags: set[str]
+    params: AnyDict = Field(repr=False)
+    metrics: dict[str, list[Metric]] = Field(repr=False)
+    schema_: AnyDict = Field(alias="schema", repr=False)
+    attributes: AnyDict = Field(repr=False)
+    resource_attributes: AnyDict = Field(repr=False)
+    events: list[SpanEvent] = Field(repr=False)
+    links: list[SpanLink] = Field(repr=False)
+
+
+class RawTask(_Task):
+    inputs: list[ObjectRef] | list[V0Object] = Field(repr=False)
+    outputs: list[ObjectRef] | list[V0Object] = Field(repr=False)
+
+
+class Task(_Task):
+    inputs: dict[str, Object] = Field(repr=False)
+    outputs: dict[str, Object] = Field(repr=False)
+
+
+class Project(BaseModel):
+    id: UUID = Field(repr=False)
+    key: str
+    name: str
+    description: str | None = Field(repr=False)
+    created_at: datetime
+    updated_at: datetime
+    run_count: int
+    last_run: RawRun | None = Field(repr=False)
+
+
+# Derived types
+
+
+class TaskTree(BaseModel):
+    task: Task
+    children: list["TaskTree"] = []
+
+
+class TraceTree(BaseModel):
+    span: Task | TraceSpan
+    children: list["TraceTree"] = []
+
+
+# User data credentials
+
+
+class UserDataCredentials(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str
+    expiration: datetime
+    region: str
+    bucket: str
+    prefix: str
+    endpoint: str | None

dreadnode-1.0.2/dreadnode/api/util.py

@@ -0,0 +1,173 @@
+from logging import getLogger
+
+from .models import (
+    Object,
+    ObjectUri,
+    ObjectVal,
+    RawObjectVal,
+    RawRun,
+    RawTask,
+    Run,
+    Task,
+    TaskTree,
+    TraceSpan,
+    TraceTree,
+    V0Object,
+)
+
+logger = getLogger(__name__)
+
+
+def convert_flat_tasks_to_tree(tasks: list[Task]) -> list[TaskTree]:
+    tree_nodes: dict[str, TaskTree] = {}
+    for task in tasks:
+        tree_nodes[task.span_id] = TaskTree(task=task)
+
+    roots: list[TaskTree] = []
+    for task in tasks:
+        if task.parent_task_span_id not in tree_nodes:
+            roots.append(tree_nodes[task.span_id])
+        else:
+            parent_node = tree_nodes.get(task.parent_task_span_id)
+            if parent_node:
+                parent_node.children.append(tree_nodes[task.span_id])
+
+    return roots
+
+
+def convert_flat_trace_to_tree(trace: list[Task | TraceSpan]) -> list[TraceTree]:
+    tree_nodes: dict[str, TraceTree] = {}
+    for span in trace:
+        tree_nodes[span.span_id] = TraceTree(span=span)
+
+    roots: list[TraceTree] = []
+    for span in trace:
+        if span.parent_span_id not in tree_nodes:
+            roots.append(tree_nodes[span.span_id])
+        else:
+            parent_node = tree_nodes.get(span.parent_span_id)
+            if parent_node:
+                parent_node.children.append(tree_nodes[span.span_id])
+
+    return roots
+
+
+def process_run(run: RawRun) -> Run:
+    inputs: dict[str, Object] = {}
+    outputs: dict[str, Object] = {}
+
+    for references, converted in ((run.inputs, inputs), (run.outputs, outputs)):
+        for ref in references:
+            if (_object := run.objects.get(ref.hash)) is None:
+                logger.error("Object %s not found in run %s", ref.hash, run.id)
+                continue
+
+            if (_schema := run.object_schemas.get(_object.schema_hash)) is None:
+                logger.error("Schema for object %s not found in run %s", ref.hash, run.id)
+                continue
+
+            if isinstance(_object, RawObjectVal):
+                converted[ref.name] = ObjectVal(
+                    name=ref.name,
+                    label=ref.label,
+                    hash=ref.hash,
+                    schema_=_schema,
+                    schema_hash=_object.schema_hash,
+                    value=_object.value,
+                )
+            else:
+                converted[ref.name] = ObjectUri(
+                    name=ref.name,
+                    label=ref.label,
+                    hash=ref.hash,
+                    schema_=_schema,
+                    schema_hash=_object.schema_hash,
+                    uri=_object.uri,
+                    size=_object.size,
+                )
+
+    return Run(
+        id=run.id,
+        name=run.name,
+        span_id=run.span_id,
+        trace_id=run.trace_id,
+        timestamp=run.timestamp,
+        duration=run.duration,
+        status=run.status,
+        exception=run.exception,
+        tags=run.tags,
+        params=run.params,
+        metrics=run.metrics,
+        inputs=inputs,
+        outputs=outputs,
+        artifacts=run.artifacts,
+        schema=run.schema_,
+    )
+
+
+def process_task(task: RawTask, run: RawRun) -> Task:
+    inputs: dict[str, Object] = {}
+    outputs: dict[str, Object] = {}
+
+    for references, converted in ((task.inputs, inputs), (task.outputs, outputs)):
+        for ref in references:
+            if isinstance(ref, V0Object):
+                converted[ref.name] = ObjectVal(
+                    name=ref.name,
+                    label=ref.label,
+                    hash="",
+                    schema_={},
+                    schema_hash="",
+                    value=ref.value,
+                )
+                continue
+
+            if (_object := run.objects.get(ref.hash)) is None:
+                logger.error("Object %s not found in run %s", ref.hash, run.id)
+                continue
+
+            if (_schema := run.object_schemas.get(_object.schema_hash)) is None:
+                logger.error("Schema for object %s not found in run %s", ref.hash, run.id)
+                continue
+
+            if isinstance(_object, RawObjectVal):
+                converted[ref.name] = ObjectVal(
+                    name=ref.name,
+                    label=ref.label,
+                    hash=ref.hash,
+                    schema_=_schema,
+                    schema_hash=_object.schema_hash,
+                    value=_object.value,
+                )
+            else:
+                converted[ref.name] = ObjectUri(
+                    name=ref.name,
+                    label=ref.label,
+                    hash=ref.hash,
+                    schema_=_schema,
+                    schema_hash=_object.schema_hash,
+                    uri=_object.uri,
+                    size=_object.size,
+                )
+
+    return Task(
+        name=task.name,
+        span_id=task.span_id,
+        trace_id=task.trace_id,
+        parent_span_id=task.parent_span_id,
+        parent_task_span_id=task.parent_task_span_id,
+        timestamp=task.timestamp,
+        duration=task.duration,
+        status=task.status,
+        exception=task.exception,
+        tags=task.tags,
+        params=task.params,
+        metrics=task.metrics,
+        inputs=inputs,
+        outputs=outputs,
+        schema=task.schema_,
+        attributes=task.attributes,
+        resource_attributes=task.resource_attributes,
+        events=task.events,
+        links=task.links,
+    )

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/integrations/transformers.py

@@ -123,7 +123,7 @@ class DreadnodeCallback(TrainerCallback):
         if self._run is None or state.epoch is None:
             return
 
-        dn.log_metric("epoch", state.epoch)
+        dn.log_metric("epoch", state.epoch, to="run")
 
         self._epoch_span = dn.task_span(f"Epoch {state.epoch}")
         self._epoch_span.__enter__()
@@ -149,7 +149,7 @@ class DreadnodeCallback(TrainerCallback):
         if self._run is None:
             return
 
-        dn.log_metric("step", state.global_step)
+        dn.log_metric("step", state.global_step, to="run")
 
         self._step_span = dn.span(f"Step {state.global_step}")
         self._step_span.__enter__()
@@ -178,6 +178,6 @@ class DreadnodeCallback(TrainerCallback):
 
         for key, value in _clean_keys(logs).items():
             if isinstance(value, float | int):
-                dn.log_metric(key, value, step=state.global_step)
+                dn.log_metric(key, value, step=state.global_step, to="run")
 
         dn.push_update()

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/main.py

@@ -7,7 +7,7 @@ import typing as t
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from urllib.parse import urljoin
+from urllib.parse import urljoin, urlparse, urlunparse
 
 import coolname  # type: ignore [import-untyped]
 import logfire
@@ -47,7 +47,10 @@ from dreadnode.tracing.span import (
     current_task_span,
 )
 from dreadnode.types import (
+    INHERITED,
     AnyDict,
+    Inherited,
+    JsonDict,
     JsonValue,
 )
 from dreadnode.util import handle_internal_errors
@@ -204,6 +207,16 @@ class Dreadnode:
                 category=DreadnodeConfigWarning,
             )
 
+        if self.server:
+            parsed_url = urlparse(self.server)
+            if not parsed_url.scheme:
+                netloc = parsed_url.path.split("/")[0]
+                path = "/".join(parsed_url.path.split("/")[1:])
+                parsed_new = parsed_url._replace(
+                    scheme="https", netloc=netloc, path=f"/{path}" if path else ""
+                )
+                self.server = urlunparse(parsed_new)
+
         if self.local_dir is not False:
             config = FileExportConfig(
                 base_path=self.local_dir,
@@ -411,8 +424,8 @@ class Dreadnode:
         name: str | None = None,
         label: str | None = None,
         log_params: t.Sequence[str] | bool = False,
-        log_inputs: t.Sequence[str] | bool = True,
-        log_output: bool = True,
+        log_inputs: t.Sequence[str] | bool | Inherited = INHERITED,
+        log_output: bool | Inherited = INHERITED,
         tags: t.Sequence[str] | None = None,
         **attributes: t.Any,
     ) -> TaskDecorator: ...
@@ -425,8 +438,8 @@ class Dreadnode:
         name: str | None = None,
         label: str | None = None,
         log_params: t.Sequence[str] | bool = False,
-        log_inputs: t.Sequence[str] | bool = True,
-        log_output: bool = True,
+        log_inputs: t.Sequence[str] | bool | Inherited = INHERITED,
+        log_output: bool | Inherited = INHERITED,
         tags: t.Sequence[str] | None = None,
         **attributes: t.Any,
     ) -> ScoredTaskDecorator[R]: ...
@@ -438,8 +451,8 @@ class Dreadnode:
         name: str | None = None,
         label: str | None = None,
         log_params: t.Sequence[str] | bool = False,
-        log_inputs: t.Sequence[str] | bool = True,
-        log_output: bool = True,
+        log_inputs: t.Sequence[str] | bool | Inherited = INHERITED,
+        log_output: bool | Inherited = INHERITED,
         tags: t.Sequence[str] | None = None,
         **attributes: t.Any,
     ) -> TaskDecorator:
@@ -621,6 +634,7 @@ class Dreadnode:
         tags: t.Sequence[str] | None = None,
         params: AnyDict | None = None,
         project: str | None = None,
+        autolog: bool = True,
         **attributes: t.Any,
     ) -> RunSpan:
         """
@@ -646,6 +660,7 @@ class Dreadnode:
             project: The project name to associate the run with. If not provided,
                 the project passed to `configure()` will be used, or the
                 run will be associated with a default project.
+            autolog: Whether to automatically log task inputs, outputs, and execution metrics if unspecified.
             **attributes: Additional attributes to attach to the run span.
         """
         if not self._initialized:
@@ -663,6 +678,7 @@ class Dreadnode:
             tags=tags,
             file_system=self._fs,
             prefix_path=self._fs_prefix,
+            autolog=autolog,
         )
 
     @handle_internal_errors()
@@ -759,6 +775,7 @@ class Dreadnode:
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        attributes: JsonDict | None = None,
         to: ToObject = "task-or-run",
     ) -> None:
         """
@@ -788,6 +805,7 @@ class Dreadnode:
                 - avg: the average of all reported values for this metric
                 - sum: the cumulative sum of all reported values for this metric
                 - count: increment every time this metric is logged - disregard value
+            attributes: A dictionary of additional attributes to attach to the metric.
             to: The target object to log the metric to. Can be "task-or-run" or "run".
                 Defaults to "task-or-run". If "task-or-run", the metric will be logged
                 to the current task or run, whichever is the nearest ancestor.
@@ -842,6 +860,7 @@ class Dreadnode:
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        attributes: JsonDict | None = None,
         to: ToObject = "task-or-run",
     ) -> None:
         task = current_task_span.get()
@@ -854,7 +873,9 @@ class Dreadnode:
         metric = (
             value
             if isinstance(value, Metric)
-            else Metric(float(value), step, timestamp or datetime.now(timezone.utc))
+            else Metric(
+                float(value), step, timestamp or datetime.now(timezone.utc), attributes or {}
+            )
         )
         target.log_metric(key, metric, origin=origin, mode=mode)
 

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/object.py

@@ -1,12 +1,15 @@
 import typing as t
 from dataclasses import dataclass
 
+from dreadnode.types import JsonDict
+
 
 @dataclass
 class ObjectRef:
     name: str
     label: str
     hash: str
+    attributes: JsonDict
 
 
 @dataclass

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/task.py

@@ -9,6 +9,7 @@ from opentelemetry.trace import Tracer
 
 from dreadnode.metric import Scorer, ScorerCallable
 from dreadnode.tracing.span import TaskSpan, current_run_span
+from dreadnode.types import INHERITED, Inherited
 
 P = t.ParamSpec("P")
 R = t.TypeVar("R")
@@ -114,9 +115,9 @@ class Task(t.Generic[P, R]):
 
     log_params: t.Sequence[str] | bool = False
     "Whether to log all, or specific, incoming arguments to the function as parameters."
-    log_inputs: t.Sequence[str] | bool = True
+    log_inputs: t.Sequence[str] | bool | Inherited = INHERITED
     "Whether to log all, or specific, incoming arguments to the function as inputs."
-    log_output: bool = True
+    log_output: bool | Inherited = INHERITED
     "Whether to automatically log the result of the function as an output."
 
     def __post_init__(self) -> None:
@@ -128,6 +129,25 @@ class Task(t.Generic[P, R]):
         self.__name__ = getattr(self.func, "__name__", self.name)
         self.__doc__ = getattr(self.func, "__doc__", None)
 
+    def __get__(self, obj: t.Any, objtype: t.Any) -> "Task[P, R]":
+        if obj is None:
+            return self
+
+        bound_func = self.func.__get__(obj, objtype)
+
+        return Task(
+            tracer=self.tracer,
+            name=self.name,
+            label=self.label,
+            attributes=self.attributes,
+            func=bound_func,
+            scorers=[scorer.clone() for scorer in self.scorers],
+            tags=self.tags.copy(),
+            log_params=self.log_params,
+            log_inputs=self.log_inputs,
+            log_output=self.log_output,
+        )
+
     def _bind_args(self, *args: P.args, **kwargs: P.kwargs) -> dict[str, t.Any]:
         signature = inspect.signature(self.func)
         bound_args = signature.bind(*args, **kwargs)
@@ -220,6 +240,9 @@ class Task(t.Generic[P, R]):
         if run is None or not run.is_recording:
             raise RuntimeError("Tasks must be executed within a run")
 
+        log_inputs = run.autolog if isinstance(self.log_inputs, Inherited) else self.log_inputs
+        log_output = run.autolog if isinstance(self.log_output, Inherited) else self.log_output
+
         bound_args = self._bind_args(*args, **kwargs)
 
         params_to_log = (
@@ -231,9 +254,9 @@ class Task(t.Generic[P, R]):
         )
         inputs_to_log = (
             bound_args
-            if self.log_inputs is True
-            else {k: v for k, v in bound_args.items() if k in self.log_inputs}
-            if self.log_inputs is not False
+            if log_inputs is True
+            else {k: v for k, v in bound_args.items() if k in log_inputs}
+            if log_inputs is not False
             else {}
         )
 
@@ -246,13 +269,16 @@ class Task(t.Generic[P, R]):
             run_id=run.run_id,
             tracer=self.tracer,
         ) as span:
-            span.run.log_metric(f"{self.label}.exec.count", 1, mode="count")
+            if run.autolog:
+                span.run.log_metric(
+                    "count", 1, prefix=f"{self.label}.exec", mode="count", attributes={"auto": True}
+                )
 
             for name, value in params_to_log.items():
                 span.log_param(name, value)
 
             input_object_hashes: list[str] = [
-                span.log_input(name, value, label=f"{self.label}.input.{name}")
+                span.log_input(name, value, label=f"{self.label}.input.{name}", auto=True)
                 for name, value in inputs_to_log.items()
             ]
 
@@ -261,17 +287,29 @@ class Task(t.Generic[P, R]):
                 if inspect.isawaitable(output):
                     output = await output
             except Exception:
-                span.run.log_metric(f"{self.label}.exec.success_rate", 0, mode="avg")
+                if run.autolog:
+                    span.run.log_metric(
+                        "success_rate",
+                        0,
+                        prefix=f"{self.label}.exec",
+                        mode="avg",
+                        attributes={"auto": True},
+                    )
                 raise
 
-            span.run.log_metric(f"{self.label}.exec.success_rate", 1, mode="avg")
+            if run.autolog:
+                span.run.log_metric(
+                    "success_rate",
+                    1,
+                    prefix=f"{self.label}.exec",
+                    mode="avg",
+                    attributes={"auto": True},
+                )
             span.output = output
 
-            if self.log_output:
+            if log_output:
                 output_object_hash = span.log_output(
-                    "output",
-                    output,
-                    label=f"{self.label}.output",
+                    "output", output, label=f"{self.label}.output", auto=True
                 )
 
                 # Link the output to the inputs

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/tracing/span.py

@@ -250,11 +250,15 @@ class RunSpan(Span):
         tracer: Tracer,
         file_system: AbstractFileSystem,
         prefix_path: str,
+        *,
         params: AnyDict | None = None,
         metrics: MetricDict | None = None,
         run_id: str | None = None,
         tags: t.Sequence[str] | None = None,
+        autolog: bool = True,
     ) -> None:
+        self.autolog = autolog
+
         self._params = params or {}
         self._metrics = metrics or {}
         self._objects: dict[str, Object] = {}
@@ -301,6 +305,7 @@ class RunSpan(Span):
     ) -> None:
         self.set_attribute(SPAN_ATTRIBUTE_PARAMS, self._params)
         self.set_attribute(SPAN_ATTRIBUTE_INPUTS, self._inputs, schema=False)
+        self.set_attribute(SPAN_ATTRIBUTE_OUTPUTS, self._outputs, schema=False)
         self.set_attribute(SPAN_ATTRIBUTE_METRICS, self._metrics, schema=False)
         self.set_attribute(SPAN_ATTRIBUTE_OBJECTS, self._objects, schema=False)
         self.set_attribute(
@@ -485,9 +490,8 @@ class RunSpan(Span):
             value,
             label=label,
             event_name=EVENT_NAME_OBJECT_INPUT,
-            **attributes,
         )
-        self._inputs.append(ObjectRef(name, label=label, hash=hash_))
+        self._inputs.append(ObjectRef(name, label=label, hash=hash_, attributes=attributes))
 
     def log_artifact(
         self,
@@ -527,6 +531,8 @@ class RunSpan(Span):
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        prefix: str | None = None,
+        attributes: JsonDict | None = None,
     ) -> None: ...
 
     @t.overload
@@ -537,6 +543,7 @@ class RunSpan(Span):
         *,
         origin: t.Any | None = None,
         mode: MetricAggMode | None = None,
+        prefix: str | None = None,
     ) -> None: ...
 
     def log_metric(
@@ -548,13 +555,21 @@ class RunSpan(Span):
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        prefix: str | None = None,
+        attributes: JsonDict | None = None,
     ) -> None:
         metric = (
             value
             if isinstance(value, Metric)
-            else Metric(float(value), step, timestamp or datetime.now(timezone.utc))
+            else Metric(
+                float(value), step, timestamp or datetime.now(timezone.utc), attributes or {}
+            )
         )
 
+        key = re.sub(r"[^\w/]+", "_", key.lower())
+        if prefix is not None:
+            key = f"{prefix}.{key}"
+
         if origin is not None:
             origin_hash = self.log_object(
                 origin,
@@ -585,9 +600,8 @@ class RunSpan(Span):
             value,
             label=label,
             event_name=EVENT_NAME_OBJECT_OUTPUT,
-            **attributes,
         )
-        self._outputs.append(ObjectRef(name, label=label, hash=hash_))
+        self._outputs.append(ObjectRef(name, label=label, hash=hash_, attributes=attributes))
 
 
 class TaskSpan(Span, t.Generic[R]):
@@ -689,9 +703,8 @@ class TaskSpan(Span, t.Generic[R]):
             value,
             label=label,
             event_name=EVENT_NAME_OBJECT_OUTPUT,
-            **attributes,
         )
-        self._outputs.append(ObjectRef(name, label=label, hash=hash_))
+        self._outputs.append(ObjectRef(name, label=label, hash=hash_, attributes=attributes))
         return hash_
 
     @property
@@ -721,9 +734,8 @@ class TaskSpan(Span, t.Generic[R]):
             value,
             label=label,
             event_name=EVENT_NAME_OBJECT_INPUT,
-            **attributes,
         )
-        self._inputs.append(ObjectRef(name, label=label, hash=hash_))
+        self._inputs.append(ObjectRef(name, label=label, hash=hash_, attributes=attributes))
         return hash_
 
     @property
@@ -740,6 +752,7 @@ class TaskSpan(Span, t.Generic[R]):
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        attributes: JsonDict | None = None,
     ) -> None: ...
 
     @t.overload
@@ -761,13 +774,18 @@ class TaskSpan(Span, t.Generic[R]):
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
         mode: MetricAggMode | None = None,
+        attributes: JsonDict | None = None,
     ) -> None:
         metric = (
             value
             if isinstance(value, Metric)
-            else Metric(float(value), step, timestamp or datetime.now(timezone.utc))
+            else Metric(
+                float(value), step, timestamp or datetime.now(timezone.utc), attributes or {}
+            )
         )
 
+        key = re.sub(r"[^\w/]+", "_", key.lower())
+
         if origin is not None:
             origin_hash = self.run.log_object(
                 origin,
@@ -786,7 +804,7 @@ class TaskSpan(Span, t.Generic[R]):
         #
         # Don't include `source` and `mode` as we handled it here.
         if (run := current_run_span.get()) is not None:
-            run.log_metric(f"{self._label}.{key}", metric)
+            run.log_metric(key, metric, prefix=self._label)
 
     def get_average_metric_value(self, key: str | None = None) -> float:
         metrics = (

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/dreadnode/types.py

@@ -23,3 +23,11 @@ class Unset:
 
 
 UNSET: Unset = Unset()
+
+
+class Inherited:
+    def __repr__(self) -> str:
+        return "Inherited"
+
+
+INHERITED: Inherited = Inherited()

{dreadnode-1.0.0rc3 → dreadnode-1.0.2}/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "dreadnode"
-version = "1.0.0-rc.3"
+version = "1.0.2"
 description = "Dreadnode SDK"
 requires-python = ">=3.10,<3.14"
 
 [tool.poetry]
 name = "dreadnode"
-version = "1.0.0-rc.3"
+version = "1.0.2"
 description = "Dreadnode SDK"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 repository = "https://github.com/dreadnode/sdk"
@@ -18,7 +18,6 @@ pydantic = "^2.9.2"
 httpx = "^0.28.0"
 logfire = "^3.5.3"
 python-ulid = "^3.0.0"
-fast-depends = "^2.4.12"
 coolname = "^2.2.0"
 pandas = "^2.2.3"
 fsspec = {version = ">=2023.1.0,<=2025.3.0", extras = ["s3"]} # Pinned for datasets compatibility
@@ -31,7 +30,7 @@ training = ["transformers"]
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.8.0"
 ruff = "^0.11.6"
-pre-commit = "^3.8.0"
+pre-commit = "^4.0.0"
 pytest = "^8.3.3"
 pytest-asyncio = "^0.24.0"
 types-protobuf = "^5.29.1.20250208"

dreadnode-1.0.0rc3/dreadnode/api/models.py

@@ -1,210 +0,0 @@
-import typing as t
-from datetime import datetime
-from uuid import UUID
-
-from pydantic import BaseModel, Field
-from ulid import ULID
-
-AnyDict = dict[str, t.Any]
-
-# User
-
-
-class UserAPIKey(BaseModel):
-    key: str
-
-
-class UserResponse(BaseModel):
-    id: UUID
-    email_address: str
-    username: str
-    api_key: UserAPIKey
-
-
-# Strikes
-
-SpanStatus = t.Literal[
-    "pending",  # A pending span has been created
-    "completed",  # The span has been finished
-    "failed",  # The raised an exception
-]
-
-ExportFormat = t.Literal["csv", "json", "jsonl", "parquet"]
-StatusFilter = t.Literal["all", "completed", "failed"]
-TimeAxisType = t.Literal["wall", "relative", "step"]
-TimeAggregationType = t.Literal["max", "min", "sum", "count"]
-MetricAggregationType = t.Literal[
-    "avg",
-    "median",
-    "min",
-    "max",
-    "sum",
-    "first",
-    "last",
-    "count",
-    "std",
-    "var",
-]
-
-
-class SpanException(BaseModel):
-    type: str
-    message: str
-    stacktrace: str
-
-
-class SpanEvent(BaseModel):
-    timestamp: datetime
-    name: str
-    attributes: AnyDict
-
-
-class SpanLink(BaseModel):
-    trace_id: str
-    span_id: str
-    attributes: AnyDict
-
-
-class TraceLog(BaseModel):
-    timestamp: datetime
-    body: str
-    severity: str
-    service: str | None
-    trace_id: str | None
-    span_id: str | None
-    attributes: AnyDict
-    container: str | None
-
-
-class TraceSpan(BaseModel):
-    timestamp: datetime
-    duration: int
-    trace_id: str
-    span_id: str
-    parent_span_id: str | None
-    service_name: str | None
-    status: SpanStatus
-    exception: SpanException | None
-    name: str
-    attributes: AnyDict
-    resource_attributes: AnyDict
-    events: list[SpanEvent]
-    links: list[SpanLink]
-
-
-class Metric(BaseModel):
-    value: float
-    step: int
-    timestamp: datetime
-    attributes: AnyDict
-
-
-class ObjectRef(BaseModel):
-    name: str
-    label: str
-    hash: str
-
-
-class ObjectUri(BaseModel):
-    hash: str
-    schema_hash: str
-    uri: str
-    size: int
-    type: t.Literal["uri"]
-
-
-class ObjectVal(BaseModel):
-    hash: str
-    schema_hash: str
-    value: t.Any
-    type: t.Literal["val"]
-
-
-Object = ObjectUri | ObjectVal
-
-
-class V0Object(BaseModel):
-    name: str
-    label: str
-    value: t.Any
-
-
-class Run(BaseModel):
-    id: ULID
-    name: str
-    span_id: str
-    trace_id: str
-    timestamp: datetime
-    duration: int
-    status: SpanStatus
-    exception: SpanException | None
-    tags: set[str]
-    params: AnyDict
-    metrics: dict[str, list[Metric]]
-    inputs: list[ObjectRef]
-    outputs: list[ObjectRef]
-    objects: dict[str, Object]
-    object_schemas: AnyDict
-    schema_: AnyDict = Field(alias="schema")
-
-
-class Task(BaseModel):
-    name: str
-    span_id: str
-    trace_id: str
-    parent_span_id: str | None
-    parent_task_span_id: str | None
-    timestamp: datetime
-    duration: int
-    status: SpanStatus
-    exception: SpanException | None
-    tags: set[str]
-    params: AnyDict
-    metrics: dict[str, list[Metric]]
-    inputs: list[ObjectRef] | list[V0Object]  # v0 compat
-    outputs: list[ObjectRef] | list[V0Object]  # v0 compat
-    schema_: AnyDict = Field(alias="schema")
-    attributes: AnyDict
-    resource_attributes: AnyDict
-    events: list[SpanEvent]
-    links: list[SpanLink]
-
-
-class Project(BaseModel):
-    id: UUID
-    key: str
-    name: str
-    description: str | None
-    created_at: datetime
-    updated_at: datetime
-    run_count: int
-    last_run: Run | None
-
-
-# Derived types
-
-
-class TaskTree(BaseModel):
-    task: Task
-    children: list["TaskTree"] = []
-
-
-class SpanTree(BaseModel):
-    """Tree representation of a trace span with its children"""
-
-    span: Task | TraceSpan
-    children: list["SpanTree"] = []
-
-
-# User data credentials
-
-
-class UserDataCredentials(BaseModel):
-    access_key_id: str
-    secret_access_key: str
-    session_token: str
-    expiration: datetime
-    region: str
-    bucket: str
-    prefix: str
-    endpoint: str | None