Fix automatic step values (#9)

Author: justusschock

src/litlogger/api/metrics_api.py

@@ -103,6 +103,21 @@ def _to_v1_metrics_tracker(tracker: MetricsTracker) -> V1MetricsTracker:
     return V1MetricsTracker(**kwargs)
 
 
+def _from_v1_metrics_tracker(v1_tracker: V1MetricsTracker) -> MetricsTracker:
+    """Convert V1MetricsTracker from API response to user-facing MetricsTracker."""
+    return MetricsTracker(
+        name=v1_tracker.name,
+        num_rows=v1_tracker.num_rows or 0,
+        min_value=v1_tracker.min_value,
+        max_value=v1_tracker.max_value,
+        min_index=v1_tracker.min_index,
+        max_index=v1_tracker.max_index,
+        last_value=v1_tracker.last_value,
+        last_index=v1_tracker.last_index,
+        max_user_step=v1_tracker.max_user_step,
+    )
+
+
 def _to_v1_phase_type(phase: PhaseType) -> str:
     """Convert user-facing PhaseType to V1PhaseType string.
 
@@ -332,3 +347,17 @@ def update_experiment_metrics(
                 trackers=v1_trackers,
             ),
         )
+
+    def get_trackers_from_metrics_store(self, metrics_store: Any) -> Dict[str, MetricsTracker] | None:
+        """Extract and convert trackers from a metrics store object.
+
+        Args:
+            metrics_store: The metrics store object from the API.
+
+        Returns:
+            Dictionary of MetricsTracker objects, or None if no trackers exist.
+        """
+        if not hasattr(metrics_store, "trackers") or not metrics_store.trackers:
+            return None
+
+        return {name: _from_v1_metrics_tracker(v1_tracker) for name, v1_tracker in metrics_store.trackers.items()}

src/litlogger/background.py

@@ -69,6 +69,7 @@ def __init__(
         store_created_at: bool,
         rate_limiting_interval: int = 1,
         max_batch_size: int = 1000,
+        trackers_init: Dict[str, MetricsTracker] | None = None,
     ) -> None:
         super().__init__(daemon=True)
         self.teamspace_id = teamspace_id
@@ -98,7 +99,7 @@ def __init__(
             client=metrics_api.client,
         )
 
-        self.trackers: Dict[str, MetricsTracker] = {}
+        self.trackers: Dict[str, MetricsTracker] = trackers_init if trackers_init is not None else {}
 
     def run(self) -> None:
         self._run()
@@ -210,6 +211,10 @@ def _update_tracker(self, name: str, values: Metrics) -> None:
 
         # Increment the number of rows
         for value_obj in values.values:
+            # Augment with step from tracker if not provided
+            if value_obj.step is None:
+                value_obj.step = tracker.num_rows
+
             value = float(value_obj.value)
 
             if tracker.started_at is None and self.store_created_at and value_obj.created_at:

src/litlogger/experiment.py

@@ -18,7 +18,6 @@
 import os
 import signal
 import sys
-from collections.abc import Mapping
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 from multiprocessing import JoinableQueue
@@ -165,6 +164,7 @@ def __init__(
             store_created_at=store_created_at,
             rate_limiting_interval=rate_limiting_interval,
             max_batch_size=max_batch_size,
+            trackers_init=self._metrics_api.get_trackers_from_metrics_store(self._metrics_store),
         )
 
         self._manager.start()
@@ -201,7 +201,7 @@ def teamspace(self) -> "Teamspace":
         """
         return self._teamspace
 
-    def log_metrics(self, metrics: Mapping[str, float], step: int | None = None) -> None:
+    def log_metrics(self, metrics: Dict[str, float], step: int | None = None, **kwargs: float) -> None:
         """Log metrics to the experiment with background uploading.
 
         Metrics are buffered locally and uploaded to the cloud in batches to optimize performance.
@@ -211,6 +211,8 @@ def log_metrics(self, metrics: Mapping[str, float], step: int | None = None) ->
             metrics: Dictionary mapping metric names to numeric values. Example: {"loss": 0.5, "accuracy": 0.95}.
             step: Optional step number for this data point (e.g., training step, epoch).
                 If None and store_step=True, no step is recorded.
+            kwargs: Additional metric values. Can be used to provide metrics more natural.
+                Example: loss=0.5, accuracy: 0.95.
 
         Raises:
             RuntimeError: If the background thread encountered an error.
@@ -219,6 +221,8 @@ def log_metrics(self, metrics: Mapping[str, float], step: int | None = None) ->
             raise self._manager.exception
 
         batch: Dict[str, Metrics] = {}
+
+        metrics.update(kwargs)
         for name, value in metrics.items():
             created_at = None
             if self.store_created_at:

tests/integrations/test_standalone.py

@@ -644,6 +644,73 @@ def test_console_output():
     )
 
 
+@pytest.mark.cloud()
+def test_resume_experiment_with_tracker_initialization():
+    """Test that resuming an experiment initializes trackers and augments steps correctly."""
+    experiment_name = f"standalone_resume_tracker-{uuid.uuid4().hex}"
+
+    # First experiment run - log metrics with explicit steps
+    exp1 = litlogger.init(name=experiment_name, teamspace="oss-litlogger")
+
+    for i in range(10):
+        litlogger.log_metrics({"loss": 1.0 - i * 0.1}, step=i)
+
+    litlogger.finalize()
+
+    # Store info for verification
+    project_id = exp1._teamspace.id
+    stream_id = exp1._metrics_store.id
+
+    # Wait for metrics to be available
+    client = LitRestClient()
+    for _ in range(30):
+        response = client.lit_logger_service_get_logger_metrics(project_id=project_id, ids=[stream_id])
+        if response.named_metrics != {}:
+            metrics = response.named_metrics
+            if len(metrics.get("loss", {}).ids_metrics.get(stream_id, {}).metrics_values or []) == 10:
+                break
+        sleep(1)
+
+    # Second experiment run (resume) - log metrics WITHOUT explicit steps
+    # The steps should be augmented from tracker's num_rows (which should be 10)
+    exp2 = litlogger.init(name=experiment_name, teamspace="oss-litlogger")
+
+    # Verify that the experiment resumed (same stream ID means same experiment)
+    assert exp2._metrics_store.id == stream_id, "Expected to resume the same experiment"
+
+    # Log 5 more metrics WITHOUT explicit steps - they should get steps 10-14
+    for i in range(5):
+        litlogger.log_metrics({"loss": 0.05 - i * 0.01})  # No step parameter
+
+    litlogger.finalize()
+
+    # Wait for all metrics to be available
+    for _ in range(30):
+        response = client.lit_logger_service_get_logger_metrics(project_id=project_id, ids=[stream_id])
+        if response.named_metrics != {}:
+            metrics = response.named_metrics
+            loss_values = metrics.get("loss", {}).ids_metrics.get(stream_id, {}).metrics_values or []
+            if len(loss_values) == 15:  # 10 from first run + 5 from second
+                break
+        sleep(1)
+
+    # Verify we have all 15 metrics
+    loss_metrics = response.named_metrics["loss"].ids_metrics[stream_id].metrics_values
+    assert len(loss_metrics) == 15, f"Expected 15 loss metrics, got {len(loss_metrics)}"
+
+    # Verify the steps are sequential (0-9 from first run, 10-14 from second run)
+    # Steps may come back as strings from the API, so convert to int for comparison
+    steps = sorted([int(m.step) for m in loss_metrics])
+    expected_steps = list(range(15))
+    assert steps == expected_steps, f"Expected steps {expected_steps}, got {steps}"
+
+    # Cleanup
+    client.lit_logger_service_delete_metrics_stream(
+        project_id=project_id,
+        body=LitLoggerServiceDeleteMetricsStreamBody(ids=[stream_id]),
+    )
+
+
 @pytest.mark.cloud()
 def test_get_or_create_experiment_metrics():
     """Test get_or_create_experiment_metrics returns existing experiment on second call."""

tests/unittests/api/test_metrics_api.py

@@ -8,10 +8,11 @@
 
 from lightning_sdk.lightning_cloud.openapi import (
     V1Metrics,
+    V1MetricsTracker,
     V1MetricValue,
     V1PhaseType,
 )
-from litlogger.api.metrics_api import MetricsApi
+from litlogger.api.metrics_api import MetricsApi, _from_v1_metrics_tracker
 from litlogger.types import MetricsTracker, PhaseType
 
 
@@ -428,3 +429,136 @@ def test_update_experiment_metrics_custom_phase(self):
         call_args = mock_client.lit_logger_service_update_metrics_stream.call_args
         assert call_args.kwargs["body"].persisted is False
         assert call_args.kwargs["body"].phase == V1PhaseType.RUNNING
+
+
+class TestFromV1MetricsTracker:
+    """Test the _from_v1_metrics_tracker helper function."""
+
+    def test_converts_full_tracker(self):
+        """Test converting a V1MetricsTracker with all fields set."""
+        v1_tracker = V1MetricsTracker(
+            name="loss",
+            num_rows=100,
+            min_value=0.1,
+            max_value=1.0,
+            min_index=50,
+            max_index=0,
+            last_value=0.2,
+            last_index=99,
+            max_user_step=99,
+        )
+
+        result = _from_v1_metrics_tracker(v1_tracker)
+
+        assert isinstance(result, MetricsTracker)
+        assert result.name == "loss"
+        assert result.num_rows == 100
+        assert result.min_value == 0.1
+        assert result.max_value == 1.0
+        assert result.min_index == 50
+        assert result.max_index == 0
+        assert result.last_value == 0.2
+        assert result.last_index == 99
+        assert result.max_user_step == 99
+
+    def test_converts_minimal_tracker(self):
+        """Test converting a V1MetricsTracker with only required fields."""
+        v1_tracker = MagicMock()
+        v1_tracker.name = "accuracy"
+        v1_tracker.num_rows = None
+        v1_tracker.min_value = None
+        v1_tracker.max_value = None
+        v1_tracker.min_index = None
+        v1_tracker.max_index = None
+        v1_tracker.last_value = None
+        v1_tracker.last_index = None
+        v1_tracker.max_user_step = None
+
+        result = _from_v1_metrics_tracker(v1_tracker)
+
+        assert isinstance(result, MetricsTracker)
+        assert result.name == "accuracy"
+        assert result.num_rows == 0  # Defaults to 0 when None
+        assert result.min_value is None
+        assert result.max_value is None
+
+    def test_converts_tracker_with_zero_num_rows(self):
+        """Test converting a tracker with explicit zero num_rows."""
+        v1_tracker = V1MetricsTracker(
+            name="metric",
+            num_rows=0,
+        )
+
+        result = _from_v1_metrics_tracker(v1_tracker)
+
+        assert result.num_rows == 0
+
+
+class TestGetTrackersFromMetricsStore:
+    """Test the get_trackers_from_metrics_store method."""
+
+    def test_returns_none_when_no_trackers_attribute(self):
+        """Test returns None when metrics store has no trackers attribute."""
+        mock_client = MagicMock()
+        api = MetricsApi(client=mock_client)
+
+        mock_metrics_store = MagicMock(spec=[])  # No attributes
+
+        result = api.get_trackers_from_metrics_store(mock_metrics_store)
+
+        assert result is None
+
+    def test_returns_none_when_trackers_is_none(self):
+        """Test returns None when metrics store trackers is None."""
+        mock_client = MagicMock()
+        api = MetricsApi(client=mock_client)
+
+        mock_metrics_store = MagicMock()
+        mock_metrics_store.trackers = None
+
+        result = api.get_trackers_from_metrics_store(mock_metrics_store)
+
+        assert result is None
+
+    def test_returns_none_when_trackers_is_empty(self):
+        """Test returns None when metrics store trackers is empty dict."""
+        mock_client = MagicMock()
+        api = MetricsApi(client=mock_client)
+
+        mock_metrics_store = MagicMock()
+        mock_metrics_store.trackers = {}
+
+        result = api.get_trackers_from_metrics_store(mock_metrics_store)
+
+        assert result is None
+
+    def test_converts_trackers_from_metrics_store(self):
+        """Test successfully converts trackers from metrics store."""
+        mock_client = MagicMock()
+        api = MetricsApi(client=mock_client)
+
+        mock_metrics_store = MagicMock()
+        mock_metrics_store.trackers = {
+            "loss": V1MetricsTracker(
+                name="loss",
+                num_rows=100,
+                min_value=0.1,
+                max_value=1.0,
+            ),
+            "accuracy": V1MetricsTracker(
+                name="accuracy",
+                num_rows=50,
+                min_value=0.8,
+                max_value=0.99,
+            ),
+        }
+
+        result = api.get_trackers_from_metrics_store(mock_metrics_store)
+
+        assert result is not None
+        assert len(result) == 2
+        assert "loss" in result
+        assert "accuracy" in result
+        assert isinstance(result["loss"], MetricsTracker)
+        assert result["loss"].num_rows == 100
+        assert result["accuracy"].num_rows == 50