hotfix: Metric modes

Author: monoxgas

dreadnode/main.py

@@ -32,7 +32,7 @@
     ENV_SERVER,
     ENV_SERVER_URL,
 )
-from dreadnode.metric import Metric, MetricMode, Scorer, ScorerCallable, T
+from dreadnode.metric import Metric, MetricAggMode, Scorer, ScorerCallable, T
 from dreadnode.task import P, R, Task
 from dreadnode.tracing.exporters import (
     FileExportConfig,
@@ -757,7 +757,7 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
+        mode: MetricAggMode | None = None,
         to: ToObject = "task-or-run",
     ) -> None:
         """
@@ -799,7 +799,7 @@ def log_metric(
         value: Metric,
         *,
         origin: t.Any | None = None,
-        mode: MetricMode = "direct",
+        mode: MetricAggMode | None = None,
         to: ToObject = "task-or-run",
     ) -> None:
         """
@@ -821,7 +821,6 @@ def log_metric(
                 as an input or output anywhere in the run.
             mode: The aggregation mode to use for the metric. Helpful when you want to let
                 the library take care of translating your raw values into better representations.
-                - direct: do not modify the value at all (default)
                 - min: always report the lowest ovbserved value for this metric
                 - max: always report the highest observed value for this metric
                 - avg: report the average of all values for this metric
@@ -841,7 +840,7 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
+        mode: MetricAggMode | None = None,
         to: ToObject = "task-or-run",
     ) -> None:
         task = current_task_span.get()

dreadnode/metric.py

@@ -3,14 +3,19 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 
+from logfire._internal.stack_info import warn_at_user_stacklevel
 from logfire._internal.utils import safe_repr
 from opentelemetry.trace import Tracer
 
 from dreadnode.types import JsonDict, JsonValue
 
 T = t.TypeVar("T")
 
-MetricMode = t.Literal["direct", "avg", "sum", "min", "max", "count"]
+MetricAggMode = t.Literal["avg", "sum", "min", "max", "count"]
+
+
+class MetricWarning(UserWarning):
+    pass
 
 
 @dataclass
@@ -57,7 +62,7 @@ def from_many(
         score_attributes = {name: value for name, value, _ in values}
         return cls(value=total / weight, step=step, attributes={**attributes, **score_attributes})
 
-    def apply_mode(self, mode: MetricMode, others: "list[Metric]") -> "Metric":
+    def apply_mode(self, mode: MetricAggMode, others: "list[Metric]") -> "Metric":
         """
         Apply an aggregation mode to the metric.
         This will modify the metric in place.
@@ -69,15 +74,13 @@ def apply_mode(self, mode: MetricMode, others: "list[Metric]") -> "Metric":
         Returns:
             self
         """
-        previous_mode = next((m.attributes.get("mode") for m in others), mode) or "direct"
-        if mode != previous_mode:
-            raise ValueError(
-                f"Cannot mix metric modes {mode} != {previous_mode}",
+        previous_mode = next((m.attributes.get("mode") for m in others), mode)
+        if previous_mode is not None and mode != previous_mode:
+            warn_at_user_stacklevel(
+                f"Metric logged with different modes ({mode} != {previous_mode}). This may result in unexpected behavior.",
+                MetricWarning,
             )
 
-        if mode == "direct":
-            return self
-
         self.attributes["original"] = self.value
         self.attributes["mode"] = mode
 

dreadnode/tracing/span.py

@@ -28,7 +28,7 @@
 from dreadnode.artifact.storage import ArtifactStorage
 from dreadnode.artifact.tree_builder import ArtifactTreeBuilder, DirectoryNode
 from dreadnode.constants import MAX_INLINE_OBJECT_BYTES
-from dreadnode.metric import Metric, MetricDict, MetricMode
+from dreadnode.metric import Metric, MetricAggMode, MetricDict
 from dreadnode.object import Object, ObjectRef, ObjectUri, ObjectVal
 from dreadnode.serialization import Serialized, serialize
 from dreadnode.types import UNSET, AnyDict, JsonDict, JsonValue, Unset
@@ -526,9 +526,8 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
-    ) -> None:
-        ...
+        mode: MetricAggMode | None = None,
+    ) -> None: ...
 
     @t.overload
     def log_metric(
@@ -537,9 +536,8 @@ def log_metric(
         value: Metric,
         *,
         origin: t.Any | None = None,
-        mode: MetricMode = "direct",
-    ) -> None:
-        ...
+        mode: MetricAggMode | None = None,
+    ) -> None: ...
 
     def log_metric(
         self,
@@ -549,7 +547,7 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
+        mode: MetricAggMode | None = None,
     ) -> None:
         metric = (
             value
@@ -566,7 +564,9 @@ def log_metric(
             metric.attributes[METRIC_ATTRIBUTE_SOURCE_HASH] = origin_hash
 
         metrics = self._metrics.setdefault(key, [])
-        metrics.append(metric.apply_mode(mode, metrics))
+        if mode is not None:
+            metric = metric.apply_mode(mode, metrics)
+        metrics.append(metric)
 
     @property
     def outputs(self) -> AnyDict:
@@ -739,9 +739,8 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
-    ) -> None:
-        ...
+        mode: MetricAggMode | None = None,
+    ) -> None: ...
 
     @t.overload
     def log_metric(
@@ -750,9 +749,8 @@ def log_metric(
         value: Metric,
         *,
         origin: t.Any | None = None,
-        mode: MetricMode = "direct",
-    ) -> None:
-        ...
+        mode: MetricAggMode | None = None,
+    ) -> None: ...
 
     def log_metric(
         self,
@@ -762,7 +760,7 @@ def log_metric(
         step: int = 0,
         origin: t.Any | None = None,
         timestamp: datetime | None = None,
-        mode: MetricMode = "direct",
+        mode: MetricAggMode | None = None,
     ) -> None:
         metric = (
             value
@@ -779,7 +777,9 @@ def log_metric(
             metric.attributes[METRIC_ATTRIBUTE_SOURCE_HASH] = origin_hash
 
         metrics = self._metrics.setdefault(key, [])
-        metrics.append(metric.apply_mode(mode, metrics))
+        if mode is not None:
+            metric = metric.apply_mode(mode, metrics)
+        metrics.append(metric)
 
         # For every metric we log, also log it to the run
         # with our `label` as a prefix.