Orchestrate the full pipeline with b2luigi (cache-warm → train → rollout → analysis) #83
Reference in New Issue
Block a user
Delete Branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
b2luigi orchestration for GIANT
Context
Every multi-step GIANT experiment is chained by hand today:
dwarf warm-cache→dwarf build-geometry-oracle→giant train→giant rollout→giant analyze prep/submit→ wait →giant analyze render. Nothing tracks what has already been produced, nothing waits for anything, and the HTCondor layer is a hand-rolled submit-file generator (giant/analysis/condor.py:write_submit,jobs.txt,analyze.sub) that is fire-and-forget —analyze submitreturns immediately and the user has to pollcondor_qthemselves and remember to runrenderafterwards. A second, parallel copy of the same idea for GPU jobs is half-finished on thecondor-gpu-train-rolloutbranch (giant/condor.py,train-submit/rollout-submit).This replaces all of it with b2luigi: a luigi wrapper that already does dependency resolution, target-based idempotency ("done" = output file exists), HTCondor submission with status polling, per-task resource settings, logging, and a central scheduler UI. After this change b2luigi is the only sanctioned way to run a multi-step pipeline; the
giant/dwarfCLIs are reduced to single-step primitives that the tasks invoke, and the bespoke condor code is deleted.Design
New package
giant/workflow/spec.pyWorkflowSpec,TrainSpec,RolloutSpec,AnalysisSpec,CondorSpec), name-uniqueness + cross-reference validation, andspec_hash()per sub-spechtcondor.pyhtcondor_settingsdicts (docker universe,+RemoteJob,+RequestWalltime, GPU requirement expressions)tasks.pyTaskclasses belowrun.py--spec, settings wiring,b2luigi.process(WorkflowTask(...), workers=N, batch=...)giant workflow run <spec.toml> [--batch] [--workers N] [--mode dry-run|show-output|remove] [--scheduler-host/--scheduler-port](new sub-app ingiant/cli.py) is a thin exec ofrun.pyso there is one documented entry point;python giant/workflow/run.py …stays valid because b2luigi's executable wrapper needs a real script path, notpython -m.Workflow spec (one file, the only place a pipeline is parameterised)
Every task carries two luigi parameters: its
nameand a shortspec_hash(8 hex of the canonicalised resolved sub-spec, including transitive parents). The hash is what makes an edited spec produce a fresh result directory instead of silently reusing stale outputs. Setuse_parameter_name_in_output = Trueso paths read…/name=baseline/spec_hash=1a2b3c4d/best.pt.Task graph
DatasetTask—b2luigi.ExternalTaskover the steps parquet path; fails fast with a clear message if/cephis not mounted.WarmCacheTask— callsgiant.tools.warm_setup_cache.run_warm_setup_cachein-process. Its real product (<data>.giant_train_cache.json) lives next to the dataset, not underresult_dir, so the b2luigi target is a small stamp JSON recording the sidecar path + its mtime/size; CPU condor job (high memory, no GPU).GeometryOracleTask— calls thedwarf build-geometry-oracleimplementation (giant/tools/geometry_oracle.py); outputoracle.pklunderresult_dir. CPU job.TrainEpochTask(name, epoch)— one short GPU condor job per epoch, chained: epoch k requires epoch k−1 (epoch 1 requiresWarmCacheTask). Each job callsrun_train_jobwithout_dir= its own output dir,epochs = k, andresume = <epoch k−1 dir>/last.pt. This needs no change to the training loop:giant/training/loop.py:160-164already setsstart_epoch = ckpt["epoch"] + 1and returns cleanly if the checkpoint already covers--epochs, so--epochs k --resume <k−1>/last.ptruns exactly epoch k.--outalready wins overresume.parent(giant/cli.py:716-732), so the per-epoch output dirs work as-is.last.pt.best.ptis written by the loop only when that epoch improved, andbest_val_losstravels inside the checkpoint, so the global best comparison stays correct across jobs — "best.ptexists in epoch dir k" means exactly "epoch k was the best so far".WarmCacheTaskguarantees a cache hit — pass--cache-setup) and one queue wait per epoch;epochs_per_jobin the spec (default 1) trades those back if the queue turns out to dominate.--seedalone, in which case every epoch job would replay the same batch order. If so, derive the loader's shuffle seed fromseed + epochwhile leaving the val-split seed untouched (the split must stay identical across jobs).TrainTask(name)— cheap local task requiring the finalTrainEpochTask. Picks the highest-numbered epoch dir containing abest.pt, and publishes the run's canonical outputs into one directory:best.pt,last.pt,config.toml, and ametrics.csvconcatenated from the per-epoch ones. Everything downstream (RolloutTask, humans,analyze metrics) points here and never has to know about the epoch fan-out.RolloutTask(name)— GPU condor job; requiresTrainTask+GeometryOracleTask. Targets arerollout.parquetandrollout.yaml.AnalysisPrepTask(name)— local (cheap, streaming); requires everyRolloutTaskit names. Callsgiant.analysis.prepwithrun_dir= its own output dir; targetsshared.json+run_meta.json.AnalysisComputeTask(name, plot_id, chunk)— one CPU condor job per (plot, chunk), replacingjobs.txt/analyze.subentirely. The job set is enumerable ahead of time fromcatalog_ids()×chunks, collapsing to one chunk for the fivechunkable=Falsespecs (same rule as today's_job_walltimes,giant/analysis/condor.py:487).output()is an explicitLocalTargeton<prep_dir>/reduced_partial/<id>__<chunk>.jsonrather thanadd_to_output, socompute-one's existing on-disk contract is untouched andmerge_onekeeps working.htcondor_settingsis a property, evaluated at submit time — i.e. after prep has run — so it can still readrun_meta.jsonand set+RequestWalltimefromgiant/analysis/runtime_estimate.py:estimate_runtime_s.AnalysisRenderTask(name)— always local (the only step importing plotstyle/LaTeX). Runsmerge_allthenrender_run, plusgallery generatewhengallery = true. Target:<prep_dir>/plots/metadata.yaml.WorkflowTask—b2luigi.WrapperTaskrequiring oneAnalysisRenderTaskper[[analysis]].Settings wiring (
run.py)result_dir,log_dir,task_file_dirfrom the spec;batch_system = "htcondor";working_dir = repo_dir;env_script;executable = [".venv/bin/python"]./cephis shared between submit host and workers, so notransfer_files— result and log dirs must be on/ceph.AnalysisRenderTaskoverridesbatch_system = "local"as a class property, which b2luigi honours per task.GPU requirement strings are ported from
condor-gpu-train-rollout:giant/condor.py::_gpu_requirements(TARGET.ProvidesEtpCeph =?= TrueANDed withGPUs_DeviceName/GPUs_GlobalMemoryMbpins) rather than rewritten.Deletions and CLI reduction
giant/analysis/condor.py: deleteSubmitConfig,_WRAPPER,_submit_description,_job_walltimes,_resolve_giant_executable,write_submit. Keepprep,derive_run_dir,RunMeta,load_rollout_yaml(s),compute_reduced/compute_one,merge_one/merge_all— that is the real logic — and rename the module togiant/analysis/run.pysince nothing in it submits any more. Update the re-exports ingiant/analysis/__init__.pyand its module docstring.giant/cli.py: delete theanalyze submitcommand (cli.py:1680-1756).prep,compute-one,merge-one,list,render,metricsstay as primitives.train-submit/rollout-submitfromcondor-gpu-train-rollout— the workflow supersedes them, and that branch'sgiant/condor.pyis reduced to the requirement-string helpers moved intogiant/workflow/htcondor.py. This is a decision that branch's eventual merge must respect; note it in CLAUDE.md.tests/test_condor.py: drop thewrite_submit/submit-description cases, keep theprep/mergeones.Required change to
giant rollout_write_prediction_ref(giant/cli.py:188) writes the sidecar to<checkpoint.parent>/<random-uuid>.yaml, which is not a deterministic target. Change it so that when--outis passed explicitly, the sidecar goes toout.with_suffix(".yaml"); the existing uuid-under-the-checkpoint behaviour is kept for the no---outcase so ad-hoc runs and the/cephpredictions convention are unaffected. Apply the same rule togiant predictfor consistency.Dependency
Add
b2luigi>=1.0,<2under a newworkflowoptional-dependency extra inpyproject.toml(it pullsluigi+tenacity), and includegiant[workflow]in thedevextra. Documentuv sync --extra cpu --extra workflowin CLAUDE.md.Files
giant/workflow/{__init__,spec,htcondor,tasks,run}.py,tests/test_workflow_spec.py,tests/test_workflow_tasks.py, an exampleconfigs/workflow_example.tomlgiant/cli.py(deleteanalyze submit, addworkflowsub-app, sidecar path rule),giant/analysis/condor.py→giant/analysis/run.py,giant/analysis/__init__.py,tests/test_condor.py,pyproject.toml,CLAUDE.md,README.mdVerification
uv sync --extra cpu --extra workflow --extra dev, thenuv run pytest,uv run ruff check .,uv run ty check ..giant workflow run configs/workflow_example.toml --mode dry-runon a local machine — asserts the whole DAG resolves and prints the tasks that would run; exit code 1 means work pending, 0 means everything already done.--mode show-outputon the same spec — eyeball that every target path is where the plan says it is (green = exists, red = missing).chunks = 1,batch_system = "local"): confirm three epoch dirs each with alast.pt, the publishedbest.pt/concatenatedmetrics.csv,rollout.parquet/rollout.yaml,reduced_partial/*.json, andplots/metadata.yaml, and that a second invocation is a no-op.giant traininvocation with a fixed seed and diff itsmetrics.csvagainst the workflow's concatenated one — the per-epoch losses should match (modulo the shuffle-seed question above).reduced_partial/<id>__<chunk>.jsonand confirm exactly that one compute task re-runs, then render.giant workflow run <real spec> --batch --workers 20against the real dataset, optionally withluigidrunning for the progress UI; verify withcondor_q -batch <job_name>that the (plot, chunk) job count matcheslen(catalog_ids()) × chunks(minus the chunk collapse for the five non-chunkable specs), and that render only fires after the last compute job succeeds.