feat(compressors): add JSONL training metrics logging with CLI controls

- Add write_training_metrics() in new compressors/training_metrics.py
  for appending epoch/step/lr/component rows as JSON Lines
- Wire --metrics-path and --log-every CLI options into train.py, passing
  them to the training loop so metrics rows are written every N steps
- Accept absolute metrics paths or paths relative to output directory
- Add quantization component to loss log alongside existing distill/contrastive
- Replace inline torch.device() with get_device() utility
- Add test_hash_training_metrics.py covering multi-row JSONL append

Infrastructure:
- Pin torch 2.7.1 + CUDA 12.8 index for Linux/Windows in pyproject.toml
- Add .justfile rsync upload recipe with .stignore exclusion
- Exclude **/__marimo__ from rsync in .stignore

Dependencies updated: numpy 2.4.5, pandas 3.0.3, black 26.5.0,
click 8.4.0, contourpy, etc.

.justfile

@@ -4,6 +4,12 @@ export MSYS2_ARG_CONV_EXCL := "*"
 export MSYS2_ENV_CONV_EXCL := "*"
 remote_ssh_target := env("REMOTE_SSH_TARGET")
 remote_docker_container := env("REMOTE_DOCKER_CONTAINER")
+remote_root := "$REMOTE_SSH_TARGET:$REMOTE_WORKDIR"
+rsync_flags := "-avLh --progress --stats --itemize-changes"
+upload_excludes := "--exclude-from=.stignore"
+
+upload:
+    rsync {{ rsync_flags }} {{ upload_excludes }} . {{ remote_root }}/; \
 
 sync-pkgs:
     uv sync --inexact

.stignore

@@ -14,3 +14,4 @@ outputs
 .sisyphus
 **/__pycache__
 **/sim_build
+**/__marimo__

mini-nav/commands/train.py

@@ -11,6 +11,17 @@ def train(
     checkpoint_path: str = typer.Option(
         "hash_checkpoint.pt", "--checkpoint", "-c", help="Checkpoint path"
     ),
+    metrics_path: str = typer.Option(
+        "hash_training_metrics.jsonl",
+        "--metrics-path",
+        "-m",
+        help="JSONL metrics path for loss curves; relative to output directory unless absolute",
+    ),
+    log_every: int = typer.Option(
+        1,
+        "--log-every",
+        help="Write one metrics row every N global steps; <=0 disables metrics logging",
+    ),
 ):
     from compressors import train as train_module
 
@@ -19,4 +30,6 @@ def train(
         batch_size=batch_size,
         lr=lr,
         checkpoint_path=checkpoint_path,
+        metrics_path=metrics_path,
+        log_every=log_every,
     )

mini-nav/compressors/train.py

@@ -5,13 +5,14 @@ import os
 import torch
 import torch.nn.functional as F
 from compressors import HashCompressor, HashLoss
+from compressors.training_metrics import write_training_metrics
 from configs import cfg_manager
+from datasets import load_dataset
+from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 from torch import nn
 from torch.utils.data import DataLoader
-from rich.progress import Progress, BarColumn, TextColumn, TimeRemainingColumn
 from transformers import AutoImageProcessor, AutoModel
-
+from utils import get_device
-from datasets import load_dataset
 
 
 def save_checkpoint(model: nn.Module, optimizer, epoch, step, path="checkpoint.pt"):
@@ -48,6 +49,8 @@ def train(
     batch_size: int = 64,
     lr: float = 1e-4,
     checkpoint_path: str = "hash_checkpoint.pt",
+    metrics_path: str = "hash_training_metrics.jsonl",
+    log_every: int = 1,
 ):
     """Train hash compressor with batch-level retrieval loss.
 
@@ -56,9 +59,11 @@ def train(
         batch_size: Batch size for training
         lr: Learning rate
         checkpoint_path: Path to save/load checkpoints
+        metrics_path: JSONL metrics path, relative to output directory unless absolute
+        log_every: Write metrics every N global steps; values <= 0 disable metrics logging
     """
     # Auto detect device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = get_device()
 
     # Global variables
     save_every = 500
@@ -98,6 +103,10 @@ def train(
 
     # Auto load checkpoint
     output_dir = cfg_manager.get().output.directory
+    metrics_file = output_dir / metrics_path
+    if os.path.isabs(metrics_path):
+        metrics_file = metrics_path
+
     if os.path.exists(output_dir / checkpoint_path):
         start_epoch, global_step = load_checkpoint(
             compressor, optimizer, output_dir / checkpoint_path
@@ -160,9 +169,19 @@ def train(
                         description=f"Epoch [{epoch + 1}/{epoch_size}] "
                         f"loss={components['total']:.4f} "
                         f"cont={components['contrastive']:.2f} "
-                        f"distill={components['distill']:.3f}",
+                        f"distill={components['distill']:.3f} "
+                        f"quant={components['quantization']:.3f}",
                     )
 
+                    if log_every > 0 and global_step % log_every == 0:
+                        write_training_metrics(
+                            metrics_file,
+                            epoch=epoch + 1,
+                            step=global_step,
+                            lr=optimizer.param_groups[0]["lr"],
+                            components=components,
+                        )
+
                     # ---- periodic save ----
                     if global_step % save_every == 0:
                         save_checkpoint(

mini-nav/compressors/training_metrics.py

@@ -0,0 +1,37 @@
+"""Utilities for recording hash compressor training metrics."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Mapping
+
+
+def write_training_metrics(
+    path: str | Path,
+    *,
+    epoch: int,
+    step: int,
+    lr: float,
+    components: Mapping[str, float],
+) -> None:
+    """Append one training metrics row as JSON Lines.
+
+    The JSONL format keeps training logging cheap and easy to resume: every
+    training step is an independent row that plotting scripts can stream later.
+    """
+    metrics_path = Path(path)
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+
+    row = {
+        "epoch": int(epoch),
+        "step": int(step),
+        "lr": float(lr),
+        "total": float(components["total"]),
+        "contrastive": float(components["contrastive"]),
+        "distill": float(components["distill"]),
+        "quantization": float(components["quantization"]),
+    }
+
+    with metrics_path.open("a", encoding="utf-8") as f:
+        f.write(json.dumps(row, ensure_ascii=False) + "\n")

pyproject.toml

@@ -22,6 +22,9 @@ dependencies = [
     "scikit-learn>=1.7.2",
     "transformers>=5.0.0",
     "typer>=0.24.1",
+    "torch==2.7.1",
+    "torchvision==0.22.1",
+    "torchaudio==2.7.1",
 ]
 
 [dependency-groups]
@@ -38,6 +41,23 @@ dev = [
     "cocotb-tools>=0.1.0",
 ]
 
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchvision = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchaudio = [
+    { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+
 [tool.ty.environment]
 python = "$UV_PROJECT_ENVIRONMENT"
 root = ["./mini-nav", "./notebooks", "./hw/sim", "./scripts"]

tests/test_hash_training_metrics.py

@@ -0,0 +1,61 @@
+import json
+import sys
+from pathlib import Path
+
+
+COMPRESSORS_DIR = Path(__file__).resolve().parents[1] / "mini-nav" / "compressors"
+sys.path.insert(0, str(COMPRESSORS_DIR))
+
+from training_metrics import write_training_metrics  # noqa: E402
+
+
+def test_write_training_metrics_appends_jsonl_rows(tmp_path):
+    metrics_path = tmp_path / "hash_training_metrics.jsonl"
+
+    write_training_metrics(
+        metrics_path,
+        epoch=2,
+        step=17,
+        lr=1e-4,
+        components={
+            "total": 1.25,
+            "contrastive": 0.75,
+            "distill": 0.4,
+            "quantization": 0.1,
+        },
+    )
+    write_training_metrics(
+        metrics_path,
+        epoch=2,
+        step=18,
+        lr=1e-4,
+        components={
+            "total": 1.0,
+            "contrastive": 0.6,
+            "distill": 0.3,
+            "quantization": 0.1,
+        },
+    )
+
+    rows = [json.loads(line) for line in metrics_path.read_text().splitlines()]
+
+    assert rows == [
+        {
+            "epoch": 2,
+            "step": 17,
+            "lr": 1e-4,
+            "total": 1.25,
+            "contrastive": 0.75,
+            "distill": 0.4,
+            "quantization": 0.1,
+        },
+        {
+            "epoch": 2,
+            "step": 18,
+            "lr": 1e-4,
+            "total": 1.0,
+            "contrastive": 0.6,
+            "distill": 0.3,
+            "quantization": 0.1,
+        },
+    ]