fix(tracking): mlflow tracking URI and example id hardening #2

Merged
fnoble merged 9 commits from v2.4.0-PR-tracking-logs into main 2026-09-07 15:21:50 +00:00
Owner
No description provided.
--dtype on the train CLI flows through TrainJob into build_run_context and
RunContext.dtype, replacing the pinned fp32. fp16 warns: the training loop
has no GradScaler, so raw fp16 underflows in AdamW.

VectorRetriever keeps its db buffer in the compute dtype so the stage-2 pool
gather (B, N, C, D) halves/quarters; stage-1 scores and the small (B*N, C)
CLaRa score matrices upcast to fp32 for exact top-k and softmax stability.

Per plans/proposals/v2.4.0-expanded-query-search-strat.md-compute-dtype.md.
VectorStoreDatasetConfig gains t_start/t_end, applied to the time column at
load time and bounding the store independently of the example datasets'
windows. The retrieval_forecast config splits its examples filter (Medicine)
from the corpus filter (11 fields), pins the store window to 1920-2018, and
retunes batch sizes and the stage-1 pool size.

Tests for the new bounds are still owed on this branch.

Per plans/proposals/v2.4.0-expanded-query-search-strat.md-corpus-scope.md.
Frees the 2.4 number for the expanded query search strategy: the drafted
Modal tidy-up chain moves to 2.5.0-2.8 (2.9 added), 2.1/2.2 plan docs move
to their archives, and the releases README table follows.

Adds the proposal docs: v2.3.1 mlflow logging fixes, v2.4.0 expanded query
search strategy (candidate strategy, compute dtype, corpus scope), and the
bf16 metric fix, renumbered v2.3.2 -> v2.4.1 per the owner's call recorded
in its Open decisions (standalone, merged after v2.4.0). Statuses of the
compute-dtype and corpus-scope proposals updated to implemented on branch
2.4.
training_step never fed train_ids while every tracker gathers {prefix}_ids
for both prefixes, so _gather_store logged an empty-store ERROR for the
train id store on every run. Both id feeds now share _feed_ids, which skips
a batch whose ids are all NaN, applied to the new train feed and to the
existing val feed (design/tracking-stores.md rule 2): return_id=False
configs emit a scalar-NaN fallback per item, so an unguarded feed would
trade the empty-store ERROR spam for NaN-rejection ERROR spam (AD2). The
guard is per batch (torch.isnan(batch.id).all()), not per dataset
assumption.

Id survey (NaN-mix trigger check; return_id is decided once per dataset
config, so a loader's batches are homogeneous, and the NaN fallback sites
emit scalar NaN per item when the flag is false):

- graph_embed_class: train return_id=False
  (src/config/experiments/graph_embed_class.py:81, CitationGraphDatasetConfig
  'train-dataset') -> train id tensors fully NaN per batch; val
  return_id=True (:91, 'test-dataset') -> fully real. Cross-loader
  asymmetry, never mixed within one batch.
- retrieval_forecast: return_id True for train and val via
  base_dataset_kwargs (src/config/experiments/retrieval_forecast.py:141)
  -> fully real per batch on both paths.
- transformer_class: return_id True for train and val via
  base_dataset_kwargs (src/config/experiments/transformer_class.py:84)
  -> fully real per batch on both paths.
- Fallback sites emit per-item scalar NaN, never a real/NaN row mix:
  src/data/datasets/polars_dataset.py:229, src/data/formaters/graph.py:45,
  src/data/formaters/citation_graph.py:53,
  src/data/datasets/text_token_dataset.py:184.

No train config mixes real and fallback ids inside one batch, so the
NaN-mix trigger does not fire.

NaN-rejection level review: process_values and _process_value keep their
ERROR level. After the guard, a NaN reaching them is a genuine bad value,
not a config fact.

With export enabled, train_ids.npy joins the debug export payload
(release doc friction): consumers of the exported arrays gain one int64
column per example. No metric names or values change.

Tests: tests/test_id_feed.py drives a minimal strategy with a recording
tracker stub; real ids land in train_ids, an all-NaN batch skips the feed
on both paths while non-id feeds continue, and a mixed batch still feeds
(pins the .all() guard granularity).
_gather_store logged an ERROR every time a store came up empty, once per
store per epoch. The message now logs at WARNING the first time a store
name is empty and at DEBUG afterwards, latched by an instance attribute
set[str] on MetricTracker (one line per store name per tracker instance,
surviving clear(); no counters), following the _warned_no_candidates
pattern in src/models/retrieval_forecast.py per AD3, as an instance
attribute instead of the class attribute that pattern uses. The returned
NaN tensor and the gather contract are unchanged.

Verified with a stdlib harness stubbing torch/mlflow/pandas/sklearn
(venv for py3.13 is unavailable in this environment; pytest/pyright not
run): first empty gather warns once, second logs DEBUG, fed stores do
not warn, the latch survives clear(), and a second tracker instance
warns independently. py_compile passes on both files.
Every logger.error(e) and logger.error(str(e)) site in
src/training/tracking/ becomes logger.exception("<context string>"),
so swallowed failures carry a traceback and name their site (AD5:
mechanical, confined to the package). 45 sites across the six tracker
modules; context strings name the plot paths (e.g. "plot rendering
failed: ROC curve in _log_plots") so a plot-path revert for the
plot-spam trigger can find them instantly.

binary_classification_tracker.py calc_metrics PR_AUC handler: the
nested try building mssg = f"{e}..." existed only to pre-format the
message for logger.error(mssg); it collapses to one logger.exception
call, which preserves the traceback the f-string discarded.

Untouched: contextual f-string logger.error sites (shape-mismatch
messages), the NaN-rejection ERROR sites in metric_tracker, and the
_gather_store level logic (owned by v2.3.1-store-latch; hunks
disjoint).

Tests: test_tracker_exception_logging.py - caplog test on poisoned
metric computation asserting the contextual message and exc_info
presence, plus an AST scan asserting no logger.error(<bare name>)
call remains in the package.

Plot-spam trigger: pending - checked at the first real eval run, not
in this commit.
run_eval_pipeline called mlflow.set_tracking_uri inside the
mlflow.start_run block, so the eval run was created against whatever
tracking URI the process had (the default store) instead of
job.env.tracking_uri. Move the call above the start_run block, mirroring
run_train_pipeline (plan v2.3.1-eval-uri; proposal AD4). Two-line
reorder, no API change.

Tests: an AST order check in tests/test_pipeline_jobs.py asserting
mlflow.set_tracking_uri precedes mlflow.start_run in both
run_eval_pipeline and run_train_pipeline, so the correct order is pinned
on both paths.
Two test-authoring bugs found in adversarial review made four latch tests
and the exception-logging behavioral test fail on a correct implementation:

- test_store_latch filtered caplog records by type(tracker).__module__
  (tests.test_store_latch), but _gather_store logs through
  training.tracking.metric_tracker, so every assert saw 0 records. Filter
  by MetricTracker.__module__ instead. test_instances_latch_independently
  also captured both instances' identical records; clear caplog between
  gathers.

- test_tracker_exception_logging poisoned preds with NaN through
  process_values, whose NaN rejection empties the store; calc_metrics then
  died on IndexError at preds.size(0) before reaching the MAE handler.
  Feed valid values and raise inside mean_absolute_error via monkeypatch
  instead, which is exactly what the handler reports. Renamed the test
  accordingly and reworded its docstring.

Also reworded test_partially_nan_batch_still_feeds: a mixed batch still
reaches process_values, which rejects it at ERROR — that is the plan's
NaN-mix descope trigger, discharged by the id survey in the v2.3.1-id-feed
commit, not evidence that mixed batches feed junk.

Tests not executed in this environment (no torch/mlflow); py_compile only.
The venv run remains an acceptance gate.
fnoble merged commit defa431a21 into main 2026-09-07 15:21:50 +00:00
Sign in to join this conversation.
No reviewers
No labels
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference
neurogenesis/forecite!2
No description provided.