support custom conditions

Author: Saba9

docs/source/alerts.md

@@ -152,6 +152,34 @@ for alert in alerts:
 | `AlertReason.MIN_EXCEEDED` | `"min_exceeded"` | Metric dropped below `min_value` |
 | `AlertReason.SPIKE` | `"spike"` | Spike detected vs. recent average |
 | `AlertReason.STAGNATION` | `"stagnation"` | No improvement for `patience` steps |
+| `AlertReason.CUSTOM` | `"custom"` | Custom condition returned `True` |
+
+### Custom Conditions
+
+Pass `fn` to [`watch`] to define your own condition. The function receives `(value, step)` and should return `True` to fire a default WARN alert, a list of alert dicts for full control, or a falsy value for no alert:
+
+```python
+def check_divergence(value, step):
+    if value > 50.0:
+        return [
+            {
+                "title": "Loss diverged",
+                "level": trackio.AlertLevel.ERROR,
+                "text": f"val_loss={value:.2f} at step {step}",
+                "data": {"reason": "diverged", "threshold": 50.0, "value": value},
+                "stop": True,
+            }
+        ]
+    return None
+
+trackio.watch("val/loss", fn=check_divergence)
+```
+
+Include `"stop": True` in a returned dict to set `should_stop()` to `True`. Custom conditions can be combined with built-in ones — both run independently on every `log()` call:
+
+```python
+trackio.watch("train/loss", nan=True, fn=check_divergence)
+```
 
 ---
 

tests/e2e-local/test_watchers.py

@@ -1,4 +1,10 @@
-from trackio.watchers import AlertReason, MetricWatcher, WatcherManager
+from trackio.watchers import (
+    AlertLevel,
+    AlertReason,
+    CustomMetricWatcher,
+    MetricWatcher,
+    WatcherManager,
+)
 
 
 def test_nan_inf_triggers_stop():
@@ -108,3 +114,50 @@ def test_manager_should_stop_and_clear():
 def test_non_numeric_ignored():
     w = MetricWatcher("loss", max_value=10.0)
     assert len(w.check("not a number", step=0)) == 0
+
+
+def test_custom_watcher_bool_return():
+    w = CustomMetricWatcher("loss", fn=lambda v, s: v > 10.0)
+    assert len(w.check(5.0, step=0)) == 0
+    alerts = w.check(15.0, step=1)
+    assert len(alerts) == 1
+    assert alerts[0]["level"] == AlertLevel.WARN
+    assert alerts[0]["data"]["reason"] == AlertReason.CUSTOM
+    assert not w.should_stop
+
+
+def test_custom_watcher_dict_return_with_stop():
+    def fn(value, step):
+        if value > 10.0:
+            return [
+                {
+                    "title": "too high",
+                    "level": AlertLevel.ERROR,
+                    "stop": True,
+                    "data": {},
+                }
+            ]
+        return None
+
+    w = CustomMetricWatcher("loss", fn=fn)
+    assert len(w.check(5.0, step=0)) == 0
+    assert not w.should_stop
+    alerts = w.check(15.0, step=1)
+    assert len(alerts) == 1
+    assert alerts[0]["title"] == "too high"
+    assert w.should_stop
+
+
+def test_custom_watcher_none_return():
+    w = CustomMetricWatcher("loss", fn=lambda v, s: None)
+    assert len(w.check(99.0, step=0)) == 0
+
+
+def test_manager_mixes_builtin_and_custom():
+    mgr = WatcherManager()
+    mgr.add(MetricWatcher("loss", max_value=20.0))
+    mgr.add(CustomMetricWatcher("loss", fn=lambda v, s: v > 10.0))
+    alerts = mgr.check({"loss": 15.0}, step=0)
+    assert len(alerts) == 1
+    alerts2 = mgr.check({"loss": 25.0}, step=1)
+    assert len(alerts2) == 2

trackio/__init__.py

@@ -7,7 +7,7 @@
 import warnings
 import webbrowser
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 
 import huggingface_hub
 from gradio_client import handle_file
@@ -41,7 +41,12 @@
 from trackio.trace import Trace
 from trackio.typehints import UploadEntry
 from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR, _emit_nonfatal_warning
-from trackio.watchers import AlertReason, MetricWatcher, WatcherManager
+from trackio.watchers import (
+    AlertReason,
+    CustomMetricWatcher,
+    MetricWatcher,
+    WatcherManager,
+)
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
@@ -843,6 +848,7 @@ def watch(
     min_value: float | None = None,
     window: int = 5,
     mode: str = "min",
+    fn: Callable | None = None,
 ) -> None:
     """
     Register a metric watcher that automatically fires alerts when conditions
@@ -872,6 +878,12 @@ def watch(
         mode (`str`, *optional*, defaults to ``"min"``):
             Whether lower (``"min"``) or higher (``"max"``) values are better.
             Affects patience-based stagnation detection.
+        fn (`Callable[[float, int | None], bool | list[dict] | None]`, *optional*):
+            A custom condition called as ``fn(value, step)`` on every
+            ``trackio.log()`` call. Return ``True`` to fire a default WARN
+            alert, a list of alert dicts for full control, or a falsy value
+            for no alert. Include ``"stop": True`` in a returned dict to
+            also set ``should_stop()`` to ``True``.
     """
     watcher = MetricWatcher(
         metric_name=metric,
@@ -885,6 +897,8 @@ def watch(
         mode=mode,
     )
     _watcher_manager.add(watcher)
+    if fn is not None:
+        _watcher_manager.add(CustomMetricWatcher(metric, fn))
 
 
 def should_stop() -> bool:

trackio/watchers.py

@@ -23,7 +23,7 @@
 
 import math
 from collections import deque
-from typing import Literal
+from typing import Callable, Literal
 
 from trackio.alerts import AlertLevel
 
@@ -43,6 +43,7 @@ class AlertReason:
     MIN_EXCEEDED = "min_exceeded"
     SPIKE = "spike"
     STAGNATION = "stagnation"
+    CUSTOM = "custom"
 
 
 class MetricWatcher:
@@ -212,11 +213,45 @@ def should_stop(self) -> bool:
         return self._triggered_stop
 
 
+class CustomMetricWatcher:
+    def __init__(self, metric_name: str, fn: Callable):
+        self.metric_name = metric_name
+        self._fn = fn
+        self._triggered_stop = False
+
+    def check(self, value, step: int | None = None) -> list[dict]:
+        result = self._fn(value, step)
+        if not result:
+            return []
+        if result is True:
+            return [
+                {
+                    "title": f"Custom condition triggered for {self.metric_name}",
+                    "level": AlertLevel.WARN,
+                    "data": {
+                        "metric": self.metric_name,
+                        "value": value,
+                        "step": step,
+                        "reason": AlertReason.CUSTOM,
+                    },
+                }
+            ]
+        alerts = list(result)
+        for a in alerts:
+            if a.get("stop"):
+                self._triggered_stop = True
+        return alerts
+
+    @property
+    def should_stop(self) -> bool:
+        return self._triggered_stop
+
+
 class WatcherManager:
     def __init__(self):
-        self._watchers: list[MetricWatcher] = []
+        self._watchers: list[MetricWatcher | CustomMetricWatcher] = []
 
-    def add(self, watcher: MetricWatcher):
+    def add(self, watcher: MetricWatcher | CustomMetricWatcher):
         self._watchers.append(watcher)
 
     def check(self, metrics: dict, step: int | None = None) -> list[dict]: