diff --git a/docs/operations/joint-crash-recovery.md b/docs/operations/joint-crash-recovery.md new file mode 100644 index 000000000..ab01aaa2c --- /dev/null +++ b/docs/operations/joint-crash-recovery.md @@ -0,0 +1,11 @@ +# Strict operational crash recovery for the constant-LR image run + +The original image-only discrete attempt412056 was a300k full allocation, not a scheduled allocation boundary. Its failed workers cannot be admitted through a fabricated scheduled completion. The explicit crash path binds its native terminal record and admitted request, immutable failure logs, and the highest healthy full-state checkpoint inspection. It rejects another live writer and requires a fresh immutable output/log namespace and same W&B identity with resume=must. + +The supported scientific identity is exactly `mot-shared-libero-image-discrete-lr1e-4-constant-s42-300k-20261021-v1`, scientific source `535d4fc26c4c85cf57bbd2a776cc79cfec085c4e`, four ranks/global32, constant AdamW1e-4 without warmup or scheduler. The budget remains300000 total updates, checkpoint and validation every20000. Scientific files are unchanged. The unlaunched image/proprio pair retains its separate first-launch admission; future crash recovery for that pair needs explicit generalization and actual matching proof. + +`crash_inspect.py` runs only on allocated compute, strictly reloads full wrapper/context/FP32 weights/AdamW/sampler/all-rank RNG and reconciles the highest checkpoint across hash-bound retained attempts. `crash_smoke.py` restores the inspected production checkpoint on four actual CUDA ranks into a distinct diagnostic namespace, executes two optimizer updates and100 scheduled validation anchors, reloads the diagnostic full state and performs a finite20-step hidden-label-NaN sample. A separate finalizer requires hostexit0 and nativeCOMPLETED0:0. Diagnostic artifacts never initialize production. + +`CrashRestoreProof` observes actual model/optimizer/sampler/RNG restoration and the first two updates on every rank. It rejects replayed boundary validation/saves and off-cadence saves. Production resume preserves all historical checkpoints. `crash_history.py` aggregates only each attempt's owned checkpoint/validation interval, preserving failed unsaved CSV history while rejecting duplicates. `verify_full.py` requires all15 immutable milestones, finite scheduled validation, final full-state reload/sample and actual restoration evidence. No extra production save is introduced. + +The101 new CPU tests and49 existing temporal-admission regressions verify the local implementation only. They do not certify production checkpoint health, destination storage, CUDA/DDP behavior or actual Lambda restoration. Before recovery, require current persistent storage read/write/fsync and space proof, parent terminal/resource-release evidence, immutable service deployment/readback, actual allocated inspection and the exact four-rank diagnostic. Do not submit while Lambda storage or node health is unverified. Never revive intentionally retired identities or change the scientific contract. diff --git a/docs/operations/joint-service-import-20261021/README.md b/docs/operations/joint-service-import-20261021/README.md new file mode 100644 index 000000000..385e3b6b9 --- /dev/null +++ b/docs/operations/joint-service-import-20261021/README.md @@ -0,0 +1,7 @@ +# Verified operational service import + +This import records the exact 23-file deployed operational service named `9a2982a8a611b9d377f7719dbd710161b39e4322` in the Lambda task `joint-modality-ablations-20261020/services/`. Its `FILES.sha256` bytes have SHA256 `1714c1b459dcef0dbe372651f0087e59391231f39f729896b7d75df5ce08d4ae`; every file was verified against that manifest before import. + +The name is a deployment identifier. Its Git ancestry was not established, and this commit does not assert that it descends from that identifier as a Git commit. The actual Git parent is the published modality-ablation branch at `255fdaf5137f7043b35df7ef48a3c1b8adc70815`. This creates a reviewable source home for verified operational bytes; it does not alter existing deployed releases or retrofit their provenance. + +Only orchestration/verification scripts and their declared dependency closure are imported. No `egomimic` model, data or configuration file changes. Existing scientific checkpoints and new image/proprio requests retain the separately verified frozen source `535d4fc26c4c85cf57bbd2a776cc79cfec085c4e` and its own archive/manifest bindings. diff --git a/docs/operations/joint-service-import-20261021/deployed-9a-FILES.sha256 b/docs/operations/joint-service-import-20261021/deployed-9a-FILES.sha256 new file mode 100644 index 000000000..1e4f346a4 --- /dev/null +++ b/docs/operations/joint-service-import-20261021/deployed-9a-FILES.sha256 @@ -0,0 +1,23 @@ +af44a2641d79754615746fe9a6c708b04b81ea807318cf9700a0ee2465bc7796 scripts/clusters/pace/__init__.py +bae908608d15f55ff8c0eab6f145baa8d6317112245d44b82b170831739f8b6f scripts/clusters/pace/gpu_admission.py +57861cbe471dfb352cc3a8818ec47dae76b062bfa5fe4eeb549d629b2e83e962 scripts/clusters/pace/gpu_ownership.py +13035d263f431a77e9d90bdf1b1e596787b19d913cbde767b1a2ca0c627bf23e scripts/clusters/pace/pace_podman.py +86ae32e41945db33a6d97a4a783efceb019bf817714508f4759db322fdbd75c0 scripts/ice/ice_checkpoint_mirror.py +64ed05603c3785266cd2e5b268dff71da1daac717f70971ee05b3de0579b70a7 scripts/train/joint/allocation.py +5ada8f1b544727719e2264c20edbb45e33d29f711455eb725bc7c9c94ee39b08 scripts/train/joint/allocation_boundary.py +2243cf846feb316d299e243db0a6a7ac378e5fcbffb8cddb934aca7d281c9ff9 scripts/train/joint/candidate_gate.py +df47a1cc2d73e59f68d2819aa01012e631dc6706f091ae2fea1655c51738214b scripts/train/joint/capture_native.sbatch +26b8c6c8f4db2753fb1f4996c603acda0cf8fb14797e7e17f37b1a8f2140ed47 scripts/train/joint/check_wandb.py +1953c4ebc6c25cc1180ed550671e648cc913f1f1ccde266731ae473491b372d2 scripts/train/joint/continuation.py +645dade981af90af9ca93851ce21a5f37e6dc54290b481c4b1db2046910c3794 scripts/train/joint/contract.py +a4feb1eb7db08a40d039942c96a3ab21da7a7db160abd3c0770090193e658c48 scripts/train/joint/lambda_full.sbatch +67523a129e6aa0b8af436ca97961f50ef3b1fb571c9bf23734bb85e5485e808e scripts/train/joint/lambda_inside.sh +4d5695db4e703a2ac1ea29418d34261ffe06c040fc367e5af960788715944d19 scripts/train/joint/lambda_runtime.py +50d3f31641ebaf494542d63cf43380486ca9ed50bc315c846cee354a2e91ce32 scripts/train/joint/lambda_stage_source.sh +884299fc0dc876147c32cda032d304b0abde261586148b7ebe18f01706cedac7 scripts/train/joint/lr_policy.py +3ca92ad5392c536cd9f47c715c6c6b2d3e8ceea0e69ecaf9c04eb427d72dfcd0 scripts/train/joint/recovery_check.py +caffa078e53eb022a265b0a79be973f37070568e787f64330a9bdcb7ecd061b9 scripts/train/joint/retention.py +792515e221bbf387c33ba136bdec08b48c7ad563c5fe783e97d869e9ee5d1235 scripts/train/joint/run.py +88a8e20ef93a958ede13203e1867d64734b05d2aea8cb08ffdf0176f562ce68b scripts/train/joint/service_bootstrap.py +e43cfe8483949a6c1c0c7bebf8ff52e589ee1e442564992b57e8f545b3d691e1 scripts/train/joint/verify_factorial.py +deb15f0a4a48b1f2843c2dfc2fcc87f431da30ca4e9bd61ff75dd953ccdbbade scripts/train/joint/verify_full.py diff --git a/docs/operations/joint-service-import-20261021/provenance.json b/docs/operations/joint-service-import-20261021/provenance.json new file mode 100644 index 000000000..e7da4fa90 --- /dev/null +++ b/docs/operations/joint-service-import-20261021/provenance.json @@ -0,0 +1,12 @@ +{ + "deployment_identifier_kind": "Recorded deployed service directory name; Git commit ancestry unverified", + "git_parent": "255fdaf5137f7043b35df7ef48a3c1b8adc70815", + "import_description": "Exact 23-file exported closure: 22 files match the earlier 8e service and contract.py differs; all manifest hashes were verified.", + "operational_manifest_sha256": "1714c1b459dcef0dbe372651f0087e59391231f39f729896b7d75df5ce08d4ae", + "operational_revision": "9a2982a8a611b9d377f7719dbd710161b39e4322", + "schema": "joint-operational-source-import-v1", + "scientific_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "scientific_files_modified": false, + "scientific_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "status": "VERIFIED_OPERATIONAL_IMPORT_NOT_DEPLOYED" +} diff --git a/docs/operations/joint-temporal-pair-admission.md b/docs/operations/joint-temporal-pair-admission.md new file mode 100644 index 000000000..7e6aae390 --- /dev/null +++ b/docs/operations/joint-temporal-pair-admission.md @@ -0,0 +1,13 @@ +# Native image/proprio temporal-pair admission + +Fresh image/proprio runs with discrete and continuous temporal representations do not satisfy the existing modality-factorial or optimizer-parent comparison contracts. `verify_temporal_pair.py` adds an explicit contract for this pair without fabricating either parent proof. + +The v1 profile pins scientific source `535d4fc`, seed 42, four ranks with eight examples each, global batch 32, constant AdamW LR `1e-4`, and 300k updates. Checkpointing and validation occur every 20k updates, retaining all 15 immutable checkpoints. An exact override allowlist prevents unrequested scientific changes, including changes shared by both arms. Discrete/continuous pose dimensions 36/9 and token limits 3,216/12,864 differ; every other composed scientific value must match after checked output and identity substitutions. + +Both actual GPU smoke proofs are independently revalidated, including 50 optimizer updates, 100 scheduled validation anchors, every rank's LR and precision checks, strict reload, and sampling proofs. All eight ranks must share the same initial image weights. Training-draw hashes, whole-model initial weights and data-context hashes are retained separately; they are not claimed to be paired across temporal representations. + +Two immutable full-request prototypes omit `temporal_comparison`. The verifier's receipt binds their original files and canonical complete request hashes. Final requests add only that receipt binding. Admission revalidates both prototypes, both smoke proof sets, both configurations and the complete receipt before requiring exact request membership. The driver requires exactly one comparison contract. + +This profile supports fresh single-allocation requests with identical runtime and data-cache bindings. It rejects continuations and allocation-boundary requests. The temporal binding is included in continuation stable fields for future extensions, but recovery and different node/cache bindings still require separately reviewed admission support. + +The 49 focused tests use configurations composed from frozen source `535d4fc` and synthetic smoke evidence fixtures. They do not certify GPU execution, runtime readiness, deployment or production training. Before deployment, bind a published service export; complete allocated runtime, data and model preflight; obtain two new exact four-rank smoke passes at LR `1e-4`; report measured timings; and refresh ownership, queues, claims, storage, W&B identities and output-collision checks. The scientific archive remains independent and unchanged. diff --git a/scripts/clusters/pace/gpu_admission.py b/scripts/clusters/pace/gpu_admission.py index 81e972862..58a852fb1 100644 --- a/scripts/clusters/pace/gpu_admission.py +++ b/scripts/clusters/pace/gpu_admission.py @@ -99,6 +99,8 @@ def worker(control): import torch.distributed as distributed rank = int(os.environ["LOCAL_RANK"]) + world = len(control["expected_uuids"]) + require(world in (2, 4), "Need two or four assigned GPUs") require( torch.__version__ == control["expected_torch"], "Scientific Torch version changed", @@ -114,7 +116,7 @@ def worker(control): "Frozen disabled-TF32 policy was not applied", ) require( - torch.cuda.device_count() == 4, + torch.cuda.device_count() == world, "Container does not expose exactly four CUDA GPUs", ) require(torch.cuda.is_available(), "CUDA unavailable") @@ -148,7 +150,7 @@ def worker(control): distributed.init_process_group("nccl", timeout=timedelta(seconds=120)) value = torch.tensor([float(rank + 1)], device="cuda") distributed.all_reduce(value) - require(value.item() == 10.0, "Four-rank NCCL reduction failed") + require(value.item() == world*(world+1)/2, "Four-rank NCCL reduction failed") distributed.barrier() torch.cuda.synchronize() distributed.destroy_process_group() @@ -157,35 +159,36 @@ def worker(control): { "status": "PASS", "rank": rank, - "world_size": 4, + "world_size": world, "uuid": actual_uuid, "name": properties.name, "compute_capability": [properties.major, properties.minor], "torch": torch.__version__, "torch_file": torch.__file__, "bf16": True, - "nccl_all_reduce": 10.0, + "nccl_all_reduce": world*(world+1)/2, }, ) def controller(control): base = validate_base(control) - require(len(control["expected_uuids"]) == 4, "Four assigned UUIDs required") + world = len(control["expected_uuids"]) + require(world in (2, 4), "Need two or four assigned UUIDs") with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) port = probe.getsockname()[1] children = [] started = time.monotonic() try: - for rank in range(4): + for rank in range(world): environment = dict(os.environ) environment.update( { "RANK": str(rank), "LOCAL_RANK": str(rank), - "WORLD_SIZE": "4", - "LOCAL_WORLD_SIZE": "4", + "WORLD_SIZE": str(world), + "LOCAL_WORLD_SIZE": str(world), "MASTER_ADDR": "127.0.0.1", "MASTER_PORT": str(port), "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", @@ -222,7 +225,7 @@ def controller(control): json.loads( (Path(control["receipt_directory"]) / f"gpu-rank-{rank}.json").read_text() ) - for rank in range(4) + for rank in range(world) ] require( [receipt["uuid"].lower() for receipt in receipts] diff --git a/scripts/clusters/pace/gpu_ownership.py b/scripts/clusters/pace/gpu_ownership.py index 3f01f5327..c833c163a 100644 --- a/scripts/clusters/pace/gpu_ownership.py +++ b/scripts/clusters/pace/gpu_ownership.py @@ -20,7 +20,8 @@ def require(condition, message): raise ValueError(message) -def indices(value): +def indices(value, world=4): + require(world in (2, 4), "Need two or four owned GPUs") result = [] for token in value.split(","): token = token.strip() @@ -35,7 +36,7 @@ def indices(value): ) result.append(token) require( - len(result) == len(set(result)) == 4, + len(result) == len(set(result)) == world, "Need exactly four distinct Slurm assignments", ) require( @@ -51,6 +52,8 @@ def bdf(value): def verify(evidence, expected): + world = expected.get("world_size", 4) + require(world in (2, 4), "Need two or four owned GPUs") cluster_lines = re.findall( r"^\s*ClusterName\s*=\s*(\S+)\s*$", evidence["scheduler_config"], re.M ) @@ -83,12 +86,12 @@ def verify(evidence, expected): fields.get("Restarts") == "0" and env.get("SLURM_RESTART_COUNT", "0") == "0", "Fresh allocation required", ) - allocations = indices(env.get("SLURM_JOB_GPUS", "")) + allocations = indices(env.get("SLURM_JOB_GPUS", ""), world) native = re.findall( - r"GRES=gpu:(?:h100|h200):4\(IDX:([0-9,-]+)\)", evidence["native"] + rf"GRES=gpu(?::h100|:h200)?:{world}\(IDX:([0-9,-]+)\)", evidence["native"] ) require(len(native) == 1, "Need one native four-Hopper GRES allocation") - native_ids = indices(native[0]) + native_ids = indices(native[0], world) require(all(x.isdigit() for x in native_ids), "Native IDX malformed") if allocations[0].isdigit(): require(set(native_ids) == set(allocations), "Slurm env and native IDX differ") @@ -134,8 +137,8 @@ def verify(evidence, expected): len(inventory) == len({r["uuid"] for r in inventory}) == len({r["index"] for r in inventory}) - == 4, - "Require four bijective filtered NVML GPUs", + == world, + "Require the exact bijective filtered NVML GPUs", ) physical = {} for path, text in evidence["proc"].items(): @@ -151,7 +154,7 @@ def verify(evidence, expected): "Ambiguous /proc GPU identity", ) physical[uuid] = dict(minor=minor, bdf=bus) - require(len(physical) >= 4, "Incomplete physical /proc inventory") + require(len(physical) >= world, "Incomplete physical /proc inventory") devices = evidence["devices"] require( len(devices) == len({r["minor"] for r in devices}) == len(physical), @@ -166,7 +169,7 @@ def verify(evidence, expected): "Invalid GPU character device", ) allowed = {r["minor"] for r in devices if r["read_open_permitted"]} - require(len(allowed) == 4, "Kernel must grant exactly four physical GPU minors") + require(len(allowed) == world, "Kernel must grant exactly four physical GPU minors") uuids = {r["uuid"] for r in inventory} require( {u for u, p in physical.items() if p["minor"] in allowed} == uuids, @@ -180,7 +183,7 @@ def verify(evidence, expected): xml = ET.fromstring(evidence["xml"]) xml_nodes = xml.findall("gpu") require( - len(xml_nodes) == 4 and {n.findtext("uuid") for n in xml_nodes} == uuids, + len(xml_nodes) == world and {n.findtext("uuid") for n in xml_nodes} == uuids, "XML/NVML UUID inventories differ", ) for n in xml_nodes: diff --git a/scripts/clusters/pace/pace_podman.py b/scripts/clusters/pace/pace_podman.py index 2b0eb4358..aa405747d 100644 --- a/scripts/clusters/pace/pace_podman.py +++ b/scripts/clusters/pace/pace_podman.py @@ -23,7 +23,10 @@ import xml.etree.ElementTree as ET from pathlib import Path, PurePosixPath -import gpu_ownership +try: + from . import gpu_ownership +except ImportError: + import gpu_ownership SCHEMA = "pace-podman-boundary-v2" GIB = 1024**3 @@ -774,6 +777,8 @@ def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--request", required=True) parser.add_argument("--request-sha256", required=True) + parser.add_argument("--check-request-only", action="store_true", + help="Verify request/helper hashes and namespace before GPU submission; no runtime or GPU claim") arguments = parser.parse_args() request_path = verified_file( {"path": arguments.request, "sha256": arguments.request_sha256} @@ -797,6 +802,10 @@ def main(): "Boundary helper is outside the reviewed directory", ) require(bool(NAME.fullmatch(request["namespace"])), "Invalid job namespace") + if arguments.check_request_only: + print(json.dumps(dict(status="BOUND_REQUEST_STATIC_PASS",namespace=request["namespace"], + request_sha256=arguments.request_sha256,bound_helpers=len(bound_helpers)))) + return 0 job_id = os.environ.get("SLURM_JOB_ID", "") require( job_id.isdigit() and os.environ.get("SLURMD_NODENAME"), @@ -984,7 +993,7 @@ def main(): "attestation": { "job_id": job_id, "node": os.environ["SLURMD_NODENAME"], - "world_size": 4, + "world_size": len(ownership["gpus"]), "ownership": ownership, "cluster": ownership["cluster"], "payload_workspace": str(local / "payload" / "workspace"), diff --git a/scripts/train/joint/allocation.py b/scripts/train/joint/allocation.py index a91e6ea5d..6e0f5f028 100644 --- a/scripts/train/joint/allocation.py +++ b/scripts/train/joint/allocation.py @@ -5,26 +5,30 @@ """ import os +import json +import subprocess +import sys from pathlib import Path -from scripts.train.joint.contract import bound, need, sha +from scripts.train.joint.contract import bound, need, sha, write_once PACE_ASSOCIATIONS = { - "ice-slurm": ("coc", "coc-ice"), - "phoenix-slurm": ("gts-dxu345-rl2", "inferno"), + "ice-slurm": (("coc", "coc-ice"), ("ece", "coe-ice")), + "phoenix-slurm": (("gts-dxu345-rl2", "inferno"),), } def verify_pace(request, attestation, environment): policy = request["allocation"] cluster = policy["cluster"] + world = request["world_size"] need(cluster in PACE_ASSOCIATIONS, "Unknown PACE cluster") need( - (policy["account"], policy["qos"]) == PACE_ASSOCIATIONS[cluster], + (policy["account"], policy["qos"]) in PACE_ASSOCIATIONS[cluster], "Wrong cluster association", ) need( - request["world_size"] == 4 and request["per_rank_batch"] == 8, + world in (2, 4) and request["per_rank_batch"] * world == 32, "PACE requires the matched four-rank global32 recipe", ) need( @@ -33,7 +37,7 @@ def verify_pace(request, attestation, environment): ) need( attestation["status"] == "PASS" - and attestation["world_size"] == 4 + and attestation["world_size"] == world and attestation["job_id"] == environment["SLURM_JOB_ID"] and attestation["node"] == environment["SLURMD_NODENAME"] and attestation["cluster"] == cluster, @@ -53,15 +57,18 @@ def verify_pace(request, attestation, environment): "Wrong current allocation owner/association", ) resources = ownership["native_resources"] + cpus = policy.get("cpus_per_task", 32) + need(type(cpus) is int and (cpus == 32 or (cluster == "phoenix-slurm" and world == 2 and cpus == 16)), + "Unsupported explicit CPU resource shape") need( - resources["NumCPUs"] == "32" + resources["NumCPUs"] == str(policy.get("cpus_per_task", 32)) and resources["MinMemoryNode"] == "192G" and resources["Requeue"] == "0", "Wrong dedicated CPU/memory/requeue contract", ) partitions = set(resources["Partition"].split(",")) expected_partitions = ( - {"ice-gpu"} if cluster == "ice-slurm" else {"gpu-h200", "gpu-h100"} + ({"coe-gpu"} if policy["account"] == "ece" else {"ice-gpu"}) if cluster == "ice-slurm" else {"gpu-h200", "gpu-h100"} ) need( partitions and partitions <= expected_partitions, "Unintended native partition" @@ -81,9 +88,9 @@ def verify_pace(request, attestation, environment): ) gpus, ranks = ownership["gpus"], probe["ranks"] need( - len(gpus) == len(ranks) == 4 - and len({g["uuid"].lower() for g in gpus}) == 4 - and len({g["minor"] for g in gpus}) == 4 + len(gpus) == len(ranks) == world + and len({g["uuid"].lower() for g in gpus}) == world + and len({g["minor"] for g in gpus}) == world and [r["uuid"].lower() for r in ranks] == [g["uuid"].lower() for g in gpus], "Wrong physical GPU/rank bijection", ) @@ -91,30 +98,88 @@ def verify_pace(request, attestation, environment): need( record["status"] == "PASS" and record["rank"] == rank - and record["world_size"] == 4 + and record["world_size"] == world and record["bf16"] - and record["nccl_all_reduce"] == 10.0 + and record["nccl_all_reduce"] == world*(world+1)/2 and record["torch"] == "2.9.1+cu126" and record["compute_capability"][0] == 9, "Incomplete rank precision/collective proof", ) - need( - attestation["payload_archive_sha256"] == policy["payload_archive_sha256"] - and attestation["base_archive_sha256"] == policy["base_archive_sha256"], - "Portable runtime/base binding changed", - ) + if policy.get("kind") == "native_slurm": + need(attestation["runtime_binding_sha256"] == policy["runtime_binding"]["sha256"], + "Native runtime binding changed") + else: + need( + attestation["payload_archive_sha256"] == policy["payload_archive_sha256"] + and attestation["base_archive_sha256"] == policy["base_archive_sha256"], + "Portable runtime/base binding changed", + ) return attestation -def admit(request): +def admit_native(request, status_root, bootstrap): + """Reuse exact ownership and real collective workers in an existing native runtime.""" + from scripts.clusters.pace import gpu_ownership + policy = request["allocation"] + need(status_root is not None and bootstrap is not None, "Native admission paths missing") + bound(policy["runtime_binding"]) + world = request["world_size"] + identity = {k: policy[k] for k in ("cluster", "user", "uid", "account", "qos")} + identity["world_size"] = world + ownership = gpu_ownership.verify(gpu_ownership.collect(), identity) + directory = Path(status_root)/("native-admission-"+os.environ["SLURM_JOB_ID"]) + directory.mkdir() + write_once(directory/"ownership.json", ownership) + from scripts.clusters.pace.pace_podman import gpu_health + xml = subprocess.run(["nvidia-smi", "-q", "-x", "-i", + ",".join(g["uuid"] for g in ownership["gpus"])], + check=True, capture_output=True, text=True, timeout=30).stdout + (directory/"gpu-health.xml").write_text(xml) + write_once(directory/"gpu-health.json", dict(status="PASS", devices=gpu_health(xml,ownership["gpus"]))) + uuids = [g["uuid"] for g in ownership["gpus"]] + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(uuids) + control = dict(expected_uuids=uuids, expected_torch="2.9.1+cu126", + receipt_directory=str(directory)) + control_path = directory/"control.json" + write_once(control_path, control) + environment = dict(os.environ, PHOENIX_BOUNDARY_CONTROL=str(control_path), + NVIDIA_TF32_OVERRIDE="0", TORCH_ALLOW_TF32_CUBLAS_OVERRIDE="0") + # Resolve the admitted operational module; old model snapshots need not + # contain the current cluster boundary at all. + from scripts.clusters.pace import gpu_admission + worker = Path(gpu_admission.__file__).resolve() + if request.get("operational_service"): + need(worker.is_relative_to(Path(request["operational_service"]["root"]).resolve()), + "GPU worker escaped the bound operational service") + subprocess.run([sys.executable, str(bootstrap), "-m", "torch.distributed.run", + "--standalone", "--nnodes=1", "--nproc-per-node="+str(world), + "--max-restarts=0", str(bootstrap), str(worker), "--worker"], + check=True, env=environment) + ranks = [json.loads((directory/f"gpu-rank-{i}.json").read_text()) for i in range(world)] + attestation = dict(status="PASS", kind="native_slurm", world_size=world, + job_id=os.environ["SLURM_JOB_ID"], node=os.environ["SLURMD_NODENAME"], + cluster=policy["cluster"], ownership=ownership, + runtime_binding_sha256=policy["runtime_binding"]["sha256"], + gpu_probe=dict(status="PASS", bf16_passed=True, nccl_passed=True, ranks=ranks)) + verify_pace(request, attestation, os.environ) + path = directory/"attestation.json" + write_once(path, attestation) + return dict(path=str(path), sha256=sha(path), job_id=attestation["job_id"], + node=attestation["node"], cluster=attestation["cluster"]) + + +def admit(request, status_root=None, bootstrap=None): policy = request.get("allocation") if policy is None: + node = request.get("lambda_runtime", {}).get("node", "gpu21") need( - os.environ["SLURMD_NODENAME"] == "gpu21" + os.environ["SLURMD_NODENAME"] == node and os.environ.get("SLURM_RESTART_COUNT", "0") == "0", "Wrong Lambda allocation", ) return None + if policy.get("kind") == "native_slurm": + return admit_native(request, status_root, bootstrap) path = Path("/run/pace-proof/attestation.json") need(path.is_file() and not path.is_symlink(), "Missing current PACE attestation") attestation = bound(dict(path=str(path), sha256=sha(path))) diff --git a/scripts/train/joint/allocation_boundary.py b/scripts/train/joint/allocation_boundary.py new file mode 100644 index 000000000..03db77318 --- /dev/null +++ b/scripts/train/joint/allocation_boundary.py @@ -0,0 +1,90 @@ +"""End an allocation at a scheduled validation boundary without changing its budget. + +The existing Lightning full-state loader owns continuation. This callback only +requests a graceful stop and publishes an additional recovery checkpoint after +validation, so restoring it also restores the post-validation RNG state. It +never submits jobs, signals supervisors, removes checkpoints or edits max_steps. +""" + +from pathlib import Path + +from lightning import Callback + +from scripts.train.joint.contract import need, sha, write_once + + +def check_interval(start_step, stop_step, target_step, cadence): + values = (start_step, stop_step, target_step, cadence) + need(all(type(v) is int for v in values), "Integer allocation steps required") + need( + cadence > 0 and 0 <= start_step < stop_step < target_step + and start_step % cadence == stop_step % cadence == target_step % cadence == 0, + "Allocation boundary must preserve the scheduled cadence and final budget", + ) + + +class ScheduledAllocationBoundary(Callback): + def __init__(self, *, start_step, stop_step, target_step, cadence, output): + check_interval(start_step, stop_step, target_step, cadence) + self.start_step, self.stop_step = start_step, stop_step + self.target_step, self.cadence = target_step, cadence + self.output = Path(output) + need(self.output.is_absolute(), "Absolute allocation evidence path required") + self.receipt = self.output / f"allocation-step{stop_step:09d}.json" + + def on_train_start(self, trainer, pl_module): + need(trainer.max_steps == self.target_step, "Scientific budget changed") + need(trainer.global_step == self.start_step, "Wrong restored optimizer step") + need( + trainer.accumulate_grad_batches == 1 + and trainer.val_check_interval == self.cadence + and trainer.check_val_every_n_epoch is None + and not trainer.sanity_checking, + "Allocation must retain the configured optimizer/validation cadence", + ) + checkpoint = trainer.checkpoint_callback + need( + checkpoint is not None + and checkpoint._every_n_train_steps == self.cadence + and checkpoint.save_top_k == -1 + and checkpoint.save_last == "link", + "Immutable scheduled-checkpoint retention is required", + ) + need(not self.receipt.exists(), "Allocation boundary already completed") + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + need(trainer.global_step <= self.stop_step, "Allocation overshot its boundary") + if trainer.global_step == self.stop_step: + trainer.should_stop = True + + def on_train_end(self, trainer, pl_module): + need(trainer.global_step == self.stop_step, "Allocation ended early") + need( + pl_module.evaluator.has_completed_validation(self.stop_step), + "Boundary validation must complete before recovery publication", + ) + need( + trainer.datamodule.consumed_batches == self.stop_step, + "Sampler cursor differs from completed optimizer updates", + ) + self.output.mkdir(parents=True, exist_ok=True) + recovery = self.output / ( + f"allocation-recovery-epoch{trainer.current_epoch:04d}" + f"-step{self.stop_step:09d}.ckpt" + ) + need(not recovery.exists(), "Recovery checkpoint must never be overwritten") + # All ranks participate in the native distributed RNG checkpoint hook. + trainer.save_checkpoint(recovery) + trainer.strategy.barrier("allocation-recovery-published") + if trainer.is_global_zero: + write_once(self.receipt, dict( + schema="joint-allocation-boundary-v1", status="PASS", + start_step=self.start_step, global_step=self.stop_step, + target_step=self.target_step, cadence=self.cadence, + world_size=trainer.world_size, validation_completed=True, + consumed_batches=trainer.datamodule.consumed_batches, + checkpoint=dict(path=str(recovery), sha256=sha(recovery), + size_bytes=recovery.stat().st_size), + checkpoint_saved_after_validation=True, + )) + trainer.strategy.barrier("allocation-receipt-published") diff --git a/scripts/train/joint/candidate_gate.py b/scripts/train/joint/candidate_gate.py new file mode 100644 index 000000000..f125adb33 --- /dev/null +++ b/scripts/train/joint/candidate_gate.py @@ -0,0 +1,87 @@ +"""Fail-closed allocation gate for coordinator-selected cross-cluster candidates. + +Run inside the allocated batch before invoking the existing training launcher. +The coordinator must persist one exclusive winner and reconcile every ungated +candidate before publishing this job/request-specific grant. This module never +self-elects, launches, releases, cancels, or edits another scheduler job. +""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import time + + +def write_once(path, value): + with Path(path).open('x') as stream: + json.dump(value, stream, indent=2, sort_keys=True) + stream.write('\n') + + +def identity(request_path, environment): + path = Path(request_path) + raw = path.read_bytes() + request = json.loads(raw) + policy = request['candidate_gate'] + if policy['schema'] != 'coordinator-candidate-v1': + raise ValueError('Unknown candidate gate contract') + job = environment['SLURM_JOB_ID'] + if not job.isdecimal() or environment.get('SLURM_RESTART_COUNT', '0') != '0': + raise ValueError('A fresh native allocation is required') + if environment['USER'] != policy['user']: + raise ValueError('Candidate owner mismatch') + directory = Path(policy['directory']) + if not directory.is_absolute() or directory.is_symlink(): + raise ValueError('Gate directory must be an absolute non-symlink path') + expected = dict(schema='coordinator-candidate-v1', job_id=job, + host=policy['host'], user=policy['user'], cell=policy['cell'], + round_id=policy['round_id'], request_sha256=hashlib.sha256(raw).hexdigest()) + return policy, expected + + +def validate_grant(grant, expected): + if any(grant.get(k) != v for k, v in expected.items()): + raise ValueError('Grant does not identify this exact candidate') + if grant.get('decision') != 'TRAIN' or grant.get('losers_reconciled') is not True: + raise ValueError('Training is not authorized by this grant') + winner_sha = grant.get('winner_record_sha256', '') + if len(winner_sha) != 64 or any(c not in '0123456789abcdef' for c in winner_sha): + raise ValueError('Missing immutable coordinator winner receipt') + + +def wait_for_grant(request_path, environment=None): + environment = os.environ if environment is None else environment + policy, expected = identity(request_path, environment) + timeout = policy.get('wait_seconds', 600) + if type(timeout) is not int or not 1 <= timeout <= 900: + raise ValueError('Gate wait must be bounded to at most 900 seconds') + directory = Path(policy['directory']) + directory.mkdir(parents=True, exist_ok=True) + prefix = directory / expected['job_id'] + write_once(str(prefix)+'.ready.json', dict(expected, node=environment['SLURMD_NODENAME'], + optimizer_started=False)) + deadline = time.monotonic() + timeout + while True: + reject = Path(str(prefix)+'.reject.json') + grant_path = Path(str(prefix)+'.grant.json') + if reject.exists() or reject.is_symlink(): + raise RuntimeError('Candidate was rejected before training') + if grant_path.is_symlink(): + raise ValueError('Symlink grant rejected') + if grant_path.exists(): + raw = grant_path.read_bytes() + validate_grant(json.loads(raw), expected) + write_once(str(prefix)+'.accepted.json', dict(expected, + grant_sha256=hashlib.sha256(raw).hexdigest())) + return + if time.monotonic() >= deadline: + raise TimeoutError('No exact coordinator grant; training never started') + time.sleep(min(2, max(0, deadline-time.monotonic()))) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('request') + args = parser.parse_args() + wait_for_grant(args.request) diff --git a/scripts/train/joint/continuation.py b/scripts/train/joint/continuation.py new file mode 100644 index 000000000..6b5908c20 --- /dev/null +++ b/scripts/train/joint/continuation.py @@ -0,0 +1,125 @@ +"""Strict request chain for scheduled native joint-training allocations. + +This only extends the shared full service's allocation lifetime. The scientific +target, optimizer schedule, model and scheduled validation/save cadence stay fixed. +""" + +from pathlib import Path + +from scripts.train.joint.contract import bound, need, sha + + +STABLE_FIELDS = ( + "run_id", "arm", "source", "source_revision", "source_manifest", + "source_manifest_sha256", "data_manifest", "data_manifest_sha256", "data_cache", + "source_archive", "source_archive_sha256", "parameters", "expert_parameters", + "expert_widths", "world_size", "global_batch", "per_rank_batch", "lr_policy", + "vector_expert_parameter_budget", "active_streams", "experiment", "target_steps", + "output", "allocation", "smoke", "factorial_controls", "factorial_experiments", + "allocation_boundaries", "archive_policy", "quota", "reserved_checkpoint_bytes", + "allocation_recovery", + "operational_service", "temporal_comparison", +) + + +def allocation_steps(request): + """No continuation is inferred from a directory or W&B's status.""" + if request.get("crash_recovery") is not None: + from scripts.train.joint.crash_recovery import admission + + _, _, start = admission(request) + return start, request["target_steps"] + target = request["target_steps"] + stop = request.get("allocation_stop_step", target) + start = 0 + continuation = request.get("continuation") + if continuation is not None: + parent = bound(continuation["parent_request"]) + receipt = bound(continuation["completion"]) + need(parent["schema"] == "joint-training-v2" and parent["phase"] == "full", + "Continuation parent is not a full request") + need(all(request.get(k) == parent.get(k) for k in STABLE_FIELDS), + "Continuation changed its scientific identity or output") + need(receipt["schema"] == "joint-mot-allocation-v1" and receipt["status"] == "PASS" + and receipt["request_sha256"] == continuation["parent_request"]["sha256"] + and receipt["run_id"] == request["run_id"] + and receipt["world_size"] == request["world_size"] + and receipt["global_step"] == parent["allocation_stop_step"] + and receipt["global_step"] < target + and receipt["strict_full_wrapper_optimizer_scheduler_reload"] + and receipt["sampler_all_rank_rng_verified"] and receipt["sample_finite"], + "Missing strict prior-allocation completion") + checkpoint = continuation["checkpoint"] + need(checkpoint["path"] == receipt["checkpoint"] + and checkpoint["sha256"] == receipt["checkpoint_sha256"] + and sha(checkpoint["path"]) == checkpoint["sha256"], + "Continuation recovery checkpoint changed") + start = receipt["global_step"] + need(type(start) is int and type(stop) is int + and 0 <= start < stop <= target + and start % 20000 == stop % 20000 == target % 20000 == 0, + "Allocation boundaries must preserve20k cadence and total budget") + if stop < target or continuation is not None: + plan = request["allocation_boundaries"] + need(plan == sorted(set(plan)) and plan[-1] == target + and stop == next((v for v in plan if v > start), None), + "Allocation skipped or changed its declared boundaries") + proof = bound(request["allocation_recovery"]) + need(proof["schema"] == "joint-recovery-check-v1" and proof["status"] == "PASS" + and proof["source_revision"] == request["source_revision"] + and proof["data_manifest_sha256"] == request["data_manifest_sha256"] + and proof["active_streams"] == request["active_streams"] + and proof["world_size"] == request["world_size"] and proof["global_batch"] == 32 + and proof["smoke_completion_sha256"] == request["smoke"]["completion"]["sha256"] + and proof["actual_updates"] == 6 and proof["all_rank_full_state_restored"] + and proof["model_data_unchanged"] and proof["no_extra_validation"], + "Actual matching BF16/DDP allocation recovery is not proven") + return start, stop + + +def configure_allocation(cfg, request, start, stop, *, request_sha256=None): + """Apply operational additions only AFTER the existing science comparison.""" + if request.get("crash_recovery") is not None: + from scripts.train.joint.crash_recovery import admission, configure + + _, inspection, checked_start = admission(request) + need(checked_start == start and stop == request["target_steps"], + "Crash recovery boundary changed after admission") + return configure(cfg, request, inspection, start, request_sha256=request_sha256) + from omegaconf import open_dict + + with open_dict(cfg): + cfg.logger.csv.save_dir = request["local_logs"] + cfg.logger.csv.version = f"step{start:09d}" + if start: + cfg.ckpt_path = request["continuation"]["checkpoint"]["path"] + cfg.logger.wandb.resume = "must" + if stop < request["target_steps"]: + cfg.callbacks.allocation_boundary = dict( + _target_="scripts.train.joint.allocation_boundary.ScheduledAllocationBoundary", + start_step=start, stop_step=stop, target_step=request["target_steps"], + cadence=20000, output=str(Path(request["output"])/"allocation-evidence"), + ) + return cfg + + +def log_segments(request): + """Keep each allocation's CSV events; reject cycles and repeated log roots.""" + if request.get("crash_recovery") is not None: + from scripts.train.joint.crash_history import history_segments + + return [segment["csv_path"] for segment in history_segments(request)] + result, seen = [], set() + for _ in range(16): + log = request["local_logs"] + need(log not in seen, "Repeated continuation log root or request cycle") + seen.add(log) + previous = request.get("continuation") + start = 0 if previous is None else bound(previous["completion"])["global_step"] + relative = Path("csv")/f"step{start:09d}"/"metrics.csv" + current = Path(log)/relative + result.append(current if current.is_file() else Path(request["output"])/"log-mirror"/relative) + if previous is None: + return result[::-1] + request = bound(previous["parent_request"]) + raise ValueError("Too many allocation attempts for the retained20k cadence") diff --git a/scripts/train/joint/contract.py b/scripts/train/joint/contract.py index e36b5939f..558c6e842 100644 --- a/scripts/train/joint/contract.py +++ b/scripts/train/joint/contract.py @@ -4,6 +4,8 @@ import hashlib import json import math +import uuid +from scripts.train.joint.lr_policy import validate, update_lr from pathlib import Path @@ -12,6 +14,93 @@ def need(condition, message): raise ValueError(message) + +def verify_model_execution(request, cfg): + """Bind native historical image/state execution without changing legacy admission.""" + stages = cfg["model"]["pipeline"]["stages"] + ids = cfg["model"]["pipeline"]["stage_ids"] + model = stages[ids["joint_denoiser"]]["model"] + native = model["_target_"] == "egomimic.models.joint_image_proprio.JointImageProprioModel" + policy = request.get("model_execution") + if not native: + need(policy is None, "Native execution policy requires the native image wrapper") + return "visual_state_action" + need(policy == {"implementation": "native_image_joint", "compile_mode": "default", + "compile_dynamic": True, "compile_backend": "inductor"}, + "Native image execution requires an explicit historical request binding") + image = model["image_model"] + need(image["_target_"] == "egomimic.models.image_proprio.ImageProprioModel" + and image["compile_mode"] == policy["compile_mode"] + and image["compile_dynamic"] == policy["compile_dynamic"] + and image["compile_backend"] == policy["compile_backend"], + "Native image compilation differs from bound historical policy") + streams = request["active_streams"] + action = streams == ["image", "proprio", "action"] + need(streams in (["image"], ["image", "proprio"], ["image", "proprio", "action"]) + and model["active_streams"] == streams, + "Native request has unbound streams") + need([stage["stream"] for stage in stages if "stream" in stage] == streams, + "Native loss graph differs from active experts") + if not action: + need(request.get("action_contract") is None + and not any("action" in name for name in ids), "Action stage in native state graph") + need(stages[ids["joint_noising"]]["noising"]["_target_"] + == "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "Native image corruption path changed") + need(stages[ids["joint_noising"]]["noising"].get("active_streams") == streams, + "Noising must exclude inactive vector draws") + need(cfg["data"].get("image_only") == (streams == ["image"]), + "Image-only data must use the historical image-only path") + sampler = stages[ids["joint_sampler"]]["sampler"] + need(sampler["_target_"] == "egomimic.models.joint_image_proprio.JointImageProprioSampler" + and sampler["state_steps"] == 20 + and (sampler.get("action_steps") == 20 if action else "action_steps" not in sampler), + "Native sampling path changed") + if action: + contract = request.get("action_contract") + need(isinstance(contract, dict) and contract.get("schema") == "native-joint-command-v1", + "Missing explicit native action contract") + native, encoded = contract.get("native_dimensions"), contract.get("encoded_dimensions") + group = cfg["data"]["stride"] + need(isinstance(native, dict) and native and all(type(v) is int and v > 0 for v in native.values()) + and cfg["data"]["mode"] == "discrete" and type(group) is int and group > 0 + and contract.get("group_size") == group + and encoded == {k: v * group for k, v in native.items()} + and model.get("action_dims") == encoded + and model.get("action_width") == contract.get("hidden_width") == 256 + and contract.get("noise") == "independent_command" + and contract.get("unknown_past_queries") == "pure_noise_no_loss", + "Action dimensions/window/noising contract changed") + need(cfg["data"].get("couple_state_observations") is True + and cfg["data"]["task"] == "visual_state_action" + and cfg["evaluator"]["objective"] == "visual_state_action" + and all(s.get("weight") == 1.0 for s in stages if "stream" in s), + "Three-stream data/objective/loss contract changed") + return "visual_state_action" + need(cfg["data"]["task"] == "visual_diffusion", "Native state data must exclude action streams") + return "visual_diffusion" + + +def verify_sampling_proof(request, completion): + policy = validate(request.get("lr_policy")) + native = request.get("model_execution") is not None + expected_action = (20,) if "action" in request["active_streams"] else ((0,) if native else (0, 20)) + need(completion["sample_state_steps"] == 20 + and completion["sample_action_steps"] in expected_action + and completion["lr_warmup_steps"] == policy["warmup_optimizer_steps"], + "Wrong sampling/scheduler proof") + if native: + need(completion.get("model_execution") == request["model_execution"], + "Missing native historical execution proof") + +def canonical_uuid(value): + """Compare complete CUDA/NVML identities without accepting GPU ordinals.""" + text = str(value) + if text[:4].lower() == "gpu-": + text = text[4:] + return str(uuid.UUID(text)) + + def sha(path): with Path(path).open("rb") as stream: return hashlib.file_digest(stream, "sha256").hexdigest() @@ -32,6 +121,38 @@ def bound(binding, json_file=True): return json.loads(content) if json_file else content +def verify_allocation_relocation(request, smoke, smoke_request_sha256): + """Allow explicit account/storage relocation on the same admitted cluster. + + Runtime, topology and hardware policy remain exact; the destination still + undergoes its own native allocation/UUID admission. Historical evidence is + always validated against the original smoke owner. + """ + old, new = smoke.get("allocation"), request.get("allocation") + if old == new: + need(request.get("allocation_relocation") is None, + "Unexpected relocation for an unchanged allocation policy") + return + need(isinstance(old, dict) and isinstance(new, dict) + and old.get("kind") == new.get("kind") == "native_slurm", + "Relocation requires native Slurm on both sides") + allowed = {"user", "uid", "account", "qos"} + need({k: v for k, v in old.items() if k not in allowed} + == {k: v for k, v in new.items() if k not in allowed} + and old.get("cluster") and old.get("runtime_binding"), + "Relocation changed cluster, runtime or hardware policy") + need(request.get("allocation_relocation") == dict( + schema="same-cluster-account-relocation-v1", + smoke_request_sha256=smoke_request_sha256, + original_allocation=old, destination_allocation=new), + "Missing exact original/destination relocation binding") + need(all(isinstance(new.get(k), str) and new[k] for k in ("user", "account", "qos")) + and type(new.get("uid")) is int and new["uid"] > 0, + "Invalid destination ownership") + for key in ("source", "world_size", "active_streams", "parameters"): + need(request.get(key) == smoke.get(key), "Recovery science changed: " + key) + + def verify_evidence(request): need( request["schema"] == "joint-training-v2" and request["phase"] == "full", @@ -54,22 +175,21 @@ def verify_evidence(request): native["JobId"] == str(evidence["job_id"]) and native["JobState"] == "COMPLETED" and native["ExitCode"] == "0:0" - and native["UserId"].startswith("paphiwetsa3(") + and native["UserId"].startswith((smoke.get("allocation") or {}).get("user", "paphiwetsa3")+"(") and native["JobName"] == smoke["job_name"], "Wrong native smoke proof", ) if smoke.get("allocation") is None: need( - request.get("allocation") is None and native["NodeList"] == "gpu21", + request.get("allocation") is None + and smoke.get("lambda_runtime") == request.get("lambda_runtime") + and native["NodeList"] == request.get("lambda_runtime", {}).get("node", "gpu21"), "Wrong Lambda smoke allocation", ) else: from scripts.train.joint.allocation import verify_pace - need( - request.get("allocation") == smoke["allocation"], - "Destination policy changed after smoke", - ) + verify_allocation_relocation(request, smoke, evidence["request"]["sha256"]) allocation = bound(evidence["allocation"]) reference = preflight.get("allocation_attestation") need( @@ -115,6 +235,8 @@ def verify_evidence(request): "experiment", ): need(smoke[key] == request[key], "Production differs from smoke: " + key) + need(smoke.get("model_execution") == request.get("model_execution"), + "Production model execution differs from smoke") need( preflight["status"] == "PREFLIGHT_PASS" and preflight["request_sha256"] == evidence["request"]["sha256"], @@ -147,11 +269,7 @@ def verify_evidence(request): "sample_finite", ): need(completion[key] is True, "Missing strict proof: " + key) - need( - completion["sample_state_steps"] == completion["sample_action_steps"] == 20 - and completion["lr_warmup_steps"] == 3000, - "Wrong sampling/scheduler proof", - ) + verify_sampling_proof(request, completion) world = request["world_size"] metric_names = { f"val/libero_gripper/{kind}/{metric}" @@ -172,14 +290,14 @@ def verify_evidence(request): ) ranks = completion["ranks"] need( - len(ranks) == world and len({r["gpu_uuid"] for r in ranks}) == world, + len(ranks) == world and len({canonical_uuid(r["gpu_uuid"]) for r in ranks}) == world, "Wrong GPU topology", ) need(len({r["weights_sha256"] for r in ranks}) == 1, "Ranks diverged") if smoke.get("allocation") is not None: need( - [r["gpu_uuid"].lower() for r in ranks] - == [r["uuid"].lower() for r in allocation["gpu_probe"]["ranks"]], + [canonical_uuid(r["gpu_uuid"]) for r in ranks] + == [canonical_uuid(r["uuid"]) for r in allocation["gpu_probe"]["ranks"]], "Training ranks differ from admitted physical GPUs", ) for rank, record in enumerate(ranks): @@ -216,7 +334,7 @@ def verify_evidence(request): need( len(record["optimizer_update_lrs"]) == 50 and all( - math.isclose(lr, 5e-6 * (i + 1) / 3000, rel_tol=1e-12) + math.isclose(lr, update_lr(request["lr_policy"], i), rel_tol=1e-12) for i, lr in enumerate(record["optimizer_update_lrs"]) ), "Wrong actual LR warmup", @@ -235,7 +353,7 @@ def verify_evidence(request): and capacity["frames_per_example"] == (64 if request["arm"] == "continuous" else 16) and capacity["per_rank_batch"] == request["per_rank_batch"] - and math.isclose(capacity["optimizer_update_lr"], 5e-6 / 3000, rel_tol=1e-12), + and math.isclose(capacity["optimizer_update_lr"], update_lr(request["lr_policy"], 0), rel_tol=1e-12), "Missing maximum-shape capacity proof", ) bound(evidence["resolved"], False) @@ -339,3 +457,107 @@ def relocate(value): and production["data"]["batch_size"] * trainer["devices"] == 32, "Fresh step0/globalbatch required", ) + + +def verify_optimizer_comparison(request, completion): + """Admit an independent optimizer-recipe trial against an immutable parent. + + Topology differences are explicit confounders, never a claim of paired RNG. + All source-owned model/data/training/config files must be byte-identical; + only smoke observation and verification tools may change. + """ + comparison = request["optimizer_comparison"] + parent = bound(comparison["parent_request"]) + need(comparison["schema"] == "optimizer-recipe-comparison-v1" + and comparison["topology_control"] == "global_batch_matched_rank_rng_unpaired" + and parent["phase"] == "full" and request["run_id"] != parent["run_id"], + "Invalid fresh optimizer comparison") + for field in ("data_manifest_sha256", "active_streams", "parameters", "expert_parameters", + "expert_widths", "global_batch", "target_steps", "model_execution", "arm", "experiment"): + need(request[field] == parent[field], "Optimizer comparison changed science: " + field) + need(request["active_streams"] in (["image"], ["image", "proprio"]) + and request["world_size"] * request["per_rank_batch"] == 32, + "Optimizer comparison requires native image/state streams and global32") + need(validate(request["lr_policy"]) != validate(parent["lr_policy"]), "Identical optimizer recipes") + def scientific_manifest(binding): + return {line.split(" ", 1)[1]: line.split(" ", 1)[0] + for line in Path(binding).read_text().splitlines() + if line.split(" ", 1)[1].startswith("egomimic/") + and line.split(" ", 1)[1] != "egomimic/pl_utils/joint_smoke.py"} + need(sha(parent["source_manifest"]) == parent["source_manifest_sha256"], "Changed parent manifest") + old, new = scientific_manifest(parent["source_manifest"]), scientific_manifest(request["source_manifest"]) + need(old and old == new, "Model/data/learning-loop/config source changed") + baseline = bound(parent["smoke"]["completion"]) + need(baseline["status"] == "PASS" and completion["status"] == "PASS" + and baseline["ranks"][0]["initial_image_weights_sha256"] + == completion["ranks"][0]["initial_image_weights_sha256"], + "Optimizer trial did not preserve actual initial image weights") + if "proprio" in request["active_streams"]: + old_initial = baseline["ranks"][0].get("initial_weights") + new_initial = completion["ranks"][0].get("initial_weights") + need(isinstance(old_initial, list) and isinstance(new_initial, list) + and old_initial and new_initial + and all(isinstance(x, str) and len(x) == 64 for x in old_initial + new_initial) + and len(set(old_initial + new_initial)) == 1, + "Optimizer trial did not preserve actual initial image/proprio weights") + return parent + + +def verify_modality_extension(request, completion): + """Bind the single optional command bank to an actual paired state smoke. + + This is a modality comparison, never an optimizer-only comparison. Existing + image/state GPU evidence is reused; the new arm must supply its own full + smoke/restart evidence and actual projected initialization and RGB draws. + """ + comparison = request["modality_comparison"] + need(comparison["schema"] == "native-command-extension-v1", "Wrong extension comparison") + parent = bound(comparison["parent_request"]) + baseline, _ = verify_evidence(parent) + need(parent["active_streams"] == ["image", "proprio"] + and request["active_streams"] == ["image", "proprio", "action"] + and request["experiment"] == "image_proprio_action" + and parent["run_id"] != request["run_id"], "Wrong command extension parent") + for key in ("data_manifest_sha256", "world_size", "per_rank_batch", "global_batch", + "target_steps", "arm", "model_execution", "lr_policy"): + need(request[key] == parent[key], "Command comparison changed control: " + key) + for name in parent["active_streams"]: + need(request["expert_parameters"][name] == parent["expert_parameters"][name], + "Existing expert parameter count changed") + need(sha(parent["source_manifest"]) == parent["source_manifest_sha256"], "Changed parent source") + def scientific(path): + return {name: digest for digest, name in (line.split(" ", 1) + for line in Path(path).read_text().splitlines()) + if name.startswith("egomimic/") and name != "egomimic/pl_utils/joint_smoke.py"} + old, new = scientific(parent["source_manifest"]), scientific(request["source_manifest"]) + changed = {k for k in old.keys() | new.keys() if old.get(k) != new.get(k)} + need(changed == {"egomimic/models/joint_image_proprio.py", + "egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio_action.yaml"}, + "Command extension changed unrelated model/data/config/learning-loop source") + old_cfg = json.loads(json.dumps(bound(parent["smoke"]["resolved"], False))) + new_cfg = bound(request["smoke"]["resolved"], False) + from omegaconf import OmegaConf + old_cfg = OmegaConf.to_container(OmegaConf.create(old_cfg), resolve=True) + new_cfg = OmegaConf.to_container(OmegaConf.create(new_cfg), resolve=True) + for key in ("seed", "world_model_attention"): + need(old_cfg[key] == new_cfg[key], "Command comparison changed config: " + key) + runtimes = [] + for cfg, reference in ((old_cfg, parent), (new_cfg, request)): + runtime = copy.deepcopy(cfg["runtime"]) + smoke_output = bound(reference["smoke"]["request"])["output"] + need(runtime.pop("slurm_signal_checkpoint_dir", None) + == str(Path(smoke_output)/"checkpoints"), + "Command comparison has an unbound signal checkpoint directory") + runtimes.append(runtime) + need(runtimes[0] == runtimes[1], "Command comparison changed runtime semantics") + old_data, new_data = copy.deepcopy(old_cfg["data"]), copy.deepcopy(new_cfg["data"]) + need(old_data.pop("task") == "visual_diffusion" + and new_data.pop("task") == "visual_state_action" and old_data == new_data, + "Command comparison changed data beyond command labels") + need(len(completion["ranks"]) == len(baseline["ranks"]) == request["world_size"], "Wrong paired rank count") + for old_rank, new_rank in zip(baseline["ranks"], completion["ranks"]): + need(new_rank.get("initial_state_control_weights_sha256") in old_rank["initial_weights"] + and new_rank["initial_image_weights_sha256"] == old_rank["initial_image_weights_sha256"] + and new_rank["paired_rgb_batch_sha256"] == old_rank["paired_rgb_batch_sha256"], + "Unpaired actual image/state initialization or RGB/mask/noise draws") + return parent diff --git a/scripts/train/joint/crash_history.py b/scripts/train/joint/crash_history.py new file mode 100644 index 000000000..9e002811f --- /dev/null +++ b/scripts/train/joint/crash_history.py @@ -0,0 +1,184 @@ +"""Read retained crash-attempt history without rewriting any historical bytes. + +This module does no checkpoint loading and never changes a file. A parent whose +workers continued beyond the restored checkpoint contributes events only through +that checkpoint. Its later rows remain preserved as lost, unsaved work. +""" + +import csv +import math +from pathlib import Path +import re + +from scripts.train.joint.contract import bound, need + + +CADENCE = 20000 + + +def _start_and_parent(request): + crash, continuation = request.get("crash_recovery"), request.get("continuation") + need(not (crash is not None and continuation is not None), + "A request cannot be both crash recovery and scheduled continuation") + if crash is not None: + need("checkpoint_inspection" not in crash, "Ambiguous crash inspection binding") + parent = bound(crash["parent_request"]) + receipt = bound(crash["inspection"]) + need(receipt["schema"] == "joint-crash-checkpoint-inspection-v1" + and receipt["status"] == "PASS" + and receipt["parent_request_sha256"] == crash["parent_request"]["sha256"] + and receipt["run_id"] == request["run_id"] + and receipt["source_revision"] == request["source_revision"], + "Unbound crash history inspection") + start = receipt["global_step"] + elif continuation is not None: + parent = bound(continuation["parent_request"]) + receipt = bound(continuation["completion"]) + need(receipt["status"] == "PASS" and receipt["run_id"] == request["run_id"] + and receipt["request_sha256"] == continuation["parent_request"]["sha256"], + "Unbound scheduled history completion") + start = receipt["global_step"] + else: + return 0, None + need(type(start) is int and 0 < start < request["target_steps"] + and start % CADENCE == 0, "Invalid retained history boundary") + need(parent["run_id"] == request["run_id"] + and parent["source_revision"] == request["source_revision"] + and parent["target_steps"] == request["target_steps"] + and parent["world_size"] == request["world_size"] + and parent["active_streams"] == request["active_streams"], + "Historical attempt belongs to another scientific identity") + return start, parent + + +def history_segments(request, target=None): + """Return chronological immutable CSV/output segments with inclusive cutoffs. + + A segment owns optimizer milestones ``start_step < step <= through_step``. + CSV rows at zero-based ``step-1`` map to their actual optimizer milestone. + Empty segments are retained so a zero-progress failed attempt cannot hide an + unexpected checkpoint or duplicated boundary validation. + """ + target = request["target_steps"] if target is None else target + need(type(target) is int and 0 < target <= request["target_steps"] + and target % CADENCE == 0, "Invalid history target") + rows, seen_outputs, seen_logs = [], set(), set() + current, through = request, target + for _ in range(32): + output, log = Path(current["output"]), Path(current["local_logs"]) + need(output.is_absolute() and log.is_absolute() + and ".." not in output.parts and ".." not in log.parts, + "Unsafe retained attempt path") + resolved_output, resolved_log = output.resolve(), log.resolve() + need(all(resolved_output != previous and resolved_output not in previous.parents + and previous not in resolved_output.parents for previous in seen_outputs) + and all(resolved_log != previous and resolved_log not in previous.parents + and previous not in resolved_log.parents for previous in seen_logs), + "Repeated history output/log root, nested attempt or request cycle") + seen_outputs.add(resolved_output); seen_logs.add(resolved_log) + start, parent = _start_and_parent(current) + need(start <= through, "History cutoff precedes its restored state") + relative = Path("csv") / f"step{start:09d}" / "metrics.csv" + direct, mirrored = log / relative, output / "log-mirror" / relative + path = direct if direct.is_file() else mirrored + rows.append(dict(request=current, start_step=start, through_step=through, + csv_path=path, output=output)) + if parent is None: + return rows[::-1] + current, through = parent, start + raise ValueError("Retained crash history is cyclic or exceeds32 attempts") + + +def collect_milestones(request, target=None): + """Collect exactly one immutable checkpoint per owned optimizer milestone. + + An unexpected newer parent checkpoint is not silently ignored. Its existence + requires the highest-checkpoint selection/inspection to be reconciled first. + ``last.ckpt`` is an alias, never a second retained milestone. + """ + result = {} + for segment in history_segments(request, target): + directory = segment["output"] / "checkpoints" + if not directory.is_dir(): + need(segment["start_step"] == segment["through_step"] and not directory.exists(), + "Missing retained checkpoint directory") + # A failed resumed process can end before its first new20k save. + # Such a zero-owned segment contributes no immutable milestone. + continue + for path in directory.glob("*.ckpt"): + if path.name == "last.ckpt": + need(path.is_symlink(), "Last checkpoint must be a symlink") + continue + match = re.search(r"([0-9]{9})\.ckpt$", path.name) + need(match is not None and "epoch" in path.name and "step" in path.name + and path.is_file() and not path.is_symlink(), + "Unexpected or mutable retained checkpoint") + step = int(match[1]) + need(step % CADENCE == 0 and segment["start_step"] < step + <= segment["through_step"], + "Unreconciled newer checkpoint or repeated resume-boundary save") + need(step not in result, "Duplicate immutable checkpoint across attempts") + result[step] = path + expected = set(range(CADENCE, (request["target_steps"] if target is None else target)+1, + CADENCE)) + need(set(result) == expected, "Missing retained checkpoint history") + return result + + +def _true_validation_step(raw): + number = float(raw) + need(math.isfinite(number) and number.is_integer() and number >= 0, + "Invalid validation CSV step") + number = int(number) + if number > 0 and number % CADENCE == 0: + return number + if (number + 1) % CADENCE == 0: + return number + 1 + raise ValueError("Validation does not occur at the20k optimizer cadence") + + +def collect_validation(request, target=None, streams=None): + """Merge full finite events and reject actual repeated boundary validation.""" + streams = request["active_streams"] if streams is None else streams + metric = {"image": "image_normalized_mse", "proprio": "proprio_normalized_mse", + "action": "action_command_mse"} + expected_keys = {f"val/libero_gripper/{kind}/{metric[name]}" + for name in streams for kind in ("forecast", "completion")} + result = {} + for segment in history_segments(request, target): + local, raw_steps = {}, {} + if not segment["csv_path"].is_file(): + need(segment["start_step"] == segment["through_step"] + and not segment["csv_path"].exists(), "Missing retained validation CSV") + # No scheduled optimizer milestone is owned by this failed attempt. + # Existing files, if any, are still scanned below for replayed events. + continue + with segment["csv_path"].open() as stream: + for row in csv.DictReader(stream): + values = {key: float(value) for key, value in row.items() + if key.startswith("val/libero_gripper/") and value} + if not values: + continue + step = _true_validation_step(row["step"]) + if step > segment["through_step"]: + # Preserve failed-attempt data beyond the actual restored + # checkpoint without incorporating it into completed history. + continue + need(step > segment["start_step"], + "Repeated validation at or before a resumed boundary") + need(all(math.isfinite(value) for value in values.values()) + and set(values) <= expected_keys, "Nonfinite or unexpected validation metric") + if step in raw_steps: + need(raw_steps[step] == row["step"], "Ambiguous validation step convention") + raw_steps[step] = row["step"] + event = local.setdefault(step, {}) + need(not (event.keys() & values.keys()), "Repeated validation metric/event") + event.update(values) + need(not (result.keys() & local.keys()), "Repeated validation across crash attempts") + need(all(set(event) == expected_keys for event in local.values()), + "Incomplete scheduled validation event") + result.update(local) + expected = set(range(CADENCE, (request["target_steps"] if target is None else target)+1, + CADENCE)) + need(set(result) == expected, "Missing finite scheduled validation history") + return result diff --git a/scripts/train/joint/crash_inspect.py b/scripts/train/joint/crash_inspect.py new file mode 100644 index 000000000..58854badf --- /dev/null +++ b/scripts/train/joint/crash_inspect.py @@ -0,0 +1,461 @@ +"""Inspect an actual highest production checkpoint on allocated compute only. + +This is read-only for source/data/checkpoints. The sole write is an exclusive +inspection receipt. No training, GPU allocation, scheduler mutation, RNG draw, +validation or checkpoint save is performed. A PASS verifies saved/restored CPU +state; the actual four-rank recovery smoke must still certify live restoration. +""" + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import random +import re +import subprocess +from collections.abc import Mapping + +from scripts.train.joint.contract import bound, need, sha, write_once + + +SCHEMA = "joint-crash-checkpoint-inspection-v1" +FROZEN_SOURCE = "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" +RUN_ID = "mot-shared-libero-image-discrete-lr1e-4-constant-s42-300k-20261021-v1" +MILESTONE = re.compile(r"epoch_epoch=(\d+)-step_step=(\d{9})\.ckpt") + + +def state_digest(value): + """Content digest preserving Python container/scalar and tensor identities.""" + import numpy as np + import torch + + digest = hashlib.sha256() + + def add(item): + if isinstance(item, torch.Tensor): + tensor = item.detach().cpu().contiguous() + need(tensor.layout == torch.strided, "Unsupported sparse state tensor") + digest.update(b"tensor\0" + str(tensor.dtype).encode() + b"\0") + digest.update(json.dumps(list(tensor.shape)).encode() + b"\0") + digest.update(tensor.numpy().tobytes()) + elif isinstance(item, np.ndarray): + digest.update(b"ndarray\0" + item.dtype.str.encode() + b"\0") + digest.update(json.dumps(list(item.shape)).encode() + b"\0") + digest.update(np.ascontiguousarray(item).tobytes()) + elif isinstance(item, dict): + digest.update(b"dict\0") + for key in sorted(item, key=lambda key: (type(key).__name__, repr(key))): + add(key) + add(item[key]) + digest.update(b"enddict\0") + elif isinstance(item, (list, tuple)): + digest.update(type(item).__name__.encode() + b"\0") + for child in item: + add(child) + digest.update(b"endsequence\0") + elif item is None or type(item) in (str, bool, int, float): + if type(item) is float: + need(math.isfinite(item), "Nonfinite scalar in saved state") + digest.update(type(item).__name__.encode() + b"\0") + digest.update(json.dumps(item, allow_nan=False).encode() + b"\0") + else: + raise ValueError("Unsupported state value: " + type(item).__name__) + + add(value) + return digest.hexdigest() + + +def verify_allocation(job_id, node, *, native_command="/usr/local/bin/scontrol", env=None, + uid=None, native_reader=None): + """Check native ownership before importing any ML library or reading a ckpt.""" + env = os.environ if env is None else env + uid = os.getuid() if uid is None else uid + need(str(job_id).isdigit() and re.fullmatch(r"gpu\d+", node), "Unsafe allocation identity") + need(uid == 41003 and env.get("SLURM_JOB_ID") == str(job_id) + and env.get("SLURMD_NODENAME") == node + and env.get("SLURM_RESTART_COUNT", "0") == "0", + "Inspector requires its exact fresh allocated compute job/node/owner") + if native_reader is None: + result = subprocess.run([native_command, "show", "job", str(job_id), "-o"], + text=True, capture_output=True, timeout=12, check=True) + text = result.stdout + else: + text = native_reader() + tokens = dict(field.split("=", 1) for field in text.split() if "=" in field) + need(tokens.get("JobId") == str(job_id) and tokens.get("JobState") == "RUNNING" + and tokens.get("UserId") == "paphiwetsa3(41003)" + and tokens.get("NodeList") == node and tokens.get("NumNodes") == "1" + and tokens.get("Restarts") == "0", "Native inspector allocation identity mismatch") + return dict(job_id=str(job_id), node=node, uid=uid, native_sha256=hashlib.sha256(text.encode()).hexdigest(), + native_state="RUNNING", native_command=native_command) + + +def retained_attempts(parent): + """Follow only hash-bound crash parents, assigning each its owned step range.""" + identity = ("run_id", "source_revision", "target_steps", "world_size", "global_batch", + "per_rank_batch", "active_streams", "arm", "parameters", "lr_policy") + attempts, seen = [], set() + current, through = parent, parent["target_steps"] + for _ in range(32): + output = Path(current["output"]) + need(output.is_absolute() and ".." not in output.parts and str(output) not in seen, + "Unsafe or repeated crash history output/cycle") + need(all(current.get(key) == parent.get(key) for key in identity), + "Crash ancestor has a foreign scientific identity") + seen.add(str(output)) + recovery = current.get("crash_recovery") + need(current.get("continuation") is None, + "Checkpoint inspector does not mix scheduled and crash ancestry") + if recovery is None: + attempts.append(dict(output=output, start_step=0, through_step=through, + resume_checkpoint=None)) + return attempts[::-1] + need(recovery.get("schema") == "joint-crash-recovery-v1", "Unknown crash parent binding") + prior = bound(recovery["parent_request"]) + receipt = bound(recovery["inspection"]) + need(receipt.get("schema") == SCHEMA and receipt.get("status") == "PASS" + and receipt.get("parent_request_sha256") == recovery["parent_request"]["sha256"] + and receipt.get("run_id") == parent["run_id"] + and receipt.get("source_revision") == parent["source_revision"], + "Crash history inspection belongs to another parent") + start = receipt.get("global_step") + need(type(start) is int and 0 < start <= through and start % 20000 == 0, + "Invalid crash history step range") + checkpoint = receipt.get("checkpoint", {}) + need(Path(checkpoint.get("path", "")).is_absolute() + and re.fullmatch(r"[a-f0-9]{64}", checkpoint.get("sha256", "")), + "Invalid ancestor checkpoint content binding") + attempts.append(dict(output=output, start_step=start, through_step=through, + resume_checkpoint=dict(checkpoint, global_step=start))) + current, through = prior, start + raise ValueError("Crash checkpoint history exceeds32 attempts") + + +def inventory_highest(directory, parent): + """Inspect bounded retained lineage; newest corruption never selects older state. + + A failed child may own no checkpoint before its next20k event. Its missing or + empty directory is valid only if prior retained history contains exactly its + restored milestone. No ancestor may introduce a file beyond its owned cutoff. + """ + directory = Path(directory) + need(directory.is_absolute() and directory == Path(parent["output"])/"checkpoints", + "Checkpoint inventory must be the exact failed parent's bounded directory") + rows, steps, indexed = [], set(), {} + attempts = retained_attempts(parent) + for attempt in attempts: + output, start, through = attempt["output"], attempt["start_step"], attempt["through_step"] + selected_parent = attempt["resume_checkpoint"] + if selected_parent is not None: + path = selected_parent["path"] + need(path in indexed and indexed[path]["global_step"] == start, + "Crash ancestor checkpoint is outside prior retained history") + known = indexed[path].get("expected_sha256") + need(known is None or known == selected_parent["sha256"], + "Conflicting historical checkpoint content bindings") + indexed[path]["expected_sha256"] = selected_parent["sha256"] + current = output/"checkpoints" + need(not output.is_symlink() and not current.is_symlink(), + "Symlink checkpoint history directory") + if not current.exists(): + need(start > 0 and steps and max(steps) == start, + "Missing checkpoint directory without a retained ancestor milestone") + continue + need(current.is_dir(), "Invalid checkpoint history directory") + local = [] + for path in current.iterdir(): + if path.name == "last.ckpt": + continue + need(path.name.endswith(".ckpt"), "Unexpected file in production checkpoint directory") + match = MILESTONE.fullmatch(path.name) + need(match and path.is_file() and not path.is_symlink(), "Invalid immutable checkpoint candidate") + step = int(match[2]) + need(start < step <= through and step < parent["target_steps"] + and step % 20000 == 0 and step not in steps, + "Invalid, duplicate, unscheduled or newer-than-cutoff production milestone") + steps.add(step) + status = path.stat() + row = dict(path=str(path), global_step=step, epoch=int(match[1]), + bytes=status.st_size, mtime_ns=status.st_mtime_ns, + ctime_ns=status.st_ctime_ns, inode=status.st_ino, device=status.st_dev, + owner_output=str(output), owner_start_step=start, owner_through_step=through) + rows.append(row); local.append(row); indexed[str(path)] = row + latest = current/"last.ckpt" + if not local: + need(start > 0 and steps and max(steps) == start and not latest.exists() + and not latest.is_symlink(), "Empty attempt has no retained ancestor or unexpected last alias") + continue + local_highest = max(local, key=lambda row: row["global_step"]) + need(latest.is_symlink() + and latest.resolve(strict=True) == Path(local_highest["path"]).resolve(strict=True), + "last.ckpt does not bind the highest immutable milestone") + need(rows, "No complete production checkpoint is available") + rows.sort(key=lambda row: row["global_step"]) + selected = rows[-1] + need(set(steps) == set(range(20000, selected["global_step"]+1, 20000)), + "Missing retained production milestone in crash history") + need(all(row["bytes"] > 0 for row in rows), + "Newest checkpoint is empty or retained checkpoint is empty; never fall back silently") + return Path(selected["path"]), rows + + +def verify_inventory_stable(directory, parent, selected, inventory): + """A new/changed milestone during slow CPU restoration invalidates selection.""" + current_selected, current_inventory = inventory_highest(directory, parent) + need(current_selected == selected and current_inventory == inventory, + "Checkpoint inventory changed during inspection; newest candidates are unresolved") + + +def validate_loop_state(checkpoint, step): + """Validate saved Lightning counters while refusing to predict replay behavior.""" + loops = checkpoint.get("loops") + need(isinstance(loops, dict) and isinstance(loops.get("fit_loop"), dict), + "Missing Lightning fit loop state") + try: + raw = json.dumps(loops, sort_keys=True, allow_nan=False, separators=(",", ":")) + except (TypeError, ValueError) as error: + raise ValueError("Malformed non-JSON Lightning loop state") from error + need(len(raw) <= 128_000, "Malformed or oversized Lightning loop state") + fit = loops["fit_loop"] + try: + optimizer = fit["epoch_loop.automatic_optimization.optim_progress"]["optimizer"]["step"]["total"] + batch = fit["epoch_loop.batch_progress"]["total"] + stepped = fit["epoch_loop.state_dict"]["_batches_that_stepped"] + except (KeyError, TypeError) as error: + raise ValueError("Unknown Lightning loop counter structure") from error + need(type(optimizer.get("completed")) is int and optimizer["completed"] == step + and type(optimizer.get("ready")) is int and optimizer["ready"] == step, + "Saved optimizer loop counter disagrees with checkpoint step") + need(type(stepped) is int and stepped in (step-1, step) + and all(type(batch.get(key)) is int for key in ("completed", "processed", "ready", "started")) + and batch["completed"] in (step-1, step) + and batch["processed"] == batch["ready"] == batch["started"] == step, + "Malformed Lightning training batch counters") + return dict(raw=loops, sha256=hashlib.sha256(raw.encode()).hexdigest(), + optimizer_steps_completed=optimizer["completed"], batches_that_stepped=stepped, + batch_completed=batch["completed"], + resume_loop_policy="reject_replayed_boundary_validation", + loop_training_resume_supported=False, + interpretation="Saved counters verified; actual recovery smoke must prove no boundary validation/save replay") + + +def validate_rng(rng, world): + import numpy as np + import torch + + need(isinstance(rng, dict) and rng.get("schema") == "mot_rng_state.distributed.v1" + and rng.get("world_size") == world and isinstance(rng.get("ranks"), list) + and len(rng["ranks"]) == world, "Missing exact distributed RNG payload") + digests = [] + for state in rng["ranks"]: + need(isinstance(state, dict) + and set(state) == {"torch_cpu", "torch_cuda", "numpy", "python", "cuda_scope"} + and state["cuda_scope"] == "current_device", "Malformed rank RNG schema") + cpu, cuda = state["torch_cpu"], state["torch_cuda"] + need(isinstance(cpu, torch.Tensor) and cpu.dtype == torch.uint8 and cpu.ndim == 1 + and cpu.numel() > 0 and isinstance(cuda, list) and len(cuda) == 1 + and isinstance(cuda[0], torch.Tensor) and cuda[0].dtype == torch.uint8 + and cuda[0].ndim == 1 and cuda[0].numel() > 0, "Malformed Torch rank RNG states") + # Validate using private CPU/NumPy/Python generators, preserving global RNG. + try: + torch.Generator(device="cpu").set_state(cpu) + np.random.RandomState().set_state(state["numpy"]) + random.Random().setstate(state["python"]) + except (RuntimeError, ValueError, TypeError) as error: + raise ValueError("Unrestorable saved CPU/NumPy/Python RNG") from error + digests.append(state_digest(state)) + return digests + + +def validate_payload(checkpoint, parent, resolved, *, expected_context_sha256, context_fingerprint): + """Structural/content inspection shared by the real CLI and real-tensor CPU tests.""" + import torch + + need(isinstance(checkpoint, dict), "Checkpoint must be a full-state mapping") + step = checkpoint.get("global_step") + need(type(step) is int and 0 < step < parent["target_steps"] and step % 20000 == 0, + "Not a recoverable scheduled production checkpoint") + need(type(checkpoint.get("epoch")) is int and checkpoint["epoch"] >= 0, "Invalid checkpoint epoch") + model_cfg = checkpoint.get("hyper_parameters", {}).get("config_tree") + from omegaconf import OmegaConf + if OmegaConf.is_config(model_cfg): + model_cfg = OmegaConf.to_container(model_cfg, resolve=True, throw_on_missing=True) + need(isinstance(model_cfg, Mapping) and model_cfg.get("model") == resolved["model"], + "Saved model configuration differs from parent") + provenance = model_cfg.get("run_provenance", {}) + expected = {key: parent[key] for key in ("source_revision", "source_manifest_sha256", "data_manifest_sha256", "run_id")} + expected["data_context_sha256"] = expected_context_sha256 + need(all(provenance.get(key) == value for key, value in expected.items()), + "Saved run/source/data provenance differs from parent") + need(isinstance(checkpoint.get("data_context"), dict) + and context_fingerprint(checkpoint["data_context"]) == expected_context_sha256 + and checkpoint.get("inference_binding", {}).get("data_context_sha256") == expected_context_sha256, + "Saved data context differs from its parent binding") + state = checkpoint.get("state_dict") + need(isinstance(state, dict) and state and all(isinstance(key, str) and isinstance(value, torch.Tensor) + and value.dtype == torch.float32 and value.layout == torch.strided + and bool(torch.isfinite(value).all()) for key, value in state.items()), + "Nonfinite, non-FP32 or invalid stored wrapper weights") + optimizers = checkpoint.get("optimizer_states") + need(isinstance(optimizers, list) and len(optimizers) == 1 + and isinstance(optimizers[0], dict) and optimizers[0].get("state") + and optimizers[0].get("param_groups"), "Missing full AdamW state") + need(parent["lr_policy"] == {"schedule": "constant", "peak_lr": 1e-4, "warmup_optimizer_steps": 0} + and resolved["model"]["scheduler"] is None + and checkpoint.get("lr_schedulers", []) == [], "Unexpected LR schedule") + optimizer = optimizers[0] + ids = [] + for group in optimizer["param_groups"]: + need(group.get("lr") == 1e-4 and group.get("weight_decay") == .01 + and isinstance(group.get("params"), list) and group["params"], "Wrong saved optimizer groups/LR") + ids.extend(group["params"]) + need(len(ids) == len(set(ids)) and set(ids) == set(optimizer["state"]), + "Missing, duplicated or unmatched optimizer parameter state") + for saved in optimizer["state"].values(): + count = saved.get("step") + if isinstance(count, torch.Tensor): + need(count.numel() == 1 and bool(torch.isfinite(count).all()), "Invalid optimizer step tensor") + count = float(count) + need(type(count) in (int, float) and count == step, "AdamW step disagrees with global step") + need(all(isinstance(saved.get(key), torch.Tensor) and saved[key].dtype == torch.float32 + and saved[key].layout == torch.strided and bool(torch.isfinite(saved[key]).all()) + for key in ("exp_avg", "exp_avg_sq")) + and saved["exp_avg"].shape == saved["exp_avg_sq"].shape + and bool((saved["exp_avg_sq"] >= 0).all()), "Invalid FP32 AdamW moments") + sampler = checkpoint.get("PixelProprioDataModule", {}) + topology = sampler.get("distributed_sampling", {}) + need(sampler.get("schema") == "mot-sampling-state-v1" and sampler.get("consumed_batches") == step + and topology.get("schema") == "pixel-distributed-sampling-v1" + and topology.get("world_size") == parent["world_size"] == 4 + and topology.get("global_batch_size") == parent["global_batch"] == 32 + and topology.get("per_rank_batch_size") == parent["per_rank_batch"] == 8 + and topology.get("training_sharding") == "global_draw_rank_slots" + and topology.get("validation_sharding") == "contiguous_global_batches_no_padding", + "Saved sampler cursor/topology differs from parent") + rank_rng = validate_rng(checkpoint.get("mot_rng_state"), parent["world_size"]) + loops = validate_loop_state(checkpoint, step) + return dict(global_step=step, epoch=checkpoint["epoch"], data_context_sha256=expected_context_sha256, + optimizer_state_sha256=state_digest(optimizer), sampler_state_sha256=state_digest(sampler), + wrapper_state_sha256=state_digest(state), rank_rng_sha256=rank_rng, + optimizer_steps=[step], sampling_consumed_batches=step, loop_state=loops) + + +def strict_cpu_restore(checkpoint, cfg, parent, expected_context_sha256): + """Actually construct/reload the frozen wrapper, optimizer and DataModule.""" + from hydra.utils import instantiate + from omegaconf import OmegaConf + from egomimic.pipeline.checkpoint_binding import validate_checkpoint_binding + from egomimic.pipeline.construction import checkpoint_construction + from egomimic.pl_utils.pl_model import ModelWrapper + + model_cfg = OmegaConf.create(checkpoint["hyper_parameters"]["config_tree"]) + data = instantiate(cfg.data) + context = data.prepare_context(mode="train", normalization="prepared", normalizer=None, + restored_state=checkpoint["data_context"]) + need(context.fingerprint() == expected_context_sha256, "Actual restored data context changed") + binding = validate_checkpoint_binding(checkpoint, model_cfg, context) + data.load_state_dict(checkpoint["PixelProprioDataModule"]) + need(state_digest(data.state_dict()) == state_digest(checkpoint["PixelProprioDataModule"]), + "Actual DataModule did not restore exact sampling state") + with checkpoint_construction(): + wrapper = ModelWrapper(config_tree=model_cfg, enable_grad_norm=False).cpu().eval() + context.bind(wrapper.model) + wrapper.data_context = context + wrapper.load_state_dict(checkpoint["state_dict"], strict=True) + need(state_digest(wrapper.state_dict()) == state_digest(checkpoint["state_dict"]), + "Strict wrapper reload differs from stored state") + network = wrapper.model.pipeline.stage_by_id("joint_denoiser").model + counts = network.expert_parameter_counts() + need(counts == parent["expert_parameters"] + and sum(parameter.numel() for parameter in network.parameters()) == parent["parameters"] + and all(parameter.requires_grad for parameter in network.parameters()), + "Reloaded expert count/trainability differs") + optimizer = instantiate(cfg.model.optimizer, params=wrapper.parameters()) + if callable(optimizer): + optimizer = optimizer() + optimizer.load_state_dict(checkpoint["optimizer_states"][0]) + need(state_digest(optimizer.state_dict()) == state_digest(checkpoint["optimizer_states"][0]), + "Actual AdamW groups/moments did not restore exactly") + need(instantiate(cfg.model.scheduler, optimizer=optimizer) is None, "Unexpected instantiated scheduler") + for parameter, saved in optimizer.state.items(): + need(saved["exp_avg"].shape == parameter.shape == saved["exp_avg_sq"].shape, + "Restored moment shape differs from parameter") + return dict(strict_full_wrapper_optimizer_reload=True, actual_datamodule_restore=True, + actual_expert_parameters=counts, inference_binding=dict(binding)) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--parent-request", required=True, type=Path) + parser.add_argument("--parent-request-sha256", required=True) + parser.add_argument("--resolved-config", required=True, type=Path) + parser.add_argument("--resolved-config-sha256", required=True) + parser.add_argument("--checkpoint-directory", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--expected-job-id", required=True) + parser.add_argument("--expected-node", required=True) + parser.add_argument("--scontrol", default="/usr/local/bin/scontrol") + args = parser.parse_args(argv) + allocation = verify_allocation(args.expected_job_id, args.expected_node, native_command=args.scontrol) + need(not args.output.exists(), "Refusing inspection receipt overwrite") + parent = bound(dict(path=str(args.parent_request), sha256=args.parent_request_sha256)) + need(parent["schema"] == "joint-training-v2" and parent["phase"] == "full" + and parent["source_revision"] == FROZEN_SOURCE and parent["run_id"] == RUN_ID + and parent["active_streams"] == ["image"] and parent["arm"] == "discrete", + "Wrong frozen production parent") + need(sha(args.resolved_config) == args.resolved_config_sha256, "Parent resolved configuration changed") + need(sha(parent["source_manifest"]) == parent["source_manifest_sha256"], "Source manifest changed") + for line in Path(parent["source_manifest"]).read_text().splitlines(): + expected, relative = line.split(" ", 1) + member = Path(relative) + need(not member.is_absolute() and ".." not in member.parts, "Unsafe source manifest member") + need(sha(Path(parent["source"])/member) == expected, "Frozen source changed: "+relative) + need(sha(parent["data_manifest"]) == parent["data_manifest_sha256"], "Data manifest changed") + selected, inventory = inventory_highest(args.checkpoint_directory, parent) + before = selected.stat() + identity = (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns, before.st_ctime_ns) + checkpoint_sha = sha(selected) + selected_inventory = next(row for row in inventory if row["path"] == str(selected)) + need(selected_inventory.get("expected_sha256", checkpoint_sha) == checkpoint_sha, + "Selected ancestor checkpoint changed from its prior inspection binding") + # All heavyweight imports and checkpoint reads occur after allocation checks. + import torch + from omegaconf import OmegaConf + from egomimic.pl_utils.data_context import state_fingerprint + + cfg = OmegaConf.load(args.resolved_config) + resolved = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + preflight = bound(parent["smoke"]["preflight"]) + expected_context = preflight["provenance"]["data_context_sha256"] + checkpoint = torch.load(selected, map_location="cpu", weights_only=False, mmap=True) + checked = validate_payload(checkpoint, parent, resolved, expected_context_sha256=expected_context, + context_fingerprint=state_fingerprint) + need(checked["global_step"] == inventory[-1]["global_step"], "Filename and saved optimizer step differ") + restored = strict_cpu_restore(checkpoint, cfg, parent, expected_context) + after = selected.stat() + need(identity == (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns, after.st_ctime_ns), + "Checkpoint changed during content inspection") + verify_inventory_stable(args.checkpoint_directory, parent, selected, inventory) + record = dict(schema=SCHEMA, status="PASS", parent_request_sha256=args.parent_request_sha256, + parent_resolved_sha256=args.resolved_config_sha256, job_id=args.expected_job_id, + source_revision=parent["source_revision"], run_id=parent["run_id"], + checkpoint=dict(path=str(selected), sha256=checkpoint_sha), + world_size=parent["world_size"], global_batch=parent["global_batch"], + context_sha256=expected_context, lr_policy=parent["lr_policy"], lr_schedulers=[], + sampler_all_rank_rng_verified=True, highest_complete_checkpoint_inspected=True, + all_newer_candidates_reconciled=True, + newer_candidate_resolution="Highest numeric file inspected; no older fallback allowed", + inventory=inventory, allocation=allocation, + resume_loop_policy="reject_replayed_boundary_validation", + loop_training_resume_supported=False, actual_gpu_rng_restoration_proved=False, + production_resume_authorized_by_this_receipt=False, + checkpoint_modified=False, **checked, **restored) + write_once(args.output, record) + print(json.dumps({key: record[key] for key in ("schema", "status", "global_step", "checkpoint", + "loop_training_resume_supported")})) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/crash_recovery.py b/scripts/train/joint/crash_recovery.py new file mode 100644 index 000000000..c5497e079 --- /dev/null +++ b/scripts/train/joint/crash_recovery.py @@ -0,0 +1,276 @@ +"""Explicit, evidence-bound recovery after failure; scheduled continuations stay separate. + +No submission lives here. A full request cannot enter this path until actual +allocated checkpoint inspection and matching recovery smoke have passed. +""" + +import os +import json +from pathlib import Path +import shutil +import subprocess + +from scripts.train.joint.contract import bound, need, sha + + +SOURCE = "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" +RUN_ID = "mot-shared-libero-image-discrete-lr1e-4-constant-s42-300k-20261021-v1" +TERMINAL_FAILURES = {"CANCELLED", "FAILED", "TIMEOUT", "NODE_FAIL", "PREEMPTED"} +STABLE = ( + "schema", "phase", "run_id", "arm", "source", "source_revision", + "source_manifest", "source_manifest_sha256", "source_archive", + "source_archive_sha256", "source_parent", "data_manifest", + "data_manifest_sha256", "data_cache", "lambda_runtime", "runtime_localization", + "allocation", "model_execution", "parameters", "expert_parameters", + "expert_widths", "world_size", "global_batch", "per_rank_batch", "lr_policy", + "vector_expert_parameter_budget", "active_streams", "experiment", "target_steps", + "smoke", "optimizer_comparison", "modality_comparison", "temporal_comparison", + "factorial_controls", "factorial_experiments", "archive_policy", "quota", + "reserved_checkpoint_bytes", +) + + +def parse_native(text): + return dict(part.split("=", 1) for part in text.split() if "=" in part) + + +def foundation_request(request): + """Original immutable fresh admission remains the comparison foundation.""" + run_id, source = request["run_id"], request["source_revision"] + seen = set() + for _ in range(32): + previous = request.get("crash_recovery") or request.get("continuation") + if previous is None: + return request + binding = previous["parent_request"] + need(binding["sha256"] not in seen, "Cyclic recovery admission foundation") + seen.add(binding["sha256"]) + request = bound(binding) + need(request["run_id"] == run_id and request["source_revision"] == source, + "Recovery foundation belongs to another scientific identity") + raise ValueError("Too many recovery admission ancestors") + + +def _separate(selected, previous, label): + a, b = Path(selected), Path(previous) + need(a.is_absolute() and b.is_absolute() and ".." not in a.parts, + "Unsafe " + label) + need(a != b and a not in b.parents and b not in a.parents, + "Recovery " + label + " overlaps retained history") + + +def admission(request): + """Read genuine failure/inspection/smoke evidence, without inventing boundary PASS.""" + recovery = request["crash_recovery"] + need(request.get("continuation") is None + and request.get("allocation_stop_step", request["target_steps"]) == request["target_steps"], + "Crash and scheduled-boundary recovery cannot be combined") + need(recovery["schema"] == "joint-crash-recovery-v1", "Wrong crash recovery binding") + parent = bound(recovery["parent_request"]) + need(parent["phase"] == request["phase"] == "full" + and request["source_revision"] == SOURCE + and request["run_id"] == RUN_ID and request["active_streams"] == ["image"] + and request["arm"] == "discrete", + "This crash service currently covers only the exact image-only1e-4 run") + need(all(parent.get(key) == request.get(key) for key in STABLE), + "Crash recovery changed scientific identity or runtime") + need(request["lr_policy"] == dict(schedule="constant", peak_lr=1e-4, + warmup_optimizer_steps=0) + and request["world_size"] == 4 and request["global_batch"] == 32 + and request["per_rank_batch"] == 8 and request["target_steps"] == 300000, + "Changed authorized constant-LR recovery contract") + _separate(request["output"], parent["output"], "output") + _separate(request["local_logs"], parent["local_logs"], "log namespace") + expected_overrides = [value.replace(parent["output"], request["output"]) + .replace(parent["local_logs"], request["local_logs"]) + for value in parent["overrides"]] + need(request["overrides"] == expected_overrides, + "Recovery overrides changed beyond output and log relocation") + failure = bound(recovery["failure"]) + need(failure["schema"] == "joint-crash-failure-v1" + and failure["parent_request_sha256"] == recovery["parent_request"]["sha256"] + and failure["run_id"] == request["run_id"] + and failure["classification"] in ("infrastructure", "worker_crash", "unexpected_scheduler_cancel") + and failure["recovery_authorized"] is True and failure["human_retired"] is False, + "Scientific failure or human retirement is not an automatic recovery") + native = parse_native(bound(failure["native"], False)) + need(native["JobId"] == str(failure["job_id"]) + and native["JobName"] == parent["job_name"] + and native["UserId"] == "paphiwetsa3(41003)" + and native["JobState"] in TERMINAL_FAILURES + and native["Requeue"] == "0" and native["Restarts"] == "0", + "Parent failure is not a verified terminal owned allocation") + need(bool(failure["logs"]), "Missing immutable failure logs") + for item in failure["logs"]: + need(bool(bound(item, False).strip()), "Empty failure log") + preflight = bound(failure["parent_preflight"]) + need(preflight["status"] == "FULL_ADMISSION_PASS" + and preflight["job_id"] == str(failure["job_id"]) + and preflight["request_sha256"] == recovery["parent_request"]["sha256"] + and preflight["provenance"]["run_id"] == request["run_id"] + and preflight["provenance"]["source_revision"] == SOURCE, + "Failed allocation was not admitted with this exact parent request") + inspection = bound(recovery["inspection"]) + need(inspection["schema"] == "joint-crash-checkpoint-inspection-v1" + and inspection["status"] == "PASS" + and inspection["parent_request_sha256"] == recovery["parent_request"]["sha256"] + and inspection["run_id"] == request["run_id"] + and inspection["source_revision"] == SOURCE + and inspection["world_size"] == 4 and inspection["global_batch"] == 32 + and inspection["strict_full_wrapper_optimizer_reload"] is True + and inspection["sampler_all_rank_rng_verified"] is True, + "Missing actual strict checkpoint inspection") + step = inspection["global_step"] + need(type(step) is int and 0 < step < request["target_steps"] and step % 20000 == 0, + "Recovery must use a healthy retained production milestone") + need(inspection["highest_complete_checkpoint_inspected"] is True + and inspection["all_newer_candidates_reconciled"] is True, + "Newer checkpoint candidates are unresolved") + checkpoint = inspection["checkpoint"] + need(Path(checkpoint["path"]).is_absolute(), "Absolute inspected checkpoint required") + from scripts.train.joint.crash_inspect import inventory_highest + + selected, inventory = inventory_highest(Path(parent["output"]) / "checkpoints", parent) + need(selected.resolve() == Path(checkpoint["path"]).resolve() + and any(Path(row["path"]).resolve() == selected.resolve() + and row["global_step"] == step for row in inventory), + "Inspection is not the current highest retained ancestor checkpoint") + need(sha(checkpoint["path"]) == checkpoint["sha256"], "Inspected checkpoint changed") + proof = bound(recovery["recovery_smoke"]) + need(proof["schema"] == "joint-crash-recovery-smoke-v1" and proof["status"] == "PASS" + and proof["source_revision"] == SOURCE + and proof["parent_request_sha256"] == recovery["parent_request"]["sha256"] + and proof["inspection_sha256"] == recovery["inspection"]["sha256"] + and proof["checkpoint_sha256"] == checkpoint["sha256"] + and proof["service_manifest_sha256"] == request["operational_service"]["manifest_sha256"] + and proof["world_size"] == 4 and proof["global_batch"] == 32 + and proof["all_rank_full_state_restored"] is True + and proof["finite_optimizer_updates"] >= 2 + and proof["scheduled_validation_anchors"] == 100 + and proof["no_replayed_boundary_validation_or_save"] is True + and proof["production_science_unchanged"] is True + and proof["sample_finite"] is True and proof["hidden_query_labels_nan"] is True + and proof["sample_steps"] == 20, + "Actual matching four-rank crash recovery smoke is missing") + diagnostic = bound(proof["request"]) + need(proof["request_sha256"] == proof["request"]["sha256"] + and diagnostic["schema"] == "joint-crash-recovery-diagnostic-v1" + and diagnostic["parent_request"] == recovery["parent_request"] + and diagnostic["inspection"] == recovery["inspection"] + and diagnostic["operational_service"] == request["operational_service"] + and all(diagnostic.get(key) == request.get(key) for key in ( + "source", "source_revision", "source_manifest_sha256", "data_manifest_sha256", + "lambda_runtime", "runtime_localization", "data_cache", "world_size", + "global_batch", "per_rank_batch", "active_streams", "lr_policy", "run_id")), + "Recovery smoke request is not bound to this exact checkpoint/science/service") + smoke_native = parse_native(bound(proof["native"], False)) + need(smoke_native["JobId"] == str(proof["job_id"]) + and smoke_native["JobState"] == "COMPLETED" and smoke_native["ExitCode"] == "0:0" + and smoke_native["JobName"] == proof["job_name"] == diagnostic["job_name"] + and smoke_native["UserId"] == "paphiwetsa3(41003)" + and smoke_native["NodeList"] == request["lambda_runtime"]["node"] + and smoke_native["Requeue"] == "0" and smoke_native["Restarts"] == "0" + and bound(proof["host_exit"], False).strip() == "0", + "Recovery smoke host/native did not finish successfully") + bound(proof["candidate"]) + diagnostic_output = Path(diagnostic["output"]) + expected = {str(diagnostic_output/"restore-proof"/f"restore-step{step:09d}-rank{rank}.json") + for rank in range(4)} + expected |= {str(diagnostic_output/"restore-proof"/f"restore-validation-rank{rank}.json") + for rank in range(4)} + expected |= {str(diagnostic_output/"validation-proof"/f"rank{rank}.json") for rank in range(4)} + need(len(proof["rank_proofs"]) == 12 + and {item["path"] for item in proof["rank_proofs"]} == expected, + "Recovery smoke is missing its actual rank proof inventory") + for item in proof["rank_proofs"]: + row = bound(item) + need(row["status"] == "PASS" and str(row["job_id"]) == str(proof["job_id"]) + and row["world_size"] == 4 and row["rank"] in range(4) + and row["start_step"] == step and row["global_step"] == step + 2, + "Recovery smoke rank proof changed") + final = proof["final_diagnostic_checkpoint"] + need(Path(final["path"]).parent == diagnostic_output/"checkpoints" + and final["strict_full_wrapper_optimizer_reload"] is True + and sha(final["path"]) == final["sha256"], + "Recovery diagnostic full-state checkpoint changed") + return parent, inspection, step + + +def assert_native_writer_quiescence(request): + """Repeat parent/writer checks in the new allocation immediately before training.""" + recovery = request["crash_recovery"] + failure = bound(recovery["failure"]) + current = os.environ.get("SLURM_JOB_ID", "") + need(current.isdecimal() and current != str(failure["job_id"]), + "Fresh allocated recovery identity required") + result = subprocess.run(["/usr/local/bin/scontrol", "show", "job", str(failure["job_id"]), "-o"], + capture_output=True, text=True, timeout=15) + if result.returncode == 0: + parent = parse_native(result.stdout) + need(parent.get("JobState") in TERMINAL_FAILURES, + "Parent allocation is still active or completing") + else: + need("Invalid job id" in result.stderr, + "Parent native state unavailable for an unexplained reason") + # Its hash-bound terminal native proof was validated by admission(). + result = subprocess.run(["/usr/local/bin/squeue", "-u", "paphiwetsa3", "-h", "-o", "%i|%j|%T"], + capture_output=True, text=True, timeout=15) + need(result.returncode == 0, "Current writer queue inspection failed") + prefix = {"image_only": "joint-image-lr1e4", + "image_proprio": "joint-" + request["arm"] + "-image-proprio-lr1e4"}.get(request["experiment"]) + need(prefix is not None and request["job_name"].startswith(prefix), + "Missing unique campaign writer name prefix") + for line in result.stdout.splitlines(): + job, name, state = line.split("|", 2) + need(job == current or not name.startswith(prefix), + "Another campaign writer is active/pending/completing: " + job) + + +def prove_destination_storage(request): + """One retained exclusive fsync/readback probe in the actual destination FS. + + Call only from the admitted allocation after prerequisite infrastructure + recovery. The allocation wrapper bounds startup; no retries or cleanup. + """ + job = os.environ.get("SLURM_JOB_ID", "") + need(job.isdecimal(), "Allocated storage proof required") + destination = Path(request["output"]) + parent = destination.parent + need(parent.is_dir() and not destination.exists(), + "Fresh recovery output parent missing or destination already exists") + free = shutil.disk_usage(parent).free + need(free > max(request["reserved_checkpoint_bytes"], 40_000_000_000), + "Insufficient actual destination checkpoint space") + payload = json.dumps(dict(schema="joint-crash-storage-probe-v1", job_id=job, + run_id=request["run_id"], output=str(destination), + free_bytes=free), sort_keys=True).encode() + path = parent / (".crash-recovery-storage-" + job + ".json") + with path.open("xb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + need(path.read_bytes() == payload, "Destination fsync/readback failed") + return dict(path=str(path), sha256=sha(path), free_bytes=free, + write_fsync_readback_pass=True) + + +def configure(cfg, request, inspection, start, *, request_sha256): + """Only operational config edits, after unchanged science has been compared.""" + from omegaconf import open_dict + need(isinstance(request_sha256, str) and len(request_sha256) == 64, + "Actual recovery request content identity required") + with open_dict(cfg): + cfg.ckpt_path = inspection["checkpoint"]["path"] + cfg.logger.wandb.resume = "must" + cfg.logger.csv.save_dir = request["local_logs"] + cfg.logger.csv.version = f"step{start:09d}" + cfg.callbacks.crash_restore_proof = dict( + _target_="scripts.train.joint.crash_restore.CrashRestoreProof", + checkpoint=inspection["checkpoint"]["path"], + inspection_path=request["crash_recovery"]["inspection"]["path"], + inspection_sha256=request["crash_recovery"]["inspection"]["sha256"], + output=str(Path(request["output"]) / "restore-proof"), + world_size=request["world_size"], start_step=start, proof_updates=2, + request_sha256=request_sha256, + ) + return cfg diff --git a/scripts/train/joint/crash_restore.py b/scripts/train/joint/crash_restore.py new file mode 100644 index 000000000..86622c95b --- /dev/null +++ b/scripts/train/joint/crash_restore.py @@ -0,0 +1,387 @@ +"""Observe real crash restoration and early updates without saving or stopping. + +Actual distributed/CUDA execution is mandatory; CPU tests do not certify the +production path. The scientific RNG restore function runs unchanged. A temporary +observer verifies its resulting state before the first training computation. +""" + +import functools +import hashlib +import importlib +import math +import os +from pathlib import Path + +import numpy as np +import torch +from lightning.pytorch.callbacks import Callback + +from scripts.train.joint.contract import bound, canonical_uuid, need, sha, write_once + + +def equal_state(actual, saved): + if isinstance(saved, torch.Tensor): + return (isinstance(actual, torch.Tensor) and actual.dtype == saved.dtype + and actual.shape == saved.shape + and torch.equal(actual.detach().cpu(), saved.detach().cpu())) + if isinstance(saved, np.ndarray): + return (isinstance(actual, np.ndarray) and actual.dtype == saved.dtype + and np.array_equal(actual, saved)) + if isinstance(saved, dict): + return (isinstance(actual, dict) and actual.keys() == saved.keys() + and all(equal_state(actual[key], value) for key, value in saved.items())) + if isinstance(saved, (list, tuple)): + return (type(actual) is type(saved) and len(actual) == len(saved) + and all(equal_state(a, b) for a, b in zip(actual, saved))) + return type(actual) is type(saved) and actual == saved + + +def weights_digest(state): + digest = hashlib.sha256() + for name, tensor in sorted(state.items()): + digest.update(name.encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(tensor.detach().cpu().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def verify_restored_state(trainer, module, saved, *, start_step, world_size): + """Strict loaded-object equality, not just constructor/config compatibility.""" + need(trainer.world_size == world_size == 4 and trainer.precision == "bf16-mixed" + and trainer.global_step == saved["global_step"] == start_step, + "Crash restore topology, precision or global step changed") + need(equal_state(module.state_dict(), saved["state_dict"]), "Weights were not exactly restored") + need(len(trainer.optimizers) == 1 and isinstance(trainer.optimizers[0], torch.optim.AdamW), + "Expected one actual restored AdamW optimizer") + optimizer = trainer.optimizers[0] + need(equal_state(optimizer.state_dict(), saved["optimizer_states"][0]), + "AdamW moments or parameter groups were not exactly restored") + need(saved.get("lr_schedulers", []) == [] and not trainer.lr_scheduler_configs, + "Constant1e-4 recovery cannot introduce a scheduler") + need(all(group["lr"] == 1e-4 and group["weight_decay"] == .01 + for group in optimizer.param_groups), "Restored AdamW LR or decay changed") + need(all(float(value["step"]) == start_step and all( + value[name].dtype == torch.float32 and torch.isfinite(value[name]).all() + for name in ("exp_avg", "exp_avg_sq")) + for value in optimizer.state.values()) and bool(optimizer.state), + "AdamW state is not finiteFP32 at the restored update") + need(equal_state(trainer.datamodule.state_dict(), saved["PixelProprioDataModule"]) + and trainer.datamodule.consumed_batches == start_step, + "Exact sampling cursor/context was not restored") + need(equal_state(module.training_behavior._pending_rng, saved["mot_rng_state"]), + "Native per-rank RNG restoration was not armed") + return optimizer + + +class CrashRestoreProof(Callback): + """Emit evidence after actual restored updates; do not add saves/validation.""" + + def __init__(self, *, checkpoint, inspection_path, inspection_sha256, output, + world_size, start_step, proof_updates=2, diagnostic=False, + validation_interval=20000, request_sha256=None): + need(world_size == 4 and type(start_step) is int and start_step > 0 + and start_step % 20000 == 0 and type(proof_updates) is int + and 2 <= proof_updates < 20000, "Invalid actual crash-restore proof range") + need(type(diagnostic) is bool and ( + (not diagnostic and validation_interval == 20000) + or (diagnostic and validation_interval == proof_updates == 2)), + "Only an explicit two-update diagnostic may use diagnostic validation cadence") + self.checkpoint, self.output = str(checkpoint), Path(output) + self.inspection_binding = dict(path=str(inspection_path), sha256=inspection_sha256) + self.world_size, self.start_step, self.proof_updates = world_size, start_step, proof_updates + self.diagnostic, self.validation_interval = diagnostic, validation_interval + self.request_sha256 = request_sha256 + self._rng_original = self._rng_observer = self._rng_module = None + self.precision_hook = None + self.done = False + self.observed_saves = [] + self.validation_steps = [] + + def on_train_start(self, trainer, pl_module): + need(torch.distributed.is_available() and torch.distributed.is_initialized() + and torch.distributed.get_world_size() == self.world_size + and torch.distributed.get_rank() == trainer.global_rank + and torch.cuda.is_available(), "Actual four-rank CUDA/DDP is required") + inspection = bound(self.inspection_binding) + need(inspection["schema"] == "joint-crash-checkpoint-inspection-v1" + and inspection["status"] == "PASS" + and inspection["global_step"] == self.start_step + and inspection["world_size"] == self.world_size + and inspection["checkpoint"]["path"] == self.checkpoint + and sha(self.checkpoint) == inspection["checkpoint"]["sha256"], + "Crash checkpoint differs from its exact allocated inspection") + saved = torch.load(self.checkpoint, map_location="cpu", weights_only=False, mmap=True) + self.optimizer = verify_restored_state( + trainer, pl_module, saved, start_step=self.start_step, world_size=self.world_size) + provenance = pl_module.hparams.config_tree.run_provenance + need(provenance.run_id == inspection["run_id"] + and provenance.source_revision == inspection["source_revision"] + and pl_module.data_context.fingerprint() == inspection["data_context_sha256"], + "Restored runtime source/run/data context differs from inspection") + self.network = pl_module.model.pipeline.stage_by_id("joint_denoiser").model + need(all(p.requires_grad and p.dtype == torch.float32 for p in self.network.parameters()), + "Every active parameter must remain trainableFP32") + self.initial_weights = weights_digest(pl_module.state_dict()) + gathered = [None] * self.world_size + torch.distributed.all_gather_object(gathered, self.initial_weights) + need(len(set(gathered)) == 1, "Initial restored weights differ across ranks") + self.initial_all_rank_weights = gathered + self.run_id, self.source_revision = inspection["run_id"], inspection["source_revision"] + self.checkpoint_sha256 = inspection["checkpoint"]["sha256"] + from scripts.clusters.pace.gpu_admission import cuda_uuid + self.gpu_uuid = cuda_uuid(torch.cuda.get_device_properties(pl_module.device)) + self.fp32_losses = self.bf16_outputs = self.gradients_checked = 0 + self.optimizer_update_lrs, self.validation_steps = [], [] + self.rng_restore_observed = False + self._install_rng_observer(pl_module, saved["mot_rng_state"]["ranks"][trainer.global_rank]) + del saved + self.output.mkdir(parents=True, exist_ok=True) + + def observe(module, args, outputs): + if not pl_module.training or not torch.is_grad_enabled() or self.done: + return + need(isinstance(outputs, dict) and set(outputs) == set(self.network.streams) + and all(value.dtype == torch.bfloat16 and torch.isfinite(value).all() + for value in outputs.values()), + "Actual resumed active-stream forward must be finiteBF16") + self.bf16_outputs += 1 + self.precision_hook = self.network.register_forward_hook(observe) + + def _install_rng_observer(self, module, expected): + implementation = importlib.import_module(type(module.training_behavior).__module__) + need(implementation.__name__ == "egomimic.pl_utils.world_model_behavior", + "Unexpected scientific RNG restoration implementation") + original = implementation._restore_rng_state + need(not getattr(original, "_crash_restore_observer", False), "Duplicate RNG restore observer") + self._rng_module, self._rng_original = implementation, original + + @functools.wraps(original) + def observe(state): + need(equal_state(state, expected), "Native RNG selected another saved rank") + original(state) + actual = implementation._capture_rng_state(current_cuda_only=True) + need(equal_state(actual, expected), "Actual native RNG state was not exactly restored") + self.rng_restore_observed = True + self._remove_rng_observer() + observe._crash_restore_observer = True + self._rng_observer = observe + implementation._restore_rng_state = observe + + def _remove_rng_observer(self): + if self._rng_module is not None: + need(self._rng_module._restore_rng_state is self._rng_observer, + "Scientific RNG restore function changed while observed") + self._rng_module._restore_rng_state = self._rng_original + self._rng_module = self._rng_original = self._rng_observer = None + + def on_before_backward(self, trainer, pl_module, loss): + if self.done: + return + need(self.rng_restore_observed and pl_module.training_behavior._pending_rng is None, + "Actual native per-rank RNG restore was not observed before training") + need(loss.dtype == torch.float32 and bool(torch.isfinite(loss)), + "Actual resumed MSE must be finiteFP32") + self.fp32_losses += 1 + + def on_before_optimizer_step(self, trainer, pl_module, optimizer): + if self.done: + return + need(optimizer is self.optimizer and all(group["lr"] == 1e-4 + and group["weight_decay"] == .01 for group in optimizer.param_groups), + "Actual resumed optimizer policy changed") + parameters = list(self.network.parameters()) + need(all(p.grad is not None and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() for p in parameters), + "Every resumed parameter needs a finiteFP32 gradient") + self.gradients_checked += 1 + self.optimizer_update_lrs.append(1e-4) + + def on_validation_start(self, trainer, pl_module): + need(not trainer.sanity_checking and trainer.global_step > self.start_step, + "Crash resume replayed boundary/sanity validation") + need(trainer.global_step % self.validation_interval == 0, + "Crash resume changed validation cadence") + need(trainer.global_step not in self.validation_steps, "Crash resume repeated validation") + self.validation_steps.append(trainer.global_step) + + def on_save_checkpoint(self, trainer, pl_module, checkpoint): + """Observe every native save, including ones before the early proof ends.""" + step = trainer.global_step + need(type(step) is int and type(checkpoint["global_step"]) is int + and checkpoint["global_step"] == step + and step > self.start_step, "Crash resume replayed a boundary checkpoint save") + expected = step == self.start_step + self.proof_updates if self.diagnostic else step % 20000 == 0 + need(expected, "Crash resume attempted an extra off-cadence checkpoint save") + need(step not in self.observed_saves, "Crash resume repeated a checkpoint save") + self.observed_saves.append(step) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + if self.done: + return + count = trainer.global_step - self.start_step + need(0 <= count <= self.proof_updates, "Skipped actual crash-restore proof boundary") + if count != self.proof_updates: + return + need(count == self.fp32_losses == self.bf16_outputs == self.gradients_checked + == len(self.optimizer_update_lrs) + and trainer.datamodule.consumed_batches == trainer.global_step + and all(float(value["step"]) == trainer.global_step + for value in self.optimizer.state.values()) + and not trainer.lr_scheduler_configs and not self.validation_steps, + "Missing actual optimizer/precision/sampler proof or repeated validation") + final = weights_digest(pl_module.state_dict()) + gathered = [None] * self.world_size + torch.distributed.all_gather_object(gathered, final) + need(len(set(gathered)) == 1 and final != self.initial_weights, + "Restored parameters did not update equally on all ranks") + record = dict( + schema="joint-crash-restore-rank-v1", status="PASS", + job_id=os.environ["SLURM_JOB_ID"], rank=trainer.global_rank, + run_id=self.run_id, source_revision=self.source_revision, + gpu_uuid=self.gpu_uuid, diagnostic=self.diagnostic, + request_sha256=self.request_sha256, + world_size=self.world_size, start_step=self.start_step, + global_step=trainer.global_step, actual_updates=count, + inspection_sha256=self.inspection_binding["sha256"], + checkpoint=self.checkpoint, checkpoint_sha256=self.checkpoint_sha256, + strict_weights_optimizer_sampler_rng=True, + actual_native_rank_rng_restored=True, fp32_losses=self.fp32_losses, + finite_bf16_outputs=self.bf16_outputs, finite_fp32_gradients=self.gradients_checked, + optimizer_update_lrs=self.optimizer_update_lrs, + initial_weights_sha256=self.initial_weights, initial_all_rank_weights=self.initial_all_rank_weights, + weights_sha256=final, all_rank_weights=gathered, + no_extra_save=True, no_replayed_validation=True, + save_cadence_guard_installed=True, observed_saves_through_proof=list(self.observed_saves), + ) + write_once(self.output / f"restore-step{self.start_step:09d}-rank{trainer.global_rank}.json", record) + self.done = True + self._remove_hooks() + + def _remove_hooks(self): + if self.precision_hook is not None: + self.precision_hook.remove() + self.precision_hook = None + self._remove_rng_observer() + + def on_exception(self, trainer, pl_module, exception): + self._remove_hooks() + + def on_train_end(self, trainer, pl_module): + self._remove_hooks() + need(self.done, "Training ended before actual crash restoration was proved") + if self.diagnostic: + expected = self.start_step + self.proof_updates + evaluator = pl_module.evaluator + metrics = {str(key): float(value) for key, value in trainer.callback_metrics.items() + if str(key).startswith("val/libero_gripper/")} + need(trainer.global_step == expected and self.validation_steps == [expected] + and self.observed_saves == [expected] + and evaluator.has_completed_validation(expected) + and evaluator.global_validation_batches == 100 + and evaluator._validation_batches == 100 // self.world_size + and len(metrics) == 2 * len(self.network.streams) + and all(math.isfinite(value) for value in metrics.values()), + "Actual recovery diagnostic did not finish all100 scheduled validation anchors") + write_once(self.output / f"restore-validation-rank{trainer.global_rank}.json", dict( + schema="joint-crash-restore-validation-rank-v1", status="PASS", + job_id=os.environ["SLURM_JOB_ID"], rank=trainer.global_rank, + world_size=self.world_size, start_step=self.start_step, + global_step=expected, validation_steps=self.validation_steps, + global_validation_batches=100, local_validation_batches=100 // self.world_size, + inspection_sha256=self.inspection_binding["sha256"], metrics=metrics, + observed_saves=self.observed_saves, + actual_scheduled_validation_after_optimizer=True, + )) + + +def validate_rank_proofs(output, *, inspection, inspection_sha256, start_step, + world_size=4, expected_job_id=None, diagnostic=False, + expected_request_sha256=None, expected_updates=2): + """Verify actual per-rank artifacts; synthetic tests never create runtime PASS.""" + import json + + rows = [] + directory = Path(output) + for rank in range(world_size): + path = directory / f"restore-step{start_step:09d}-rank{rank}.json" + need(path.is_file() and not path.is_symlink() and path.stat().st_size <= 2_000_000, + "Missing immutable actual rank restore proof") + row = json.loads(path.read_text()) + need(row["schema"] == "joint-crash-restore-rank-v1" and row["status"] == "PASS" + and row["rank"] == rank and row["world_size"] == world_size == 4 + and row["start_step"] == start_step + and row["global_step"] == start_step + row["actual_updates"] + and row["actual_updates"] == expected_updates + and row["inspection_sha256"] == inspection_sha256 + and row["checkpoint"] == inspection["checkpoint"]["path"] + and row["checkpoint_sha256"] == inspection["checkpoint"]["sha256"] + and row["run_id"] == inspection["run_id"] + and row["source_revision"] == inspection["source_revision"] + and row["diagnostic"] is diagnostic, + "Actual rank proof differs from the current inspected recovery") + count = row["actual_updates"] + need(row["strict_weights_optimizer_sampler_rng"] is True + and row["actual_native_rank_rng_restored"] is True + and row["fp32_losses"] == row["finite_bf16_outputs"] + == row["finite_fp32_gradients"] == count + and row["optimizer_update_lrs"] == [1e-4] * count + and row["no_extra_save"] is True and row["no_replayed_validation"] is True + and row["save_cadence_guard_installed"] is True + and (row["observed_saves_through_proof"] in ([], [start_step + expected_updates]) + if diagnostic else row["observed_saves_through_proof"] == []), + "Actual optimizer/precision/RNG proof is incomplete") + need(expected_job_id is None or str(row["job_id"]) == str(expected_job_id), + "Restore proof belongs to another allocation") + need(expected_request_sha256 is None + or row["request_sha256"] == expected_request_sha256, + "Restore proof belongs to another request") + rows.append(row) + need(len({str(row["job_id"]) for row in rows}) == 1 + and len({canonical_uuid(row["gpu_uuid"]) for row in rows}) == world_size, + "Restore ranks do not share one allocation with unique physical GPUs") + initial = [row["initial_weights_sha256"] for row in rows] + final = [row["weights_sha256"] for row in rows] + need(len(set(initial)) == len(set(final)) == 1 and initial[0] != final[0] + and all(row["initial_all_rank_weights"] == initial + and row["all_rank_weights"] == final for row in rows) + and len({row["actual_updates"] for row in rows}) == 1, + "Actual restored and updated all-rank weights disagree") + if diagnostic: + for rank, row in enumerate(rows): + value = json.loads((directory / f"restore-validation-rank{rank}.json").read_text()) + need(value["schema"] == "joint-crash-restore-validation-rank-v1" + and value["status"] == "PASS" and value["rank"] == rank + and value["world_size"] == world_size + and str(value["job_id"]) == str(row["job_id"]) + and value["start_step"] == start_step + and value["global_step"] == start_step + expected_updates + and value["validation_steps"] == [start_step + expected_updates] + and value["global_validation_batches"] == 100 + and value["local_validation_batches"] == 100 // world_size + and value["inspection_sha256"] == inspection_sha256 + and value["actual_scheduled_validation_after_optimizer"] is True + and value["observed_saves"] == [start_step + expected_updates] + and value["metrics"] and all(math.isfinite(float(v)) + for v in value["metrics"].values()), + "Missing actual all-rank scheduled recovery validation") + return rows + + +def assert_actual_restore(request, *, expected_job_id=None, expected_request_sha256=None, + diagnostic=False): + """Production postflight entry point used before full completion is written.""" + recovery = request["crash_recovery"] + inspection = bound(recovery["inspection"]) + need(inspection["schema"] == "joint-crash-checkpoint-inspection-v1" + and inspection["status"] == "PASS" + and inspection["run_id"] == request["run_id"] + and inspection["source_revision"] == request["source_revision"], + "Wrong crash inspection for actual restoration proof") + return validate_rank_proofs( + Path(request["output"]) / "restore-proof", inspection=inspection, + inspection_sha256=recovery["inspection"]["sha256"], + start_step=inspection["global_step"], world_size=request["world_size"], + expected_job_id=(os.environ.get("SLURM_JOB_ID") + if expected_job_id is None else expected_job_id), + diagnostic=diagnostic, expected_request_sha256=expected_request_sha256) diff --git a/scripts/train/joint/crash_smoke.py b/scripts/train/joint/crash_smoke.py new file mode 100644 index 000000000..8eeef3748 --- /dev/null +++ b/scripts/train/joint/crash_smoke.py @@ -0,0 +1,481 @@ +"""Allocated two-update crash-resume diagnostic and native-proof finalizer. + +This driver never submits work and never writes the production run. Its short +save/validation cadence belongs only to a fresh, W&B-free diagnostic directory. +PASS is withheld until separate host and terminal native evidence is bound. +""" + +import argparse +import copy +import csv +import json +import math +import os +from pathlib import Path +import shutil +import subprocess +import sys + +from scripts.train.joint.contract import bound, need, sha, write_once + + +SCHEMA = "joint-crash-recovery-diagnostic-v1" +SOURCE = "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" + + +def separate(path, retained): + path, retained = Path(path), Path(retained) + need(path.is_absolute() and ".." not in path.parts + and path != retained and path not in retained.parents + and retained not in path.parents, "Diagnostic path overlaps retained production") + + +def validate_request(q): + need(q["schema"] == SCHEMA, "Wrong crash diagnostic request") + parent, inspection = bound(q["parent_request"]), bound(q["inspection"]) + need(parent["schema"] == "joint-training-v2" and parent["phase"] == "full" + and parent["source_revision"] == SOURCE and parent["active_streams"] == ["image"] + and parent["arm"] == "discrete" and parent["parameters"] == 88178688 + and parent["world_size"] == 4 and parent["global_batch"] == 32 + and parent["per_rank_batch"] == 8 + and parent["lr_policy"] == dict(schedule="constant", peak_lr=1e-4, + warmup_optimizer_steps=0), + "Diagnostic differs from the exact image-only1e-4 production parent") + need(inspection["schema"] == "joint-crash-checkpoint-inspection-v1" + and inspection["status"] == "PASS" + and inspection["parent_request_sha256"] == q["parent_request"]["sha256"] + and inspection["parent_resolved_sha256"] == q["resolved"]["sha256"] + and inspection["source_revision"] == SOURCE + and inspection["run_id"] == parent["run_id"] + and inspection["world_size"] == 4 and inspection["global_batch"] == 32 + and inspection["strict_full_wrapper_optimizer_reload"] is True + and inspection["sampler_all_rank_rng_verified"] is True, + "Missing exact allocated production checkpoint inspection") + start = inspection["global_step"] + need(type(start) is int and start > 0 and start % 20000 == 0 + and start + 2 < parent["target_steps"], "Invalid diagnostic recovery start") + for key in ("source", "source_revision", "source_manifest", "source_manifest_sha256", + "source_archive", "source_archive_sha256", "data_manifest", + "data_manifest_sha256", "data_cache", "lambda_runtime", + "runtime_localization", "allocation", "world_size", "global_batch", + "per_rank_batch", "active_streams", "lr_policy", "run_id", + "parameters", "expert_parameters", "expert_widths"): + need(q.get(key) == parent.get(key), "Diagnostic changed bound infrastructure/science: " + key) + need(q.get("operational_service") is not None, "Hash-bound operational service required") + separate(q["output"], parent["output"]) + separate(q["local_logs"], parent["local_logs"]) + need(q["reserved_checkpoint_bytes"] >= 4_000_000_000, "Diagnostic checkpoint reserve too small") + need(q["operational_service"]["manifest_sha256"] + == sha(q["operational_service"]["manifest"]), "Changed diagnostic service manifest") + need(sha(inspection["checkpoint"]["path"]) == inspection["checkpoint"]["sha256"], + "Inspected production checkpoint changed") + bound(q["resolved"], False) + return parent, inspection + + +def assert_parent_terminal(q, parent): + """A diagnostic must not race an active production writer of this identity.""" + failure = bound(q["failure"]) + need(failure["schema"] == "joint-crash-failure-v1" + and failure["parent_request_sha256"] == q["parent_request"]["sha256"] + and failure["run_id"] == parent["run_id"] + and failure["recovery_authorized"] is True and failure["human_retired"] is False, + "Missing honest diagnostic parent failure/authorization") + native = dict(item.split("=", 1) for item in bound(failure["native"], False).split() if "=" in item) + terminal = {"CANCELLED", "FAILED", "PREEMPTED", "TIMEOUT", "NODE_FAIL"} + need(native["JobId"] == str(failure["job_id"]) and native["JobName"] == parent["job_name"] + and native["UserId"] == "paphiwetsa3(41003)" and native["JobState"] in terminal, + "Diagnostic parent is not a terminal owned failure") + current = subprocess.run(["/usr/local/bin/scontrol", "show", "job", str(failure["job_id"]), "-o"], + capture_output=True, text=True, timeout=15) + if current.returncode == 0: + state = dict(item.split("=", 1) for item in current.stdout.split() if "=" in item) + need(state.get("JobState") in terminal, "Parent is still active or completing") + else: + need("Invalid job id" in current.stderr, "Parent scheduler read failed") + queue = subprocess.run(["/usr/local/bin/squeue", "-u", "paphiwetsa3", "-h", "-o", "%i|%j|%T"], + capture_output=True, text=True, timeout=15) + need(queue.returncode == 0, "Diagnostic duplicate writer inspection failed") + for row in queue.stdout.splitlines(): + job, name, _ = row.split("|", 2) + need(job == os.environ["SLURM_JOB_ID"] or not name.startswith("joint-image-lr1e4"), + "Another image-only1e-4 campaign writer is active/pending/completing") + + +def phase_config(original, q, parent, inspection, *, request_sha256): + """The model/data remain exact; only explicit diagnostic operations change.""" + from omegaconf import OmegaConf, open_dict + + cfg = OmegaConf.create(copy.deepcopy(OmegaConf.to_container(original, resolve=True))) + original_values = OmegaConf.to_container(cfg, resolve=True) + preserved = {key: original_values[key] for key in + ("model", "data", "seed", "runtime", "joint", "world_model_attention") + if key in cfg} + start, output = inspection["global_step"], str(q["output"]) + with open_dict(cfg): + cfg.ckpt_path = inspection["checkpoint"]["path"] + cfg.paths.output_dir = output + cfg.val_at_start = cfg.val_at_end = False + cfg.trainer.max_steps = start + 2 + cfg.trainer.val_check_interval = 2 + cfg.trainer.check_val_every_n_epoch = None + cfg.trainer.num_sanity_val_steps = 0 + cfg.logger.pop("wandb", None) + need(set(cfg.logger) == {"csv"}, "Diagnostic must use only local CSV logging") + cfg.logger.csv.save_dir = q["local_logs"] + cfg.logger.csv.name = "csv" + cfg.logger.csv.version = f"step{start:09d}" + cfg.logger.csv.flush_logs_every_n_steps = 1 + cfg.evaluator.pixel_diagnostics.log_wandb = False + cfg.evaluator.pixel_diagnostics.artifact_dir = output + "/validation-artifacts" + cfg.callbacks.pop("allocation_boundary", None) + cfg.callbacks.pop("joint_smoke", None) + cb = cfg.callbacks.model_checkpoint + cb.dirpath = output + "/checkpoints" + cb.filename = "diagnostic-epoch{epoch:04d}-step{step:09d}" + cb.every_n_train_steps, cb.every_n_epochs = 2, 0 + cb.save_top_k, cb.save_last, cb.save_on_train_epoch_end = -1, "link", False + cb.train_time_interval = None + cfg.callbacks.crash_restore_proof = dict( + _target_="scripts.train.joint.crash_restore.CrashRestoreProof", + checkpoint=inspection["checkpoint"]["path"], + inspection_path=q["inspection"]["path"], + inspection_sha256=q["inspection"]["sha256"], output=output + "/restore-proof", + world_size=4, start_step=start, proof_updates=2, + diagnostic=True, validation_interval=2, request_sha256=request_sha256) + cfg.callbacks.crash_validation_proof = dict( + _target_="scripts.train.joint.crash_smoke.diagnostic_validation_callback", + output=output + "/validation-proof", start_step=start, + stop_step=start + 2, world_size=4) + configured_values = OmegaConf.to_container(cfg, resolve=True) + for key, value in preserved.items(): + need(configured_values[key] == value, + "Diagnostic changed scientific configuration: " + key) + need(cfg.trainer.devices == cfg.data.distributed_world_size == 4 + and cfg.data.batch_size == 8 and cfg.trainer.precision == "bf16-mixed" + and cfg.model.optimizer.lr == 1e-4 and cfg.model.optimizer.weight_decay == .01 + and cfg.model.scheduler is None, "Diagnostic topology/optimizer changed") + return cfg + + +def diagnostic_validation_callback(*, output, start_step, stop_step, world_size): + """Hydra factory keeps ML imports behind actual allocated runtime admission.""" + import torch + from lightning.pytorch.callbacks import Callback + + class DiagnosticValidationProof(Callback): + def __init__(self): + need(world_size == 4 and stop_step == start_step + 2, + "Exactly two resumed optimizer updates required") + self.starts = [] + + def on_validation_start(self, trainer, pl_module): + need(not trainer.sanity_checking and trainer.global_step == stop_step, + "Diagnostic repeated boundary validation or ran an extra event") + self.starts.append(trainer.global_step) + + def on_train_end(self, trainer, pl_module): + need(trainer.world_size == world_size and trainer.global_step == stop_step + and self.starts == [stop_step] + and trainer.datamodule.consumed_batches == stop_step + and pl_module.evaluator.has_completed_validation(stop_step) + and pl_module.evaluator.global_validation_batches == 100, + "Diagnostic did not consume exactly100 scheduled validation anchors") + metrics = {str(k): float(v) for k, v in trainer.callback_metrics.items() + if str(k).startswith("val/libero_gripper/")} + need(set(metrics) == {"val/libero_gripper/forecast/image_normalized_mse", + "val/libero_gripper/completion/image_normalized_mse"} + and all(math.isfinite(v) for v in metrics.values()), + "Missing finite actual image validation metrics") + need(torch.distributed.is_initialized() + and torch.distributed.get_world_size() == world_size, + "Actual four-rank validation required") + directory = Path(output) + directory.mkdir(parents=True, exist_ok=True) + write_once(directory / f"rank{trainer.global_rank}.json", dict( + schema="joint-crash-diagnostic-validation-rank-v1", status="PASS", + job_id=os.environ["SLURM_JOB_ID"], rank=trainer.global_rank, + world_size=world_size, start_step=start_step, global_step=stop_step, + scheduled_validation_steps=self.starts, scheduled_validation_anchors=100, + local_validation_anchors=pl_module.evaluator._validation_batches, + metrics=metrics, no_replayed_boundary_validation=True)) + + return DiagnosticValidationProof() + + +def verify_validation(q, inspection, job_id): + output, stop = Path(q["output"]), inspection["global_step"] + 2 + rows = [] + for rank in range(4): + record = json.loads((output / "validation-proof" / f"rank{rank}.json").read_text()) + need(record["schema"] == "joint-crash-diagnostic-validation-rank-v1" + and record["status"] == "PASS" and record["job_id"] == job_id + and record["rank"] == rank and record["world_size"] == 4 + and record["global_step"] == stop + and record["scheduled_validation_steps"] == [stop] + and record["scheduled_validation_anchors"] == 100 + and record["local_validation_anchors"] == 25 + and record["no_replayed_boundary_validation"] is True, + "Missing actual per-rank scheduled validation proof") + rows.append(record) + path = Path(q["local_logs"]) / "csv" / f"step{inspection['global_step']:09d}" / "metrics.csv" + events = {} + with path.open() as stream: + for row in csv.DictReader(stream): + metrics = {key: float(value) for key, value in row.items() + if key.startswith("val/libero_gripper/") and value} + if metrics: + raw_step = float(row["step"]) + need(raw_step.is_integer() and all(math.isfinite(v) for v in metrics.values()), + "Invalid diagnostic validation history") + events.setdefault(int(raw_step), {}).update(metrics) + need(len(events) == 1 and next(iter(events)) in (stop-1, stop) + and len(next(iter(events.values()))) == 2, + "Diagnostic CSV contains missing or replayed validation events") + return rows + + +def verify_final_checkpoint(q, parent, inspection, cfg): + """Strictly reload the diagnostic output, preserving every production byte.""" + import torch + from scripts.train.joint.crash_inspect import strict_cpu_restore + + directory = Path(q["output"]) / "checkpoints" + files = [path for path in directory.iterdir() if path.name != "last.ckpt"] + stop = inspection["global_step"] + 2 + need(len(files) == 1 and files[0].is_file() and not files[0].is_symlink() + and files[0].name.endswith(f"step={stop:09d}.ckpt"), + "Diagnostic checkpoint cadence repeated or skipped a save") + checkpoint_path = files[0] + need((directory / "last.ckpt").is_symlink() + and (directory / "last.ckpt").resolve() == checkpoint_path.resolve(), + "Diagnostic last alias changed") + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False, mmap=True) + need(checkpoint["global_step"] == stop and checkpoint.get("lr_schedulers", []) == [] + and all(t.dtype == torch.float32 and torch.isfinite(t).all() + for t in checkpoint["state_dict"].values()) + and checkpoint["PixelProprioDataModule"]["consumed_batches"] == stop, + "Invalid final diagnostic full state") + states = checkpoint["optimizer_states"] + need(len(states) == 1 and states[0]["state"] + and all(group["lr"] == 1e-4 and group["weight_decay"] == .01 + for group in states[0]["param_groups"]) + and all(float(value["step"]) == stop and all( + value[key].dtype == torch.float32 and torch.isfinite(value[key]).all() + for key in ("exp_avg", "exp_avg_sq")) for value in states[0]["state"].values()), + "Diagnostic optimizer did not continue in finite FP32") + rng = checkpoint["mot_rng_state"] + need(rng["schema"] == "mot_rng_state.distributed.v1" and rng["world_size"] == 4 + and len(rng["ranks"]) == 4, "Diagnostic lost all-rank RNG state") + restored = strict_cpu_restore(checkpoint, cfg, parent, inspection["data_context_sha256"]) + from hydra.utils import instantiate + from omegaconf import OmegaConf + from egomimic.pipeline.construction import checkpoint_construction + from egomimic.pipeline.checkpoint_binding import validate_checkpoint_binding + from egomimic.pl_utils.pl_model import ModelWrapper + data = instantiate(cfg.data) + context = data.prepare_context(mode="eval", normalization="prepared", normalizer=None, + restored_state=checkpoint["data_context"]) + model_cfg = OmegaConf.create(checkpoint["hyper_parameters"]["config_tree"]) + validate_checkpoint_binding(checkpoint, model_cfg, context) + with checkpoint_construction(): + wrapper = ModelWrapper(config_tree=model_cfg, enable_grad_norm=False).cuda().eval() + context.bind(wrapper.model) + wrapper.data_context = context + wrapper.load_state_dict(checkpoint["state_dict"], strict=True) + sample = next(iter(data.val_dataloader()))["libero_gripper"] + for stream in sample["world"].values(): + stream["value"] = torch.where(stream["observed"][..., None], stream["value"], float("nan")) + groups = wrapper.model.process_batch_for_training({"libero_gripper": sample}) + wrapper.model.device = next(wrapper.parameters()).device + torch.manual_seed(42) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + predictions = wrapper.model.forward_eval(groups)["libero_gripper"]["image_proprio_prediction"] + need(set(predictions) == {"image"}, "Diagnostic sampler changed active streams") + for name, value in predictions.items(): + stream = groups["libero_gripper"]["world"][name] + need(value.dtype == torch.float32 and value.shape == stream["value"].shape + and torch.isfinite(value).all() + and torch.equal(value[stream["observed"]], stream["value"][stream["observed"]]), + "Restored diagnostic hidden-label-NaN sample failed") + return dict(path=str(checkpoint_path), sha256=sha(checkpoint_path), + strict_full_wrapper_optimizer_reload=restored["strict_full_wrapper_optimizer_reload"], + sample_finite=True, hidden_query_labels_nan=True, sample_steps=20) + + +def run(request_path, q): + from scripts.train.joint.crash_inspect import verify_allocation + + verify_allocation(os.environ.get("SLURM_JOB_ID", ""), q["lambda_runtime"]["node"]) + parent, inspection = validate_request(q) + assert_parent_terminal(q, parent) + from scripts.train.joint.contract import verify_evidence + verify_evidence(parent) + for line in Path(parent["source_manifest"]).read_text().splitlines(): + expected, relative = line.split(" ", 1) + need(sha(Path(parent["source"]) / relative) == expected, + "Frozen diagnostic source changed: " + relative) + need(sha(parent["source_manifest"]) == parent["source_manifest_sha256"] + and sha(parent["data_manifest"]) == parent["data_manifest_sha256"], + "Frozen diagnostic source/data manifest changed") + from scripts.train.joint.lambda_runtime import filesystem_recovery_reserve + filesystem_recovery_reserve(q) + output = Path(q["output"]) + need(not output.exists() and shutil.disk_usage(output.parent).free > q["reserved_checkpoint_bytes"], + "Diagnostic output collision or insufficient space") + root = request_path.parent.parent + from scripts.train.joint.allocation import admit + allocation = admit(q, status_root=root / "status", bootstrap=root / "runtime_bootstrap.py") + if allocation is None: + from gpu_health_set import admit_gpu_set + admit_gpu_set(root / "status" / f"crash-diagnostic-{os.environ['SLURM_JOB_ID']}-gpu.json", 4) + from omegaconf import OmegaConf + original = OmegaConf.load(q["resolved"]["path"]) + cfg = phase_config(original, q, parent, inspection, request_sha256=sha(request_path)) + output.mkdir() + probe = output / "diagnostic-storage-probe.json" + payload = json.dumps(dict(job_id=os.environ["SLURM_JOB_ID"], request_sha256=sha(request_path))).encode() + with probe.open("xb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + need(probe.read_bytes() == payload, "Diagnostic persistent fsync/readback failed") + directory = output / "config" + directory.mkdir() + config_path = directory / "diagnostic.yaml" + with config_path.open("x") as stream: + stream.write(OmegaConf.to_yaml(cfg, resolve=True)) + from scripts.train.joint.service_bootstrap import command + bootstrap = root / "runtime_bootstrap.py" + subprocess.run([ + sys.executable, str(bootstrap), "-m", "torch.distributed.run", "--standalone", + "--nnodes=1", "--nproc-per-node=4", "--max-restarts=0", + *command(q, request_path, bootstrap, ["-m", "egomimic.trainHydra"])[1:], + "--config-path", str(directory), "--config-name", "diagnostic", + "hydra.run.dir=" + str(output / "hydra"), "hydra.job.chdir=false", + "+run_provenance={" + ",".join(key + ":" + json.dumps(value) + for key, value in dict(source_revision=parent["source_revision"], + source_manifest_sha256=parent["source_manifest_sha256"], + data_manifest_sha256=parent["data_manifest_sha256"], + data_context_sha256=inspection["data_context_sha256"], run_id=parent["run_id"]).items()) + "}", + ], check=True, cwd=parent["source"]) + from scripts.train.joint.crash_restore import assert_actual_restore + execution = dict(q, crash_recovery=dict(inspection=q["inspection"])) + assert_actual_restore(execution, diagnostic=True, expected_request_sha256=sha(request_path)) + validation = verify_validation(q, inspection, os.environ["SLURM_JOB_ID"]) + final = verify_final_checkpoint(q, parent, inspection, original) + proof_paths = sorted((output / "restore-proof").glob("*.json")) + proof_paths += sorted((output / "validation-proof").glob("*.json")) + need(len(proof_paths) == 12, "Expected all startup/update and diagnostic validation rank proofs") + write_once(output / "crash-recovery-smoke.candidate.json", dict( + schema="joint-crash-recovery-smoke-candidate-v1", status="ACTUAL_DIAGNOSTIC_PASS_NATIVE_PENDING", + request_sha256=sha(request_path), job_id=os.environ["SLURM_JOB_ID"], + request=dict(path=str(request_path), sha256=sha(request_path)), job_name=q["job_name"], + source_revision=SOURCE, parent_request_sha256=q["parent_request"]["sha256"], + inspection_sha256=q["inspection"]["sha256"], checkpoint_sha256=inspection["checkpoint"]["sha256"], + service_manifest_sha256=q["operational_service"]["manifest_sha256"], + world_size=4, global_batch=32, all_rank_full_state_restored=True, + finite_optimizer_updates=2, scheduled_validation_anchors=100, + no_replayed_boundary_validation_or_save=True, production_science_unchanged=True, + final_diagnostic_checkpoint=final, validation_ranks=validation, + sample_finite=final["sample_finite"], hidden_query_labels_nan=True, sample_steps=20, + rank_proofs=[dict(path=str(path), sha256=sha(path)) for path in proof_paths], + parent_checkpoint_modified=False, production_wandb_written=False)) + + +def finalize(request_path, q, *, native_binding, host_binding, job_id): + """Only a terminal successful native job and host zero can produce PASS.""" + from scripts.train.joint.crash_inspect import verify_allocation + # Finalization may rehash a full checkpoint, so it too must run on compute. + verify_allocation(os.environ.get("SLURM_JOB_ID", ""), q["lambda_runtime"]["node"]) + need(q["schema"] == SCHEMA and str(job_id).isdigit(), "Wrong diagnostic finalization") + parent, inspection = validate_request(q) + candidate_path = Path(q["output"]) / "crash-recovery-smoke.candidate.json" + candidate = json.loads(candidate_path.read_text()) + native = dict(item.split("=", 1) for item in bound(native_binding, False).split() if "=" in item) + need(candidate["schema"] == "joint-crash-recovery-smoke-candidate-v1" + and candidate["status"] == "ACTUAL_DIAGNOSTIC_PASS_NATIVE_PENDING" + and candidate["request_sha256"] == sha(request_path) + and candidate["job_id"] == str(job_id) + and native["JobId"] == str(job_id) and native["JobState"] == "COMPLETED" + and native["ExitCode"] == "0:0" and native["UserId"] == "paphiwetsa3(41003)" + and native["NodeList"] == q["lambda_runtime"]["node"] + and native["JobName"] == q["job_name"] + and native["Requeue"] == "0" and native["Restarts"] == "0" + and bound(host_binding, False).strip() == "0", + "Diagnostic host/native success is not proved") + need(candidate["source_revision"] == parent["source_revision"] + and candidate["parent_request_sha256"] == q["parent_request"]["sha256"] + and candidate["inspection_sha256"] == q["inspection"]["sha256"] + and candidate["checkpoint_sha256"] == inspection["checkpoint"]["sha256"] + and candidate["service_manifest_sha256"] == q["operational_service"]["manifest_sha256"] + and candidate["world_size"] == 4 and candidate["global_batch"] == 32 + and candidate["all_rank_full_state_restored"] is True + and candidate["finite_optimizer_updates"] == 2 + and candidate["scheduled_validation_anchors"] == 100 + and candidate["no_replayed_boundary_validation_or_save"] is True + and candidate["production_science_unchanged"] is True + and candidate["sample_finite"] is True and candidate["hidden_query_labels_nan"] is True + and candidate["sample_steps"] == 20 + and candidate["parent_checkpoint_modified"] is False + and candidate["production_wandb_written"] is False, + "Actual diagnostic candidate is incomplete or bound to other inputs") + output, start = Path(q["output"]), inspection["global_step"] + expected_proofs = {str(output / "restore-proof" / f"restore-step{start:09d}-rank{rank}.json") + for rank in range(4)} + expected_proofs |= {str(output / "restore-proof" / f"restore-validation-rank{rank}.json") + for rank in range(4)} + expected_proofs |= {str(output / "validation-proof" / f"rank{rank}.json") for rank in range(4)} + need(len(candidate["rank_proofs"]) == 12 + and {item["path"] for item in candidate["rank_proofs"]} == expected_proofs, + "Missing exact diagnostic rank proof inventory") + for item in candidate["rank_proofs"]: + row = bound(item) + need(row["status"] == "PASS" and str(row["job_id"]) == str(job_id) + and row["world_size"] == 4 and row["rank"] in range(4) + and row["start_step"] == start and row["global_step"] == start + 2, + "Diagnostic rank proof changed or belongs to another allocation") + final = candidate["final_diagnostic_checkpoint"] + need(Path(final["path"]).parent == output / "checkpoints" + and final["strict_full_wrapper_optimizer_reload"] is True + and sha(final["path"]) == final["sha256"], "Diagnostic final full checkpoint changed") + from scripts.train.joint.crash_restore import assert_actual_restore + execution = dict(q, crash_recovery=dict(inspection=q["inspection"])) + assert_actual_restore(execution, expected_job_id=str(job_id), diagnostic=True, + expected_request_sha256=sha(request_path)) + verify_validation(q, inspection, str(job_id)) + record = {key: value for key, value in candidate.items() + if key not in ("schema", "status", "request_sha256")} + record.update(schema="joint-crash-recovery-smoke-v1", status="PASS", + request_sha256=sha(request_path), native=native_binding, host_exit=host_binding, + request=dict(path=str(request_path), sha256=sha(request_path)), job_name=q["job_name"], + candidate=dict(path=str(candidate_path), sha256=sha(candidate_path))) + write_once(Path(q["output"]) / "crash-recovery-smoke.json", record) + return record + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("request", type=Path) + parser.add_argument("--finalize", action="store_true") + parser.add_argument("--job-id") + parser.add_argument("--native-proof", type=Path) + parser.add_argument("--native-sha256") + parser.add_argument("--host-exit", type=Path) + parser.add_argument("--host-exit-sha256") + args = parser.parse_args(argv) + q = json.loads(args.request.read_text()) + if args.finalize: + need(all((args.job_id, args.native_proof, args.native_sha256, + args.host_exit, args.host_exit_sha256)), "Exact finalization bindings required") + finalize(args.request, q, job_id=args.job_id, + native_binding=dict(path=str(args.native_proof), sha256=args.native_sha256), + host_binding=dict(path=str(args.host_exit), sha256=args.host_exit_sha256)) + else: + need(not any((args.job_id, args.native_proof, args.native_sha256, + args.host_exit, args.host_exit_sha256)), "Unexpected finalization flags") + run(args.request, q) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/lambda_full.sbatch b/scripts/train/joint/lambda_full.sbatch new file mode 100644 index 000000000..c8f430f41 --- /dev/null +++ b/scripts/train/joint/lambda_full.sbatch @@ -0,0 +1,73 @@ +#!/bin/bash +set -euo pipefail +root=${1:?Operations root} +request=${2:?Exact request} +request_sha=${3:?Bound request hash} +files_sha=${4:?Bound operational manifest} +[[ "${SLURM_RESTART_COUNT:-0}" == 0 ]] +[[ "$request" == "$root"/requests/*.json && "$request_sha" =~ ^[0-9a-f]{64}$ && "$files_sha" =~ ^[0-9a-f]{64}$ ]] +printf '%s %s\n' "$request_sha" "$request" | sha256sum --check --status +printf '%s %s\n' "$files_sha" "$root/FILES.sha256" | sha256sum --check --status +(cd "$root" && sha256sum --check --status FILES.sha256) +source "$root/launcher_common.sh" +export PIXEL_SOURCE_ROOT +PIXEL_SOURCE_ROOT=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source"])' "$request") +world=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["world_size"])' "$request") +output=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["output"])' "$request") +local_logs=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["local_logs"])' "$request") +archive=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source_archive"])' "$request") +archive_sha=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source_archive_sha256"])' "$request") +source_manifest=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source_manifest"])' "$request") +source_manifest_sha=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source_manifest_sha256"])' "$request") +[[ "$world" == 2 || "$world" == 4 ]] +[[ ! -e "$local_logs" ]] +resume=$(python3 -c 'import json,sys;print(int(json.load(open(sys.argv[1])).get("continuation") is not None))' "$request") +if [[ "$resume" == 0 ]]; then [[ ! -e "$output" ]]; else [[ "$resume" == 1 && -f "$output/.training-admitted" ]]; fi +# The shared native service checks the complete prior-request/checkpoint chain. +runtime_exports=$(python3 "$root/lambda_runtime.py" "$request") +eval "$runtime_exports" +mounts=$(cat "$JOINT_RUNTIME_MOUNTS") +credentials_mount= +if [[ -f "$HOME/.netrc" ]]; then credentials_mount=",$HOME/.netrc:/run/pixel-wandb.netrc:ro"; fi +[[ -n "${WANDB_API_KEY:-}" || -n "$credentials_mount" ]] || fail "W&B credentials unavailable" +mkdir -p "$JOINT_NODE_SCRATCH/jit" "$JOINT_NODE_SCRATCH/logs" "$JOINT_NODE_SCRATCH/source" +name="$(basename "$request" .json)-$SLURM_JOB_ID" +receipt="$root/status/$name.exit" +attempt="$root/status/$name.attempt" +mkdir "$attempt" +child= mirror= +finish() { + code=$? + trap - EXIT + if [[ -n "$child" ]] && kill -0 "$child" 2>/dev/null; then kill -TERM "$child" 2>/dev/null || true; wait "$child" 2>/dev/null || true; fi + if [[ -n "$mirror" ]]; then kill "$mirror" 2>/dev/null || true; wait "$mirror" 2>/dev/null || true; fi + if [[ -d "$local_logs" ]]; then + destination="$output/log-mirror" + [[ -f "$output/.training-admitted" ]] || destination="$attempt/startup-log-mirror" + mirror_once "$local_logs" "$destination" || code=1 + fi + (set -o noclobber; printf '%s\n' "$code" > "$receipt") + exit "$code" +} +trap finish EXIT +trap 'exit 143' TERM +python3 "$root/verify_joint_data_cache.py" "$request" "$attempt/data-cache-stage.json" +cache_root=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["data_cache"]["root"])' "$request") +data_root=$(python3 -c 'import json,sys,pathlib;print(pathlib.Path(json.load(open(sys.argv[1]))["data_manifest"]).parent)' "$request") +mounts+=",$cache_root:$data_root:ro" +source_cache="$JOINT_NODE_SCRATCH/source/$SLURM_JOB_ID" +bash "$root/lambda_stage_source.sh" "$archive" "$archive_sha" "$source_manifest" "$source_manifest_sha" "$source_cache" +cp -- "$source_cache/source-stage.json" "$attempt/source-stage.json" +srun --jobid="$SLURM_JOB_ID" --exclusive --exact --unbuffered --nodes=1 --ntasks=1 \ + --cpus-per-task=32 --gres="gpu:h100:$world" --mem=192G --gpu-bind=none \ + --container-name="joint-$name" --container-image="$JOINT_CONTAINER_IMAGE" --container-writable \ + --container-mounts="/workspace:/workspace:rw,$mounts,$source_cache/source:$PIXEL_SOURCE_ROOT:ro,$JOINT_NODE_SCRATCH:$JOINT_NODE_SCRATCH:rw,$JOINT_JOB_TMP:$JOINT_JOB_TMP:rw$credentials_mount" \ + --container-workdir="$PIXEL_SOURCE_ROOT" --output="$root/logs/$name.out" --error="$root/logs/$name.err" \ + bash "$root/lambda_inside.sh" "$root" "$request" & +child=$! +watch_log_mirror "$child" "$local_logs" "$output/log-mirror" "$output/.training-admitted" "$attempt/unused" & +mirror=$! +status=0 +wait "$child" || status=$? +child= +exit "$status" diff --git a/scripts/train/joint/lambda_inside.sh b/scripts/train/joint/lambda_inside.sh new file mode 100644 index 000000000..34b327331 --- /dev/null +++ b/scripts/train/joint/lambda_inside.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -euo pipefail +root=${1:?Operations root} +request=${2:?Exact request} +source "$root/launcher_common.sh" +runtime_exports=$(python3 "$root/lambda_runtime.py" "$request" --inside) +eval "$runtime_exports" +export PIXEL_FA4_ROOT="$JOINT_FA4_ROOT" PIXEL_FA3_ROOT="$JOINT_FA3_ROOT" +export PIXEL_TASK="$JOINT_RUNTIME_TASK" +export PYTHONPATH="$root:$PIXEL_SOURCE_ROOT:$PIXEL_FA3_ROOT:$PIXEL_FA4_ROOT/runtime/nvidia_cutlass_dsl/dsl_packages:$PIXEL_FA4_ROOT/torch-runtime:$PIXEL_FA4_ROOT/runtime:$JOINT_EXTRA_PYTHONPATH" +export PATH="$(dirname "$JOINT_ACTUAL_PYTHON"):$PATH" +export PYTHONDONTWRITEBYTECODE=1 NVIDIA_TF32_OVERRIDE=0 +export OMP_NUM_THREADS=8 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=8 +export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 +cache="$JOINT_NODE_SCRATCH/jit/${SLURM_JOB_ID}" +export CUTE_DSL_CACHE_DIR="$cache/cute" TORCHINDUCTOR_CACHE_DIR="$cache/inductor" TRITON_CACHE_DIR="$cache/triton" +export TMPDIR="$JOINT_JOB_TMP" XDG_CACHE_HOME="$cache/cache" +export WANDB_CACHE_DIR="$cache/cache/wandb" WANDB_DATA_DIR="$cache/data/wandb" +mkdir -p "$CUTE_DSL_CACHE_DIR" "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR" "$TMPDIR" "$XDG_CACHE_HOME" "$WANDB_CACHE_DIR" "$WANDB_DATA_DIR" +phase=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["phase"])' "$request") +[[ "$phase" == smoke || "$phase" == full || "$phase" == recovery ]] +export PIXEL_WANDB_NETRC=/run/pixel-wandb.netrc +exec "$JOINT_BOOTSTRAP_PYTHON" "$root/runtime_bootstrap.py" "$root/run_full.py" "$request" diff --git a/scripts/train/joint/lambda_runtime.py b/scripts/train/joint/lambda_runtime.py new file mode 100644 index 000000000..1ae806093 --- /dev/null +++ b/scripts/train/joint/lambda_runtime.py @@ -0,0 +1,132 @@ +"""Bind the maintained Lambda launcher to explicit runtime and scratch inputs.""" + +import hashlib +import json +import os +from pathlib import Path +import shlex +import shutil +import socket +import sys +import tempfile + + +FIELDS = { + "scratch_root": "JOINT_NODE_SCRATCH", + "container_image": "JOINT_CONTAINER_IMAGE", + "fa4_root": "JOINT_FA4_ROOT", + "fa3_root": "JOINT_FA3_ROOT", + "runtime_task": "JOINT_RUNTIME_TASK", + "extra_pythonpath": "JOINT_EXTRA_PYTHONPATH", + "actual_python": "JOINT_ACTUAL_PYTHON", + "bootstrap_python": "JOINT_BOOTSTRAP_PYTHON", +} + + +def validate_policy(request, environment): + runtime = request["lambda_runtime"] + if (runtime.get("schema") != "joint-lambda-runtime-v1" + or runtime["node"] != environment.get("SLURMD_NODENAME") + or not environment.get("SLURM_JOB_ID", "").isdigit() + or environment.get("SLURM_RESTART_COUNT", "0") != "0"): + raise ValueError("Exact fresh Lambda node/allocation required") + scratch = Path(runtime["scratch_root"]) + if (not scratch.is_absolute() or ".." in scratch.parts + or scratch.parts[1:2] != ("ephemeral",) or len(scratch.parts) < 4): + raise ValueError("Specific user-owned /ephemeral scratch required") + for field in FIELDS.keys() - {"container_image"}: + path = Path(runtime[field]) + if not path.is_absolute() or ".." in path.parts: + raise ValueError("Absolute runtime path required: " + field) + temporary = Path(runtime["temporary_root"]) + if (not temporary.is_absolute() or ".." in temporary.parts + or temporary.parts[1:2] != ("ephemeral",) or len(temporary.parts) < 4 + or len(os.fsencode(temporary / ("j" + environment["SLURM_JOB_ID"]) / "tmp")) > 60): + raise ValueError("Short explicit /ephemeral temporary root required for AF_UNIX sockets") + if runtime["container_image"] != "nvcr.io#nvidia/pytorch:25.06-py3": + raise ValueError("Changed pinned container recipe") + if type(runtime["minimum_free_scratch_bytes"]) is not int or runtime["minimum_free_scratch_bytes"] < 10_000_000_000: + raise ValueError("Explicit compiler/source scratch headroom required") + return runtime + + +def socket_probe(directory): + """Exercise the actual data-loader resource-sharer socket naming pattern.""" + with tempfile.TemporaryDirectory(prefix="pymp-", dir=directory) as child: + address = str(Path(child) / "listener-0123456789abcdef") + if len(os.fsencode(address)) >= 104: + raise ValueError("Temporary path exceeds portable AF_UNIX socket budget") + with socket.socket(socket.AF_UNIX) as listener: + listener.bind(address) + + +def exports(request, environment, *, inside=False): + runtime = validate_policy(request, environment) + temporary = Path(runtime["temporary_root"]) / ("j" + environment["SLURM_JOB_ID"]) / "tmp" + mounts = runtime["mounts"] + path = Path(mounts["path"]) + if (not path.is_absolute() + or hashlib.sha256(path.read_bytes()).hexdigest() != mounts["sha256"]): + raise ValueError("Runtime mount manifest changed") + if not inside: + for item in path.read_text().strip().split(","): + fields = item.split(":") + if len(fields) != 3 or fields[2] != "ro" or not Path(fields[0]).is_dir(): + raise ValueError("Missing exact node-local runtime mount: " + item) + scratch = Path(runtime["scratch_root"]) + scratch.mkdir(parents=True, exist_ok=True) + if scratch.resolve() != scratch or scratch.stat().st_uid != os.getuid(): + raise ValueError("Scratch must be an owned physical directory") + if shutil.disk_usage(scratch).free < runtime["minimum_free_scratch_bytes"]: + raise ValueError("Insufficient node-local scratch") + with tempfile.TemporaryFile(dir=scratch) as stream: + stream.write(b"joint-runtime-write-check\n") + stream.flush() + os.fsync(stream.fileno()) + temporary.mkdir(parents=True, exist_ok=True) + if temporary.resolve() != temporary or temporary.stat().st_uid != os.getuid(): + raise ValueError("Temporary directory must be owned and physical") + socket_probe(temporary) + result = {name: runtime[field] for field, name in FIELDS.items()} + result["JOINT_RUNTIME_MOUNTS"] = str(path) + result["JOINT_JOB_TMP"] = str(temporary) + return "\n".join("export " + name + "=" + shlex.quote(str(value)) + for name, value in result.items()) + + +def filesystem_recovery_reserve(request): + """Check Lambda's actual persistent destination, without a fabricated Lustre quota.""" + quota = request["quota"] + output, mount = Path(request["output"]), Path(quota["mount"]) + if (quota.get("kind") != "filesystem" or quota["uid"] != os.getuid() + or quota["user"] != os.environ.get("SLURM_JOB_USER") + or quota["node"] != os.environ.get("SLURMD_NODENAME") + or not mount.is_absolute() or len(mount.parts) < 2 + or not output.is_relative_to(mount) or not output.parent.is_dir()): + raise ValueError("Persistent recovery filesystem/owner binding changed") + free = shutil.disk_usage(output.parent).free + if free <= request["reserved_checkpoint_bytes"]: + raise ValueError("Insufficient persistent recovery storage") + return json.dumps(dict(kind="filesystem", user=quota["user"], uid=quota["uid"], + node=quota["node"], mount=str(mount), free_bytes=free)) + + +def main(): + if len(sys.argv) not in (2, 3) or (len(sys.argv) == 3 and sys.argv[2] not in ("--inside", "--check-only")): + raise ValueError("Expected request [--inside|--check-only]") + request = json.loads(Path(sys.argv[1]).read_text()) + result = exports(request, os.environ, inside=sys.argv[2:] == ["--inside"]) + if sys.argv[2:] == ["--check-only"]: + scratch = Path(request["lambda_runtime"]["scratch_root"]) + print(json.dumps(dict(status="PASS", scratch_root=str(scratch), uid=os.getuid(), + node=os.environ["SLURMD_NODENAME"], + free_bytes=shutil.disk_usage(scratch).free, + temporary_root=str(Path(request["lambda_runtime"]["temporary_root"])), + runtime_mounts_verified=True, write_fsync_verified=True, + unix_socket_verified=True))) + else: + print(result) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/lambda_stage_source.sh b/scripts/train/joint/lambda_stage_source.sh new file mode 100644 index 000000000..e4bae4d43 --- /dev/null +++ b/scripts/train/joint/lambda_stage_source.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Run only on the allocated node. Preserve exact source bytes and all file checks. +set -euo pipefail +archive=${1:?Pinned source archive} +archive_sha=${2:?Archive SHA256} +manifest=${3:?Pinned source manifest} +manifest_sha=${4:?Manifest SHA256} +cache=${5:?Unique allocated node-local destination} +[[ "${SLURM_JOB_ID:-}" =~ ^[0-9]+$ ]] +[[ "${SLURMD_NODENAME:-}" =~ ^[a-zA-Z0-9._-]+$ ]] +[[ "$archive_sha" =~ ^[0-9a-f]{64}$ && "$manifest_sha" =~ ^[0-9a-f]{64}$ ]] +: "${JOINT_NODE_SCRATCH:?explicit admitted node-local scratch required}" +[[ "$cache" == "$JOINT_NODE_SCRATCH/source/$SLURM_JOB_ID" ]] +[[ ! -e "$cache" ]] +mkdir -p "$JOINT_NODE_SCRATCH/source" +mkdir "$cache" "$cache/source" +cp -- "$archive" "$cache/source.tar.gz" +cp -- "$manifest" "$cache/source.files.sha256" +printf '%s %s\n' "$archive_sha" "$cache/source.tar.gz" | sha256sum --check --status +printf '%s %s\n' "$manifest_sha" "$cache/source.files.sha256" | sha256sum --check --status +tar -xzf "$cache/source.tar.gz" -C "$cache/source" +(cd "$cache/source" && sha256sum --check --status "$cache/source.files.sha256") +chmod -R a-w "$cache/source" +(set -o noclobber; printf '{"schema":"joint-node-local-source-v2","status":"PASS","job_id":%s,"node":"%s","archive_sha256":"%s","manifest_sha256":"%s","source":"%s/source","all_manifest_files_verified":true}\n' \ + "$SLURM_JOB_ID" "$SLURMD_NODENAME" "$archive_sha" "$manifest_sha" "$cache" > "$cache/source-stage.json") +cat "$cache/source-stage.json" diff --git a/scripts/train/joint/lr_policy.py b/scripts/train/joint/lr_policy.py new file mode 100644 index 000000000..d52108580 --- /dev/null +++ b/scripts/train/joint/lr_policy.py @@ -0,0 +1,50 @@ +"""Declared optimizer policy shared by admission, smoke and restart checks. + +This validates configuration and observed state; it never changes an optimizer. +""" +import math + + +def validate(policy): + if not isinstance(policy, dict) or set(policy) != {"peak_lr", "warmup_optimizer_steps", "schedule"}: + raise ValueError("Explicit complete LR policy required") + rate, warmup = policy["peak_lr"], policy["warmup_optimizer_steps"] + if type(rate) not in (float, int) or not math.isfinite(rate) or rate <= 0 or type(warmup) is not int: + raise ValueError("Invalid LR or warmup") + if not ((policy["schedule"] == "constant" and warmup == 0) + or (policy["schedule"] == "linear_warmup_then_constant" and warmup > 0)): + raise ValueError("Unsupported LR schedule") + return dict(policy) + + +def from_config(model): + scheduler = model.get("scheduler") + warmup = 0 + if scheduler is not None: + if (scheduler.get("_target_") != "egomimic.utils.schedulers.linear_warmup_constant_scheduler" + or model.get("scheduler_interval") != "step" or model.get("scheduler_frequency") != 1): + raise ValueError("Unsupported configured scheduler") + warmup = scheduler.get("warmup_steps") + return validate(dict(peak_lr=model["optimizer"]["lr"], warmup_optimizer_steps=warmup, + schedule="constant" if scheduler is None else "linear_warmup_then_constant")) + + +def update_lr(policy, completed_steps): + policy = validate(policy) + if type(completed_steps) is not int or completed_steps < 0: + raise ValueError("Actual nonnegative optimizer step required") + warmup = policy["warmup_optimizer_steps"] + return policy["peak_lr"] * (min((completed_steps + 1) / warmup, 1.0) if warmup else 1.0) + + +def verify_scheduler(policy, states, completed_steps): + policy = validate(policy) + warmup = policy["warmup_optimizer_steps"] + if not warmup: + if states != []: + raise ValueError("Constant LR must not restore a scheduler") + elif (len(states) != 1 or states[0]["last_epoch"] != completed_steps + or states[0]["base_lrs"] != [policy["peak_lr"]] + or states[0]["lr_lambdas"] != [{"warmup_steps": warmup}]): + raise ValueError("Wrong exact scheduler continuation state") + return True diff --git a/scripts/train/joint/recovery_check.py b/scripts/train/joint/recovery_check.py new file mode 100644 index 000000000..3ae682774 --- /dev/null +++ b/scripts/train/joint/recovery_check.py @@ -0,0 +1,319 @@ +"""Prove real native BF16/DDP continuation from an already verified joint smoke. + +This diagnostic uses the checkpoint-bound model and the existing trainHydra +loader. It does not submit jobs, change model/data settings or authorize full +training by itself. Its short cadence is diagnostic, never a production recipe. +""" + +import copy +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys + +import torch +from omegaconf import OmegaConf, open_dict + +from egomimic.pl_utils.joint_smoke import JointSmokeProof, weights_digest +from scripts.train.joint.lr_policy import from_config, verify_scheduler +from scripts.train.joint.contract import bound, canonical_uuid, need, sha, write_once, verify_allocation_relocation + + +def equal_state(actual, saved): + if isinstance(saved, torch.Tensor): + return (isinstance(actual, torch.Tensor) and actual.dtype == saved.dtype + and actual.shape == saved.shape and torch.equal(actual.detach().cpu(), saved.cpu())) + if isinstance(saved, dict): + return (isinstance(actual, dict) and actual.keys() == saved.keys() + and all(equal_state(actual[k], v) for k, v in saved.items())) + if isinstance(saved, (list, tuple)): + return (type(actual) is type(saved) and len(actual) == len(saved) + and all(equal_state(a, b) for a, b in zip(actual, saved))) + # NumPy's RNG state includes arrays, without requiring a pickle workaround. + import numpy as np + if isinstance(saved, np.ndarray): + return isinstance(actual, np.ndarray) and np.array_equal(actual, saved) + return type(actual) is type(saved) and actual == saved + + +def quota_free_bytes(text, mount): + """Read the aggregate Lustre user row; filesystem free alone ignores quota.""" + rows = [line.split() for line in text.splitlines() if line.split()[:1] == [mount]] + need(len(rows) == 1 and len(rows[0]) >= 4, "Missing unambiguous live Lustre quota") + used, soft, hard = [int(v.rstrip("*")) for v in rows[0][1:4]] + limits = [v for v in (soft, hard) if v > 0] + need(limits, "Lustre quota limit unavailable") + return max(0, min(limits) - used) * 1024 + + +class ResumeOptimizerProof(JointSmokeProof): + """Reuse smoke precision/LR/gradient hooks, checking restored native state.""" + def __init__(self, *, output, world_size, start_step, stop_step, checkpoint): + super().__init__(output, world_size, expected_steps=stop_step, warmup_steps=start_step) + self.start_step, self.checkpoint = start_step, checkpoint + + def on_fit_start(self, trainer, pl_module): + # Lightning restores optimizer/loop state after on_fit_start. + pass + + def on_train_start(self, trainer, pl_module): + need(trainer.world_size == self.world_size and trainer.precision == "bf16-mixed", + "Recovery must use the original BF16/DDP topology") + need(trainer.global_step == self.start_step, "Wrong restored step") + saved = torch.load(self.checkpoint, map_location="cpu", weights_only=False, mmap=True) + need(saved["global_step"] == self.start_step, "Wrong input checkpoint") + need(equal_state(pl_module.state_dict(), saved["state_dict"]), "Weights not strictly restored") + need(equal_state(trainer.optimizers[0].state_dict(), saved["optimizer_states"][0]), + "AdamW moments/groups not strictly restored") + self.lr_policy = from_config(pl_module.hparams.config_tree.model) + restored_schedulers = [c.scheduler.state_dict() for c in trainer.lr_scheduler_configs] + need(equal_state(restored_schedulers, saved.get("lr_schedulers", [])), + "LR scheduler not strictly restored") + verify_scheduler(self.lr_policy, restored_schedulers, self.start_step) + need(equal_state(pl_module.training_behavior._pending_rng, saved["mot_rng_state"]), + "Native per-rank RNG restoration was not armed") + need(trainer.datamodule.consumed_batches == self.start_step, "Sampler cursor not restored") + self.initial_weights = weights_digest(pl_module.state_dict()) + self.network = pl_module.model.pipeline.stage_by_id("joint_denoiser").model + need(all(p.requires_grad and p.dtype == torch.float32 for p in self.network.parameters()), + "Stored weight dtype/trainability changed") + self.output.mkdir(parents=True, exist_ok=True) + self.fp32_losses = self.bf16_outputs = 0 + self.optimizer_lrs, self.validation_steps = [], [] + self.gradients_checked = False + + def observe(module, args, outputs): + need(set(outputs) == set(self.network.streams) + and all(v.dtype == torch.bfloat16 for v in outputs.values()), + "Native active-stream forward did not use BF16") + if pl_module.training and torch.is_grad_enabled(): + self.bf16_outputs += 1 + self.precision_hook = self.network.register_forward_hook(observe) + + def on_before_optimizer_step(self, trainer, pl_module, optimizer): + # Check every resumed update, retaining the maintained smoke checks. + self.gradients_checked = False + super().on_before_optimizer_step(trainer, pl_module, optimizer) + need(pl_module.training_behavior._pending_rng is None, "Native RNG was not restored on training") + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + pass + + def on_validation_start(self, trainer, pl_module): + self.validation_steps.append(trainer.global_step) + + def on_validation_batch_end(self, *args, **kwargs): + pass + + def on_train_end(self, trainer, pl_module): + self.precision_hook.remove() + count = self.expected_steps - self.start_step + need(trainer.global_step == self.expected_steps + and self.fp32_losses == self.bf16_outputs == len(self.optimizer_lrs) == count, + "Missing actual resumed optimizer/precision updates") + need(self.validation_steps == list(range(self.start_step + 2, self.expected_steps + 1, 2)), + "Continuation repeated or skipped validation events") + need(trainer.datamodule.consumed_batches == self.expected_steps + and pl_module.evaluator.has_completed_validation(self.expected_steps) + and pl_module.evaluator.global_validation_batches == 100, + "Missing resumed sampling cursor or all100 validation anchors") + metrics = {str(k): float(v) for k, v in trainer.callback_metrics.items() if str(k).startswith("val/")} + import math + need(metrics and all(math.isfinite(v) for v in metrics.values()), "Nonfinite resumed validation") + final = weights_digest(pl_module.state_dict()) + gathered = [None] * self.world_size + torch.distributed.all_gather_object(gathered, final) + need(len(set(gathered)) == 1 and final != self.initial_weights, "Resumed weights did not update identically") + need(all(float(s["step"]) == self.expected_steps for s in trainer.optimizers[0].state.values()) + and verify_scheduler(self.lr_policy, [c.scheduler.state_dict() for c in trainer.lr_scheduler_configs], self.expected_steps), + "Optimizer/scheduler restarted instead of continuing") + from scripts.clusters.pace.gpu_admission import cuda_uuid + + write_once(self.output / f"resume-{self.start_step}-{self.expected_steps}-rank{trainer.global_rank}.json", dict( + status="RESUME_PASS", job_id=os.environ["SLURM_JOB_ID"], rank=trainer.global_rank, + world_size=self.world_size, start_step=self.start_step, global_step=self.expected_steps, + optimizer_updates=count, fp32_losses=self.fp32_losses, bf16_outputs=self.bf16_outputs, + optimizer_update_lrs=self.optimizer_lrs, validation_steps=self.validation_steps, + strict_start_weights_optimizer_scheduler_rng_cursor=True, + fp32_gradients_verified=self.gradients_checked, weights_sha256=final, + initial_weights_sha256=self.initial_weights, validation_metrics=metrics, + gpu_uuid=cuda_uuid(torch.cuda.get_device_properties(pl_module.device)), + )) + + +def phase_config(original, *, output, checkpoint, start_step, stop_step, service_output, world): + cfg = OmegaConf.create(copy.deepcopy(OmegaConf.to_container(original, resolve=True))) + original_model, original_data = copy.deepcopy(cfg.model), copy.deepcopy(cfg.data) + with open_dict(cfg): + cfg.ckpt_path = checkpoint + cfg.paths.output_dir = output + cfg.val_at_start = cfg.val_at_end = False + cfg.trainer.max_steps = 56 + cfg.trainer.val_check_interval = 2 + cfg.callbacks.pop("joint_smoke") + cb = cfg.callbacks.model_checkpoint + cb.every_n_train_steps, cb.every_n_epochs = 2, 0 + cb.save_top_k, cb.save_last = -1, "link" + cb.dirpath = output + "/checkpoints" + cb.filename = "allocation-epoch{epoch:04d}-step{step:09d}" + cfg.callbacks.resume_proof = dict( + _target_="scripts.train.joint.recovery_check.ResumeOptimizerProof", + output=service_output, world_size=world, start_step=start_step, + stop_step=stop_step, checkpoint=checkpoint, + ) + if stop_step < 56: + cfg.callbacks.allocation_boundary = dict( + _target_="scripts.train.joint.allocation_boundary.ScheduledAllocationBoundary", + start_step=start_step, stop_step=stop_step, target_step=56, + cadence=2, output=service_output, + ) + cfg.logger.csv.save_dir = output + "/logs" + cfg.evaluator.pixel_diagnostics.artifact_dir = service_output + f"/validation-{stop_step}" + need(OmegaConf.to_container(cfg.model, resolve=True) == OmegaConf.to_container(original_model, resolve=True) + and OmegaConf.to_container(cfg.data, resolve=True) == OmegaConf.to_container(original_data, resolve=True), + "Recovery changed model/data semantics") + return cfg + + +def finish_proof(request_path, q, original, admission, training_job_id, repaired=False): + output = Path(q["output"]) + evidence = output/"proof" + boundary = json.loads((evidence/"allocation-step000000052.json").read_text()) + checkpoint = boundary["checkpoint"]["path"] + need(boundary["status"] == "PASS" and boundary["checkpoint_saved_after_validation"] + and sha(checkpoint) == boundary["checkpoint"]["sha256"], "Changed boundary recovery") + phases = [] + for start, stop in ((50,52),(52,56)): + ranks = [json.loads((evidence/f"resume-{start}-{stop}-rank{r}.json").read_text()) + for r in range(original["world_size"])] + need(all(r["status"] == "RESUME_PASS" and r["job_id"] == training_job_id + and r["rank"] == i and r["world_size"] == original["world_size"] + and r["start_step"] == start and r["global_step"] == stop + and r["optimizer_updates"] == r["fp32_losses"] == r["bf16_outputs"] == stop-start + and r["fp32_gradients_verified"] and r["strict_start_weights_optimizer_scheduler_rng_cursor"] + and r["validation_steps"] == list(range(start+2,stop+1,2)) + for i,r in enumerate(ranks)), "Missing all-rank actual optimizer recovery proof") + phases.append(ranks) + need(all(a["weights_sha256"] == b["initial_weights_sha256"] for a,b in zip(*phases)), + "Second process did not restore first-process weights") + if admission is not None: + attestation = bound(admission) + need(all([canonical_uuid(r["gpu_uuid"]) for r in ranks] + == [canonical_uuid(r["uuid"]) for r in attestation["gpu_probe"]["ranks"]] + for ranks in phases), "Resumed optimizer ranks differ from admitted physical GPUs") + write_once(output/"recovery-completion.json", dict( + schema="joint-recovery-check-v1", status="PASS", job_id=training_job_id, + postflight_job_id=os.environ["SLURM_JOB_ID"], postflight_repaired=repaired, + request_sha256=sha(request_path), smoke_completion_sha256=q["smoke_completion"]["sha256"], + source_revision=original["source_revision"], data_manifest_sha256=original["data_manifest_sha256"], + active_streams=original["active_streams"], world_size=original["world_size"], global_batch=32, + resumed_from=50, boundary_step=52, resumed_to=56, actual_updates=6, + model_data_unchanged=True, all_rank_full_state_restored=True, + allocation_policy=q["allocation"], allocation_relocation=q.get("allocation_relocation"), + no_extra_validation=True, allocation_attestation=admission, + boundary_checkpoint_sha256=sha(checkpoint), + native_gpu_status="FAILED_POSTFLIGHT_UUID_FORMAT" if repaired else "postflight_complete", + )) + + +def main(): + request_path = Path(sys.argv[1]) + q = json.loads(request_path.read_text()) + need(q["schema"] == "joint-recovery-check-v1", "Wrong recovery request") + original = bound(q["smoke_request"]) + verify_allocation_relocation(q, original, q["smoke_request"]["sha256"]) + execution = copy.deepcopy(original) + execution["allocation"] = q["allocation"] + if q.get("operational_service"): + execution["operational_service"] = q["operational_service"] + completion = bound(q["smoke_completion"]) + bound(q["resolved"], False) + cfg = OmegaConf.load(q["resolved"]["path"]) + need(completion["status"] == "PASS" and completion["global_step"] == 50 + and completion["world_size"] == original["world_size"], "Missing exact GPU smoke") + need(sha(q["checkpoint"]["path"]) == completion["checkpoint_sha256"] == q["checkpoint"]["sha256"], + "Changed smoke checkpoint") + for line in Path(original["source_manifest"]).read_text().splitlines(): + digest, relative = line.split(" ", 1) + need(sha(Path(original["source"]) / relative) == digest, "Changed checkpoint-bound source") + need(sha(original["source_manifest"]) == original["source_manifest_sha256"] + and sha(original["data_manifest"]) == original["data_manifest_sha256"], "Changed source/data manifest") + if len(sys.argv) == 4 and sys.argv[2] == "--postflight": + job = sys.argv[3] + need(job.isdecimal(), "Exact native training job required") + root = request_path.parent.parent + raw = subprocess.check_output(["sacct","-X","-n","-P","-j",job, + "-o","JobIDRaw,State,ExitCode,User,UID,NodeList,Account,QOS"],text=True) + rows=[line.split("|") for line in raw.splitlines() if line.split("|")[0] == job] + need(len(rows)==1 and rows[0][1:3]==["FAILED","1:0"], "Unexpected original job termination") + native=rows[0] + policy=execution["allocation"] + need(native[3:5]==[policy["user"],str(policy["uid"])] + and native[6:8]==[policy["account"],policy["qos"]], "Wrong native GPU job owner") + error=(root/"logs"/f"recovery-{job}.err").read_text() + need(error.count("ValueError: Resumed optimizer ranks differ from admitted physical GPUs")==1 + and error.count("ValueError:")==1, "Failure was not solely UUID-format postflight") + attestation_path=root/"status"/f"native-admission-{job}"/"attestation.json" + admission=dict(path=str(attestation_path),sha256=sha(attestation_path)) + attestation=bound(admission) + need(attestation["job_id"]==job and attestation["node"]==native[5], "Wrong recorded allocation") + from scripts.train.joint.allocation import verify_pace + + verify_pace(execution,attestation,dict(SLURM_JOB_ID=job,SLURMD_NODENAME=native[5],SLURM_RESTART_COUNT="0")) + (Path(q["output"])/"proof"/f"native-{job}.txt").write_text(raw) + finish_proof(request_path,q,original,admission,job,repaired=True) + return + output = Path(q["output"]) + required = q["reserved_checkpoint_bytes"] + need(type(required) is int and required >= 4 * Path(q["checkpoint"]["path"]).stat().st_size + 500000000, + "Four retained recovery checkpoints plus artifact headroom are required") + quota = q["quota"] + if execution.get("allocation") is None: + from scripts.train.joint.lambda_runtime import filesystem_recovery_reserve + + live_quota = filesystem_recovery_reserve(q) + else: + need(output.is_relative_to(quota["mount"]) and quota["user"] == execution["allocation"]["user"], + "Quota is not bound to the actual output owner") + live_quota = subprocess.check_output(["lfs", "quota", "-u", quota["user"], quota["mount"]], text=True) + need(quota_free_bytes(live_quota, quota["mount"]) > required, "Insufficient live user quota") + need(not output.exists() and shutil.disk_usage(output.parent).free > q["reserved_checkpoint_bytes"], + "Recovery output collision or insufficient retention space") + output.mkdir() + (output/"quota-before.txt").write_text(live_quota) + root = request_path.parent.parent + from scripts.train.joint.allocation import admit + admission = admit(execution, status_root=root/"status", bootstrap=root/"runtime_bootstrap.py") + evidence = output / "proof" + evidence.mkdir() + config_dir = output / "configs" + config_dir.mkdir() + checkpoint = q["checkpoint"]["path"] + from scripts.train.joint.service_bootstrap import command as service_command + for start, stop in ((50, 52), (52, 56)): + phase = phase_config(cfg, output=str(output), checkpoint=checkpoint, + start_step=start, stop_step=stop, + service_output=str(evidence), world=original["world_size"]) + path = config_dir / f"resume-{start}-{stop}.yaml" + path.write_text(OmegaConf.to_yaml(phase, resolve=True)) + command = [sys.executable, str(root/"runtime_bootstrap.py"), "-m", "torch.distributed.run", + "--standalone", "--nnodes=1", "--nproc-per-node="+str(original["world_size"]), + "--max-restarts=0", + *service_command(q, request_path, root/"runtime_bootstrap.py", + ["-m", "egomimic.trainHydra"])[1:], + "--config-path", str(config_dir), "--config-name", path.stem, + "hydra.run.dir="+str(output/f"hydra-{stop}"), "hydra.job.chdir=false"] + subprocess.run(command, check=True, cwd=original["source"]) + if stop == 52: + boundary = json.loads((evidence/"allocation-step000000052.json").read_text()) + need(boundary["status"] == "PASS" and boundary["checkpoint_saved_after_validation"], + "Missing post-validation boundary recovery") + checkpoint = boundary["checkpoint"]["path"] + need(sha(checkpoint) == boundary["checkpoint"]["sha256"], "Changed boundary checkpoint") + finish_proof(request_path,q,original,admission,os.environ["SLURM_JOB_ID"]) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/retention.py b/scripts/train/joint/retention.py new file mode 100644 index 000000000..60523c57d --- /dev/null +++ b/scripts/train/joint/retention.py @@ -0,0 +1,142 @@ +"""Read the maintained mirror's strict receipts for remotely retained milestones. + +Local pruning is still owned by ice_checkpoint_mirror. This verifier does not +copy or delete files and never substitutes a filename for remote SHA evidence. +""" + +import json +import hashlib +from pathlib import Path +import re +import shutil +import subprocess + +from scripts.train.joint.contract import bound, need + + +def mirror_row(request): + policy = request["archive_policy"] + manifest = bound(policy["manifest"]) + rows = [r for r in manifest["runs"] if r["id"] == request["run_id"]] + need(len(rows) == 1, "Mirror manifest does not identify this run uniquely") + row = rows[0] + need(Path(row["local_root"]) == Path(request["output"]) + and row["checkpoint_glob"] == "checkpoints/*.ckpt" + and row["retain_local"] == 2, "Mirror retention policy changed") + bound(policy["validator"],False) + need(row["validator"] == policy["validator"]["path"], "Mirror validator changed") + return row + + +def check_storage(request, completion): + """Reduced local reserve requires an actual strict zero-error mirror proof.""" + smoke = bound(request["smoke"]["request"]) + size = (Path(smoke["output"])/completion["ranks"][0]["checkpoint_name"]).stat().st_size + stops = request.get("allocation_boundaries", [request["target_steps"]]) + need(stops == sorted(set(stops)) and stops[-1] == request["target_steps"] + and all(type(s) is int and s > 0 and s % 20000 == 0 for s in stops) + and request.get("allocation_stop_step", request["target_steps"]) in stops, + "Unbound allocation/checkpoint storage plan") + count = request["target_steps"]//20000 + policy = request.get("archive_policy") + if policy is not None: + # The probe binds the exact already-GPU-verified checkpoint and native + # mirror receipts; a running mirror job or empty cycle cannot pass. + probe = policy["probe"] + state, pressure = bound(probe["state"]), bound(probe["pressure"]) + native = dict(s.split("=",1) for s in bound(probe["native"],False).split() if "=" in s) + owner = (request.get("allocation") or {}).get("user","paphiwetsa3") + need(native["JobId"] == str(probe["job_id"]) and native["JobState"] == "COMPLETED" + and native["ExitCode"] == "0:0" and native["UserId"].startswith(owner+"(") + and pressure["cycle_errors"] == 0, + "Mirror probe did not complete a zero-error native cycle") + matches = [v for v in state["files"].values() if v.get("sha256") == completion["checkpoint_sha256"]] + need(len(matches) == 1 and matches[0]["remote_verified"] is True + and matches[0]["global_step"] == matches[0]["metadata"]["global_step"] == 50 + and matches[0]["metadata"]["strict_model_reload"] is True + and matches[0]["metadata"]["strict_model_parameter_count"] == request["parameters"] + and pressure["updated_at_unix"] >= matches[0]["verified_at_unix"], + "Mirror probe lacks the exact strict model/checkpoint archive") + mirror_row(request) + bound(policy["ssh"],False) + need(policy["ssh"]["sha256"] == probe["ssh_sha256"] + and policy["remote_host"] == probe["remote_host"] + and policy["validator"]["sha256"] == probe["validator_sha256"], + "Archive SSH route or strict validator changed") + from scripts.ice.ice_checkpoint_mirror import remote_command + + remote_command(policy["ssh"]["path"],policy["remote_host"],["true"]) + # Two verified recovery copies remain local while the next immutable + # save is written and archived. Never spend that in-flight save space + # on boundary recoveries or artifact headroom. + count = 3 + required = (count+len(stops)-1)*size + 1000000000 + need(type(request["reserved_checkpoint_bytes"]) is int + and request["reserved_checkpoint_bytes"] >= required, + "Reservation does not retain every scheduled and boundary checkpoint") + output = Path(request["output"]) + existing = output + while not existing.exists(): + existing = existing.parent + retained = sum(p.stat().st_size for pattern in ("checkpoints/*.ckpt","allocation-evidence/allocation-recovery*.ckpt") + for p in output.glob(pattern) if p.is_file() and not p.is_symlink()) + remaining = max(0,request["reserved_checkpoint_bytes"]-retained) + need(shutil.disk_usage(existing).free > remaining, + "Insufficient filesystem storage") + allocation = request.get("allocation") or {} + if allocation.get("cluster") == "ice-slurm": + from scripts.train.joint.recovery_check import quota_free_bytes + + quota = request["quota"] + need(quota["user"] == allocation["user"] and output.is_relative_to(quota["mount"]), + "Quota does not identify this output owner") + text = subprocess.check_output(["lfs","quota","-u",quota["user"],quota["mount"]],text=True) + need(quota_free_bytes(text,quota["mount"]) > remaining, + "Insufficient live user quota for retained training") + return required + + +def archived_milestones(request, missing): + if not missing: + return {}, None + policy = request.get("archive_policy") + need(policy is not None, "Missing scheduled local milestones and no archive policy") + row = mirror_row(request) + bound(policy["ssh"], False) + directory = Path(request["output"])/"checkpoints" + state_path = Path(policy["state"]) + state_bytes = state_path.read_bytes() + state = json.loads(state_bytes) + found, verified = {}, [] + for local_name, entry in state["files"].items(): + if entry.get("run_id") != request["run_id"]: + continue + local = Path(local_name) + need(local.parent == directory, "Mirror checkpoint escapes this run") + match = re.search(r"([0-9]{9})\.ckpt$", local.name) + need(match is not None and local.name != "last.ckpt", "Unqualified archived checkpoint name") + step = int(match[1]) + if step not in missing: + continue + metadata = entry["metadata"] + digest = entry["sha256"] + remote = entry["remote_path"] + need(step not in found and entry["global_step"] == metadata["global_step"] == step + and entry["remote_verified"] is True + and metadata["strict_model_reload"] is True + and metadata["strict_model_parameter_count"] == request["parameters"] + and metadata["run_id"] == request["run_id"] + and re.fullmatch(r"[0-9a-f]{64}", digest) + and remote == row["remote_dir"].rstrip("/")+"/"+local.stem+"."+digest+".ckpt", + "Missing strict checkpoint identity or immutable archive receipt") + from scripts.ice.ice_checkpoint_mirror import remote_sha + + need(remote_sha(policy["ssh"]["path"], policy["remote_host"], remote) == digest, + "Archived milestone failed current remote SHA verification") + found[step] = Path(remote) + verified.append(dict(global_step=step, sha256=digest, remote_path=remote, + strict_model_parameter_count=metadata["strict_model_parameter_count"])) + need(set(found) == set(missing), "Scheduled checkpoints are absent from the verified archive") + return found, dict(state_path=str(state_path), state_sha256=hashlib.sha256(state_bytes).hexdigest(), + verified_steps=sorted(found), verified_checkpoints=verified, + manifest_sha256=policy["manifest"]["sha256"]) diff --git a/scripts/train/joint/run.py b/scripts/train/joint/run.py index 7435b981a..4896a01bb 100644 --- a/scripts/train/joint/run.py +++ b/scripts/train/joint/run.py @@ -11,12 +11,14 @@ import sys from pathlib import Path +from scripts.train.joint.lr_policy import from_config, validate from scripts.train.joint.contract import ( bound, compare_configs, need, sha, verify_evidence, + verify_model_execution, write_once, ) @@ -46,13 +48,15 @@ def main(): "Wrong training request", ) smoke = request["phase"] == "smoke" + crash = request.get("crash_recovery") is not None + need(not crash or not smoke, "Crash recovery cannot enter the fresh smoke runner") need( - request["target_steps"] == (50 if smoke else 100000), + request["target_steps"] in ((50,) if smoke else (100000, 300000)), "Wrong authorized training budget", ) from scripts.train.joint.allocation import admit - allocation = admit(request) + allocation = admit(request, status_root=root/"status", bootstrap=root/"runtime_bootstrap.py") need( sha(request["source_manifest"]) == request["source_manifest_sha256"], "Changed source manifest", @@ -65,32 +69,63 @@ def main(): "Changed data manifest", ) preflight = None + start, stop = 0, request["target_steps"] + crash_storage = None if not smoke: - _, preflight = verify_evidence(request) - controls = bound(request["factorial_controls"]) - need( - controls["status"] == "PASS" - and controls["source_revision"] == request["source_revision"] - and controls["experiments"] == request["factorial_experiments"], - "Missing paired four-arm GPU controls", - ) - need( - any( - row["experiment"] == request["experiment"] - and row["smoke_request_sha256"] == request["smoke"]["request"]["sha256"] - and row["smoke_completion_sha256"] - == request["smoke"]["completion"]["sha256"] - for row in controls["evidence"] - ), - "Paired evidence belongs to another smoke", - ) + completion, preflight = verify_evidence(request) + from scripts.train.joint.retention import check_storage + + from scripts.train.joint.continuation import allocation_steps + + start, stop = allocation_steps(request) + if crash: + from scripts.train.joint.crash_recovery import ( + assert_native_writer_quiescence, prove_destination_storage, + ) + + assert_native_writer_quiescence(request) + crash_storage = prove_destination_storage(request) + check_storage(request, completion) + from scripts.train.joint.verify_temporal_pair import admission_mode + + comparison_request = request + if crash: + from scripts.train.joint.crash_recovery import foundation_request + + comparison_request = foundation_request(request) + mode = admission_mode(comparison_request) + if mode == "temporal_comparison": + from scripts.train.joint.verify_temporal_pair import admit as admit_temporal_pair + admit_temporal_pair(comparison_request) + elif mode == "modality_comparison": + from scripts.train.joint.contract import verify_modality_extension + verify_modality_extension(request, completion) + elif mode == "optimizer_comparison": + from scripts.train.joint.contract import verify_optimizer_comparison + verify_optimizer_comparison(request, completion) + else: + controls = bound(request["factorial_controls"]) + need( + controls["status"] == "PASS" + and controls["source_revision"] == request["source_revision"] + and controls["experiments"] == request["factorial_experiments"], + "Missing paired four-arm GPU controls", + ) + need( + any( + row["experiment"] == request["experiment"] + and row["smoke_request_sha256"] == request["smoke"]["request"]["sha256"] + and row["smoke_completion_sha256"] + == request["smoke"]["completion"]["sha256"] + for row in controls["evidence"] + ), + "Paired evidence belongs to another smoke", + ) credentials() from scripts.train.joint.check_wandb import inspect_run - need( - inspect_run(request["run_id"])["exists"] is False, - "Fresh W&B identity exists; inspect rather than reset", - ) + need(inspect_run(request["run_id"])["exists"] == bool(start), + "W&B identity does not match fresh/continued allocation") import torch from hydra import compose, initialize_config_dir from hydra.utils import instantiate @@ -103,6 +138,7 @@ def main(): config_name="train_zarr_cartesian", overrides=request["overrides"] ) resolved = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + expected_data_task = verify_model_execution(request, resolved) world = request["world_size"] need( cfg.seed == 42 @@ -112,27 +148,11 @@ def main(): need( cfg.data.batch_size == request["per_rank_batch"] and cfg.data.batch_size * world == 32 - and cfg.model.optimizer.lr == 5e-6 + and cfg.model.optimizer.lr == validate(request["lr_policy"])["peak_lr"] and cfg.model.optimizer.weight_decay == 0.01, "Batch/optimizer mismatch", ) - need( - cfg.model.scheduler._target_ - == "egomimic.utils.schedulers.linear_warmup_constant_scheduler" - and cfg.model.scheduler.warmup_steps == 3000 - and cfg.model.scheduler_interval == "step" - and cfg.model.scheduler_frequency == 1, - "Scheduler contract changed", - ) - need( - request["lr_policy"] - == dict( - peak_lr=5e-6, - warmup_optimizer_steps=3000, - schedule="linear_warmup_then_constant", - ), - "Unbound LR policy", - ) + need(from_config(cfg.model) == validate(request["lr_policy"]), "Unbound LR policy") need( cfg.ckpt_path is None and not cfg.model.pipeline.get("initialization") @@ -140,9 +160,13 @@ def main(): "Fresh step0/trainable experts required", ) need( - cfg.data.couple_state_observations - and not cfg.data.image_only - and cfg.data.task == "visual_state_action", + ( + (expected_data_task == "visual_diffusion" + and request["active_streams"] == ["image"] + and cfg.data.image_only and not cfg.data.couple_state_observations) + or (cfg.data.couple_state_observations and not cfg.data.image_only) + ) + and cfg.data.task == expected_data_task, "Paired full data/noising contract changed", ) need( @@ -168,6 +192,9 @@ def main(): OmegaConf.create(bound(request["smoke"]["resolved"], False)), resolve=True ) compare_configs(resolved, smoke_cfg, request) + from scripts.train.joint.continuation import configure_allocation + + configure_allocation(cfg, request, start, stop, request_sha256=sha(path)) with torch.device("meta"): network = instantiate( cfg.model.pipeline.stages[cfg.model.pipeline.stage_ids.joint_denoiser].model @@ -202,13 +229,14 @@ def main(): and len(data.val_dataloader()) == 100, "Wrong split or validation anchors", ) - need( - shutil.disk_usage(root).free > request["reserved_checkpoint_bytes"], - "Insufficient full retention space", - ) + if smoke: + need(shutil.disk_usage(root).free > request["reserved_checkpoint_bytes"], + "Insufficient smoke retention space") output = Path(request["output"]) need( - str(cfg.paths.output_dir) == str(output) and not output.exists(), + str(cfg.paths.output_dir) == str(output) + and output.exists() == bool(start and not crash) + and (not start or crash or (output / ".training-admitted").is_file()), "Output collision/mismatch", ) provenance = dict( @@ -242,6 +270,7 @@ def main(): world_size=world, global_batch=32, allocation_attestation=allocation, + crash_destination_storage=crash_storage, ), ) with path.with_suffix(".resolved.yaml").open("x") as stream: @@ -250,14 +279,16 @@ def main(): from scripts.world_model.check_joint_capacity import check_capacity check_capacity(cfg, data, path.with_suffix(".capacity.json")) - output.mkdir(parents=True) - (output / ".training-admitted").write_text(sha(path) + "\n") + if not start or crash: + output.mkdir(parents=True) + (output / ".training-admitted").write_text(sha(path) + "\n") extra = ( "+run_provenance={" + ",".join(k + ":" + json.dumps(v) for k, v in provenance.items()) + "}" ) bootstrap = root / "runtime_bootstrap.py" + from scripts.train.joint.service_bootstrap import command print( json.dumps( dict( @@ -268,6 +299,16 @@ def main(): ), flush=True, ) + training_arguments = request["overrides"] + if not smoke: + # The fully resolved file includes only the strictly compared science + # plus the guarded native checkpoint/callback additions above. + training_arguments = [ + "--config-path", str(path.parent), + "--config-name", path.with_suffix(".resolved.yaml").stem, + "hydra.run.dir="+str(output/"attempts"/f"step{start:09d}"), + "hydra.job.chdir=false", + ] subprocess.run( [ sys.executable, @@ -278,22 +319,18 @@ def main(): "--nnodes=1", "--nproc-per-node=" + str(world), "--max-restarts=0", - str(bootstrap), - str(source / "egomimic/trainHydra.py"), - *request["overrides"], + *command(request, path, bootstrap, ["-m", "egomimic.trainHydra"])[1:], + *training_arguments, extra, ], check=True, cwd=source, ) - verifier = source / ( - "scripts/world_model/verify_joint_smoke.py" - if smoke - else "scripts/train/joint/verify_full.py" - ) + verifier = (source / "scripts/world_model/verify_joint_smoke.py" + if smoke else Path(__file__).with_name("verify_full.py")) arguments = [str(output), "--world-size", str(world)] if smoke else [str(path)] subprocess.run( - [sys.executable, str(bootstrap), str(verifier), *arguments], + command(request, path, bootstrap, [str(verifier), *arguments]), check=True, cwd=source, ) diff --git a/scripts/train/joint/service_bootstrap.py b/scripts/train/joint/service_bootstrap.py new file mode 100644 index 000000000..958319642 --- /dev/null +++ b/scripts/train/joint/service_bootstrap.py @@ -0,0 +1,117 @@ +"""Load a hash-bound operational service beside an immutable scientific source. + +The runtime bootstrap still owns Torch and dependency admission. This second +stage selects operational tools and their explicit cluster dependencies, never +egomimic or its configs. +Use it in every child process so the old source cannot shadow newer recovery +and archive admission code through Python's namespace-package search order. +""" + +import hashlib +import importlib +import json +import runpy +import sys +from pathlib import Path + + +# Dependencies must ship with the service, not be assumed present in historical +# scientific checkouts. Keep this allowlist narrow; model/data code stays there. +DEPENDENCIES = { + "scripts/clusters/pace/__init__.py", + "scripts/clusters/pace/gpu_ownership.py", + "scripts/clusters/pace/gpu_admission.py", + "scripts/clusters/pace/pace_podman.py", + "scripts/ice/ice_checkpoint_mirror.py", +} + + +def service_files(root): + """Exact export closure shared by packaging and admission regression checks.""" + root = Path(root) + names = {str(p.relative_to(root)) for p in + (root / "scripts/train/joint").iterdir() + if p.is_file() and p.suffix in (".py", ".sbatch", ".sh")} + names |= DEPENDENCIES + for name in names: + if not (root / name).is_file() or (root / name).is_symlink(): + raise ValueError("Missing operational dependency: " + name) + return sorted(names) + + +def activate(request): + binding = request["operational_service"] + root = Path(binding["root"]) + manifest = Path(binding["manifest"]) + if not root.is_absolute() or manifest.parent != root: + raise ValueError("Absolute content-bound service root required") + if hashlib.sha256(manifest.read_bytes()).hexdigest() != binding["manifest_sha256"]: + raise ValueError("Operational service manifest changed") + names = set() + for line in manifest.read_text().splitlines(): + expected, name = line.split(" ", 1) + relative = Path(name) + if relative.is_absolute() or ".." in relative.parts or (relative.parts[:3] != ("scripts", "train", "joint") and name not in DEPENDENCIES): + raise ValueError("Service contains an unapproved operational path") + path = root / relative + if path.is_symlink() or hashlib.sha256(path.read_bytes()).hexdigest() != expected: + raise ValueError("Operational service changed: " + name) + names.add(name) + required = {"scripts/train/joint/"+n+".py" for n in ( + "service_bootstrap", "run", "contract", "continuation", "allocation_boundary", "retention", "verify_full", + )} + if not (required | DEPENDENCIES) <= names: + raise ValueError("Incomplete joint service") + # A namespace may already be imported by run_module. Select this audited + # service first, while preserving source-owned modules not included here. + sys.path.insert(0, str(root)) + for name in ("scripts", "scripts.train", "scripts.train.joint", + "scripts.clusters", "scripts.clusters.pace", "scripts.ice"): + package = importlib.import_module(name) + local = str(root.joinpath(*name.split("."))) + package.__path__ = [local] + [str(p) for p in package.__path__ if str(p) != local] + owned_modules = {str(Path(n).with_suffix("")).replace("/", ".") for n in names + if n.endswith(".py")} + for name, module in list(sys.modules.items()): + if name in owned_modules: + filename = getattr(module, "__file__", None) + if filename and not Path(filename).resolve().is_relative_to(root.resolve()): + raise ValueError("Service imported before operational admission: " + name) + return root + + +def command(request, request_path, runtime_bootstrap, arguments): + prefix = [sys.executable, str(runtime_bootstrap)] + if request.get("operational_service"): + prefix += [str(Path(request["operational_service"]["root"])/"scripts/train/joint/service_bootstrap.py"), + str(request_path), "--"] + return prefix + list(arguments) + + +def main(): + request_path, delimiter, *arguments = sys.argv[1:] + if delimiter != "--" or not arguments: + raise ValueError("Expected request -- module-or-script") + root = activate(json.loads(Path(request_path).read_text())) + if arguments in (["--check-only"], ["--check-dependencies"]): + tools = (("contract",) if arguments == ["--check-dependencies"] else + ("contract", "continuation", "retention", "allocation_boundary")) + paths = {name: importlib.import_module("scripts.train.joint."+name).__file__ for name in tools} + for name in DEPENDENCIES: + module = name.removesuffix(".py").replace("/", ".") + paths[module] = importlib.import_module(module).__file__ + if not all(Path(p).resolve().is_relative_to(root.resolve()) for p in paths.values()): + raise ValueError("Operational tool import drift") + print(json.dumps(dict(status="PASS", operational_imports=paths, + model_package=importlib.import_module("egomimic").__file__))) + return + if arguments[0] == "-m": + sys.argv = [arguments[1], *arguments[2:]] + runpy.run_module(arguments[1], run_name="__main__", alter_sys=True) + else: + sys.argv = arguments + runpy.run_path(arguments[0], run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/verify_factorial.py b/scripts/train/joint/verify_factorial.py index 743fe6c8c..95ef0a879 100644 --- a/scripts/train/joint/verify_factorial.py +++ b/scripts/train/joint/verify_factorial.py @@ -7,16 +7,42 @@ from omegaconf import OmegaConf -from scripts.train.joint.contract import bound, need, sha, verify_evidence, write_once +from scripts.train.joint.contract import bound, need, sha, verify_evidence, verify_model_execution, write_once CELLS = ("image_only", "image_proprio", "image_action", "image_proprio_action") +def image_controls(request, cfg, controls): + """Project only the declared native proprio consequences out of paired controls. + + The full context remains independently bound by verify_evidence for each run. + Manifest content, every remaining data setting and actual RGB/mask/noise + digests must still match. Legacy joint requests retain exact context equality. + """ + if request.get("model_execution") is None: + return False + verify_model_execution(request, cfg) + proprio = request["active_streams"] == ["image", "proprio"] + need(request["experiment"] == ("image_proprio" if proprio else "image_only"), + "Native experiment/stream mismatch") + data = controls["data"] + need(data["image_only"] is (not proprio) + and data["couple_state_observations"] is proprio, + "Unexpected native modality data flags") + data.pop("image_only") + data.pop("couple_state_observations") + frames = data["discrete_frames"] if data["mode"] == "discrete" else data["max_frames"] + need(isinstance(frames, int) and frames > 0, "Invalid native token budget") + controls["joint"]["max_tokens"] -= frames * int(proprio) + need(controls["joint"]["max_tokens"] > 0, "Missing image token budget") + return True + + def verify(requests): - need( - len(requests) == 4 and {r["experiment"] for r in requests} == set(CELLS), - "Need exactly four factorial arms", - ) + cells = set(CELLS) if len(requests) == 4 else {"image_only", "image_proprio"} + need(len(requests) in (2, 4) and {r["experiment"] for r in requests} == cells, + "Need all four factorial arms or the exact image/proprio pair") + need(len({r.get("arm", "discrete") for r in requests}) == 1, "Paired controls require one time mode") reference = None rows = [] for request in requests: @@ -41,10 +67,14 @@ def verify(requests): .replace(smoke["output"], "") .replace(smoke["local_logs"], "") ) + native = image_controls(request, cfg, controls) paired = dict( controls=controls, source_revision=request["source_revision"], - data_context=completion["ranks"][0]["data_context_sha256"], + data_context=(dict(scope="native-modality-specific-contexts", + data_manifest_sha256=request["data_manifest_sha256"], + model_execution=request["model_execution"]) + if native else completion["ranks"][0]["data_context_sha256"]), initial_image=[ r["initial_image_weights_sha256"] for r in completion["ranks"] ], @@ -62,13 +92,14 @@ def verify(requests): smoke_completion_sha256=request["smoke"]["completion"]["sha256"], smoke_request_sha256=request["smoke"]["request"]["sha256"], smoke_job=request["smoke"]["job_id"], + data_context_sha256=completion["ranks"][0]["data_context_sha256"], ) ) return dict( schema="joint-factorial-gpu-controls-v1", status="PASS", source_revision=requests[0]["source_revision"], - experiments=sorted(CELLS), + experiments=sorted(cells), paired=reference, evidence=rows, ) @@ -76,7 +107,7 @@ def verify(requests): def main(): parser = argparse.ArgumentParser() - parser.add_argument("requests", nargs=4, type=Path) + parser.add_argument("requests", nargs="+", type=Path) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() result = verify([json.loads(path.read_text()) for path in args.requests]) diff --git a/scripts/train/joint/verify_full.py b/scripts/train/joint/verify_full.py index f7166e679..bbe7fddad 100644 --- a/scripts/train/joint/verify_full.py +++ b/scripts/train/joint/verify_full.py @@ -2,6 +2,7 @@ import csv import json +from scripts.train.joint.lr_policy import validate, update_lr, verify_scheduler import math import os import re @@ -12,11 +13,15 @@ def main(): - request = json.loads(Path(sys.argv[1]).read_text()) + request_path = Path(sys.argv[1]) + request = json.loads(request_path.read_text()) _, preflight = verify_evidence(request) output = Path(request["output"]) - target = request["target_steps"] - need(target in (100000, 300000), "Unsupported retained training budget") + from scripts.train.joint.continuation import allocation_steps + + start, target = allocation_steps(request) + partial = target < request["target_steps"] + need(request["target_steps"] in (100000, 300000), "Unsupported retained training budget") streams = request["active_streams"] count = target // 20000 import torch @@ -40,18 +45,41 @@ def main(): step = int(match[1]) need(step not in milestones, "Duplicate checkpoint step") milestones[step] = path - need( - sorted(milestones) == list(range(20000, target + 1, 20000)), - "Missing or extra production milestones", - ) + expected = set(range(20000, target + 1, 20000)) + need(set(milestones) <= expected and target in milestones, "Unexpected milestones or missing local recovery") + from scripts.train.joint.retention import archived_milestones + + if request.get("crash_recovery") is not None: + from scripts.train.joint.crash_history import collect_milestones + + # Parent checkpoints stay in their original immutable directories. + # The local attempt may legitimately begin after a restored milestone. + need(request.get("archive_policy") is None, + "Crash recovery requires all immutable checkpoints retained locally") + milestones = collect_milestones(request, target) + archive_proof = dict(kind="retained_crash_attempt_directories", + checkpoints={str(k): str(v) for k, v in milestones.items()}) + else: + archived, archive_proof = archived_milestones(request, expected - set(milestones)) + milestones.update(archived) latest = directory / "last.ckpt" need( latest.is_symlink() and latest.resolve() == milestones[target].resolve(), "Last must link to immutable terminal milestone", ) - checkpoint = torch.load(latest, map_location="cpu", weights_only=False, mmap=True) + selected = latest.resolve() + if partial: + receipt_path = output/"allocation-evidence"/f"allocation-step{target:09d}.json" + boundary = json.loads(receipt_path.read_text()) + need(boundary["status"] == "PASS" and boundary["start_step"] == start + and boundary["global_step"] == target + and boundary["target_step"] == request["target_steps"] + and boundary["checkpoint_saved_after_validation"], "Missing graceful boundary proof") + selected = Path(boundary["checkpoint"]["path"]) + need(sha(selected) == boundary["checkpoint"]["sha256"], "Changed post-validation recovery") + checkpoint = torch.load(selected, map_location="cpu", weights_only=False, mmap=True) need(checkpoint["global_step"] == target, "Training ended before requested budget") - cfg = OmegaConf.load(output / ".hydra/config.yaml") + cfg = OmegaConf.load(request_path.with_suffix(".resolved.yaml")) model_config = OmegaConf.create(checkpoint["hyper_parameters"]["config_tree"]) need( OmegaConf.to_container(cfg.model, resolve=True) @@ -74,19 +102,15 @@ def main(): checkpoint.get("lr_schedulers", []), ) need( - len(optimizers) == len(schedulers) == 1 and bool(optimizers[0]["state"]), + len(optimizers) == 1 and bool(optimizers[0]["state"]), "Missing AdamW/scheduler state", ) + policy = validate(request["lr_policy"]) need( - all(g["lr"] == g["initial_lr"] == 5e-6 for g in optimizers[0]["param_groups"]), + all(g["lr"] == update_lr(policy, target) and (not policy["warmup_optimizer_steps"] or g["initial_lr"] == policy["peak_lr"]) for g in optimizers[0]["param_groups"]), "Wrong terminal LR", ) - need( - schedulers[0]["last_epoch"] == target - and schedulers[0]["base_lrs"] == [5e-6] - and schedulers[0]["lr_lambdas"] == [{"warmup_steps": 3000}], - "Wrong scheduler state", - ) + verify_scheduler(policy, schedulers, target) for state in optimizers[0]["state"].values(): need( float(state["step"]) == target @@ -141,10 +165,11 @@ def main(): optimizer = optimizer() scheduler = instantiate(cfg.model.scheduler, optimizer=optimizer) optimizer.load_state_dict(optimizers[0]) - scheduler.load_state_dict(schedulers[0]) + if scheduler is not None: + scheduler.load_state_dict(schedulers[0]) need( - scheduler.state_dict() == schedulers[0] - and all(g["lr"] == 5e-6 for g in optimizer.param_groups), + (scheduler.state_dict() == schedulers[0] if policy["warmup_optimizer_steps"] else scheduler is None) + and all(g["lr"] == update_lr(policy, target) for g in optimizer.param_groups), "Strict AdamW/scheduler reload failed", ) batch = next(iter(data.val_dataloader())) @@ -173,20 +198,29 @@ def main(): "Observed values changed", ) validation = {} - csv_path = Path(request["local_logs"]) / "csv/version_0/metrics.csv" - with csv_path.open() as stream: - for row in csv.DictReader(stream): - metrics = { - k: float(v) - for k, v in row.items() - if k.startswith("val/libero_gripper/") and v - } - if metrics: - need( - all(math.isfinite(v) for v in metrics.values()), - "Nonfinite validation history", - ) - validation[int(float(row["step"]))] = metrics + from scripts.train.joint.continuation import log_segments + + for csv_path in ([] if request.get("crash_recovery") is not None else log_segments(request)): + segment = {} + with csv_path.open() as stream: + for row in csv.DictReader(stream): + metrics = { + k: float(v) + for k, v in row.items() + if k.startswith("val/libero_gripper/") and v + } + if metrics: + step = int(float(row["step"])) + need(all(math.isfinite(v) for v in metrics.values()), "Nonfinite validation history") + segment.setdefault(step, {}).update(metrics) + need(not validation.keys() & segment.keys(), "Repeated validation across allocations") + validation.update(segment) + if request.get("crash_recovery") is not None: + from scripts.train.joint.crash_history import collect_validation + from scripts.train.joint.crash_restore import assert_actual_restore + + validation = collect_validation(request, target, streams) + assert_actual_restore(request, expected_request_sha256=sha(request_path)) need( len(validation) == count and all(len(v) == 2 * len(streams) for v in validation.values()) @@ -197,21 +231,24 @@ def main(): "Missing scheduled validation events", ) write_once( - output / "full-completion.json", + (output/"allocation-evidence"/f"completion-step{target:09d}.json" + if partial else output / "full-completion.json"), dict( - schema="joint-mot-full-v1", + schema="joint-mot-allocation-v1" if partial else "joint-mot-full-v1", status="PASS", + request_sha256=sha(request_path), job_id=os.environ["SLURM_JOB_ID"], run_id=request["run_id"], global_step=target, world_size=request["world_size"], - checkpoint_sha256=sha(latest), + checkpoint_sha256=sha(selected), strict_full_wrapper_optimizer_scheduler_reload=True, sampler_all_rank_rng_verified=True, sample_finite=True, retained_milestones=count, validation_events=count, - checkpoint=str(latest.resolve()), + checkpoint=str(selected), + archive_inventory=archive_proof, ), ) diff --git a/scripts/train/joint/verify_temporal_pair.py b/scripts/train/joint/verify_temporal_pair.py new file mode 100644 index 000000000..3ec638be0 --- /dev/null +++ b/scripts/train/joint/verify_temporal_pair.py @@ -0,0 +1,319 @@ +"""Admit two fresh native image/proprio temporal-representation trials. + +This v1 contract is deliberately narrow: frozen535d, constant1e-4, four ranks, +300k updates and all15 scheduled checkpoints. It validates both real GPU smokes; +a PASS string alone is never an admission. Different D/C training draws and +whole-model initial hashes are retained, not presented as paired controls. +""" + +import argparse +import copy +import hashlib +import json +from pathlib import Path +import re + +from scripts.train.joint.contract import ( + bound, compare_configs, need, sha, verify_evidence, verify_model_execution, + write_once, +) + +SCHEMA = "native-image-proprio-temporal-pair-v1" +SOURCE = "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" +SOURCE_MANIFEST = "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184" +SOURCE_ARCHIVE = "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325" +DATA_MANIFEST = "8c40d792a5cedd4105ff8270ff3141767b147513c52f11292f888d9b1527bfe2" +MODES = ("temporal_comparison", "optimizer_comparison", "modality_comparison", + "factorial_controls") +PROFILES = { + "discrete": dict(strategy="discrete_window", pose=36, tokens=3216, + parameters=102135076, proprio_parameters=13956388), + "continuous": dict(strategy="full_episode_uniform", pose=9, tokens=12864, + parameters=102121225, proprio_parameters=13942537), +} + + +def digest(value): + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"), + allow_nan=False).encode()).hexdigest() + + +def canonical_request(request): + """The receipt binding is the only field excluded from the request identity.""" + value = copy.deepcopy(request) + value.pop("temporal_comparison", None) + return digest(value) + + +def admission_mode(request): + modes = [name for name in MODES if request.get(name) is not None] + need(len(modes) == 1, "Exactly one production comparison contract required") + return modes[0] + + +def valid_hash(value): + return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None + + +def evidence_binding(value): + need(isinstance(value, dict) and set(value) == {"path", "sha256"} + and isinstance(value["path"], str) and Path(value["path"]).is_absolute() + and valid_hash(value["sha256"]), "Invalid immutable evidence binding") + return value + + +def request_shape(request): + need(request.get("schema") == "joint-training-v2" and request.get("phase") == "full" + and request.get("arm") in PROFILES, "Temporal pair requires D/C full requests") + need(not any(k.startswith("_") for k in request), "Request still contains preparation fields") + need(not any(name in request for name in MODES), "Prototype already has an admission contract") + need(not any(name in request for name in ( + "continuation", "allocation_stop_step", "allocation_boundaries", "allocation_recovery")), + "Temporal v1 requires fresh single-allocation requests") + need(request.get("source_revision") == SOURCE + and request.get("source_manifest_sha256") == SOURCE_MANIFEST + and request.get("source_archive_sha256") == SOURCE_ARCHIVE + and request.get("data_manifest_sha256") == DATA_MANIFEST, + "Temporal v1 source/data identity changed") + need(request.get("active_streams") == ["image", "proprio"] + and request.get("experiment") == "image_proprio" + and request.get("action_contract") is None, + "Temporal pair supports image/proprio only") + profile = PROFILES[request["arm"]] + need(request.get("parameters") == profile["parameters"] + and request.get("expert_parameters") == { + "image": 88178688, "proprio": profile["proprio_parameters"]} + and request.get("expert_widths") == [768, 256] + and request.get("vector_expert_parameter_budget") == 20000000, + "Temporal expert architecture changed") + need(request.get("world_size") == 4 and request.get("per_rank_batch") == 8 + and request.get("global_batch") == 32 and request.get("target_steps") == 300000 + and request.get("lr_policy") == { + "peak_lr": 1e-4, "schedule": "constant", "warmup_optimizer_steps": 0}, + "Temporal topology, budget or LR recipe changed") + for name in ("source", "source_manifest", "source_archive", "data_manifest", "output", "local_logs"): + need(isinstance(request.get(name), str) and Path(request[name]).is_absolute(), + "Absolute bound path required: " + name) + need(isinstance(request.get("run_id"), str) and bool(request["run_id"]), "Missing run identity") + need(request.get("intended_full_run_id", request["run_id"]) == request["run_id"] + and request.get("intended_full_output", request["output"]) == request["output"] + and request.get("intended_full_target_steps", 300000) == 300000, + "Full request contradicts intended identity") + return profile + + + +def verify_overrides(request): + """The narrow v1 recipe permits no additional scientific override on both arms.""" + arm, p = request["arm"], PROFILES[request["arm"]] + expected = [ + f"+experiment=pixel_proprio/native_joint/{arm}_image_proprio", "hydra/launcher=basic", + "data.manifest_path=" + request["data_manifest"], "paths.work_dir=" + request["source"], + "paths.output_dir=" + request["output"], "hydra.run.dir=" + request["output"], + "hydra.sweep.dir=" + request["output"], "hydra.job.chdir=false", + "+trainer.enable_progress_bar=false", "+trainer.enable_model_summary=false", + "logger=[wandb,csv]", "logger.wandb.id=" + request["run_id"], "logger.wandb.resume=never", + "logger.csv.name=csv", "logger.csv.flush_logs_every_n_steps=1", + "logger.csv.save_dir=" + request["local_logs"], "evaluator.pixel_diagnostics.log_wandb=true", + "modality_ablation.screening_budget_updates=300000", "trainer.devices=4", + "data.distributed_world_size=4", "data.batch_size=8", "model.optimizer.lr=1e-4", + "model.scheduler=null", "data.mode=" + arm, "data.sampling_strategy=" + p["strategy"], + "data.min_frames=2", "data.max_frames=64", "data.discrete_frames=16", "data.stride=4", + "joint.max_tokens=" + str(p["tokens"]), "joint.pose_dims.libero_gripper=" + str(p["pose"]), + "trainer.max_steps=300000", "callbacks.model_checkpoint.save_last=link", + ] + actual = request.get("overrides") + need(isinstance(actual, list) and all(isinstance(v, str) for v in actual) + and sorted(actual) == sorted(expected), "Additional or changed temporal v1 override") + + +def compose_request(request): + from hydra import compose, initialize_config_dir + from omegaconf import OmegaConf + with initialize_config_dir(config_dir=str(Path(request["source"])/"egomimic/hydra_configs"), + version_base=None): + cfg = compose(config_name="train_zarr_cartesian", overrides=request["overrides"]) + return OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + + +def normalize_config(request, cfg): + """Check every allowed substitution at its exact executed/resolved path.""" + profile, arm = PROFILES[request["arm"]], request["arm"] + cfg = copy.deepcopy(cfg) + need(verify_model_execution(request, cfg) == "visual_diffusion", "Wrong native graph") + need(cfg["seed"] == 42 and cfg["data"]["seed"] == 42 + and cfg["model"]["optimizer"]["lr"] == 1e-4 + and cfg["model"]["optimizer"]["weight_decay"] == 0.01 + and cfg["model"]["scheduler"] is None and not cfg["model"]["enable_grad_norm"] + and cfg["ckpt_path"] is None + and not cfg["model"]["pipeline"].get("initialization") + and not cfg["model"]["pipeline"].get("trainability"), "Wrong fresh optimizer recipe") + need(cfg["runtime"]["slurm_requeue_owner"] == "none" and cfg["runtime"]["disable_tf32"] + and cfg["trainer"]["precision"] == "bf16-mixed" + and cfg["trainer"]["devices"] == cfg["data"]["distributed_world_size"] == 4 + and cfg["data"]["batch_size"] == 8 and cfg["trainer"]["accumulate_grad_batches"] == 1, + "Wrong precision/topology/runtime") + need(set(cfg["callbacks"]) == {"model_checkpoint", "learning_rate_monitor", "wandb_profiler"} + and cfg["trainer"]["limit_val_batches"] == 1.0 + and cfg["trainer"]["limit_train_batches"] == 1.0 + and cfg["val_at_end"] is True, + "Unexpected checkpoint callback or partial/final-validation policy") + data = cfg["data"] + need(data["mode"] == arm and data["sampling_strategy"] == profile["strategy"] + and data["discrete_frames"] == 16 and data["stride"] == 4 + and data["min_frames"] == 2 and data["max_frames"] == 64 + and data["max_past_seconds"] is None and data["max_future_seconds"] is None + and data["checkpointable_sampling"] is True and data["couple_state_observations"] is True + and data["image_only"] is False and data["manifest_path"] == request["data_manifest"], + "Unexpected temporal sampling contract") + need(data["validation_sampling_strategy"] == "discrete_window" + and data["validation_discrete_frames"] == 16 and data["validation_stride"] == 4 + and data["validation_anchors_per_episode"] == 2 + and data["validation_batch_size"] == 1 + and data["validation_batches_per_embodiment"] is None + and cfg["evaluator"]["objective"] == "visual_proprio" + and cfg["evaluator"]["sample_seed"] == 42, + "Scheduled held-out validation changed") + data["mode"], data["sampling_strategy"] = "", "" + need(cfg["name"] == f"native_mot_{arm}_image_proprio_s42_300k", "Unbound recipe name") + cfg["name"] = "" + need(cfg["joint"]["pose_dims"] == {"libero_gripper": profile["pose"]} + and cfg["joint"]["max_tokens"] == profile["tokens"] + and cfg["pixel_video"]["max_tokens"] == profile["tokens"], "Wrong temporal dimensions") + cfg["joint"]["pose_dims"], cfg["joint"]["max_tokens"] = "", "" + cfg["pixel_video"]["max_tokens"] = "" + ids = cfg["model"]["pipeline"]["stage_ids"] + need(ids == {"joint_noising": 0, "joint_denoiser": 1, "image_loss": 2, + "proprio_loss": 3, "joint_sampler": 4} + and cfg["modality_ablation"]["stage_ids"] == ids, + "Unexpected native stage graph") + models = [cfg["model"]["pipeline"]["stages"][1]["model"], + cfg["modality_ablation"]["nodes"]["joint_denoiser"]["model"], + cfg["modality_ablation"]["graph"][1]["model"]] + image_models = [cfg["native_joint"]["reference"]["pipeline"]["stages"][1]["model"]] + for model in models: + need(model["pose_dims"] == {"libero_gripper": profile["pose"]} + and model["active_streams"] == ["image", "proprio"] + and model["proprio_width"] == 256, "Wrong resolved proprio model") + model["pose_dims"] = "" + image_models.append(model["image_model"]) + for model in image_models: + need(model["time_mode"] == arm and model["max_tokens"] == profile["tokens"], + "Unexpected image temporal representation") + model["time_mode"], model["max_tokens"] = "", "" + need(cfg["paths"]["output_dir"] == request["output"] + and cfg["logger"]["csv"]["save_dir"] == request["local_logs"] + and cfg["logger"]["wandb"]["id"] == request["run_id"], "Unbound output/log identity") + cfg["logger"]["wandb"]["id"] = "" + def paths(value): + if isinstance(value, dict): + return {key: paths(item) for key, item in value.items()} + if isinstance(value, list): + return [paths(item) for item in value] + if isinstance(value, str): + for old, replacement in ((request["output"], ""), + (request["local_logs"], "")): + if value == old or value.startswith(old + "/"): + return replacement + value[len(old):] + return value + return paths(cfg) + + +def disjoint(paths): + normalized = [Path(value) for value in paths] + need(all(p.is_absolute() and ".." not in p.parts for p in normalized), "Unsafe artifact path") + for i, left in enumerate(normalized): + need(all(left != right and left not in right.parents and right not in left.parents + for right in normalized[i+1:]), "Overlapping experiment output/log paths") + + +def verify(prototypes): + """Return reproducible evidence from two immutable full-request prototypes.""" + need(isinstance(prototypes, list) and len(prototypes) == 2, "Exactly two prototypes required") + pairs = [(evidence_binding(binding), bound(binding)) for binding in prototypes] + need({request.get("arm") for _, request in pairs} == set(PROFILES), "Missing or duplicate temporal arm") + pairs.sort(key=lambda row: row[1]["arm"]) + requests = [request for _, request in pairs] + for request in requests: + request_shape(request) + verify_overrides(request) + need(sha(request["source_manifest"]) == SOURCE_MANIFEST + and sha(request["data_manifest"]) == DATA_MANIFEST, "Source/data manifest bytes changed") + for key in ("source", "source_revision", "source_manifest", "source_manifest_sha256", + "source_archive", "source_archive_sha256", "data_manifest", "data_manifest_sha256", + "active_streams", "experiment", "world_size", "per_rank_batch", "global_batch", + "target_steps", "lr_policy", "model_execution", "operational_service", + "lambda_runtime", "runtime_localization", "data_cache", "allocation"): + need(requests[0].get(key) == requests[1].get(key), "Temporal control changed: " + key) + need(requests[0]["run_id"] != requests[1]["run_id"], "Duplicate full W&B identity") + need(requests[0]["smoke"]["job_id"] != requests[1]["smoke"]["job_id"], "Shared GPU smoke job") + outputs, image_hashes, records, normalized = [], [], [], [] + for binding, request in pairs: + completion, preflight = verify_evidence(request) + smoke_request = bound(request["smoke"]["request"]) + need(smoke_request.get("intended_full_run_id") == request["run_id"] + and smoke_request.get("intended_full_output") == request["output"] + and smoke_request.get("intended_full_target_steps") == request["target_steps"], + "Smoke belongs to another full experiment") + outputs.extend((request["output"], request["local_logs"], smoke_request["output"], smoke_request["local_logs"])) + from omegaconf import OmegaConf + smoke_cfg = OmegaConf.to_container(OmegaConf.create(bound(request["smoke"]["resolved"], False)), resolve=True) + cfg = compose_request(request) + compare_configs(cfg, smoke_cfg, request) + normalized.append(normalize_config(request, cfg)) + ranks = [] + for rank in completion["ranks"]: + value = rank["initial_image_weights_sha256"] + need(valid_hash(value) and valid_hash(rank["data_context_sha256"]) + and all(valid_hash(v) for v in rank["initial_weights"] + rank["paired_rgb_batch_sha256"]), + "Missing actual initial/batch/context hash") + image_hashes.append(value) + ranks.append({key: rank[key] for key in ( + "rank", "initial_image_weights_sha256", "initial_weights", "paired_rgb_batch_sha256", "data_context_sha256")}) + records.append(dict(arm=request["arm"], prototype=binding, + canonical_request_sha256=canonical_request(request), + resolved_sha256=digest(cfg), smoke_job_id=request["smoke"]["job_id"], + smoke_completion_sha256=request["smoke"]["completion"]["sha256"], + smoke_resolved_sha256=request["smoke"]["resolved"]["sha256"], + data_context_sha256=preflight["provenance"]["data_context_sha256"], + ranks=ranks)) + disjoint(outputs) + for request in requests: + disjoint([request["output"], request["local_logs"], request["source"], request["data_cache"]["root"]]) + need(len(image_hashes) == 8 and len(set(image_hashes)) == 1, "Actual initial image weights differ across ranks/arms") + need(normalized[0] == normalized[1], "Unexpected D/C scientific configuration difference") + return dict(schema=SCHEMA, status="PASS", source_revision=SOURCE, + source_manifest_sha256=SOURCE_MANIFEST, data_manifest_sha256=DATA_MANIFEST, + initial_image_weights_sha256=image_hashes[0], normalized_config_sha256=digest(normalized[0]), + fairness="FAIR WITH CAVEATS", + caveats=["Equal optimizer updates and global batch; temporal representation, frame/token exposure and compute differ.", + "Actual training draws, full-model initialization and data-context hashes are recorded separately, not paired."], + evidence=records) + + +def admit(request): + need(admission_mode(request) == "temporal_comparison", "Wrong temporal admission mode") + need(not any(name in request for name in ("continuation", "allocation_stop_step", "allocation_boundaries", "allocation_recovery")), + "Temporal v1 supports fresh single-allocation requests only") + receipt = bound(evidence_binding(request["temporal_comparison"])) + need(receipt.get("schema") == SCHEMA and receipt.get("status") == "PASS" + and len(receipt.get("evidence", [])) == 2, "Missing temporal GPU evidence") + expected = verify([row["prototype"] for row in receipt["evidence"]]) + need(receipt == expected, "Temporal receipt differs from freshly verified evidence") + members = [row for row in receipt["evidence"] + if row["arm"] == request["arm"] and row["canonical_request_sha256"] == canonical_request(request)] + need(len(members) == 1, "Full request is not an exact admitted prototype member") + return receipt + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("requests", nargs=2, type=Path, help="Two immutable full prototypes without temporal_comparison") + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + prototypes = [dict(path=str(path.resolve()), sha256=sha(path)) for path in args.requests] + write_once(args.output, verify(prototypes)) + + +if __name__ == "__main__": + main() diff --git a/tests/fixtures/joint_crash_recovery_535d/image-only-parent.request.json b/tests/fixtures/joint_crash_recovery_535d/image-only-parent.request.json new file mode 100644 index 000000000..5778461d2 --- /dev/null +++ b/tests/fixtures/joint_crash_recovery_535d/image-only-parent.request.json @@ -0,0 +1,55 @@ +{ + "fixture_only": "Synthetic path/hash bindings; historical image-only constant-LR scientific identity, not launch evidence", + "schema": "joint-training-v2", + "phase": "full", + "source_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "run_id": "mot-shared-libero-image-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "arm": "discrete", + "experiment": "image_only", + "active_streams": [ + "image" + ], + "parameters": 88178688, + "expert_parameters": { + "image": 88178688 + }, + "expert_widths": [ + 768 + ], + "world_size": 4, + "global_batch": 32, + "per_rank_batch": 8, + "target_steps": 300000, + "lr_policy": { + "schedule": "constant", + "peak_lr": 0.0001, + "warmup_optimizer_steps": 0 + }, + "source": "/synthetic/source", + "source_manifest": "/synthetic/source/FILES.sha256", + "source_manifest_sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "source_archive": "/synthetic/source.tar.gz", + "source_archive_sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "data_manifest": "/synthetic/data/manifest.json", + "data_manifest_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "data_cache": { + "root": "/synthetic/cache" + }, + "lambda_runtime": { + "node": "gpu19" + }, + "runtime_localization": null, + "allocation": null, + "model_execution": { + "precision": "bf16-mixed" + }, + "job_name": "joint-image-lr1e4-full-4gpu", + "output": "/synthetic/parent", + "local_logs": "/synthetic/parent-logs", + "overrides": [ + "model.optimizer.lr=1e-4", + "paths.output_dir=/synthetic/parent", + "logger.csv.save_dir=/synthetic/parent-logs" + ], + "reserved_checkpoint_bytes": 40000000000 +} diff --git a/tests/fixtures/joint_temporal_pair_535d/continuous-full.composed.json b/tests/fixtures/joint_temporal_pair_535d/continuous-full.composed.json new file mode 100644 index 000000000..40b3b6cb2 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/continuous-full.composed.json @@ -0,0 +1,598 @@ +{ + "callbacks": { + "learning_rate_monitor": { + "_target_": "lightning.pytorch.callbacks.LearningRateMonitor", + "logging_interval": "step" + }, + "model_checkpoint": { + "_target_": "lightning.pytorch.callbacks.ModelCheckpoint", + "dirpath": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu/checkpoints", + "every_n_epochs": 0, + "every_n_train_steps": 20000, + "filename": "epoch_{epoch:04d}-step_{step:09d}", + "save_last": "link", + "save_on_train_epoch_end": false, + "save_top_k": -1 + }, + "wandb_profiler": { + "_target_": "egomimic.utils.timing_callback.WandbProfilerLogger", + "log_every_n_steps": 100 + } + }, + "ckpt_path": null, + "data": { + "_target_": "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule", + "balance_training_work": true, + "batch_size": 8, + "checkpointable_sampling": true, + "couple_state_observations": true, + "discrete_frames": 16, + "distributed_world_size": 4, + "forecasting_probability": 0.5, + "image_only": false, + "manifest_path": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "max_frames": 64, + "max_future_seconds": null, + "max_past_seconds": null, + "min_frames": 2, + "mode": "continuous", + "num_workers": 2, + "persistent_workers": true, + "pin_memory": true, + "prefetch_factor": 2, + "sampling_strategy": "full_episode_uniform", + "seed": 42, + "stride": 4, + "task": "visual_diffusion", + "validation_anchors_per_episode": 2, + "validation_batch_size": 1, + "validation_batches_per_embodiment": null, + "validation_discrete_frames": 16, + "validation_sampling_strategy": "discrete_window", + "validation_stride": 4 + }, + "data_context_loader": { + "_target_": "egomimic.rldb.zarr.data_module.load_data_context" + }, + "description": "Historical image core with minimal joint state extension, historical compilation preserved", + "evaluator": { + "_target_": "egomimic.eval.image_proprio_eval.ImageProprioEval", + "objective": "visual_proprio", + "pixel_diagnostics": { + "artifact_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu/pixel_diagnostics", + "log_wandb": true, + "max_samples_per_source_kind": 1, + "playback_fps": 20, + "write_video": true + }, + "sample_seed": 42 + }, + "inference_config": { + "enabled": false, + "output_path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu/checkpoints/inference-config.yaml" + }, + "joint": { + "compile_mode": "default", + "max_tokens": 12864, + "pose_dims": { + "libero_gripper": 9 + } + }, + "launch_params": { + "gpus_per_node": 1, + "nodes": 1 + }, + "logger": { + "csv": { + "_target_": "lightning.pytorch.loggers.CSVLogger", + "flush_logs_every_n_steps": 1, + "name": "csv", + "save_dir": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "version": 0 + }, + "wandb": { + "_target_": "lightning.pytorch.loggers.wandb.WandbLogger", + "anonymous": null, + "entity": "rl2-group", + "group": "native-image-joint-state", + "id": "mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1", + "job_type": "", + "log_model": false, + "offline": false, + "prefix": "", + "project": "mot-world-model", + "resume": "never", + "save_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "tags": [ + "historical-image-core", + "native-compiled", + "no-action", + "paired-initialization" + ] + } + }, + "modality_ablation": { + "active_streams": [ + "image", + "proprio" + ], + "graph": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ], + "initialization": "Native image construction first; vector initialization preserves CPU RNG", + "nodes": { + "image_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + "joint_denoiser": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + "joint_noising": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + "joint_sampler": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + }, + "proprio_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + } + }, + "screening_budget_updates": 300000, + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "training_noise": "Native image draws unchanged; optional tied proprio draws preserve image RNG" + }, + "mode": "train", + "model": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 0.0001, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ] + }, + "scheduler": null, + "scheduler_frequency": 1, + "scheduler_interval": "step", + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior", + "distributed_noise_seed": 42 + } + }, + "name": "native_mot_continuous_image_proprio_s42_300k", + "native_joint": { + "reference": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 5e-05, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_proprio": 1, + "pixel_denoiser": 1, + "pixel_loss": 2, + "pixel_noising": 0, + "pixel_sampler": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage", + "noising": { + "_target_": "egomimic.models.pixel_diffusion.LinearGaussianImageNoising", + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoDenoiserStage", + "model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoSamplerStage", + "sampler": { + "_target_": "egomimic.models.pixel_diffusion.LinearDDIMImageSampler", + "steps": 20 + } + } + ] + }, + "scheduler": null, + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior" + } + } + }, + "norm_stats": null, + "norm_stats_only": false, + "normalizer": null, + "paths": { + "dataset_dir": "/coc/flash7/scratch/egoverseS3ZarrDataset", + "output_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "work_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" + }, + "pixel_video": { + "max_tokens": 12864 + }, + "runtime": { + "cluster_environment": "torch_elastic", + "disable_tf32": true, + "slurm_requeue_owner": "none", + "slurm_save_signal": "SIGUSR2", + "slurm_signal_checkpoint_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu/checkpoints", + "slurm_signal_train_batch_only": false, + "strict_fp32": false + }, + "seed": 42, + "trainer": { + "_target_": "lightning.pytorch.trainer.Trainer", + "accelerator": "gpu", + "accumulate_grad_batches": 1, + "check_val_every_n_epoch": null, + "default_root_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "deterministic": false, + "devices": 4, + "enable_model_summary": false, + "enable_progress_bar": false, + "limit_train_batches": 1.0, + "limit_val_batches": 1.0, + "log_every_n_steps": 1, + "max_epochs": -1, + "max_steps": 300000, + "min_epochs": 0, + "num_nodes": 1, + "num_sanity_val_steps": 0, + "precision": "bf16-mixed", + "strategy": "ddp", + "use_distributed_sampler": false, + "val_check_interval": 20000 + }, + "val_at_end": true, + "world_model_attention": { + "activation_checkpointing": false, + "backend": "flash4", + "disable_tf32": true, + "dtype": "bfloat16", + "query_chunk_size": 0, + "strict_fp32": false, + "trainer_precision": "bf16-mixed" + }, + "world_model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "loss_reduction": "mean_per_example", + "token_execution": "packed_valid" + } +} diff --git a/tests/fixtures/joint_temporal_pair_535d/continuous-full.template.json b/tests/fixtures/joint_temporal_pair_535d/continuous-full.template.json new file mode 100644 index 000000000..b773a3726 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/continuous-full.template.json @@ -0,0 +1,183 @@ +{ + "_pending_bindings": [ + "Pin maintained service closure and manifest; confirm temporal-pair full admission implementation.", + "Fresh native gpu03 allocation/cache/runtime readback; existing bound producer evidence is a reuse candidate, not fresh node certification.", + "Verify both new W&B identities absent, exact queues/claims/outputs empty, destination retention space.", + "Resolve both exact smoke/full configs in scheduled CPU allocation and record actual meta-model counts.", + "Both exact new4-rank smoke strict PASS + host0 + nativeCOMPLETED0:0, including50 actual1e-4 LR updates/all ranks,100 scheduled validation anchors,FP32/fullstate/hiddenNaN/capacity proofs.", + "New maintained temporal-pair admission receipt, common initial image hash and allowed D/C config differences. No optimizer parent or modality factorial surrogate.", + "Report measured per-arm smoke timings before either full submission." + ], + "_template_status": "NON_SUBMITTABLE_PENDING_EXACT_BINDINGS", + "active_streams": [ + "image", + "proprio" + ], + "allocation": null, + "arm": "continuous", + "data_cache": { + "job_id": 411978, + "native_receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-native.txt", + "native_receipt_sha256": "ad08aa8efbd7ad8e74587aea8e839650d10e7757890c715bfd03f55dc79df3a0", + "node": "gpu03", + "parent": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data", + "receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-receipt.json", + "receipt_sha256": "008d9947db54e14d7fee2d194e5e89786892d305b1edecd4e9ea209b5d9ba4a5", + "root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data/411978" + }, + "data_manifest": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "data_manifest_sha256": "8c40d792a5cedd4105ff8270ff3141767b147513c52f11292f888d9b1527bfe2", + "experiment": "image_proprio", + "expert_parameters": { + "image": 88178688, + "proprio": 13942537 + }, + "expert_widths": [ + 768, + 256 + ], + "global_batch": 32, + "intended_full_output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "intended_full_run_id": "mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1", + "intended_full_target_steps": 300000, + "job_name": "joint-continuous-image-proprio-lr1e4-full-4gpu", + "lambda_runtime": { + "actual_python": "/workspace/users/paphiwetsa3/tasks/yam-arc400-lambda-20260921/runtime/venv-pinned/bin/python", + "bootstrap_python": "/workspace/users/paphiwetsa3/tasks/arc-global-basis-training-20260924/runtime-r6a/bin/python", + "container_image": "nvcr.io#nvidia/pytorch:25.06-py3", + "extra_pythonpath": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/runtime-overlays/dino4571/site-packages", + "fa3_root": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1/runtime/fa3", + "fa4_root": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/training-runtime/profile/flash-attention", + "minimum_free_scratch_bytes": 10000000000, + "mounts": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-mounts.csv", + "sha256": "b2835189da9504de5263e565bc5a1e827cfa6ce9b788a42e5aed02daca9bc41f" + }, + "node": "gpu03", + "runtime_task": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1", + "schema": "joint-lambda-runtime-v1", + "scratch_root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020", + "temporary_root": "/ephemeral/enroot/41003" + }, + "local_logs": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "lr_policy": { + "peak_lr": 0.0001, + "schedule": "constant", + "warmup_optimizer_steps": 0 + }, + "model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "implementation": "native_image_joint" + }, + "operational_change": "Fresh continuous image+proprio constant AdamW1e-4 experiment from source535d: four dedicated GPUs,8/rank,global32,seed42,300000 optimizer updates,20k checkpoint/validation cadence,retain all15 immutable checkpoints. No action expert, checkpoint initialization, continuation or allocation boundary.", + "operational_service": { + "manifest": "/UNBOUND/operational-service/FILES.sha256", + "manifest_sha256": "UNBOUND_SERVICE_MANIFEST_SHA256", + "revision": "UNBOUND_MAINTAINED_SERVICE_REVISION", + "root": "/UNBOUND/operational-service" + }, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "overrides": [ + "+experiment=pixel_proprio/native_joint/continuous_image_proprio", + "hydra/launcher=basic", + "data.manifest_path=/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "paths.work_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "paths.output_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "hydra.run.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "hydra.sweep.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "hydra.job.chdir=false", + "+trainer.enable_progress_bar=false", + "+trainer.enable_model_summary=false", + "logger=[wandb,csv]", + "logger.wandb.id=mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1", + "logger.wandb.resume=never", + "logger.csv.name=csv", + "logger.csv.flush_logs_every_n_steps=1", + "logger.csv.save_dir=/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "evaluator.pixel_diagnostics.log_wandb=true", + "modality_ablation.screening_budget_updates=300000", + "trainer.devices=4", + "data.distributed_world_size=4", + "data.batch_size=8", + "model.optimizer.lr=1e-4", + "model.scheduler=null", + "data.mode=continuous", + "data.sampling_strategy=full_episode_uniform", + "data.min_frames=2", + "data.max_frames=64", + "data.discrete_frames=16", + "data.stride=4", + "joint.max_tokens=12864", + "joint.pose_dims.libero_gripper=9", + "trainer.max_steps=300000", + "callbacks.model_checkpoint.save_last=link" + ], + "parameters": 102121225, + "per_rank_batch": 8, + "phase": "full", + "quota": { + "kind": "filesystem", + "mount": "/workspace", + "node": "gpu03", + "uid": 41003, + "user": "paphiwetsa3" + }, + "reserved_checkpoint_bytes": 40000000000, + "run_id": "mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1", + "runtime_localization": { + "admission": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-admission.json", + "sha256": "f5dbe6abf5aa0bdab9da7b96dbf65dd84c9ed7dc261e2bffc28694d882056670" + }, + "archive_sha256": "92e461b28c8659b31b664e8899ae2e390bd6fdf6d85d575cef77fec5a6956fa8", + "destination_node": "gpu03", + "files_manifest_sha256": "c8bddcf4b2f998adae085f96a263356ad1e6392ce93f7c23a3c6d392e2214289", + "schema": "readonly-runtime-export-v1", + "source_gpu_job": "411963", + "source_node": "gpu21" + }, + "schema": "joint-training-v2", + "smoke": { + "capacity": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/continuous-image-proprio-lr1e4-smoke.capacity.json", + "sha256": "UNBOUND_EXACT_SMOKE_CAPACITY_SHA256" + }, + "completion": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/smoke-completion.json", + "sha256": "UNBOUND_EXACT_SMOKE_COMPLETION_SHA256" + }, + "host_exit": { + "path": "/UNBOUND/exact-smoke-host-exit", + "sha256": "UNBOUND_EXACT_SMOKE_HOST_EXIT_SHA256" + }, + "job_id": "UNBOUND_EXACT_CONTINUOUS_SMOKE_JOB", + "native": { + "path": "/UNBOUND/exact-smoke-native-completion", + "sha256": "UNBOUND_EXACT_SMOKE_NATIVE_SHA256" + }, + "preflight": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/continuous-image-proprio-lr1e4-smoke.preflight.json", + "sha256": "UNBOUND_EXACT_SMOKE_PREFLIGHT_SHA256" + }, + "request": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/continuous-image-proprio-lr1e4-smoke.json", + "sha256": "UNBOUND_EXACT_SMOKE_REQUEST_SHA256" + }, + "resolved": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/continuous-image-proprio-lr1e4-smoke.resolved.yaml", + "sha256": "UNBOUND_EXACT_SMOKE_RESOLVED_SHA256" + } + }, + "source": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "source_archive": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.tar.gz", + "source_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "source_manifest": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.files.sha256", + "source_manifest_sha256": "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184", + "source_parent": "3785f94392b9d900e21d9948fbc5bf94a6d975b8", + "source_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "target_steps": 300000, + "vector_expert_parameter_budget": 20000000, + "world_size": 4 +} diff --git a/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.composed.json b/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.composed.json new file mode 100644 index 000000000..e672ee2c9 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.composed.json @@ -0,0 +1,585 @@ +{ + "callbacks": { + "joint_smoke": { + "_target_": "egomimic.pl_utils.joint_smoke.JointSmokeProof", + "expected_steps": 50, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "warmup_steps": 10, + "world_size": 4 + }, + "learning_rate_monitor": { + "_target_": "lightning.pytorch.callbacks.LearningRateMonitor", + "logging_interval": "step" + }, + "model_checkpoint": { + "_target_": "lightning.pytorch.callbacks.ModelCheckpoint", + "dirpath": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints", + "every_n_epochs": 0, + "every_n_train_steps": 0, + "filename": "epoch_{epoch:04d}-step_{step:09d}", + "save_last": false, + "save_on_train_epoch_end": false, + "save_top_k": -1 + }, + "wandb_profiler": { + "_target_": "egomimic.utils.timing_callback.WandbProfilerLogger", + "log_every_n_steps": 100 + } + }, + "ckpt_path": null, + "data": { + "_target_": "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule", + "balance_training_work": true, + "batch_size": 8, + "checkpointable_sampling": true, + "couple_state_observations": true, + "discrete_frames": 16, + "distributed_world_size": 4, + "forecasting_probability": 0.5, + "image_only": false, + "manifest_path": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "max_frames": 64, + "max_future_seconds": null, + "max_past_seconds": null, + "min_frames": 2, + "mode": "continuous", + "num_workers": 2, + "persistent_workers": true, + "pin_memory": true, + "prefetch_factor": 2, + "sampling_strategy": "full_episode_uniform", + "seed": 42, + "stride": 4, + "task": "visual_diffusion", + "validation_anchors_per_episode": 2, + "validation_batch_size": 1, + "validation_batches_per_embodiment": null, + "validation_discrete_frames": 16, + "validation_sampling_strategy": "discrete_window", + "validation_stride": 4 + }, + "data_context_loader": { + "_target_": "egomimic.rldb.zarr.data_module.load_data_context" + }, + "description": "Historical image core with minimal joint state extension, historical compilation preserved", + "evaluator": { + "_target_": "egomimic.eval.image_proprio_eval.ImageProprioEval", + "objective": "visual_proprio", + "pixel_diagnostics": { + "artifact_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/pixel_diagnostics", + "log_wandb": false, + "max_samples_per_source_kind": 1, + "playback_fps": 20, + "write_video": true + }, + "sample_seed": 42 + }, + "inference_config": { + "enabled": false, + "output_path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints/inference-config.yaml" + }, + "joint": { + "compile_mode": "default", + "max_tokens": 12864, + "pose_dims": { + "libero_gripper": 9 + } + }, + "launch_params": { + "gpus_per_node": 1, + "nodes": 1 + }, + "logger": { + "csv": { + "_target_": "lightning.pytorch.loggers.CSVLogger", + "flush_logs_every_n_steps": 1, + "name": "csv", + "save_dir": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-continuous-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "version": 0 + } + }, + "modality_ablation": { + "active_streams": [ + "image", + "proprio" + ], + "graph": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ], + "initialization": "Native image construction first; vector initialization preserves CPU RNG", + "nodes": { + "image_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + "joint_denoiser": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + "joint_noising": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + "joint_sampler": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + }, + "proprio_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + } + }, + "screening_budget_updates": 300000, + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "training_noise": "Native image draws unchanged; optional tied proprio draws preserve image RNG" + }, + "mode": "train", + "model": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 0.0001, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 9 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ] + }, + "scheduler": null, + "scheduler_frequency": 1, + "scheduler_interval": "step", + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior", + "distributed_noise_seed": 42 + } + }, + "name": "native_mot_continuous_image_proprio_s42_300k", + "native_joint": { + "reference": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 5e-05, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_proprio": 1, + "pixel_denoiser": 1, + "pixel_loss": 2, + "pixel_noising": 0, + "pixel_sampler": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage", + "noising": { + "_target_": "egomimic.models.pixel_diffusion.LinearGaussianImageNoising", + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoDenoiserStage", + "model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 12864, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "continuous", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoSamplerStage", + "sampler": { + "_target_": "egomimic.models.pixel_diffusion.LinearDDIMImageSampler", + "steps": 20 + } + } + ] + }, + "scheduler": null, + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior" + } + } + }, + "norm_stats": null, + "norm_stats_only": false, + "normalizer": null, + "paths": { + "dataset_dir": "/coc/flash7/scratch/egoverseS3ZarrDataset", + "output_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "work_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" + }, + "pixel_video": { + "max_tokens": 12864 + }, + "runtime": { + "cluster_environment": "torch_elastic", + "disable_tf32": true, + "slurm_requeue_owner": "none", + "slurm_save_signal": "SIGUSR2", + "slurm_signal_checkpoint_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints", + "slurm_signal_train_batch_only": false, + "strict_fp32": false + }, + "seed": 42, + "trainer": { + "_target_": "lightning.pytorch.trainer.Trainer", + "accelerator": "gpu", + "accumulate_grad_batches": 1, + "check_val_every_n_epoch": null, + "default_root_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "deterministic": false, + "devices": 4, + "enable_model_summary": false, + "enable_progress_bar": false, + "limit_train_batches": 1.0, + "limit_val_batches": 1.0, + "log_every_n_steps": 1, + "max_epochs": -1, + "max_steps": 50, + "min_epochs": 0, + "num_nodes": 1, + "num_sanity_val_steps": 0, + "precision": "bf16-mixed", + "strategy": "ddp", + "use_distributed_sampler": false, + "val_check_interval": 50 + }, + "val_at_end": true, + "world_model_attention": { + "activation_checkpointing": false, + "backend": "flash4", + "disable_tf32": true, + "dtype": "bfloat16", + "query_chunk_size": 0, + "strict_fp32": false, + "trainer_precision": "bf16-mixed" + }, + "world_model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "loss_reduction": "mean_per_example", + "token_execution": "packed_valid" + } +} diff --git a/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.template.json b/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.template.json new file mode 100644 index 000000000..2263101dd --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/continuous-smoke.template.json @@ -0,0 +1,139 @@ +{ + "_pending_bindings": [ + "Pin maintained service closure and manifest; confirm temporal-pair full admission implementation.", + "Fresh native gpu03 allocation/cache/runtime readback; existing bound producer evidence is a reuse candidate, not fresh node certification.", + "Verify both new W&B identities absent, exact queues/claims/outputs empty, destination retention space.", + "Resolve both exact smoke/full configs in scheduled CPU allocation and record actual meta-model counts." + ], + "_template_status": "NON_SUBMITTABLE_PENDING_EXACT_BINDINGS", + "active_streams": [ + "image", + "proprio" + ], + "allocation": null, + "arm": "continuous", + "data_cache": { + "job_id": 411978, + "native_receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-native.txt", + "native_receipt_sha256": "ad08aa8efbd7ad8e74587aea8e839650d10e7757890c715bfd03f55dc79df3a0", + "node": "gpu03", + "parent": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data", + "receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-receipt.json", + "receipt_sha256": "008d9947db54e14d7fee2d194e5e89786892d305b1edecd4e9ea209b5d9ba4a5", + "root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data/411978" + }, + "data_manifest": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "data_manifest_sha256": "8c40d792a5cedd4105ff8270ff3141767b147513c52f11292f888d9b1527bfe2", + "experiment": "image_proprio", + "expert_parameters": { + "image": 88178688, + "proprio": 13942537 + }, + "expert_widths": [ + 768, + 256 + ], + "global_batch": 32, + "intended_full_output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_continuous_image_proprio_lr1e4constant_4gpu", + "intended_full_run_id": "mot-shared-libero-image-proprio-continuous-lr1e-4-constant-s42-300k-20261021-v1", + "intended_full_target_steps": 300000, + "job_name": "joint-continuous-image-proprio-lr1e4-smoke-4gpu", + "lambda_runtime": { + "actual_python": "/workspace/users/paphiwetsa3/tasks/yam-arc400-lambda-20260921/runtime/venv-pinned/bin/python", + "bootstrap_python": "/workspace/users/paphiwetsa3/tasks/arc-global-basis-training-20260924/runtime-r6a/bin/python", + "container_image": "nvcr.io#nvidia/pytorch:25.06-py3", + "extra_pythonpath": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/runtime-overlays/dino4571/site-packages", + "fa3_root": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1/runtime/fa3", + "fa4_root": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/training-runtime/profile/flash-attention", + "minimum_free_scratch_bytes": 10000000000, + "mounts": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-mounts.csv", + "sha256": "b2835189da9504de5263e565bc5a1e827cfa6ce9b788a42e5aed02daca9bc41f" + }, + "node": "gpu03", + "runtime_task": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1", + "schema": "joint-lambda-runtime-v1", + "scratch_root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020", + "temporary_root": "/ephemeral/enroot/41003" + }, + "local_logs": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-continuous-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "lr_policy": { + "peak_lr": 0.0001, + "schedule": "constant", + "warmup_optimizer_steps": 0 + }, + "model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "implementation": "native_image_joint" + }, + "operational_change": "Fresh continuous image+proprio constant AdamW1e-4 experiment from source535d: four dedicated GPUs,8/rank,global32,seed42,300000 optimizer updates,20k checkpoint/validation cadence,retain all15 immutable checkpoints. No action expert, checkpoint initialization, continuation or allocation boundary.", + "operational_service": { + "manifest": "/UNBOUND/operational-service/FILES.sha256", + "manifest_sha256": "UNBOUND_SERVICE_MANIFEST_SHA256", + "revision": "UNBOUND_MAINTAINED_SERVICE_REVISION", + "root": "/UNBOUND/operational-service" + }, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "overrides": [ + "+experiment=pixel_proprio/native_joint/continuous_image_proprio", + "+joint_execution=smoke", + "hydra/launcher=basic", + "data.manifest_path=/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "paths.work_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "paths.output_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.run.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.sweep.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-continuous-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.job.chdir=false", + "+trainer.enable_progress_bar=false", + "+trainer.enable_model_summary=false", + "logger=[csv]", + "~logger.wandb", + "logger.csv.name=csv", + "logger.csv.flush_logs_every_n_steps=1", + "logger.csv.save_dir=/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-continuous-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "evaluator.pixel_diagnostics.log_wandb=false", + "modality_ablation.screening_budget_updates=300000", + "trainer.devices=4", + "data.distributed_world_size=4", + "data.batch_size=8", + "model.optimizer.lr=1e-4", + "model.scheduler=null", + "data.mode=continuous", + "data.sampling_strategy=full_episode_uniform", + "data.min_frames=2", + "data.max_frames=64", + "data.discrete_frames=16", + "data.stride=4", + "joint.max_tokens=12864", + "joint.pose_dims.libero_gripper=9" + ], + "parameters": 102121225, + "per_rank_batch": 8, + "phase": "smoke", + "reserved_checkpoint_bytes": 1600000000, + "runtime_localization": { + "admission": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-admission.json", + "sha256": "f5dbe6abf5aa0bdab9da7b96dbf65dd84c9ed7dc261e2bffc28694d882056670" + }, + "archive_sha256": "92e461b28c8659b31b664e8899ae2e390bd6fdf6d85d575cef77fec5a6956fa8", + "destination_node": "gpu03", + "files_manifest_sha256": "c8bddcf4b2f998adae085f96a263356ad1e6392ce93f7c23a3c6d392e2214289", + "schema": "readonly-runtime-export-v1", + "source_gpu_job": "411963", + "source_node": "gpu21" + }, + "schema": "joint-training-v2", + "source": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "source_archive": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.tar.gz", + "source_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "source_manifest": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.files.sha256", + "source_manifest_sha256": "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184", + "source_parent": "3785f94392b9d900e21d9948fbc5bf94a6d975b8", + "source_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "target_steps": 50, + "vector_expert_parameter_budget": 20000000, + "world_size": 4 +} diff --git a/tests/fixtures/joint_temporal_pair_535d/discrete-full.composed.json b/tests/fixtures/joint_temporal_pair_535d/discrete-full.composed.json new file mode 100644 index 000000000..1f059eb8d --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/discrete-full.composed.json @@ -0,0 +1,598 @@ +{ + "callbacks": { + "learning_rate_monitor": { + "_target_": "lightning.pytorch.callbacks.LearningRateMonitor", + "logging_interval": "step" + }, + "model_checkpoint": { + "_target_": "lightning.pytorch.callbacks.ModelCheckpoint", + "dirpath": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu/checkpoints", + "every_n_epochs": 0, + "every_n_train_steps": 20000, + "filename": "epoch_{epoch:04d}-step_{step:09d}", + "save_last": "link", + "save_on_train_epoch_end": false, + "save_top_k": -1 + }, + "wandb_profiler": { + "_target_": "egomimic.utils.timing_callback.WandbProfilerLogger", + "log_every_n_steps": 100 + } + }, + "ckpt_path": null, + "data": { + "_target_": "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule", + "balance_training_work": true, + "batch_size": 8, + "checkpointable_sampling": true, + "couple_state_observations": true, + "discrete_frames": 16, + "distributed_world_size": 4, + "forecasting_probability": 0.5, + "image_only": false, + "manifest_path": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "max_frames": 64, + "max_future_seconds": null, + "max_past_seconds": null, + "min_frames": 2, + "mode": "discrete", + "num_workers": 2, + "persistent_workers": true, + "pin_memory": true, + "prefetch_factor": 2, + "sampling_strategy": "discrete_window", + "seed": 42, + "stride": 4, + "task": "visual_diffusion", + "validation_anchors_per_episode": 2, + "validation_batch_size": 1, + "validation_batches_per_embodiment": null, + "validation_discrete_frames": 16, + "validation_sampling_strategy": "discrete_window", + "validation_stride": 4 + }, + "data_context_loader": { + "_target_": "egomimic.rldb.zarr.data_module.load_data_context" + }, + "description": "Historical image core with minimal joint state extension, historical compilation preserved", + "evaluator": { + "_target_": "egomimic.eval.image_proprio_eval.ImageProprioEval", + "objective": "visual_proprio", + "pixel_diagnostics": { + "artifact_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu/pixel_diagnostics", + "log_wandb": true, + "max_samples_per_source_kind": 1, + "playback_fps": 20, + "write_video": true + }, + "sample_seed": 42 + }, + "inference_config": { + "enabled": false, + "output_path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu/checkpoints/inference-config.yaml" + }, + "joint": { + "compile_mode": "default", + "max_tokens": 3216, + "pose_dims": { + "libero_gripper": 36 + } + }, + "launch_params": { + "gpus_per_node": 1, + "nodes": 1 + }, + "logger": { + "csv": { + "_target_": "lightning.pytorch.loggers.CSVLogger", + "flush_logs_every_n_steps": 1, + "name": "csv", + "save_dir": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "version": 0 + }, + "wandb": { + "_target_": "lightning.pytorch.loggers.wandb.WandbLogger", + "anonymous": null, + "entity": "rl2-group", + "group": "native-image-joint-state", + "id": "mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "job_type": "", + "log_model": false, + "offline": false, + "prefix": "", + "project": "mot-world-model", + "resume": "never", + "save_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "tags": [ + "historical-image-core", + "native-compiled", + "no-action", + "paired-initialization" + ] + } + }, + "modality_ablation": { + "active_streams": [ + "image", + "proprio" + ], + "graph": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ], + "initialization": "Native image construction first; vector initialization preserves CPU RNG", + "nodes": { + "image_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + "joint_denoiser": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + "joint_noising": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + "joint_sampler": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + }, + "proprio_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + } + }, + "screening_budget_updates": 300000, + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "training_noise": "Native image draws unchanged; optional tied proprio draws preserve image RNG" + }, + "mode": "train", + "model": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 0.0001, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ] + }, + "scheduler": null, + "scheduler_frequency": 1, + "scheduler_interval": "step", + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior", + "distributed_noise_seed": 42 + } + }, + "name": "native_mot_discrete_image_proprio_s42_300k", + "native_joint": { + "reference": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 5e-05, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_proprio": 1, + "pixel_denoiser": 1, + "pixel_loss": 2, + "pixel_noising": 0, + "pixel_sampler": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage", + "noising": { + "_target_": "egomimic.models.pixel_diffusion.LinearGaussianImageNoising", + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoDenoiserStage", + "model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoSamplerStage", + "sampler": { + "_target_": "egomimic.models.pixel_diffusion.LinearDDIMImageSampler", + "steps": 20 + } + } + ] + }, + "scheduler": null, + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior" + } + } + }, + "norm_stats": null, + "norm_stats_only": false, + "normalizer": null, + "paths": { + "dataset_dir": "/coc/flash7/scratch/egoverseS3ZarrDataset", + "output_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "work_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" + }, + "pixel_video": { + "max_tokens": 3216 + }, + "runtime": { + "cluster_environment": "torch_elastic", + "disable_tf32": true, + "slurm_requeue_owner": "none", + "slurm_save_signal": "SIGUSR2", + "slurm_signal_checkpoint_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu/checkpoints", + "slurm_signal_train_batch_only": false, + "strict_fp32": false + }, + "seed": 42, + "trainer": { + "_target_": "lightning.pytorch.trainer.Trainer", + "accelerator": "gpu", + "accumulate_grad_batches": 1, + "check_val_every_n_epoch": null, + "default_root_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "deterministic": false, + "devices": 4, + "enable_model_summary": false, + "enable_progress_bar": false, + "limit_train_batches": 1.0, + "limit_val_batches": 1.0, + "log_every_n_steps": 1, + "max_epochs": -1, + "max_steps": 300000, + "min_epochs": 0, + "num_nodes": 1, + "num_sanity_val_steps": 0, + "precision": "bf16-mixed", + "strategy": "ddp", + "use_distributed_sampler": false, + "val_check_interval": 20000 + }, + "val_at_end": true, + "world_model_attention": { + "activation_checkpointing": false, + "backend": "flash4", + "disable_tf32": true, + "dtype": "bfloat16", + "query_chunk_size": 0, + "strict_fp32": false, + "trainer_precision": "bf16-mixed" + }, + "world_model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "loss_reduction": "mean_per_example", + "token_execution": "packed_valid" + } +} diff --git a/tests/fixtures/joint_temporal_pair_535d/discrete-full.template.json b/tests/fixtures/joint_temporal_pair_535d/discrete-full.template.json new file mode 100644 index 000000000..c14ae7ad4 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/discrete-full.template.json @@ -0,0 +1,183 @@ +{ + "_pending_bindings": [ + "Pin maintained service closure and manifest; confirm temporal-pair full admission implementation.", + "Fresh native gpu03 allocation/cache/runtime readback; existing bound producer evidence is a reuse candidate, not fresh node certification.", + "Verify both new W&B identities absent, exact queues/claims/outputs empty, destination retention space.", + "Resolve both exact smoke/full configs in scheduled CPU allocation and record actual meta-model counts.", + "Both exact new4-rank smoke strict PASS + host0 + nativeCOMPLETED0:0, including50 actual1e-4 LR updates/all ranks,100 scheduled validation anchors,FP32/fullstate/hiddenNaN/capacity proofs.", + "New maintained temporal-pair admission receipt, common initial image hash and allowed D/C config differences. No optimizer parent or modality factorial surrogate.", + "Report measured per-arm smoke timings before either full submission." + ], + "_template_status": "NON_SUBMITTABLE_PENDING_EXACT_BINDINGS", + "active_streams": [ + "image", + "proprio" + ], + "allocation": null, + "arm": "discrete", + "data_cache": { + "job_id": 411978, + "native_receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-native.txt", + "native_receipt_sha256": "ad08aa8efbd7ad8e74587aea8e839650d10e7757890c715bfd03f55dc79df3a0", + "node": "gpu03", + "parent": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data", + "receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-receipt.json", + "receipt_sha256": "008d9947db54e14d7fee2d194e5e89786892d305b1edecd4e9ea209b5d9ba4a5", + "root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data/411978" + }, + "data_manifest": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "data_manifest_sha256": "8c40d792a5cedd4105ff8270ff3141767b147513c52f11292f888d9b1527bfe2", + "experiment": "image_proprio", + "expert_parameters": { + "image": 88178688, + "proprio": 13956388 + }, + "expert_widths": [ + 768, + 256 + ], + "global_batch": 32, + "intended_full_output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "intended_full_run_id": "mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "intended_full_target_steps": 300000, + "job_name": "joint-discrete-image-proprio-lr1e4-full-4gpu", + "lambda_runtime": { + "actual_python": "/workspace/users/paphiwetsa3/tasks/yam-arc400-lambda-20260921/runtime/venv-pinned/bin/python", + "bootstrap_python": "/workspace/users/paphiwetsa3/tasks/arc-global-basis-training-20260924/runtime-r6a/bin/python", + "container_image": "nvcr.io#nvidia/pytorch:25.06-py3", + "extra_pythonpath": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/runtime-overlays/dino4571/site-packages", + "fa3_root": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1/runtime/fa3", + "fa4_root": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/training-runtime/profile/flash-attention", + "minimum_free_scratch_bytes": 10000000000, + "mounts": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-mounts.csv", + "sha256": "b2835189da9504de5263e565bc5a1e827cfa6ce9b788a42e5aed02daca9bc41f" + }, + "node": "gpu03", + "runtime_task": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1", + "schema": "joint-lambda-runtime-v1", + "scratch_root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020", + "temporary_root": "/ephemeral/enroot/41003" + }, + "local_logs": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "lr_policy": { + "peak_lr": 0.0001, + "schedule": "constant", + "warmup_optimizer_steps": 0 + }, + "model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "implementation": "native_image_joint" + }, + "operational_change": "Fresh discrete image+proprio constant AdamW1e-4 experiment from source535d: four dedicated GPUs,8/rank,global32,seed42,300000 optimizer updates,20k checkpoint/validation cadence,retain all15 immutable checkpoints. No action expert, checkpoint initialization, continuation or allocation boundary.", + "operational_service": { + "manifest": "/UNBOUND/operational-service/FILES.sha256", + "manifest_sha256": "UNBOUND_SERVICE_MANIFEST_SHA256", + "revision": "UNBOUND_MAINTAINED_SERVICE_REVISION", + "root": "/UNBOUND/operational-service" + }, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "overrides": [ + "+experiment=pixel_proprio/native_joint/discrete_image_proprio", + "hydra/launcher=basic", + "data.manifest_path=/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "paths.work_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "paths.output_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "hydra.run.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "hydra.sweep.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "hydra.job.chdir=false", + "+trainer.enable_progress_bar=false", + "+trainer.enable_model_summary=false", + "logger=[wandb,csv]", + "logger.wandb.id=mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "logger.wandb.resume=never", + "logger.csv.name=csv", + "logger.csv.flush_logs_every_n_steps=1", + "logger.csv.save_dir=/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1-step000000000", + "evaluator.pixel_diagnostics.log_wandb=true", + "modality_ablation.screening_budget_updates=300000", + "trainer.devices=4", + "data.distributed_world_size=4", + "data.batch_size=8", + "model.optimizer.lr=1e-4", + "model.scheduler=null", + "data.mode=discrete", + "data.sampling_strategy=discrete_window", + "data.min_frames=2", + "data.max_frames=64", + "data.discrete_frames=16", + "data.stride=4", + "joint.max_tokens=3216", + "joint.pose_dims.libero_gripper=36", + "trainer.max_steps=300000", + "callbacks.model_checkpoint.save_last=link" + ], + "parameters": 102135076, + "per_rank_batch": 8, + "phase": "full", + "quota": { + "kind": "filesystem", + "mount": "/workspace", + "node": "gpu03", + "uid": 41003, + "user": "paphiwetsa3" + }, + "reserved_checkpoint_bytes": 40000000000, + "run_id": "mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "runtime_localization": { + "admission": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-admission.json", + "sha256": "f5dbe6abf5aa0bdab9da7b96dbf65dd84c9ed7dc261e2bffc28694d882056670" + }, + "archive_sha256": "92e461b28c8659b31b664e8899ae2e390bd6fdf6d85d575cef77fec5a6956fa8", + "destination_node": "gpu03", + "files_manifest_sha256": "c8bddcf4b2f998adae085f96a263356ad1e6392ce93f7c23a3c6d392e2214289", + "schema": "readonly-runtime-export-v1", + "source_gpu_job": "411963", + "source_node": "gpu21" + }, + "schema": "joint-training-v2", + "smoke": { + "capacity": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/discrete-image-proprio-lr1e4-smoke.capacity.json", + "sha256": "UNBOUND_EXACT_SMOKE_CAPACITY_SHA256" + }, + "completion": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/smoke-completion.json", + "sha256": "UNBOUND_EXACT_SMOKE_COMPLETION_SHA256" + }, + "host_exit": { + "path": "/UNBOUND/exact-smoke-host-exit", + "sha256": "UNBOUND_EXACT_SMOKE_HOST_EXIT_SHA256" + }, + "job_id": "UNBOUND_EXACT_DISCRETE_SMOKE_JOB", + "native": { + "path": "/UNBOUND/exact-smoke-native-completion", + "sha256": "UNBOUND_EXACT_SMOKE_NATIVE_SHA256" + }, + "preflight": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/discrete-image-proprio-lr1e4-smoke.preflight.json", + "sha256": "UNBOUND_EXACT_SMOKE_PREFLIGHT_SHA256" + }, + "request": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/discrete-image-proprio-lr1e4-smoke.json", + "sha256": "UNBOUND_EXACT_SMOKE_REQUEST_SHA256" + }, + "resolved": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/lambda-shared-ops/image-proprio-lr1e4-dc-20261021-v1/requests/discrete-image-proprio-lr1e4-smoke.resolved.yaml", + "sha256": "UNBOUND_EXACT_SMOKE_RESOLVED_SHA256" + } + }, + "source": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "source_archive": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.tar.gz", + "source_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "source_manifest": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.files.sha256", + "source_manifest_sha256": "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184", + "source_parent": "3785f94392b9d900e21d9948fbc5bf94a6d975b8", + "source_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "target_steps": 300000, + "vector_expert_parameter_budget": 20000000, + "world_size": 4 +} diff --git a/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.composed.json b/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.composed.json new file mode 100644 index 000000000..78a969256 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.composed.json @@ -0,0 +1,585 @@ +{ + "callbacks": { + "joint_smoke": { + "_target_": "egomimic.pl_utils.joint_smoke.JointSmokeProof", + "expected_steps": 50, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "warmup_steps": 10, + "world_size": 4 + }, + "learning_rate_monitor": { + "_target_": "lightning.pytorch.callbacks.LearningRateMonitor", + "logging_interval": "step" + }, + "model_checkpoint": { + "_target_": "lightning.pytorch.callbacks.ModelCheckpoint", + "dirpath": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints", + "every_n_epochs": 0, + "every_n_train_steps": 0, + "filename": "epoch_{epoch:04d}-step_{step:09d}", + "save_last": false, + "save_on_train_epoch_end": false, + "save_top_k": -1 + }, + "wandb_profiler": { + "_target_": "egomimic.utils.timing_callback.WandbProfilerLogger", + "log_every_n_steps": 100 + } + }, + "ckpt_path": null, + "data": { + "_target_": "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule", + "balance_training_work": true, + "batch_size": 8, + "checkpointable_sampling": true, + "couple_state_observations": true, + "discrete_frames": 16, + "distributed_world_size": 4, + "forecasting_probability": 0.5, + "image_only": false, + "manifest_path": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "max_frames": 64, + "max_future_seconds": null, + "max_past_seconds": null, + "min_frames": 2, + "mode": "discrete", + "num_workers": 2, + "persistent_workers": true, + "pin_memory": true, + "prefetch_factor": 2, + "sampling_strategy": "discrete_window", + "seed": 42, + "stride": 4, + "task": "visual_diffusion", + "validation_anchors_per_episode": 2, + "validation_batch_size": 1, + "validation_batches_per_embodiment": null, + "validation_discrete_frames": 16, + "validation_sampling_strategy": "discrete_window", + "validation_stride": 4 + }, + "data_context_loader": { + "_target_": "egomimic.rldb.zarr.data_module.load_data_context" + }, + "description": "Historical image core with minimal joint state extension, historical compilation preserved", + "evaluator": { + "_target_": "egomimic.eval.image_proprio_eval.ImageProprioEval", + "objective": "visual_proprio", + "pixel_diagnostics": { + "artifact_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/pixel_diagnostics", + "log_wandb": false, + "max_samples_per_source_kind": 1, + "playback_fps": 20, + "write_video": true + }, + "sample_seed": 42 + }, + "inference_config": { + "enabled": false, + "output_path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints/inference-config.yaml" + }, + "joint": { + "compile_mode": "default", + "max_tokens": 3216, + "pose_dims": { + "libero_gripper": 36 + } + }, + "launch_params": { + "gpus_per_node": 1, + "nodes": 1 + }, + "logger": { + "csv": { + "_target_": "lightning.pytorch.loggers.CSVLogger", + "flush_logs_every_n_steps": 1, + "name": "csv", + "save_dir": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-discrete-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "version": 0 + } + }, + "modality_ablation": { + "active_streams": [ + "image", + "proprio" + ], + "graph": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ], + "initialization": "Native image construction first; vector initialization preserves CPU RNG", + "nodes": { + "image_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + "joint_denoiser": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + "joint_noising": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + "joint_sampler": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + }, + "proprio_loss": { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + } + }, + "screening_budget_updates": 300000, + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "training_noise": "Native image draws unchanged; optional tied proprio draws preserve image RNG" + }, + "mode": "train", + "model": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 0.0001, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_loss": 2, + "joint_denoiser": 1, + "joint_noising": 0, + "joint_sampler": 4, + "proprio_loss": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointNoisingStage", + "noising": { + "_target_": "egomimic.models.joint_image_proprio.CoupledImageProprioNoising", + "active_streams": [ + "image", + "proprio" + ], + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointDenoiserStage", + "model": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioModel", + "active_streams": [ + "image", + "proprio" + ], + "image_model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + }, + "mlp_ratio": 4, + "pose_dims": { + "libero_gripper": 36 + }, + "proprio_width": 256 + } + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "image", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "stream": "proprio", + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_joint_state_action.JointSamplerStage", + "sampler": { + "_target_": "egomimic.models.joint_image_proprio.JointImageProprioSampler", + "state_steps": 20 + } + } + ] + }, + "scheduler": null, + "scheduler_frequency": 1, + "scheduler_interval": "step", + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior", + "distributed_noise_seed": 42 + } + }, + "name": "native_mot_discrete_image_proprio_s42_300k", + "native_joint": { + "reference": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "enable_grad_norm": false, + "optimizer": { + "_partial_": true, + "_target_": "torch.optim.AdamW", + "lr": 5e-05, + "weight_decay": 0.01 + }, + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "stage_ids": { + "image_proprio": 1, + "pixel_denoiser": 1, + "pixel_loss": 2, + "pixel_noising": 0, + "pixel_sampler": 3 + }, + "stages": [ + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage", + "noising": { + "_target_": "egomimic.models.pixel_diffusion.LinearGaussianImageNoising", + "maximum": 1.0 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoDenoiserStage", + "model": { + "_target_": "egomimic.models.image_proprio.ImageProprioModel", + "action_dims": null, + "activation_checkpointing": false, + "allow_patch_compression": false, + "attention_backend": "flash4", + "attention_dtype": "bfloat16", + "attention_query_chunk_size": 0, + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "depth": 8, + "heads": 6, + "image_dim": 768, + "loss_reduction": "mean_per_example", + "max_cameras": 2, + "max_tokens": 3216, + "mlp_ratio": 4, + "noise_conditioning": "adaln_zero", + "objective": "visual_only", + "phase_a_condition_on_proprio": false, + "pose_dims": null, + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [ + 64, + 32, + 32 + ], + "spatial_grid_size": [ + 10, + 10 + ], + "strict_fp32": false, + "time_mode": "discrete", + "token_execution": "packed_valid", + "use_camera_embedding": true, + "use_embodiment_embedding": false, + "use_modality_embedding": false, + "use_phase_embedding": false, + "width": 768 + } + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoLossStage", + "loss": { + "_target_": "egomimic.models.pixel_diffusion.PlainPixelMSE", + "reduction": "mean_per_example" + }, + "weight": 1.0 + }, + { + "_target_": "egomimic.pipeline.stages_pixel_video.PixelVideoSamplerStage", + "sampler": { + "_target_": "egomimic.models.pixel_diffusion.LinearDDIMImageSampler", + "steps": 20 + } + } + ] + }, + "scheduler": null, + "train_log_on_step": true, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior" + } + } + }, + "norm_stats": null, + "norm_stats_only": false, + "normalizer": null, + "paths": { + "dataset_dir": "/coc/flash7/scratch/egoverseS3ZarrDataset", + "output_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "work_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e" + }, + "pixel_video": { + "max_tokens": 3216 + }, + "runtime": { + "cluster_environment": "torch_elastic", + "disable_tf32": true, + "slurm_requeue_owner": "none", + "slurm_save_signal": "SIGUSR2", + "slurm_signal_checkpoint_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1/checkpoints", + "slurm_signal_train_batch_only": false, + "strict_fp32": false + }, + "seed": 42, + "trainer": { + "_target_": "lightning.pytorch.trainer.Trainer", + "accelerator": "gpu", + "accumulate_grad_batches": 1, + "check_val_every_n_epoch": null, + "default_root_dir": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "deterministic": false, + "devices": 4, + "enable_model_summary": false, + "enable_progress_bar": false, + "limit_train_batches": 1.0, + "limit_val_batches": 1.0, + "log_every_n_steps": 1, + "max_epochs": -1, + "max_steps": 50, + "min_epochs": 0, + "num_nodes": 1, + "num_sanity_val_steps": 0, + "precision": "bf16-mixed", + "strategy": "ddp", + "use_distributed_sampler": false, + "val_check_interval": 50 + }, + "val_at_end": true, + "world_model_attention": { + "activation_checkpointing": false, + "backend": "flash4", + "disable_tf32": true, + "dtype": "bfloat16", + "query_chunk_size": 0, + "strict_fp32": false, + "trainer_precision": "bf16-mixed" + }, + "world_model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "loss_reduction": "mean_per_example", + "token_execution": "packed_valid" + } +} diff --git a/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.template.json b/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.template.json new file mode 100644 index 000000000..61b7cdbe6 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/discrete-smoke.template.json @@ -0,0 +1,139 @@ +{ + "_pending_bindings": [ + "Pin maintained service closure and manifest; confirm temporal-pair full admission implementation.", + "Fresh native gpu03 allocation/cache/runtime readback; existing bound producer evidence is a reuse candidate, not fresh node certification.", + "Verify both new W&B identities absent, exact queues/claims/outputs empty, destination retention space.", + "Resolve both exact smoke/full configs in scheduled CPU allocation and record actual meta-model counts." + ], + "_template_status": "NON_SUBMITTABLE_PENDING_EXACT_BINDINGS", + "active_streams": [ + "image", + "proprio" + ], + "allocation": null, + "arm": "discrete", + "data_cache": { + "job_id": 411978, + "native_receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-native.txt", + "native_receipt_sha256": "ad08aa8efbd7ad8e74587aea8e839650d10e7757890c715bfd03f55dc79df3a0", + "node": "gpu03", + "parent": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data", + "receipt": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/data-recovery-receipt.json", + "receipt_sha256": "008d9947db54e14d7fee2d194e5e89786892d305b1edecd4e9ea209b5d9ba4a5", + "root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/complete-data/411978" + }, + "data_manifest": "/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "data_manifest_sha256": "8c40d792a5cedd4105ff8270ff3141767b147513c52f11292f888d9b1527bfe2", + "experiment": "image_proprio", + "expert_parameters": { + "image": 88178688, + "proprio": 13956388 + }, + "expert_widths": [ + 768, + 256 + ], + "global_batch": 32, + "intended_full_output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/runs/shared_discrete_image_proprio_lr1e4constant_4gpu", + "intended_full_run_id": "mot-shared-libero-image-proprio-discrete-lr1e-4-constant-s42-300k-20261021-v1", + "intended_full_target_steps": 300000, + "job_name": "joint-discrete-image-proprio-lr1e4-smoke-4gpu", + "lambda_runtime": { + "actual_python": "/workspace/users/paphiwetsa3/tasks/yam-arc400-lambda-20260921/runtime/venv-pinned/bin/python", + "bootstrap_python": "/workspace/users/paphiwetsa3/tasks/arc-global-basis-training-20260924/runtime-r6a/bin/python", + "container_image": "nvcr.io#nvidia/pytorch:25.06-py3", + "extra_pythonpath": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/runtime-overlays/dino4571/site-packages", + "fa3_root": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1/runtime/fa3", + "fa4_root": "/workspace/users/paphiwetsa3/tasks/continuous-time-mot-world-model-20261016/preparation/training-runtime/profile/flash-attention", + "minimum_free_scratch_bytes": 10000000000, + "mounts": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-mounts.csv", + "sha256": "b2835189da9504de5263e565bc5a1e827cfa6ce9b788a42e5aed02daca9bc41f" + }, + "node": "gpu03", + "runtime_task": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1", + "schema": "joint-lambda-runtime-v1", + "scratch_root": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020", + "temporary_root": "/ephemeral/enroot/41003" + }, + "local_logs": "/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-discrete-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "lr_policy": { + "peak_lr": 0.0001, + "schedule": "constant", + "warmup_optimizer_steps": 0 + }, + "model_execution": { + "compile_backend": "inductor", + "compile_dynamic": true, + "compile_mode": "default", + "implementation": "native_image_joint" + }, + "operational_change": "Fresh discrete image+proprio constant AdamW1e-4 experiment from source535d: four dedicated GPUs,8/rank,global32,seed42,300000 optimizer updates,20k checkpoint/validation cadence,retain all15 immutable checkpoints. No action expert, checkpoint initialization, continuation or allocation boundary.", + "operational_service": { + "manifest": "/UNBOUND/operational-service/FILES.sha256", + "manifest_sha256": "UNBOUND_SERVICE_MANIFEST_SHA256", + "revision": "UNBOUND_MAINTAINED_SERVICE_REVISION", + "root": "/UNBOUND/operational-service" + }, + "output": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "overrides": [ + "+experiment=pixel_proprio/native_joint/discrete_image_proprio", + "+joint_execution=smoke", + "hydra/launcher=basic", + "data.manifest_path=/workspace/users/paphiwetsa3/tasks/pixel-proprio-gripper-20261017-v1/data/manifest.json", + "paths.work_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "paths.output_dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.run.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.sweep.dir=/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/smokes/shared-discrete-image-proprio-lr1e4-lambda-4gpu-gpu03-v1", + "hydra.job.chdir=false", + "+trainer.enable_progress_bar=false", + "+trainer.enable_model_summary=false", + "logger=[csv]", + "~logger.wandb", + "logger.csv.name=csv", + "logger.csv.flush_logs_every_n_steps=1", + "logger.csv.save_dir=/ephemeral/enroot/paphiwetsa3/joint-modality-ablations-20261020/logs/shared-discrete-image-proprio-lr1e4-smoke-4gpu-gpu03-v1", + "evaluator.pixel_diagnostics.log_wandb=false", + "modality_ablation.screening_budget_updates=300000", + "trainer.devices=4", + "data.distributed_world_size=4", + "data.batch_size=8", + "model.optimizer.lr=1e-4", + "model.scheduler=null", + "data.mode=discrete", + "data.sampling_strategy=discrete_window", + "data.min_frames=2", + "data.max_frames=64", + "data.discrete_frames=16", + "data.stride=4", + "joint.max_tokens=3216", + "joint.pose_dims.libero_gripper=36" + ], + "parameters": 102135076, + "per_rank_batch": 8, + "phase": "smoke", + "reserved_checkpoint_bytes": 1600000000, + "runtime_localization": { + "admission": { + "path": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/provenance/replacement/high-lr-gpu03-stage/local-runtime-admission.json", + "sha256": "f5dbe6abf5aa0bdab9da7b96dbf65dd84c9ed7dc261e2bffc28694d882056670" + }, + "archive_sha256": "92e461b28c8659b31b664e8899ae2e390bd6fdf6d85d575cef77fec5a6956fa8", + "destination_node": "gpu03", + "files_manifest_sha256": "c8bddcf4b2f998adae085f96a263356ad1e6392ce93f7c23a3c6d392e2214289", + "schema": "readonly-runtime-export-v1", + "source_gpu_job": "411963", + "source_node": "gpu21" + }, + "schema": "joint-training-v2", + "source": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "source_archive": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.tar.gz", + "source_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "source_manifest": "/workspace/users/paphiwetsa3/tasks/joint-modality-ablations-20261020/inputs/replacement/source-535d4fc26c4c85cf57bbd2a776cc79cfec085c4e.files.sha256", + "source_manifest_sha256": "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184", + "source_parent": "3785f94392b9d900e21d9948fbc5bf94a6d975b8", + "source_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "target_steps": 50, + "vector_expert_parameter_budget": 20000000, + "world_size": 4 +} diff --git a/tests/fixtures/joint_temporal_pair_535d/provenance.json b/tests/fixtures/joint_temporal_pair_535d/provenance.json new file mode 100644 index 000000000..1f17a5821 --- /dev/null +++ b/tests/fixtures/joint_temporal_pair_535d/provenance.json @@ -0,0 +1,20 @@ +{ + "files": { + "continuous-full.composed.json": "5c572989c978e1d2a916a2dcc2c3f9e774395bd58c81f80f80c66640e70cb16f", + "continuous-full.template.json": "cbba858f3d832b208eb100b29274a138a88e2214f30fbe29fd213d9971828654", + "continuous-smoke.composed.json": "fbfdee559c1343026dd529735e095ae825d58e886679c692f19e13c9bc8a3295", + "continuous-smoke.template.json": "154190a6ba03b58a164c34b0a3999144f29744c385eadceeebe7669845657951", + "discrete-full.composed.json": "9d23ca57ce397b2cfde455f2ef6397cd396ca1140682021a14e5b98499e23a35", + "discrete-full.template.json": "5422386326fc53658ed672406def430e271f4482ed0f2c877ccf987b1bd68578", + "discrete-smoke.composed.json": "9d35f3a742d907070f642bfb805e1f9b3db708810359eb2582a76ebc3a9e37d3", + "discrete-smoke.template.json": "c97c1b8507c5bd37f78f9624b12bd473d57904e4294db84167f49bb31800c792" + }, + "gpu_certified": false, + "hydra": "1.3.2", + "kind": "Frozen535d composed config/request fixtures; synthetic smoke evidence generated only at test time", + "omegaconf": "2.3.1", + "python": "3.12.14", + "scientific_archive_sha256": "f23b041bb2e8f250e406cbc24ef7be06b463221e8a004125d175f006661d8325", + "scientific_revision": "535d4fc26c4c85cf57bbd2a776cc79cfec085c4e", + "source_manifest_sha256": "f29118056a2f7824b5215e503bed99b77264b1598bfe154ba20c5412be106184" +} diff --git a/tests/test_joint_crash_admission.py b/tests/test_joint_crash_admission.py new file mode 100644 index 000000000..fcdc67e50 --- /dev/null +++ b/tests/test_joint_crash_admission.py @@ -0,0 +1,192 @@ +"""Synthetic evidence rejection tests; no checkpoint/runtime/GPU certification.""" +import copy +import hashlib +import json +import os +from pathlib import Path +from types import SimpleNamespace +import tempfile +import unittest +from unittest.mock import patch + +from scripts.train.joint import crash_recovery as c +from scripts.train.joint.continuation import allocation_steps + + +class AdmissionTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.root = Path(self.temp.name) + self.counter = 0 + original = Path(__file__).parent / 'fixtures/joint_crash_recovery_535d/image-only-parent.request.json' + self.parent = json.loads(original.read_text()) + self.parent['output'] = str(self.root / 'parent') + self.parent['local_logs'] = str(self.root / 'parent-logs') + self.parent['overrides'] = ['model.optimizer.lr=1e-4', 'paths.output_dir=' + self.parent['output'], + 'logger.csv.save_dir=' + self.parent['local_logs']] + self.q = copy.deepcopy(self.parent) + self.q['output'] = str(self.root / 'attempt') + self.q['local_logs'] = str(self.root / 'attempt-logs') + self.q['overrides'] = [v.replace(self.parent['output'],self.q['output']).replace(self.parent['local_logs'],self.q['local_logs']) for v in self.parent['overrides']] + self.q['job_name'] = 'joint-image-lr1e4-recovery-test' + self.q['operational_service'] = {'manifest_sha256':'a'*64} + parent = self.binding(self.parent) + native = self.binding('JobId=412056 JobName=joint-image-lr1e4-full-4gpu UserId=paphiwetsa3(41003) JobState=CANCELLED Requeue=0 Restarts=0', text=True) + preflight = self.binding({'status':'FULL_ADMISSION_PASS','job_id':'412056','request_sha256':parent['sha256'], + 'provenance':{'run_id':self.q['run_id'],'source_revision':c.SOURCE}}) + self.failure = {'schema':'joint-crash-failure-v1','parent_request_sha256':parent['sha256'], + 'run_id':self.q['run_id'],'classification':'worker_crash', 'recovery_authorized':True, + 'human_retired':False,'job_id':412056,'native':native, 'logs':[self.binding('NCCL timeout/SIGABRT',text=True)], + 'parent_preflight':preflight} + ckpt = self.root/'parent/checkpoints/epoch_epoch=0011-step_step=000020000.ckpt' + ckpt.parent.mkdir(parents=True);ckpt.write_bytes(b'synthetic-not-a-real-checkpoint') + (ckpt.parent/'last.ckpt').symlink_to(ckpt.name) + self.inspection = {'schema':'joint-crash-checkpoint-inspection-v1','status':'PASS', + 'parent_request_sha256':parent['sha256'],'run_id':self.q['run_id'], + 'source_revision':c.SOURCE,'world_size':4,'global_batch':32,'global_step':20000, + 'strict_full_wrapper_optimizer_reload':True,'sampler_all_rank_rng_verified':True, + 'highest_complete_checkpoint_inspected':True,'all_newer_candidates_reconciled':True, + 'checkpoint':{'path':str(ckpt),'sha256':hashlib.sha256(ckpt.read_bytes()).hexdigest()}} + inspect = self.binding(self.inspection) + self.proof = {'schema':'joint-crash-recovery-smoke-v1','status':'PASS','source_revision':c.SOURCE, + 'parent_request_sha256':parent['sha256'],'inspection_sha256':inspect['sha256'], + 'checkpoint_sha256':self.inspection['checkpoint']['sha256'], + 'service_manifest_sha256':'a'*64,'world_size':4,'global_batch':32, + 'all_rank_full_state_restored':True,'finite_optimizer_updates':2, + 'scheduled_validation_anchors':100,'no_replayed_boundary_validation_or_save':True, + 'production_science_unchanged':True,'sample_finite':True,'hidden_query_labels_nan':True, + 'sample_steps':20,'job_id':'999','job_name':'joint-image-lr1e4-crash-diagnostic', + 'native':self.binding('JobId=999 JobName=joint-image-lr1e4-crash-diagnostic UserId=paphiwetsa3(41003) NodeList=gpu19 Requeue=0 Restarts=0 JobState=COMPLETED ExitCode=0:0',text=True), + 'host_exit':self.binding('0\n',text=True)} + self.q['crash_recovery']={'schema':'joint-crash-recovery-v1','parent_request':parent, + 'failure':self.binding(self.failure),'inspection':inspect, + 'recovery_smoke':None} + self.bind_proof() + + def tearDown(self): + self.temp.cleanup() + + def binding(self, value, text=False): + self.counter += 1;p=self.root/f'evidence-{self.counter}.txt' + p.write_text(value if text else json.dumps(value)) + return {'path':str(p),'sha256':hashlib.sha256(p.read_bytes()).hexdigest()} + + def update_failure(self, **values): + self.failure.update(values);self.q['crash_recovery']['failure']=self.binding(self.failure) + + def update_inspection(self, **values): + self.inspection.update(values);b=self.binding(self.inspection) + self.q['crash_recovery']['inspection']=b;self.proof['inspection_sha256']=b['sha256'] + self.bind_proof() + + def bind_proof(self): + diagnostic = copy.deepcopy(self.q) + diagnostic.pop('crash_recovery',None) + diagnostic.update(schema='joint-crash-recovery-diagnostic-v1', + parent_request=self.q['crash_recovery']['parent_request'], + inspection=self.q['crash_recovery']['inspection'], + job_name='joint-image-lr1e4-crash-diagnostic',output=str(self.root/'diagnostic')) + request=self.binding(diagnostic);self.proof['request']=request + self.proof['request_sha256']=request['sha256'] + self.proof['candidate']=self.binding({'status':'SYNTHETIC_TEST_ONLY'}) + records=[] + for group in ['restore','validation','anchors']: + for rank in range(4): + name = (f'restore-proof/restore-step000020000-rank{rank}.json' if group=='restore' else + f'restore-proof/restore-validation-rank{rank}.json' if group=='validation' else + f'validation-proof/rank{rank}.json') + path=self.root/'diagnostic'/name;path.parent.mkdir(parents=True,exist_ok=True) + path.write_text(json.dumps({'status':'PASS','job_id':'999','world_size':4, + 'rank':rank,'start_step':20000,'global_step':20002})) + records.append({'path':str(path),'sha256':hashlib.sha256(path.read_bytes()).hexdigest()}) + self.proof['rank_proofs']=records + final=self.root/'diagnostic/checkpoints/diagnostic.ckpt';final.parent.mkdir(exist_ok=True) + final.write_bytes(b'SYNTHETIC_TEST_ONLY') + self.proof['final_diagnostic_checkpoint']={'path':str(final), + 'sha256':hashlib.sha256(final.read_bytes()).hexdigest(),'strict_full_wrapper_optimizer_reload':True} + self.q['crash_recovery']['recovery_smoke']=self.binding(self.proof) + + def test_bound_failure_and_actual_proof_contract(self): + self.assertEqual(c.admission(self.q)[2],20000) + self.assertEqual(allocation_steps(self.q),(20000,300000)) + + def test_no_recovery_smoke_no_admission(self): + del self.q['crash_recovery']['recovery_smoke'] + with self.assertRaises((KeyError,ValueError)):c.admission(self.q) + + def test_wrong_native_exit_rejected(self): + for state in ('RUNNING','PENDING','COMPLETING','COMPLETED'): + with self.subTest(state=state): + self.update_failure(native=self.binding('JobId=412056 JobName=joint-image-lr1e4-full-4gpu UserId=paphiwetsa3(41003) JobState='+state+' Requeue=0 Restarts=0',text=True)) + with self.assertRaises(ValueError):c.admission(self.q) + + def test_deliberate_human_retirement_rejected(self): + self.update_failure(human_retired=True) + with self.assertRaises(ValueError):c.admission(self.q) + + def test_changed_science_or_lr_rejected(self): + for field,value in [('world_size',2),('source_revision','0'*40),('lr_policy',dict(schedule='constant',peak_lr=5e-5,warmup_optimizer_steps=0)),('target_steps',320000)]: + q=copy.deepcopy(self.q);q[field]=value + with self.subTest(field=field),self.assertRaises(ValueError):c.admission(q) + + def test_unresolved_newer_checkpoint_rejected(self): + self.update_inspection(all_newer_candidates_reconciled=False) + with self.assertRaises(ValueError):c.admission(self.q) + + def test_corrupted_checkpoint_rejected(self): + Path(self.inspection['checkpoint']['path']).write_bytes(b'changed') + with self.assertRaises(ValueError):c.admission(self.q) + + def test_overlapping_or_reused_output_rejected(self): + self.q['output']=self.parent['output']+'/recovery' + with self.assertRaises(ValueError):c.admission(self.q) + + def test_override_only_log_output_relocation(self): + self.q['overrides'].append('data.max_frames=2') + with self.assertRaises(ValueError):c.admission(self.q) + + def test_smoke_host_native_and_service_closure_required(self): + for field,value in [('host_exit',self.binding('1\n',text=True)),('service_manifest_sha256','b'*64),('scheduled_validation_anchors',0),('no_replayed_boundary_validation_or_save',False)]: + proof=copy.deepcopy(self.proof);proof[field]=value;q=copy.deepcopy(self.q) + q['crash_recovery']['recovery_smoke']=self.binding(proof) + with self.subTest(field=field),self.assertRaises(ValueError):c.admission(q) + + def test_accepted_parent_request_preflight_required(self): + self.update_failure(parent_preflight=self.binding({'status':'FULL_ADMISSION_PASS','job_id':'123','request_sha256':'bad','provenance':{}})) + with self.assertRaises(ValueError):c.admission(self.q) + + def test_scheduled_path_is_unchanged(self): + self.assertEqual(allocation_steps(self.parent),(0,300000)) + q=copy.deepcopy(self.q);q['continuation']={} + with self.assertRaises(ValueError):c.admission(q) + + def test_original_comparison_foundation_is_retained(self): + self.assertEqual(c.foundation_request(self.q),self.parent) + other=copy.deepcopy(self.parent);other['run_id']='foreign' + self.q['crash_recovery']['parent_request']=self.binding(other) + with self.assertRaises(ValueError):c.foundation_request(self.q) + + def test_writer_quiescence_rejects_old_active_parent(self): + run=SimpleNamespace(returncode=0,stdout='JobState=COMPLETING',stderr='') + with patch.dict(os.environ,SLURM_JOB_ID='123'),patch.object(c.subprocess,'run',return_value=run): + with self.assertRaises(ValueError):c.assert_native_writer_quiescence(self.q) + + def test_writer_quiescence_rejects_unknown_transport(self): + run=SimpleNamespace(returncode=1,stdout='',stderr='connection unavailable') + with patch.dict(os.environ,SLURM_JOB_ID='123'),patch.object(c.subprocess,'run',return_value=run): + with self.assertRaises(ValueError):c.assert_native_writer_quiescence(self.q) + + def test_writer_quiescence_checks_entire_owned_queue(self): + runs=[SimpleNamespace(returncode=1,stdout='',stderr='Invalid job id specified'), + SimpleNamespace(returncode=0,stdout='123|joint-image-lr1e4-recovery-test|RUNNING\n789|joint-image-lr1e4-other|PENDING',stderr='')] + with patch.dict(os.environ,SLURM_JOB_ID='123'),patch.object(c.subprocess,'run',side_effect=runs): + with self.assertRaises(ValueError):c.assert_native_writer_quiescence(self.q) + + def test_exclusive_actual_destination_fsync_probe(self): + with patch.dict(os.environ,SLURM_JOB_ID='123'),patch.object(c.shutil,'disk_usage',return_value=SimpleNamespace(free=100_000_000_000)): + result=c.prove_destination_storage(self.q) + self.assertTrue(result['write_fsync_readback_pass']) + with self.assertRaises(FileExistsError):c.prove_destination_storage(self.q) + + +if __name__=='__main__':unittest.main() diff --git a/tests/test_joint_crash_history.py b/tests/test_joint_crash_history.py new file mode 100644 index 000000000..513b0611b --- /dev/null +++ b/tests/test_joint_crash_history.py @@ -0,0 +1,211 @@ +"""Real temporary-file history counterfactuals; no checkpoint content is loaded.""" + +import copy +import csv +import hashlib +import json +from pathlib import Path +import sys +import tempfile +import unittest + + +from scripts.train.joint.crash_history import collect_milestones, collect_validation, history_segments + + +METRICS = ["val/libero_gripper/forecast/image_normalized_mse", + "val/libero_gripper/completion/image_normalized_mse"] + + +class HistoryTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.root = Path(self.temp.name).resolve() + self.parent = self.request("parent") + self.write_checkpoint(self.parent, 20000) + self.write_csv(self.parent, 0, [dict(step=19999, **dict.fromkeys(METRICS, .01))]) + parent_binding = self.binding("parent.json", self.parent) + self.inspection = dict(schema="joint-crash-checkpoint-inspection-v1", status="PASS", + parent_request_sha256=parent_binding["sha256"], run_id="same-run", + source_revision="source", global_step=20000) + self.child = self.request("child") + self.child["crash_recovery"] = dict( + parent_request=parent_binding, inspection=self.binding("inspection.json", self.inspection)) + self.write_checkpoint(self.child, 40000) + self.write_csv(self.child, 20000, [dict(step=39999, **dict.fromkeys(METRICS, .02))]) + + def tearDown(self): + self.temp.cleanup() + + def request(self, name): + q = dict(run_id="same-run", source_revision="source", world_size=4, + active_streams=["image"], target_steps=40000, + output=str(self.root / name), local_logs=str(self.root / (name + "-logs"))) + (Path(q["output"]) / "checkpoints").mkdir(parents=True) + return q + + def binding(self, name, data): + path = self.root / name + path.write_text(json.dumps(data)) + return dict(path=str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + def write_checkpoint(self, q, step): + path = Path(q["output"]) / "checkpoints" / f"epoch_epoch=0001-step_step={step:09d}.ckpt" + path.write_text("synthetic file, never loaded as a checkpoint") + return path + + def write_csv(self, q, start, rows): + path = Path(q["local_logs"]) / "csv" / f"step{start:09d}" / "metrics.csv" + path.parent.mkdir(parents=True, exist_ok=True) + keys = list(dict.fromkeys(["step", *METRICS, *(key for row in rows for key in row)])) + with path.open("w") as stream: + writer = csv.DictWriter(stream, fieldnames=keys) + writer.writeheader(); writer.writerows(rows) + return path + + def test_merges_two_immutable_checkpoint_and_validation_segments(self): + before = {path: path.read_bytes() for path in self.root.rglob("*") if path.is_file()} + self.assertEqual(set(collect_milestones(self.child)), {20000, 40000}) + self.assertEqual(set(collect_validation(self.child)), {20000, 40000}) + self.assertEqual([(s["start_step"], s["through_step"]) + for s in history_segments(self.child)], [(0, 20000), (20000, 40000)]) + self.assertTrue(all(path.read_bytes() == data for path, data in before.items())) + + def test_ignores_preserved_parent_events_beyond_restored_checkpoint(self): + self.write_csv(self.parent, 0, [dict(step=19999, **dict.fromkeys(METRICS, .01)), + dict(step=39999, **dict.fromkeys(METRICS, .07))]) + self.assertEqual(collect_validation(self.child)[40000][METRICS[0]], .02) + + def test_rejects_unreconciled_newer_parent_checkpoint(self): + self.write_checkpoint(self.parent, 40000) + with self.assertRaisesRegex(ValueError, "Unreconciled"): + collect_milestones(self.child) + + def test_rejects_repeated_boundary_save_in_new_tree(self): + self.write_checkpoint(self.child, 20000) + with self.assertRaisesRegex(ValueError, "repeated resume-boundary save"): + collect_milestones(self.child) + + def test_rejects_checkpoint_symlink(self): + path = self.write_checkpoint(self.child, 40000) + target = self.root / "external.ckpt"; target.write_text("retained") + path.unlink(); path.symlink_to(target) + with self.assertRaisesRegex(ValueError, "mutable retained"): + collect_milestones(self.child) + + def test_last_alias_does_not_replace_immutable_milestone(self): + path = self.write_checkpoint(self.child, 40000) + path.rename(path.parent / "other-file.bin") + (path.parent / "last.ckpt").symlink_to(path.parent / "other-file.bin") + with self.assertRaisesRegex(ValueError, "Missing retained checkpoint"): + collect_milestones(self.child) + + def test_rejects_repeated_boundary_validation(self): + self.write_csv(self.child, 20000, [dict(step=19999, **dict.fromkeys(METRICS, .01)), + dict(step=39999, **dict.fromkeys(METRICS, .02))]) + with self.assertRaisesRegex(ValueError, "Repeated validation at"): + collect_validation(self.child) + + def test_rejects_nonfinite_metric(self): + self.write_csv(self.child, 20000, [dict(step=39999, **{METRICS[0]: "nan", METRICS[1]: .01})]) + with self.assertRaisesRegex(ValueError, "Nonfinite"): + collect_validation(self.child) + + def test_rejects_missing_stream_metric(self): + self.write_csv(self.child, 20000, [dict(step=39999, **{METRICS[0]: .01})]) + with self.assertRaisesRegex(ValueError, "Incomplete scheduled"): + collect_validation(self.child) + + def test_rejects_unscheduled_validation(self): + self.write_csv(self.child, 20000, [dict(step=29999, **dict.fromkeys(METRICS, .02))]) + with self.assertRaisesRegex(ValueError, "20k optimizer cadence"): + collect_validation(self.child) + + def test_rejects_duplicate_validation_even_if_metric_equal(self): + row = dict(step=39999, **dict.fromkeys(METRICS, .02)) + self.write_csv(self.child, 20000, [row, row]) + with self.assertRaisesRegex(ValueError, "Repeated validation metric"): + collect_validation(self.child) + + def test_merges_disjoint_metric_rows_at_one_event(self): + self.write_csv(self.child, 20000, [dict(step=39999, **{METRICS[0]: .01}), + dict(step=39999, **{METRICS[1]: .02})]) + self.assertEqual(set(collect_validation(self.child)[40000]), set(METRICS)) + + def test_mixed_csv_step_conventions_fail_in_one_event(self): + self.write_csv(self.child, 20000, [dict(step=39999, **{METRICS[0]: .01}), + dict(step=40000, **{METRICS[1]: .02})]) + with self.assertRaisesRegex(ValueError, "Ambiguous"): + collect_validation(self.child) + + def test_zero_progress_attempt_kept_in_chain(self): + child, _ = self.with_zero_progress_parent() + self.assertEqual(len(history_segments(child)), 3) + self.assertEqual(set(collect_milestones(child)), {20000, 40000}) + + def with_zero_progress_parent(self): + failed = self.request("zero-progress") + failed["crash_recovery"] = copy.deepcopy(self.child["crash_recovery"]) + self.write_csv(failed, 20000, []) + child = copy.deepcopy(self.child) + binding = self.binding("zero.json", failed) + inspection = {**self.inspection, "parent_request_sha256": binding["sha256"]} + child["crash_recovery"] = dict(parent_request=binding, + inspection=self.binding("zero-inspect.json", inspection)) + return child, failed + + def test_zero_owned_attempt_may_have_no_checkpoint_directory(self): + child, failed = self.with_zero_progress_parent() + (Path(failed["output"]) / "checkpoints").rmdir() + self.assertEqual(set(collect_milestones(child)), {20000, 40000}) + + def test_zero_owned_attempt_may_have_no_csv_or_checkpoint_directory(self): + child, failed = self.with_zero_progress_parent() + (Path(failed["output"]) / "checkpoints").rmdir() + (Path(failed["local_logs"]) / "csv" / "step000020000" / "metrics.csv").unlink() + self.assertEqual(set(collect_milestones(child)), {20000, 40000}) + self.assertEqual(set(collect_validation(child)), {20000, 40000}) + + def test_zero_owned_existing_checkpoint_is_not_silently_ignored(self): + child, failed = self.with_zero_progress_parent() + self.write_checkpoint(failed, 20000) + with self.assertRaisesRegex(ValueError, "repeated resume-boundary save"): + collect_milestones(child) + + def test_zero_owned_existing_validation_is_not_silently_ignored(self): + child, failed = self.with_zero_progress_parent() + self.write_csv(failed, 20000, [dict(step=19999, **dict.fromkeys(METRICS, .01))]) + with self.assertRaisesRegex(ValueError, "Repeated validation at"): + collect_validation(child) + + def test_rejects_overlapping_attempt_output_and_wrong_identity(self): + repeated = copy.deepcopy(self.child); repeated["output"] = self.parent["output"] + with self.assertRaisesRegex(ValueError, "Repeated history"): + history_segments(repeated) + repeated = copy.deepcopy(self.child); repeated["run_id"] = "another-run" + with self.assertRaisesRegex(ValueError, "Unbound crash"): + history_segments(repeated) + + def test_rejects_nested_or_aliased_attempt_output(self): + repeated = copy.deepcopy(self.child) + repeated["output"] = str(Path(self.parent["output"]) / "nested-attempt") + with self.assertRaisesRegex(ValueError, "nested attempt"): + history_segments(repeated) + alias = self.root / "alias" + alias.symlink_to(Path(self.parent["output"]), target_is_directory=True) + repeated["output"] = str(alias) + with self.assertRaisesRegex(ValueError, "Repeated history"): + history_segments(repeated) + + def test_rejects_ambiguous_continuation_and_corrupt_binding(self): + ambiguous = copy.deepcopy(self.child); ambiguous["continuation"] = {} + with self.assertRaisesRegex(ValueError, "both crash recovery"): + history_segments(ambiguous) + binding = self.child["crash_recovery"]["inspection"] + Path(binding["path"]).write_text("{}") + with self.assertRaises(ValueError): + history_segments(self.child) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_joint_crash_inspect.py b/tests/test_joint_crash_inspect.py new file mode 100644 index 000000000..f274bae90 --- /dev/null +++ b/tests/test_joint_crash_inspect.py @@ -0,0 +1,435 @@ +"""CPU tensor/checkpoint rejection tests; no production checkpoint is read.""" + +import copy +import contextlib +import hashlib +import json +import os +from pathlib import Path +import random +import sys +import tempfile +import types +import unittest +from unittest.mock import patch + + +import numpy as np +import torch +from omegaconf import OmegaConf + +from scripts.train.joint import crash_inspect as target + + +def fingerprint(value): + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +class CheckpointInspectionTest(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.model = torch.nn.Linear(3, 2) + optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-4, weight_decay=.01) + self.model(torch.ones(1, 3)).sum().backward() + optimizer.step() + for state in optimizer.state.values(): + state["step"] = torch.tensor(20000.) + self.parent = dict(source_revision=target.FROZEN_SOURCE, run_id=target.RUN_ID, + source_manifest_sha256="a"*64, data_manifest_sha256="b"*64, + world_size=4, global_batch=32, per_rank_batch=8, + lr_policy=dict(schedule="constant", peak_lr=1e-4, warmup_optimizer_steps=0), + target_steps=300000, output=str(self.root/"run")) + self.resolved = dict(model=dict(optimizer=dict(_target_="torch.optim.AdamW", lr=1e-4, + weight_decay=.01), scheduler=None)) + context = {"schema": "synthetic-context", "normalizer": {"mean": [0., 0., 0.]}} + self.context_hash = fingerprint(context) + provenance = {key: self.parent[key] for key in ("source_revision", "source_manifest_sha256", + "data_manifest_sha256", "run_id")} + provenance["data_context_sha256"] = self.context_hash + self.checkpoint = dict( + global_step=20000, epoch=11, state_dict=self.model.state_dict(), + hyper_parameters=dict(config_tree=dict(model=copy.deepcopy(self.resolved["model"]), + run_provenance=provenance)), + optimizer_states=[optimizer.state_dict()], lr_schedulers=[], data_context=context, + inference_binding=dict(data_context_sha256=self.context_hash), + PixelProprioDataModule=dict(schema="mot-sampling-state-v1", consumed_batches=20000, + distributed_sampling=dict(schema="pixel-distributed-sampling-v1", world_size=4, + per_rank_batch_size=8, global_batch_size=32, + training_sharding="global_draw_rank_slots", + validation_sharding="contiguous_global_batches_no_padding")), + mot_rng_state=dict(schema="mot_rng_state.distributed.v1", world_size=4, + ranks=[dict(torch_cpu=torch.get_rng_state(), torch_cuda=[torch.get_rng_state()], + numpy=np.random.RandomState(rank).get_state(), python=random.Random(rank).getstate(), + cuda_scope="current_device") for rank in range(4)]), + loops=dict(fit_loop={ + "epoch_loop.automatic_optimization.optim_progress": { + "optimizer": {"step": {"total": {"completed": 20000, "ready": 20000}}}}, + "epoch_loop.batch_progress": {"total": { + "completed": 19999, "processed": 20000, "ready": 20000, "started": 20000}}, + "epoch_loop.state_dict": {"_batches_that_stepped": 19999}, + "epoch_loop.val_loop.batch_progress": {"total": {"completed": 0}}, + })) + + def inspect(self, checkpoint=None): + return target.validate_payload(self.checkpoint if checkpoint is None else checkpoint, + self.parent, self.resolved, expected_context_sha256=self.context_hash, + context_fingerprint=fingerprint) + + def reject(self, checkpoint=None, match=None): + with self.assertRaisesRegex(ValueError, match or ".+"): + self.inspect(checkpoint) + + def first_optimizer_state(self): + return next(iter(self.checkpoint["optimizer_states"][0]["state"].values())) + + def test_actual_torch_save_mmap_reload_has_expected_state_digests(self): + path = self.root/"synthetic.ckpt" + torch.save(self.checkpoint, path) + loaded = torch.load(path, map_location="cpu", mmap=True, weights_only=False) + result = self.inspect(loaded) + self.assertEqual(result["global_step"], 20000) + self.assertEqual(result["optimizer_steps"], [20000]) + self.assertEqual(result["optimizer_state_sha256"], target.state_digest(self.checkpoint["optimizer_states"][0])) + self.assertEqual(len(result["rank_rng_sha256"]), 4) + self.assertFalse(result["loop_state"]["loop_training_resume_supported"]) + + def test_actual_saved_dictconfig_shape_is_supported(self): + self.checkpoint["hyper_parameters"]["config_tree"] = OmegaConf.create( + self.checkpoint["hyper_parameters"]["config_tree"]) + self.assertEqual(self.inspect()["global_step"], 20000) + + def test_no_global_rng_draw_or_mutation_during_validation(self): + before = target.state_digest(dict(torch=torch.get_rng_state(), numpy=np.random.get_state(), + python=random.getstate())) + self.inspect() + after = target.state_digest(dict(torch=torch.get_rng_state(), numpy=np.random.get_state(), + python=random.getstate())) + self.assertEqual(before, after) + + def test_state_digest_preserves_tensor_dtype_shape_and_container_type(self): + self.assertNotEqual(target.state_digest(torch.ones(2)), target.state_digest(torch.ones(2, dtype=torch.float64))) + self.assertNotEqual(target.state_digest(torch.ones(2)), target.state_digest(torch.ones(1, 2))) + self.assertNotEqual(target.state_digest([1, 2]), target.state_digest((1, 2))) + + def test_nonfinite_or_nonfp32_wrapper_weights_rejected(self): + for value in [torch.tensor([float("nan")]), torch.tensor([float("inf")]), torch.ones(3, dtype=torch.float16)]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["state_dict"]["weight"] = value + self.reject(checkpoint, "stored wrapper weights") + + def test_missing_full_optimizer_rejected(self): + self.checkpoint["optimizer_states"] = [] + self.reject(match="full AdamW") + + def test_optimizer_parameter_identity_must_be_complete_and_unique(self): + for params in [[0], [0, 0], [0, 1, 9]]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["optimizer_states"][0]["param_groups"][0]["params"] = params + self.reject(checkpoint, "optimizer parameter state") + + def test_nonfinite_nonfp32_negative_or_mismatched_moments_rejected(self): + for key, value in [("exp_avg", torch.tensor([float("nan")])), + ("exp_avg", torch.ones(2, 3, dtype=torch.float16)), + ("exp_avg_sq", -torch.ones(2, 3)), + ("exp_avg_sq", torch.ones(1))]: + checkpoint = copy.deepcopy(self.checkpoint) + next(iter(checkpoint["optimizer_states"][0]["state"].values()))[key] = value + self.reject(checkpoint, "AdamW moments") + + def test_optimizer_step_must_equal_global_step(self): + self.first_optimizer_state()["step"] = torch.tensor(0.) + self.reject(match="step disagrees") + + def test_lr_weight_decay_and_scheduler_are_exact(self): + for key, value in [("lr", 5e-5), ("weight_decay", 0.)]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["optimizer_states"][0]["param_groups"][0][key] = value + self.reject(checkpoint, "groups/LR") + self.checkpoint["lr_schedulers"] = [{"last_epoch": 20000}] + self.reject(match="LR schedule") + + def test_saved_model_config_and_provenance_match(self): + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["hyper_parameters"]["config_tree"]["model"]["optimizer"]["lr"] = 5e-5 + self.reject(checkpoint, "model configuration") + for key in ["run_id", "source_revision", "data_context_sha256"]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["hyper_parameters"]["config_tree"]["run_provenance"][key] = "wrong" + self.reject(checkpoint, "provenance") + + def test_context_content_and_binding_both_match(self): + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["data_context"]["normalizer"]["mean"][0] = .1 + self.reject(checkpoint, "data context") + self.checkpoint["inference_binding"]["data_context_sha256"] = "c"*64 + self.reject(match="data context") + + def test_sampler_cursor_and_all_topology_fields_checked(self): + self.checkpoint["PixelProprioDataModule"]["consumed_batches"] = 0 + self.reject(match="sampler cursor") + self.checkpoint["PixelProprioDataModule"]["consumed_batches"] = 20000 + for key, value in [("world_size", 2), ("global_batch_size", 16), + ("per_rank_batch_size", 4), ("training_sharding", "wrong")]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["PixelProprioDataModule"]["distributed_sampling"][key] = value + self.reject(checkpoint, "sampler cursor") + + def test_missing_rank_or_old_rng_schema_rejected(self): + self.checkpoint["mot_rng_state"]["ranks"].pop() + self.reject(match="distributed RNG") + + def test_cpu_cuda_numpy_python_rng_payloads_checked(self): + for key, value in [("torch_cpu", torch.ones(3, dtype=torch.float32)), + ("torch_cpu", torch.ones(3, dtype=torch.uint8)), + ("torch_cuda", []), ("numpy", ("invalid",)), + ("python", (0, (), None)), ("cuda_scope", "all_devices")]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["mot_rng_state"]["ranks"][0][key] = value + self.reject(checkpoint, "RNG") + + def test_no_smoke_or_invalid_production_step(self): + for step in [True, 0, 50, 38863, 300000]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["global_step"] = step + self.reject(checkpoint, "scheduled production") + + def test_missing_malformed_or_nonfinite_loops_rejected(self): + for loops in [None, {}, {"fit_loop": {}}, {"fit_loop": {"bad": float("nan")}}]: + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["loops"] = loops + self.reject(checkpoint) + + def test_optimizer_loop_counter_mismatch_rejected(self): + self.checkpoint["loops"]["fit_loop"]["epoch_loop.automatic_optimization.optim_progress"]["optimizer"]["step"]["total"]["completed"] = 0 + self.reject(match="loop counter") + + def test_unknown_batch_loop_structure_rejected(self): + self.checkpoint["loops"]["fit_loop"]["epoch_loop.batch_progress"]["total"]["completed"] = 1 + self.reject(match="batch counters") + + def test_boundary_state_is_recorded_without_inventing_replay_safety(self): + result = self.inspect()["loop_state"] + self.assertEqual(result["batch_completed"], 19999) + self.assertFalse(result["loop_training_resume_supported"]) + self.assertEqual(result["resume_loop_policy"], "reject_replayed_boundary_validation") + + def add_checkpoint(self, epoch, step, content=b"checkpoint"): + directory = Path(self.parent["output"])/"checkpoints" + directory.mkdir(parents=True, exist_ok=True) + path = directory/f"epoch_epoch={epoch:04d}-step_step={step:09d}.ckpt" + path.write_bytes(content) + latest = directory/"last.ckpt" + if latest.is_symlink(): + latest.unlink() + latest.symlink_to(path.name) + return directory, path + + def test_inventory_selects_numeric_highest_before_loading(self): + directory, _ = self.add_checkpoint(9999, 20000) + _, highest = self.add_checkpoint(1, 40000) + selected, rows = target.inventory_highest(directory, self.parent) + self.assertEqual(selected, highest) + self.assertEqual([row["global_step"] for row in rows], [20000, 40000]) + + def test_empty_newest_prevents_silent_fallback(self): + directory, _ = self.add_checkpoint(1, 20000) + self.add_checkpoint(2, 40000, b"") + with self.assertRaisesRegex(ValueError, "Newest checkpoint is empty"): + target.inventory_highest(directory, self.parent) + + def test_new_milestone_during_restore_invalidates_inspection(self): + directory, _ = self.add_checkpoint(1, 20000) + selected, inventory = target.inventory_highest(directory, self.parent) + target.verify_inventory_stable(directory, self.parent, selected, inventory) + self.add_checkpoint(2, 40000) + with self.assertRaisesRegex(ValueError, "inventory changed"): + target.verify_inventory_stable(directory, self.parent, selected, inventory) + + def test_historical_file_changed_during_restore_invalidates_inspection(self): + directory, first = self.add_checkpoint(1, 20000) + self.add_checkpoint(2, 40000) + selected, inventory = target.inventory_highest(directory, self.parent) + first.write_bytes(b"changed old milestone") + with self.assertRaisesRegex(ValueError, "inventory changed"): + target.verify_inventory_stable(directory, self.parent, selected, inventory) + + def test_wrong_last_symlink_or_duplicate_step_rejected(self): + directory, first = self.add_checkpoint(1, 20000) + self.add_checkpoint(2, 40000) + (directory/"last.ckpt").unlink() + (directory/"last.ckpt").symlink_to(first.name) + with self.assertRaisesRegex(ValueError, "last.ckpt"): + target.inventory_highest(directory, self.parent) + self.add_checkpoint(9, 20000) + with self.assertRaisesRegex(ValueError, "duplicate"): + target.inventory_highest(directory, self.parent) + + def test_no_symlink_checkpoint_or_unbounded_inventory(self): + directory, first = self.add_checkpoint(1, 20000) + other = directory/"epoch_epoch=0002-step_step=000040000.ckpt" + other.symlink_to(first.name) + with self.assertRaisesRegex(ValueError, "immutable checkpoint"): + target.inventory_highest(directory, self.parent) + with self.assertRaisesRegex(ValueError, "bounded directory"): + target.inventory_highest(directory.parent, self.parent) + + def save_json_binding(self, name, value): + path = self.root/name + path.write_text(json.dumps(value, sort_keys=True)) + return dict(path=str(path), sha256=target.sha(path)) + + def crash_child(self, parent, checkpoint, number): + binding = self.save_json_binding(f"request-{number}.json", parent) + receipt = dict(schema=target.SCHEMA, status="PASS", parent_request_sha256=binding["sha256"], + run_id=parent["run_id"], source_revision=parent["source_revision"], + global_step=int(target.MILESTONE.fullmatch(checkpoint.name)[2]), + checkpoint=dict(path=str(checkpoint), sha256=target.sha(checkpoint))) + child = copy.deepcopy(parent) + child["output"] = str(self.root/f"recovery-{number}") + child["crash_recovery"] = dict(schema="joint-crash-recovery-v1", parent_request=binding, + inspection=self.save_json_binding(f"inspection-{number}.json", receipt)) + return child + + def test_failed_recovery_before40k_uses_retained_ancestor20k(self): + _, original = self.add_checkpoint(1, 20000) + child = self.crash_child(self.parent, original, 1) + directory = Path(child["output"])/"checkpoints" + self.assertFalse(directory.exists()) + selected, rows = target.inventory_highest(directory, child) + self.assertEqual(selected, original) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["expected_sha256"], target.sha(original)) + target.verify_inventory_stable(directory, child, selected, rows) + + def test_two_repeated_early_failures_still_use_original20k(self): + _, original = self.add_checkpoint(1, 20000) + first = self.crash_child(self.parent, original, 1) + second = self.crash_child(first, original, 2) + # One empty checkpoint directory and one absent directory are both valid + # because neither failed recovery owns any20k milestone yet. + (Path(first["output"])/"checkpoints").mkdir(parents=True) + selected, rows = target.inventory_highest(Path(second["output"])/"checkpoints", second) + self.assertEqual(selected, original) + self.assertEqual([row["global_step"] for row in rows], [20000]) + + def test_foreign_or_repeated_output_ancestor_rejected(self): + _, original = self.add_checkpoint(1, 20000) + child = self.crash_child(self.parent, original, 1) + child["run_id"] = "foreign-run" + with self.assertRaisesRegex(ValueError, "inspection belongs|foreign scientific"): + target.inventory_highest(Path(child["output"])/"checkpoints", child) + child = self.crash_child(self.parent, original, 2) + child["output"] = self.parent["output"] + with self.assertRaisesRegex(ValueError, "output/cycle"): + target.inventory_highest(Path(child["output"])/"checkpoints", child) + + def test_ancestor_with_newer_unreconciled_checkpoint_rejected(self): + _, original = self.add_checkpoint(1, 20000) + child = self.crash_child(self.parent, original, 1) + self.add_checkpoint(2, 40000) + with self.assertRaisesRegex(ValueError, "newer-than-cutoff"): + target.inventory_highest(Path(child["output"])/"checkpoints", child) + + def test_resume_binding_cannot_select_checkpoint_outside_lineage(self): + _, original = self.add_checkpoint(1, 20000) + child = self.crash_child(self.parent, original, 1) + receipt = target.bound(child["crash_recovery"]["inspection"]) + receipt["checkpoint"]["path"] = str(self.root/"foreign"/original.name) + child["crash_recovery"]["inspection"] = self.save_json_binding("foreign-inspection.json", receipt) + with self.assertRaisesRegex(ValueError, "outside prior retained history"): + target.inventory_highest(Path(child["output"])/"checkpoints", child) + + def test_missing_retained_earlier_milestone_cannot_be_hidden(self): + directory, _ = self.add_checkpoint(2, 40000) + with self.assertRaisesRegex(ValueError, "Missing retained production milestone"): + target.inventory_highest(directory, self.parent) + + def test_native_allocation_exact_owner_node_job_and_restart(self): + env = dict(SLURM_JOB_ID="123", SLURMD_NODENAME="gpu19", SLURM_RESTART_COUNT="0") + native = "JobId=123 JobState=RUNNING UserId=paphiwetsa3(41003) NodeList=gpu19 NumNodes=1 Restarts=0" + result = target.verify_allocation("123", "gpu19", env=env, uid=41003, native_reader=lambda: native) + self.assertEqual(result["job_id"], "123") + for bad in [native.replace("RUNNING", "COMPLETING"), native.replace("gpu19", "gpu03"), + native.replace("41003", "41004"), native.replace("Restarts=0", "Restarts=1")]: + with self.assertRaisesRegex(ValueError, "Native inspector"): + target.verify_allocation("123", "gpu19", env=env, uid=41003, native_reader=lambda: bad) + with self.assertRaisesRegex(ValueError, "allocated compute"): + target.verify_allocation("123", "gpu19", env={}, uid=41003, native_reader=lambda: native) + + def toy_actual_restore(self, checkpoint): + """Exercise real torch strict load/AdamW; stub source integration only.""" + captured = self + + class Context: + def fingerprint(self): + return captured.context_hash + + def bind(self, model): + pass + + class Data: + def prepare_context(self, **kwargs): + return Context() + + def load_state_dict(self, state): + self.state = copy.deepcopy(state) + + def state_dict(self): + return self.state + + class Wrapper(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.layer = torch.nn.Linear(3, 2) + self.layer.expert_parameter_counts = lambda: {"image": 8} + self.model = types.SimpleNamespace(pipeline=types.SimpleNamespace( + stage_by_id=lambda stage: types.SimpleNamespace(model=self.layer))) + + def state_dict(self, *args, **kwargs): + return self.layer.state_dict(*args, **kwargs) + + def load_state_dict(self, state, strict=True): + return self.layer.load_state_dict(state, strict=strict) + + def instantiate(cfg, **kwargs): + if cfg is None: + return None + if cfg.get("_target_") == "test.Data": + return Data() + return torch.optim.AdamW(kwargs["params"], lr=cfg.lr, weight_decay=cfg.weight_decay) + + modules = { + "egomimic.pipeline.checkpoint_binding": types.SimpleNamespace( + validate_checkpoint_binding=lambda checkpoint, config, context: checkpoint["inference_binding"]), + "egomimic.pipeline.construction": types.SimpleNamespace(checkpoint_construction=contextlib.nullcontext), + "egomimic.pl_utils.pl_model": types.SimpleNamespace(ModelWrapper=Wrapper), + } + cfg = OmegaConf.create({**self.resolved, "data": {"_target_": "test.Data"}}) + parent = {**self.parent, "expert_parameters": {"image": 8}, "parameters": 8} + with patch.dict(sys.modules, modules), patch("hydra.utils.instantiate", side_effect=instantiate): + return target.strict_cpu_restore(checkpoint, cfg, parent, self.context_hash) + + def test_real_cpu_wrapper_and_adamw_restore_executes(self): + before = target.state_digest(self.checkpoint) + result = self.toy_actual_restore(self.checkpoint) + self.assertTrue(result["strict_full_wrapper_optimizer_reload"]) + self.assertTrue(result["actual_datamodule_restore"]) + self.assertEqual(result["actual_expert_parameters"], {"image": 8}) + self.assertEqual(target.state_digest(self.checkpoint), before) + + def test_actual_strict_wrapper_key_or_adamw_parameter_shape_failure(self): + checkpoint = copy.deepcopy(self.checkpoint) + checkpoint["state_dict"].pop("bias") + with self.assertRaisesRegex(RuntimeError, "Missing key"): + self.toy_actual_restore(checkpoint) + checkpoint = copy.deepcopy(self.checkpoint) + state = next(iter(checkpoint["optimizer_states"][0]["state"].values())) + state["exp_avg"] = state["exp_avg_sq"] = torch.ones(1) + with self.assertRaisesRegex(ValueError, "moment shape"): + self.toy_actual_restore(checkpoint) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_joint_crash_restore.py b/tests/test_joint_crash_restore.py new file mode 100644 index 000000000..d281188ed --- /dev/null +++ b/tests/test_joint_crash_restore.py @@ -0,0 +1,278 @@ +"""Real CPU Torch object/callback guard tests, not CUDA/DDP certification.""" + +import copy +import hashlib +import json +from pathlib import Path +import random +import sys +import tempfile +from types import SimpleNamespace +import unittest +from unittest.mock import patch +import uuid + +import numpy as np +import torch + + +from scripts.train.joint.crash_restore import ( + CrashRestoreProof, assert_actual_restore, equal_state, validate_rank_proofs, + verify_restored_state, weights_digest, +) + + +def capture(): + return dict(torch_cpu=torch.get_rng_state(), torch_cuda=[], numpy=np.random.get_state(), + python=random.getstate(), cuda_scope="current_device") + + +def fixture(): + network = torch.nn.Linear(2, 2) + optimizer = torch.optim.AdamW(network.parameters(), lr=1e-4, weight_decay=.01) + network(torch.ones(1, 2)).square().mean().backward(); optimizer.step(); optimizer.zero_grad() + for value in optimizer.state.values(): + value["step"] = torch.tensor(20000.) + sampling = dict(schema="mot-sampling-state-v1", consumed_batches=20000, + distributed_sampling=dict(world_size=4, global_batch_size=32)) + rng = dict(schema="mot_rng_state.distributed.v1", world_size=4, + ranks=[capture() for _ in range(4)]) + module = SimpleNamespace(state_dict=network.state_dict, + training_behavior=SimpleNamespace(_pending_rng=rng)) + data = SimpleNamespace(state_dict=lambda: sampling, consumed_batches=20000) + trainer = SimpleNamespace(world_size=4, precision="bf16-mixed", global_step=20000, + optimizers=[optimizer], lr_scheduler_configs=[], datamodule=data) + saved = copy.deepcopy(dict(global_step=20000, state_dict=network.state_dict(), + optimizer_states=[optimizer.state_dict()], lr_schedulers=[], + PixelProprioDataModule=sampling, mot_rng_state=rng)) + return trainer, module, saved, network + + +class LoadedStateTests(unittest.TestCase): + def test_exact_real_torch_model_and_adamw_state(self): + trainer, module, saved, _ = fixture() + self.assertIs(verify_restored_state(trainer, module, saved, start_step=20000, + world_size=4), trainer.optimizers[0]) + + def test_rejects_unrestored_or_wrong_dtype_weights(self): + trainer, module, saved, network = fixture() + with torch.no_grad(): + network.weight.add_(1.) + with self.assertRaisesRegex(ValueError, "Weights"): + verify_restored_state(trainer, module, saved, start_step=20000, world_size=4) + trainer, module, saved, network = fixture(); network.double() + with self.assertRaisesRegex(ValueError, "Weights"): + verify_restored_state(trainer, module, saved, start_step=20000, world_size=4) + + def test_rejects_changed_optimizer_moments_lr_and_step(self): + for field in ("exp_avg", "step", "lr"): + with self.subTest(field=field): + trainer, module, saved, _ = fixture() + if field == "lr": trainer.optimizers[0].param_groups[0]["lr"] = 5e-5 + else: next(iter(trainer.optimizers[0].state.values()))[field].add_(1.) + with self.assertRaisesRegex(ValueError, "AdamW"): + verify_restored_state(trainer, module, saved, start_step=20000, world_size=4) + + def test_rejects_matching_but_nonfinite_adamw_moments(self): + trainer, module, saved, _ = fixture() + next(iter(trainer.optimizers[0].state.values()))["exp_avg"].fill_(float("inf")) + saved["optimizer_states"][0] = copy.deepcopy(trainer.optimizers[0].state_dict()) + with self.assertRaisesRegex(ValueError, "finiteFP32"): + verify_restored_state(trainer, module, saved, start_step=20000, world_size=4) + + def test_rejects_scheduler_cursor_rng_or_world_reset(self): + for field in ("scheduler", "cursor", "rng", "world"): + with self.subTest(field=field): + trainer, module, saved, _ = fixture() + if field == "scheduler": trainer.lr_scheduler_configs = [object()] + elif field == "cursor": trainer.datamodule.consumed_batches = 0 + elif field == "rng": module.training_behavior._pending_rng = None + else: trainer.world_size = 2 + with self.assertRaises(ValueError): + verify_restored_state(trainer, module, saved, start_step=20000, world_size=4) + + def test_cpu_cannot_claim_four_rank_callback_start(self): + callback = CrashRestoreProof(checkpoint="none", inspection_path="none", + inspection_sha256="x", output="none", world_size=4, + start_step=20000) + with self.assertRaisesRegex(ValueError, "Actual four-rank CUDA"): + callback.on_train_start(None, None) + + def test_no_implicit_diagnostic_validation_cadence(self): + with self.assertRaisesRegex(ValueError, "explicit two-update"): + CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000, + validation_interval=2) + callback = CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000, + diagnostic=True, validation_interval=2) + self.assertTrue(callback.diagnostic) + + def test_rejects_replayed_validation_at_the_checkpoint_step(self): + callback = CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000) + with self.assertRaisesRegex(ValueError, "replayed"): + callback.on_validation_start(SimpleNamespace(sanity_checking=False, global_step=20000), None) + + def test_weights_digest_handles_scalar_buffers_and_dtypes(self): + first = weights_digest(dict(value=torch.tensor(1., dtype=torch.float32))) + self.assertNotEqual(first, weights_digest(dict(value=torch.tensor(1., dtype=torch.float64)))) + + def test_actual_save_hook_rejects_boundary_offcadence_and_duplicate_save(self): + callback = CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000) + for step in (20000, 20001, 20002, 39999): + with self.subTest(step=step), self.assertRaises(ValueError): + callback.on_save_checkpoint(SimpleNamespace(global_step=step), None, dict(global_step=step)) + callback.on_save_checkpoint(SimpleNamespace(global_step=40000), None, dict(global_step=40000)) + self.assertEqual(callback.observed_saves, [40000]) + with self.assertRaisesRegex(ValueError, "repeated a checkpoint"): + callback.on_save_checkpoint(SimpleNamespace(global_step=40000), None, dict(global_step=40000)) + + def test_actual_save_hook_checks_serialized_step_and_diagnostic_terminal(self): + callback = CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000, + diagnostic=True, validation_interval=2) + with self.assertRaisesRegex(ValueError, "boundary checkpoint"): + callback.on_save_checkpoint(SimpleNamespace(global_step=20002), None, dict(global_step=20000)) + with self.assertRaisesRegex(ValueError, "off-cadence"): + callback.on_save_checkpoint(SimpleNamespace(global_step=20004), None, dict(global_step=20004)) + callback.on_save_checkpoint(SimpleNamespace(global_step=20002), None, dict(global_step=20002)) + self.assertEqual(callback.observed_saves, [20002]) + + +class RNGObserverTests(unittest.TestCase): + def setUp(self): + self.previous = capture() + + def tearDown(self): + self.restore(self.previous) + + @staticmethod + def restore(state): + torch.set_rng_state(state["torch_cpu"]); np.random.set_state(state["numpy"]) + random.setstate(state["python"]) + + def callback(self): + return CrashRestoreProof(checkpoint="none", inspection_path="none", inspection_sha256="x", + output="none", world_size=4, start_step=20000) + + def test_observes_actual_cpu_numpy_python_rng_restore_without_consuming_rng(self): + torch.manual_seed(42); np.random.seed(42); random.seed(42); expected = capture() + torch.rand(2); np.random.random(); random.random() + implementation = SimpleNamespace(__name__="egomimic.pl_utils.world_model_behavior", + _restore_rng_state=self.restore, + _capture_rng_state=lambda **kw: capture()) + original = implementation._restore_rng_state + callback = self.callback() + with patch("scripts.train.joint.crash_restore.importlib.import_module", return_value=implementation): + callback._install_rng_observer(SimpleNamespace(training_behavior=object()), expected) + implementation._restore_rng_state(expected) + self.assertTrue(callback.rng_restore_observed) + self.assertTrue(equal_state(capture(), expected)) + self.assertIs(implementation._restore_rng_state, original) + + def test_rejects_broken_restore_and_restores_observer_on_exception(self): + expected = capture(); torch.rand(1) + implementation = SimpleNamespace(__name__="egomimic.pl_utils.world_model_behavior", + _restore_rng_state=lambda state: None, + _capture_rng_state=lambda **kw: capture()) + original = implementation._restore_rng_state; callback = self.callback() + with patch("scripts.train.joint.crash_restore.importlib.import_module", return_value=implementation): + callback._install_rng_observer(SimpleNamespace(training_behavior=object()), expected) + with self.assertRaisesRegex(ValueError, "not exactly restored"): + implementation._restore_rng_state(expected) + callback.on_exception(None, None, ValueError()) + self.assertIs(implementation._restore_rng_state, original) + + def test_rejects_another_saved_rank_rng(self): + expected = capture(); torch.rand(1); wrong = capture() + implementation = SimpleNamespace(__name__="egomimic.pl_utils.world_model_behavior", + _restore_rng_state=self.restore, + _capture_rng_state=lambda **kw: capture()) + callback = self.callback() + with patch("scripts.train.joint.crash_restore.importlib.import_module", return_value=implementation): + callback._install_rng_observer(SimpleNamespace(training_behavior=object()), expected) + with self.assertRaisesRegex(ValueError, "another saved rank"): + implementation._restore_rng_state(wrong) + callback.on_exception(None, None, ValueError()) + + +class RankProofTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory(); self.root = Path(self.temp.name) + self.inspection = dict(schema="joint-crash-checkpoint-inspection-v1", status="PASS", + global_step=20000, world_size=4, run_id="same-run", + source_revision="source", checkpoint=dict(path="/actual-checkpoint", sha256="f"*64)) + self.inspection_sha = "e"*64 + self.paths = [] + for rank in range(4): + row = dict(schema="joint-crash-restore-rank-v1", status="PASS", rank=rank, + world_size=4, start_step=20000, global_step=20002, actual_updates=2, + inspection_sha256=self.inspection_sha, checkpoint="/actual-checkpoint", + checkpoint_sha256="f"*64, + run_id="same-run", source_revision="source", diagnostic=False, + job_id="123", gpu_uuid="GPU-"+str(uuid.UUID(int=rank+1)), + strict_weights_optimizer_sampler_rng=True, actual_native_rank_rng_restored=True, + fp32_losses=2, finite_bf16_outputs=2, finite_fp32_gradients=2, + optimizer_update_lrs=[1e-4]*2, no_extra_save=True, no_replayed_validation=True, + save_cadence_guard_installed=True, observed_saves_through_proof=[], + initial_weights_sha256="a"*64, initial_all_rank_weights=["a"*64]*4, + weights_sha256="b"*64, all_rank_weights=["b"*64]*4) + path = self.root / f"restore-step000020000-rank{rank}.json" + path.write_text(json.dumps(row)); self.paths.append(path) + + def tearDown(self): + self.temp.cleanup() + + def check(self, **kwargs): + return validate_rank_proofs(self.root, inspection=self.inspection, + inspection_sha256=self.inspection_sha, + start_step=20000, expected_job_id="123", **kwargs) + + def mutate(self, rank, **values): + path = self.paths[rank]; row = json.loads(path.read_text()); row.update(values) + path.write_text(json.dumps(row)) + + def test_complete_counterfactual_shape_is_accepted_only_as_unit_fixture(self): + self.assertEqual(len(self.check()), 4) + + def test_missing_rank_wrong_job_uuid_or_rng_proof_fails(self): + mutations = [dict(job_id="124"), dict(gpu_uuid="GPU-"+str(uuid.UUID(int=1))), + dict(actual_native_rank_rng_restored=False), dict(optimizer_update_lrs=[5e-5]*2), + dict(weights_sha256="c"*64), dict(actual_updates=3, global_step=20003)] + original = self.paths[1].read_text() + for values in mutations: + with self.subTest(values=values): + self.paths[1].write_text(original); self.mutate(1, **values) + with self.assertRaises(ValueError): self.check() + self.paths[1].unlink() + with self.assertRaisesRegex(ValueError, "Missing immutable"): + self.check() + + def test_diagnostic_flag_cannot_substitute_for_production_restore(self): + self.mutate(0, diagnostic=True) + with self.assertRaisesRegex(ValueError, "current inspected recovery"): + self.check() + + def test_diagnostic_requires_actual_postfit_validation_proof(self): + for rank in range(4): self.mutate(rank, diagnostic=True) + with self.assertRaises(FileNotFoundError): self.check(diagnostic=True) + for rank in range(4): + (self.root / f"restore-validation-rank{rank}.json").write_text(json.dumps(dict( + schema="joint-crash-restore-validation-rank-v1", status="PASS", rank=rank, + world_size=4, job_id="123", start_step=20000, global_step=20002, + validation_steps=[20002], global_validation_batches=100, + local_validation_batches=25, inspection_sha256=self.inspection_sha, + observed_saves=[20002], + actual_scheduled_validation_after_optimizer=True, metrics={"image_mse": .01}))) + self.assertEqual(len(self.check(diagnostic=True)), 4) + path = self.root / "restore-validation-rank3.json" + row = json.loads(path.read_text()); row["global_validation_batches"] = 99 + path.write_text(json.dumps(row)) + with self.assertRaisesRegex(ValueError, "scheduled recovery validation"): + self.check(diagnostic=True) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_joint_crash_smoke.py b/tests/test_joint_crash_smoke.py new file mode 100644 index 000000000..96c8f583b --- /dev/null +++ b/tests/test_joint_crash_smoke.py @@ -0,0 +1,207 @@ +"""Real config/receipt rejection tests; no allocation or runtime PASS is claimed.""" + +import copy +import hashlib +import importlib +import json +from pathlib import Path +import sys +import tempfile +import unittest +from unittest.mock import patch +from types import SimpleNamespace + +from omegaconf import OmegaConf + +smoke = importlib.import_module("scripts.train.joint.crash_smoke") + + +def binding(path, value, text=False): + path.write_text(value if text else json.dumps(value)) + return dict(path=str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +class CrashSmokeTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.base = Path(self.temp.name) + self.parent = dict( + schema="joint-training-v2", phase="full", source_revision=smoke.SOURCE, + run_id="mot-shared-libero-image-discrete-lr1e-4-constant-s42-300k-20261021-v1", + active_streams=["image"], arm="discrete", parameters=88178688, + expert_parameters=dict(image=88178688), expert_widths=[768], + world_size=4, global_batch=32, per_rank_batch=8, target_steps=300000, + lr_policy=dict(schedule="constant", peak_lr=1e-4, warmup_optimizer_steps=0), + source=str(self.base/"source"), source_manifest=str(self.base/"manifest"), + source_manifest_sha256="a"*64, source_archive=str(self.base/"source.tar"), + source_archive_sha256="b"*64, data_manifest=str(self.base/"data"), + data_manifest_sha256="c"*64, data_cache=dict(root="/tmp/data"), + lambda_runtime=dict(node="gpu19"), runtime_localization=None, allocation=None, + output=str(self.base/"production"), local_logs=str(self.base/"oldlogs")) + self.parent["job_name"] = "joint-image-lr1e4-full-4gpu" + self.q = copy.deepcopy(self.parent) + self.q.update(schema=smoke.SCHEMA, output=str(self.base/"diagnostic"), + local_logs=str(self.base/"newlogs"), reserved_checkpoint_bytes=4_000_000_000, + job_name="joint-image-lr1e4-crash-diagnostic") + self.q["parent_request"] = binding(self.base/"parent.json", self.parent) + self.q["resolved"] = binding(self.base/"resolved.yaml", "model: {}\n", text=True) + self.checkpoint = self.base/"input.ckpt" + self.checkpoint.write_bytes(b"retained production bytes") + self.inspection = dict(schema="joint-crash-checkpoint-inspection-v1", status="PASS", + parent_request_sha256=self.q["parent_request"]["sha256"], + parent_resolved_sha256=self.q["resolved"]["sha256"], source_revision=smoke.SOURCE, + run_id=self.parent["run_id"], world_size=4, global_batch=32, global_step=20000, + strict_full_wrapper_optimizer_reload=True, sampler_all_rank_rng_verified=True, + checkpoint=dict(path=str(self.checkpoint), sha256=hashlib.sha256(self.checkpoint.read_bytes()).hexdigest())) + self.q["inspection"] = binding(self.base/"inspection.json", self.inspection) + manifest = binding(self.base/"FILES.sha256", "test content", text=True) + self.q["operational_service"] = dict(root=str(self.base), manifest=manifest["path"], + manifest_sha256=manifest["sha256"]) + self.allocation = patch("scripts.train.joint.crash_inspect.verify_allocation", return_value={}) + self.allocation.start() + self.restore = patch("scripts.train.joint.crash_restore.assert_actual_restore", return_value=[]) + self.restore.start() + + def tearDown(self): + self.allocation.stop() + self.restore.stop() + self.temp.cleanup() + + def test_request_exact_bindings_and_no_input_write(self): + before = self.checkpoint.read_bytes() + parent, inspected = smoke.validate_request(self.q) + self.assertEqual(parent, self.parent) + self.assertEqual(inspected["global_step"], 20000) + self.assertEqual(self.checkpoint.read_bytes(), before) + self.assertFalse(Path(self.q["output"]).exists()) + + def test_reject_changed_identity_and_runtime(self): + for key, bad in (("run_id", "other"), ("active_streams", ["image", "proprio"]), + ("world_size", 2), ("lambda_runtime", {"node": "gpu21"})): + q = copy.deepcopy(self.q);q[key] = bad + with self.subTest(key=key), self.assertRaises(ValueError): + smoke.validate_request(q) + + def test_reject_checkpoint_or_receipt_tamper(self): + self.checkpoint.write_bytes(b"changed") + with self.assertRaises(ValueError): smoke.validate_request(self.q) + + def test_reject_production_path_overlap(self): + for path in (self.parent["output"], self.parent["output"]+"/child", str(self.base)): + q = copy.deepcopy(self.q);q["output"] = path + with self.subTest(path=path), self.assertRaises(ValueError): smoke.validate_request(q) + + def test_config_unchanged_model_data_and_local_only_logging(self): + original = OmegaConf.create(dict( + model=dict(optimizer=dict(lr=1e-4, weight_decay=.01), scheduler=None, + pipeline=dict(science="unchanged")), + data=dict(batch_size=8, distributed_world_size=4, science="same draws"), + seed=42, runtime=dict(disable_tf32=True), joint=dict(max_tokens=3216), + world_model_attention=dict(backend="fa4"), paths=dict(output_dir=self.parent["output"]), + val_at_start=False, val_at_end=False, + trainer=dict(devices=4,precision="bf16-mixed",max_steps=300000,val_check_interval=20000), + logger=dict(wandb=dict(id=self.parent["run_id"], resume="never"),csv=dict()), + evaluator=dict(pixel_diagnostics=dict(log_wandb=True,artifact_dir="old")), + callbacks=dict(model_checkpoint=dict(), allocation_boundary=dict(), joint_smoke=dict()))) + before = OmegaConf.to_container(original, resolve=True) + cfg = smoke.phase_config(original, self.q, self.parent, self.inspection, + request_sha256="d"*64) + self.assertEqual(before, OmegaConf.to_container(original, resolve=True)) + self.assertEqual(cfg.model, original.model) + self.assertEqual(cfg.data, original.data) + self.assertEqual(set(cfg.logger), {"csv"}) + self.assertEqual(cfg.trainer.max_steps, 20002) + self.assertEqual(cfg.trainer.val_check_interval, 2) + self.assertFalse(cfg.evaluator.pixel_diagnostics.log_wandb) + self.assertTrue(cfg.callbacks.crash_restore_proof.diagnostic) + self.assertEqual(cfg.callbacks.model_checkpoint.save_top_k, -1) + self.assertFalse(Path(self.q["output"]).exists()) + + def _finalization_fixture(self, state="COMPLETED", exit_code="0:0", host="0", job="123"): + qpath = self.base/"request.json";qpath.write_text(json.dumps(self.q)) + output = Path(self.q["output"]);output.mkdir(exist_ok=True) + candidate = dict(schema="joint-crash-recovery-smoke-candidate-v1", + status="ACTUAL_DIAGNOSTIC_PASS_NATIVE_PENDING", request_sha256=smoke.sha(qpath), + job_id="123", finite_optimizer_updates=2, all_rank_full_state_restored=True, + source_revision=self.parent["source_revision"],parent_request_sha256=self.q["parent_request"]["sha256"], + inspection_sha256=self.q["inspection"]["sha256"],checkpoint_sha256=self.inspection["checkpoint"]["sha256"], + service_manifest_sha256=self.q["operational_service"]["manifest_sha256"],world_size=4,global_batch=32, + scheduled_validation_anchors=100,no_replayed_boundary_validation_or_save=True, + production_science_unchanged=True,parent_checkpoint_modified=False,production_wandb_written=False, + sample_finite=True,hidden_query_labels_nan=True,sample_steps=20) + proofs=[] + for rank in range(4): + for relative in (f"restore-proof/restore-step000020000-rank{rank}.json", + f"restore-proof/restore-validation-rank{rank}.json", + f"validation-proof/rank{rank}.json"): + path=output/relative;path.parent.mkdir(exist_ok=True) + row=dict(status="PASS",job_id="123",world_size=4,rank=rank, + start_step=20000,global_step=20002) + if relative.startswith("validation-proof"): + row.update(schema="joint-crash-diagnostic-validation-rank-v1", + scheduled_validation_steps=[20002],scheduled_validation_anchors=100, + local_validation_anchors=25,no_replayed_boundary_validation=True) + proofs.append(binding(path,row)) + logs=Path(self.q["local_logs"])/"csv/step000020000";logs.mkdir(parents=True,exist_ok=True) + (logs/"metrics.csv").write_text("step,val/libero_gripper/forecast/image_normalized_mse,val/libero_gripper/completion/image_normalized_mse\n20001,0.01,0.02\n") + candidate["rank_proofs"]=proofs + final=output/"checkpoints/diagnostic-step=000020002.ckpt";final.parent.mkdir(exist_ok=True) + final.write_bytes(b"synthetic full state") + candidate["final_diagnostic_checkpoint"]={"path":str(final),"sha256":smoke.sha(final), + "strict_full_wrapper_optimizer_reload":True} + (output/"crash-recovery-smoke.candidate.json").write_text(json.dumps(candidate)) + native = binding(self.base/"native.txt", f"JobId={job} JobState={state} ExitCode={exit_code} " + f"UserId=paphiwetsa3(41003) NodeList=gpu19 JobName={self.q['job_name']} Requeue=0 Restarts=0", text=True) + hostbind = binding(self.base/"host.txt", host, text=True) + return qpath, native, hostbind + + def test_native_running_is_not_diagnostic_pass(self): + path, native, host = self._finalization_fixture(state="RUNNING") + with self.assertRaises(ValueError): + smoke.finalize(path, self.q, native_binding=native,host_binding=host,job_id="123") + self.assertFalse((Path(self.q["output"])/"crash-recovery-smoke.json").exists()) + + def test_host_nonzero_and_wrong_job_cannot_finalize(self): + for opts in (dict(host="1"),dict(job="124"),dict(exit_code="1:0")): + with self.subTest(opts=opts): + path,native,host=self._finalization_fixture(**opts) + with self.assertRaises(ValueError): + smoke.finalize(path,self.q,native_binding=native,host_binding=host,job_id="123") + + def test_finalization_is_exclusive_and_binds_native(self): + path,native,host=self._finalization_fixture() + record=smoke.finalize(path,self.q,native_binding=native,host_binding=host,job_id="123") + self.assertEqual(record["status"],"PASS") + self.assertEqual(record["native"],native) + with self.assertRaises(FileExistsError): + smoke.finalize(path,self.q,native_binding=native,host_binding=host,job_id="123") + + def test_parent_current_state_and_duplicate_writer_guards(self): + native=binding(self.base/"parent-native.txt",f"JobId=122 JobState=CANCELLED " + f"UserId=paphiwetsa3(41003) JobName={self.parent['job_name']}",text=True) + self.q["failure"]=binding(self.base/"failure.json",dict(schema="joint-crash-failure-v1", + parent_request_sha256=self.q["parent_request"]["sha256"],run_id=self.parent["run_id"], + recovery_authorized=True,human_retired=False,job_id=122,native=native)) + with patch.dict(smoke.os.environ,SLURM_JOB_ID="123"): + with patch.object(smoke.subprocess,"run",return_value=SimpleNamespace( + returncode=0,stdout="JobState=COMPLETING",stderr="")): + with self.assertRaises(ValueError):smoke.assert_parent_terminal(self.q,self.parent) + results=[SimpleNamespace(returncode=0,stdout="JobState=CANCELLED",stderr=""), + SimpleNamespace(returncode=0,stdout="124|joint-image-lr1e4-other|PENDING\n",stderr="")] + with patch.object(smoke.subprocess,"run",side_effect=results): + with self.assertRaises(ValueError):smoke.assert_parent_terminal(self.q,self.parent) + results=[SimpleNamespace(returncode=0,stdout="JobState=CANCELLED",stderr=""), + SimpleNamespace(returncode=0,stdout="123|joint-image-lr1e4-this|RUNNING\n",stderr="")] + with patch.object(smoke.subprocess,"run",side_effect=results): + smoke.assert_parent_terminal(self.q,self.parent) + + def test_finalization_rejects_changed_rank_proof(self): + path,native,host=self._finalization_fixture() + proof=Path(self.q["output"])/"restore-proof/restore-step000020000-rank0.json" + proof.write_text("{}") + with self.assertRaises(ValueError): + smoke.finalize(path,self.q,native_binding=native,host_binding=host,job_id="123") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_joint_temporal_pair.py b/tests/test_joint_temporal_pair.py new file mode 100644 index 000000000..b2b808aa1 --- /dev/null +++ b/tests/test_joint_temporal_pair.py @@ -0,0 +1,252 @@ +"""Synthetic evidence tests + actual535d composed-config comparison; no GPU claim.""" +import copy +import json +from pathlib import Path +import uuid + +import pytest + +from scripts.train.joint import verify_temporal_pair as tp +from scripts.train.joint.contract import compare_configs +from scripts.train.joint.continuation import STABLE_FIELDS + +BASE = Path(__file__).resolve().parent / "fixtures/joint_temporal_pair_535d" + + +def template(arm, phase): + q = json.loads((BASE/f"{arm}-{phase}.template.json").read_text()) + return {k: v for k, v in q.items() if not k.startswith("_")} + + +def composed(arm, phase): + return json.loads((BASE/f"{arm}-{phase}.composed.json").read_text()) + + +@pytest.fixture +def pair(tmp_path, monkeypatch): + requests, states = [], {} + def put(name, obj, text=False): + path = tmp_path/name + path.write_text(obj if text else json.dumps(obj, sort_keys=True)) + return {"path": str(path), "sha256": tp.sha(path)} + for idx, arm in enumerate(("discrete", "continuous")): + q, smoke = template(arm, "full"), template(arm, "smoke") + service = {"root": "/exact/service", "manifest": "/exact/service/FILES.sha256", "manifest_sha256": "1"*64, "revision": "unit-fixture"} + q["operational_service"] = smoke["operational_service"] = service + job, context = str(900+idx), ("b" if idx else "c")*64 + smoke_binding = put(f"{arm}-smoke.json", smoke) + ranks = [] + for rank in range(4): + ranks.append(dict(status="FIT_PASS", rank=rank, world_size=4, job_id=job, + global_step=50, fp32_losses=50, bf16_training_outputs=50, + fp32_gradients_verified=True, validation_batches=100, local_validation_batches=25, + data_context_sha256=context, active_streams=["image", "proprio"], + metrics={f"val/libero_gripper/{kind}/{metric}": 0.01 for kind in ("forecast", "completion") for metric in ("image_normalized_mse", "proprio_normalized_mse")}, + paired_rgb_batch_sha256=[tp.digest([arm, rank, draw]) for draw in range(3)], + initial_weights=[("d" if idx else "e")*64]*4, all_rank_weights=["f"*64]*4, + weights_sha256="f"*64, optimizer_update_lrs=[1e-4]*50, + measured_updates=40, warm_wall_seconds=1.0, + initial_image_weights_sha256="a"*64, + gpu_uuid="GPU-"+str(uuid.UUID(int=rank+1)))) + completion = dict(schema="joint-mot-smoke-v1", status="PASS", job_id=job, world_size=4, + global_step=50, validation_batches=100, strict_full_wrapper_reload=True, + optimizer_sampler_all_rank_rng_verified=True, lr_scheduler_strictly_restored=True, + all_rank_weights_agree=True, sample_finite=True, sample_state_steps=20, + sample_action_steps=0, lr_warmup_steps=0, model_execution=q["model_execution"], + active_streams=q["active_streams"], expert_parameters=q["expert_parameters"], ranks=ranks) + preflight = dict(status="PREFLIGHT_PASS", request_sha256=smoke_binding["sha256"], + parameters=q["parameters"], expert_parameters=q["expert_parameters"], expert_widths=q["expert_widths"], + world_size=4, global_batch=32, + provenance={**{key:q[key] for key in ("source_revision", "source_manifest_sha256", "data_manifest_sha256")}, "data_context_sha256":context}) + capacity = dict(status="CAPACITY_PASS", active_streams=q["active_streams"], sample_all_streams_finite=True, + frames_per_example=16 if arm=="discrete" else 64, per_rank_batch=8, optimizer_update_lr=1e-4) + native = f"JobId={job} JobState=COMPLETED ExitCode=0:0 UserId=paphiwetsa3(41003) JobName={smoke['job_name']} NodeList=gpu03" + q["smoke"] = dict(job_id=job, request=smoke_binding, + completion=put(f"{arm}-completion.json", completion), preflight=put(f"{arm}-preflight.json", preflight), + capacity=put(f"{arm}-capacity.json", capacity), native=put(f"{arm}-native.txt", native, True), + host_exit=put(f"{arm}-exit.txt", "0", True), resolved=put(f"{arm}-resolved.json", composed(arm,"smoke"))) + requests.append(q);states[arm]=(completion, preflight, capacity, native) + real_sha = tp.sha + manifest_digests = {requests[0]["source_manifest"]:tp.SOURCE_MANIFEST, requests[0]["data_manifest"]:tp.DATA_MANIFEST} + monkeypatch.setattr(tp, "sha", lambda p: manifest_digests.get(str(p)) or real_sha(p)) + monkeypatch.setattr(tp, "compose_request", lambda q: composed(q["arm"], "full")) + def bind(): + return [put(f"{q['arm']}-prototype.json",q) for q in requests] + def replace_evidence(arm,key,obj,text=False): + q=next(q for q in requests if q['arm']==arm) + q['smoke'][key]=put(f"{arm}-{key}-changed",obj,text) + return requests, states, bind, put, replace_evidence + + +def test_actual_configs_compare_and_normalize(): + values=[] + for arm in tp.PROFILES: + q=template(arm,"full") + compare_configs(composed(arm,"full"),composed(arm,"smoke"),q) + values.append(tp.normalize_config(q,composed(arm,"full"))) + assert values[0]==values[1] + + +def test_pair_keeps_unpaired_proofs_and_exact_receipt_membership(pair): + requests,_,bind,put,_=pair + receipt=tp.verify(bind()) + assert receipt['status']=='PASS' and receipt['fairness']=='FAIR WITH CAVEATS' + a,b=receipt['evidence'] + assert a['ranks'][0]['initial_weights']!=b['ranks'][0]['initial_weights'] + assert a['ranks'][0]['paired_rgb_batch_sha256']!=b['ranks'][0]['paired_rgb_batch_sha256'] + assert a['data_context_sha256']!=b['data_context_sha256'] + request={**requests[0], 'temporal_comparison':put('receipt.json',receipt)} + assert tp.admit(request)==receipt + + +@pytest.mark.parametrize('change',[ + lambda q:q.update(arm='continuous'), + lambda q:q.update(lr_policy=dict(peak_lr=5e-5,schedule='constant',warmup_optimizer_steps=0)), + lambda q:q.update(world_size=2), + lambda q:q.update(parameters=1), + lambda q:q.update(active_streams=['image','proprio','action']), + lambda q:q.update(source_revision='0'*40), + lambda q:q.update(source_manifest_sha256='0'*64), + lambda q:q.update(data_manifest_sha256='0'*64), + lambda q:q.update(operational_service=dict(root='/different')), + lambda q:q.update(_template_status='PENDING'), + lambda q:q.update(continuation={}), + lambda q:q.update(allocation_boundaries=[100000,300000]), + lambda q:q.update(optimizer_comparison={}), +]) +def test_reject_changed_identity_and_controls(pair,change): + requests,_,bind,_,_=pair + change(requests[0]) + with pytest.raises(ValueError):tp.verify(bind()) + + +def test_reject_duplicate_jobs(pair): + requests,_,bind,_,_=pair + requests[1]['smoke']['job_id']=requests[0]['smoke']['job_id'] + with pytest.raises(ValueError,match='Shared GPU smoke'):tp.verify(bind()) + + +def test_reject_duplicate_ids(pair): + requests,_,bind,_,_=pair + requests[1]['run_id']=requests[1]['intended_full_run_id']=requests[0]['run_id'] + requests[1]['overrides']=[('logger.wandb.id='+requests[0]['run_id']) if v.startswith('logger.wandb.id=') else v for v in requests[1]['overrides']] + with pytest.raises(ValueError,match='Duplicate full'):tp.verify(bind()) + + +@pytest.mark.parametrize('native_change',[ + lambda t:t.replace('COMPLETED','FAILED'), + lambda t:t.replace('0:0','1:0'), + lambda t:t.replace('JobId=900','JobId=999'), + lambda t:t.replace('gpu03','gpu19'), +]) +def test_native_proof_revalidated(pair,native_change): + _,states,bind,_,replace=pair + replace('discrete','native',native_change(states['discrete'][3]),True) + with pytest.raises(ValueError):tp.verify(bind()) + + +def test_host_failure_rejected(pair): + _,_,bind,_,replace=pair + replace('discrete','host_exit','1',True) + with pytest.raises(ValueError,match='host failed'):tp.verify(bind()) + + +@pytest.mark.parametrize('change',[ + lambda c:c['ranks'][3].update(initial_image_weights_sha256='1'*64), + lambda c:c['ranks'][2].update(initial_image_weights_sha256='bad'), + lambda c:c['ranks'][0].update(optimizer_update_lrs=[5e-5]*50), + lambda c:c.update(strict_full_wrapper_reload=False), + lambda c:c.update(global_step=49), + lambda c:c['ranks'][0].update(metrics={'bad':float('nan')}), +]) +def test_gpu_evidence_revalidated(pair,change): + _,states,bind,_,replace=pair + c=copy.deepcopy(states['discrete'][0]);change(c) + replace('discrete','completion',c) + with pytest.raises(ValueError):tp.verify(bind()) + + +@pytest.mark.parametrize('path,value',[ + (('model','scheduler'),{'kind':'warmup'}), + (('model','optimizer','weight_decay'),0.1), + (('data','max_frames'),63), + (('data','validation_discrete_frames'),12), + (('joint','pose_dims','libero_gripper'),18), + (('evaluator','sample_seed'),99), + (('runtime','slurm_requeue_owner'),'runner'), +]) +def test_exact_config_policy(path,value): + q=template('continuous','full');cfg=composed('continuous','full');cur=cfg + for key in path[:-1]:cur=cur[key] + cur[path[-1]]=value + with pytest.raises(ValueError):tp.normalize_config(q,cfg) + + +def test_unexpected_difference_rejected(pair,monkeypatch): + _,_,bind,_,_=pair + original=tp.normalize_config + monkeypatch.setattr(tp,'normalize_config',lambda q,cfg:{**original(q,cfg),'unexpected':q['arm']}) + with pytest.raises(ValueError,match='Unexpected D/C'):tp.verify(bind()) + + +def test_checkpoint_policy_rejected(pair,monkeypatch): + requests,_,bind,_,replace=pair + full=composed('discrete','full');smoke=composed('discrete','smoke') + full['callbacks']['model_checkpoint']['save_top_k']=1 + smoke['callbacks']['model_checkpoint']['save_top_k']=1 + replace('discrete','resolved',smoke) + monkeypatch.setattr(tp,'compose_request',lambda q:full if q['arm']=='discrete' else composed('continuous','full')) + with pytest.raises(ValueError,match='retention'):tp.verify(bind()) + + +def test_receipt_tampering_rejected_even_when_rehashed(pair): + requests,_,bind,put,_=pair + receipt=tp.verify(bind());receipt['initial_image_weights_sha256']='0'*64 + request={**requests[0],'temporal_comparison':put('receipt.json',receipt)} + with pytest.raises(ValueError,match='freshly verified'):tp.admit(request) + + +def test_changed_member_rejected(pair): + requests,_,bind,put,_=pair + receipt=tp.verify(bind()) + request={**requests[0],'temporal_comparison':put('receipt.json',receipt),'extra_unbound_field':True} + with pytest.raises(ValueError,match='exact admitted'):tp.admit(request) + + +def test_prototype_tampering_rejected(pair): + _,_,bind,_,_=pair + bindings=bind();p=Path(bindings[0]['path']);p.write_text(p.read_text()+' ') + with pytest.raises(ValueError,match='Changed evidence'):tp.verify(bindings) + + +@pytest.mark.parametrize('second',['optimizer_comparison','modality_comparison','factorial_controls']) +def test_exclusive_dispatch(second): + with pytest.raises(ValueError,match='Exactly one'):tp.admission_mode({'temporal_comparison':{},second:{}}) + + +def test_no_implicit_dispatch(): + with pytest.raises(ValueError):tp.admission_mode({}) + assert tp.admission_mode({'optimizer_comparison':{}})=='optimizer_comparison' + + +def test_continuation_is_bound_and_fail_closed(pair): + requests,_,_,_,_=pair + assert 'temporal_comparison' in STABLE_FIELDS + with pytest.raises(ValueError,match='single-allocation'): + tp.admit({**requests[0],'temporal_comparison':{},'continuation':{}}) + + +def test_path_collision_guard(): + with pytest.raises(ValueError,match='Overlapping'):tp.disjoint(['/out/x','/out/x/logs']) + with pytest.raises(ValueError,match='Unsafe'):tp.disjoint(['/out/../data','/other']) + + +@pytest.mark.parametrize('override',[ + 'data.forecasting_probability=0.8', + 'modality_ablation.nodes.image_loss.weight=2', + 'modality_ablation.nodes.joint_noising.noising.maximum=0.5', +]) +def test_same_unrequested_change_on_both_arms_is_rejected(pair,override): + requests,_,bind,_,_=pair + for q in requests:q['overrides'].append(override) + with pytest.raises(ValueError,match='Additional or changed'):tp.verify(bind())