feat(mlflow): auto log config leaves as mlflow parameters #1

Merged
fnoble merged 12 commits from v2.4.0-PR-train-params into main 2026-09-07 15:21:25 +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.
collect_scalars recurses dataclass and pydantic BaseModel instances
field by field, treats Mapping[str, scalar] roots as explicit leaves,
keeps str/int/float/bool, and stringifies other objects (StrategyConfig's
live model/tracker/stream/device fields) as their class name so run start
never raises on a config that holds them. None stringifies as "None".
Two bare leaves with differing values raise KeyError naming the colliding
paths.

log_params_guarded logs to the active MLflow run, truncating keys and
values to the MLflow caps (250/6000, mlflow 3.x) instead of raising.
Identical values are a no-op and changed values warn while keeping the
original, because MLflow params are immutable per run and the train
pipeline reuses job.run_id on resume (AD9). mlflow imports lazily inside
the function.

Tested against synthetic dataclasses, pydantic models, live-object fields,
mapping roots, and a mocked MLflow client; no server. Nothing imports the
new module yet.
SUFFIX_RULES owns every non-bare key (AD10): a new ambiguous leaf or a
key rename is one table row, with no caller edits. resolve_keys emits each
leaf's bare key plus "<leaf>-<suffix>" for each tabled suffix appearing as
a segment of the leaf's origin, so the key set is a pure function of the
leaves and the table and is identical across runs. A tabled leaf with
differing values keeps the first origin's declaration-order value on the
bare key and carries each origin's value on its suffixed key; unknown
bare-key collisions keep surfacing as the walker's KeyError, so
collect_scalars now exempts tabled names from its collision raise.

Tests cover bare keys, bare-plus-suffix duplication, the differing-value
collision raise, suffixes absent from an origin, and cross-run identity.
log_params(train_dataset, test_dataset, scheduler, config) and
_get_scheduler_params have no callers anywhere in src/ or tests/; they
made param logging look handled when it was not. Deleting first clears
the name for the new walker-based scheme (plan v2.3.1-train-params,
commit 1).

Grep evidence (grep -rn "log_params\|_get_scheduler_params"
--include="*.py" src tests, excluding log_params_guarded hits):

- src/training/tracking/__init__.py:6,18 - re-export and __all__ entry
  (removed in this commit)
- src/training/tracking/log_params.py - definition, deleted
- src/training/tracking/_get_scheduler_params.py - definition, deleted
- src/training/tracking/params.py:152 - mlflow.log_params call in
  log_params_guarded (new walker module, unrelated name)
- src/training/pipeline/train.py:122, eval.py:154 - mlflow.log_params
  dict calls, unrelated method name
- tests/test_param_walker.py:129 - stub attribute on a mocked client,
  unrelated

No caller of the dead helpers exists; no imports break.
compose_train_params(exp, runtime) assembles the walker roots per the
plan's binding composition: experiment/train/val explicit mappings (the
four legacy keys fold in as experiment_name, model_class, examples-train,
examples-val per AD7), runtime and strategy.config as instance roots
(live objects stringify as class names without raising), and the model's
self.config as a root only when it is a dataclass or pydantic instance.
graph_recurrent keeps config as a class attribute, so it rides model_class
only with no special case (AD8). The four-key mlflow.log_params dict in
the start_run block is replaced by log_params_guarded(resolve_keys(
collect_scalars(*compose_train_params(exp, runtime)))), which also guards
resume against MLflow param immutability (AD9).

One deviation from the plan's compose_train_params(exp) signature: runtime
is passed explicitly, because RunContext lives in run_train_pipeline and
is not an Experiment attribute; the composition is otherwise exactly as
specified.

tests/test_train_params.py stubs the experiment and asserts the emitted
key set, including the ambiguous-leaf duplication (embed_dim-embedder/
embed_dim-forecaster, batch_size-train/batch_size-val,
examples-train/examples-val), the live-object strategy config and the
class-attribute model config surviving without raise, and the legacy four
keys in the new scheme.

Static trigger check (walker-raise / model-config): the walker cannot
raise on the real shapes - StrategyConfig's live fields (model, loss_fn,
tracker, stream, device) stringify as class names, runtime.device and
strategy.device are both torch.device so their stringified values are
equal and non-colliding, scheduler spec tuples stringify under a single
origin each, and the only model-config leaf named in more than one
component is embed_dim, which SUFFIX_RULES tables. The runtime param
smoke (one single-epoch train of graph_embed_class, retrieval_forecast,
transformer_class) remains the acceptance gate and needs an MLflow-backed
environment.
Adversarial review of this PR found the runtime scalars recording no
information: the walker stringifies live objects as type(obj).__name__, so
torch.device logged as "device", torch.dtype as "dtype" (its __name__),
and WarmupCosineSpec.milestones as "tuple". bf16 vs fp32 runs and cuda vs
cpu runs were indistinguishable in the param table, contradicting the
ratified feature list. The stub tests hid this by faking classes named
"device"/"dtype".

- params.py: _visit records str() for torch.device/torch.dtype and repr()
  for containers of scalars; other live objects keep the class-name
  fallback. log_params_guarded skips empty values with a warning: some
  MLflow backends reject them ("Param value cannot be empty"), and a
  rejected log_params would raise out of run start — the one failure mode
  the guard does not cover. compile_mode defaults to "".
- train.py: the experiment root now logs eval_interval and
  checkpoint_interval (named in the proposal's feature list, previously
  unlogged). epochs deliberately stays on the scheduler spec, which
  derives it from the same constant; a second epochs source would be a
  differing-value collision the walker raises on.
- tests: real torch.device/torch.bfloat16 in the RunContext stub with
  value assertions ("cpu", "torch.bfloat16"), milestones pinned to
  "(6,)", the new experiment keys in the emitted key set, walker tests
  for the device/dtype/container branches, and an empty-value guard test.
- docs: training.md no longer lists the deleted log_params.py (it now
  names params.py/param_keys.py); the release README's 2.2 row links the
  renamed implementation plan (v2.2.md -> v2.5.0.md).

Known limitation, deferred: the renumbering commit left the release
README's 2.2/2.5.0 rows inconsistent (2.5.0.md's heading still says
"Plan 2.2"; the 2.5.0 row links the CLI-launch implementation plan).
That is renumbering intent, not a link fix, and needs the owner's call.

Tests not executed in this environment (no torch/mlflow); py_compile plus
a stubbed-torch smoke of the walker and guard pass. The venv pytest and
the plan's param smoke remain acceptance gates.
Replaces the blanket str() on the torch branch with the admission rule we
settled on: a type joins the branch only together with its known string
form and the transform for it. torch.dtype is always "torch.<name>", so
it records the bare name ("bfloat16"); torch.device has no module prefix
("cpu", "cuda:0") and passes through unchanged. Anything not in the
branch keeps the class-name fallback, loudly and manually, until someone
admits it with its transform.

The value logged is the resolved RunContext.dtype, not the --dtype flag:
the object is what the model parameters and retriever buffers are built
from, and build_run_context is where dtype decisions live (it already
warns on fp16), so it is the value that stays true if a fallback is ever
added there.

Tests updated: dtype pins to "bfloat16", device passthrough ("cuda:0")
and the collision guard covered in test_param_walker; test_train_params
asserts the shortened form.

The corresponding manual dtype logging experiment on the v2.4.1 branch
was reverted; this walker is now the only dtype mechanism. py_compile and
a stubbed-torch smoke pass; pytest on a torch box remains the gate.
fnoble merged commit edc07ee5ac into main 2026-09-07 15:21:25 +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!1
No description provided.