Skip to content

V0.2_FITUP.md

Latest commit

 

History

History
227 lines (181 loc) · 10.3 KB

V0.2_FITUP.md

File metadata and controls

227 lines (181 loc) · 10.3 KB
Raw

MIME v0.2 fit-up — task list

Bringing MIME into full compatibility with, and adoption of, MADDENING v0.2.0, including all existing experiments. This document is a handoff for a dedicated agent.

Context & ground rules

  • MADDENING v0.2.0 is released (tag v0.2.0). New surfaces: the static_data channel + StaticArray wrapper; halo_width() (replacing the deprecated requires_halo); compile-time edge-validation warnings (ShapeMismatchWarning / DtypeMismatchWarning / UnitMismatchWarning); BinaryStateEncoder field-subscriptions + zstd; AWS/GCP cloud providers; preempt/resume snapshots; profiler endpoints (/sim/profile); checkpoint schema versioning; GraphManager.validate_sharding(); MigrationError.
  • Work from MIME main (HEAD 1432ac5 when this was written). The repo may be checked out on feat/fvm-float64-audit — that branch predates the v0.2 integration; do not base work on it. Another agent may be active in this repo — check before switching branches; use a git worktree if in doubt.
  • main already has: the maddening>=0.2.0,<0.3 pin, the requires_halo→halo_width node migration, and a green suite (626 passed against MADDENING v0.2).
  • Test: cd MIME && source ../.venv/bin/activate && python -m pytest -q (shared venv /home/nick/MSF/msf/.venv; -m 'not slow' is the default lane).
  • Scope boundary: this is compatibility + v0.2-surface adoption + MIME's own v0.2 release. MIME node physics improvements and new node types are a SEPARATE later cycle — not in this list.

Suggested order: §1 → §2 (clean baseline) → §3 (the big one) → §4 → §5 → §6 → §7.

1 — Finish the halo-migration cleanup · P1 · small

The node migration is done; residue remains:

  • Migrate the two test-stub classes _FieldStub (tests/verification/test_coupling_group_field_body.py:64) and _BodyStub (:102) from requires_halo to halo_width(). They currently emit a FutureWarning on subclassing (MADDENING v0.2's __init_subclass__); v0.3 turns that into a MigrationError.
  • Rewrite tests that assert the deprecated property to assert halo_width() instead: tests/nodes/lbm/test_iblbm_fluid_node.py TestRequiresHalo (:236-238) → halo_width() == {0:1,1:1,2:1}; tests/nodes/test_external_magnetic_field.py:109 and tests/nodes/test_phase_tracking.py:127 → halo_width() == {}.
  • Update stale docstrings/comments mentioning requires_halo: src/mime/core/node.py:7,40, src/mime/nodes/environment/fvm/fluid_node.py:287, src/mime/nodes/environment/lbm/fluid_node.py:99.

Done when: grep -rn requires_halo src/ tests/ is empty (or only intentional compat references); suite green; subclassing a MIME node emits no FutureWarning.

2 — Verify against released MADDENING, not a branch · P1 · small

  • .github/workflows/ci.yml checks out MADDENING with no ref: — it builds against MADDENING main/HEAD, not a released tag. Pin it to ref: v0.2.0 (or, once MADDENING is on PyPI, install maddening from PyPI and drop the source checkout).
  • Reconcile the CI jax==0.5.1 / jaxlib==0.5.1 pin against pyproject.toml's jax>=0.4 — pick one consistent constraint.
  • Run the full suite (-m 'slow or not slow') against the v0.2.0 tag; confirm 0 failures and no escalating warnings.

Done when: CI is green pinned to MADDENING v0.2.0.

3 — Adopt the static_data / StaticArray channel · P2 · LARGE

The main engineering item. MIME nodes hold large non-evolving arrays as plain instance attributes closed over in update() (re-baked into the JIT'd HLO every step) or — worse — thread them through state. v0.2's static_data channel is purpose-built for exactly these (meshes, wall masks, lookup tables, basis functions, NN weights): wrap each in StaticArray and expose it via the node's static_data property. Contract: static_data is not checkpointed — keep the provider/config in self.params so it rebuilds on resume.

Candidates:

  • FVM — the FVMMesh pytree (owner / neighbour / Sf / n / area / d / w / V / x / …). src/mime/nodes/environment/fvm/mesh.py:100-112 literally calls it a "closed-over pytree … potentially a 1 GB FVM mesh". make_piso_step() is rebuilt every update() (fvm/fluid_node.py:367), re-baking the mesh into HLO each step.
  • LBM — _pipe_wall / _pipe_missing masks (lbm/fluid_node.py:181,184). Also: stop threading the static pipe-wall through state — initial_state() puts solid_mask in state (:226) and update() returns a fresh one every step (:323) though only the UMR portion moves. Split static vs dynamic.
  • Stokeslet — _R, _lu, _piv (BEM LU factors, stokeslet/fluid_node.py:265-307), _body_mesh / _wall_mesh, and WallTable.data (the (n_rho,n_rho,n_dphi,n_dz,3,3) Green's-function tensor, stokeslet/cylinder_wall_table.py:54).
  • MLPResistanceNode — MLP weights + normalization stats (stokeslet/mlp_resistance_node.py:217-225).
  • GNN flux corrector — EdgeMLPParams weight tuples (fvm/gnn.py:116,238).
  • DefectCorrection — pipe masks (defect_correction/fluid_node.py:164-171).

Done when: these arrays no longer re-bake into HLO per step; static_data_hash() is stable across steps and changes when geometry changes; suite green; a before/after recompile or step-time measurement shows the win.

4 — Confirm edge-validation cleanliness · P2 · small

MIME builds unit-tagged EdgeSpecs (lbm/fluid_node.py:445-484, stokeslet/fluid_node.py:486-520, stokeslet/mlp_resistance_node.py), so v0.2's compile-time edge validation now inspects them automatically.

  • Confirm no MIME graph emits a spurious ShapeMismatchWarning / DtypeMismatchWarning / UnitMismatchWarning at compile(). (A prior smoke check found zero — verify it still holds.)
  • Add a test that compiles the representative experiment graphs with warnings recorded and asserts no EdgeValidationWarning.

Done when: a test pins MIME's graphs as edge-validation-clean.

5 — Experiments & verification audit · P2 · medium

Every existing experiment must run correctly on v0.2.

  • Run each experiment end-to-end against v0.2.0: experiments/ar4_helical_drive/, experiments/dejongh_confined/, experiments/umr_confinement/, plus the GraphManager-based src/mime/experiments/dejongh.py / dejongh_new_chain.py.
  • Re-run the MIME-VER-* verification benchmarks (tests/ verification/, ~20 files incl. the FVM / LBM / stokeslet suites) and confirm the reports under docs/validation/benchmark_reports/ still hold.
  • The long de Jongh / GNN cloud sweeps (scripts/run_confinement_sweep.py, sky_tasks/gnn_sweep.yaml, jobs/production_h100.yaml, jobs/rehearsal_a100.yaml) are multi-hour — wire them to v0.2's preempt/resume snapshots + checkpoint schema versioning so a reclaimed spot VM resumes instead of restarting from zero.

Done when: every experiment + verification benchmark is confirmed green on v0.2.0; the cloud sweeps survive a simulated preemption.

6 — MIME's own v0.2 release · P3 · medium

MIME is still version = "0.1.0", tagged only v0.1.0. Mirror what MADDENING did for its v0.2.0 release:

  • Bump pyproject.toml version 0.1.0 → 0.2.0; review classifiers (still Development Status :: 2 - Pre-Alpha), keywords, description.
  • CHANGELOG.md currently has only an [Unreleased] section — write a real [0.2.0] entry (and ideally backfill [0.1.0]).
  • Add release notes under docs/ (there is no docs/release_notes/) — a v0.2 "what's new" page plus a migration note.
  • docs/_static/switcher.json has a single latest entry — add a v0.2 versioned row, and bump MIME_VERSIONS in MICROROBOTICA/docs/Makefile (currently latest:origin/main) to v0.2:v0.2.0 latest:origin/main.
  • Tag v0.2.0 — publish.yml (PyPI Trusted Publishing) fires on a GitHub Release.
  • After tagging, bump the docs/projects/mime submodule pointer in MICROROBOTICA (still at the pre-v0.2 27c2bbd).

Done when: mime-engine 0.2.0 is tagged, on PyPI, and the docs site shows a MIME version switcher.

7 — Optional v0.2 feature adoption — evaluate · P4

Decide per item whether it earns a place now or is deferred:

  • Profiler — MIME's runner/server.py exposes ZMQ endpoints; evaluate adopting MADDENING's /sim/profile Perfetto export.
  • BinaryStateEncoder — MIME has its own data/hdf5_schema.py (SweepDataWriter) and viz/streaming_observer.py; field subscriptions + zstd could shrink the wire / disk footprint.
  • validate_sharding() — call it if/when MIME shards graphs.
  • Surrogates — MIME has its own mime/surrogates/cholesky_mlp.py (and src/mime/surrogates/__init__.py is empty). Evaluate consuming MADDENING's restructured surrogates/primitives/ — this also intersects MADDENING's deferred "decoder-zoo pull-over from MIME" item.
  • AWS/GCP cloud providers — the cloud job specs (jobs/*.yaml, sky_tasks/) currently target RunPod; the v0.2 providers add AWS/GCP.

Out of scope (separate later cycle)

MIME node physics improvements, accuracy work, and new node types. This list is compatibility + v0.2-surface adoption + the v0.2 release prep only.

Note that these physics improvements will also be shipped with mime-engine v0.2 however the primary goal is to get these things ready first.

DO NOT release mime v0.2 after these things are done. As more things will have to be done after that are out of scope. These include:

  1. Migrating all fluid nodes to follow a shared contract to make them truly swappable.
  2. Updating node specific stuff in runner/server.py to a better design pattern
  3. Improving the robot arm ar4 robot arm simulation experiment.

Structural notes for the agent

  • DefectCorrectionFluidNode, LBMFarFieldNode, and StokesletFluidNode extend MADDENING's SimulationNode directly, bypassing MimeNode — keep that in mind when touching base-class behaviour.
  • MIME's nodes carry maddening compliance metadata (NodeMeta, stability, ValidatedRegime) — preserve it through any refactor.