From ad834001ba06fd70dc9584060f1e6e808885fcf4 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Fri, 10 Jul 2026 23:39:44 +0000 Subject: [PATCH] =?UTF-8?q?fix(run):=20exit=20after=20profile.async=5Ftest?= =?UTF-8?q?=20=E2=80=94=20donation=20has=20consumed=20state?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The async_test block calls the donating step_fn (donate=all-except-first) on `state` without rebinding it, then fell through to the mem_profile path and the main loop holding deleted buffers — guaranteed "Array has been deleted" whenever async_test was combined with a normal run. Mirror the mem_profile guard: os._exit(0) at the end of the block, plus a ProfileConfig validator refusing async_test+mem_profile together (each exits before the other could run). Audit finding M2 (2026-07-10 compile audit). Crew-Address: task/fix-async-test-donated-state Co-Authored-By: Claude Fable 5 --- param_decomp/configs.py | 13 +++++++++++-- param_decomp/run.py | 1 + param_decomp/tests/test_config.py | 10 ++++++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/param_decomp/configs.py b/param_decomp/configs.py index 455375223..826bbe9b9 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -625,8 +625,9 @@ class ProfileConfig(BaseConfig): `mem_profile` (static + runtime memory analysis, then exits), `time_steps` (per-step wall breakdown), `trace` (perfetto trace over `trace_start`..`trace_start+trace_steps`), - `profile_max_events` (raise the perfetto GPU-activity event cap), `async_test`, - `leaf_bench`, `no_checkpoint` (skip ALL saves — throwaway profiling only). + `profile_max_events` (raise the perfetto GPU-activity event cap), `async_test` + (dispatch-asynchrony probe, then exits), `leaf_bench`, `no_checkpoint` (skip ALL + saves — throwaway profiling only). """ mem_profile: bool = False @@ -645,6 +646,14 @@ class ProfileConfig(BaseConfig): leaf_bench: bool = False no_checkpoint: bool = False + @model_validator(mode="after") + def validate_exclusive_exiting_hooks(self) -> Self: + assert not (self.async_test and self.mem_profile), ( + "async_test and mem_profile each consume `state` via step_fn donation and exit " + "before the loop — run them one at a time" + ) + return self + class LaunchEnv(BaseConfig): """The process-environment surface a SLURM-launched rank runs with — the XLA *client* diff --git a/param_decomp/run.py b/param_decomp/run.py index 17ebca0eb..43c0f78dd 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -599,6 +599,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: f"sync/host-bound => each call big)", flush=True, ) + os._exit(0) # profiling-only path; donation has consumed `state`, so don't enter the loop if profile.mem_profile: _gib = 1024**3 diff --git a/param_decomp/tests/test_config.py b/param_decomp/tests/test_config.py index 38a0519dd..3d396d0a2 100644 --- a/param_decomp/tests/test_config.py +++ b/param_decomp/tests/test_config.py @@ -16,6 +16,7 @@ ImportanceMinimalityLossConfig, PDConfig, PersistentPGDReconLossConfig, + ProfileConfig, ) from param_decomp.recon import ( build_loss_terms, @@ -89,6 +90,15 @@ def test_legacy_top_level_n_mask_samples_pushes_onto_stochastic_terms(): assert by_type["StochasticReconSubsetLoss"].n_mask_samples == 7 # pyright: ignore[reportAttributeAccessIssue] +def test_profile_exiting_hooks_are_mutually_exclusive(): + # Both hooks consume `state` via step_fn donation and os._exit(0) before the loop, + # so a config combining them would silently skip whichever runs second. + ProfileConfig(async_test=True) + ProfileConfig(mem_profile=True) + with pytest.raises(ValidationError, match="one at a time"): + ProfileConfig(async_test=True, mem_profile=True) + + def test_b128_config_converts(): converted, raw = load_config(CONFIGS / "llama8b_l18_b128_cmp32.yaml", RUN_ID) assert raw["pd"]["batch_size"] == 128