diff --git a/docs/integration/ELMO_SNAPSHOT_IMPORT_20261026.md b/docs/integration/ELMO_SNAPSHOT_IMPORT_20261026.md new file mode 100644 index 000000000..09784c656 --- /dev/null +++ b/docs/integration/ELMO_SNAPSHOT_IMPORT_20261026.md @@ -0,0 +1,39 @@ +# Elmo deployed source import (2026-10-26) + +This change imports the verified source snapshot used by the Elmo image, +proprioception and action model into maintained Git history. It has the real +parent `255fdaf5137f7043b35df7ef48a3c1b8adc70815`, published on +`codex/joint-modality-ablation-20261019-v1` in `GaTech-RL2/EgoVerse-graph`. + +The deployment recorded source label +`153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920`. Its original Git object and ancestry +could not be recovered. **This import does not claim that label is a parent or +ancestor, or reconstruct its missing author/commit history.** The new import +receives its own Git identity. Existing runs and checkpoints retain their +original source records. + +The independently retained inputs are: + +| Input | SHA-256 | +| --- | --- | +| `source-153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920.tar.gz` (2,061,765 bytes) | `5e3f11a43ae33411591d4149c77432ad355c0671f9764a68e18fdddbf180aa2f` | +| `source-153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920.files.sha256` | `208ba5fa51d0514fd27d00f0f202cbf73f6e31ab7f3537e26ca55004e4fe3f99` | + +All 1,030 manifest paths are reproduced byte-for-byte with archive executable +bits: 988 already matched this parent, 25 are updated, and 17 are added. Every +parent path outside the export is retained. This change imports the deployed +baseline only; the larger-model configuration, fused objective and boundary +sampling experiments belong in separate changes. + +Runtime packaging must use the exact 1,030-path snapshot allowlist, plus only +explicitly reviewed experimental additions. Retaining older recipes and tests +in Git does not select them for execution or certify their compatibility. The +native model is `egomimic.models.joint_image_proprio.JointImageProprioModel`; +the Elmo selectors live under `experiment/pixel_proprio/native_joint/`. +Operational services remain separately hash-bound; do not accidentally use +older operational modules retained from this parent. + +Byte equivalence and local CPU checks do not certify cluster readiness. A new +immutable source export/readback, explicit runtime/data bindings and matching +allocated optimization, validation and full-state/RNG smoke are required for +new experimental runs. No historical run identity is relabeled by this import. diff --git a/egomimic/eval/image_proprio_eval.py b/egomimic/eval/image_proprio_eval.py index 7e647d8f8..b1bbf1475 100644 --- a/egomimic/eval/image_proprio_eval.py +++ b/egomimic/eval/image_proprio_eval.py @@ -9,11 +9,7 @@ class ImageProprioEval(Eval): prediction_key = "image_proprio_prediction" def __init__( - self, - objective="visual_proprio", - sample_seed=42, - pixel_diagnostics=None, - prediction_streams=None, + self, objective="visual_proprio", sample_seed=42, pixel_diagnostics=None, active_streams_from_data=False ): if objective not in ( "visual_only", @@ -23,16 +19,7 @@ def __init__( ): raise ValueError("Unknown image/proprio evaluation objective") self.objective, self.sample_seed = objective, int(sample_seed) - self.prediction_streams = ( - None if prediction_streams is None else tuple(prediction_streams) - ) - if self.prediction_streams is not None and ( - not self.prediction_streams - or len(set(self.prediction_streams)) != len(self.prediction_streams) - or set(self.prediction_streams) - {"image", "proprio", "action"} - or (objective != "direct_action" and self.prediction_streams[0] != "image") - ): - raise ValueError("Declare unique supported prediction_streams") + self.active_streams_from_data = bool(active_streams_from_data) self.totals = {} self._completed_validation_step = None self._validation_step = None @@ -96,21 +83,13 @@ def _sampling_index(batch, batch_idx): if all(value is None for value in indexed): return int(batch_idx) if any(value is None for value in indexed): - raise ValueError( - "Global validation batch identities must cover every source" - ) + raise ValueError("Global validation batch identities must cover every source") indices = [] for item, values in zip(batch.values(), indexed): - if not isinstance(values, (tuple, list)) or len(values) != len( - item["task_kind"] - ): - raise ValueError( - "One global validation batch identity is required per sample" - ) + if not isinstance(values, (tuple, list)) or len(values) != len(item["task_kind"]): + raise ValueError("One global validation batch identity is required per sample") if any(type(value) is not int or value < 0 for value in values): - raise ValueError( - "Global validation batch identities must be nonnegative integers" - ) + raise ValueError("Global validation batch identities must be nonnegative integers") indices.extend(values) if len(set(indices)) != 1: raise ValueError("A validation batch must share one global batch identity") @@ -122,29 +101,26 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0): ) with torch.random.fork_rng(devices=devices): torch.manual_seed( - self.sample_seed - + self._sampling_index(batch, batch_idx) + self.sample_seed + self._sampling_index(batch, batch_idx) + 100000 * dataloader_idx ) output = self.model.forward_eval(batch) - names = ( - self.prediction_streams - or { - "visual_only": ("image",), - "visual_proprio": ("image", "proprio"), - "visual_state_action": ("image", "proprio", "action"), - "direct_action": ("action",), - }[self.objective] - ) + names = { + "visual_only": ("image",), + "visual_proprio": ("image", "proprio"), + "visual_state_action": ("image", "proprio", "action"), + "direct_action": ("action",), + }[self.objective] for source, item in batch.items(): - if self.prediction_streams is not None and set( - output[source][self.prediction_key] - ) != set(names): - raise ValueError( - "Evaluator prediction_streams differ from model output" - ) kinds = item["task_kind"] - for name in names: + selected_names = names + if self.active_streams_from_data: + selected_names = tuple(item['world']) + if (not {'image','proprio','camera'} <= set(selected_names) + or not set(selected_names) <= {'image','proprio','camera','action'} + or set(output[source][self.prediction_key]) != set(selected_names)): + raise ValueError('Actual multimodal prediction keys must match native world streams') + for name in selected_names: stream = item["world"][name] mask = ( stream["valid"] @@ -194,16 +170,12 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0): def _reduce_validation_totals(self): """Sum target-weighted errors, including keys absent on some ranks.""" self.global_validation_batches = self._validation_batches - if not ( - torch.distributed.is_available() and torch.distributed.is_initialized() - ): + if not (torch.distributed.is_available() and torch.distributed.is_initialized()): return self.totals keys_by_rank = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(keys_by_rank, sorted(self.totals)) keys = sorted({key for rank_keys in keys_by_rank for key in rank_keys}) - totals = torch.zeros( - len(keys), 2, dtype=torch.float64, device=self.model.device - ) + totals = torch.zeros(len(keys), 2, dtype=torch.float64, device=self.model.device) for index, key in enumerate(keys): if key in self.totals: numerator, denominator = self.totals[key] @@ -212,19 +184,13 @@ def _reduce_validation_totals(self): valid = (torch.isfinite(totals).all() & (totals >= 0).all()).to(torch.int64) torch.distributed.all_reduce(valid, op=torch.distributed.ReduceOp.MIN) if not valid.item(): - raise ValueError( - "Invalid accumulated validation metric on a distributed rank" - ) + raise ValueError("Invalid accumulated validation metric on a distributed rank") if keys: torch.distributed.all_reduce(totals, op=torch.distributed.ReduceOp.SUM) - batches = torch.tensor( - self._validation_batches, dtype=torch.int64, device=self.model.device - ) + batches = torch.tensor(self._validation_batches, dtype=torch.int64, device=self.model.device) torch.distributed.all_reduce(batches, op=torch.distributed.ReduceOp.SUM) self.global_validation_batches = int(batches.item()) - return { - key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys) - } + return {key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys)} def log_epoch_metrics(self): self._metrics_logged = False @@ -278,9 +244,7 @@ def on_validation_end(self): and any(denominator > 0 for _, denominator in self._global_totals.values()) ) if torch.distributed.is_available() and torch.distributed.is_initialized(): - success = torch.tensor( - int(complete), dtype=torch.int64, device=self.model.device - ) + success = torch.tensor(int(complete), dtype=torch.int64, device=self.model.device) torch.distributed.all_reduce(success, op=torch.distributed.ReduceOp.MIN) complete = bool(success.item()) if publication_error is not None: diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/common.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/common.yaml new file mode 100644 index 000000000..b87f2f684 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/common.yaml @@ -0,0 +1,98 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/video_continuous +- /model/pixel_proprio/video@native_joint.reference +- override /model: pixel_proprio/native_joint +- _self_ +# Reference the historical image constructor; do not reconstruct its defaults. +joint: + pose_dims: {libero_gripper: 9} + max_tokens: 12864 + compile_mode: default +world_model_execution: + compile_mode: ${joint.compile_mode} + compile_dynamic: true + compile_backend: inductor + token_execution: packed_valid + loss_reduction: mean_per_example +pixel_video: + max_tokens: ${joint.max_tokens} +modality_ablation: + active_streams: null + nodes: + joint_noising: + _target_: egomimic.pipeline.stages_joint_state_action.JointNoisingStage + noising: + _target_: egomimic.models.joint_image_proprio.CoupledImageProprioNoising + maximum: 1.0 + active_streams: ${modality_ablation.active_streams} + joint_denoiser: + _target_: egomimic.pipeline.stages_joint_state_action.JointDenoiserStage + model: + _target_: egomimic.models.joint_image_proprio.JointImageProprioModel + image_model: ${native_joint.reference.pipeline.stages.1.model} + pose_dims: ${joint.pose_dims} + proprio_width: 256 + mlp_ratio: 4 + active_streams: ${modality_ablation.active_streams} + image_loss: + _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage + stream: image + weight: 1.0 + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: mean_per_example + proprio_loss: + _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage + stream: proprio + weight: 1.0 + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: mean_per_example + joint_sampler: + _target_: egomimic.pipeline.stages_joint_state_action.JointSamplerStage + sampler: + _target_: egomimic.models.joint_image_proprio.JointImageProprioSampler + state_steps: 20 + stage_ids: null + graph: null + initialization: Native image construction first; vector initialization preserves CPU RNG + training_noise: Native image draws unchanged; optional tied proprio draws preserve image RNG + screening_budget_updates: 300000 +data: + image_only: false + task: visual_diffusion + couple_state_observations: true + max_frames: 64 + batch_size: 8 + distributed_world_size: 4 + validation_batch_size: 1 + validation_anchors_per_episode: 2 + num_workers: 2 + pin_memory: true + persistent_workers: true + prefetch_factor: 2 + balance_training_work: true +trainer: + strategy: ddp + devices: 4 + num_nodes: 1 + accumulate_grad_batches: 1 + use_distributed_sampler: false + max_steps: 300000 + val_check_interval: 20000 +callbacks: + model_checkpoint: + every_n_train_steps: 20000 + save_top_k: -1 + learning_rate_monitor: + _target_: lightning.pytorch.callbacks.LearningRateMonitor + logging_interval: step +runtime: + cluster_environment: torch_elastic + slurm_requeue_owner: none +logger: + wandb: + group: native-image-joint-state + resume: never + tags: [historical-image-core, native-compiled, no-action, paired-initialization] diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_only.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_only.yaml new file mode 100644 index 000000000..fa6092ffc --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_only.yaml @@ -0,0 +1,22 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/common +- _self_ +name: native_mot_continuous_image_only_s42_300k +description: Historical image core with minimal joint state extension, historical compilation preserved +modality_ablation: + active_streams: [image] + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, joint_sampler: 3} + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.joint_sampler} +evaluator: + objective: visual_only + +joint: + max_tokens: 12800 +data: + image_only: true + couple_state_observations: false diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_proprio.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_proprio.yaml new file mode 100644 index 000000000..97fae1ecf --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_proprio.yaml @@ -0,0 +1,17 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/common +- _self_ +name: native_mot_continuous_image_proprio_s42_300k +description: Historical image core with minimal joint state extension, historical compilation preserved +modality_ablation: + active_streams: [image, proprio] + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, joint_sampler: 4} + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.proprio_loss} + - ${modality_ablation.nodes.joint_sampler} +evaluator: + objective: visual_proprio diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_only.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_only.yaml new file mode 100644 index 000000000..90dbf3193 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_only.yaml @@ -0,0 +1,30 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/common +- _self_ +name: native_mot_discrete_image_only_s42_300k +description: Historical image core with minimal joint state extension, historical compilation preserved +modality_ablation: + active_streams: [image] + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, joint_sampler: 3} + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.joint_sampler} +evaluator: + objective: visual_only +joint: + pose_dims: {libero_gripper: 36} + max_tokens: 3200 +data: + image_only: true + couple_state_observations: false + mode: discrete + sampling_strategy: discrete_window + discrete_frames: 16 + stride: 4 + batch_size: 16 + distributed_world_size: 2 +trainer: + devices: 2 diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio.yaml new file mode 100644 index 000000000..3764477f5 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio.yaml @@ -0,0 +1,29 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/common +- _self_ +name: native_mot_discrete_image_proprio_s42_300k +description: Historical image core with minimal joint state extension, historical compilation preserved +modality_ablation: + active_streams: [image, proprio] + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, joint_sampler: 4} + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.proprio_loss} + - ${modality_ablation.nodes.joint_sampler} +evaluator: + objective: visual_proprio +joint: + pose_dims: {libero_gripper: 36} + max_tokens: 3216 +data: + mode: discrete + sampling_strategy: discrete_window + discrete_frames: 16 + stride: 4 + batch_size: 16 + distributed_world_size: 2 +trainer: + devices: 2 diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio_action.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio_action.yaml new file mode 100644 index 000000000..0e4a452dd --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio_action.yaml @@ -0,0 +1,41 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/discrete_image_proprio +- _self_ +name: native_mot_discrete_image_proprio_action_s42_300k +description: Historical shared image-state core with an optional real-command expert +joint: + action_dims: {libero_gripper: 28} + max_tokens: 3232 +modality_ablation: + active_streams: [image, proprio, action] + nodes: + joint_denoiser: + model: + action_dims: ${joint.action_dims} + action_width: 256 + action_loss: + _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage + stream: action + weight: 1.0 + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: mean_per_example + joint_sampler: + sampler: + action_steps: 20 + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, action_loss: 4, joint_sampler: 5} + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.proprio_loss} + - ${modality_ablation.nodes.action_loss} + - ${modality_ablation.nodes.joint_sampler} +data: + task: visual_state_action +evaluator: + objective: visual_state_action +logger: + wandb: + tags: [historical-image-core, native-compiled, joint-image-proprio-action, paired-initialization] diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous.yaml new file mode 100644 index 000000000..ab6e44181 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous.yaml @@ -0,0 +1,19 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/subtask_language_continuous +- _self_ +name: native_mot_elmo_pick_place_subtask_continuous +joint: + pose_dims: {eva_bimanual: 20, human_bimanual: 126} + action_dims: {eva_bimanual: 20} + camera_dims: {eva_bimanual: 9, human_bimanual: 9} + max_tokens: 19520 +pixel_video: + max_cameras: 3 +# Multiple private experts plus grad/no-grad validation need bounded decode +# specializations. Mandatory fullgraph compilation remains active, never eager. +world_model_execution: + recompile_limit: 64 + accumulated_recompile_limit: 256 +evaluator: + active_streams_from_data: true diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_boundary.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_boundary.yaml new file mode 100644 index 000000000..a1e67ee5c --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_boundary.yaml @@ -0,0 +1,15 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous +- _self_ +name: native_mot_elmo_continuous_subtask_boundary_uncertainty +# The held-out loader disables this policy and retains the original annotations. +# Training still uses independent frame noise and the original plain MSE. +data: + boundary_sampling: + schema: annotation-boundary-uncertainty-sampling-v1 + policy: annotation_boundary_uncertainty + boundary_probability: 0.5 + maximum_radius_seconds: 0.25 + duration_fraction: 0.1 + rng_namespace: 1112425561 diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_fused.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_fused.yaml new file mode 100644 index 000000000..4290e4081 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous_fused.yaml @@ -0,0 +1,32 @@ +# @package _global_ +# UNAPPLIED CANDIDATE: source publication, scale decision and actual smoke pending. +defaults: +- /experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_continuous +- _self_ +name: native_mot_continuous_subtask_language_time_fused_snr +fused_objective: + cap: 5.0 + tau_seconds: 1.0 + # Raw capped objective. Final experiment scale requires an explicit decision; + # any calibrated constant must come from the actual TRAIN sampling population. + fixed_scale: 1.0 + loss: + _target_: egomimic.models.time_fused_snr.TimeAwareFusedSNRMSE + cap: ${fused_objective.cap} + tau_seconds: ${fused_objective.tau_seconds} + fixed_scale: ${fused_objective.fixed_scale} + reduction: mean_per_example +modality_ablation: + nodes: + image_loss: + log_objective: true + loss: ${fused_objective.loss} + proprio_loss: + log_objective: true + loss: ${fused_objective.loss} + action_loss: + log_objective: true + loss: ${fused_objective.loss} + camera_loss: + log_objective: true + loss: ${fused_objective.loss} diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_discrete.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_discrete.yaml new file mode 100644 index 000000000..3f2897055 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/elmo_pick_place_subtask_discrete.yaml @@ -0,0 +1,19 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/subtask_language +- _self_ +name: native_mot_elmo_pick_place_subtask_discrete +joint: + pose_dims: {eva_bimanual: 80, human_bimanual: 504} + action_dims: {eva_bimanual: 80} + camera_dims: {eva_bimanual: 36, human_bimanual: 36} + max_tokens: 4880 +pixel_video: + max_cameras: 3 +# Multiple private experts plus grad/no-grad validation need bounded decode +# specializations. Mandatory fullgraph compilation remains active, never eager. +world_model_execution: + recompile_limit: 64 + accumulated_recompile_limit: 256 +evaluator: + active_streams_from_data: true diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language.yaml new file mode 100644 index 000000000..fef096d2f --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language.yaml @@ -0,0 +1,69 @@ +# @package _global_ +# Shared subtask profile; physical data/coordinates/dimensions must be bound by +# a verified preparation manifest. No LIBERO dimensions are inherited silently. +defaults: +- /experiment/pixel_proprio/native_joint/discrete_image_proprio_action +- _self_ +name: native_mot_subtask_language +description: Independent per-frame noise with reviewed subtask-only language conditioning +joint: + # Null explicitly clears inherited LIBERO maps before dataset leaves merge. + pose_dims: null + action_dims: null + camera_dims: null + max_tokens: ??? +pixel_video: + require_observed_context: false + max_cameras: ??? +modality_ablation: + active_streams: [image, proprio, action, camera] + nodes: + joint_denoiser: + language_key: language_features + model: + separate_proprio_experts: true + language_feature_dim: 1024 + language_width: 128 + camera_dims: ${joint.camera_dims} + camera_width: 256 + camera_max_count: ${pixel_video.max_cameras} + joint_sampler: + language_key: language_features + action_loss: + optional: true + camera_loss: + _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage + stream: camera + weight: 1.0 + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: mean_per_example + graph: + - ${modality_ablation.nodes.joint_noising} + - ${modality_ablation.nodes.joint_denoiser} + - ${modality_ablation.nodes.image_loss} + - ${modality_ablation.nodes.proprio_loss} + - ${modality_ablation.nodes.action_loss} + - ${modality_ablation.nodes.camera_loss} + - ${modality_ablation.nodes.joint_sampler} + stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, action_loss: 4, camera_loss: 5, joint_sampler: 6} +data: + training_conditioning: independent_frame_noise + camera_pose: true + action_embodiments: [eva_bimanual] + language_features_path: ??? + language_features_sha256: ??? + subtask_manifest_path: ??? + subtask_manifest_sha256: ??? +logger: + wandb: + tags: [native-image-proprio-action, independent-frame-noise, subtask-language] +# Homogeneous robot/human batches leave private banks inactive on each draw. +# Reuse the maintained conditional-BC DDP route; do not fabricate human actions +# or attach zero-valued dummy paths just to satisfy the reducer. +trainer: + strategy: ddp_find_unused_parameters_true +# Track actual updates of optional private banks in every retained checkpoint. +callbacks: + conditional_optimizer_usage: + _target_: egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language_continuous.yaml new file mode 100644 index 000000000..d0c0a5b53 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/native_joint/subtask_language_continuous.yaml @@ -0,0 +1,14 @@ +# @package _global_ +defaults: +- /experiment/pixel_proprio/native_joint/subtask_language +- _self_ +name: native_mot_continuous_subtask_language +data: + mode: continuous + sampling_strategy: full_episode_uniform + min_frames: 2 + max_frames: 64 + # New subtask evaluation uses the declared continuous sampling contract. + validation_sampling_strategy: null + validation_discrete_frames: null + validation_stride: null diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml index a87083315..2bcc30729 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml @@ -20,14 +20,6 @@ pixel_readout: data: image_only: false task: visual_state_action - max_frames: 128 # Preserve deferred stage2 sampling independently of stage1. -trainer: - max_steps: 1000 # Deferred readout keeps its original pilot budget. - val_check_interval: 250 -callbacks: - model_checkpoint: - every_n_train_steps: 250 - save_last: true evaluator: objective: visual_state_action logger: diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml index c18a53122..7173be946 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml @@ -4,7 +4,3 @@ defaults: - _self_ name: pixel_proprio_video_continuous_pilot description: gripper_pixel160_patch16_continuous_full_episode_s42 -trainer: - max_steps: 300000 -data: - max_frames: 64 # Uniform timestamp selection still spans the full episode. diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml index 3e38073e5..1c5690ade 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml @@ -5,8 +5,6 @@ defaults: name: pixel_proprio_video_discrete_pilot description: gripper_pixel160_patch16_discrete16_stride4_s42 -trainer: - max_steps: 300000 data: mode: discrete sampling_strategy: discrete_window diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml index 0a60b565e..63141a7c9 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml @@ -37,7 +37,7 @@ trainer: max_epochs: -1 limit_train_batches: 1.0 limit_val_batches: 1.0 - val_check_interval: 20000 + val_check_interval: 250 check_val_every_n_epoch: null max_steps: 1000 num_sanity_val_steps: 0 @@ -46,10 +46,10 @@ callbacks: model_checkpoint: filename: "epoch_{epoch:04d}-step_{step:09d}" every_n_epochs: 0 - every_n_train_steps: 20000 + every_n_train_steps: 250 save_on_train_epoch_end: false save_top_k: -1 - save_last: link + save_last: true runtime: strict_fp32: ${world_model_attention.strict_fp32} disable_tf32: ${world_model_attention.disable_tf32} diff --git a/egomimic/hydra_configs/model/pixel_proprio/native_joint.yaml b/egomimic/hydra_configs/model/pixel_proprio/native_joint.yaml new file mode 100644 index 000000000..f4532af91 --- /dev/null +++ b/egomimic/hydra_configs/model/pixel_proprio/native_joint.yaml @@ -0,0 +1,15 @@ +defaults: +- /model/pixel_proprio/video@_here_ +- _self_ +optimizer: + lr: 0.000005 +scheduler: + _target_: egomimic.utils.schedulers.linear_warmup_constant_scheduler + warmup_steps: 3000 +scheduler_interval: step +scheduler_frequency: 1 +training_behavior: + distributed_noise_seed: ${seed} +pipeline: + stage_ids: ${modality_ablation.stage_ids} + stages: ${modality_ablation.graph} diff --git a/egomimic/hydra_configs/model/pixel_proprio/video.yaml b/egomimic/hydra_configs/model/pixel_proprio/video.yaml index fbc5d875b..be62cca49 100644 --- a/egomimic/hydra_configs/model/pixel_proprio/video.yaml +++ b/egomimic/hydra_configs/model/pixel_proprio/video.yaml @@ -26,6 +26,7 @@ pipeline: model: _target_: egomimic.models.image_proprio.ImageProprioModel objective: visual_only + require_observed_context: ${oc.select:pixel_video.require_observed_context,true} phase_a_condition_on_proprio: false pose_dims: null action_dims: null @@ -34,7 +35,7 @@ pipeline: depth: 8 heads: 6 mlp_ratio: 4 - max_cameras: 2 + max_cameras: ${oc.select:pixel_video.max_cameras,2} use_camera_embedding: true use_modality_embedding: false use_embodiment_embedding: false @@ -56,6 +57,8 @@ pipeline: compile_mode: ${oc.select:world_model_execution.compile_mode,none} compile_dynamic: ${oc.select:world_model_execution.compile_dynamic,false} compile_backend: ${oc.select:world_model_execution.compile_backend,inductor} + compile_recompile_limit: ${oc.select:world_model_execution.recompile_limit,null} + compile_accumulated_recompile_limit: ${oc.select:world_model_execution.accumulated_recompile_limit,null} - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoLossStage weight: 1.0 loss: diff --git a/egomimic/models/attention_execution.py b/egomimic/models/attention_execution.py index 596870947..eb056c6a2 100644 --- a/egomimic/models/attention_execution.py +++ b/egomimic/models/attention_execution.py @@ -30,9 +30,7 @@ class AttentionLayout: max_seqlen_k: int = 0 -def prepare_attention_layout( - valid, *, query_valid=None, query_tokens=None, all_valid=None -): +def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_valid=None): """Prepare dense or packed attention without assuming prefix-only padding. ``query_valid=None`` retains every query, including for query chunks or @@ -46,12 +44,8 @@ def prepare_attention_layout( if batch < 1 or key_tokens < 1: raise ValueError("Attention requires nonempty batch and key dimensions") if query_valid is not None: - if ( - query_valid.ndim != 2 - or query_valid.dtype != torch.bool - or query_valid.shape[0] != batch - or query_valid.device != valid.device - ): + if (query_valid.ndim != 2 or query_valid.dtype != torch.bool + or query_valid.shape[0] != batch or query_valid.device != valid.device): raise ValueError("Query validity must be boolean [B,Q] on the key device") if query_tokens is not None and query_tokens != query_valid.shape[1]: raise ValueError("Query token count differs from query validity") @@ -71,8 +65,7 @@ def prepare_attention_layout( return AttentionLayout("dense", batch, query_tokens, key_tokens, key_mask) key_counts = valid.sum(-1, dtype=torch.int32) query_counts = ( - query_valid.sum(-1, dtype=torch.int32) - if query_valid is not None + query_valid.sum(-1, dtype=torch.int32) if query_valid is not None else torch.full((batch,), query_tokens, device=valid.device, dtype=torch.int32) ) # This one eager transfer also verifies the optional kernel preconditions. @@ -80,26 +73,16 @@ def prepare_attention_layout( if min(counts[1]) < 1: raise ValueError("Fused attention requires at least one valid key per example") if min(counts[0]) < 1: - raise ValueError( - "Packed attention requires at least one valid query per example" - ) + raise ValueError("Packed attention requires at least one valid query per example") zero = key_counts.new_zeros(1) return AttentionLayout( - "varlen", - batch, - query_tokens, - key_tokens, - key_mask, - query_indices=( - query_valid.reshape(-1).nonzero(as_tuple=True)[0] - if query_valid is not None - else None - ), + "varlen", batch, query_tokens, key_tokens, key_mask, + query_indices=(query_valid.reshape(-1).nonzero(as_tuple=True)[0] + if query_valid is not None else None), key_indices=valid.reshape(-1).nonzero(as_tuple=True)[0], cu_seqlens_q=torch.cat((zero, query_counts.cumsum(0, dtype=torch.int32))), cu_seqlens_k=torch.cat((zero, key_counts.cumsum(0, dtype=torch.int32))), - max_seqlen_q=max(counts[0]), - max_seqlen_k=max(counts[1]), + max_seqlen_q=max(counts[0]), max_seqlen_k=max(counts[1]), ) @@ -108,11 +91,9 @@ def prepare_temporal_rope(times, head_dim): if times.ndim != 2 or type(head_dim) is not int or head_dim < 2 or head_dim % 2: raise ValueError("RoPE requires [B,N] times and a positive even head width") half = head_dim // 2 - freq = torch.exp( - -math.log(10000) - * torch.arange(half, device=times.device, dtype=torch.float32) - / half - ) + freq = torch.exp(-math.log(10000) * torch.arange( + half, device=times.device, dtype=torch.float32 + ) / half) phase = times.float()[:, None, :, None] * freq return phase.cos(), phase.sin() @@ -120,40 +101,22 @@ def prepare_temporal_rope(times, head_dim): def resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size): """Validate time/x/y channel allocation and explicit image patch grid H/W.""" if type(head_dim) is not int or head_dim < 6 or head_dim % 2: - raise ValueError( - "Spatial-temporal RoPE requires an even head width of at least 6" - ) + raise ValueError("Spatial-temporal RoPE requires an even head width of at least 6") if axis_dims is None: if head_dim % 8: - raise ValueError( - "Automatic RoPE axis dimensions require head width divisible by 8" - ) + raise ValueError("Automatic RoPE axis dimensions require head width divisible by 8") axis_dims = (head_dim // 2, head_dim // 4, head_dim // 4) - if ( - not isinstance(axis_dims, Sequence) - or isinstance(axis_dims, (str, bytes)) - or len(axis_dims) != 3 - or any(type(value) is not int or value < 2 or value % 2 for value in axis_dims) - or sum(axis_dims) != head_dim - ): - raise ValueError( - "rope_axis_dims must be positive-even time/x/y dimensions summing to head width" - ) - if ( - not isinstance(spatial_grid_size, Sequence) - or isinstance(spatial_grid_size, (str, bytes)) - or len(spatial_grid_size) != 2 - or any(type(value) is not int or value < 1 for value in spatial_grid_size) - ): - raise ValueError( - "spatial_grid_size must explicitly name positive integer patch-grid height/width" - ) + if (not isinstance(axis_dims, Sequence) or isinstance(axis_dims, (str, bytes)) or len(axis_dims) != 3 + or any(type(value) is not int or value < 2 or value % 2 for value in axis_dims) + or sum(axis_dims) != head_dim): + raise ValueError("rope_axis_dims must be positive-even time/x/y dimensions summing to head width") + if (not isinstance(spatial_grid_size, Sequence) or isinstance(spatial_grid_size, (str, bytes)) or len(spatial_grid_size) != 2 + or any(type(value) is not int or value < 1 for value in spatial_grid_size)): + raise ValueError("spatial_grid_size must explicitly name positive integer patch-grid height/width") return tuple(axis_dims), tuple(spatial_grid_size) -def prepare_spatiotemporal_rope( - times, xy, head_dim, *, axis_dims=None, spatial_grid_size=None -): +def prepare_spatiotemporal_rope(times, xy, head_dim, *, axis_dims=None, spatial_grid_size=None): """FP32 axial rotations for time plus normalized image-patch center XY. Normalized centers scale by(grid_width/2,grid_height/2), producing centered @@ -166,33 +129,21 @@ def prepare_spatiotemporal_rope( concatenates axis frequency vectors in that order. A64/32/32 allocation means32 time,16 x and16 y *pairs*, not a contiguous-slice rotation. """ - axes, grid = resolve_spatiotemporal_rope_geometry( - head_dim, axis_dims, spatial_grid_size - ) + axes, grid = resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size) if times.ndim != 2 or times.dtype == torch.bool or times.is_complex(): raise ValueError("Spatial-temporal RoPE requires real-valued times[B,N]") - if xy is not None and ( - xy.shape != (*times.shape, 2) - or xy.device != times.device - or not xy.is_floating_point() - ): - raise ValueError( - "Spatial RoPE coordinates must be floating-point[B,N,2] on the time device" - ) + if xy is not None and (xy.shape != (*times.shape, 2) or xy.device != times.device + or not xy.is_floating_point()): + raise ValueError("Spatial RoPE coordinates must be floating-point[B,N,2] on the time device") positions = [times.float()] - positions.extend( - (torch.zeros_like(times, dtype=torch.float32),) * 2 - if xy is None - else (xy[..., 0].float() * (grid[1] / 2), xy[..., 1].float() * (grid[0] / 2)) - ) + positions.extend((torch.zeros_like(times, dtype=torch.float32),) * 2 if xy is None else + (xy[..., 0].float() * (grid[1] / 2), xy[..., 1].float() * (grid[0] / 2))) phases = [] for coordinate, dimension in zip(positions, axes): pairs = dimension // 2 - frequency = torch.exp( - -math.log(10000) - * torch.arange(pairs, device=times.device, dtype=torch.float32) - / pairs - ) + frequency = torch.exp(-math.log(10000) * torch.arange( + pairs, device=times.device, dtype=torch.float32 + ) / pairs) phases.append(coordinate[:, None, :, None] * frequency) phase = torch.cat(phases, dim=-1) return phase.cos(), phase.sin() @@ -201,9 +152,7 @@ def prepare_spatiotemporal_rope( def apply_temporal_rope(x, rope): cos, sin = rope left, right = x.float().chunk(2, dim=-1) - return torch.cat((left * cos - right * sin, left * sin + right * cos), -1).to( - x.dtype - ) + return torch.cat((left * cos - right * sin, left * sin + right * cos), -1).to(x.dtype) @dataclass(frozen=True) @@ -225,16 +174,10 @@ def load_flash_attention_kernels(backend, device): raise RuntimeError(f"{backend} requires Hopper or newer CUDA hardware") if backend == "flash4": from packaging.version import Version - if Version(torch.__version__.split("+")[0]) < Version("2.9"): - raise RuntimeError( - "flash4 requires the separately verified Torch>=2.9 runtime" - ) - module_names = ( - ("flash_attn.cute",) - if backend == "flash4" - else ("flash_attn_3.flash_attn_interface", "flash_attn_interface") - ) + raise RuntimeError("flash4 requires the separately verified Torch>=2.9 runtime") + module_names = (("flash_attn.cute",) if backend == "flash4" + else ("flash_attn_3.flash_attn_interface", "flash_attn_interface")) try: for index, module_name in enumerate(module_names): try: @@ -244,10 +187,8 @@ def load_flash_attention_kernels(backend, device): # Older official FA3 builds expose a top-level alias. Only # missing module/package names permit that compatibility path; # an installed backend with a missing dependency must fail. - if index == len(module_names) - 1 or error.name not in { - module_name, - module_name.split(".")[0], - }: + if (index == len(module_names) - 1 + or error.name not in {module_name, module_name.split(".")[0]}): raise dense, varlen = module.flash_attn_func, module.flash_attn_varlen_func except (ImportError, AttributeError) as error: @@ -265,12 +206,8 @@ def _output_tensor(result, expected): if len(result) != 2: raise RuntimeError("Unsupported fused-attention result tuple") result = result[0] - if ( - not torch.is_tensor(result) - or result.shape != expected.shape - or result.dtype != expected.dtype - or result.device != expected.device - ): + if (not torch.is_tensor(result) or result.shape != expected.shape + or result.dtype != expected.dtype or result.device != expected.device): raise RuntimeError("Unsupported fused-attention output shape/dtype/device") return result @@ -280,15 +217,10 @@ def execute_flash_attention(query, key, value, layout, kernels): if any(x.dtype != torch.bfloat16 for x in (query, key, value)): raise ValueError("Fused attention requires bfloat16 Q/K/V") b, h, nq, d = query.shape - if ( - b != layout.batch_size - or nq != layout.query_tokens - or key.shape != (b, h, layout.key_tokens, d) - or value.shape != key.shape - or key.device != query.device - or value.device != query.device - or layout.key_mask.device != query.device - ): + if (b != layout.batch_size or nq != layout.query_tokens + or key.shape != (b, h, layout.key_tokens, d) or value.shape != key.shape + or key.device != query.device or value.device != query.device + or layout.key_mask.device != query.device): raise ValueError("Prepared attention layout differs from Q/K/V shape/device") q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value)) if layout.kind == "dense": @@ -301,19 +233,11 @@ def execute_flash_attention(query, key, value, layout, kernels): q = q.index_select(0, layout.query_indices) k = k.reshape(-1, h, d).index_select(0, layout.key_indices) v = v.reshape(-1, h, d).index_select(0, layout.key_indices) - out = _output_tensor( - kernels.varlen( - q, - k, - v, - cu_seqlens_q=layout.cu_seqlens_q, - cu_seqlens_k=layout.cu_seqlens_k, - max_seqlen_q=layout.max_seqlen_q, - max_seqlen_k=layout.max_seqlen_k, - causal=False, - ), - q, - ) + out = _output_tensor(kernels.varlen( + q, k, v, cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), q) if layout.query_indices is not None: out = out.new_zeros(b * nq, h, d).index_copy(0, layout.query_indices, out) return out.reshape(b, nq, h, d).transpose(1, 2) @@ -329,32 +253,16 @@ def execute_packed_flash_attention(query, key, value, layout, kernels): raise ValueError("Direct packed attention requires a varlen layout") if any(x.ndim != 3 or x.dtype != torch.bfloat16 for x in (query, key, value)): raise ValueError("Packed fused attention requires bfloat16 [tokens,H,D] Q/K/V") - query_count = ( - layout.query_indices.numel() - if layout.query_indices is not None - else layout.batch_size * layout.query_tokens - ) - if ( - query.shape[0] != query_count - or key.shape != (layout.key_indices.numel(), *query.shape[1:]) - or value.shape != key.shape - or key.device != query.device - or value.device != query.device - or layout.key_mask.device != query.device - ): - raise ValueError( - "Prepared packed attention layout differs from Q/K/V shape/device" - ) - return _output_tensor( - kernels.varlen( - query.contiguous(), - key.contiguous(), - value.contiguous(), - cu_seqlens_q=layout.cu_seqlens_q, - cu_seqlens_k=layout.cu_seqlens_k, - max_seqlen_q=layout.max_seqlen_q, - max_seqlen_k=layout.max_seqlen_k, - causal=False, - ), - query, - ) + query_count = (layout.query_indices.numel() if layout.query_indices is not None + else layout.batch_size * layout.query_tokens) + if (query.shape[0] != query_count + or key.shape != (layout.key_indices.numel(), *query.shape[1:]) + or value.shape != key.shape or key.device != query.device + or value.device != query.device or layout.key_mask.device != query.device): + raise ValueError("Prepared packed attention layout differs from Q/K/V shape/device") + return _output_tensor(kernels.varlen( + query.contiguous(), key.contiguous(), value.contiguous(), + cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), query) diff --git a/egomimic/models/compilation.py b/egomimic/models/compilation.py index 1b2c17194..1341d2ebe 100644 --- a/egomimic/models/compilation.py +++ b/egomimic/models/compilation.py @@ -6,19 +6,27 @@ from torch._functorch import config as _aot_config -def compile_region(function, *, mode="none", dynamic=False, backend="inductor"): +def compile_region(function, *, mode="none", dynamic=False, backend="inductor", + recompile_limit=None, accumulated_recompile_limit=None): if mode not in ("none", "default", "reduce-overhead", "max-autotune"): raise ValueError("Unknown compile mode") if type(dynamic) is not bool: raise ValueError("compile dynamic must be a boolean") if not isinstance(backend, str) or not backend: raise ValueError("compile backend must be a nonempty string") + limits = {} + if recompile_limit is not None or accumulated_recompile_limit is not None: + if (type(recompile_limit) is not int or type(accumulated_recompile_limit) is not int + or not 1 <= recompile_limit <= accumulated_recompile_limit <= 4096): + raise ValueError("Compile limits require positive bounded integers (recompile <= accumulated <=4096)") + limits = dict(recompile_limit=recompile_limit, + accumulated_recompile_limit=accumulated_recompile_limit) + if mode != "none" and not all(hasattr(torch._dynamo.config, key) for key in limits): + raise RuntimeError("Opt-in compilation limits require the verified modern Torch runtime") if mode == "none": return function if torch._dynamo.config.suppress_errors: - raise RuntimeError( - "Compilation requires suppress_errors=False; no eager fallback" - ) + raise RuntimeError("Compilation requires suppress_errors=False; no eager fallback") compiled = torch.compile( function, backend=backend, mode=mode, dynamic=dynamic, fullgraph=True ) @@ -33,7 +41,7 @@ def execute(*args, **kwargs): # whole model's gradient buckets (Torch 2.9 can emit invalid SymInt # subgraph outputs). This affects compilation only; DDP's reducer and # gradient synchronization remain active. Restore the caller's policy. - with torch._dynamo.config.patch(optimize_ddp=False): + with torch._dynamo.config.patch(optimize_ddp=False, **limits): if hasattr(_aot_config, "backward_pass_autocast"): with _aot_config.patch(backward_pass_autocast="off"): return compiled(*args, **kwargs) diff --git a/egomimic/models/image_proprio.py b/egomimic/models/image_proprio.py index 678f03b2b..217234a11 100644 --- a/egomimic/models/image_proprio.py +++ b/egomimic/models/image_proprio.py @@ -109,14 +109,22 @@ def __init__( compile_mode="none", compile_dynamic=False, compile_backend="inductor", + compile_recompile_limit=None, + compile_accumulated_recompile_limit=None, token_execution="padded", use_phase_embedding=True, noise_conditioning="additive", position_encoding="linear_xy_temporal_rope", rope_axis_dims=None, spatial_grid_size=None, + require_observed_context=True, ): super().__init__() + if type(require_observed_context) is not bool: + raise ValueError("require_observed_context must be boolean") + if not require_observed_context and objective == "direct_action": + raise ValueError("Direct action BC requires observed context") + self.require_observed_context = require_observed_context if type(strict_fp32) is not bool: raise ValueError("strict_fp32 must be a boolean") if strict_fp32 and attention_dtype != "native": @@ -207,23 +215,17 @@ def __init__( raise ValueError("Direct action BC has no diffusion noise conditioning") self.noise_conditioning = noise_conditioning if position_encoding not in ("linear_xy_temporal_rope", "spatiotemporal_rope"): - raise ValueError( - "position_encoding must be linear_xy_temporal_rope or spatiotemporal_rope" - ) + raise ValueError("position_encoding must be linear_xy_temporal_rope or spatiotemporal_rope") self.position_encoding = position_encoding if position_encoding == "spatiotemporal_rope": if width % heads: raise ValueError("Model width must be divisible by attention heads") - self.rope_axis_dims, self.spatial_grid_size = ( - resolve_spatiotemporal_rope_geometry( - width // heads, rope_axis_dims, spatial_grid_size - ) + self.rope_axis_dims, self.spatial_grid_size = resolve_spatiotemporal_rope_geometry( + width // heads, rope_axis_dims, spatial_grid_size ) else: if rope_axis_dims is not None or spatial_grid_size is not None: - raise ValueError( - "Spatial RoPE geometry requires position_encoding='spatiotemporal_rope'" - ) + raise ValueError("Spatial RoPE geometry requires position_encoding='spatiotemporal_rope'") self.rope_axis_dims = self.spatial_grid_size = None self.image_dim, self.width = image_dim, width self.allow_patch_compression = allow_patch_compression @@ -241,6 +243,8 @@ def __init__( self.loss_reduction = loss_reduction self.compile_mode, self.compile_dynamic = compile_mode, compile_dynamic self.compile_backend = compile_backend + self.compile_recompile_limit = compile_recompile_limit + self.compile_accumulated_recompile_limit = compile_accumulated_recompile_limit if token_execution not in ("padded", "packed_valid"): raise ValueError("token_execution must be padded or packed_valid") self.token_execution = token_execution @@ -282,8 +286,7 @@ def __init__( ) self.noise_level = nn.Sequential( nn.Linear(256 if noise_conditioning == "adaln_zero" else 1, width), - nn.SiLU(), - nn.Linear(width, width), + nn.SiLU(), nn.Linear(width, width) ) if use_phase_embedding: self.phase = nn.Embedding(1 if objective == "visual_only" else 2, width) @@ -304,13 +307,10 @@ def __init__( ) for _ in range(depth) ) - self.norm = nn.ModuleList( - [ - nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) - if noise_conditioning == "adaln_zero" - else nn.LayerNorm(width) - ] - ) + self.norm = nn.ModuleList([ + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if noise_conditioning == "adaln_zero" else nn.LayerNorm(width) + ]) if self.has_proprio_branch: self.norm.append( nn.Identity() if objective == "visual_only" else nn.LayerNorm(width) @@ -329,7 +329,9 @@ def __init__( # OptimizedModule: historical checkpoint parameter names stay intact. for block in self.blocks: block.configure_compilation( - mode=compile_mode, dynamic=compile_dynamic, backend=compile_backend + mode=compile_mode, dynamic=compile_dynamic, backend=compile_backend, + recompile_limit=compile_recompile_limit, + accumulated_recompile_limit=compile_accumulated_recompile_limit, ) compiled_names = ("_embed_image", "_decode_image") if noise_conditioning == "adaln_zero": @@ -343,6 +345,8 @@ def __init__( mode=compile_mode, dynamic=compile_dynamic, backend=compile_backend, + recompile_limit=compile_recompile_limit, + accumulated_recompile_limit=compile_accumulated_recompile_limit, ), ) @@ -452,15 +456,19 @@ def validate(self, batch, embodiment): ) if total > self.max_tokens: raise ValueError(f"Token budget exceeded: {total} > {self.max_tokens}") - # A query without any clean observation is not a conditional BC example. - conditioning = MODALITIES if self.phase_a_condition_on_proprio else ("image",) - context = torch.cat([batch[n]["observed"] for n in conditioning], 1) - checks.append( - ( - ~context.any(1).all(), - "Every sample requires an observed image or proprio token", + # Independent per-frame diffusion training uses noisy context at every + # timestamp. Preserve the historical clean-context guard by default; + # the explicit profile removes only this requirement, never shape, + # validity, finite-observation or padding checks. + if self.require_observed_context: + conditioning = MODALITIES if self.phase_a_condition_on_proprio else ("image",) + context = torch.cat([batch[n]["observed"] for n in conditioning], 1) + checks.append( + ( + ~context.any(1).all(), + "Every sample requires an observed image or proprio token", + ) ) - ) # One device-to-host synchronization for all value checks. Boolean # indexing/nonzero used here previously introduced additional syncs. for failed, (_, message) in zip( @@ -515,11 +523,8 @@ def forward( def _prepare_rope(self, times, xy=None): if self.position_encoding == "spatiotemporal_rope": return prepare_spatiotemporal_rope( - times, - xy, - self.blocks[0].head_dim, - axis_dims=self.rope_axis_dims, - spatial_grid_size=self.spatial_grid_size, + times, xy, self.blocks[0].head_dim, + axis_dims=self.rope_axis_dims, spatial_grid_size=self.spatial_grid_size, ) return prepare_temporal_rope(times, self.blocks[0].head_dim) @@ -537,16 +542,10 @@ def _prepare_execution(self, times, valid, image_xy=None): load_flash_attention_kernels(backend, joint_valid.device) layout = prepare_attention_layout(joint_valid, query_valid=joint_valid) if self.position_encoding == "spatiotemporal_rope" and image_xy is None: - raise ValueError( - "Image spatial-temporal attention requires patch coordinates" - ) - safe_xy = ( - None if image_xy is None else torch.where(valid[0][..., None], image_xy, 0) - ) - rope = tuple( - self._prepare_rope(t, safe_xy if index == 0 else None) - for index, t in enumerate(times) - ) + raise ValueError("Image spatial-temporal attention requires patch coordinates") + safe_xy = None if image_xy is None else torch.where(valid[0][..., None], image_xy, 0) + rope = tuple(self._prepare_rope(t, safe_xy if index == 0 else None) + for index, t in enumerate(times)) return layout, rope @staticmethod @@ -564,65 +563,39 @@ def _prepare_noise_frame_map(image): positions = valid[batch_index].nonzero(as_tuple=False).flatten() if positions.numel() == 0: continue - keys = torch.stack( - ( - image["time"][batch_index, positions].double(), - image["observed"][batch_index, positions].double(), - ), - dim=1, - ) + keys = torch.stack((image["time"][batch_index, positions].double(), + image["observed"][batch_index, positions].double()), dim=1) unique, inverse = torch.unique(keys, dim=0, return_inverse=True) - first = torch.full( - (unique.shape[0],), - positions.numel(), - device=positions.device, - dtype=torch.long, - ) - first.scatter_reduce_( - 0, - inverse, - torch.arange(positions.numel(), device=positions.device), - reduce="amin", - ) + first = torch.full((unique.shape[0],), positions.numel(), device=positions.device, dtype=torch.long) + first.scatter_reduce_(0, inverse, torch.arange(positions.numel(), device=positions.device), reduce="amin") representatives.append(batch_index * valid.shape[1] + positions[first]) groups[batch_index, positions] = inverse + offset offset += unique.shape[0] - indices = ( - torch.cat(representatives) - if representatives - else torch.zeros(1, device=valid.device, dtype=torch.long) - ) + indices = (torch.cat(representatives) if representatives else + torch.zeros(1, device=valid.device, dtype=torch.long)) return _NoiseFrameMap(indices, groups) def _embed_noise_frames(self, levels): # Embedding convention only: tau remains the original linear-corruption # strength in [0,1]. Multiplying by1000 does not discretize/change it. - frequency = torch.exp( - -math.log(10000) - * torch.arange(128, device=levels.device, dtype=torch.float32) - / 128 - ) + frequency = torch.exp(-math.log(10000) * torch.arange( + 128, device=levels.device, dtype=torch.float32 + ) / 128) angles = (1000 * levels.float())[:, None] * frequency[None] return self.noise_level(torch.cat((angles.cos(), angles.sin()), dim=-1)) def _image_noise_condition(self, image, levels, frame_map): effective = levels.masked_fill(image["observed"] | ~image["valid"], 0) grouped = effective.reshape(-1).index_select(0, frame_map.representatives) - reconstructed = grouped.index_select( - 0, frame_map.token_groups.reshape(-1) - ).reshape_as(effective) + reconstructed = grouped.index_select(0, frame_map.token_groups.reshape(-1)).reshape_as(effective) if ((effective != reconstructed) & image["valid"]).any(): - raise ValueError( - "Image noise levels must agree within each timestamp and observed/query group" - ) + raise ValueError("Image noise levels must agree within each timestamp and observed/query group") return self._embed_noise_frames_execute(grouped), frame_map.token_groups def _prepare_image_execution(self, image, valid): execution = self._prepare_image_attention_execution(image, valid) if self.noise_conditioning == "adaln_zero": - return _ConditionedImageExecution( - execution, self._prepare_noise_frame_map(image) - ) + return _ConditionedImageExecution(execution, self._prepare_noise_frame_map(image)) return execution def _prepare_image_attention_execution(self, image, valid): @@ -646,11 +619,7 @@ def _prepare_image_attention_execution(self, image, valid): if any(block.attention_query_chunk_size for block in self.blocks): raise ValueError("Packed token execution requires unchunked attention") indices = layout.key_indices - metadata_keys = ( - ("xy", "observed", "camera") - if self.use_camera_embedding - else ("xy", "observed") - ) + metadata_keys = ("xy", "observed", "camera") if self.use_camera_embedding else ("xy", "observed") packed_image = { key: image[key].flatten(0, 1).index_select(0, indices).unsqueeze(0) for key in metadata_keys @@ -658,24 +627,17 @@ def _prepare_image_attention_execution(self, image, valid): packed_times = times.reshape(-1).index_select(0, indices).unsqueeze(0) packed_valid = valid.new_ones((1, indices.numel())) return _PackedImageExecution( - layout, - packed_image, - packed_valid, + layout, packed_image, packed_valid, self._prepare_rope(packed_times, packed_image["xy"]), ) - def _forward_packed_image( - self, values, levels, embodiment, execution, image_condition=None - ): + def _forward_packed_image(self, values, levels, embodiment, execution, image_condition=None): indices = execution.layout.key_indices packed_values = values.flatten(0, 1).index_select(0, indices).unsqueeze(0) packed_levels = levels.reshape(-1).index_select(0, indices).unsqueeze(0) if image_condition is not None: condition, frame_groups = image_condition - image_condition = ( - condition, - frame_groups.reshape(-1).index_select(0, indices).unsqueeze(0), - ) + image_condition = (condition, frame_groups.reshape(-1).index_select(0, indices).unsqueeze(0)) stream = self._embed_image_execute( execution.image, packed_values, packed_levels, execution.valid, "A", False ) @@ -686,38 +648,27 @@ def _forward_packed_image( for block in self.blocks: if self.activation_checkpointing and torch.is_grad_enabled(): stream = checkpoint( - block.forward_packed, - stream, - execution.valid, - attention_layout=execution.layout, - rope=execution.rope, + block.forward_packed, stream, execution.valid, + attention_layout=execution.layout, rope=execution.rope, image_condition=image_condition, use_reentrant=False, ) else: stream = block.forward_packed( - stream, - execution.valid, - attention_layout=execution.layout, - rope=execution.rope, + stream, execution.valid, + attention_layout=execution.layout, rope=execution.rope, image_condition=image_condition, ) - predicted = self._decode_image_execute( - stream, execution.valid, image_condition - ).squeeze(0) + predicted = self._decode_image_execute(stream, execution.valid, image_condition).squeeze(0) batch, tokens = execution.layout.batch_size, execution.layout.key_tokens - return ( - predicted.new_zeros(batch * tokens, self.image_dim) - .index_copy(0, indices, predicted) - .reshape(batch, tokens, self.image_dim) - ) + return predicted.new_zeros(batch * tokens, self.image_dim).index_copy( + 0, indices, predicted + ).reshape(batch, tokens, self.image_dim) def _embed_image(self, image, values, levels, image_valid, phase, direct): image_x = self.image_in(torch.where(image_valid[..., None], values, 0)) if self.position_encoding == "linear_xy_temporal_rope": - image_x = image_x + self.xy( - torch.where(image_valid[..., None], image["xy"], 0) - ) + image_x = image_x + self.xy(torch.where(image_valid[..., None], image["xy"], 0)) if self.use_camera_embedding: image_x = image_x + self.camera( image["camera"].masked_fill(~image_valid, 0) @@ -736,9 +687,7 @@ def _decode_image(self, stream, valid, image_condition=None): if self.noise_conditioning == "adaln_zero": condition, frame_groups = image_condition shift, scale = self.final_adaln(condition).chunk(2, dim=-1) - normalized = normalized * ( - 1 + MoTBlock._frame_values(scale, frame_groups) - ) + MoTBlock._frame_values(shift, frame_groups) + normalized = normalized * (1 + MoTBlock._frame_values(scale, frame_groups)) + MoTBlock._frame_values(shift, frame_groups) return self.image_out(normalized).masked_fill(~valid[..., None], 0) def _forward_validated( @@ -770,12 +719,7 @@ def _forward_validated( image["observed"][..., None], image["value"], image_values ) frame_map = None - if ( - self.noise_conditioning == "adaln_zero" - and phase == "A" - and not include_pose - and execution is None - ): + if self.noise_conditioning == "adaln_zero" and phase == "A" and not include_pose and execution is None: execution = self._prepare_image_execution(image, image_valid) if isinstance(execution, _ConditionedImageExecution): frame_map = execution.noise_frames @@ -787,17 +731,12 @@ def _forward_validated( # Phase B retains the existing multi-expert route. Packing here changes # execution only; no token, camera, or per-example boundary is removed. if self.token_execution == "packed_valid" and phase == "A" and not include_pose: - execution = ( - self._prepare_image_execution(image, image_valid) - if execution is None - else execution - ) + execution = (self._prepare_image_execution(image, image_valid) + if execution is None else execution) if isinstance(execution, _PackedImageExecution): - return { - "image": self._forward_packed_image( - values, levels, embodiment, execution, image_condition - ) - } + return {"image": self._forward_packed_image( + values, levels, embodiment, execution, image_condition + )} image_x = self._embed_image_execute( image, values, levels, image_valid, phase, direct ) @@ -864,17 +803,11 @@ def _forward_validated( ) else: streams = block( - streams, - times, - valid, - attention_layout=attention_layout, - rope=rope, + streams, times, valid, attention_layout=attention_layout, rope=rope, image_condition=image_condition, ) image_result = ( - None - if direct - else self._decode_image_execute(streams[0], image_valid, image_condition) + None if direct else self._decode_image_execute(streams[0], image_valid, image_condition) ) if include_pose: streams[1] = self.norm[1](streams[1]) diff --git a/egomimic/models/joint_image_proprio.py b/egomimic/models/joint_image_proprio.py new file mode 100644 index 000000000..6a17c1b7e --- /dev/null +++ b/egomimic/models/joint_image_proprio.py @@ -0,0 +1,596 @@ +"""Joint state denoising built around the unchanged native image model. + +Both variants run one active-stream layer loop using the historical image +embedding, compiled tensor regions and decoder. An optional narrow proprio +expert participates in the same shared attention operation. Optional command +tokens use another narrow expert in that same loop, without phase-B regression. +""" + +import math + +import torch +from torch import nn + +from egomimic.models.attention_execution import ( + apply_temporal_rope, prepare_attention_layout, load_flash_attention_kernels, + execute_packed_flash_attention, +) +from egomimic.models.compilation import compile_region +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, +) +from egomimic.models.world_model import MoTBlock + +ACTIVE_STREAMS = (("image",), ("image", "proprio"), + ("image", "action"), ("image", "proprio", "action")) +SUPPORTED_STREAMS = ACTIVE_STREAMS + (("image", "proprio", "camera"), + ("image", "proprio", "action", "camera")) + + +def paired_state_groups(world, name="proprio"): + """Bind one pose to each image timestamp, with identical observation roles.""" + image, pose = world["image"], world[name] + image_groups = torch.zeros_like(image["valid"], dtype=torch.long) + pose_groups = torch.zeros_like(pose["valid"], dtype=torch.long) + representatives, offset = [], 0 + if image["valid"].shape[0] != pose["valid"].shape[0]: + raise ValueError("Image/proprio batch sizes differ") + for b in range(image["valid"].shape[0]): + iv, pv = image["valid"][b], pose["valid"][b] + positions = iv.nonzero(as_tuple=True)[0] + times, imap = torch.unique(image["time"][b, iv], return_inverse=True) + ptimes, pmap = torch.unique(pose["time"][b, pv], return_inverse=True) + if not torch.equal(times, ptimes) or pmap.numel() != times.numel(): + raise ValueError("Image/proprio require one pose per valid timestamp") + observed = torch.empty_like(times, dtype=torch.bool) + observed[pmap] = pose["observed"][b, pv] + if not torch.equal(image["observed"][b, iv], observed[imap]): + raise ValueError("Image/proprio observation masks must be tied") + first = torch.full_like(times, positions.numel(), dtype=torch.long) + first.scatter_reduce_( + 0, imap, torch.arange(positions.numel(), device=positions.device), + reduce="amin", + ) + if "index" in image and "index" in pose: + indices = image["index"][b, positions[first]] + if not torch.equal(image["index"][b, iv], indices[imap]) or not torch.equal( + pose["index"][b, pv], indices[pmap] + ): + raise ValueError("Image/proprio discrete indices must be tied") + representatives.append(b * iv.numel() + positions[first]) + image_groups[b, iv], pose_groups[b, pv] = imap + offset, pmap + offset + offset += times.numel() + if offset == 0: + raise ValueError("Paired state requires valid timestamps") + return image_groups, pose_groups, torch.cat(representatives) + + +class CoupledImageProprioNoising: + """Native image corruption followed by tied-strength proprio corruption. + +Image-only corruption delegates exactly to the historical noiser and consumes +no vector draws. Vector corruption forks the RNG so it cannot advance the +subsequent image sequence. Commands draw independent strengths and noise after +the optional proprio draw; inactive streams consume no random numbers. + """ + + def __init__(self, maximum=1.0, active_streams=("image", "proprio")): + self.streams = tuple(active_streams) + if self.streams not in SUPPORTED_STREAMS: + raise ValueError("Expected image with optional proprio/action streams") + self.image_noising = LinearGaussianImageNoising(maximum) + + def __call__(self, world): + image = self.image_noising(world["image"]) + result = { + "values": {"image": image["image_values"]}, + "levels": {"image": image["noise_levels"]}, + "targets": {"image": image["target"]}, + "masks": {"image": image["mask"]}, + "metadata": {"image": image["metadata"]}, + } + clean = world["image"]["value"] + devices = [clean.device.index] if clean.is_cuda else [] + with torch.random.fork_rng(devices=devices): + for name in self.streams[1:]: + if name not in world: + continue + if name == "proprio": + self._add_proprio(world, image, result, name) + else: + self._add_action(world, result, name) + return result + + def _add_proprio(self, world, image, result, name="proprio"): + image_groups, pose_groups, representatives = paired_state_groups(world, name) + grouped = image["noise_levels"].reshape(-1).index_select(0, representatives) + if ((image["noise_levels"] != grouped[image_groups]) & world["image"]["valid"]).any(): + raise ValueError("Image noise strengths must agree per timestamp") + pose = world[name] + known = pose["valid"] & pose.get("target_valid", pose["valid"]) + clean = torch.where((known | pose["observed"])[..., None], pose["value"], 0) + if not torch.isfinite(clean).all(): + raise ValueError("Nonfinite measured proprioception target") + mask = known & ~pose["observed"] + levels = grouped[pose_groups].masked_fill(~pose["valid"], 0) + # Missing tracked state is an unsupervised pure-noise query. Its cache + # placeholder is never a clean observation or a denoising target. + known = pose.get("target_valid", pose["valid"]) + levels = levels.masked_fill(pose["valid"] & ~known & ~pose["observed"], 1) + sigma = levels[..., None].to(clean.dtype) + pose_noise = torch.randn_like(clean) + noisy = ((1 - sigma) * clean + sigma * pose_noise).masked_fill( + ~pose["valid"][..., None], 0 + ) + result["values"][name] = noisy + result["levels"][name] = levels + result["targets"][name] = clean + result["masks"][name] = mask + result["metadata"][name] = { + "timestamps": pose["time"], "valid": pose["valid"], "observed": pose["observed"] + } + + def _add_action(self, world, result, name="action"): + stream = world[name] + valid, observed = stream["valid"], stream["observed"] + known = valid & stream.get("target_valid", valid) + clean = torch.where((known | observed)[..., None], stream["value"], 0) + if not torch.isfinite(clean).all(): + raise ValueError("Nonfinite recorded action command") + levels = self.image_noising.frame_noise_levels(stream, self.image_noising.maximum) + # Unsupervised command queries stay pure noise and never leak a clean label. + levels = levels.masked_fill(valid & ~known & ~observed, 1) + sigma = levels[..., None].to(clean.dtype) + noisy = (1 - sigma) * clean + sigma * torch.randn_like(clean) + result["values"][name] = noisy.masked_fill(~valid[..., None], 0) + result["levels"][name] = levels + result["targets"][name] = clean + result["masks"][name] = known & ~observed + result["metadata"][name] = { + "timestamps": stream["time"], "valid": valid, "observed": observed, + } + + +class _ProprioBlock(nn.Module): + """One private residual bank projected into the native image attention space.""" + + def __init__(self, image_block, width, mlp_ratio, compilation): + super().__init__() + self.heads, self.head_dim = image_block.heads, image_block.head_dim + attention_width = self.heads * self.head_dim + self.norm1 = nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + self.norm2 = nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + self.qkv = nn.Linear(width, 3 * attention_width) + self.out = nn.Linear(attention_width, width) + self.ffn = nn.Sequential( + nn.Linear(width, width * mlp_ratio), nn.GELU(), + nn.Linear(width * mlp_ratio, width), + ) + self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(width, 6 * width)) + nn.init.zeros_(self.adaln[-1].weight) + nn.init.zeros_(self.adaln[-1].bias) + for name in ("modulation", "project", "finish"): + object.__setattr__(self, name + "_execute", compile_region( + getattr(self, name), **compilation + )) + + def modulation(self, condition): + return self.adaln(condition).chunk(6, -1) + + def project(self, x, rope, mods, groups): + shift, scale = (MoTBlock._frame_values(v, groups) for v in mods[:2]) + normalized = self.norm1(x) * (1 + scale) + shift + q, k, v = self.qkv(normalized).reshape( + *x.shape[:2], 3, self.heads, self.head_dim + ).permute(2, 0, 3, 1, 4) + return apply_temporal_rope(q, rope), apply_temporal_rope(k, rope), v + + def finish(self, x, attended, valid, mods, groups): + gate, shift, scale, mlp_gate = ( + MoTBlock._frame_values(v, groups) for v in mods[2:] + ) + x = x + gate * self.out(attended.transpose(1, 2).flatten(2)) + x = x + mlp_gate * self.ffn(self.norm2(x) * (1 + scale) + shift) + return x.masked_fill(~valid[..., None], 0) + + +class _ProprioExpert(nn.Module): + def __init__(self, image_model, pose_dims, width, mlp_ratio): + super().__init__() + self.inputs = nn.ModuleDict({k: nn.Linear(d, width) for k, d in pose_dims.items()}) + self.outputs = nn.ModuleDict({k: nn.Linear(width, d) for k, d in pose_dims.items()}) + self.noise = nn.Sequential(nn.Linear(256, width), nn.SiLU(), nn.Linear(width, width)) + self.norm = nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + self.final_adaln = nn.Sequential(nn.SiLU(), nn.Linear(width, 2 * width)) + compilation = dict(mode=image_model.compile_mode, dynamic=image_model.compile_dynamic, + backend=image_model.compile_backend, + recompile_limit=image_model.compile_recompile_limit, + accumulated_recompile_limit=image_model.compile_accumulated_recompile_limit) + self.blocks = nn.ModuleList( + _ProprioBlock(block, width, mlp_ratio, compilation) for block in image_model.blocks + ) + for layer in self.outputs.values(): + nn.init.zeros_(layer.weight) + nn.init.zeros_(layer.bias) + nn.init.zeros_(self.final_adaln[-1].weight) + nn.init.zeros_(self.final_adaln[-1].bias) + for layer in (self.noise[0], self.noise[2]): + nn.init.normal_(layer.weight, std=0.02) + nn.init.zeros_(layer.bias) + for name in ("embed_noise", "decode"): + object.__setattr__(self, name + "_execute", compile_region( + getattr(self, name), **compilation + )) + + def embed_noise(self, levels): + frequency = torch.exp(-math.log(10000) * torch.arange( + 128, device=levels.device, dtype=torch.float32 + ) / 128) + angles = (1000 * levels.float())[:, None] * frequency[None] + return self.noise(torch.cat((angles.cos(), angles.sin()), -1)) + + def decode(self, x, valid, condition, groups, embodiment): + shift, scale = self.final_adaln(condition).chunk(2, -1) + normalized = self.norm(x) * (1 + MoTBlock._frame_values(scale, groups)) + normalized = normalized + MoTBlock._frame_values(shift, groups) + return self.outputs[embodiment](normalized).masked_fill(~valid[..., None], 0) + + +class JointImageProprioModel(nn.Module): + """Native image core and optional state/control experts with AdaLN language. + + Cached per-example language features add no tokens or attention edges. Each + expert fuses the shared projected feature with its own frame noise vector. + """ + + def __init__(self, image_model, pose_dims, proprio_width=256, mlp_ratio=4, + active_streams=("image", "proprio"), action_dims=None, action_width=256, + language_feature_dim=0, language_width=128, camera_dims=None, camera_width=256, camera_max_count=3, separate_proprio_experts=False): + super().__init__() + self.streams = tuple(active_streams) + if self.streams not in SUPPORTED_STREAMS: + raise ValueError("Expected image with optional proprio/action streams") + if not isinstance(image_model, ImageProprioModel) or image_model.has_proprio_branch: + raise ValueError("Provide the native image-only ImageProprioModel") + if (image_model.objective != "visual_only" or image_model.noise_conditioning != "adaln_zero" + or image_model.use_phase_embedding or image_model.use_modality_embedding + or image_model.use_embodiment_embedding or image_model.activation_checkpointing): + raise ValueError("Joint state requires native AdaLN image-only execution without extra conditioning") + if any(type(x) is not int or x < 1 for x in (proprio_width, action_width, mlp_ratio)): + raise ValueError("Proprio width and MLP ratio must be positive integers") + if type(language_feature_dim) is not int or language_feature_dim not in (0, 1024): + raise ValueError("language_feature_dim must be 0 (disabled) or 1024") + if type(language_width) is not int or language_width != 128: + raise ValueError("language_width must be 128") + if type(camera_max_count) is not int or camera_max_count < 1: + raise ValueError("camera_max_count must be positive integer") + if type(separate_proprio_experts) is not bool: + raise ValueError("separate_proprio_experts must be boolean") + self.separate_proprio_experts = separate_proprio_experts + self.camera_max_count = camera_max_count + self.language_feature_dim = language_feature_dim + self.language_width = language_width + self.image_model = image_model + self.pose_dims = ImageProprioModel._dimensions(pose_dims, "pose_dims") + self.action_dims = (ImageProprioModel._dimensions(action_dims, "action_dims") + if "action" in self.streams else None) + self.camera_dims = (ImageProprioModel._dimensions(camera_dims, "camera_dims") + if "camera" in self.streams else None) + self.expert_widths = (image_model.width,) + if len(self.streams) > 1: + # The native image model was initialized first. Vector initialization + # cannot perturb its weights or the caller's following CPU RNG draws. + # Construction is CPU/meta only, before Lightning moves the module. + if any(p.device.type not in ("cpu", "meta") for p in image_model.parameters()): + raise ValueError("Construct the joint model on CPU/meta before device placement") + with torch.random.fork_rng(devices=[]): + for name in self.streams[1:]: + dims = self._stream_dimensions(name) + width = {"proprio": proprio_width, "action": action_width, "camera": camera_width}[name] + if name == "proprio" and self.separate_proprio_experts: + self.proprio_banks = nn.ModuleDict({ + emb: _ProprioExpert(image_model, {emb: dim}, width, mlp_ratio) + for emb, dim in dims.items() + }) + else: + setattr(self, name, _ProprioExpert(image_model, dims, width, mlp_ratio)) + self.expert_widths += (width,) + if "camera" in self.streams: + with torch.random.fork_rng(devices=[]): + self.camera_identity = nn.Embedding(camera_max_count, camera_width) + if self.language_feature_dim: + # Preserve the complete native initialization and subsequent RNG + # sequence. The disabled path constructs no additional parameters. + if any(p.device.type not in ("cpu", "meta") for p in image_model.parameters()): + raise ValueError("Construct language conditioning on CPU/meta before device placement") + with torch.random.fork_rng(devices=[]): + self.language_projection = nn.Linear(language_feature_dim, language_width) + self.language_fusions = nn.ModuleDict({ + name: nn.Sequential(nn.Linear(width + language_width, width), + nn.SiLU(), nn.Linear(width, width)) + for name, width in zip(self.streams, self.expert_widths) + for name in ([f"proprio__{emb}" for emb in self.pose_dims] + if name == "proprio" and self.separate_proprio_experts else [name]) + }) + for module in (self.language_projection, self.language_fusions): + for layer in module.modules(): + if isinstance(layer, nn.Linear): + nn.init.normal_(layer.weight, std=0.02) + nn.init.zeros_(layer.bias) + + def _expert(self, name, embodiment): + if name == "proprio" and self.separate_proprio_experts: + return self.proprio_banks[embodiment] + return getattr(self, name) + + def _stream_dimensions(self, name): + return {"proprio": self.pose_dims, "action": self.action_dims, + "camera": self.camera_dims}[name] + + def active_streams(self, embodiment): + """Private interfaces can be absent, e.g. humans have no command expert.""" + return tuple(name for name in self.streams if name == "image" or + embodiment in self._stream_dimensions(name)) + + def image_parameter_state(self): + return self.image_model.state_dict() + + def expert_parameter_counts(self): + result = {"image": sum(p.numel() for p in self.image_model.parameters())} + for name in self.streams[1:]: + result[name] = sum(p.numel() for p in ( + self.proprio_banks.parameters() if name == "proprio" and self.separate_proprio_experts + else getattr(self, name).parameters())) + if "camera" in self.streams: + result["camera"] += sum(p.numel() for p in self.camera_identity.parameters()) + if self.language_feature_dim: + for name in self.streams: + keys = ([f"proprio__{emb}" for emb in self.pose_dims] + if name == "proprio" and self.separate_proprio_experts else [name]) + result[name] += sum(p.numel() for key in keys for p in self.language_fusions[key].parameters()) + result["language_shared"] = sum(p.numel() for p in self.language_projection.parameters()) + return result + + def _validate_language_features(self, world, language_features): + if not self.language_feature_dim: + if language_features is not None: + raise ValueError("language_features supplied to a language-disabled model") + return + image = world["image"]["value"] + if (not isinstance(language_features, torch.Tensor) + or language_features.shape != (image.shape[0], self.language_feature_dim) + or not language_features.is_floating_point()): + raise ValueError("language_features must be a floating B x 1024 tensor") + if language_features.device != image.device: + raise ValueError("language_features must be on the same device as the world") + features = language_features.to(self.language_projection.weight.dtype) + if not torch.isfinite(features).all(): + raise ValueError("language_features must be finite") + return features + + def _language_condition(self, name, condition, projected, representatives, token_count, embodiment=None): + if projected is None: + return condition + # Representatives index the original B x tokens layout, not the + # expert-major packed sequence. Each frame receives its own sample's + # instruction even when examples have different valid frame counts. + sample = torch.div(representatives, token_count, rounding_mode="floor") + features = projected.index_select(0, sample) + key = (f"proprio__{embodiment}" if name == "proprio" and self.separate_proprio_experts else name) + return condition + self.language_fusions[key](torch.cat((condition, features), -1)) + + def validate(self, world, embodiment): + streams = self.active_streams(embodiment) + if not set(streams).issubset(world): + raise ValueError("Required embodiment streams absent") + if "action" in self.streams and embodiment not in self.action_dims and "action" in world: + raise ValueError("Embodiment has no declared action interface") + self.image_model.validate({"image": world["image"]}, embodiment) + for name in streams[1:]: + pose = world[name] + dims = self._stream_dimensions(name) + value, valid, observed = pose["value"], pose["valid"], pose["observed"] + if name == "camera": + camera = pose.get("camera") + if (camera is None or camera.shape != valid.shape or camera.dtype != torch.long + or ((camera < 0) | (camera >= self.camera_max_count))[valid].any()): + raise ValueError("Camera pose stream requires valid explicit camera identities") + if embodiment not in dims or value.ndim != 3 or value.shape[-1] != dims[embodiment]: + raise ValueError(f"Invalid {name} embodiment/value shape") + if not value.is_floating_point() or valid.dtype != torch.bool or observed.dtype != torch.bool: + raise ValueError(f"Invalid {name} value/mask dtype") + if any(pose[k].shape != value.shape[:2] for k in ("valid", "observed", "time")): + raise ValueError(f"{name} metadata shape mismatch") + if (observed & ~valid).any() or not torch.isfinite(pose["time"].masked_fill(~valid, 0)).all(): + raise ValueError(f"Invalid {name} observation/timestamp") + if not torch.isfinite(torch.where(observed[..., None], value, 0)).all(): + raise ValueError(f"Nonfinite observed {name}") + if self.image_model.time_mode == "discrete": + if "index" not in pose or pose["index"].shape != valid.shape: + raise ValueError(f"Discrete {name} requires grid indices") + indices = pose["index"].masked_fill(~valid, 0) + if not torch.isfinite(indices).all() or (indices != indices.round()).any(): + raise ValueError("Discrete indices must be finite integers") + if sum(world[n]["value"].shape[1] for n in streams) > self.image_model.max_tokens: + raise ValueError("Joint token budget exceeded") + return paired_state_groups(world) if "proprio" in streams else None + + def _prepare_execution(self, world, streams): + """One layout policy for any active expert count; no cross-example keys.""" + native = self.image_model + valid = [world[name]["valid"] for name in streams] + times = [native._clock(world[name], v) for name, v in zip(streams, valid)] + layout, rope = native._prepare_execution(times, valid, world["image"]["xy"]) + packed = native.token_execution == "packed_valid" and not bool(torch.cat(valid, 1).all()) + if not packed: + return valid, times, layout, rope, None + if any(b.attention_backend not in ("flash3", "flash4") or + b.attention_dtype != "bfloat16" or b.attention_query_chunk_size + for b in native.blocks): + raise ValueError("Ragged packed_valid requires unchunked BF16 flash3 or flash4") + layout = prepare_attention_layout(torch.cat(valid, 1), query_valid=torch.cat(valid, 1)) + indices = [v.reshape(-1).nonzero(as_tuple=True)[0] for v in valid] + # Experts store their own packed tokens. Attention consumes example-major + # order, then its inverse restores expert-major order before private FFNs. + ids = torch.arange(valid[0].shape[0], device=valid[0].device) + sample_ids = torch.cat([ids[:, None].expand_as(v)[v] for v in valid]) + order = torch.argsort(sample_ids, stable=True) + packed_times = [t.reshape(-1).index_select(0, i)[None] for t, i in zip(times, indices)] + xy = world["image"]["xy"].flatten(0, 1).index_select(0, indices[0])[None] + rope = tuple(native._prepare_rope(t, xy if j == 0 else None) + for j, t in enumerate(packed_times)) + return valid, packed_times, layout, rope, (indices, order, torch.argsort(order)) + + @staticmethod + def _pack(value, indices): + return value.flatten(0, 1).index_select(0, indices).unsqueeze(0) + + def _attention(self, block, qkv, valid, layout, packing): + if packing is not None: + _, order, inverse = packing + dtype = qkv[0].dtype + kernels = load_flash_attention_kernels(block.attention_backend, qkv[0].device) + with torch.autocast(device_type=qkv[0].device.type, enabled=False): + qkv = [x.squeeze(0).transpose(0, 1).index_select(0, order).to(torch.bfloat16) + for x in qkv] + attended = execute_packed_flash_attention(*qkv, layout, kernels).to(dtype) + return attended.index_select(0, inverse).transpose(0, 1).unsqueeze(0) + mask = torch.cat(valid, 1)[:, None, None, :] if layout is None else layout.key_mask + if block.attention_query_chunk_size: + raise ValueError("Shared MoT execution requires unchunked attention") + return (block._apply_attention(*qkv, mask) if layout is None else + block._apply_attention(*qkv, mask, attention_layout=layout)) + + def forward(self, world, embodiment, values, levels, language_features=None): + streams = self.active_streams(embodiment) + groups = self.validate(world, embodiment) + features = self._validate_language_features(world, language_features) + projected = self.language_projection(features) if features is not None else None + native, image = self.image_model, world["image"] + valid, times, layout, rope, packing = self._prepare_execution(world, streams) + frame_map = native._prepare_noise_frame_map(image) + condition, image_groups = native._image_noise_condition(image, levels["image"], frame_map) + condition = self._language_condition("image", condition, projected, + frame_map.representatives, image["valid"].shape[1]) + conditions = {"image": (condition, image_groups)} + if "proprio" in streams: + paired_image, pose_groups, representatives = groups + grouped = levels["image"].reshape(-1).index_select(0, representatives) + for name, index in (("image", paired_image), ("proprio", pose_groups)): + tied = world[name]["valid"] + if name == "proprio": + tied = tied & (world[name].get("target_valid", tied) | world[name]["observed"]) + if ((levels[name] != grouped[index]) & tied).any(): + raise ValueError("Image/proprio noise strengths must be tied") + condition = self._language_condition("proprio", self._expert("proprio", embodiment).embed_noise_execute(grouped), + projected, representatives, image["valid"].shape[1], embodiment) + conditions["proprio"] = (condition, pose_groups) + for name in (n for n in streams if n in ("action", "camera")): + stream = world[name] + mapping = native._prepare_noise_frame_map(stream) + grouped = levels[name].reshape(-1).index_select(0, mapping.representatives) + reconstructed = grouped[mapping.token_groups] + if ((levels[name] != reconstructed) & stream["valid"]).any(): + raise ValueError("Action strengths must agree within each command timestamp") + condition = self._language_condition(name, self._expert(name, embodiment).embed_noise_execute(grouped), + projected, mapping.representatives, stream["valid"].shape[1]) + conditions[name] = (condition, mapping.token_groups) + tokens, execution_valid = [], [] + for i, name in enumerate(streams): + stream, level = world[name], levels[name] + if values[name].shape != stream["value"].shape or level.shape != stream["valid"].shape: + raise ValueError("Noised state shape mismatch") + if ((~torch.isfinite(level) | (level < 0) | (level > 1)) & stream["valid"]).any(): + raise ValueError("Noise strengths must be finite in [0,1]") + if (level.masked_select(stream["observed"]) != 0).any(): + raise ValueError("Observed state must be clean") + value = torch.where(stream["observed"][..., None], stream["value"], values[name]) + v, metadata = valid[i], stream + if packing is not None: + index = packing[0][i] + value, level = self._pack(value, index), self._pack(level, index) + v = v.new_ones(value.shape[:2]) + metadata = {k: self._pack(stream[k], index) for k in ("xy", "observed", "camera") if k in stream} + c, g = conditions[name] + conditions[name] = (c, self._pack(g, index)) + token = (native._embed_image_execute(metadata, value, level, v, "A", False) + if name == "image" else self._expert(name, embodiment).inputs[embodiment](value.masked_fill(~v[..., None], 0))) + if name == "camera": + token = token + self.camera_identity(metadata["camera"].masked_fill(~v, 0)) + tokens.append(token.masked_fill(~v[..., None], 0)) + execution_valid.append(v) + for depth, ib in enumerate(native.blocks): + projected, modulations = [], [] + for i, name in enumerate(streams): + c, g = conditions[name] + if name == "image": + mods = ib._image_modulation_executor(c) + qkv = ib._project_inputs_executor([tokens[i]], [times[i]], (rope[i],), mods, g) + else: + pb = self._expert(name, embodiment).blocks[depth] + mods = pb.modulation_execute(c) + qkv = pb.project_execute(tokens[i], rope[i], mods, g) + projected.append(qkv) + modulations.append(mods) + qkv = [torch.cat(parts, 2) for parts in zip(*projected)] + attended = self._attention(ib, qkv, execution_valid, layout, packing) + pieces = attended.split([x.shape[1] for x in tokens], dim=2) + for i, name in enumerate(streams): + x, v, g = tokens[i], execution_valid[i], conditions[name][1] + tokens[i] = (ib._finish_outputs_executor([x], pieces[i], [v], [x.shape[1]], modulations[i], g)[0] + if name == "image" else self._expert(name, embodiment).blocks[depth].finish_execute(x, pieces[i], v, modulations[i], g)) + result = {} + for i, name in enumerate(streams): + value = (native._decode_image_execute(tokens[i], execution_valid[i], conditions[name]) + if name == "image" else self._expert(name, embodiment).decode_execute(tokens[i], execution_valid[i], *conditions[name], embodiment)) + if packing is not None: + batch, count = valid[i].shape + value = value.new_zeros(batch * count, value.shape[-1]).index_copy( + 0, packing[0][i], value.squeeze(0)).reshape(batch, count, -1) + result[name] = value + return result + + +class JointImageProprioSampler: + """Native linear-DDIM image integration; joint states use the same schedule.""" + + def __init__(self, state_steps=20, action_steps=None): + LinearDDIMImageSampler._check_steps(state_steps) + self.state_steps = state_steps + if action_steps is not None: + if action_steps != state_steps: + raise ValueError("Shared native sampler requires equal state/action integration steps") + self.action_steps = action_steps + + @torch.no_grad() + def __call__(self, model, world, embodiment, language_features=None): + streams = model.active_streams(embodiment) + model.validate(world, embodiment) + model._validate_language_features(world, language_features) + if "action" in streams and getattr(self, "action_steps", None) != self.state_steps: + raise ValueError("Action sampling needs an explicit matching action-step contract") + current = {} + s = world["image"] + noise = torch.randn(s["value"].shape, device=s["value"].device, dtype=s["value"].dtype) + current["image"] = torch.where(s["observed"][..., None], s["value"], noise).masked_fill(~s["valid"][..., None], 0) + devices = [s["value"].device.index] if s["value"].is_cuda else [] + with torch.random.fork_rng(devices=devices): + for name in streams[1:]: + s = world[name] + noise = torch.randn_like(s["value"]) + current[name] = torch.where(s["observed"][..., None], s["value"], noise).masked_fill(~s["valid"][..., None], 0) + for j in range(self.state_steps, 0, -1): + level, next_level = j / self.state_steps, (j - 1) / self.state_steps + levels = {name: torch.full_like(world[name]["time"], level, dtype=torch.float32).masked_fill( + world[name]["observed"] | ~world[name]["valid"], 0 + ) for name in streams} + language = {} if language_features is None else {"language_features": language_features} + prediction = model(world, embodiment, current, levels, **language) + for name in streams: + s = world[name] + x0 = prediction[name].to(current[name].dtype) + candidate = (1 - next_level) * x0 + next_level / level * (current[name] - (1 - level) * x0) + current[name] = torch.where(s["observed"][..., None], s["value"], candidate).masked_fill(~s["valid"][..., None], 0) + return current diff --git a/egomimic/models/time_fused_snr.py b/egomimic/models/time_fused_snr.py new file mode 100644 index 000000000..de8bb80e8 --- /dev/null +++ b/egomimic/models/time_fused_snr.py @@ -0,0 +1,178 @@ +"""Parameter-free time-aware strongest-anchor weighting for continuous x0. + +This is an explicit fused-SNR-style heuristic, not Diffusion Forcing's EMA. +It changes the training objective only: corruption, model inputs, sampling RNG, +and unweighted evaluation stay with their existing owners. +""" + +import math +from collections.abc import Mapping + +import torch + + +class TimeAwareFusedSNRMSE: + """Weight x0 squared error using actual same-stream noisy observations. + + For x_s=(1-s)*x0+s*epsilon, let u=min(((1-s)/s)^2,cap)/cap. + Another temporal noise group contributes exp(-|dt|/tau)*u; h is the + strongest such contribution, or zero without another group. The numerator + weight is fixed_scale*cap*(u+(1-u)*h). The denominator remains the ORIGINAL + supervision count, including at zero weight. No batch normalization occurs. + + Grouping is within each sample/stream, by timestamp and actual noise role: + observed input, supervised query, or unknown pure-noise query. Repeated + camera/patch tokens cannot become extra witnesses. Same-time observed and + queried tokens may differ only because they are genuinely distinct roles. + Unknown queries must be pure noise; image strengths cannot supply their SNR. + + cap=5 and tau_seconds=1 are provisional experiment settings. fixed_scale=1 + means the raw capped objective; any alternative calibration must be frozen + externally from the training population, never recomputed by this class. + """ + + requires_objective_metrics = True + + def __init__(self, cap=5.0, tau_seconds=1.0, fixed_scale=1.0, + reduction="mean_per_example"): + if reduction not in ("mean_per_example", "token_mean"): + raise ValueError("Unknown masked MSE reduction") + for name, value in (("cap", cap), ("tau_seconds", tau_seconds), + ("fixed_scale", fixed_scale)): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(name + " must be a positive finite FP32 number") + # Constructor must also work under the graph's meta-device audit; + # these pure scalar checks must not instantiate/read a Tensor. + if not math.isfinite(value) or not 2.0 ** -149 <= value <= torch.finfo(torch.float32).max: + raise ValueError(name + " must be a positive finite FP32 number") + maximum = cap * fixed_scale + if not math.isfinite(maximum) or not 2.0 ** -149 <= maximum <= torch.finfo(torch.float32).max: + raise ValueError("Scaled weight cap must be positive and finite in FP32") + self.cap, self.tau_seconds = float(cap), float(tau_seconds) + self.fixed_scale, self.reduction = float(fixed_scale), reduction + + @staticmethod + def _tensor(value, shape, device, *, boolean=False, name): + if (not torch.is_tensor(value) or value.shape != shape + or value.device != device + or (value.dtype != torch.bool if boolean else not value.is_floating_point())): + raise ValueError(name + " has an incompatible shape, device or dtype") + return value + + @torch.no_grad() + def _weights_and_support(self, noise_levels, metadata, mask): + if not torch.is_tensor(mask) or mask.ndim != 2 or mask.dtype != torch.bool: + raise ValueError("The original BxN boolean supervision mask is required") + if not isinstance(metadata, Mapping): + raise ValueError("Actual timestamps/valid/observed metadata is required") + shape, device = mask.shape, mask.device + valid = self._tensor(metadata.get("valid"), shape, device, + boolean=True, name="valid") + observed = self._tensor(metadata.get("observed"), shape, device, + boolean=True, name="observed") + times = self._tensor(metadata.get("timestamps"), shape, device, + name="timestamps") + levels = self._tensor(noise_levels, shape, device, name="noise_levels") + if (observed & ~valid).any() or (mask & (~valid | observed)).any(): + raise ValueError("Supervision/observed roles disagree with valid tokens") + with torch.autocast(device_type=device.type, enabled=False): + levels, times = levels.detach().float(), times.detach().float() + if (not torch.isfinite(levels[valid]).all() + or not torch.isfinite(times[valid]).all() + or ((levels[valid] < 0) | (levels[valid] > 1)).any()): + raise ValueError("Valid noise strengths/timestamps must be finite and ranged") + unknown = valid & ~observed & ~mask + if (levels[observed] != 0).any() or (levels[unknown] != 1).any(): + raise ValueError("Observed inputs must be clean and unknown queries pure noise") + weights, support = torch.zeros_like(levels), torch.zeros_like(levels) + threshold = 1.0 / (1.0 + math.sqrt(self.cap)) + for b in range(shape[0]): + take = valid[b] + if not take.any(): + continue + unique_times, time_index = torch.unique(times[b, take], return_inverse=True) + role = torch.where(observed[b, take], 0, + torch.where(mask[b, take], 1, 2)) + groups, inverse = torch.unique(time_index * 3 + role, return_inverse=True) + count = groups.numel() + strength = levels[b, take] + low = torch.full((count,), float("inf"), device=device, dtype=torch.float32) + high = torch.full((count,), -float("inf"), device=device, dtype=torch.float32) + low.scatter_reduce_(0, inverse, strength, reduce="amin") + high.scatter_reduce_(0, inverse, strength, reduce="amax") + if not torch.equal(low, high): + raise ValueError("Repeated tokens in a temporal noise group disagree") + # Normalize BEFORE square to avoid overflow, and clamp the + # divisor at the clipping threshold so s=0 never divides by 0. + # Clipping the ratio also gives exact u(1)=0 for very small caps + # whose threshold rounds to 1. This equals clipped SNR / cap. + ratio = ((1.0 - low) / low.clamp_min(threshold)) / math.sqrt(self.cap) + u = ratio.clamp(0.0, 1.0).square() + u = torch.where(low == 0, torch.ones_like(u), + torch.where(low == 1, torch.zeros_like(u), u)) + group_times = unique_times[torch.div(groups, 3, rounding_mode="floor")] + distance = (group_times[:, None] - group_times[None, :]).abs() + candidates = torch.exp(-distance / self.tau_seconds) * u[None, :] + candidates.fill_diagonal_(0.0) + h = candidates.amax(dim=1) + raw_weight = self.cap * (u + (1.0 - u) * h) + if (not torch.isfinite(raw_weight).all() + or (raw_weight < 0).any() or (raw_weight > self.cap).any()): + raise ValueError("Nonfinite or out-of-range fused-SNR weight") + weights[b, take] = raw_weight[inverse] * self.fixed_scale + support[b, take] = h[inverse] + if not torch.isfinite(weights).all(): + raise ValueError("Nonfinite scaled fused-SNR weight") + return weights, support + + def weights(self, noise_levels, metadata, mask): + """Return detached FP32 token weights without reading predictions/labels.""" + return self._weights_and_support(noise_levels, metadata, mask)[0] + + def _reduce(self, values, mask): + count = mask.float() + if self.reduction == "mean_per_example": + denominator = count.sum(1).clamp_min(1.0) + return (values.sum(1) / denominator).mean() + return values.sum() / count.sum().clamp_min(1.0) + + def loss_and_metrics(self, prediction, target, mask, noise_levels=None, + metadata=None, *, check_target=True): + if check_target is not True: + raise ValueError("Fused-SNR loss always checks all supervised targets") + if (not torch.is_tensor(prediction) or prediction.ndim != 3 + or not prediction.is_floating_point() or prediction.shape[0] == 0 + or prediction.shape[2] == 0): + raise ValueError("Prediction must be a nonempty BxNxC floating tensor") + self._tensor(target, prediction.shape, prediction.device, name="target") + self._tensor(mask, prediction.shape[:2], prediction.device, + boolean=True, name="mask") + weights, support = self._weights_and_support(noise_levels, metadata, mask) + with torch.autocast(device_type=prediction.device.type, enabled=False): + # Sanitize BOTH operands outside the original mask, BEFORE square. + # Otherwise 0*NaN can poison loss or gradients. Zero SNR weight does + # not remove an originally supervised error from these checks. + safe_prediction = torch.where(mask[..., None], prediction.float(), 0.0) + safe_target = torch.where(mask[..., None], target.float(), 0.0) + if not torch.isfinite(safe_prediction).all() or not torch.isfinite(safe_target).all(): + raise ValueError("Nonfinite originally supervised prediction or target") + error = (safe_prediction - safe_target).square().mean(-1) + if not torch.isfinite(error).all(): + raise ValueError("Nonfinite originally supervised squared error") + objective = self._reduce(error * weights, mask) + plain = self._reduce(error, mask) + if not torch.isfinite(objective) or not torch.isfinite(plain): + raise ValueError("Nonfinite reduced fused objective or unweighted MSE") + metrics = { + "mse": plain.detach(), + "fused_weight_mean": self._reduce(weights * mask, mask).detach(), + "fused_raw_weight_mean": self._reduce(weights * mask, mask).detach() / self.fixed_scale, + "fused_support_mean": self._reduce(support * mask, mask).detach(), + "fused_zero_weight_fraction": self._reduce(((weights == 0) & mask).float(), mask).detach(), + } + return objective, metrics + + def __call__(self, prediction, target, mask, noise_levels=None, + metadata=None, *, check_target=True): + return self.loss_and_metrics(prediction, target, mask, noise_levels, + metadata, check_target=check_target)[0] diff --git a/egomimic/models/world_model.py b/egomimic/models/world_model.py index 3d361c0d7..d950ddf11 100644 --- a/egomimic/models/world_model.py +++ b/egomimic/models/world_model.py @@ -32,10 +32,7 @@ def masked_mse(pred, target, weight, *, reduction="token_mean", check_target=Tru if reduction not in ("token_mean", "mean_per_example"): raise ValueError("Unknown masked MSE reduction") active = weight > 0 - if ( - check_target - and not torch.isfinite(torch.where(active[..., None], target, 0)).all() - ): + if check_target and not torch.isfinite(torch.where(active[..., None], target, 0)).all(): raise ValueError("Non-finite supervised target") safe_target = torch.where(active[..., None], target, torch.zeros_like(target)) error = (pred.float() - safe_target.float()).square().mean(-1) @@ -85,88 +82,62 @@ def __init__( attention_backend="sdpa", attention_dtype="native", image_adaln=False, - expert_widths=None, ): super().__init__() if width % heads or (width // heads) % 2: raise ValueError("RoPE requires an even attention head width") self.heads = heads self.head_dim = width // heads - if ( - type(attention_query_chunk_size) is not int - or attention_query_chunk_size < 0 - ): + if type(attention_query_chunk_size) is not int or attention_query_chunk_size < 0: raise ValueError("attention_query_chunk_size must be a nonnegative integer") self.attention_query_chunk_size = attention_query_chunk_size if attention_backend not in ("sdpa", "efficient", "flash3", "flash4"): - raise ValueError( - "attention_backend must be sdpa, efficient, flash3 or flash4" - ) + raise ValueError("attention_backend must be sdpa, efficient, flash3 or flash4") if attention_dtype not in ("native", "bfloat16"): raise ValueError("attention_dtype must be native or bfloat16") if attention_backend in ("flash3", "flash4") and attention_dtype != "bfloat16": - raise ValueError( - f"{attention_backend} requires explicit bfloat16 attention" - ) + raise ValueError(f"{attention_backend} requires explicit bfloat16 attention") self.attention_backend = attention_backend self.attention_dtype = attention_dtype if type(image_adaln) is not bool: raise ValueError("image_adaln must be bool") self.image_adaln = image_adaln - modalities = tuple(modalities) - self.expert_widths = ( - (width,) * len(modalities) - if expert_widths is None - else tuple(expert_widths) - ) - if len(self.expert_widths) != len(modalities) or any( - type(d) is not int or d < 1 for d in self.expert_widths - ): - raise ValueError("Declare one positive residual width per expert") self.norm1 = nn.ModuleList( - nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) - if image_adaln and i == 0 - else nn.LayerNorm(d) - for i, d in enumerate(self.expert_widths) + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if image_adaln and i == 0 else nn.LayerNorm(width) + for i, _ in enumerate(modalities) ) - # Attention heads share a width; private residual/FFN widths need not. - self.qkv = nn.ModuleList(nn.Linear(d, 3 * width) for d in self.expert_widths) - self.out = nn.ModuleList(nn.Linear(width, d) for d in self.expert_widths) + self.qkv = nn.ModuleList(nn.Linear(width, 3 * width) for _ in modalities) + self.out = nn.ModuleList(nn.Linear(width, width) for _ in modalities) self.norm2 = nn.ModuleList( - nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) - if image_adaln and i == 0 - else nn.LayerNorm(d) - for i, d in enumerate(self.expert_widths) + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if image_adaln and i == 0 else nn.LayerNorm(width) + for i, _ in enumerate(modalities) ) self.ffn = nn.ModuleList( nn.Sequential( - nn.Linear(d, d * mlp_ratio), + nn.Linear(width, width * mlp_ratio), nn.GELU(), - nn.Linear(d * mlp_ratio, d), + nn.Linear(width * mlp_ratio, width), ) - for d in self.expert_widths + for _ in modalities ) if image_adaln: - image_width = self.expert_widths[0] - self.adaln_image = nn.Sequential( - nn.SiLU(), nn.Linear(image_width, 6 * image_width) - ) + self.adaln_image = nn.Sequential(nn.SiLU(), nn.Linear(width, 6 * width)) nn.init.zeros_(self.adaln_image[-1].weight) nn.init.zeros_(self.adaln_image[-1].bias) self.configure_compilation() - def configure_compilation(self, mode="none", dynamic=False, backend="inductor"): + def configure_compilation(self, mode="none", dynamic=False, backend="inductor", + recompile_limit=None, accumulated_recompile_limit=None): names = ("_project_inputs", "_finish_outputs") if self.image_adaln: names += ("_image_modulation",) for name in names: - object.__setattr__( - self, - name + "_executor", - compile_region( - getattr(self, name), mode=mode, dynamic=dynamic, backend=backend - ), - ) + object.__setattr__(self, name + "_executor", compile_region( + getattr(self, name), mode=mode, dynamic=dynamic, backend=backend, + recompile_limit=recompile_limit, accumulated_recompile_limit=accumulated_recompile_limit + )) def _image_modulation(self, condition): # Project once per physical-frame/role group, not once per patch. @@ -178,47 +149,25 @@ def _frame_values(values, frame_indices): *frame_indices.shape, values.shape[-1] ) - def _project_inputs( - self, streams, times, rope, image_modulation=None, frame_indices=None - ): + def _project_inputs(self, streams, times, rope, image_modulation=None, frame_indices=None): qs, ks, vs = [], [], [] for i, x in enumerate(streams): - b, n, _ = x.shape + b, n, d = x.shape normalized = self.norm1[i](x) if i == 0 and self.image_adaln: - shift, scale = ( - self._frame_values(value, frame_indices) - for value in image_modulation[:2] - ) + shift, scale = (self._frame_values(value, frame_indices) for value in image_modulation[:2]) normalized = normalized * (1 + scale) + shift q, k, v = ( self.qkv[i](normalized) - .reshape(b, n, 3, self.heads, self.head_dim) + .reshape(b, n, 3, self.heads, d // self.heads) .permute(2, 0, 3, 1, 4) ) - qs.append( - temporal_rope(q, times[i]) - if rope is None - else apply_temporal_rope(q, rope[i]) - ) - ks.append( - temporal_rope(k, times[i]) - if rope is None - else apply_temporal_rope(k, rope[i]) - ) + qs.append(temporal_rope(q, times[i]) if rope is None else apply_temporal_rope(q, rope[i])) + ks.append(temporal_rope(k, times[i]) if rope is None else apply_temporal_rope(k, rope[i])) vs.append(v) return torch.cat(qs, 2), torch.cat(ks, 2), torch.cat(vs, 2) - def forward( - self, - streams, - times, - valid, - *, - attention_layout=None, - rope=None, - image_condition=None, - ): + def forward(self, streams, times, valid, *, attention_layout=None, rope=None, image_condition=None): sizes = [s.shape[1] for s in streams] modulation = None frame_indices = None @@ -227,16 +176,11 @@ def forward( raise ValueError("Image AdaLN requires frame conditioning") condition, frame_indices = image_condition modulation = self._image_modulation_executor(condition) - query, key, value = self._project_inputs_executor( - streams, times, rope, modulation, frame_indices - ) + query, key, value = self._project_inputs_executor(streams, times, rope, modulation, frame_indices) else: query, key, value = self._project_inputs_executor(streams, times, rope) - mask = ( - torch.cat(valid, 1)[:, None, None, :] - if attention_layout is None - else attention_layout.key_mask - ) + mask = (torch.cat(valid, 1)[:, None, None, :] if attention_layout is None + else attention_layout.key_mask) chunk = self.attention_query_chunk_size if chunk and query.shape[2] > chunk: # Split only queries. Every row retains all keys/values and the same @@ -245,42 +189,29 @@ def forward( # attention intermediates and losing the intended memory bound. parts = [] for q in query.split(chunk, dim=2): - if torch.is_grad_enabled() and any( - t.requires_grad for t in (q, key, value) - ): - part = checkpoint( - self._apply_attention, q, key, value, mask, use_reentrant=False - ) + if torch.is_grad_enabled() and any(t.requires_grad for t in (q, key, value)): + part = checkpoint(self._apply_attention, q, key, value, mask, use_reentrant=False) else: part = self._apply_attention(q, key, value, mask) parts.append(part) result = torch.cat(parts, dim=2) else: # Keep the historical four-argument call for native SDPA overrides. - result = ( - self._apply_attention(query, key, value, mask) - if attention_layout is None - or self.attention_backend not in ("flash3", "flash4") - else self._apply_attention( - query, key, value, mask, attention_layout=attention_layout - ) - ) + result = (self._apply_attention(query, key, value, mask) + if attention_layout is None or self.attention_backend not in ("flash3", "flash4") + else self._apply_attention( + query, key, value, mask, attention_layout=attention_layout)) if self.image_adaln: - return self._finish_outputs_executor( - streams, result, valid, sizes, modulation, frame_indices - ) + return self._finish_outputs_executor(streams, result, valid, sizes, modulation, frame_indices) return self._finish_outputs_executor(streams, result, valid, sizes) - def _finish_outputs( - self, streams, result, valid, sizes, image_modulation=None, frame_indices=None - ): + def _finish_outputs(self, streams, result, valid, sizes, image_modulation=None, frame_indices=None): result = result.transpose(1, 2).flatten(2).split(sizes, 1) outputs = [] for i, (x, y) in enumerate(zip(streams, result)): if i == 0 and self.image_adaln: gate_attention, shift, scale, gate_mlp = ( - self._frame_values(value, frame_indices) - for value in image_modulation[2:] + self._frame_values(value, frame_indices) for value in image_modulation[2:] ) x = x + gate_attention * self.out[i](y) x = x + gate_mlp * self.ffn[i](self.norm2[i](x) * (1 + scale) + shift) @@ -290,9 +221,7 @@ def _finish_outputs( outputs.append(x.masked_fill(~valid[i][..., None], 0)) return outputs - def forward_packed( - self, stream, valid, *, attention_layout, rope, image_condition=None - ): + def forward_packed(self, stream, valid, *, attention_layout, rope, image_condition=None): """Run expert zero on packed image tokens, preserving example boundaries. Pure tensor regions reuse the same modules and compilation policy as @@ -301,9 +230,7 @@ def forward_packed( if self.attention_backend not in ("flash3", "flash4"): raise ValueError("Packed token execution requires flash3 or flash4") if self.attention_dtype != "bfloat16": - raise ValueError( - "Packed token execution requires explicit bfloat16 attention" - ) + raise ValueError("Packed token execution requires explicit bfloat16 attention") if self.attention_query_chunk_size: raise ValueError("Packed token execution requires unchunked attention") if stream.ndim != 3 or stream.shape[0] != 1 or valid.shape != stream.shape[:2]: @@ -315,9 +242,7 @@ def forward_packed( raise ValueError("Image AdaLN requires frame conditioning") condition, frame_indices = image_condition modulation = self._image_modulation_executor(condition) - query, key, value = self._project_inputs_executor( - [stream], None, (rope,), modulation, frame_indices - ) + query, key, value = self._project_inputs_executor([stream], None, (rope,), modulation, frame_indices) else: query, key, value = self._project_inputs_executor([stream], None, (rope,)) original_dtype = query.dtype @@ -348,20 +273,15 @@ def _apply_attention(self, query, key, value, mask, *, attention_layout=None): if self.attention_dtype == "bfloat16": query, key, value = (x.to(torch.bfloat16) for x in (query, key, value)) if self.attention_backend in ("flash3", "flash4"): - kernels = load_flash_attention_kernels( - self.attention_backend, query.device - ) + kernels = load_flash_attention_kernels(self.attention_backend, query.device) layout = attention_layout or prepare_attention_layout( mask[:, 0, 0], query_tokens=query.shape[2] ) result = execute_flash_attention(query, key, value, layout, kernels) elif self.attention_backend == "efficient": if query.device.type != "cuda": - raise RuntimeError( - "efficient attention requires CUDA; no math fallback" - ) + raise RuntimeError("efficient attention requires CUDA; no math fallback") from torch.nn.attention import SDPBackend, sdpa_kernel - with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): result = self._attend(query, key, value, mask) else: @@ -424,9 +344,7 @@ def __init__( if type(strict_fp32) is not bool: raise ValueError("strict_fp32 must be a boolean") if strict_fp32 and attention_dtype != "native": - raise ValueError( - "Explicit BF16 attention conflicts with strict_fp32; select the mixed-attention execution profile" - ) + raise ValueError("Explicit BF16 attention conflicts with strict_fp32; select the mixed-attention execution profile") if image_prediction_type not in ("v", "x0"): raise ValueError("image_prediction_type must be v or x0") self.image_prediction_type = image_prediction_type @@ -471,14 +389,9 @@ def __init__( nn.Linear(2, width), nn.SiLU(), nn.Linear(width, width) ) self.blocks = nn.ModuleList( - MoTBlock( - width, - heads, - mlp_ratio, - attention_query_chunk_size, - attention_backend=attention_backend, - attention_dtype=attention_dtype, - ) + MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size, + attention_backend=attention_backend, + attention_dtype=attention_dtype) for _ in range(depth) ) # One private image adapter per embodiment per block. Zero disables the diff --git a/egomimic/pipeline/stages_image_proprio.py b/egomimic/pipeline/stages_image_proprio.py index a2cc4b37b..3d9fa1e5a 100644 --- a/egomimic/pipeline/stages_image_proprio.py +++ b/egomimic/pipeline/stages_image_proprio.py @@ -73,8 +73,8 @@ def execute(self, batch, *, mode): image = batch["world"]["image"] batch[f"log/{self.log_prefix}_image_tokens"] = image["valid"].sum().float() batch[f"log/{self.log_prefix}_image_query_tokens"] = ( - (image["valid"] & ~image["observed"]).sum().float() - ) + image["valid"] & ~image["observed"] + ).sum().float() if self.model.objective in ("visual_proprio", "visual_state_action"): batch["log/pose_weight"] = losses["total"].new_tensor(scale) if self.model.objective == "visual_state_action": diff --git a/egomimic/pipeline/stages_joint_state_action.py b/egomimic/pipeline/stages_joint_state_action.py index 4dabff0dd..7d789b99c 100644 --- a/egomimic/pipeline/stages_joint_state_action.py +++ b/egomimic/pipeline/stages_joint_state_action.py @@ -33,9 +33,16 @@ class JointDenoiserStage(Stage): "inference": ("joint/model",), } - def __init__(self, model): + def __init__(self, model, language_key=None): super().__init__() self.model = model + if language_key is not None and (not isinstance(language_key, str) or not language_key): + raise ValueError("language_key must be a nonempty batch key or None") + self.language_key = language_key + if language_key is not None: + self.reads_by_mode = { + mode: (*reads, language_key) for mode, reads in self.reads_by_mode.items() + } def execute(self, batch, *, mode): batch["joint/model"] = self.model @@ -45,6 +52,7 @@ def execute(self, batch, *, mode): _embodiment(batch), batch["joint/values"], batch["joint/levels"], + **({"language_features": batch[self.language_key]} if self.language_key is not None else {}), ) return batch @@ -59,10 +67,10 @@ class JointLossStage(Stage): "joint/metadata", ) - def __init__(self, stream, loss, weight=1.0): + def __init__(self, stream, loss, weight=1.0, optional=False, log_objective=False): super().__init__() if ( - stream not in ("image", "proprio", "action") + stream not in ("image", "proprio", "action", "camera") or not math.isfinite(weight) or weight < 0 ): @@ -70,19 +78,44 @@ def __init__(self, stream, loss, weight=1.0): "Declare a joint stream and finite nonnegative loss weight" ) self.stream, self.loss, self.weight = stream, loss, weight + self.optional = optional + if type(log_objective) is not bool: + raise ValueError("log_objective must be an explicit boolean") + if log_objective and not callable(getattr(loss, "loss_and_metrics", None)): + raise ValueError("Objective logging requires loss_and_metrics") + if getattr(loss, "requires_objective_metrics", False) and not log_objective: + raise ValueError("Weighted objective requires separate unweighted metrics") + self.log_objective = log_objective self.writes = ("loss/" + stream,) def execute(self, batch, *, mode): k = self.stream - loss = self.loss( + if k not in batch["joint/prediction"]: + if not self.optional: + raise ValueError(f"Required stream {k} absent") + # Graph bookkeeping only: no fabricated human action tokens/targets. + loss = batch["joint/prediction"]["image"].sum() * 0 + batch[self.writes[0]] = loss + return batch + evaluate = self.loss.loss_and_metrics if self.log_objective else self.loss + result = evaluate( batch["joint/prediction"][k], batch["joint/targets"][k], batch["joint/masks"][k], noise_levels=batch["joint/levels"][k], metadata=batch["joint/metadata"][k], ) + if self.log_objective: + loss, metrics = result + if "mse" not in metrics: + raise ValueError("Objective logging requires an unweighted MSE metric") + batch["log/joint_" + k + "_objective"] = loss.detach() + for name, value in metrics.items(): + batch["log/joint_" + k + "_" + name] = value.detach() + else: + loss = result + batch["log/joint_" + k + "_mse"] = loss.detach() batch[self.writes[0]] = loss * self.weight - batch["log/joint_" + k + "_mse"] = loss.detach() return batch @@ -91,12 +124,18 @@ class JointSamplerStage(Stage): reads = ("world", "embodiment", "joint/model") writes = ("image_proprio_prediction",) - def __init__(self, sampler): + def __init__(self, sampler, language_key=None): super().__init__() self.sampler = sampler + if language_key is not None and (not isinstance(language_key, str) or not language_key): + raise ValueError("language_key must be a nonempty batch key or None") + self.language_key = language_key + if language_key is not None: + self.reads = (*self.reads, language_key) def execute(self, batch, *, mode): batch[self.writes[0]] = self.sampler( - batch["joint/model"], batch["world"], _embodiment(batch) + batch["joint/model"], batch["world"], _embodiment(batch), + **({"language_features": batch[self.language_key]} if self.language_key is not None else {}), ) return batch diff --git a/egomimic/pl_utils/joint_smoke.py b/egomimic/pl_utils/joint_smoke.py index 2e3c0ec5e..b4f650603 100644 --- a/egomimic/pl_utils/joint_smoke.py +++ b/egomimic/pl_utils/joint_smoke.py @@ -10,6 +10,8 @@ import torch from lightning.pytorch.callbacks import Callback from omegaconf import OmegaConf +from scripts.train.joint.lr_policy import from_config, update_lr, verify_scheduler +from egomimic.rldb.boundary_sampling import validate_boundary_record, validate_boundary_sampling def require(condition, message): @@ -31,9 +33,29 @@ def weights_digest(state): return digest.hexdigest() +def state_control_parameter_state(state): + """Project only the optional named action bank out of a native wrapper. + + Retaining full wrapper keys lets the actual three-stream initialization be + compared with an immutable image/proprio smoke without constructing it again. + """ + result = {key: value for key, value in state.items() if "action" not in key.split(".")} + require(result and len(result) < len(state), "Missing optional action bank") + return result + + def image_parameter_state(network): - """Select the common image expert, whose bank index is always zero.""" - return { + """Select the model-owned image state or the existing joint bank zero. + + An unsupported namespace must fail, never certify an empty image digest. + """ + provider = getattr(network, "image_parameter_state", None) + if callable(provider): + state = provider() + require(bool(state) and all(torch.is_tensor(v) for v in state.values()), + "Model-owned image state is empty or invalid") + return state + state = { key: value for key, value in network.state_dict().items() if key.startswith( @@ -54,6 +76,8 @@ def image_parameter_state(network): ) ) } + require(bool(state), "Unsupported image parameter namespace; empty digest rejected") + return state def paired_image_digest(world, values, levels): @@ -104,12 +128,144 @@ def load_smoke_configs(output): return cfg, model_config +def check_conditional_adam_state(optimizer_state, usage_record, expected_step): + """Verify actual bank-wise Adam counts using the preserved parameter order.""" + require(usage_record.get('schema')=='native-conditional-optimizer-usage-v1' + and usage_record.get('global_step')==expected_step, + 'Missing exact conditional optimizer usage checkpoint') + names=usage_record['parameter_names'];counts=usage_record['parameter_update_counts'] + ids=[key for group in optimizer_state['param_groups'] for key in group['params']] + require(len(ids)==len(names)==len(set(names)) and set(counts)==set(names), + 'Conditional optimizer parameter order/identity differs') + for name,key in zip(names,ids): + count=counts[name];state=optimizer_state['state'].get(key) + require(type(count) is int and 0<=count<=expected_step, + 'Invalid conditional optimizer count') + require((state is None)==(count==0),'Unexpected Adam state for unused parameter') + if state is not None: + require(float(state['step'])==count,'Conditional Adam update count differs') + require(all(state[k].dtype==torch.float32 and torch.isfinite(state[k]).all() + for k in ('exp_avg','exp_avg_sq')),'Invalid conditional Adam moments') + + +class ConditionalOptimizerUsage(Callback): + """Opt-in production checkpoint accounting for genuinely absent branches.""" + def __init__(self): + self.counts=None + self.restored=None + + def on_load_checkpoint(self, trainer, pl_module, checkpoint): + record=checkpoint.get('native_conditional_optimizer_usage') + require(record is not None,'Conditional training resume requires exact optimizer usage') + require(len(checkpoint.get('optimizer_states',[]))==1,'Expected one Adam optimizer') + check_conditional_adam_state(checkpoint['optimizer_states'][0],record,checkpoint['global_step']) + self.restored=record + + def on_fit_start(self, trainer, pl_module): + names=[name for name,_ in pl_module.named_parameters()] + if self.restored is None: + require(trainer.global_step==0,'Missing restored conditional optimizer usage') + self.counts={name:0 for name in names} + else: + require(names==self.restored['parameter_names'],'Conditional parameter order changed on resume') + self.counts=dict(self.restored['parameter_update_counts']) + + def on_before_optimizer_step(self, trainer, pl_module, optimizer): + require(self.counts is not None and len(trainer.optimizers)==1,'Conditional usage callback not initialized') + for name,param in pl_module.named_parameters(): + self.counts[name]+=int(param.grad is not None) + + def on_save_checkpoint(self, trainer, pl_module, checkpoint): + names=[name for name,_ in pl_module.named_parameters()] + record=dict(schema='native-conditional-optimizer-usage-v1',global_step=int(trainer.global_step), + parameter_names=names,parameter_update_counts=dict(self.counts)) + require(len(checkpoint.get('optimizer_states',[]))==1,'Expected full Adam state on save') + check_conditional_adam_state(checkpoint['optimizer_states'][0],record,trainer.global_step) + checkpoint['native_conditional_optimizer_usage']=record + + +def _check_multimodal_sample(sample, *, boundary_policy=None): + boundary_policy=validate_boundary_sampling(boundary_policy) + features=sample.get('language_features') + require(torch.is_tensor(features) and features.ndim==2 and features.shape[1]==1024 + and torch.isfinite(features).all(),'Missing finite actual Qwen language condition') + metadata=sample.get('pixel_metadata',[]) + require(len(metadata)==features.shape[0],'One metadata record required per sample') + for i,record in enumerate(metadata): + unit=record.get('subtask',{}) + require(unit.get('interval_semantics')=='half_open' and unit.get('annotation_provenance'), + 'Missing reviewed subtask provenance') + frames=record.get('frame_indices',[]) + chunks=record.get('proprio_frame_indices',[]) + indices=frames+[v for chunk in chunks for v in chunk] + for key in ('action_frame_indices','camera_frame_indices'): + if key in record: + for item in record[key]:indices.extend(item if isinstance(item,list) else [item]) + if boundary_policy is None: + require('boundary_sampling' not in record,'Legacy/held-out smoke cannot admit boundary uncertainty metadata') + else: + require(record.get('boundary_sampling',{}).get('policy')==boundary_policy, + 'Boundary sample differs from explicit expected training policy') + if 'boundary_sampling' in record: + packet=validate_boundary_record(record,unit,indices) + require(type(sample.get('sampling_cursor')) is int + and packet['rng_identity']['cursor']==sample['sampling_cursor']-1, + 'Boundary RNG cursor differs from checkpointed sample draw') + require(record['chunk_size']==1 and chunks==[[v] for v in frames], + 'Boundary image/state selection must share native timestamps') + for name in ('action','camera'): + if name in sample['world']: + require(record.get(name+'_frame_indices')==[[v] for v in frames], + 'Boundary action/camera selection must share native timestamps') + physical=torch.tensor(record['timestamps_seconds'],dtype=torch.float64) + relative=(physical-float(record['time_origin_seconds'])).float() + repeats={'image':len(record['camera_names'])*(record['image_size'][0]//record['patch_size'])*(record['image_size'][1]//record['patch_size']), + 'proprio':1,'action':1,'camera':len(record['camera_names'])} + for name,stream in sample['world'].items(): + expected=relative.repeat_interleave(repeats[name]) + actual=stream['time'][i,:len(expected)].detach().cpu() + require(actual.shape==expected.shape and torch.equal(actual,expected), + 'Boundary stream timestamp differs from shared recorded selection') + else: + require(indices and all(unit['start_idx']<=v0);counts['start_outward']+=int(start<0) + counts['end_inward']+=int(end<0);counts['end_outward']+=int(end>0) + return counts + + class JointSmokeProof(Callback): - def __init__(self, output, world_size, expected_steps=50, warmup_steps=10): + def __init__(self, output, world_size, expected_steps=50, warmup_steps=10, multimodal=False): require(world_size in (2, 4), "Joint pilots use two/four dedicated GPUs") require(0 < warmup_steps < expected_steps, "Need warmed optimization timing") self.output, self.world_size = Path(output), world_size self.expected_steps, self.warmup_steps = expected_steps, warmup_steps + self.multimodal = bool(multimodal) + self.boundary_policy = None def on_fit_start(self, trainer, pl_module): require( @@ -136,21 +292,32 @@ def on_fit_start(self, trainer, pl_module): ), "All active experts must be trainable FP32 stored weights", ) - require( - all( - math.isclose(group["lr"], 5e-6 / 3000, rel_tol=1e-12) - for optimizer in trainer.optimizers - for group in optimizer.param_groups - ), - "Expected first update LR5e-6/3000", - ) - require( - pl_module.hparams.config_tree.model.scheduler.warmup_steps == 3000 - and len(trainer.lr_scheduler_configs) == 1 - and trainer.lr_scheduler_configs[0].interval == "step" - and trainer.lr_scheduler_configs[0].frequency == 1, - "Both pilots require a3000 optimizer-update warmup", - ) + if self.multimodal: + _check_pilot_interfaces(self.network,set(self.network.pose_dims),set(self.network.action_dims)) + require('camera' in self.network.streams and len(self.network.pose_dims)==2, + 'Multimodal smoke requires robot/human plus camera') + require(torch.distributed.get_backend()=='nccl', 'Multimodal production smoke requires NCCL') + prop=torch.cuda.get_device_properties(pl_module.device) + require('H100' in prop.name and prop.major>=9,'Production smoke requires actual H100') + image=self.network.image_model + require(image.token_execution=='packed_valid' and image.blocks[0].attention_backend=='flash4', + 'Production smoke requires Flash4 packed_valid') + self.parameter_usage={name:0 for name,_ in pl_module.named_parameters()} + self.embodiment_batches={emb:0 for emb in self.network.pose_dims} + self.current_embodiment=None + self.boundary_policy=validate_boundary_sampling(trainer.datamodule.settings.get('boundary_sampling')) + if self.boundary_policy is not None: + self.boundary_counts={emb:_boundary_sample_counts({'pixel_metadata':[]}) for emb in self.network.pose_dims} + require(len(trainer.datamodule.val_dataloader())>=100//self.world_size, + 'Held-out loader cannot provide100 global validation batches') + require(trainer.limit_val_batches==100//self.world_size, + 'Distributed smoke must cap exactly100 global held-out batches') + self.lr_policy = from_config(pl_module.hparams.config_tree.model) + require(all(math.isclose(group["lr"], update_lr(self.lr_policy, 0), rel_tol=1e-12) + for optimizer in trainer.optimizers for group in optimizer.param_groups), + "Wrong first optimizer update LR") + verify_scheduler(self.lr_policy, + [c.scheduler.state_dict() for c in trainer.lr_scheduler_configs], 0) self.optimizer_lrs = [] digest = weights_digest(pl_module.state_dict()) initial = [None] * self.world_size @@ -158,16 +325,21 @@ def on_fit_start(self, trainer, pl_module): require(len(set(initial)) == 1, "Fresh DDP initial weights differ") self.initial_weights = initial self.initial_image_weights = weights_digest(image_parameter_state(self.network)) + self.initial_state_control_weights = ( + weights_digest(state_control_parameter_state(pl_module.state_dict())) + if "action" in self.network.streams else None + ) self.paired_rgb_batches = [] def observe(module, args, outputs): require( - set(outputs) == set(self.network.streams) + set(outputs) == set(self.network.active_streams(args[1]) if self.multimodal else self.network.streams) and all(value.dtype == torch.bfloat16 for value in outputs.values()), "Actual active-stream predictions must execute in BF16", ) if pl_module.training and torch.is_grad_enabled(): self.bf16_outputs += 1 + if self.multimodal:self.current_embodiment=args[1] if len(self.paired_rgb_batches) < 3: self.paired_rgb_batches.append( paired_image_digest(args[0], args[2], args[3]) @@ -175,6 +347,21 @@ def observe(module, args, outputs): self.precision_hook = self.network.register_forward_hook(observe) + def on_train_batch_start(self, trainer, pl_module, batch, batch_idx): + if self.multimodal: + require(len(batch)==1,'Native training batches must be homogeneous by embodiment') + for emb,sample in batch.items(): + _check_multimodal_sample(sample,boundary_policy=self.boundary_policy) + if self.boundary_policy is not None: + for record in sample['pixel_metadata']: + require(record.get('boundary_sampling',{}).get('policy')==self.boundary_policy, + 'Actual boundary policy differs from context-bound opt-in') + counts=_boundary_sample_counts(sample) + for key,value in counts.items():self.boundary_counts[emb][key]+=value + else: + require(all('boundary_sampling' not in record for record in sample['pixel_metadata']), + 'Legacy smoke cannot silently admit boundary uncertainty') + def on_before_backward(self, trainer, pl_module, loss): require( loss.dtype == torch.float32 and bool(torch.isfinite(loss)), @@ -183,7 +370,7 @@ def on_before_backward(self, trainer, pl_module, loss): self.fp32_losses += 1 def on_before_optimizer_step(self, trainer, pl_module, optimizer): - expected_lr = 5e-6 * min((trainer.global_step + 1) / 3000, 1.0) + expected_lr = update_lr(self.lr_policy, trainer.global_step) require( all( math.isclose(g["lr"], expected_lr, rel_tol=1e-12) @@ -192,6 +379,20 @@ def on_before_optimizer_step(self, trainer, pl_module, optimizer): "Actual optimizer LR does not follow the3000-update warmup", ) self.optimizer_lrs.append(expected_lr) + if self.multimodal: + emb=self.current_embodiment + require(emb in self.embodiment_batches,'No actual training embodiment recorded') + self.embodiment_batches[emb]+=1 + active=[p for p in self.network.parameters() if p.grad is not None] + require(active and all(p.grad.dtype==torch.float32 and torch.isfinite(p.grad).all() for p in active), + 'Active multimodal gradients must be finite FP32') + for other,bank in self.network.proprio_banks.items(): + if other!=emb:require(all(p.grad is None for p in bank.parameters()),'Inactive private proprio gradient') + if emb not in self.network.action_dims: + require(all(p.grad is None for p in self.network.action.parameters()),'Human action bank must be unused') + for name,param in pl_module.named_parameters():self.parameter_usage[name]+=int(param.grad is not None) + self.gradients_checked=True + return if self.gradients_checked: return parameters = list(self.network.parameters()) @@ -225,6 +426,37 @@ def on_validation_start(self, trainer, pl_module): def on_validation_batch_end( self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0 ): + if self.multimodal: + # Exact contiguous validation shards may put an embodiment entirely + # on nonzero ranks. Elect one actual owner per current embodiment; + # every rank participates even after its sample is already saved. + rank_embodiments = [None] * self.world_size + torch.distributed.all_gather_object(rank_embodiments, sorted(batch)) + writers = {} + for rank, names in enumerate(rank_embodiments): + for name in names: + writers.setdefault(name, rank) + for emb,sample in batch.items(): + if writers[emb] != trainer.global_rank:continue + _check_multimodal_sample(sample) + path=self.output/f'sample-input-{emb}.pt' + if path.exists():continue + def cpu(value): + if torch.is_tensor(value):return value.detach().cpu() + if isinstance(value,dict):return {k:cpu(v) for k,v in value.items()} + if isinstance(value,list):return [cpu(v) for v in value] + if isinstance(value,tuple):return tuple(cpu(v) for v in value) + return value + sample=cpu(sample) + require('language_features' in sample and 'pixel_metadata' in sample, + 'Need real subtask/language camera sample provenance') + for stream in sample['world'].values(): + stream['value']=torch.where(stream['observed'][...,None],stream['value'],float('nan')) + torch.save(sample,path) + # Publish completed files before any rank checks both embodiments + # at train end. This never moves a held-out example between ranks. + torch.distributed.barrier() + return if trainer.global_rank != 0 or batch_idx != 0: return sample = batch["libero_gripper"] @@ -269,10 +501,18 @@ def on_train_end(self, trainer, pl_module): "image": "image_normalized_mse", "proprio": "proprio_normalized_mse", "action": "action_command_mse", - }[name] + }.get(name,name+"_normalized_mse") for name in self.network.streams ) } + if self.multimodal: + expected={f'val/{emb}/{kind}/{"action_command_mse" if stream=="action" else stream+"_normalized_mse"}' + for emb in self.network.pose_dims for kind in ('forecast','completion') + for stream in self.network.active_streams(emb)} + require(all(count>0 for count in self.embodiment_batches.values()),'Both embodiments must optimize') + _check_pilot_adam_usage(trainer.optimizers[0],pl_module.named_parameters(),self.parameter_usage) + require(all((self.output/f'sample-input-{emb}.pt').is_file() for emb in self.network.pose_dims), + 'Both embodiments need actual validation samples') require( set(metrics) == expected and all(math.isfinite(v) for v in metrics.values()), @@ -313,11 +553,13 @@ def on_train_end(self, trainer, pl_module): initial_weights=self.initial_weights, active_streams=list(self.network.streams), initial_image_weights_sha256=self.initial_image_weights, + initial_state_control_weights_sha256=self.initial_state_control_weights, paired_rgb_batch_sha256=self.paired_rgb_batches, bf16_training_outputs=self.bf16_outputs, fp32_losses=self.fp32_losses, fp32_gradients_verified=self.gradients_checked, - lr_warmup_steps=3000, + lr_warmup_steps=self.lr_policy["warmup_optimizer_steps"], + lr_policy=self.lr_policy, optimizer_update_lrs=self.optimizer_lrs, metrics=metrics, gpu_uuid=str(torch.cuda.get_device_properties(pl_module.device).uuid), @@ -329,4 +571,377 @@ def on_train_end(self, trainer, pl_module): data_context_sha256=pl_module.data_context.fingerprint(), checkpoint_name=checkpoint_name, ) + if self.multimodal: + record.update(schema='joint-multimodal-fit-v1',multimodal=True, + parameter_update_counts=self.parameter_usage,embodiment_optimizer_batches=self.embodiment_batches, + active_streams_by_embodiment={e:list(self.network.active_streams(e)) for e in self.network.pose_dims}, + private_bank_optimizer_steps_verified=True,human_action_absent=True, + subtask_boundaries_verified=True,language_conditioning_verified=True,camera_targets_verified=True, + gpu_name=torch.cuda.get_device_properties(pl_module.device).name, + gpu_capability=list(torch.cuda.get_device_capability(pl_module.device)), + attention_backend='flash4',token_execution='packed_valid') + if self.boundary_policy is not None: + all_counts=[None]*self.world_size + torch.distributed.all_gather_object(all_counts,self.boundary_counts) + totals={emb:{key:sum(rank_counts[emb][key] for rank_counts in all_counts) + for key in self.boundary_counts[emb]} for emb in self.boundary_counts} + require(all(counts['samples']>0 for counts in totals.values()) + and sum(counts['outward_samples'] for counts in totals.values())>0, + 'Changed sampling smoke must exercise both embodiments and actual bounded outward frames') + # Explicitly false: this arm does not claim every target is within + # the original reviewed interval. New admission must read the policy. + record.update(subtask_boundaries_verified=False, + original_subtask_annotation_identity_verified=True, + subtask_boundary_verification_policy='annotation_boundary_uncertainty', + boundary_uncertainty=dict(schema='boundary-uncertainty-smoke-realizations-v1', + policy=self.boundary_policy,all_rank_counts=all_counts,global_counts_by_embodiment=totals, + synchronized_native_streams_verified=True,original_instruction_retained=True, + neighbor_instruction_equivalence_asserted=False,validation_sampling_unchanged=True)) write_once(self.output / f"fit-rank{trainer.global_rank}.json", record) + + +def _check_pilot_interfaces(net, expected, action_embodiments): + require(set(net.pose_dims) == set(expected) and set(net.camera_dims) == set(expected), + 'Extra/missing embodiment banks in resolved model') + require(set(net.action_dims) == set(action_embodiments), + 'Extra/missing action banks in resolved model') + require(net.separate_proprio_experts, 'Full private proprio banks required') + + +def _check_pilot_adam_usage(optimizer, named_parameters, usage): + for name,p in named_parameters: + item=optimizer.state.get(p) + require((item is None)==(usage[name]==0),'Adam state exists for unused parameter') + if item is not None: + require(int(item['step'])==usage[name], 'Private-bank Adam update count mismatch') + require(all(item[k].dtype==torch.float32 and torch.isfinite(item[k]).all() for k in ('exp_avg','exp_avg_sq')), + 'Adam moments must be finite FP32') + + +def validate_multimodal_pilot(config_path, output, *, optimizer_updates=4, postflight_only=False): + # Multiple embodiment/dtype/autograd decode specializations are expected. + # Keep compilation mandatory and bounded; never fall back to eager execution. + with torch._dynamo.config.patch(recompile_limit=64, accumulated_recompile_limit=256, + suppress_errors=False): + return _validate_multimodal_pilot(config_path, output, + optimizer_updates=optimizer_updates, postflight_only=postflight_only) + + +def _validate_multimodal_pilot(config_path, output, *, optimizer_updates=4, postflight_only=False): + """Explicit one-GPU code validation on the reviewed TRAIN pilot. + + Reuses the declared pipeline, optimizer, sampler and prepared-data loader. + This is not a production training launcher, held-out score, maximum-shape + capacity claim, or the historical 50-update/100-anchor DDP smoke contract. + """ + import inspect + import random + import numpy as np + from hydra.utils import instantiate + from egomimic.rldb.pixel_proprio import PreparedPixelProprioDataset, collate_pixel_proprio + + require(optimizer_updates == 4, 'Pilot contract is exactly four diagnostic optimizer updates') + require(torch.cuda.is_available() and torch.cuda.device_count() == 1, + 'Pilot validation requires exactly one visible allocated GPU') + root = Path(output) + root.mkdir(parents=True, exist_ok=True) + require(not (root/'pilot-validation.json').exists(), 'Completed pilot validation is immutable') + checkpoint=root/'validation-epoch0000-step000000004.ckpt' + if postflight_only: + require(checkpoint.is_file(), 'Postflight requires the existing optimizer checkpoint') + else: + require(not checkpoint.exists(), + 'Completed optimizer checkpoint exists; resume postflight, never repeat optimizer updates') + config_bytes=Path(config_path).read_bytes() + config_sha256=hashlib.sha256(config_bytes).hexdigest() + cfg = OmegaConf.load(config_path) + data_cfg = OmegaConf.to_container(cfg.data, resolve=True) + accepted = set(inspect.signature(PreparedPixelProprioDataset).parameters) + settings = {k:v for k,v in data_cfg.items() if k in accepted} + settings.update(split='train') + data = PreparedPixelProprioDataset(**settings) + batches,offset = {},0 + for (episode,unit),count in zip(data.sample_units,data.counts): + emb = episode['embodiment'] + if emb not in batches: + sample=data[(offset,0,0)] + require(not any(s['observed'].any() for s in sample['world'].values()), + 'Pilot must use the selected independent-noise training contract') + batches[emb] = collate_pixel_proprio([sample]) + offset += count + expected = {e['embodiment'] for e in data.manifest['episodes']} + require(set(batches) == expected and len(expected) == 2, + 'Both native robot and human temporal units are required') + torch.manual_seed(int(cfg.seed));np.random.seed(int(cfg.seed));random.seed(int(cfg.seed)) + + def construct(): + algo=instantiate(cfg.model.pipeline, device='cuda') + net=algo.pipeline.stage_by_id('joint_denoiser').model + _check_pilot_interfaces(net, expected, cfg.data.action_embodiments) + require(net.image_model.compile_mode == 'default' and net.image_model.compile_dynamic, + 'Pilot must exercise native dynamic Inductor compilation') + require(all(p.dtype == torch.float32 and p.requires_grad for p in algo.nets.parameters()), + 'Pilot stored weights must be trainable FP32') + return algo,net + + algo,net=construct() + optimizer=instantiate(cfg.model.optimizer, params=algo.nets.parameters()) + if callable(optimizer):optimizer=optimizer() + scheduler=instantiate(cfg.model.scheduler, optimizer=optimizer) if cfg.model.scheduler else None + policy=from_config(cfg.model) + order=sorted(expected) + usage={name:0 for name,_ in algo.nets.named_parameters()} + rows=[] + torch.cuda.reset_peak_memory_stats() + started=time.perf_counter() + if not postflight_only: + for step in range(optimizer_updates): + emb=order[step % len(order)] + algo.nets.train();optimizer.zero_grad(set_to_none=True) + batch=algo.process_batch_for_training(batches[emb]) + with torch.autocast('cuda',dtype=torch.bfloat16): + predicted=algo.forward_training(batch) + losses=algo.compute_losses(predicted,batch) + loss=losses['loss'] + require(loss.dtype==torch.float32 and torch.isfinite(loss),'Diagnostic loss must be finite FP32') + # The model-owned denoiser stage records the actual active predictions. + sample=predicted[emb] if emb in predicted else predicted + outputs=sample['joint/prediction'] + require(set(outputs)==set(net.active_streams(emb)) and + all(v.dtype==torch.bfloat16 for v in outputs.values()), + 'Active expert outputs did not execute in BF16') + loss.backward() + active=[p for p in algo.nets.parameters() if p.grad is not None] + require(active and all(p.grad.dtype==torch.float32 and torch.isfinite(p.grad).all() for p in active), + 'Nonfinite/non-FP32 active gradients') + for other,bank in net.proprio_banks.items(): + if other != emb:require(all(p.grad is None for p in bank.parameters()),'Inactive private bank got gradients') + if emb not in net.action_dims: + require('action' not in outputs and all(p.grad is None for p in net.action.parameters()), + 'Human action path must be absent and unused') + for name,p in algo.nets.named_parameters():usage[name]+=int(p.grad is not None) + actual_lr=optimizer.param_groups[0]['lr'] + require(all(math.isclose(g['lr'],update_lr(policy,step),rel_tol=1e-12) for g in optimizer.param_groups), + 'Diagnostic optimizer uses a different LR policy') + optimizer.step() + if scheduler is not None:scheduler.step() + rows.append(dict(step=step+1,embodiment=emb,loss=float(loss.detach()),lr=actual_lr, + active_streams=list(outputs),active_gradient_parameters=sum(p.numel() for p in active))) + state=optimizer.state_dict() + _check_pilot_adam_usage(optimizer, algo.nets.named_parameters(), usage) + checkpoint=root/'validation-epoch0000-step000000004.ckpt' + require(not checkpoint.exists(),'Diagnostic checkpoint may not be overwritten') + rng=dict(torch_cpu=torch.get_rng_state(),torch_cuda=torch.cuda.get_rng_state(), + numpy=np.random.get_state(),python=random.getstate()) + torch.save(dict(schema='joint-multimodal-pilot-validation-v1',global_step=4,epoch=0, + state_dict=algo.nets.state_dict(),optimizer_state=state,scheduler_state=None if scheduler is None else scheduler.state_dict(), + rng=rng,parameter_update_counts=usage,optimizer_rows=rows,config_raw_sha256=config_sha256,config=OmegaConf.to_container(cfg,resolve=True)),checkpoint) + loaded=torch.load(checkpoint,map_location='cpu',weights_only=False,mmap=True) + require(loaded['schema']=='joint-multimodal-pilot-validation-v1' and + loaded['global_step']==4 and loaded['epoch']==0, + 'Postflight checkpoint is not the completed four-update diagnostic') + require(loaded['config']==OmegaConf.to_container(cfg,resolve=True), + 'Postflight resolved config differs from checkpoint') + if 'config_raw_sha256' in loaded: + require(loaded['config_raw_sha256']==config_sha256, 'Postflight config bytes differ') + if postflight_only: + algo.nets.load_state_dict(loaded['state_dict'],strict=True) + optimizer.load_state_dict(loaded['optimizer_state']) + if scheduler is not None:scheduler.load_state_dict(loaded['scheduler_state']) + usage=loaded['parameter_update_counts'] + require(set(usage)==set(dict(algo.nets.named_parameters())), 'Checkpoint parameter usage identity differs') + require(all(count in (0,2,4) for count in usage.values()) and 4 in usage.values(), + 'Checkpoint optimizer update history is invalid') + _check_pilot_adam_usage(optimizer,algo.nets.named_parameters(),usage) + rows=loaded.get('optimizer_rows',[]) + restored,restored_net=construct() + restored.nets.load_state_dict(loaded['state_dict'],strict=True) + require(weights_digest(algo.nets.state_dict())==weights_digest(restored.nets.state_dict()),'Strict model reload changed weights') + ro=instantiate(cfg.model.optimizer,params=restored.nets.parameters()) + if callable(ro):ro=ro() + rs=instantiate(cfg.model.scheduler,optimizer=ro) if cfg.model.scheduler else None + ro.load_state_dict(loaded['optimizer_state']) + if rs is not None:rs.load_state_dict(loaded['scheduler_state']) + require((None if rs is None else rs.state_dict())==loaded['scheduler_state'],'Scheduler reload differs') + for left,right in zip(optimizer.state.values(),ro.state.values()): + require(set(left)==set(right),'Optimizer reload keys differ') + for key,value in left.items(): + require(torch.equal(value.cpu(),right[key].cpu()) if torch.is_tensor(value) else value==right[key], + 'Optimizer reload value differs') + torch.set_rng_state(loaded['rng']['torch_cpu']);torch.cuda.set_rng_state(loaded['rng']['torch_cuda']) + np.random.set_state(loaded['rng']['numpy']);random.setstate(loaded['rng']['python']) + require(torch.equal(torch.get_rng_state(),loaded['rng']['torch_cpu']) and + torch.equal(torch.cuda.get_rng_state(),loaded['rng']['torch_cuda']) and + random.getstate()==loaded['rng']['python'] and + np.array_equal(np.random.get_state()[1],loaded['rng']['numpy'][1]),'RNG restoration differs') + validation=[] + for emb in order: + algo.nets.eval();restored.nets.eval() + batch=algo.process_batch_for_training(batches[emb]) + restored_batch=restored.process_batch_for_training(batches[emb]) + with torch.no_grad(),torch.autocast('cuda',dtype=torch.bfloat16): + torch.manual_seed(1701);original=algo.forward_training(batch) + torch.manual_seed(1701);replayed=restored.forward_training(restored_batch) + a=original[emb] if emb in original else original + b=replayed[emb] if emb in replayed else replayed + for key,value in a['joint/prediction'].items(): + require(torch.equal(value,b['joint/prediction'][key]),'Reload prediction changed under identical noise RNG') + metrics=restored.compute_losses(replayed,restored_batch) + require(all(torch.isfinite(v) for v in metrics.values() if torch.is_tensor(v)),'Nonfinite diagnostic validation') + with torch.no_grad(),torch.autocast('cuda',dtype=torch.bfloat16): + query=restored.process_batch_for_training(batches[emb]) + for stream in query[emb]['world'].values(): + stream['value']=torch.where(stream['observed'][...,None],stream['value'],float('nan')) + torch.manual_seed(1701) + sampled=restored.forward_eval(query)[emb]['image_proprio_prediction'] + require(set(sampled)==set(restored_net.active_streams(emb)) and + all(v.dtype==torch.float32 and torch.isfinite(v).all() for v in sampled.values()), + 'Strictly restored native sampler produced invalid active streams') + sampler=restored.pipeline.stage_by_id('joint_sampler').sampler + require(sampler.state_steps==20 and (emb not in net.action_dims or sampler.action_steps==20), + 'Configured native sampler schedule changed') + validation.append(dict(embodiment=emb,loss=float(metrics['loss']),split='TRAIN_pilot_diagnostic', + native_sampler_steps=20,sampled_streams=list(sampled))) + torch.cuda.synchronize() + with checkpoint.open('rb') as checkpoint_stream: + checkpoint_sha=hashlib.file_digest(checkpoint_stream,'sha256').hexdigest() + result=dict(schema='joint-multimodal-pilot-validation-v1',status='PASS',optimizer_updates=4, + production_training=False,held_out_benchmark=False,full_shape_capacity=False, + config_raw_sha256=config_sha256,config_path=str(config_path), + validation_source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + rows=rows,optimizer_telemetry_available=bool(rows),postflight_only=postflight_only, + optimizer_updates_this_invocation=0 if postflight_only else 4, + dynamo_recompile_limit=64,dynamo_accumulated_recompile_limit=256, + validation=validation,strict_model_optimizer_scheduler_rng_reload=True, + parameter_update_counts=usage,parameters_total=sum(p.numel() for p in algo.nets.parameters()), + parameters_trainable=sum(p.numel() for p in algo.nets.parameters() if p.requires_grad), + lr_policy=policy,compile_mode=net.image_model.compile_mode,compile_dynamic=net.image_model.compile_dynamic, + attention_backend=net.image_model.blocks[0].attention_backend,token_execution=net.image_model.token_execution, + elapsed_seconds=time.perf_counter()-started,peak_gpu_bytes=torch.cuda.max_memory_allocated(), + checkpoint=dict(path=str(checkpoint),bytes=checkpoint.stat().st_size, + sha256=checkpoint_sha),torch_version=torch.__version__) + write_once(root/'pilot-validation.json',result) + return result + + +def _capacity_manifest_sha256(cfg, expected): + require(isinstance(expected,str) and len(expected)==64 + and all(c in '0123456789abcdef' for c in expected), + 'Capacity requires exact request data manifest SHA256') + actual=hashlib.sha256(Path(cfg.data.manifest_path).read_bytes()).hexdigest() + require(actual==expected,'Capacity data manifest differs from bound request') + return actual + + +def _optimizer_capacity_reserve_bytes(parameters): + parameters=list(parameters) + require(parameters and all(p.dtype==torch.float32 for p in parameters), + 'Optimizer capacity reserve requires FP32 parameters') + return 3*sum(p.numel()*p.element_size() for p in parameters) + + +def validate_multimodal_capacity(config_path, output, *, source_manifest_sha256, data_manifest_sha256): + """One real maximum-shape distributed forward/backward, zero optimizer steps. + + Invoke through the existing GPU allocation with torchrun, matching the + production rank batch. Uses an actual eligible robot subtask and all three + native RGB cameras; no synthetic all-valid camera-pose labels are invented. + """ + import inspect + from hydra.utils import instantiate + from egomimic.rldb.pixel_proprio import PreparedPixelProprioDataset, collate_pixel_proprio + from torch.nn.parallel import DistributedDataParallel + world=int(os.environ['WORLD_SIZE']);rank=int(os.environ['RANK']);local=int(os.environ['LOCAL_RANK']) + require(world in (2,4),'Capacity check requires production two/four ranks') + require(torch.cuda.device_count()==world, + 'Capacity visible CUDA device count must equal WORLD_SIZE before device or model access') + torch.cuda.set_device(local) + created=not torch.distributed.is_initialized() + if created:torch.distributed.init_process_group('nccl') + destination=Path(output);root=destination.parent;root.mkdir(parents=True,exist_ok=True) + require(not destination.exists(),'Capacity proof is immutable') + require(len(source_manifest_sha256)==64 and all(c in '0123456789abcdef' for c in source_manifest_sha256), + 'Capacity requires exact source manifest SHA256') + cfg=OmegaConf.load(config_path) + data_manifest_sha256=_capacity_manifest_sha256(cfg,data_manifest_sha256) + require(cfg.data.mode in ('discrete','continuous') and cfg.data.batch_size*world==32 + and cfg.data.distributed_world_size==world,'Wrong capacity temporal/global batch contract') + settings=OmegaConf.to_container(cfg.data,resolve=True) + accepted=set(inspect.signature(PreparedPixelProprioDataset).parameters) + settings={k:v for k,v in settings.items() if k in accepted} + frames=64 if cfg.data.mode=='continuous' else 16 + require(cfg.data.mode!='discrete' or cfg.data.discrete_frames==16,'Discrete capacity requires native16anchors') + settings.update(split='train') + if cfg.data.mode=='continuous':settings.update(min_frames=64,max_frames=64) + data=PreparedPixelProprioDataset(**settings) + offset=0;sample=None + for (episode,unit),count in zip(data.sample_units,data.counts): + if episode['embodiment'] in cfg.data.action_embodiments and len(episode['cameras'])==3: + sample=collate_pixel_proprio([data[(offset,0,rank*cfg.data.batch_size+i)] + for i in range(cfg.data.batch_size)]) + break + offset+=count + require(sample is not None,'No eligible actual64-frame three-camera robot subtask') + emb=next(iter(sample));item=sample[emb] + _check_multimodal_sample(item,boundary_policy=settings.get('boundary_sampling')) + require(all(len(m['frame_indices'])==frames and len(m['camera_names'])==3 + and m['image_size']==[160,160] and m['patch_size']==16 for m in item['pixel_metadata']), + 'Capacity sample is not actual64x3x160x160 native RGB') + require(item['world']['image']['value'].shape[:2]==(cfg.data.batch_size,frames*3*100) + and item['world']['image']['valid'].all(),'Missing actual maximum RGB tokens') + torch.manual_seed(int(cfg.seed)) + algo=instantiate(cfg.model.pipeline,device=f'cuda:{local}') + network=algo.pipeline.stage_by_id('joint_denoiser').model + image=network.image_model;props=torch.cuda.get_device_properties(local) + require('H100' in props.name and props.major>=9 and image.compile_mode=='default' + and image.compile_dynamic and image.compile_backend=='inductor' + and image.token_execution=='packed_valid' and image.blocks[0].attention_backend=='flash4', + 'Capacity requires exact H100/Flash4/packed/compiled execution') + class CapacityPass(torch.nn.Module): + def __init__(self): + super().__init__();self.nets=algo.nets + def forward(self,batch): + predictions=algo.forward_training(batch) + require(all(value.dtype==torch.bfloat16 for result in predictions.values() + for value in result['joint/prediction'].values()), + 'Maximum-shape actual expert outputs must be BF16') + return algo.compute_losses(predictions,batch)['loss'] + require(all(p.dtype==torch.float32 and p.requires_grad for p in algo.nets.parameters()), + 'Maximum-shape stored expert weights must be trainableFP32') + wrapped=DistributedDataParallel(CapacityPass(),device_ids=[local],find_unused_parameters=True) + batch=algo.process_batch_for_training(sample) + # Conservative declared memory reserve: two FP32 Adam moments plus one + # FP32 foreach workspace. These are memory-only buffers, never optimizer state. + optimizer_reserve_bytes=_optimizer_capacity_reserve_bytes(wrapped.parameters()) + optimizer_memory_reserve=torch.empty(optimizer_reserve_bytes//4,device=f'cuda:{local}',dtype=torch.float32) + optimizer_memory_reserve.zero_() + torch.cuda.reset_peak_memory_stats();started=time.perf_counter() + try: + with torch.autocast('cuda',dtype=torch.bfloat16):loss=wrapped(batch) + require(loss.dtype==torch.float32 and torch.isfinite(loss),'Capacity loss must be finiteFP32') + loss.backward() + active=[p for p in wrapped.parameters() if p.grad is not None] + require(active and all(p.dtype==torch.float32 and p.grad.dtype==torch.float32 + and torch.isfinite(p.grad).all() for p in active),'Capacity active gradients must be finiteFP32') + torch.cuda.synchronize() + record=dict(schema='native-multimodal-capacity-v1',status='PASS',rank=rank,world_size=world, + job_id=os.environ['SLURM_JOB_ID'],global_batch_size=32,rank_batch_size=int(cfg.data.batch_size), + frames=frames,cameras=3,image_size=[160,160],patch_size=16,image_tokens_per_sample=frames*300, + per_rank_batch=int(cfg.data.batch_size),backend='flash4',compiled_bf16_forward_backward=True, + peak_memory_bytes=torch.cuda.max_memory_allocated(),device_memory_bytes=props.total_memory, + optimizer_memory_reserve_bytes=optimizer_reserve_bytes, + optimizer_memory_reserve_policy='2FP32AdamMoments+1FP32ForeachWorkspace; nooptimizerstateorupdates', + source_manifest_sha256=source_manifest_sha256, + actual_embodiment=emb,optimizer_updates=0,attention_backend='flash4',token_execution='packed_valid', + compile_mode='default',compile_dynamic=True,bf16_forward=True,fp32_weights_gradients=True, + loss=float(loss.detach()),gpu_uuid=str(props.uuid),gpu_name=props.name, + peak_allocated_bytes=torch.cuda.max_memory_allocated(),peak_reserved_bytes=torch.cuda.max_memory_reserved(), + seconds=time.perf_counter()-started,config_raw_sha256=hashlib.sha256(Path(config_path).read_bytes()).hexdigest(), + config_sha256=hashlib.sha256(Path(config_path).read_bytes()).hexdigest(), + data_manifest_sha256=data_manifest_sha256) + write_once(root/f'{destination.stem}-rank{rank}.json',record) + ranks=[None]*world;torch.distributed.all_gather_object(ranks,record) + require(len({r['gpu_uuid'] for r in ranks})==world,'Capacity ranks must use distinct GPUs') + if rank==0:write_once(destination,dict(record,ranks=ranks)) + finally: + if created:torch.distributed.destroy_process_group() diff --git a/egomimic/pl_utils/world_model_behavior.py b/egomimic/pl_utils/world_model_behavior.py index 1b61a3f93..74a79090f 100644 --- a/egomimic/pl_utils/world_model_behavior.py +++ b/egomimic/pl_utils/world_model_behavior.py @@ -19,16 +19,11 @@ def _distributed_world(): def _capture_rng_state(*, current_cuda_only=False): cuda = [] if torch.cuda.is_available(): - cuda = ( - [torch.cuda.get_rng_state()] - if current_cuda_only - else torch.cuda.get_rng_state_all() - ) + cuda = ([torch.cuda.get_rng_state()] if current_cuda_only + else torch.cuda.get_rng_state_all()) state = { - "torch_cpu": torch.get_rng_state(), - "torch_cuda": cuda, - "numpy": np.random.get_state(), - "python": random.getstate(), + "torch_cpu": torch.get_rng_state(), "torch_cuda": cuda, + "numpy": np.random.get_state(), "python": random.getstate(), } if current_cuda_only: state["cuda_scope"] = "current_device" @@ -134,8 +129,6 @@ def on_save_checkpoint(self, checkpoint): def on_load_checkpoint(self, checkpoint): self._pending_rng = checkpoint.get("mot_rng_state") - if ( - getattr(self.context.trainer.datamodule, "checkpointable_sampling", False) - and self._pending_rng is None - ): + if (getattr(self.context.trainer.datamodule, "checkpointable_sampling", False) + and self._pending_rng is None): raise ValueError("Exact MoT resume requires saved model RNG state") diff --git a/egomimic/rldb/boundary_sampling.py b/egomimic/rldb/boundary_sampling.py new file mode 100644 index 000000000..6d8543ce4 --- /dev/null +++ b/egomimic/rldb/boundary_sampling.py @@ -0,0 +1,298 @@ +"""Opt-in bounded annotation-uncertainty sampling of recorded native frames. + +No annotation is rewritten and no neighbor instruction equivalence is asserted. +The legacy sampler supplies K and advances its original RNG before this helper; +the component/pair draw uses a separate checkpointed cursor/slot namespace. +""" +from __future__ import annotations + +import copy +from collections.abc import Mapping +import math +import numpy as np + +POLICY = { + "schema": "annotation-boundary-uncertainty-sampling-v1", + "policy": "annotation_boundary_uncertainty", + "boundary_probability": 0.5, + "maximum_radius_seconds": 0.25, + "duration_fraction": 0.1, + "rng_namespace": 0x424E4459, +} +SCHEMA = "annotation-boundary-uncertainty-realization-v1" + + +def _need(ok, message): + if not ok: + raise ValueError(message) + + +def validate_boundary_sampling(value): + if value is None: + return None + _need(isinstance(value, Mapping) and dict(value) == POLICY, + "Require the explicit versioned annotation_boundary_uncertainty policy") + _need(type(value["boundary_probability"]) is float + and type(value["maximum_radius_seconds"]) is float + and type(value["duration_fraction"]) is float + and type(value["rng_namespace"]) is int, + "Boundary policy requires exact typed parameters") + return copy.deepcopy(POLICY) + + +def _clock(full_time, episode): + times = np.asarray(full_time, dtype=np.float64) + _need(times.shape == (episode["frames"],) and np.isfinite(times).all() + and (np.diff(times) > 0).all(), "Boundary sampling needs real increasing episode timestamps") + return times + + +def _snapshot(unit): + return {key: copy.deepcopy(unit[key]) for key in + ("start_idx", "end_idx", "task_id", "instruction", "annotation_provenance")} + + +def sample_boundary_pair(legacy_chunks, *, full_time, episode, unit, all_units, + seed, draw, policy): + """Return one synchronized continuous-frame selection and honest provenance. + + The caller invokes its unmodified legacy sampler first, even in the boundary + branch. Therefore K and the following origin/role RNG draw order are exact. + A requested but unrealizable boundary pair falls back to that same sample. + """ + policy = validate_boundary_sampling(policy) + _need(policy is not None, "Boundary helper requires explicit opt-in") + _need(isinstance(draw, tuple) and len(draw) == 2 + and all(type(value) is int and value >= 0 for value in draw), + "Boundary sampling requires checkpointed cursor/global-slot keys") + _need(type(seed) is int and seed >= 0, "Boundary sampling requires nonnegative seed") + times = _clock(full_time, episode) + a, b = unit["start_idx"], unit["end_idx"] + _need(type(a) is int and type(b) is int and 0 <= a < b <= len(times), "Invalid original half-open unit") + chunks = np.asarray(legacy_chunks) + _need(chunks.ndim == 2 and chunks.shape[1] == 1 and chunks.dtype.kind in "iu" + and 2 <= len(chunks) <= min(64, b-a) + and (np.diff(chunks[:, 0]) > 0).all() + and ((a <= chunks[:, 0]) & (chunks[:, 0] < b)).all(), + "Boundary recipe requires original distinct continuous K2..64 sample") + segments = [list(segment) for segment in episode.get("continuity_segments", [[0, len(times)]]) + if len(segment) == 2 and segment[0] <= a < b <= segment[1]] + _need(len(segments) == 1, "Original subtask needs one unambiguous recording-continuity segment") + lo, hi = segments[0] + _need(type(lo) is int and type(hi) is int and 0 <= lo <= a < b <= hi <= len(times), + "Invalid physical continuity bounds") + if "continuity_segment" in unit: + _need(list(unit["continuity_segment"]) == [lo, hi], "Subtask continuity identity differs") + relative = times - times[0] + duration = float(times[b-1] - times[a]) + radius = min(policy["maximum_radius_seconds"], policy["duration_fraction"] * duration) + _need(duration > 0 and radius > 0, "No recorded physical duration for boundary jitter") + end_index = b if b < hi else b-1 + kind = "first_excluded_recorded_frame" if b < hi else "last_included_recorded_frame_at_recording_end" + centers = relative[[a, end_index]] + eligible = np.arange(lo, hi, dtype=np.int64) + bands = [eligible[np.abs(relative[eligible] - center) <= radius] for center in centers] + left = np.repeat(bands[0], len(bands[1])) + right = np.tile(bands[1], len(bands[0])) + pair_ok = left < right + left, right = left[pair_ok], right[pair_ok] + rng = np.random.default_rng(np.random.SeedSequence([seed, *draw, policy["rng_namespace"]])) + requested = bool(rng.random() < policy["boundary_probability"]) + reason = None + pair = None + selected = chunks[:, 0].copy() + if requested: + if len(bands[0]) < 2: + reason = "fewer_than_two_recorded_start_band_candidates" + elif len(bands[1]) < 2: + reason = "fewer_than_two_recorded_end_band_candidates" + elif len(left) < 2: + reason = "fewer_than_two_distinct_ordered_boundary_pairs" + else: + choice = int(rng.integers(len(left))) + pair = [int(left[choice]), int(right[choice])] + remainder = np.setdiff1d(np.arange(a, b, dtype=np.int64), pair) + extra = rng.choice(remainder, len(chunks)-2, replace=False) + selected = np.sort(np.r_[pair, extra]).astype(np.int64) + component = "boundary_pair" if pair is not None else ("uniform_fallback" if requested else "legacy_uniform") + admitted_lo = min(a, *(int(band[0]) for band in bands if len(band))) + admitted_hi = max(b, *(int(band[-1])+1 for band in bands if len(band))) + neighbors = [] + for other in all_units: + if other is unit or (other["start_idx"], other["end_idx"], other["task_id"]) == (a, b, unit["task_id"]): + continue + if other["start_idx"] < admitted_hi and admitted_lo < other["end_idx"]: + neighbors.append({key: copy.deepcopy(other[key]) for key in + ("start_idx", "end_idx", "task_id", "event_ids") if key in other}) + metadata = { + "schema": SCHEMA, "policy": policy, + "instruction_policy": "original_instruction_with_bounded_annotation_uncertainty", + "assert_neighbor_instruction_equivalence": False, + "original_subtask": _snapshot(unit), + "original_event_ids": copy.deepcopy(unit.get("event_ids", unit["annotation_provenance"].get("event_ids", []))), + "original_bounds": [a, b], "admitted_bounds": [admitted_lo, admitted_hi], + "episode_id": episode["episode_id"], "episode_frames": len(times), + "source_identity": {key: copy.deepcopy(episode[key]) for key in + ("embodiment", "rgb_path", "rgb_sha256", "path", "sha256") if key in episode}, + "continuity_segment": [lo, hi], + "continuity_first_last_timestamps_seconds": [float(times[lo]), float(times[hi-1])], + "duration_basis": "recorded_last_minus_first", "duration_seconds": duration, + "original_first_timestamp_seconds": float(times[a]), + "original_last_timestamp_seconds": float(times[b-1]), + "center_indices": [a, end_index], "center_kind": ["first_included_recorded_frame", kind], + "centers_seconds": [float(times[a]), float(times[end_index])], + "radius_seconds": radius, + "band_candidate_indices": [band.tolist() for band in bands], + "band_candidate_timestamps_seconds": [times[band].tolist() for band in bands], + "admissible_ordered_pair_count": int(len(left)), + "component": component, "requested_boundary": requested, + "pair_indices": pair, "fallback_reason": reason, + "rng_identity": {"seed": seed, "cursor": draw[0], "global_slot": draw[1], "namespace": policy["rng_namespace"]}, + "count": len(selected), "selected_frame_indices": selected.tolist(), + "selected_timestamps_seconds": times[selected].tolist(), + "selected_outside_original": selected[(selected < a) | (selected >= b)].tolist(), + "pair_offsets_frames": None if pair is None else [pair[0]-a, pair[1]-end_index], + "pair_offsets_seconds": None if pair is None else [float(times[pair[0]]-times[a]), float(times[pair[1]]-times[end_index])], + "physical_elapsed_seconds": float(times[selected[-1]]-times[selected[0]]), + "truncation": { + "start_band_left": bool(float(times[a])-radius < float(times[lo])), + "start_band_right": bool(float(times[a])+radius > float(times[hi-1])), + "end_band_left": bool(float(times[end_index])-radius < float(times[lo])), + "end_band_right": bool(float(times[end_index])+radius > float(times[hi-1])), + "terminal_center_truncated": bool(b == hi), + }, + "mixture_counters": {"requested_boundary": int(requested), "realized_boundary": int(pair is not None), + "fallback": int(requested and pair is None), "legacy_uniform": int(not requested)}, + "neighbor_units_if_available": neighbors, + } + return selected[:, None], metadata + + +def validate_boundary_record(record, unit, indices): + """Strict opt-in smoke provenance, without weakening the legacy branch.""" + packet = record["boundary_sampling"] + _need(packet.get("schema") == SCHEMA and validate_boundary_sampling(packet.get("policy")) is not None, + "Missing explicit boundary-uncertainty provenance") + _need(packet.get("instruction_policy") == "original_instruction_with_bounded_annotation_uncertainty" + and packet.get("assert_neighbor_instruction_equivalence") is False, + "Boundary uncertainty may not assert exact neighbor label truth") + a, b = unit["start_idx"], unit["end_idx"] + _need(packet["original_subtask"] == _snapshot(unit) and packet["original_bounds"] == [a, b] + and packet["episode_id"] == record["episode_id"], "Original annotation identity was rewritten") + selected = packet["selected_frame_indices"] + times = packet["selected_timestamps_seconds"] + _need(record["frame_indices"] == selected and record["timestamps_seconds"] == times + and len(selected) == packet["count"] and 2 <= len(selected) <= min(64, b-a) + and all(type(i) is int for i in selected) and all(x < y for x, y in zip(selected, selected[1:])) + and len(times) == len(selected) and all(math.isfinite(t) for t in times) + and all(x < y for x, y in zip(times, times[1:])), "Invalid actual boundary selection/clock") + lo, hi = packet["continuity_segment"] + admitted_lo, admitted_hi = packet["admitted_bounds"] + _need(0 <= lo <= admitted_lo <= a < b <= admitted_hi <= hi <= packet["episode_frames"] + and indices and all(admitted_lo <= i < admitted_hi for i in indices), + "Boundary selection crossed actual recording/continuity bounds") + _need(set(indices) == set(selected), "Boundary modalities are not synchronized to one selection") + duration = packet["original_last_timestamp_seconds"] - packet["original_first_timestamp_seconds"] + radius = min(0.25, 0.1*duration) + _need(packet["duration_basis"] == "recorded_last_minus_first" and duration > 0 + and math.isclose(packet["duration_seconds"], duration, rel_tol=1e-12, abs_tol=1e-9) + and math.isclose(packet["radius_seconds"], radius, rel_tol=1e-12, abs_tol=1e-9), + "Boundary radius differs from physical original span") + end_index = b if b < hi else b-1 + expected_kind = "first_excluded_recorded_frame" if b < hi else "last_included_recorded_frame_at_recording_end" + _need(packet["center_indices"] == [a, end_index] + and packet["center_kind"] == ["first_included_recorded_frame", expected_kind] + and packet["centers_seconds"][0] == packet["original_first_timestamp_seconds"], + "Invalid half-open boundary center semantics") + _need(len(packet["centers_seconds"]) == 2 + and all(math.isfinite(v) for v in packet["centers_seconds"]) + and (packet["centers_seconds"][1] == packet["original_last_timestamp_seconds"] if b == hi + else packet["centers_seconds"][1] > packet["original_last_timestamp_seconds"]), + "End center is not the recorded half-open or terminal center") + bands = packet["band_candidate_indices"] + band_times = packet["band_candidate_timestamps_seconds"] + _need(len(bands) == len(band_times) == 2, "Missing recorded boundary bands") + for band, clock, center_index, center in zip(bands, band_times, packet["center_indices"], packet["centers_seconds"]): + _need(len(band) == len(clock) and band and center_index in band + and all(type(v) is int and lo <= v < hi for v in band) + and all(x < y for x,y in zip(band,band[1:])) + and all(math.isfinite(v) for v in clock) + and all(x < y for x,y in zip(clock,clock[1:])) + and clock[band.index(center_index)] == center, + "Boundary candidate indices/timestamps are not actual ordered continuous records") + for index, value in zip(band,clock): + tolerance = abs(float(np.spacing(max(abs(value),abs(center),1.0))))*4 + _need(abs(value-center) <= radius+tolerance, "Candidate band exceeds physical radius") + if index in selected: + _need(times[selected.index(index)] == value, "Pair clock differs from selected native frame") + pair_count = sum(i < j for i in bands[0] for j in bands[1]) + _need(packet["admissible_ordered_pair_count"] == pair_count + and [min(a,bands[0][0],bands[1][0]), max(b,bands[0][-1]+1,bands[1][-1]+1)] == packet["admitted_bounds"], + "Admitted bounds or pair count differs from actual candidate bands") + outside = [i for i in selected if i < a or i >= b] + _need(packet["selected_outside_original"] == outside, "Outward frames lost uncertainty provenance") + identity = packet["rng_identity"] + _need(set(identity) == {"seed", "cursor", "global_slot", "namespace"} + and all(type(identity[k]) is int and identity[k] >= 0 for k in ("seed", "cursor", "global_slot")) + and type(identity["namespace"]) is int and identity["namespace"] == POLICY["rng_namespace"], + "Invalid checkpointed boundary RNG identity") + rng = np.random.default_rng(np.random.SeedSequence([identity["seed"],identity["cursor"],identity["global_slot"],identity["namespace"]])) + requested = bool(rng.random() < POLICY["boundary_probability"]) + _need(type(packet["requested_boundary"]) is bool and packet["requested_boundary"] == requested, + "Requested mixture does not match checkpointed component draw") + component = packet["component"] + pair = packet["pair_indices"] + if component == "boundary_pair": + _need(packet["requested_boundary"] is True and packet["fallback_reason"] is None + and isinstance(pair, list) and len(pair) == 2 and pair[0] < pair[1] + and all(i in selected for i in pair) and all(i in pair for i in outside) + and all(len(band) >= 2 for band in packet["band_candidate_indices"]) + and packet["admissible_ordered_pair_count"] >= 2, "Boundary branch is not a genuine random pair") + pairs = [(i,j) for i in bands[0] for j in bands[1] if i < j] + expected_pair = list(pairs[int(rng.integers(len(pairs)))]) + remainder = np.setdiff1d(np.arange(a,b,dtype=np.int64),expected_pair) + expected_selected = np.sort(np.r_[expected_pair,rng.choice(remainder,len(selected)-2,replace=False)]).tolist() + _need(pair == expected_pair and selected == expected_selected, + "Boundary pair/remaining frames differ from deterministic cursor/slot draw") + for index, center, band in zip(pair, packet["centers_seconds"], packet["band_candidate_indices"]): + value = times[selected.index(index)] + tolerance = abs(float(np.spacing(max(abs(value), abs(center), 1.0)))) * 4 + _need(index in band and abs(value-center) <= radius+tolerance, + "Outward pair exceeds bounded physical-time jitter") + else: + _need(component in ("legacy_uniform", "uniform_fallback") and pair is None and not outside, + "Uniform component may not contain outward frames") + _need((component == "uniform_fallback") == packet["requested_boundary"], "Misreported requested/realized mixture") + _need((packet["fallback_reason"] is not None) == (component == "uniform_fallback"), "Missing fallback reason") + if component == "uniform_fallback": + reason = ("fewer_than_two_recorded_start_band_candidates" if len(bands[0]) < 2 else + "fewer_than_two_recorded_end_band_candidates" if len(bands[1]) < 2 else + "fewer_than_two_distinct_ordered_boundary_pairs" if pair_count < 2 else None) + _need(reason is not None and packet["fallback_reason"] == reason, "Fallback does not match actual insufficient candidates") + _need(type(packet["requested_boundary"]) is bool, "Requested component must be an actual boolean") + if pair is not None: + actual_offsets = [pair[0]-a,pair[1]-end_index] + time_offsets = [times[selected.index(pair[0])]-packet["centers_seconds"][0], + times[selected.index(pair[1])]-packet["centers_seconds"][1]] + _need(packet["pair_offsets_frames"] == actual_offsets and packet["pair_offsets_seconds"] == time_offsets, + "Boundary offsets differ from actual selected pair") + else: + _need(packet["pair_offsets_frames"] is None and packet["pair_offsets_seconds"] is None, + "Uniform selection invented a boundary pair") + _need(packet["physical_elapsed_seconds"] == times[-1]-times[0], "Physical span differs from recorded selection") + first,last = packet["continuity_first_last_timestamps_seconds"] + centers = packet["centers_seconds"] + expected_truncation = { + "start_band_left": centers[0]-radius < first, + "start_band_right": centers[0]+radius > last, + "end_band_left": centers[1]-radius < first, + "end_band_right": centers[1]+radius > last, + "terminal_center_truncated": b == hi, + } + _need(packet["truncation"] == expected_truncation, "Recording truncation provenance differs") + expected_counters = {"requested_boundary": int(packet["requested_boundary"]), + "realized_boundary": int(component == "boundary_pair"), + "fallback": int(component == "uniform_fallback"), + "legacy_uniform": int(component == "legacy_uniform")} + _need(packet["mixture_counters"] == expected_counters, "Misreported boundary mixture counters") + return packet diff --git a/egomimic/rldb/camera_pose.py b/egomimic/rldb/camera_pose.py new file mode 100644 index 000000000..1fa6b7499 --- /dev/null +++ b/egomimic/rldb/camera_pose.py @@ -0,0 +1,109 @@ +"""Measured camera-to-reference poses for world-model windows. + +Aria ``obs_head_pose`` stores world-from-RGB-camera, xyz then quaternion wxyz. +EVA native consumers treat stored arm extrinsics as base-from-camera (despite +the historical T_cam_base name); preserve that actual consumer convention. Wrist-camera poses require their own measured calibration. +No unknown pose is replaced by an identity measurement. +""" +import numpy as np +from scipy.spatial.transform import Rotation + + +def camera_pose_matrices(xyzwxyz): + """Convert native xyz+wxyz camera-to-world observations to SE(3).""" + pose = np.asarray(xyzwxyz, dtype=np.float64) + if pose.ndim < 2 or pose.shape[-1] != 7 or not np.isfinite(pose).all(): + raise ValueError("Expected finite native xyz+wxyz camera poses") + if not np.allclose(np.linalg.norm(pose[..., 3:], axis=-1), 1, atol=1e-3): + raise ValueError("Native camera quaternion must be unit length") + result = np.broadcast_to(np.eye(4), pose.shape[:-1] + (4, 4)).copy() + result[..., :3, :3] = Rotation.from_quat(pose[..., [4, 5, 6, 3]].reshape(-1, 4)).as_matrix().reshape(pose.shape[:-1] + (3, 3)) + result[..., :3, 3] = pose[..., :3] + return result + + +def relative_camera_pose(world_from_camera, valid, reference_index, *, observed=None): + """Return per-camera xyz+first-two-rotation-columns and validity. + + Inputs are T,C,4,4 in meters, T,C validity, and one reference time for + each camera (scalar accepted). Training selects the first sampled time of + an arbitrary window, independently of diffusion noise. Inference supplies + observed=T,C and rejects a reference not in real observed context. + Invalid rows remain NaN and false; they cannot become training targets. + """ + poses = np.asarray(world_from_camera, dtype=np.float64) + valid = np.asarray(valid) + if poses.ndim != 4 or poses.shape[-2:] != (4, 4) or valid.shape != poses.shape[:2] or valid.dtype != bool: + raise ValueError("Expected T,C,4,4 poses and boolean T,C validity") + n, c = valid.shape + references = np.broadcast_to(np.asarray(reference_index), (c,)) + if not np.issubdtype(references.dtype, np.integer) or np.any((references < 0) | (references >= n)): + raise ValueError("Reference indices outside selected window") + if observed is not None: + observed = np.asarray(observed) + if observed.shape != valid.shape or observed.dtype != bool: + raise ValueError("Observed camera roles must be boolean T,C") + out = np.full((n, c, 9), np.nan, dtype=np.float32) + usable = valid.copy() + for camera, ref in enumerate(references): + if not valid[ref, camera]: + usable[:, camera] = False + continue + if observed is not None and not observed[ref, camera]: + raise ValueError("Inference camera reference must be a real observed pose") + selected = poses[valid[:, camera], camera] + r = selected[:, :3, :3] + if (not np.isfinite(selected).all() or not np.allclose(selected[:, 3], [0, 0, 0, 1]) + or not np.allclose(r.swapaxes(-1, -2) @ r, np.eye(3), atol=1e-5) + or not np.allclose(np.linalg.det(r), 1, atol=1e-5)): + raise ValueError("Invalid measured camera rigid transform") + relative = np.linalg.inv(poses[ref, camera]) @ selected + out[valid[:, camera], camera] = np.concatenate( + [relative[:, :3, 3], relative[:, :3, 0], relative[:, :3, 1]], -1 + ) + return out, usable + + +def native_camera_world_poses(embodiment, camera_names, total_frames, *, + head_pose=None, extrinsics=None, robot_reference_arm="left"): + """Build measured camera targets from the native Zarr schema without writes. + + Returns absolute T,C,4,4 and boolean validity for preparation, not already + sample-relative targets. The window loader must select its reference after + sampling. EVA uses the chosen arm base as a fixed world coordinate system; + its stored front-camera calibration is base-from-camera in native consumers. + Human RGB world pose is ``obs_head_pose`` including its optical correction. + The schema contains no wrist-camera calibration: those cameras stay missing. + """ + if not camera_names or len(set(camera_names)) != len(camera_names) or total_frames < 1: + raise ValueError("Unique camera names and positive frame count required") + out = np.full((total_frames, len(camera_names), 4, 4), np.nan, dtype=np.float64) + valid = np.zeros((total_frames, len(camera_names)), dtype=bool) + if "front_1" not in camera_names: + return out, valid + camera = camera_names.index("front_1") + if embodiment == "eva_bimanual": + if not isinstance(extrinsics, dict) or robot_reference_arm not in extrinsics: + return out, valid + measured = np.asarray(extrinsics[robot_reference_arm], dtype=np.float64) + if measured.shape != (4, 4) or not np.isfinite(measured).all(): + raise ValueError("Malformed measured EVA camera calibration") + # Shared validator checks orthogonality/determinant before admitting it. + relative_camera_pose(measured[None, None], np.ones((1, 1), bool), 0) + out[:, camera] = measured + valid[:, camera] = True + elif embodiment in ("human", "human_bimanual", "aria_bimanual"): + if head_pose is None: + return out, valid + pose = np.asarray(head_pose, dtype=np.float64) + if pose.shape != (total_frames, 7): + raise ValueError("Native obs_head_pose must have T,7 layout") + measured = (np.isfinite(pose).all(axis=1) + & (np.abs(pose[:,:3]).max(axis=1) < 1000) + & np.isclose(np.linalg.norm(pose[:, 3:], axis=1), 1, atol=1e-3)) + if measured.any(): + out[measured, camera] = camera_pose_matrices(pose[measured]) + valid[measured, camera] = True + else: + raise ValueError("Unsupported native camera embodiment; require explicit converter") + return out, valid diff --git a/egomimic/rldb/pixel_language.py b/egomimic/rldb/pixel_language.py new file mode 100644 index 000000000..0beb70abd --- /dev/null +++ b/egomimic/rldb/pixel_language.py @@ -0,0 +1,397 @@ +"""Verified, immutable task-language features for prepared pixel datasets.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +LANGUAGE_FEATURE_DIMENSION = 1024 +LANGUAGE_FEATURE_SCHEMA = "pixel-task-language-features-v1" +LANGUAGE_VECTOR_CACHE_SCHEMA = "pixel-language-vector-cache-v1" + + +def validate_language_feature_options(path, sha256): + """Keep the opt-in explicit: neither an unbound file nor a digest alone.""" + if (path is None) != (sha256 is None): + raise ValueError( + "language_features_path and language_features_sha256 require both or neither" + ) + if path is not None and ( + not isinstance(path, (str, Path)) + or not str(path).strip() + or not isinstance(sha256, str) + or re.fullmatch(r"[0-9a-f]{64}", sha256) is None + ): + raise ValueError("Language features require a path and exact lowercase SHA256") + + +def _unique_object(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate language sidecar JSON key: {key}") + result[key] = value + return result + + +def _nonfinite_json(value): + raise ValueError(f"Nonfinite language sidecar JSON value: {value}") + + +@dataclass(frozen=True) +class PixelTaskLanguageFeatures: + """Loaded features plus the complete small identity used by checkpoints.""" + + features: dict[str, torch.Tensor] + instructions: dict[str, str] + identity: dict + + def context_identity(self): + return copy.deepcopy(self.identity) + + def tensor(self, task_id): + """Return a private FP32 feature for a real bound task.""" + return self.features[task_id].clone() + + def sample(self, task_id): + # A consumer must not be able to mutate the shared feature table. + return dict( + language_features=self.tensor(task_id), + language_instruction=self.instructions[task_id], + task_id=task_id, + ) + + +def _load_language_features(path, sha256, manifest=None, manifest_sha256=None, *, vector_cache=False): + """Bind every task in both splits to finite, unit-normalized frozen features. + + The sidecar never rewrites the prepared manifest or caches. Encoder metadata + is retained in full, including any additional generation provenance. + """ + validate_language_feature_options(path, sha256) + if path is None: + return None + raw = Path(path).read_bytes() + if hashlib.sha256(raw).hexdigest() != sha256: + raise ValueError("Language features SHA256 mismatch") + sidecar = json.loads( + raw, object_pairs_hook=_unique_object, parse_constant=_nonfinite_json + ) + if ( + not isinstance(sidecar, dict) + or sidecar.get("schema") != (LANGUAGE_VECTOR_CACHE_SCHEMA if vector_cache else LANGUAGE_FEATURE_SCHEMA) + ): + raise ValueError("Unsupported pixel task-language feature schema") + if vector_cache: + if (sidecar.get("data_admitted") is not False or + re.fullmatch(r"[0-9a-f]{64}", str(sidecar.get("source_catalog_sha256"))) is None): + raise ValueError("Vector cache requires an exact catalog identity and cannot admit data") + elif sidecar.get("prepared_manifest_sha256") != manifest_sha256: + raise ValueError("Language features prepared manifest SHA256 mismatch") + encoder = sidecar.get("encoder") + if ( + not isinstance(encoder, dict) + or any( + not isinstance(encoder.get(key), str) or not encoder[key].strip() + for key in ("model_id", "revision", "tokenizer_revision", "pooling") + ) + or type(encoder.get("dimension")) is not int + or encoder["dimension"] != LANGUAGE_FEATURE_DIMENSION + or encoder.get("frozen") is not True + or encoder.get("normalization") != "l2" + ): + raise ValueError( + "Language features require exact frozen encoder provenance and 1024D L2 normalization" + ) + episodes = ([dict(task_id=tid) for tid in sidecar.get("tasks", {})] + if vector_cache else manifest.get("episodes", [])) + if not episodes or any( + not isinstance(ep.get("task_id"), str) or not ep["task_id"].strip() + for ep in episodes + ): + raise ValueError("Every prepared episode requires an explicit language task_id") + tasks = sidecar.get("tasks") + if not isinstance(tasks, dict) or set(tasks) != {ep["task_id"] for ep in episodes}: + raise ValueError( + "Language feature tasks must exactly cover all prepared episodes in both splits" + ) + features, instructions = {}, {} + for task_id, item in tasks.items(): + if ( + not isinstance(item, dict) + or not isinstance(item.get("instruction"), str) + or not item["instruction"].strip() + ): + raise ValueError(f"Task {task_id} requires a nonempty language instruction") + values = item.get("features") + if ( + not isinstance(values, list) + or len(values) != LANGUAGE_FEATURE_DIMENSION + or any(type(value) not in (int, float) for value in values) + ): + raise ValueError( + f"Task {task_id} requires exactly 1024 numeric language features" + ) + try: + array = np.asarray(values, dtype=np.float64) + except (OverflowError, TypeError, ValueError) as error: + raise ValueError( + f"Task {task_id} language features must be finite numbers" + ) from error + if not np.isfinite(array).all() or not np.isclose( + np.linalg.norm(array), 1.0, rtol=0, atol=1e-4 + ): + raise ValueError( + f"Task {task_id} language features must be finite and unit-normalized" + ) + features[task_id] = torch.from_numpy(array.astype(np.float32)) + instructions[task_id] = item["instruction"] + return PixelTaskLanguageFeatures( + features=features, + instructions=instructions, + identity=dict( + schema=sidecar["schema"], + sha256=sha256, + source_catalog_sha256=sidecar.get("source_catalog_sha256"), + prepared_manifest_sha256=manifest_sha256, + encoder=copy.deepcopy(encoder), + tasks=copy.deepcopy(instructions), + ), + ) + + +def load_pixel_task_language_features(path, sha256, manifest, manifest_sha256): + """Admit only exact prepared-manifest-bound language sidecars.""" + return _load_language_features(path, sha256, manifest, manifest_sha256) + + +def load_reusable_language_features(path, sha256): + """Read an SHA-bound text/encoder cache without admitting dataset rows. + + Frozen text vectors depend on exact text and encoder identity, not on the + physical episode split. Final dataset sidecars are rebound separately. + """ + validate_language_feature_options(path, sha256) + raw = Path(path).read_bytes() + if hashlib.sha256(raw).hexdigest() != sha256: + raise ValueError("Reusable language cache SHA256 mismatch") + value = json.loads(raw, object_pairs_hook=_unique_object, parse_constant=_nonfinite_json) + if value.get("schema") == LANGUAGE_VECTOR_CACHE_SCHEMA: + return _load_language_features(path, sha256, vector_cache=True) + bound = value.get("prepared_manifest_sha256") + if re.fullmatch(r"[0-9a-f]{64}", str(bound)) is None: + raise ValueError("Reusable sidecar has no valid original manifest identity") + manifest = dict(episodes=[dict(task_id=tid) for tid in value.get("tasks", {})]) + return load_pixel_task_language_features(path, sha256, manifest, bound) + + +def collate_pixel_language(samples): + """Language is an all-or-none batch contract, with explicit provenance.""" + keys = {"language_features", "language_instruction", "task_id"} + if not any(keys.intersection(sample) for sample in samples): + return {} + if any(not keys.issubset(sample) for sample in samples): + raise ValueError( + "Language features and provenance must be present for every sample" + ) + for sample in samples: + features = sample["language_features"] + if ( + not isinstance(features, torch.Tensor) + or features.dtype != torch.float32 + or features.shape != (LANGUAGE_FEATURE_DIMENSION,) + or not torch.isfinite(features).all() + or any( + not isinstance(sample[key], str) or not sample[key].strip() + for key in ("language_instruction", "task_id") + ) + ): + raise ValueError("Invalid per-sample language features or provenance") + return dict( + language_features=torch.stack( + [sample["language_features"] for sample in samples] + ), + language_instruction=[sample["language_instruction"] for sample in samples], + task_id=[sample["task_id"] for sample in samples], + ) + + +QWEN_MODEL_ID = "Qwen/Qwen3-Embedding-0.6B" +QWEN_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3" + + +def encode_subtask_language(manifest_path, index_path, index_sha256, output_path, *, cache_dir, device="cpu", reuse_features_path=None, reuse_features_sha256=None): + """Encode reviewed subtask text using the unchanged native Qwen pooler. + + Exact revisions and weights are cached once. Existing verified outputs take + the fast path without loading the model or redownloading its weights. + """ + from egomimic.rldb.pixel_subtasks import load_subtask_manifest + manifest_raw = Path(manifest_path).read_bytes() + manifest = json.loads(manifest_raw) + manifest_sha = hashlib.sha256(manifest_raw).hexdigest() + units = load_subtask_manifest(index_path, index_sha256, manifest, manifest_sha) + instructions = {} + for spans in units.values(): + for span in spans: + tid, text = span["task_id"], span["instruction"] + if tid in instructions and instructions[tid] != text: + raise ValueError("Task ID aliases different subtask instructions") + instructions[tid] = text + return _encode_instructions(instructions, output_path, + binding=dict(schema=LANGUAGE_FEATURE_SCHEMA, prepared_manifest_sha256=manifest_sha, + subtask_index_sha256=index_sha256), + cache_dir=cache_dir, device=device, reuse_features_path=reuse_features_path, + reuse_features_sha256=reuse_features_sha256) + + +def encode_instruction_catalog(catalog_path, catalog_sha256, output_path, *, cache_dir, + device="cpu", reuse_features_path=None, reuse_features_sha256=None): + """Precompute exact text vectors; this never admits a training window.""" + validate_language_feature_options(catalog_path, catalog_sha256) + raw=Path(catalog_path).read_bytes() + if hashlib.sha256(raw).hexdigest() != catalog_sha256: + raise ValueError("Instruction catalog SHA256 mismatch") + catalog=json.loads(raw,object_pairs_hook=_unique_object,parse_constant=_nonfinite_json) + tasks=catalog.get("tasks") + if (catalog.get("schema") != "pixel-instruction-catalog-v1" or + catalog.get("data_admitted") is not False or catalog.get("clock_verified") is not False or + not isinstance(tasks,dict) or not tasks or + any(not isinstance(k,str) or not k.strip() or not isinstance(v,str) or not v.strip() + for k,v in tasks.items())): + raise ValueError("Expected a non-admitting exact instruction catalog") + return _encode_instructions(tasks,output_path, + binding=dict(schema=LANGUAGE_VECTOR_CACHE_SCHEMA,source_catalog_sha256=catalog_sha256, + data_admitted=False,clock_verified=False), + cache_dir=cache_dir,device=device,reuse_features_path=reuse_features_path, + reuse_features_sha256=reuse_features_sha256) + + +def _encode_instructions(instructions, output_path, *, binding, cache_dir, device, + reuse_features_path=None, reuse_features_sha256=None): + import importlib.metadata + from huggingface_hub import snapshot_download + from egomimic.scripts.embedding_process.qwen3_embedding import Qwen3TextEmbedding + language_manifest = dict(episodes=[dict(task_id=tid) for tid in instructions]) + def validate_output(digest): + if binding["schema"] == LANGUAGE_VECTOR_CACHE_SCHEMA: + cached = _load_language_features(output, digest, vector_cache=True) + if cached.identity["source_catalog_sha256"] != binding["source_catalog_sha256"]: + raise ValueError("Existing vector cache has another catalog identity") + return cached + return load_pixel_task_language_features(output, digest, language_manifest, binding["prepared_manifest_sha256"]) + output = Path(output_path) + if output.exists(): + digest = hashlib.sha256(output.read_bytes()).hexdigest() + cached = validate_output(digest) + if cached.instructions != instructions: + raise ValueError("Existing language features differ from reviewed instructions") + if cached.context_identity()["encoder"]["revision"] != QWEN_REVISION: + raise ValueError("Existing encoder output has a different revision") + return dict(status="VERIFIED_EXISTING",path=str(output),sha256=digest,tasks=len(instructions)) + reused = {} + reuse_encoder = None + if reuse_features_path is not None or reuse_features_sha256 is not None: + validate_language_feature_options(reuse_features_path,reuse_features_sha256) + cached=load_reusable_language_features(reuse_features_path,reuse_features_sha256) + reuse_encoder=cached.context_identity()['encoder'] + expected = dict(model_id=QWEN_MODEL_ID, revision=QWEN_REVISION, + tokenizer_revision=QWEN_REVISION, pooling="last_nonpadding_token", + dimension=1024, frozen=True, normalization="l2", dtype="float32", max_length=128) + if any(reuse_encoder.get(k) != v for k,v in expected.items()): + raise ValueError('Reuse requires the same exact native encoder contract') + reused={tid:cached.features[tid].numpy() for tid,text in instructions.items() + if cached.instructions.get(tid)==text} + missing_ids=sorted(set(instructions)-set(reused)) + snapshot = Path(snapshot_download(repo_id=QWEN_MODEL_ID,revision=QWEN_REVISION, + cache_dir=cache_dir,local_files_only=not bool(missing_ids),allow_patterns=["config.json","tokenizer*","special_tokens_map.json", + "vocab.json","merges.txt","model*.safetensors"])) + weights = sorted(snapshot.glob("*.safetensors")) + if not weights: + raise ValueError("Pinned Qwen snapshot has no safetensors weights") + proof_path = Path(cache_dir)/"qwen-weights-proof.json" + weight_identities=[] + if proof_path.exists(): + proof=json.loads(proof_path.read_text()) + if proof.get("revision") != QWEN_REVISION: + raise ValueError("Cached Qwen proof revision mismatch") + weight_identities=proof["weights"] + if [(x["name"],x["bytes"]) for x in weight_identities] != [(x.name,x.stat().st_size) for x in weights]: + raise ValueError("Cached Qwen weight identity mismatch") + else: + for weight in weights: + digest=hashlib.sha256() + with weight.open("rb") as handle: + for chunk in iter(lambda:handle.read(8*1024*1024),b""): + digest.update(chunk) + weight_identities.append(dict(name=weight.name,bytes=weight.stat().st_size,sha256=digest.hexdigest())) + with proof_path.open("x") as handle: + json.dump(dict(revision=QWEN_REVISION,weights=weight_identities),handle,indent=2) + weights_digest = (weight_identities[0]["sha256"] if len(weight_identities)==1 else + hashlib.sha256(json.dumps(weight_identities,sort_keys=True).encode()).hexdigest()) + if reuse_encoder is not None and reuse_encoder.get("weights_sha256") != weights_digest: + raise ValueError("Reusable language vectors use different exact encoder weights") + task_ids=sorted(instructions) + if missing_ids: + embedder=Qwen3TextEmbedding(input_keys=["annotations"],model_name=str(snapshot), + batch_size=4,max_length=128,device=device,dtype=torch.float32) + embedder.setup() + embedder.model.requires_grad_(False) + new_vectors=embedder._embed_texts([instructions[tid] for tid in missing_ids]) + else: + new_vectors=np.empty((0,1024),np.float32) + by_id=dict(reused);by_id.update(zip(missing_ids,new_vectors)) + vectors=np.stack([by_id[tid] for tid in task_ids]) + if vectors.shape != (len(task_ids),1024) or not np.isfinite(vectors).all(): + raise ValueError("Native Qwen embedding output shape/finiteness failure") + value=dict(**binding, + feature_reuse=dict(source_sha256=reuse_features_sha256,reused_tasks=len(reused),new_tasks=len(missing_ids)), + encoder=dict(model_id=QWEN_MODEL_ID,revision=QWEN_REVISION,tokenizer_revision=QWEN_REVISION, + pooling="last_nonpadding_token",dimension=1024,frozen=True,normalization="l2", + dtype="float32",max_length=128,weights_sha256=weights_digest,weights=weight_identities, + torch_version=torch.__version__,transformers_version=importlib.metadata.version("transformers")), + tasks={tid:dict(instruction=instructions[tid],features=vec.tolist()) for tid,vec in zip(task_ids,vectors)}) + with output.open("x") as handle: + json.dump(value,handle,indent=2) + handle.write("\n") + digest=hashlib.sha256(output.read_bytes()).hexdigest() + validate_output(digest) + return dict(status="VERIFIED_NEW",path=str(output),sha256=digest,tasks=len(task_ids)) + + +def main(): + import argparse + parser=argparse.ArgumentParser(description="Native immutable subtask language features") + parser.add_argument("--prepared-manifest") + parser.add_argument("--instruction-catalog") + parser.add_argument("--catalog-sha256") + parser.add_argument("--subtask-index") + parser.add_argument("--subtask-sha256") + parser.add_argument("--output",required=True) + parser.add_argument("--encoder-cache",required=True) + parser.add_argument("--device",default="cpu") + parser.add_argument("--reuse-features") + parser.add_argument("--reuse-features-sha256") + args=parser.parse_args() + if args.instruction_catalog: + if args.prepared_manifest or args.subtask_index or args.subtask_sha256 or not args.catalog_sha256: + parser.error("Catalog mode requires --catalog-sha256 and excludes prepared-data arguments") + print(json.dumps(encode_instruction_catalog(args.instruction_catalog,args.catalog_sha256, + args.output,cache_dir=args.encoder_cache,device=args.device, + reuse_features_path=args.reuse_features,reuse_features_sha256=args.reuse_features_sha256))) + return + if not all((args.prepared_manifest,args.subtask_index,args.subtask_sha256)) or args.catalog_sha256: + parser.error("Prepared mode requires manifest, subtask index and SHA") + print(json.dumps(encode_subtask_language(args.prepared_manifest,args.subtask_index, + args.subtask_sha256,args.output,cache_dir=args.encoder_cache,device=args.device,reuse_features_path=args.reuse_features,reuse_features_sha256=args.reuse_features_sha256))) + + +if __name__ == "__main__": + main() diff --git a/egomimic/rldb/pixel_proprio.py b/egomimic/rldb/pixel_proprio.py index cb9909edb..84a518f3c 100644 --- a/egomimic/rldb/pixel_proprio.py +++ b/egomimic/rldb/pixel_proprio.py @@ -7,9 +7,13 @@ from __future__ import annotations +import hashlib import json import re from contextlib import ExitStack +from egomimic.rldb.pixel_subtasks import load_subtask_manifest +from egomimic.rldb.boundary_sampling import sample_boundary_pair, validate_boundary_sampling +from egomimic.rldb.camera_pose import relative_camera_pose from dataclasses import dataclass from itertools import islice from pathlib import Path @@ -18,6 +22,11 @@ import torch from torch.utils.data import Dataset +from egomimic.rldb.pixel_language import ( + collate_pixel_language, + load_pixel_task_language_features, + validate_language_feature_options, +) from egomimic.rldb.world_model import ( HomogeneousWorldBatchSampler, _physical_times, @@ -34,7 +43,7 @@ def collate_pixel_proprio(samples): if not samples or len({s["embodiment"] for s in samples}) != 1: raise ValueError("Batch must be nonempty and homogeneous in embodiment") names = set(samples[0]["world"]) - if names not in ({"image"}, {"image", "proprio"}, {"image", "proprio", "action"}): + if names - {"camera"} not in ({"image"}, {"image", "proprio"}, {"image", "proprio", "action"}): raise ValueError("Expected image stream with optional proprio/action labels") if any(set(s["world"]) != names for s in samples): raise ValueError("Cannot mix objectives in a batch") @@ -54,13 +63,14 @@ def collate_pixel_proprio(samples): values.append(target) result[name][key] = torch.stack(values) kinds = [s.get("task_kind") for s in samples] - allowed = {"forecast", "completion", "direct_action"} + allowed = {"forecast", "completion", "direct_action", "diffusion_forcing"} if any(kind not in allowed for kind in kinds): raise ValueError( "Every sample requires its explicit observation-mask task kind" ) embodiment = samples[0]["embodiment"] item = dict(embodiment=embodiment, world=result, task_kind=kinds) + item.update(collate_pixel_language(samples)) if any("pixel_metadata" in s for s in samples): if not all("pixel_metadata" in s for s in samples): raise ValueError( @@ -216,7 +226,7 @@ def __init__(self, dataset, anchors_per_episode, batch_size=1): raise ValueError("Positive integer anchor and batch counts required") self.groups, self.batch_size = {}, batch_size offset = 0 - for ep, count in zip(dataset.episodes, dataset.counts): + for ep, count in zip(getattr(dataset, "sampling_episodes", dataset.episodes), dataset.counts): selected = np.linspace( 0, count - 1, min(anchors_per_episode, count), dtype=np.int64 ) @@ -260,6 +270,128 @@ def sample_full_episode_times(length, rng, min_frames=2, max_frames=128, anchor= return np.sort(np.r_[anchor, other]).astype(np.int64) +def validate_completion_gap_schedule(schedule): + """Normalize an opt-in optimizer-step curriculum; keep legacy identity absent.""" + if schedule is None: + return None + result = [] + for item in schedule: + if set(item) != {"start_step", "max_gap_seconds"}: + raise ValueError("Completion gap stages require start_step/max_gap_seconds") + step, gap = item["start_step"], item["max_gap_seconds"] + if ( + type(step) is not int + or step < 0 + or isinstance(gap, (bool, np.bool_)) + or not isinstance(gap, (int, float, np.integer, np.floating)) + or not np.isfinite(gap) + or gap <= 0 + ): + raise ValueError( + "Completion gap stages require nonnegative steps/positive seconds" + ) + if result and ( + step <= result[-1]["start_step"] or gap < result[-1]["max_gap_seconds"] + ): + raise ValueError( + "Completion gap stage steps must increase and caps cannot decrease" + ) + result.append({"start_step": step, "max_gap_seconds": float(gap)}) + if not result or result[0]["start_step"] != 0: + raise ValueError("Completion gap curriculum must start at optimizer step zero") + return result + + +def completion_gap_cap(schedule, optimizer_step): + if type(optimizer_step) is not int or optimizer_step < 0: + raise ValueError("Completion gap requires a nonnegative optimizer step") + return next( + stage["max_gap_seconds"] + for stage in reversed(schedule) + if optimizer_step >= stage["start_step"] + ) + + +def sample_gap_bounded_times(timestamps, rng, count, anchor, max_gap_seconds): + """Grow a recorded-frame set around the retained anchor, with bounded gaps. + + The next draw is uniform among unselected frames inside the selected span + or within the cap of either endpoint. Inserting an interior frame preserves + the bound. This is local growth, not a uniform draw over all feasible sets. + Counts, the anchor, actual timestamps, and episode identity are preserved; + no interpolated frames or independent whole-window bound are introduced. + """ + raw = np.asarray(timestamps, dtype=np.float64) + times = _physical_times(raw, len(raw)) + if ( + type(count) is not int + or not 2 <= count <= len(times) + or type(anchor) is not int + or not 0 <= anchor < len(times) + or isinstance(max_gap_seconds, (bool, np.bool_)) + or not np.isfinite(max_gap_seconds) + or max_gap_seconds <= 0 + ): + raise ValueError("Invalid completion gap sample") + # Account only for representational error in large absolute timestamps. + tolerance = max(1e-10, 4 * abs(float(np.spacing(np.max(np.abs(raw)))))) + disconnected = np.flatnonzero(np.diff(times) > max_gap_seconds + tolerance) + lo = max((int(i) + 1 for i in disconnected if i < anchor), default=0) + hi = min((int(i) + 1 for i in disconnected if i >= anchor), default=len(times)) + if hi - lo < count: + raise ValueError("Too few connected recorded frames for completion gap/count") + selected = np.zeros(len(times), bool) + selected[anchor] = True + left = right = anchor + for _ in range(count - 1): + first = max( + lo, + int( + np.searchsorted( + times, times[left] - max_gap_seconds - tolerance, side="left" + ) + ), + ) + last = min( + hi, + int( + np.searchsorted( + times, times[right] + max_gap_seconds + tolerance, side="right" + ) + ), + ) + candidates = np.flatnonzero(~selected[first:last]) + first + if not len(candidates): + raise ValueError("No feasible next recorded completion frame") + chosen = int(rng.choice(candidates)) + selected[chosen] = True + left, right = min(left, chosen), max(right, chosen) + result = np.flatnonzero(selected).astype(np.int64) + if np.any(np.diff(times[result]) > max_gap_seconds + tolerance): + raise ValueError("Completion frame gap invariant failed") + return result + + +def sample_pixel_observation_roles(task, f, rng, forecasting_probability): + """Historical task/mask RNG contract, shared with metadata-only checks.""" + if task == "direct_action": + return "direct_action", { + name: np.ones(f, bool) for name in ("image", "proprio") + } + if rng.random() < forecasting_probability: + known = np.arange(f) < int(rng.integers(1, f)) + return "forecast", { + "image": known.copy(), + "proprio": known & (rng.random(f) < 0.8), + } + observed = {name: rng.random(f) < 0.5 for name in ("image", "proprio")} + if not observed["image"].any(): + observed["image"][rng.integers(f)] = True + if observed["image"].all(): + observed["image"][rng.integers(f)] = False + return "completion", observed + + def sample_discrete_window(length, rng, frames, stride, anchor): """A complete raw-frame window containing anchor, grouped without overlap. @@ -338,12 +470,15 @@ def patchify_rgb(rgb, patch_size): channels, ) patches = patches.transpose(0, 1, 2, 4, 3, 5, 6).reshape( - frames, cameras, height // patch_size * (width // patch_size), 3 * patch_size**2 + frames, + cameras, + height // patch_size * (width // patch_size), + 3 * patch_size**2, ) return patches.astype(np.float32) / np.float32(127.5) - np.float32(1) -def load_pose_normalization(manifest, root, verifier): +def load_pose_normalization(manifest, root, verifier, camera_pose=False): contract = manifest.get("normalization", {}) if ( contract.get("scope") != "train_only" @@ -358,6 +493,8 @@ def load_pose_normalization(manifest, root, verifier): verifier.verify(path, contract["sha256"]) saved = json.loads(path.read_text()) required = {f"pose/{emb}" for emb in manifest["pose_schemas"]} + if camera_pose: + required |= {f"camera/{emb}" for emb in manifest["camera_pose_schemas"]} if ( saved.get("schema") != "pixel-proprio-normalization-v1" or saved.get("scope") != "train_only" @@ -373,7 +510,9 @@ def load_pose_normalization(manifest, root, verifier): np.asarray(item["mean"], np.float32), np.asarray(item["std"], np.float32), ) - dim = manifest["pose_schemas"][key[5:]]["dimension"] + kind, embodiment = key.split("/", 1) + schemas = manifest["pose_schemas"] if kind == "pose" else manifest["camera_pose_schemas"] + dim = schemas[embodiment]["dimension"] if ( mean.shape != (dim,) or std.shape != (dim,) @@ -412,7 +551,18 @@ def __init__( sampling_strategy="full_episode_uniform", image_only=False, couple_state_observations=False, + completion_gap_schedule=None, + training_conditioning="mixed_context", + temporal_gap_schedule=None, + language_features_path=None, + language_features_sha256=None, + subtask_manifest_path=None, + subtask_manifest_sha256=None, + camera_pose=False, + action_embodiments=None, + boundary_sampling=None, ): + validate_language_feature_options(language_features_path, language_features_sha256) if task not in ( "visual_diffusion", "visual_state_action", @@ -427,6 +577,15 @@ def __init__( if type(image_only) is not bool or (image_only and task != "visual_diffusion"): raise ValueError("image_only requires visual_diffusion and a boolean flag") self.image_only = image_only + if type(camera_pose) is not bool or (camera_pose and image_only): + raise ValueError("Camera pose requires explicit image/proprio data") + self.camera_pose = camera_pose + if action_embodiments is not None and ( + not isinstance(action_embodiments, (list, tuple)) or not action_embodiments + or any(not isinstance(x, str) or not x for x in action_embodiments) + or len(set(action_embodiments)) != len(action_embodiments)): + raise ValueError("Require explicit unique action-bearing embodiment IDs") + self.action_embodiments = action_embodiments if type(couple_state_observations) is not bool or ( couple_state_observations and image_only ): @@ -441,10 +600,43 @@ def __init__( seed, mode, ) - self.manifest = json.loads(self.path.read_text()) + manifest_bytes = self.path.read_bytes() + self.manifest = json.loads(manifest_bytes) manifest = self.manifest if manifest.get("schema") != "pixel-proprio-prepared-v1": raise ValueError("Unsupported pixel/proprio prepared schema") + manifest_sha256 = hashlib.sha256(self.path.read_bytes()).hexdigest() + self.subtasks = load_subtask_manifest( + subtask_manifest_path, subtask_manifest_sha256, self.manifest, manifest_sha256 + ) + language_manifest = self.manifest + if self.subtasks is not None: + if language_features_path is None: + raise ValueError("Subtask-only sampling requires language conditioning") + language_manifest = dict(self.manifest, episodes=[ + dict(ep, task_id=unit["task_id"]) + for ep in self.manifest["episodes"] + for unit in self.subtasks[(ep["embodiment"], ep["episode_id"])]] ) + self.language_features = load_pixel_task_language_features( + language_features_path, language_features_sha256, + language_manifest, manifest_sha256, + ) + if self.subtasks is not None: + for units in self.subtasks.values(): + for unit in units: + if self.language_features.instructions[unit["task_id"]] != unit["instruction"]: + raise ValueError("Subtask instruction differs from bound language feature") + if camera_pose: + schemas = manifest.get("camera_pose_schemas", {}) + if set(schemas) != {ep["embodiment"] for ep in manifest["episodes"]}: + raise ValueError("Every embodiment requires an explicit camera pose schema") + for schema in schemas.values(): + _coordinate_schema(schema, "Camera pose schema") + if schema.get("dimension") != 9 or schema.get("representation") != "xyz_rotation6d": + raise ValueError("Camera pose requires native 9D xyz+rotation6D") + if action_embodiments is not None and not set(action_embodiments) <= { + ep["embodiment"] for ep in manifest["episodes"]}: + raise ValueError("Unknown action-bearing embodiment ID") identity = manifest.get("pixel_identity", {}) size, patch = identity.get("image_size"), identity.get("patch_size") if ( @@ -507,8 +699,12 @@ def __init__( paths.add(path) if ( ep["embodiment"] not in manifest["pose_schemas"] - or ep.get("frame") != "fixed_reference" - or ep.get("pose_kind") != "measured_pose" + or not (ep.get("frame") == "fixed_reference" or ( + ep["embodiment"] == "human_bimanual" + and ep.get("frame") == "per_frame_rgb_camera" + and ep.get("pose_kind") == "measured_keypoints" + and manifest["pose_schemas"][ep["embodiment"]].get("coordinate_reference") == "each_frame_RGBcamera")) + or ep.get("pose_kind") not in ("measured_pose", "measured_keypoints") or not _sha(ep.get("sha256")) ): raise ValueError("Require exact measured-pose identity and frame") @@ -522,12 +718,12 @@ def __init__( ep["rgb_sha256"], trusted_fingerprint=ep.get("rgb_verified_fingerprint"), ) - if task in ("direct_action", "visual_state_action"): + if self._has_actions(ep): self._validate_actions(ep) self.normalization = ( {} if image_only - else load_pose_normalization(manifest, self.path.parent, self.verifier) + else load_pose_normalization(manifest, self.path.parent, self.verifier, camera_pose) ) self.episodes = [ep for ep in manifest["episodes"] if ep["split"] == split] if not self.episodes: @@ -565,33 +761,76 @@ def __init__( raise ValueError( "Direct action baseline requires single-frame continuous observations" ) + if training_conditioning not in ("mixed_context", "independent_frame_noise"): + raise ValueError("Unknown training conditioning profile") + if training_conditioning == "independent_frame_noise" and ( + split != "train" or task == "direct_action" or completion_gap_schedule is not None + ): + raise ValueError("Independent frame noise requires training without completion masks") + self.training_conditioning = training_conditioning + self.temporal_gap_schedule = validate_completion_gap_schedule(temporal_gap_schedule) + if self.temporal_gap_schedule is not None and ( + training_conditioning != "independent_frame_noise" + or mode != "continuous" or sampling_strategy != "full_episode_uniform" + ): + raise ValueError("Temporal gap schedule requires independent-noise continuous training") + self.completion_gap_schedule = validate_completion_gap_schedule( + completion_gap_schedule + ) + if self.completion_gap_schedule is not None and ( + split != "train" + or mode != "continuous" + or sampling_strategy != "full_episode_uniform" + or task == "direct_action" + ): + raise ValueError( + "Completion gap curriculum requires continuous full-episode training" + ) + self.boundary_sampling = validate_boundary_sampling(boundary_sampling) + if self.boundary_sampling is not None: + if (split != "train" or mode != "continuous" + or sampling_strategy != "full_episode_uniform" + or training_conditioning != "independent_frame_noise" + or task == "direct_action" or self.subtasks is None + or self.temporal_gap_schedule is not None + or self.completion_gap_schedule is not None or max_frames > 64): + raise ValueError("Boundary uncertainty requires reviewed-subtask independent-noise continuous K2..64 training only") + # Small physical-clock arrays only; prepared artifacts are already + # immutable/hash verified. No image/state cache or global RNG state. + self._boundary_time_cache = {} self.forecasting_probability = forecasting_probability self.sampling = dict(min_frames=min_frames, max_frames=max_frames) self.sampling_strategy = sampling_strategy self.discrete_frames, self.stride = discrete_frames, stride self.chunk_size = stride if mode == "discrete" else 1 - self.anchor_indices = [] + self.anchor_indices, self.sample_units = [], [] + minimum = (discrete_frames * stride if sampling_strategy == "discrete_window" + else min_frames) if task != "direct_action" else 1 + self.excluded_short_subtasks = [] for ep in self.episodes: - if task == "direct_action": - anchors = np.arange(ep["frames"]) - else: - anchors = ( - np.arange(ep["frames"]) - if ep["frames"] - >= ( - discrete_frames * stride - if sampling_strategy == "discrete_window" - else min_frames - ) - else np.array([], np.int64) - ) - self.anchor_indices.append(anchors) + units = ([dict(start_idx=0, end_idx=ep["frames"], task_id=ep.get("task_id"))] + if self.subtasks is None else + self.subtasks[(ep["embodiment"], ep["episode_id"])]) + for unit in units: + length = unit["end_idx"] - unit["start_idx"] + if length < minimum: + if self.subtasks is None: + raise ValueError("Episode too short for temporal recipe") + self.excluded_short_subtasks.append(dict(episode_id=ep["episode_id"], **unit)) + continue + self.sample_units.append((ep, unit)) + self.anchor_indices.append(np.arange(length)) + self.sampling_episodes = [ep for ep, _ in self.sample_units] self.counts = [len(indices) for indices in self.anchor_indices] - if any(c < 1 for c in self.counts): - raise ValueError("Episode too short for temporal recipe") + if not self.counts: + raise ValueError("No subtasks/episodes long enough for temporal recipe") self.ends = np.cumsum(self.counts) self.rng = self.worker_seed = None + def _has_actions(self, ep): + return self.task in ("direct_action", "visual_state_action") and ( + self.action_embodiments is None or ep["embodiment"] in self.action_embodiments) + def _validate_actions(self, ep): schema = self.manifest.get("action_schemas", {}).get(ep["embodiment"]) provenance, source = ( @@ -678,12 +917,17 @@ def _validate_actions(self, ep): def __len__(self): return int(self.ends[-1]) - def _episode_anchor(self, index): + def _unit_anchor(self, index): if not 0 <= index < len(self): raise IndexError(index) ei = int(np.searchsorted(self.ends, index, side="right")) offset = 0 if ei == 0 else int(self.ends[ei - 1]) - return self.episodes[ei], int(self.anchor_indices[ei][index - offset]) + ep, unit = self.sample_units[ei] + return ep, unit, int(self.anchor_indices[ei][index - offset]) + + def _episode_anchor(self, index): + ep, unit, anchor = self._unit_anchor(index) + return ep, unit["start_idx"] + anchor def _sample_frame_chunks(self, length, anchor, rng): """One RNG contract shared by actual data reads and metadata-only work preview.""" @@ -698,10 +942,35 @@ def _sample_frame_chunks(self, length, anchor, rng): :, None ] + def _select_frame_chunks(self, episode, unit, anchor, rng, draw, full_time=None): + """One selection for actual reads and checkpointed work preview. + + Disabled behavior calls only the exact legacy sampler. The opt-in may + read/cache the small immutable physical clock, never RGB or pose data. + """ + chunks = self._sample_frame_chunks(unit["end_idx"] - unit["start_idx"], anchor, rng) + unit["start_idx"] + if self.boundary_sampling is None: + return chunks, None + if full_time is None: + key = (episode["embodiment"], episode["episode_id"]) + if key not in self._boundary_time_cache: + with open_prepared_episode(self.path.parent / episode["rgb_path"]) as data: + clock = np.array(data["time"][:], dtype=np.float64, copy=True) + _physical_times(clock, episode["frames"]) + clock.setflags(write=False) + self._boundary_time_cache[key] = clock + full_time = self._boundary_time_cache[key] + return sample_boundary_pair( + chunks, full_time=full_time, episode=episode, unit=unit, + all_units=self.subtasks[(episode["embodiment"], episode["episode_id"])], + seed=self.seed, draw=draw, policy=self.boundary_sampling, + ) + def predicted_token_work(self, index, draw): """Pure preview of image-attention work for the exact checkpointed draw. - The preview reads manifest metadata only and never advances a worker's + The disabled preview reads manifest metadata only. Explicit boundary + sampling may cache the small physical clock; neither advances a worker's RNG. Squared image-token count estimates attention cost; it is not a measured GPU time or a reason to change the selected frames. """ @@ -711,9 +980,10 @@ def predicted_token_work(self, index, draw): or any(type(value) is not int or value < 0 for value in draw) ): raise ValueError("Work preview requires a training cursor/global-slot draw") - episode, anchor = self._episode_anchor(index) + episode, unit, anchor = self._unit_anchor(index) rng = np.random.default_rng(np.random.SeedSequence([self.seed, *draw])) - frames = len(self._sample_frame_chunks(episode["frames"], anchor, rng)) + chunks, _ = self._select_frame_chunks(episode, unit, anchor, rng, tuple(draw)) + frames = len(chunks) return (frames * len(episode["cameras"]) * self.grid[0] * self.grid[1]) ** 2 def __getitem__(self, index): @@ -726,7 +996,7 @@ def __getitem__(self, index): if isinstance(index, tuple): index, cursor, slot = index draw = int(cursor), int(slot) - ep, anchor = self._episode_anchor(index) + ep, unit, anchor = self._unit_anchor(index) worker_seed = torch.initial_seed() if self.rng is None or self.worker_seed != worker_seed: self.rng = np.random.default_rng((worker_seed + self.seed) % 2**63) @@ -767,11 +1037,44 @@ def __getitem__(self, index): raise ValueError( "Native RGB shape/dtype differs from declared pixel geometry/cameras" ) - chunk_indices = self._sample_frame_chunks(ep["frames"], anchor, rng) + start, end = unit["start_idx"], unit["end_idx"] + chunk_indices, boundary_metadata = self._select_frame_chunks(ep, unit, anchor, rng, draw, full_time) indices = chunk_indices[:, 0] + f = len(indices) + origin = 0 if self.task == "direct_action" else int(rng.integers(f)) + # Reads/normalization below consume no RNG. Choosing roles here keeps + # the historical frame/origin/task/mask draw order exactly intact. + if self.training_conditioning == "independent_frame_noise": + # No clean-context role selection: the noiser draws a fresh + # independent strength at every valid timestamp. Every stream + # remains a supervised denoising target, including past actions. + task_kind = "diffusion_forcing" + observed = {name: np.zeros(f, bool) for name in ("image", "proprio")} + else: + task_kind, observed = sample_pixel_observation_roles( + self.task, f, rng, self.forecasting_probability + ) + gap_cap = None + gap_schedule = self.temporal_gap_schedule or self.completion_gap_schedule + if gap_schedule is not None: + if draw is None: + raise ValueError( + "Completion gap curriculum requires checkpointed draw keys" + ) + if self.temporal_gap_schedule is not None or task_kind == "completion": + # Cursor is zero-based and equals the next optimizer step + # when each global batch performs exactly one update. + gap_cap = completion_gap_cap(gap_schedule, draw[0]) + gap_rng = np.random.default_rng( + np.random.SeedSequence([self.seed, *draw, 0x474150]) + ) + indices = sample_gap_bounded_times( + full_time[start:end], gap_rng, f, anchor, gap_cap + ) + indices = indices + start + chunk_indices = indices[:, None] raw_indices = chunk_indices.reshape(-1) time = np.asarray(full_time[indices], np.float64) - origin = 0 if self.task == "direct_action" else int(rng.integers(len(time))) physical_times = time.tolist() physical_origin = float(time[origin]) time = (time - time[origin]).astype(np.float32) @@ -803,27 +1106,6 @@ def __getitem__(self, index): pose[~pose_valid] = 0 pose = pose.reshape(f, self.chunk_size * dim) pose_valid = pose_valid.reshape(f, self.chunk_size).all(axis=1) - # Preserve the established RNG consumption even when proprio is - # omitted: completion's image-mask repairs and later draws must - # select identical RGB frames/masks across conditioning variants. - if self.task == "direct_action": - task_kind = "direct_action" - observed = {name: np.ones(f, bool) for name in ("image", "proprio")} - elif rng.random() < self.forecasting_probability: - task_kind = "forecast" - known = np.arange(f) < int(rng.integers(1, f)) - observed = { - "image": known.copy(), - "proprio": known & (rng.random(f) < 0.8), - } - else: - task_kind = "completion" - observed = {name: rng.random(f) < 0.5 for name in ("image", "proprio")} - if not observed["image"].any(): - observed["image"][rng.integers(f)] = True - if observed["image"].all(): - observed["image"][rng.integers(f)] = False - if self.couple_state_observations: if not pose_valid.all(): raise ValueError( @@ -851,7 +1133,29 @@ def base(value, repeats, known): streams["proprio"] = base(pose, 1, observed["proprio"]) streams["proprio"]["target_valid"] = pose_valid streams["proprio"]["observed"] &= pose_valid - if self.task in ("direct_action", "visual_state_action"): + if self.camera_pose: + if ("camera_world_pose" not in data or "camera_pose_valid" not in data + or data["camera_world_pose"].shape != (ep["frames"], cameras, 4, 4) + or data["camera_pose_valid"].shape != (ep["frames"], cameras)): + raise ValueError("Require per-frame/per-camera calibrated relative pose arrays") + raw_camera = np.asarray(data["camera_world_pose"][raw_indices], np.float64) + cv = np.asarray(data["camera_pose_valid"][raw_indices]) + cp, cv = relative_camera_pose(raw_camera, cv, 0) + if not cv.any(): + raise ValueError("Selected window has no measured camera reference at its first frame") + if cv.dtype != np.bool_ or not np.isfinite(cp[cv]).all(): + raise ValueError("Invalid calibrated camera pose values/validity") + mean, std = self.normalization[f"camera/{ep['embodiment']}"] + cp = (cp - mean) / std + cp[~cv] = 0 + cp = cp.reshape(f, self.chunk_size, cameras, 9).transpose(0,2,1,3) + cv = cv.reshape(f,self.chunk_size,cameras).all(axis=1).reshape(-1) + streams["camera"] = base(cp, cameras, observed["image"]) + streams["camera"].update( + camera=np.tile(np.arange(cameras), f), valid=cv, + target_valid=cv.copy(), observed=np.repeat(observed["image"],cameras)&cv, + ) + if self._has_actions(ep): action = np.asarray(data["action"][raw_indices], np.float32).reshape( f, -1 ) @@ -890,6 +1194,34 @@ def base(value, repeats, known): proprio_frame_indices=chunk_indices.tolist(), proprio_timestamps_seconds=full_time[chunk_indices].tolist(), ) + if self.language_features is not None: + sample.update(self.language_features.sample(unit["task_id"])) + if self.camera_pose: + sample["pixel_metadata"]["camera_reference_frame_index"] = int(raw_indices[0]) + sample["pixel_metadata"]["camera_reference_policy"] = "first_selected_native_frame_per_camera" + if self.subtasks is not None: + sample["pixel_metadata"]["subtask"] = dict( + start_idx=start, end_idx=end, task_id=unit["task_id"], + instruction=unit["instruction"], + annotation_provenance=unit["annotation_provenance"], + interval_semantics="half_open", + ) + if boundary_metadata is not None: + sample["pixel_metadata"]["boundary_sampling"] = boundary_metadata + if "action" in streams: + sample["pixel_metadata"]["action_frame_indices"] = chunk_indices.tolist() + if "camera" in streams: + sample["pixel_metadata"]["camera_frame_indices"] = chunk_indices.tolist() + if self.training_conditioning == "independent_frame_noise": + sample["pixel_metadata"]["training_conditioning"] = self.training_conditioning + if gap_cap is not None: + key = "temporal_gap_curriculum" if self.temporal_gap_schedule is not None else "completion_gap_curriculum" + sample["pixel_metadata"][key] = { + "optimizer_step": draw[0], + "max_gap_seconds": gap_cap, + "largest_selected_gap_seconds": float(np.diff(physical_times).max()), + "sampling": "anchor_retaining_local_growth_v1", + } if draw is not None: sample["sampling_cursor"] = draw[0] + 1 if validation_batch_index is not None: diff --git a/egomimic/rldb/pixel_proprio_datamodule.py b/egomimic/rldb/pixel_proprio_datamodule.py index c1e1d19c9..47ee36a5b 100644 --- a/egomimic/rldb/pixel_proprio_datamodule.py +++ b/egomimic/rldb/pixel_proprio_datamodule.py @@ -6,11 +6,14 @@ from torch.utils.data import DataLoader from egomimic.pl_utils.data_context import DataContext +from egomimic.rldb.boundary_sampling import validate_boundary_sampling +from egomimic.rldb.pixel_language import validate_language_feature_options from egomimic.rldb.pixel_proprio import ( DistributedPixelTrainingBatchSampler, DistributedPixelValidationBatchSampler, EvenlySpacedEpisodeBatchSampler, PreparedPixelProprioDataset, + validate_completion_gap_schedule, collate_pixel_proprio, ) from egomimic.rldb.world_model import HomogeneousWorldBatchSampler @@ -49,7 +52,22 @@ def __init__( prefetch_factor=None, balance_training_work=False, couple_state_observations=False, + completion_gap_schedule=None, + training_conditioning="mixed_context", + temporal_gap_schedule=None, + language_features_path=None, + language_features_sha256=None, + subtask_manifest_path=None, + subtask_manifest_sha256=None, + camera_pose=False, + action_embodiments=None, + boundary_sampling=None, ): + validate_language_feature_options(language_features_path, language_features_sha256) + self.language_features_path = language_features_path + self.subtask_manifest_path = subtask_manifest_path + if (subtask_manifest_path is None) != (subtask_manifest_sha256 is None): + raise ValueError("Subtask path and SHA256 require both or neither") if type(batch_size) is not int or batch_size < 1: raise ValueError("batch_size must be a positive integer") if type(num_workers) is not int or num_workers < 0: @@ -127,6 +145,47 @@ def __init__( batch_size if validation_batch_size is None else validation_batch_size ) self.settings.update(task=task, sampling_strategy=sampling_strategy) + if camera_pose: + self.settings["camera_pose"] = camera_pose + if action_embodiments is not None: + self.settings["action_embodiments"] = action_embodiments + if language_features_path is not None: + self.settings["language_features_sha256"] = language_features_sha256 + if subtask_manifest_path is not None: + self.settings["subtask_manifest_sha256"] = subtask_manifest_sha256 + if training_conditioning not in ("mixed_context", "independent_frame_noise"): + raise ValueError("Unknown training conditioning profile") + if training_conditioning == "independent_frame_noise": + if task == "direct_action" or completion_gap_schedule is not None: + raise ValueError("Independent frame noise excludes clean-context completion masks") + self.settings["training_conditioning"] = training_conditioning + temporal_schedule = validate_completion_gap_schedule(temporal_gap_schedule) + if temporal_schedule is not None: + if (training_conditioning != "independent_frame_noise" or not checkpointable_sampling + or mode != "continuous" or sampling_strategy != "full_episode_uniform"): + raise ValueError("Temporal gap schedule requires checkpointable independent-noise continuous sampling") + self.settings["temporal_gap_schedule"] = temporal_schedule + gap_schedule = validate_completion_gap_schedule(completion_gap_schedule) + if gap_schedule is not None: + if ( + not checkpointable_sampling + or mode != "continuous" + or sampling_strategy != "full_episode_uniform" + ): + raise ValueError( + "Completion gap curriculum requires checkpointable continuous sampling" + ) + self.settings["completion_gap_schedule"] = gap_schedule + boundary_policy = validate_boundary_sampling(boundary_sampling) + if boundary_policy is not None: + if (not checkpointable_sampling or mode != "continuous" + or sampling_strategy != "full_episode_uniform" + or training_conditioning != "independent_frame_noise" + or task == "direct_action" or subtask_manifest_path is None + or temporal_schedule is not None or gap_schedule is not None + or not 2 <= min_frames <= max_frames <= 64): + raise ValueError("Boundary uncertainty requires checkpointed reviewed-subtask independent-noise continuous K2..64 training") + self.settings["boundary_sampling"] = boundary_policy # Preserve historical conditioned checkpoint contexts when disabled. if image_only: self.settings["image_only"] = True @@ -147,39 +206,17 @@ def __init__( if value is not None } - def preflight_configuration(self): - """Check both declared pixel samplers without reading a prepared cache.""" - result = super().preflight_configuration() - for split, settings in ( - ("train", self.settings), - ("val", self.settings | self.validation_settings), - ): - strategy = settings["sampling_strategy"] - if ( - (strategy, settings["mode"]) - not in ( - ("full_episode_uniform", "continuous"), - ("discrete_window", "continuous"), - ("discrete_window", "discrete"), - ) - or settings.get("max_past_seconds") is not None - or settings.get("max_future_seconds") is not None - ): - raise ValueError(f"Invalid {split} pixel sampling strategy/time bounds") - for key, minimum in (("discrete_frames", 2), ("stride", 1)): - value = settings[key] - if type(value) is not int or value < minimum: - raise ValueError(f"{split} {key} must be an integer >= {minimum}") - if ( - settings["task"] == "direct_action" - and strategy != "full_episode_uniform" - ): - raise ValueError("Direct action requires full_episode_uniform sampling") - return result | { - "validation_sampling": dict(self.validation_settings), - "distributed_world_size": self.distributed_world_size, - "validation_batch_size": self.validation_batch_size, - } + if gap_schedule is not None: + self.validation_settings["completion_gap_schedule"] = None + if training_conditioning == "independent_frame_noise": + # Fixed held-out forecasting/completion metrics remain comparable; + # these masks are evaluation queries, never the training recipe. + self.validation_settings["training_conditioning"] = "mixed_context" + if temporal_schedule is not None: + self.validation_settings["temporal_gap_schedule"] = None + if boundary_policy is not None: + # Held-out anchor selection and query masks stay exactly comparable. + self.validation_settings["boundary_sampling"] = None def setup(self, stage=None): if not self.datasets: @@ -189,6 +226,8 @@ def setup(self, stage=None): self.manifest_path, split=split, verification_memo=verifier, + language_features_path=self.language_features_path, + subtask_manifest_path=self.subtask_manifest_path, **( self.settings | (self.validation_settings if split == "val" else {}) @@ -198,12 +237,15 @@ def setup(self, stage=None): } def configure_evaluation(self, requirements): - if requirements.complete_episodes or set(requirements.required_keys) - { + allowed = { "world", "embodiment", "task_kind", "pixel_metadata", - }: + } + if self.language_features_path is not None: + allowed.update(("language_features", "language_instruction", "task_id")) + if requirements.complete_episodes or set(requirements.required_keys) - allowed: raise ValueError("Unsupported pixel/proprio evaluation requirements") def prepare_context(self, *, mode, normalization, normalizer, restored_state=None): @@ -223,6 +265,8 @@ def prepare_context(self, *, mode, normalization, normalizer, restored_state=Non sampling=self.settings, validation_overrides=self.validation_settings, ) + if dataset.language_features is not None: + state["language_features"] = dataset.language_features.context_identity() if self.checkpointable_sampling: state["sampling_resume"] = { "schema": "mot-sampling-state-v1", @@ -261,6 +305,16 @@ def prepare_context(self, *, mode, normalization, normalizer, restored_state=Non ) def _loader(self, split): + if ( + split == "train" + and (self.settings.get("completion_gap_schedule") is not None + or self.settings.get("temporal_gap_schedule") is not None) + and self.trainer is not None + and self.trainer.accumulate_grad_batches != 1 + ): + raise ValueError( + "Completion gap optimizer-step schedule requires one update per global batch" + ) self.setup() rank = 0 if self.trainer is not None: @@ -328,9 +382,9 @@ def _loader(self, split): def state_dict(self): state = super().state_dict() if state: - state["validation_anchors_per_episode"] = ( - self.validation_anchors_per_episode - ) + state[ + "validation_anchors_per_episode" + ] = self.validation_anchors_per_episode state["validation_overrides"] = self.validation_settings if self.distributed_world_size > 1: state["distributed_sampling"] = self._distributed_sampling_identity() diff --git a/egomimic/rldb/pixel_proprio_preparation.py b/egomimic/rldb/pixel_proprio_preparation.py new file mode 100644 index 000000000..6cc3c3845 --- /dev/null +++ b/egomimic/rldb/pixel_proprio_preparation.py @@ -0,0 +1,433 @@ +"""Prepare native EVA/human Zarr recordings for the raw-pixel world model. + +Inputs explicitly select physical episodes and episode-level splits. S3 reads +reuse the native downloader; outputs are new immutable files. No annotation +inference, training, resampling, image flips, or source writes occur here. +""" +from __future__ import annotations +import argparse +import copy +import os +import shutil +import hashlib +import json +from pathlib import Path +from urllib.parse import urlsplit +import h5py +import numpy as np +from egomimic.rldb.camera_pose import native_camera_world_poses, relative_camera_pose +from egomimic.rldb.world_model_preparation import StreamingMoments, training_identities, ArtifactVerificationMemo +from egomimic.utils.pose_utils import _xyzwxyz_to_matrix, _matrix_to_xyzrot6d + + +def sha(path): + digest = hashlib.sha256() + with Path(path).open('rb') as source: + for part in iter(lambda: source.read(8 * 1024 * 1024), b''): + digest.update(part) + return digest.hexdigest() + + +def immutable_json(path, value): + with Path(path).open('x') as out: + json.dump(value, out, indent=2, sort_keys=True) + out.write('\n') + + +def _array(group, name, frames): + values = np.asarray(group[name][:frames]) + if len(values) != frames: + raise ValueError(f'Incomplete native field {name}') + return values + + +def native_numeric_arrays(group, cameras, *, maximum_source_gap_s=0.1): + attrs = dict(group.attrs) + n = int(attrs['total_frames']) + embodiment = attrs['embodiment'] + if n < 1: + raise ValueError('Empty native episode') + if embodiment == 'eva_bimanual': + poses, actions = [], [] + for arm in ('left', 'right'): + # Actual native EVA consumer takes inverse(stored calibration). + extrinsic = np.asarray(attrs.get('extrinsics', {}).get(arm), np.float64) + if extrinsic.shape != (4,4): + raise ValueError('Measured per-arm front camera calibration required') + for role, target in [('obs', poses), ('cmd', actions)]: + raw = _array(group, f'{arm}.{role}_ee_pose', n) + matrix = np.linalg.inv(extrinsic) @ _xyzwxyz_to_matrix(raw) + gripper = _array(group, f'{arm}.{role}_gripper', n).reshape(n,1) + target.append(np.concatenate([_matrix_to_xyzrot6d(matrix), gripper], -1)) + pose, action = np.concatenate(poses,-1), np.concatenate(actions,-1) + fps = float(attrs['fps']) + if not np.isfinite(fps) or fps <= 0: + raise ValueError('Positive native frame rate required') + times = np.arange(n, dtype=np.float64) / fps + clock = 'native synchronized EVA row index / recorded fps; no invented acquisition timestamps' + head = None + pose_valid = np.isfinite(pose).all(axis=1) & (np.abs(pose).max(axis=1) < 1000) + if not pose_valid.all(): + raise ValueError('Invalid recorded robot Cartesian state') + elif embodiment == 'human_bimanual': + pose = np.concatenate([_array(group, f'{arm}.obs_keypoints', n).reshape(n,63) + for arm in ('left','right')], -1) + action = None + hands = pose.reshape(n,2,63) + pose_valid = (np.isfinite(hands).all(axis=2) & (np.abs(hands).max(axis=2) < 1000) + & (np.abs(hands).sum(axis=2) > 1e-9)).all(axis=1) + pose = np.where(pose_valid[:,None],pose,0) + timestamps = _array(group, 'obs_rgb_timestamps_ns', n).reshape(n) + times = (timestamps - timestamps[0]).astype(np.float64) * 1e-9 + clock = 'recorded obs_rgb_timestamps_ns minus first physical timestamp' + head = _array(group, 'obs_head_pose', n) + else: + raise ValueError(f'Unsupported native embodiment: {embodiment}') + if not pose_valid.any(): + raise ValueError('No valid measured proprioception in selected episode') + if not np.isfinite(maximum_source_gap_s) or maximum_source_gap_s <= 0: + raise ValueError('Positive source-gap threshold required') + if action is not None and not np.isfinite(action).all(): + raise ValueError('Nonfinite real native command') + if not np.isfinite(times).all() or (np.diff(times) <= 0).any(): + raise ValueError('Native clock must strictly increase') + cuts = np.flatnonzero(np.diff(times) > maximum_source_gap_s) + 1 + boundaries = np.r_[0,cuts,n] + source_cuts = [[int(a),int(b)] for a,b in zip(boundaries[:-1],boundaries[1:])] + camera, valid = native_camera_world_poses(embodiment, cameras, n, + head_pose=head, extrinsics=attrs.get('extrinsics')) + front_valid = valid[:,cameras.index('front_1')] if 'front_1' in cameras else valid.any(axis=1) + usable = pose_valid & front_valid + if embodiment == 'human_bimanual': + # Native Human keypoint transform: inv(world_from_RGBcamera_t) @ + # world_keypoints_t. Future camera measurements construct targets only; + # the loader still exposes their unknown diffusion roles, never context. + from egomimic.rldb.zarr.action_chunk_transforms import ActionChunkCoordinateFrameTransform + transform = ActionChunkCoordinateFrameTransform('camera','keypoints','camera_keypoints',mode='xyz') + camera_pose = np.zeros_like(pose) + for t in np.flatnonzero(usable): + transformed = transform.transform(dict(camera=head[t],keypoints=pose[t].reshape(42,3))) + camera_pose[t] = transformed['camera_keypoints'].reshape(126) + pose = camera_pose + pose_valid = usable.copy() + continuity = [] + for a,b in source_cuts: + inside = usable[a:b] + starts = np.flatnonzero(inside & ~np.r_[False,inside[:-1]]) + a + ends = np.flatnonzero(inside & ~np.r_[inside[1:],False]) + a + 1 + continuity.extend([[int(lo),int(hi)] for lo,hi in zip(starts,ends)]) + if not continuity: + raise ValueError('No jointly measured state/front-camera rows for sampling') + return dict(embodiment=embodiment, frames=n, time=times, pose=pose.astype(np.float32), + action=None if action is None else action.astype(np.float32), + camera_world_pose=camera, camera_pose_valid=valid, clock=clock, pose_valid=pose_valid, + continuity_segments=continuity,maximum_source_gap_s=maximum_source_gap_s, + source_cut_segments=source_cuts, usable_frame=usable, + quality_exclusions=dict(invalid_pose_frames=int((~pose_valid).sum()), + invalid_front_camera_frames=int((~front_valid).sum()), + excluded_union_frames=int((~usable).sum()), usable_frames=int(usable.sum()))) + + +def prepare(selection_path, destination, *, source_cache, image_size=(160,160), patch_size=16, defer_normalization=False): + import cv2 + import zarr + from egomimic.rldb.zarr.zarr_dataset_multi import decode_jpeg_payload + from egomimic.utils.aws.aws_data_utils import s3_sync_to_local + selection = selection_path if isinstance(selection_path,dict) else json.loads(Path(selection_path).read_text()) + entries = selection['episodes'] + if (len(image_size) != 2 or any(type(v) is not int or v < 1 for v in image_size) + or type(patch_size) is not int or patch_size < 1 or any(v % patch_size for v in image_size)): + raise ValueError('Positive image dimensions divisible by patch size required') + if len({e['source'] for e in entries}) != len(entries): + raise ValueError('One physical source cannot appear in multiple split entries') + if not defer_normalization and {e['embodiment'] for e in entries} != {e['embodiment'] for e in entries if e['split']=='train'}: + raise ValueError('Every prepared embodiment requires training-only normalization data') + identities = [(e['embodiment'],e['episode_id']) for e in entries] + if not entries or len(set(identities)) != len(entries): + raise ValueError('Explicit nonduplicate physical episode selection required') + output = Path(destination) + output.mkdir(parents=True, exist_ok=True) + if (output/'manifest.json').exists(): + raise FileExistsError('Prepared manifest already exists; preserve immutable output') + moments, episodes, pose_schemas, action_schemas, camera_schemas = {}, [], {}, {}, {} + for entry in entries: + if entry['split'] not in ('train','val'): + raise ValueError('Explicit train/val split required') + eid, embodiment = entry['episode_id'], entry['embodiment'] + if Path(eid).name != eid or eid in ('.','..'): + raise ValueError('Episode ID must be safe single component') + source = entry['source'] + if source.startswith('s3://'): + uri = urlsplit(source) + local = Path(source_cache)/eid + s3_sync_to_local(uri.netloc, uri.path.lstrip('/'), local) + else: + local = Path(source) + if not (local/'zarr.json').is_file(): + raise ValueError('Native Zarr v3 metadata missing') + # One bounded native episode inventory binds all original bytes read. + native_files = [{'path':str(f.relative_to(local)), 'bytes':f.stat().st_size, 'sha256':sha(f)} + for f in sorted(local.rglob('*')) if f.is_file()] + native_payload = hashlib.sha256(json.dumps(native_files,sort_keys=True).encode()).hexdigest() + g = zarr.open_group(str(local), mode='r') + cameras = list(entry['cameras']) + arrays = native_numeric_arrays(g, cameras) + if arrays['embodiment'] != embodiment: + raise ValueError('Source embodiment differs from selection') + n = arrays['frames'] + numeric = output/f'{eid}-state.h5' + rgb = output/f'{eid}-rgb.h5' + with h5py.File(numeric,'x') as f: + for key in ('time','pose','camera_world_pose','camera_pose_valid'): + f[key] = arrays[key] + f['pose_valid'] = arrays['pose_valid'] + if arrays['action'] is not None: + f['action'] = arrays['action'] + f['action_valid'] = np.ones(n,bool) + with h5py.File(rgb,'x') as f: + f['time'] = arrays['time'] + frames = f.create_dataset('rgb',(n,len(cameras),*image_size,3),dtype='u1', + chunks=(1,1,*image_size,3),compression='lzf') + for c,name in enumerate(cameras): + native = g[f'images.{name}'] + for t in range(n): + payload = native[t] + if isinstance(payload,np.ndarray) and payload.ndim==3 and payload.dtype==np.uint8: + image = payload if payload.shape[-1]==3 else np.moveaxis(payload,0,-1) + else: + image = np.moveaxis(decode_jpeg_payload(payload),0,-1) + image = np.rint(image*255).astype(np.uint8) + frames[t,c] = cv2.resize(image,(image_size[1],image_size[0]),interpolation=cv2.INTER_AREA) + native_proof = output/f'{eid}-native.json' + immutable_json(native_proof,dict(source=source,files=native_files,clock=arrays['clock'], + image_processing='full_frame_INTER_AREA_no_flip',camera_reference='absolute measured camera matrices; reanchor after window sampling')) + ep=dict(embodiment=embodiment,episode_id=eid,split=entry['split'],frames=n, + path=numeric.name,sha256=sha(numeric),rgb_path=rgb.name,rgb_sha256=sha(rgb), + cameras=cameras,frame='fixed_reference' if arrays['action'] is not None else 'per_frame_rgb_camera', + pose_kind='measured_pose' if arrays['action'] is not None else 'measured_keypoints', + continuity_segments=arrays['continuity_segments'],maximum_source_gap_s=arrays['maximum_source_gap_s'], + source_cut_segments=arrays['source_cut_segments'],quality_exclusions=arrays['quality_exclusions'], + native_source=dict(source_demo_sha256=native_payload,episode=eid, + provenance_path=native_proof.name,provenance_sha256=sha(native_proof))) + dim=arrays['pose'].shape[-1] + pose_schemas[embodiment]=dict(dimension=dim,fields=[f'coordinate_{i}' for i in range(dim)], + representation='bimanual_xyz_rotation6d_gripper_camera' if arrays['action'] is not None else 'bimanual_MANO_per_frame_RGBcamera_keypoints', + coordinate_reference='fixed_front_camera' if arrays['action'] is not None else 'each_frame_RGBcamera', + units='meters') + camera_schemas[embodiment]=dict(dimension=9,fields=['x','y','z','r00','r10','r20','r01','r11','r21'],representation='xyz_rotation6d') + if arrays['action'] is not None: + proof=output/f'{eid}-alignment.json' + immutable_json(proof,dict(schema='pixel-proprio-action-alignment-v1',relation='action_t_applied_to_state_t', + source_demo_sha256=native_payload,source_episode=eid,producer_sha256=sha(__file__), + raw_export_metadata_sha256=sha(local/'zarr.json'),raw_payload_sha256=native_payload, + cache_frame_index=list(range(n)),source_state_index=list(range(n)),source_action_index=list(range(n)), + basis=dict(cache_index_source='Every native Zarr row preserved without temporal resampling', + action_index_source='Native left/right.cmd_ee_pose and cmd_gripper at same row', + state_action_relation_source='EVA extractor and eva_to_zarr preserve synchronized cmd/obs arrays with identical frame indices; command target, not finite-difference action'))) + ep['action_provenance']=dict(path=proof.name,sha256=sha(proof)) + action_schemas[embodiment]=dict(dimension=20,fields=[f'command_{i}' for i in range(20)], + semantics='action_command',value_space='native_commands',representation='native_absolute_cmd_pose_and_gripper_in_fixed_front_camera') + ep['prepared_file_identity']={key:dict(bytes=path.stat().st_size,mtime_ns=path.stat().st_mtime_ns,inode=path.stat().st_ino,uid=path.stat().st_uid) for key,path in [('path',numeric),('rgb_path',rgb)]} + episodes.append(ep) + if entry['split']=='train' and not defer_normalization: + moments.setdefault(f'pose/{embodiment}',StreamingMoments()).update(arrays['pose'][arrays['usable_frame']]) + cam_valid = arrays['camera_pose_valid'] & arrays['usable_frame'][:,None] + refs = [int(np.flatnonzero(cam_valid[:,c])[0]) if cam_valid[:,c].any() else 0 for c in range(cam_valid.shape[1])] + relative, valid = relative_camera_pose(arrays['camera_world_pose'],cam_valid,refs) + moments.setdefault(f'camera/{embodiment}',StreamingMoments()).update(relative[valid]) + if defer_normalization: + normalization=dict(scope='pending_global_train_only',arrays_normalized=False) + else: + norm=output/'normalization.json' + immutable_json(norm,dict(schema='pixel-proprio-normalization-v1',scope='train_only', + train_episodes=training_identities(episodes),statistics={k:v.result(.05) for k,v in moments.items()}, + camera_moment_reference='first usable physical episode frame per measured camera, fixed TRAIN-only statistics over usable rows; samples reanchor to first sampled frame')) + normalization=dict(scope='train_only',arrays_normalized=False,path=norm.name,sha256=sha(norm)) + manifest=dict(schema='pixel-proprio-prepared-v1',episodes=episodes,pose_schemas=pose_schemas, + action_schemas=action_schemas,camera_pose_schemas=camera_schemas, + pixel_identity=dict(image_size=list(image_size),patch_size=patch_size,channels=3,patch_order='row_major_hwc',normalization='uint8_to_minus_one_one'), + normalization=normalization) + immutable_json(output/'manifest.json',manifest) + return manifest + + + +def _native_record_source(ep, root): + proof=ep['native_source'] + path=root/proof['provenance_path'] + if sha(path)!=proof['provenance_sha256']: + raise ValueError('Prepared source provenance digest mismatch') + return json.loads(path.read_text())['source'] + + +def merge_prepared_manifests(selection_path, manifest_paths, destination, *, link_caches=False): + """Publish one train-only normalization across immutable raw episode caches. + + Physical splits come only from the authoritative selection. RGB is referenced, + never copied or decoded. Deferred manifests cannot be loaded before this gate. + """ + selection=json.loads(Path(selection_path).read_text()) + entries=selection['episodes'] + selected={(e['embodiment'],e['episode_id']):e for e in entries} + if len(selected)!=len(entries) or len({e['source'] for e in entries})!=len(entries): + raise ValueError('Duplicate physical episode/source in split selection') + if any(e['split'] not in ('train','val') for e in entries): + raise ValueError('Explicit physical train/val split required') + output=Path(destination);output.mkdir(parents=True,exist_ok=True) + episodes,schemas,pixels={},dict(pose_schemas={},action_schemas={},camera_pose_schemas={}),None + bindings=[] + for manifest_path in manifest_paths: + path=Path(manifest_path).resolve();manifest=json.loads(path.read_text());root=path.parent + if manifest.get('schema')!='pixel-proprio-prepared-v1' or manifest.get('normalization',{}).get('arrays_normalized') is not False: + raise ValueError('Raw prepared manifest required for global normalization') + if pixels is None:pixels=manifest['pixel_identity'] + if pixels!=manifest['pixel_identity']:raise ValueError('Mixed pixel processing identities') + bindings.append(dict(path=str(path),sha256=sha(path))) + for name,group in schemas.items(): + for emb,item in manifest[name].items(): + if emb in group and group[emb]!=item:raise ValueError('Mixed embodiment schema') + group[emb]=item + for original in manifest['episodes']: + key=original['embodiment'],original['episode_id'] + if key not in selected:continue + if key in episodes:raise ValueError('Duplicate prepared physical episode') + entry=selected[key] + if _native_record_source(original,root)!=entry['source']: + raise ValueError('Prepared cache source differs from authoritative selection') + if entry.get('native_metadata_sha256'): + native=json.loads((root/original['native_source']['provenance_path']).read_text()) + metadata=[f['sha256'] for f in native.get('files',[]) if f.get('path')=='zarr.json'] + if metadata!=[entry['native_metadata_sha256']]:raise ValueError('Native root metadata changed after source selection') + if 'frames' in entry and original['frames']!=entry['frames']: + raise ValueError('Prepared source length differs from selected native metadata') + ep=copy.deepcopy(original);ep['split']=entry['split'] + for name in ('path','rgb_path'):ep[name]=str((root/ep[name]).resolve()) + for name in ('action_provenance','native_source'): + field='path' if name=='action_provenance' else 'provenance_path' + if name in ep:ep[name][field]=str((root/ep[name][field]).resolve()) + if not Path(ep['path']).is_file() or not Path(ep['rgb_path']).is_file(): + raise ValueError('Prepared cache is missing') + if link_caches: + linked=output/'episodes'/ep['episode_id'];linked.mkdir(parents=True,exist_ok=True) + def link(original_path,filename): + src=Path(original_path).resolve(strict=True);dst=linked/filename + if Path(original_path).is_symlink():raise ValueError('Prepared source must be a regular immutable file') + if dst.exists(): + if not os.path.samefile(src,dst):raise ValueError('Prepared hardlink destination collision') + else:os.link(src,dst) + return str(dst.relative_to(output)) + ep['path']=link(ep['path'],'state.h5');ep['rgb_path']=link(ep['rgb_path'],'rgb.h5') + ep['native_source']['provenance_path']=link(ep['native_source']['provenance_path'],'native.json') + if 'action_provenance' in ep:ep['action_provenance']['path']=link(ep['action_provenance']['path'],'alignment.json') + episodes[key]=ep + if set(episodes)!=set(selected):raise ValueError('Prepared corpus does not cover selection exactly') + verifier=ArtifactVerificationMemo() + for ep in episodes.values(): + for field,digest,fingerprint in [('path','sha256','verified_fingerprint'),('rgb_path','rgb_sha256','rgb_verified_fingerprint')]: + proof=verifier.verify(output/ep[field],ep[digest],trusted_fingerprint=ep.get(fingerprint)) + if proof is not None:ep[fingerprint]=proof + moments={} + for key,ep in sorted(episodes.items()): + if ep['split']!='train':continue + emb=ep['embodiment'] + with h5py.File(output/ep['path'],'r') as data: + pose_valid=np.asarray(data['pose_valid'],bool) + camera_valid=np.asarray(data['camera_pose_valid'],bool) + usable=pose_valid & camera_valid[:,ep['cameras'].index('front_1')] + moments.setdefault('pose/'+emb,StreamingMoments()).update(np.asarray(data['pose'])[usable]) + camera_valid &= usable[:,None] + refs=[int(np.flatnonzero(camera_valid[:,c])[0]) if camera_valid[:,c].any() else 0 for c in range(camera_valid.shape[1])] + relative,valid=relative_camera_pose(np.asarray(data['camera_world_pose']),camera_valid,refs) + moments.setdefault('camera/'+emb,StreamingMoments()).update(relative[valid]) + required={kind+'/'+emb for kind in ('pose','camera') for emb in schemas['pose_schemas']} + if set(moments)!=required:raise ValueError('Every embodiment requires actual train rows') + output=Path(destination);output.mkdir(parents=True,exist_ok=True) + ordered=[episodes[k] for k in sorted(episodes)] + norm=output/'normalization.json' + immutable_json(norm,dict(schema='pixel-proprio-normalization-v1',scope='train_only',train_episodes=training_identities(ordered),statistics={k:v.result(.05) for k,v in moments.items()},camera_moment_reference='first usable physical episode frame per measured camera; fixed TRAIN-only statistics')) + result=dict(schema='pixel-proprio-prepared-v1',episodes=ordered,**schemas,pixel_identity=pixels,normalization=dict(scope='train_only',arrays_normalized=False,path=norm.name,sha256=sha(norm)),preparation=dict(selection_sha256=sha(selection_path),immutable_input_manifests=bindings,rgb_decoded_or_copied=False,immutable_hardlinks=link_caches)) + immutable_json(output/'manifest.json',result) + return result + + +def _verify_complete_episode(manifest_path, entry): + path=Path(manifest_path);root=path.parent;prior=json.loads(path.read_text()) + if len(prior['episodes'])!=1: + raise ValueError('One physical episode per resumable cache required') + ep=prior['episodes'][0] + if (ep['episode_id']!=entry['episode_id'] or ep['embodiment']!=entry['embodiment'] + or _native_record_source(ep,root)!=entry['source']): + raise ValueError('Existing episode cache identity mismatch') + for key in ('path','rgb_path'): + file=root/ep[key];identity=ep.get('prepared_file_identity',{}).get(key) + if not file.is_file() or file.is_symlink() or not identity: + raise ValueError('Missing verified prepared file identity') + stat=file.stat() + if identity!=dict(bytes=stat.st_size,mtime_ns=stat.st_mtime_ns,inode=stat.st_ino,uid=stat.st_uid) or stat.st_size<1: + raise ValueError('Prepared file identity changed after verification') + if 'action_provenance' in ep and sha(root/ep['action_provenance']['path'])!=ep['action_provenance']['sha256']: + raise ValueError('Prepared action provenance changed') + return prior + + +def prepare_shard(selection_path, destination, *, source_cache, image_size=(160,160), cleanup_source=False): + """Resume completed per-episode caches; remove only marked task-owned copies.""" + selection=json.loads(Path(selection_path).read_text());output=Path(destination);output.mkdir(parents=True,exist_ok=True) + cache=Path(source_cache).resolve();cache.mkdir(parents=True,exist_ok=True) + completed=[] + for entry in selection['episodes']: + eid=entry['episode_id'];target=output/eid;manifest=target/'manifest.json' + if manifest.exists(): + _verify_complete_episode(manifest,entry) + else: + owned=cache/eid + if not isinstance(eid,str) or not eid or eid in ('.','..') or Path(eid).name!=eid or owned.parent!=cache:raise ValueError('Unsafe task-owned source path') + marker=cache/(eid+'.owner.json') + ownership=dict(episode_id=eid,source=entry['source'],uid=os.getuid(),cache_root=str(cache)) + if marker.exists() and json.loads(marker.read_text())!=ownership:raise ValueError('Source cache ownership mismatch') + if not marker.exists():immutable_json(marker,ownership) + prepare(dict(episodes=[entry]),target,source_cache=cache,image_size=image_size,defer_normalization=True) + _verify_complete_episode(manifest,entry) + completed.append(str(manifest.resolve())) + (output/'progress.json').write_text(json.dumps(dict(completed=completed,total=len(selection['episodes'])),indent=2)+'\n') + if cleanup_source: + owned=cache/eid;marker=cache/(eid+'.owner.json') + # Only the new staged source copy is removable, after a complete + # immutable prepared manifest and exact matching ownership marker. + if owned.exists(): + if not marker.is_file() or owned.is_symlink() or owned.resolve().parent!=cache: + raise ValueError('Refusing unowned/escaped native source cleanup') + record=json.loads(marker.read_text()) + if record!=dict(episode_id=eid,source=entry['source'],uid=os.getuid(),cache_root=str(cache)) or owned.stat().st_uid!=os.getuid(): + raise ValueError('Native source cleanup ownership differs') + shutil.rmtree(owned) + index=output/'shard-manifests.json' + if index.exists(): + if json.loads(index.read_text())!=dict(manifests=completed):raise ValueError('Completed shard index differs') + else: + immutable_json(index,dict(manifests=completed)) + return completed + + +def main(): + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--selection',required=True) + p.add_argument('--output',required=True) + p.add_argument('--source-cache',required=True) + p.add_argument('--image-size',type=int,nargs=2,default=(160,160)) + p.add_argument('--patch-size',type=int,default=16) + p.add_argument('--defer-normalization',action='store_true') + p.add_argument('--shard',action='store_true') + p.add_argument('--cleanup-source',action='store_true') + p.add_argument('--merge-manifest',action='append',default=[]) + p.add_argument('--link-caches',action='store_true') + a=p.parse_args() + if a.merge_manifest: + merge_prepared_manifests(a.selection,a.merge_manifest,a.output,link_caches=a.link_caches) + elif a.shard: + prepare_shard(a.selection,a.output,source_cache=a.source_cache,image_size=tuple(a.image_size),cleanup_source=a.cleanup_source) + else: + prepare(a.selection,a.output,source_cache=a.source_cache,image_size=tuple(a.image_size),patch_size=a.patch_size,defer_normalization=a.defer_normalization) + +if __name__=='__main__': + main() diff --git a/egomimic/rldb/pixel_subtasks.py b/egomimic/rldb/pixel_subtasks.py new file mode 100644 index 000000000..dec188f42 --- /dev/null +++ b/egomimic/rldb/pixel_subtasks.py @@ -0,0 +1,258 @@ +"""Half-open, language-conditioned subtask units for native pixel world models. + +Physical episodes retain their original split/normalization identity. This index +only restricts sampling; it never rewrites image, state, action or source Zarr. +""" +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + + +def normalize_scale_events(document, *, fps, context_labels=()): + """Preserve original Scale object names; reject mistakes/unknown actions. + + Raw Scale clip time is the annotated video's clock, not the human robot + clock. The caller must separately verify video-row/native-frame alignment + before an episode can enter a training index. + """ + import math + from fractions import Fraction + if document.get("time_unit") != "microseconds" or document.get("time_offset") != 0: + raise ValueError("Require explicit microsecond video clock with zero offset") + if not isinstance(fps,(int,float)) or not math.isfinite(fps) or fps <= 0: + raise ValueError("Positive verified native video frame rate required") + rate=Fraction(str(fps)) + result=[] + for annotation in document.get("annotations",[]): + if annotation.get("type") != "text_annotation": + continue + label=annotation.get("label","") + hand=("left" if "left" in label.lower() else "right" if "right" in label.lower() else "both" if "both" in label.lower() else None) + role="context" if label in context_labels else "hand_action" + for clip in annotation.get("clips",[]): + attributes={} + for attribute in clip.get("attributes",[]): + attributes.setdefault(attribute.get("name"),[]).append(attribute) + def single(name): + options=attributes.get(name,[]) + if len(options)!=1 or options[0].get("static") is not True: + return None + values=options[0].get("values",[]) + return values[0] if len(values)==1 and isinstance(values[0],str) else None + action,mistake=single("Action"),single("Mistake") + text=clip.get("text","") + instruction=((action or "")+" "+(text if isinstance(text,str) else "")).strip() + reasons=[] + if role == "context": reasons.append("non_hand_context_annotation") + if mistake != "No": reasons.append("source_mistake_or_unknown") + if not action or not action.strip(): reasons.append("missing_or_ambiguous_action") + if not hand: reasons.append("missing_hand") + if not isinstance(text,str) or not text.strip(): reasons.append("empty_source_text") + start,duration=clip.get("timestamp"),clip.get("duration") + if type(start) is not int or type(duration) is not int or start<0 or duration<=0: + raise ValueError("Malformed source clip time") + result.append(dict(event_id=annotation["id"]+":"+clip["id"],hand=hand, + start_idx=math.ceil(Fraction(start,1000000)*rate),end_idx=math.ceil(Fraction(start+duration,1000000)*rate), + instruction=instruction,approved=not reasons,review_reasons=reasons, + source_label=label,source_action=action,source_mistake=mistake,event_role=role)) + return result + + +def atomic_subtasks(events, frames): + """Combine simultaneous hand instructions over disjoint atomic intervals. + + Callers supply reviewed, frame-aligned source events. Unapproved events, + invalid bounds, and contradictory same-hand overlaps are quarantined with + reasons. No clipping, synonym replacement or fabricated labels occurs. + """ + if type(frames) is not int or frames < 1: + raise ValueError("Require positive physical episode frame count") + valid, rejected, blocked = [], [], [] + seen = set() + for event in events: + reason = None + if not isinstance(event, dict): + raise ValueError("Subtask events must be objects") + a, b = event.get("start_idx"), event.get("end_idx") + if event.get("event_role") == "context": + if event.get("approved") is True: + raise ValueError("Context annotations cannot be approved hand targets") + rejected.append(dict(event=event,reason="non_hand_context_annotation")) + continue + if event.get("approved") is not True: + reason = "annotation_not_approved" + elif (type(a) is not int or type(b) is not int or not 0 <= a < b <= frames): + reason = "invalid_or_out_of_bounds_half_open_interval" + elif (event.get("hand") not in ("left", "right", "both") + or not isinstance(event.get("instruction"), str) + or not event["instruction"].strip() + or not isinstance(event.get("event_id"), str) + or not event["event_id"]): + reason = "missing_hand_instruction_or_source_event_id" + elif event["event_id"] in seen: + raise ValueError("Duplicate source event ID") + if reason: + rejected.append(dict(event=event, reason=reason)) + if type(a) is int and type(b) is int and a < b: + blocked.append(event) + else: + seen.add(event["event_id"]) + valid.append(event) + boundaries = sorted({x for e in valid + blocked + for x in (e["start_idx"], e["end_idx"]) if 0 <= x <= frames}) + units = [] + for a, b in zip(boundaries, boundaries[1:]): + active = [e for e in valid if e["start_idx"] <= a and b <= e["end_idx"]] + if not active: + continue + uncertain = [e for e in blocked if e["start_idx"] < b and a < e["end_idx"]] + if uncertain: + rejected.append(dict(start_idx=a, end_idx=b, + event_ids=sorted(e.get("event_id", "unknown") for e in uncertain), + reason="unapproved_concurrent_instruction")) + continue + per_hand = {h: {e["instruction"] for e in active if e["hand"] == h} + for h in ("left", "right", "both")} + if any(len(v) > 1 for v in per_hand.values()) or (per_hand["both"] and + (per_hand["left"] or per_hand["right"])): + rejected.append(dict(start_idx=a, end_idx=b, + event_ids=sorted(e["event_id"] for e in active), + reason="conflicting_concurrent_hand_instructions")) + continue + instruction = " ".join(f"{h.capitalize()} hand: {next(iter(per_hand[h]))}" + for h in ("left", "right", "both") if per_hand[h]) + event_ids = sorted(e["event_id"] for e in active) + if units and units[-1]["end_idx"] == a and units[-1]["event_ids"] == event_ids: + units[-1]["end_idx"] = b + else: + units.append(dict(start_idx=a, end_idx=b, instruction=instruction, + event_ids=event_ids)) + for unit in units: + unit["task_id"] = "subtask-" + hashlib.sha256(unit["instruction"].encode()).hexdigest() + return units, rejected + + +def load_subtask_manifest(path, sha256, manifest, manifest_sha256): + """Load an immutable reviewed index bound to exact prepared episode bytes. + + A generated annotation is not admitted just because a file exists: every + unit requires explicit quality/name and source-clock alignment admission. + """ + if (path is None) != (sha256 is None): + raise ValueError("Subtask path and SHA256 require both or neither") + if path is None: + return None + raw = Path(path).read_bytes() + if hashlib.sha256(raw).hexdigest() != sha256: + raise ValueError("Subtask manifest SHA256 mismatch") + index = json.loads(raw) + if (index.get("schema") != "pixel-subtask-index-v1" + or index.get("prepared_manifest_sha256") != manifest_sha256 + or index.get("interval_semantics") != "half_open" + or index.get("concurrent_policy") != "combine_active_hands_atomic_intervals"): + raise ValueError("Subtask index contract/manifest identity mismatch") + expected = {(e["embodiment"], e["episode_id"]): e for e in manifest["episodes"]} + found = {} + for record in index.get("episodes", []): + key = record.get("embodiment"), record.get("episode_id") + if key not in expected or key in found: + raise ValueError("Unknown or duplicate physical episode in subtask index") + ep = expected[key] + if record.get("split") != ep["split"] or record.get("frames") != ep["frames"]: + raise ValueError("Subtask index may not change physical episode split/length") + units, previous = record.get("subtasks"), 0 + if not isinstance(units, list) or not units: + raise ValueError("Every selected episode needs reviewed subtask intervals") + for unit in units: + a, b = unit.get("start_idx"), unit.get("end_idx") + if (type(a) is not int or type(b) is not int + or not previous <= a < b <= ep["frames"]): + raise ValueError("Require disjoint in-bounds half-open subtasks") + previous = b + segments = ep.get("continuity_segments", [[0, ep["frames"]]]) + if not any(lo <= a < b <= hi for lo,hi in segments): + raise ValueError("Subtask crosses a native recording discontinuity") + if (unit.get("quality_and_names_verified") is not True + or unit.get("source_clock_alignment_verified") is not True + or not isinstance(unit.get("task_id"), str) or not unit["task_id"] + or not isinstance(unit.get("instruction"), str) or not unit["instruction"].strip() + or not unit.get("annotation_provenance")): + raise ValueError("Subtask lacks annotation, naming or clock admission") + found[key] = units + if set(found) != set(expected): + raise ValueError("Subtask index must cover the selected prepared corpus exactly") + return found + + +def build_subtask_index(prepared_manifest_path, events_path): + """Compile reviewed frame-aligned sidecars without touching physical data.""" + manifest_raw = Path(prepared_manifest_path).read_bytes() + event_raw = Path(events_path).read_bytes() + manifest, source = json.loads(manifest_raw), json.loads(event_raw) + if source.get("schema") != "pixel-subtask-events-v1": + raise ValueError("Require pixel-subtask-events-v1 normalized source events") + records = {} + for record in source.get("episodes", []): + key = record.get("embodiment"), record.get("episode_id") + if key in records: + raise ValueError("Duplicate normalized event episode") + records[key] = record + index = dict(schema="pixel-subtask-index-v1", + prepared_manifest_sha256=hashlib.sha256(manifest_raw).hexdigest(), + annotation_events_sha256=hashlib.sha256(event_raw).hexdigest(), + interval_semantics="half_open", + concurrent_policy="combine_active_hands_atomic_intervals", episodes=[]) + quarantine = [] + for ep in manifest["episodes"]: + key = ep["embodiment"], ep["episode_id"] + record = records.get(key) + if not record or record.get("source_clock_alignment_verified") is not True: + quarantine.append(dict(episode_id=ep["episode_id"],reason="missing_verified_source_clock")) + continue + if record.get("frames") != ep["frames"] or record.get("split") != ep["split"]: + raise ValueError("Event source may not alter episode frame count or split") + units, rejected = atomic_subtasks(record.get("events", []), ep["frames"]) + quarantine.extend(dict(episode_id=ep["episode_id"], **r) for r in rejected) + segments = ep.get("continuity_segments", [[0, ep["frames"]]]) + bounded=[] + for unit in units: + for lo,hi in segments: + a,b=max(unit["start_idx"],lo),min(unit["end_idx"],hi) + if a= {minimum}") - if self.settings["mode"] not in ("continuous", "discrete"): - raise ValueError("Prepared sampling mode must be continuous or discrete") - if self.settings["max_frames"] < self.settings["min_frames"]: - raise ValueError("max_frames must be >= min_frames") - probability = self.settings["forecasting_probability"] - if ( - isinstance(probability, bool) - or not isinstance(probability, (int, float)) - or not math.isfinite(probability) - or not 0 <= probability <= 1 - ): - raise ValueError("forecasting_probability must be finite and in [0, 1]") - bounded = _bounded_time_sampling( - self.settings.get("max_past_seconds"), - self.settings.get("max_future_seconds"), - ) - if bounded and self.settings["mode"] != "continuous": - raise ValueError("Physical-time bounds apply only to continuous sampling") - return { - "scope": "configuration_only", - "artifacts_opened": False, - "sampling": dict(self.settings), - "batch_size": self.batch_size, - "num_workers": self.num_workers, - "checkpointable_sampling": self.checkpointable_sampling, - } - def setup(self, stage=None): if not self.datasets: # Reuse finalized fingerprints only for unchanged exact files; # otherwise hash once per setup and share that proof across splits. verification = ArtifactVerificationMemo() self.datasets = { - s: PreparedWorldDataset( - self.manifest_path, - split=s, - verification_memo=verification, - **self.settings, - ) + s: PreparedWorldDataset(self.manifest_path, split=s, + verification_memo=verification, **self.settings) for s in ("train", "val") } @@ -151,10 +91,7 @@ def prepare_context(self, *, mode, normalization, normalizer, restored_state=Non "sampling": self.settings, } if self.checkpointable_sampling: - state["sampling_resume"] = { - "schema": "mot-sampling-state-v1", - "batch_size": self.batch_size, - } + state["sampling_resume"] = {"schema": "mot-sampling-state-v1", "batch_size": self.batch_size} if self.validation_batches_per_embodiment is not None: state["validation_sampling"] = { "batches_per_embodiment": self.validation_batches_per_embodiment, @@ -192,15 +129,15 @@ def _loader(self, split): ) dataset = self.datasets[split] sampler = HomogeneousWorldBatchSampler( - dataset, - self.batch_size, - self.settings["seed"], - shuffle=split == "train", - checkpointable=self.checkpointable_sampling and split == "train", - validation_batches_per_embodiment=( - self.validation_batches_per_embodiment if split == "val" else None - ), - ) + dataset, + self.batch_size, + self.settings["seed"], + shuffle=split == "train", + checkpointable=self.checkpointable_sampling and split == "train", + validation_batches_per_embodiment=( + self.validation_batches_per_embodiment if split == "val" else None + ), + ) if split == "train": self.train_sampler = sampler sampler.consumed_batches = self.consumed_batches @@ -208,7 +145,6 @@ def _loader(self, split): generator = None if self.checkpointable_sampling: import torch - generator = torch.Generator().manual_seed(self.settings["seed"]) return DataLoader( dataset, @@ -230,22 +166,14 @@ def acknowledge_batch(self, cursor): def state_dict(self): if not self.checkpointable_sampling: return {} - return { - "schema": "mot-sampling-state-v1", - "consumed_batches": self.consumed_batches, - "manifest_sha256": hashlib.sha256( - Path(self.manifest_path).read_bytes() - ).hexdigest(), - "settings": self.settings, - "batch_size": self.batch_size, - } + return {"schema": "mot-sampling-state-v1", "consumed_batches": self.consumed_batches, + "manifest_sha256": hashlib.sha256(Path(self.manifest_path).read_bytes()).hexdigest(), + "settings": self.settings, "batch_size": self.batch_size} def load_state_dict(self, state_dict): if not self.checkpointable_sampling: if state_dict: - raise ValueError( - "Checkpointed sampling requires its original opt-in configuration" - ) + raise ValueError("Checkpointed sampling requires its original opt-in configuration") return expected = self.state_dict() for key in ("schema", "manifest_sha256", "settings", "batch_size"): diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 91e7a2a9e..64ed66db5 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -94,7 +94,8 @@ def _validate_episode_name_pin( def decode_jpeg_payload(payload) -> np.ndarray: """Decode either one JPEG or a temporal sequence without dropping time.""" - if isinstance(payload, np.ndarray) and payload.ndim == 0: + # Zarr v3 variable-length codecs may nest scalar object wrappers. + while isinstance(payload, np.ndarray) and payload.ndim == 0: payload = payload.item() if isinstance(payload, (bytes, bytearray, memoryview)): image = simplejpeg.decode_jpeg(bytes(payload), colorspace="RGB") diff --git a/egomimic/scripts/embedding_process/qwen3_embedding.py b/egomimic/scripts/embedding_process/qwen3_embedding.py index fd0241d66..41120678d 100644 --- a/egomimic/scripts/embedding_process/qwen3_embedding.py +++ b/egomimic/scripts/embedding_process/qwen3_embedding.py @@ -74,6 +74,7 @@ def __init__( device: str | torch.device = "cuda", chunk_timesteps: int = 100, dtype: torch.dtype = torch.float16, + revision: str | None = None, ): if output_keys is not None and len(input_keys) != len(output_keys): raise ValueError( @@ -89,6 +90,7 @@ def __init__( chunk_timesteps=chunk_timesteps, ) self.model_name = model_name + self.revision = revision self.max_length = max_length self.dtype = dtype self.tokenizer = None @@ -100,9 +102,9 @@ def setup(self) -> None: logger.info("Loading Qwen3 embedding model: %s", self.model_name) self.tokenizer = AutoTokenizer.from_pretrained( - self.model_name, padding_side="left" + self.model_name, padding_side="left", revision=self.revision ) - self.model = AutoModel.from_pretrained(self.model_name, torch_dtype=self.dtype) + self.model = AutoModel.from_pretrained(self.model_name, torch_dtype=self.dtype, revision=self.revision) self.model.to(self.device).eval() self._embed_dim = int(self.model.config.hidden_size) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 4768bf385..648fe201c 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -59,9 +59,7 @@ def _configure_strict_fp32(cfg: DictConfig) -> None: if enabled and str(precision) not in {"32", "32-true"}: raise ValueError("runtime.strict_fp32 requires trainer.precision=32-true") if os.environ.get("NVIDIA_TF32_OVERRIDE") != "0": - raise ValueError( - "FP32/disable_tf32 backend policy requires NVIDIA_TF32_OVERRIDE=0" - ) + raise ValueError("FP32/disable_tf32 backend policy requires NVIDIA_TF32_OVERRIDE=0") torch.set_float32_matmul_precision("highest") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False @@ -154,8 +152,7 @@ def _instantiate_slurm_callbacks(cfg: DictConfig) -> List[Callback]: checkpoint_dir=str(checkpoint_dir), save_signal=signal.SIGUSR2, train_batch_only=OmegaConf.select( - cfg, "runtime.slurm_signal_train_batch_only", default=False - ), + cfg, "runtime.slurm_signal_train_batch_only", default=False), ) ] @@ -198,14 +195,10 @@ def _torch_elastic_environment(cfg: DictConfig) -> TorchElasticEnvironment: except ValueError as exc: raise ValueError(f"torch_elastic requires integer {name}") from exc world, local_world = integers["WORLD_SIZE"], integers["LOCAL_WORLD_SIZE"] - if ( - world < 1 - or local_world < 1 - or world % local_world - or not 0 <= integers["RANK"] < world - or not 0 <= integers["LOCAL_RANK"] < local_world - or not 0 <= integers["GROUP_RANK"] < world // local_world - ): + if (world < 1 or local_world < 1 or world % local_world + or not 0 <= integers["RANK"] < world + or not 0 <= integers["LOCAL_RANK"] < local_world + or not 0 <= integers["GROUP_RANK"] < world // local_world): raise ValueError("torch_elastic has inconsistent world size or rank variables") if not os.environ.get("MASTER_ADDR"): raise ValueError("torch_elastic requires MASTER_ADDR") diff --git a/scripts/audit_hydra_configs.py b/scripts/audit_hydra_configs.py index 5963d7435..1ffd6d616 100644 --- a/scripts/audit_hydra_configs.py +++ b/scripts/audit_hydra_configs.py @@ -19,14 +19,6 @@ CONFIGS = ROOT / "egomimic/hydra_configs" EXPERIMENT_CONTEXTS = { - "data/world_model/prepared": "world_model/continuous", - "model/world_model/base": "world_model/continuous", - "data/dino_proprio/prepared": "dino_proprio/visual_state", - "model/dino_proprio/base": "dino_proprio/visual_state", - "data/pixel_proprio/prepared": "pixel_proprio/video_continuous", - "model/pixel_proprio/base": "pixel_proprio/video_continuous", - "model/pixel_proprio/video": "pixel_proprio/video_continuous", - "model/pixel_proprio/readout": "pixel_proprio/readout_continuous", "data/e1/yam_human_keypoints_multiplier_proportional_val01": "e1/yam_human_matched_af64_multiplier", "data/e1/yam_human_matched_noaug_proportional_val01": "e1/yam_human_matched_dp_noaug", "data/pusht/planar_chain_manual4919_standard_retimed": "pusht/planar_usocket_chain_manual4919_standard_dp_retimed_masked", @@ -194,49 +186,6 @@ def _compose_for_audit(path): cfg.hydra.runtime.output_dir = "/tmp/egoverse-config-audit" cfg.hydra.job.num = 0 cfg.hydra.job.id = "config-audit" - # Prepared manifests and readout weights are launch-owned inputs. These - # synthetic paths resolve the offline graph only; no artifact is opened - # or certified here, and production YAML keeps its mandatory bindings. - data_target = OmegaConf.select(cfg, "data._target_") - prepared_manifests = { - "egomimic.rldb.world_model_datamodule.WorldModelDataModule": ( - "/tmp/egoverse-config-audit/world-manifest.json" - ), - "egomimic.rldb.dino_proprio_datamodule.DinoProprioDataModule": ( - "/tmp/egoverse-config-audit/dino-manifest.json" - ), - "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule": ( - "/tmp/egoverse-config-audit/pixel-manifest.json" - ), - } - if data_target in prepared_manifests: - if OmegaConf.is_missing(cfg.data, "manifest_path"): - cfg.data.manifest_path = prepared_manifests[data_target] - if ( - data_target - == "egomimic.rldb.dino_proprio_datamodule.DinoProprioDataModule" - and cfg.data.task == "direct_action" - and OmegaConf.select(cfg.data, "action_manifest_path") is None - ): - cfg.data.action_manifest_path = ( - "/tmp/egoverse-config-audit/dino-actions.json" - ) - for stage in OmegaConf.select(cfg, "model.pipeline.stages", default=[]): - model = stage.get("model") - if ( - OmegaConf.is_dict(model) - and model.get("_target_") == "egomimic.models.world_model.MoTWorldModel" - and OmegaConf.is_missing(model, "pose_dims") - ): - # Unverified structural width, not a prepared-data pose schema. - model.pose_dims = {"audit_measured_pose": 9} - if "pixel_readout" in cfg: - for key, value in { - "initial_weights": "/tmp/egoverse-config-audit/pixel-readout.pt", - "initial_weights_sha256": "0" * 64, - }.items(): - if OmegaConf.is_missing(cfg.pixel_readout, key): - cfg.pixel_readout[key] = value if "pi05" in cfg and OmegaConf.is_missing(cfg.pi05, "pretrained_weights"): cfg.pi05.pretrained_weights = "/tmp/egoverse-config-audit/pi-weights" if "e1" in cfg: 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/boundary_uncertainty.py b/scripts/train/joint/boundary_uncertainty.py new file mode 100644 index 000000000..ab4559ddb --- /dev/null +++ b/scripts/train/joint/boundary_uncertainty.py @@ -0,0 +1,453 @@ +"""Pure opt-in admission for bounded annotation uncertainty; no ML imports. + +This verifies a new data policy rather than relabeling strict-boundary evidence. +The sampler/data candidate and this operational candidate must be reviewed and +source-bound together. Missing census/native proof fails closed. +""" +from collections.abc import Mapping, Sequence +from copy import deepcopy +from datetime import datetime +import hashlib +import json +import math +import re + +POLICY = dict(schema='annotation-boundary-uncertainty-sampling-v1',policy='annotation_boundary_uncertainty', + boundary_probability=0.5,maximum_radius_seconds=0.25,duration_fraction=0.1,rng_namespace=0x424E4459) +BASE_REVISION='153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920' +BASE_MANIFEST='208ba5fa51d0514fd27d00f0f202cbf73f6e31ab7f3537e26ca55004e4fe3f99' +SAMPLER_FILES={ + 'egomimic/rldb/boundary_sampling.py':'1522b1c1d9d9f7c24977d6184bc2aeca125a7dd7aea890069dced7e7384c42f0', + 'egomimic/rldb/pixel_proprio.py':'9b4a40f66ac36883553e9c2c14007167d23cdbcfed4aa0aa5afb5d1d7b83e4b1', + 'egomimic/rldb/pixel_proprio_datamodule.py':'ec3fd8280b30cce469078a3901c07811a664efe2c939d3480ed4e73312be8e2c', + 'egomimic/pl_utils/joint_smoke.py':'8da5518c058afab469b7e89ac8940c7b246e4ddb8cc6252bc1163350a36ada1c', +} +STREAMS = ['image','proprio','action','camera'] +NOISER = dict(_target_='egomimic.models.joint_image_proprio.CoupledImageProprioNoising',maximum=1.0,active_streams=STREAMS) +LOSSES = {name:dict(_target_='egomimic.pipeline.stages_joint_state_action.JointLossStage',stream=name,weight=1.0, + loss=dict(_target_='egomimic.models.pixel_diffusion.PlainPixelMSE',reduction='mean_per_example'), + **({'optional':True} if name=='action' else {})) for name in STREAMS} +OBJECTIVE = dict(schema='boundary-independent-noise-plain-objective-v1',training_conditioning='independent_frame_noise', + noising=NOISER,loss_stages=LOSSES,forced_clean_initial_frame=False) +UNCERTAINTY_FIELDS={'boundary_uncertainty','subtask_boundary_verification_policy','original_subtask_annotation_identity_verified','boundary_policy_sha256','boundary_objective_sha256'} +COUNT_KEYS = {'samples','requested_boundary','realized_boundary','fallback','legacy_uniform','outward_samples', + 'start_inward','start_outward','end_inward','end_outward'} +PROOF_FLAGS = ('synchronized_native_streams_verified','original_instruction_retained','validation_sampling_unchanged') + + +def need(ok,message): + if not ok:raise ValueError(message) + + +def digest(value): + def plain(item): + if isinstance(item,Mapping):return {k:plain(v) for k,v in item.items()} + if isinstance(item,Sequence) and not isinstance(item,(str,bytes,bytearray)): + return [plain(v) for v in item] + return item + return hashlib.sha256(json.dumps(plain(value),sort_keys=True,separators=(',',':'),allow_nan=False).encode()).hexdigest() + + +def policy(value): + if value is None:return None + need(isinstance(value,Mapping) and dict(value)==POLICY + and all(type(value[k]) is float for k in ('boundary_probability','maximum_radius_seconds','duration_fraction')) + and type(value['rng_namespace']) is int,'Exact typed boundary uncertainty policy required') + return deepcopy(POLICY) + + +def sampling_config(cfg): + data=cfg['data'];selected=policy(data.get('boundary_sampling')) + if selected is None:return None + need(data['mode']=='continuous' and data['sampling_strategy']=='full_episode_uniform' + and data['min_frames']==2 and data['max_frames']==64 and data['checkpointable_sampling'] is True + and data['training_conditioning']=='independent_frame_noise' and data['task']=='visual_state_action' + and data.get('temporal_gap_schedule') is None and data.get('completion_gap_schedule') is None + and data.get('subtask_manifest_sha256') and data.get('language_features_sha256') + and cfg['seed']==42,'Boundary arm requires exact checkpointed independent-noise continuous data') + stages=cfg['model']['pipeline']['stages'];ids=cfg['model']['pipeline']['stage_ids'] + noising=stages[ids['joint_noising']] + need(digest(noising)==digest(dict(_target_='egomimic.pipeline.stages_joint_state_action.JointNoisingStage',noising=NOISER)), + 'Boundary arm changed independent timestamp noiser') + need(all(digest(stages[ids[name+'_loss']])==digest(LOSSES[name]) for name in STREAMS),'Boundary arm changed plain FP32 objective') + need(cfg['model']['training_behavior']['distributed_noise_seed']==42,'Boundary arm changed distributed model-noise seed') + return selected + + +def request_config(request,cfg): + selected=sampling_config(cfg);contract=request.get('boundary_sampling_contract') + need(not any(k.startswith('boundary') and k!='boundary_sampling_contract' for k in request), + 'Unexpected unbound boundary request metadata') + if selected is None: + need(contract is None,'Baseline request cannot admit boundary contract/metadata') + return None + need(isinstance(contract,dict) and set(contract)=={'schema','policy','objective','eligibility'} + and contract['schema']=='boundary-uncertainty-admission-v1' + and policy(contract['policy'])==selected and digest(contract['objective'])==digest(OBJECTIVE), + 'Request does not bind exact boundary policy/noiser/plain objective') + need(request['arm']=='continuous','Boundary uncertainty is a continuous arm') + need(set(contract['eligibility'])=={'report','producer_request','producer_spec','producer_files','producer_completion','claim','native','host_exit'},'Complete allocated eligibility producer bindings required') + for binding in contract['eligibility'].values(): + need(isinstance(binding,dict) and set(binding)=={'path','sha256'} + and isinstance(binding['path'],str) and binding['path'].startswith('/') + and '..' not in binding['path'].split('/') + and re.fullmatch('[0-9a-f]{64}',str(binding['sha256'])), + 'Allocated eligibility bindings remain unbound or unsafe') + return contract + + +def context(cfg,state,sampler=None): + selected=sampling_config(cfg) + sampling=state.get('sampling',{});overrides=state.get('validation_overrides',{}) + if selected is None: + need(sampling.get('boundary_sampling') is None and 'boundary_sampling' not in overrides, + 'Baseline checkpoint context contains boundary sampling identity') + return None + need(policy(sampling.get('boundary_sampling'))==selected + and sampling.get('training_conditioning')=='independent_frame_noise' + and 'boundary_sampling' in overrides and overrides['boundary_sampling'] is None + and overrides.get('training_conditioning')=='mixed_context', + 'Checkpoint context must bind training policy and strict legacy validation') + if sampler is not None: + need(sampler.get('settings')==sampling and sampler.get('validation_overrides')==overrides, + 'Checkpoint sampler and data context differ on boundary/validation policy') + return selected + + +def strict_validation_samples(samples): + """Held-out query targets may already be NaN-hidden; inspect metadata only.""" + for sample in samples.values(): + records=sample.get('pixel_metadata',[]) + need(records,'Missing actual strict held-out sample metadata') + for i,row in enumerate(records): + unit=row.get('subtask',{});need(unit.get('interval_semantics')=='half_open','Held-out subtask semantics changed') + need('boundary_sampling' not in row,'Boundary policy leaked into held-out sampling') + indices=list(row.get('frame_indices',[])) + for key in ('proprio_frame_indices','action_frame_indices','camera_frame_indices'): + for group in row.get(key,[]):indices.extend(group if isinstance(group,list) else [group]) + need(indices and all(unit['start_idx']<=v=0 for v in counts.values()),'Invalid actual boundary realization counters') + n,r,b,f,u=(counts[k] for k in ('samples','requested_boundary','realized_boundary','fallback','legacy_uniform')) + need(n==r+u and r==b+f and counts['outward_samples']<=b + and counts['start_inward']+counts['start_outward']<=b + and counts['end_inward']+counts['end_outward']<=b, + 'Requested/realized/fallback/offset counters are inconsistent') + + +def fit_rows(ranks,*,selected,world,per_rank_batch,embodiments,context_sha256): + if selected is None: + for row in ranks: + need(not (UNCERTAINTY_FIELDS & set(row)) + and row.get('subtask_boundaries_verified') is True,'Baseline fit cannot admit uncertainty evidence') + return None + need(policy(selected)==POLICY and len(ranks)==world,'Boundary fit topology/policy changed') + proof=ranks[0].get('boundary_uncertainty') + need(isinstance(proof,dict) and proof.get('schema')=='boundary-uncertainty-smoke-realizations-v1' + and policy(proof.get('policy'))==selected and all(proof.get(k) is True for k in PROOF_FLAGS) + and proof.get('neighbor_instruction_equivalence_asserted') is False, + 'Missing distinct checked boundary uncertainty evidence') + all_counts=proof.get('all_rank_counts') + need(isinstance(all_counts,list) and len(all_counts)==world,'All actual rank boundary counters required') + expected=set(embodiments) + for rank,row in enumerate(ranks): + need(row.get('rank')==rank and row.get('world_size')==world and row.get('global_step')==50 + and row.get('boundary_uncertainty')==proof and row.get('subtask_boundaries_verified') is False + and row.get('original_subtask_annotation_identity_verified') is True + and row.get('subtask_boundary_verification_policy')==POLICY['policy'] + and row.get('data_context_sha256')==context_sha256, + 'Boundary fit policy/context/annotation identity differs across ranks') + local=all_counts[rank] + need(set(local)==expected and set(row.get('embodiment_optimizer_batches',{}))==expected, + 'Missing actual robot/human boundary training populations') + need(all(type(v) is int and v>0 for v in row['embodiment_optimizer_batches'].values()) + and sum(row['embodiment_optimizer_batches'].values())==50, + 'Every rank must optimize both real embodiment populations over50 updates') + for emb,counters in local.items(): + _counts(counters) + need(counters['samples']==row['embodiment_optimizer_batches'][emb]*per_rank_batch, + 'Boundary counts differ from actual optimized embodiment batches') + need(sum(c['samples'] for c in local.values())==50*per_rank_batch,'Boundary counters do not cover all50 updates') + totals={emb:{key:sum(rank[emb][key] for rank in all_counts) for key in COUNT_KEYS} for emb in expected} + need(proof.get('global_counts_by_embodiment')==totals + and all(v['samples']>0 for v in totals.values()) + and sum(v['outward_samples'] for v in totals.values())>0, + 'Boundary global aggregation/actual outward coverage is missing') + return deepcopy(proof) + + +def _diagnostic(unit,clock,segments): + a,b=unit['start_idx'],unit['end_idx'];segments=[s for s in segments if s[0]<=a=b) and istart and fields.get('ExitCode')=='0:0','Shared census host/native terminal proof failed') + for key in ('ReqTRES','AllocTRES'): + tres=fields.get(key,'').split(',') + need(not any('gpu' in value.lower() for value in tres),'Shared census may not allocate a GPU') + if key=='ReqTRES' or state in {'RUNNING','COMPLETED'}: + need('cpu=4' in tres and 'node=1' in tres and ('mem=24G' in tres or 'mem=24576M' in tres),'Wrong shared census exact TRES') + return fields + + + +def census_context_contract(spec): + """Pin derivation before execution without inventing an actual context.""" + return dict(schema='allocated-boundary-context-derivation-v1', + method='candidate_PixelProprioDataModule.prepare_context_on_verified_train_dataset', + datamodule=dict(relative_path='candidate/pixel_proprio_datamodule.py', + sha256=SAMPLER_FILES['egomimic/rldb/pixel_proprio_datamodule.py']), + resolved_config=dict(relative_path='inputs/resolved.yaml', + sha256='c2377e6c095974db9a753cab4a0ee2daffad0ddd0b8ceada9b3d6d6ee721efbb'), + prepared_inputs=spec['data']['hashes'],boundary_sampling=POLICY,dataset_split='train',projection_only=True, + actual_context_emitted_only_by_allocated_producer=True, + future_training_requires_exact_actual_context_fingerprint=True) + +def eligibility(request,manifest,index,reader): + contract=request.get('boundary_sampling_contract') + if contract is None:return None + need(contract.get('schema')=='boundary-uncertainty-admission-v1' and policy(contract.get('policy'))==POLICY + and digest(contract.get('objective'))==digest(OBJECTIVE),'Invalid explicit boundary census admission contract') + bindings=contract['eligibility'];report=reader(bindings['report']);producer=reader(bindings['producer_request']) + spec_raw=reader(bindings['producer_spec'],False);spec=json.loads(spec_raw) + files_raw=reader(bindings['producer_files'],False) + complete=reader(bindings['producer_completion']);claim=reader(bindings['claim']) + native=reader(bindings['native'],False);host=reader(bindings['host_exit'],False) + need(host.strip()=='0','Allocated shared census host failed') + root=producer.get('remote_root','');output=producer.get('output','') + spec_sha=bindings['producer_spec']['sha256'];request_sha=bindings['producer_request']['sha256'];files_sha=bindings['producer_files']['sha256'] + files={} + for line in files_raw.splitlines(): + if not line.strip():continue + parts=line.split(None,1);need(len(parts)==2,'Malformed shared producer FILES manifest') + checksum,name=parts;name=name.lstrip('*').removeprefix('./') + need(name not in files,'Duplicate shared producer FILES entry') + files[name]=checksum + need(files.get('request.json')==request_sha and files.get('spec.json')==spec_sha + and all(files.get(name)==wanted for name,wanted in spec.get('code',{}).items()), + 'Shared producer FILES does not bind exact request/spec/code closure') + # The CPU census certifies portable, hash-bound TRAIN metadata, not the + # future training node. Bind its own producer node throughout request/spec, + # held acceptance, running record and terminal proof. The target node keeps + # its separate actual runtime/cache admission. + census_node=producer.get('node') + need(isinstance(census_node,str) and re.fullmatch(r'gpu[0-9]{2}',census_node), + 'Explicit actual census producer node required') + resources=dict(node=census_node,cpus=4,memory_gib=24,time_limit='00:45:00',gpus=0) + need(producer.get('schema')=='train-calibration-and-boundary-census-request-v1' + and producer.get('spec')==bindings['producer_spec'] + and producer.get('binding_state')=='BOUND_READY_FOR_ROOT_REVIEW' and producer.get('root_review')=='PASS' + and root.startswith('/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/calibration-census-') + and '..' not in root.split('/') and output==root+'/output' + and producer.get('node')==resources['node'] and producer.get('uid')==41003 and producer.get('resources')==resources, + 'Shared census request is not actually bound/reviewed on the exact CPU resources') + need(spec.get('schema')=='train-calibration-and-boundary-census-spec-v1' + and spec.get('resources')==resources + and spec.get('source',{}).get('revision')==BASE_REVISION + and spec['source'].get('manifest_sha256')==BASE_MANIFEST + and spec.get('data',{}).get('root')==request['data_manifest'].rsplit('/',1)[0] + and spec.get('data',{}).get('hashes',{}).get('manifest.json')==request['data_manifest_sha256'] + and spec['data']['hashes'].get('subtasks.json')==request['subtask_contract']['subtask_manifest']['sha256'] + and spec.get('code',{}).get('candidate/boundary_sampling.py')==SAMPLER_FILES['egomimic/rldb/boundary_sampling.py'], + 'Shared census specification changed base source/exact sampler/data identity') + derivation=census_context_contract(spec) + need(spec.get('context_derivation')==derivation and producer.get('context_derivation')==derivation + and report.get('context_derivation')==derivation and spec.get('config')==derivation['resolved_config'] + and spec.get('code',{}).get('candidate/pixel_proprio_datamodule.py')==derivation['datamodule']['sha256'] + and files.get('inputs/resolved.yaml')==derivation['resolved_config']['sha256'], + 'Shared census changed exact DataModule/config/input context derivation') + for item in (derivation['datamodule'],derivation['resolved_config']): + reader(dict(path=root+'/'+item['relative_path'],sha256=item['sha256']),False) + expected_name='elmo-calib-census-'+spec_sha[:16] + command=producer.get('command_binding',{}) + need(command==dict(path=root+'/producer.sbatch',sha256=spec['code'].get('producer.sbatch')) + and re.fullmatch('[0-9a-f]{64}',str(command['sha256'])) and bool(reader(command,False).strip()), + 'Missing immutable shared census command binding') + need(report.get('schema')=='allocated-boundary-census-v1' and report.get('status')=='CENSUS_PASS' + and report.get('producer_request_sha256')==request_sha and report.get('producer_spec_sha256')==spec_sha + and report.get('source_manifest_sha256')==BASE_MANIFEST + and report.get('boundary_sampler_sha256')==SAMPLER_FILES['egomimic/rldb/boundary_sampling.py'] + and report.get('data_manifest_sha256')==request['data_manifest_sha256'] + and report.get('subtask_manifest_sha256')==request['subtask_contract']['subtask_manifest']['sha256'] + and policy(report.get('policy'))==POLICY and digest(report.get('objective'))==digest(OBJECTIVE) + and report.get('zero_optimizer_updates') is True + and report.get('every_clock_read_from_bound_prepared_artifact') is True + and report.get('source_stat_unchanged') is True,'Actual base+sampler-bound allocated clock census missing') + # Whole source manifests legitimately differ after applying the candidate. + # The eventual staged source must contain the exact reviewed sampler/guard + # bytes used to define this CPU diagnostic's projected data semantics. + production_manifest=reader(dict(path=request['source_manifest'],sha256=request['source_manifest_sha256']),False) + production_files={} + for line in production_manifest.splitlines(): + if not line.strip():continue + parts=line.split(None,1);need(len(parts)==2,'Malformed production source manifest') + checksum,name=parts;name=name.lstrip('*').removeprefix('./') + need(name not in production_files,'Duplicate production source manifest entry') + production_files[name]=checksum + for name,wanted in SAMPLER_FILES.items(): + need(production_files.get(name)==wanted,'Published production sampler differs from reviewed census semantics: '+name) + reader(dict(path=request['source'].rstrip('/')+'/'+name,sha256=wanted),False) + job=str(report.get('job_id',''));name=producer.get('job_name','') + need(job.isdecimal() and name==expected_name + and claim.get('schema')=='train-calibration-and-boundary-census-claim-v1' + and claim.get('status')=='HELD_NATIVE_ACCEPTANCE_BOUND' + and claim.get('producer_request_sha256')==request_sha and claim.get('producer_spec_sha256')==spec_sha + and claim.get('files_manifest_sha256')==files_sha + and claim.get('job_id')==job and claim.get('job_name')==name, + 'Shared census terminal proof is not bound to exact accepted request/spec/FILES claim') + claim_root=root+'/claims/'+spec_sha + need(bindings['claim']['path']==claim_root+'/claim.json' + and claim.get('accepted_held_native',{}).get('path')==claim_root+'/HELD_NATIVE.txt' + and claim.get('submit_argv',{}).get('path')==claim_root+'/SUBMIT_ARGV.json', + 'Shared census claim artifacts escaped immutable spec namespace') + held=reader(claim['accepted_held_native'],False) + argv=reader(claim['submit_argv']) + need(argv==census_argv(producer,request_sha,files_sha),'Shared census submission argv differs from exact held producer request') + held_fields=census_native(held,job,name,'PENDING',root,resources) + terminal_fields=census_native(native,job,name,'COMPLETED',root,resources) + need(complete.get('schema')=='train-calibration-and-boundary-census-v1' + and complete.get('status')=='CPU_DIAGNOSTICS_PASS_HOST_NATIVE_COMPLETION_REQUIRED' + and complete.get('job_id')==job and complete.get('node')==resources['node'] + and complete.get('producer_request_sha256')==request_sha and complete.get('producer_spec_sha256')==spec_sha + and complete.get('files_manifest_sha256')==files_sha + and complete.get('claim_sha256')==bindings['claim']['sha256'] + and complete.get('held_native_sha256')==claim['accepted_held_native']['sha256'] + and complete.get('proofs',{}).get('boundary-unit-census.json')==bindings['report']['sha256'] + and complete.get('training_or_GPU_certified') is False and complete.get('production_scale_bound') is False, + 'Shared CPU diagnostic completion does not bind the actual census/claim') + running=reader(dict(path=output+'/native-running.json',sha256=complete['proofs']['native-running.json'])) + need(running.get('job_id')==job and running.get('request_sha256')==request_sha + and running.get('spec_sha256')==spec_sha and running.get('files_manifest_sha256')==files_sha + and running.get('claim_sha256')==bindings['claim']['sha256'] + and running.get('held_native_sha256')==claim['accepted_held_native']['sha256'], + 'Shared running allocation proof differs from accepted producer claim') + running_fields=census_native(running['raw'],job,name,'RUNNING',root,resources) + need(running_fields['StartTime']==terminal_fields['StartTime'],'Terminal census belongs to another allocation interval') + state=report.get('data_context',{}) + projection=reader(dict(path=output+'/projected-boundary-data-context.json', + sha256=complete['proofs']['projected-boundary-data-context.json'])) + need(projection.get('schema')=='allocated-projected-boundary-data-context-v1' + and projection.get('status')=='ACTUAL_METADATA_CONTEXT_PASS' and projection.get('job_id')==job + and projection.get('context_derivation')==derivation and projection.get('data_context')==state + and projection.get('data_context_sha256')==report.get('data_context_sha256') + and projection.get('actual_manifest_read') is True and projection.get('actual_candidate_prepare_context_called') is True + and projection.get('future_model_or_checkpoint_context_claimed') is False, + 'Allocated actual DataModule context projection is missing or changed') + need(re.fullmatch('[0-9a-f]{64}',str(report.get('data_context_sha256'))) + and digest(state)==report['data_context_sha256'] + and state.get('manifest_sha256')==request['data_manifest_sha256'] and state.get('manifest')==manifest + and policy(state.get('sampling',{}).get('boundary_sampling'))==POLICY + and state.get('sampling',{}).get('training_conditioning')=='independent_frame_noise' + and 'boundary_sampling' in state.get('validation_overrides',{}) + and state['validation_overrides']['boundary_sampling'] is None + and state['validation_overrides'].get('training_conditioning')=='mixed_context', + 'Census data context does not bind exact training/strict-heldout policies') + physical={(e['embodiment'],e['episode_id']):e for e in manifest['episodes'] if e['split']=='train'} + subtasks={(e['embodiment'],e['episode_id']):e['subtasks'] for e in index['episodes'] if e['split']=='train'} + rows=report.get('episodes',[]);need(len(rows)==len(physical),'Census must cover every training physical episode only') + seen=set();totals={emb:dict(units=0,eligible_units=0,fallback_units=0,excluded_short_units=0,outward_possible_units=0) for emb,_ in physical} + anchor_totals={emb:dict(all=0,eligible=0,outward_possible=0) for emb,_ in physical} + for row in rows: + key=row['embodiment'],row['episode_id'];need(key in physical and key not in seen,'Census duplicated/changed physical episode') + seen.add(key);ep=physical[key];clock=row['clock_seconds'] + need(row.get('continuity_segments_explicit')==('continuity_segments' in ep) + and row.get('rgb_sha256')==ep['rgb_sha256'] and len(clock)==ep['frames'] + and all(type(t) in (int,float) and math.isfinite(t) for t in clock) + and all(x=2] + expected=[_diagnostic(u,clock,ep.get('continuity_segments',[[0,len(clock)]])) for u in eligible_units] + need(row.get('units')==expected,'Census eligibility diagnostics differ from actual clocks/original annotations') + excluded=[digest(u) for u in units if u['end_idx']-u['start_idx']<2] + need(row.get('excluded_short_unit_sha256')==excluded,'Census changed short-unit exclusion') + counts=totals[key[0]];counts['units']+=len(expected);counts['excluded_short_units']+=len(excluded) + counts['eligible_units']+=sum(x['eligible'] for x in expected) + counts['fallback_units']+=sum(not x['eligible'] for x in expected) + counts['outward_possible_units']+=sum(x['outward_pair_possible'] for x in expected) + for unit,diagnostic in zip(eligible_units,expected): + length=unit['end_idx']-unit['start_idx'];mass=anchor_totals[key[0]] + mass['all']+=length;mass['eligible']+=length*diagnostic['eligible'] + mass['outward_possible']+=length*diagnostic['outward_pair_possible'] + need(seen==set(physical) and report.get('totals_by_embodiment')==totals + and report.get('anchor_weighted_totals_by_embodiment')==anchor_totals + and sum(c['outward_possible_units'] for c in totals.values())>0, + 'Census aggregation/actual outward eligibility missing') + return report + + +def completion(request,smoke,result,preflight): + selected=request.get('boundary_sampling_contract') + need(selected==smoke.get('boundary_sampling_contract'),'Full boundary contract differs from completed smoke') + if selected is None: + need(not (UNCERTAINTY_FIELDS & set(result)) and result.get('subtask_boundaries_verified') is True, + 'Baseline completion cannot admit uncertainty evidence') + fit_rows(result['ranks'],selected=None,world=request['world_size'],per_rank_batch=request['per_rank_batch'], + embodiments=request['interfaces']['pose_dims'],context_sha256=result.get('data_context_sha256')) + return None + need(policy(selected.get('policy'))==POLICY and digest(selected.get('objective'))==digest(OBJECTIVE), + 'Full request changed boundary policy/plain objective') + context_sha=preflight['provenance']['data_context_sha256'] + need(result.get('data_context_sha256')==context_sha and result.get('subtask_boundaries_verified') is False + and result.get('original_subtask_annotation_identity_verified') is True + and result.get('subtask_boundary_verification_policy')==POLICY['policy'], + 'Completion must distinguish uncertainty from strict original-boundary proof') + proof=fit_rows(result['ranks'],selected=POLICY,world=request['world_size'],per_rank_batch=request['per_rank_batch'], + embodiments=request['interfaces']['pose_dims'],context_sha256=context_sha) + need(result.get('boundary_uncertainty')==proof + and result.get('boundary_objective_sha256')==digest(OBJECTIVE) + and result.get('boundary_policy_sha256')==digest(POLICY), + 'Postflight does not bind checked policy/objective/rank evidence') + expected=dict(boundary_policy_sha256=digest(POLICY),boundary_objective_sha256=digest(OBJECTIVE), + boundary_eligibility_report_sha256=selected['eligibility']['report']['sha256']) + need(all(preflight['provenance'].get(k)==v for k,v in expected.items()),'Smoke preflight lacks exact boundary census/policy/objective binding') + return proof 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/checkpoint_inspection.py b/scripts/train/joint/checkpoint_inspection.py new file mode 100644 index 000000000..29bd40362 --- /dev/null +++ b/scripts/train/joint/checkpoint_inspection.py @@ -0,0 +1,309 @@ +"""Allocated CPU inspection of retained native Elmo production checkpoints. + +The existing ModelWrapper, data module, conditional optimizer checker and +checkpoint binding own semantic restoration. No model or training loop is +reimplemented here. This is not a full-completion verifier or GPU restore proof. +""" +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import re +import shutil +import sys + +from scripts.train.joint.contract import bound, need, sha, write_once + + +def scheduled_filename(name): + """Exact production template, with Lightning's configurable metric insertion. + + Census accepts either setting without loading ML. Strict allocated inspection + then runs the maintained formatter against the actual resolved callback and + saved epoch; no checkpoint name is inferred solely from this parser. + """ + match = re.fullmatch(r"(?:epoch_[0-9]{4,}-step_|epoch_epoch=[0-9]{4,}-step_step=)([0-9]{9})\.ckpt", name) + need(match is not None, "Checkpoint does not use the maintained production filename template") + return int(match[1]) + + +def verify_source_binding(request, execution): + """Bind the executing import root and runtime/cache policy to the parent.""" + import importlib.util + + for key in ("source", "source_revision", "source_manifest", "source_manifest_sha256", + "data_manifest", "data_manifest_sha256", "lambda_runtime", "data_cache"): + need(execution.get(key) == request.get(key), "Recovery execution changed parent binding: " + key) + root = Path(request["source"]).resolve(strict=True) + need(Path(os.environ["PIXEL_SOURCE_ROOT"]).resolve(strict=True) == root, + "Actual scientific import environment points to another source") + spec = importlib.util.find_spec("egomimic") + need(spec is not None and spec.origin is not None and Path(spec.origin).resolve().is_relative_to(root), + "Actual scientific package escaped hash-bound source") + for name, module in sys.modules.items(): + if name == "egomimic" or name.startswith("egomimic."): + filename = getattr(module, "__file__", None) + need(filename is not None and Path(filename).resolve().is_relative_to(root), + "Scientific module imported from another source before admission") + + +def attempt_chain(request): + """Enumerate only hash-bound attempts, never guess from run-name globbing.""" + seen, result = set(), [] + identity = {k: request[k] for k in ("run_id", "source_revision", "source_manifest_sha256", + "data_manifest_sha256", "parameters", "target_steps")} + for _ in range(64): + need(request["schema"] == "joint-training-v2" and request["phase"] == "full" + and request["kind"] == "native-subtask-world-model-v1" + and all(request[k] == v for k, v in identity.items()), "Attempt chain changed scientific identity") + output = Path(request["output"]) + need(output.is_absolute() and ".." not in output.parts, "Absolute production output required") + # Scheduled allocations deliberately reuse their output; interruption + # attempts must use a fresh one. A repeated request hash is a cycle. + result.append(request) + previous = request.get("interruption") or request.get("continuation") + need(not (request.get("interruption") and request.get("continuation")), + "Ambiguous scheduled and interrupted parent") + if previous is None: + return result + binding = previous["parent_request"] + need(binding["sha256"] not in seen, "Cycle in immutable production attempt chain") + seen.add(binding["sha256"]) + request = bound(binding) + raise ValueError("Production attempt chain exceeds64; explicit audit required") + + +def checkpoint_census(request): + """Small stat-only census; unreadable/missing collections fail closed.""" + directories, rows, extras = set(), [], [] + for attempt in attempt_chain(request): + need(attempt.get("archive_policy") is None, + "Native interruption v1 requires all checkpoints retained locally; archived recovery is not implemented") + directory = Path(attempt["output"]) / "checkpoints" + if str(directory) in directories: + continue + directories.add(str(directory)) + need(directory.is_dir() and not directory.is_symlink(), "Retained checkpoint collection missing") + for path in sorted(directory.iterdir()): + if path.name == "last.ckpt": + need(path.is_symlink() and path.resolve(strict=True).parent == directory.resolve(), + "last.ckpt is not a retained local milestone link") + continue + stat = path.lstat() + if path.suffix != ".ckpt": + extras.append(dict(path=str(path), size_bytes=stat.st_size, retained=True)) + continue + need(path.is_file() and not path.is_symlink(), "Unqualified production checkpoint") + step = scheduled_filename(path.name) + need(step > 0 and step % 20000 == 0 and step <= request["target_steps"], + "Checkpoint is not an immutable scheduled production milestone") + rows.append(dict(path=str(path), global_step=step, size_bytes=stat.st_size, + mtime_ns=stat.st_mtime_ns, inode=stat.st_ino, device=stat.st_dev)) + need(rows, "No retained production milestone exists; never initialize from smoke/step0") + return dict(collections=sorted(directories), candidates=sorted(rows, key=lambda r: (-r["global_step"], r["path"])), + preserved_noncheckpoint_entries=extras) + + +def check_rng(rng, world): + """Validate CPU/Python/NumPy states using isolated RNG objects, never global draws.""" + import random + import numpy as np + import torch + + need(rng.get("schema") == "mot_rng_state.distributed.v1" and rng.get("world_size") == world + and len(rng.get("ranks", [])) == world, "Complete all-rank RNG state required") + for row in rng["ranks"]: + need(set(row) == {"torch_cpu", "torch_cuda", "numpy", "python", "cuda_scope"} + and row["cuda_scope"] == "current_device" and len(row["torch_cuda"]) == 1, + "Rank RNG scope changed") + need(all(isinstance(x, torch.Tensor) and x.dtype == torch.uint8 and x.ndim == 1 + and x.numel() > 0 for x in (row["torch_cpu"], row["torch_cuda"][0])), + "Malformed saved RNG tensor") + torch.Generator(device="cpu").set_state(row["torch_cpu"]) + np.random.RandomState().set_state(row["numpy"]) + random.Random().setstate(row["python"]) + + +def inspect_checkpoint(path, step, request, cfg, preflight): + """Use only inside an actual allocated inspection producer.""" + import torch + 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 + from egomimic.pl_utils.joint_smoke import weights_digest + from scripts.train.joint.lr_policy import validate, update_lr, verify_scheduler + from scripts.train.joint.native_recovery import checkpoint_science, conditional_record, state_digest + from scripts.train.joint.recovery_check import equal_state + + checkpoint = torch.load(path, map_location="cpu", weights_only=False, mmap=True) + need(checkpoint["global_step"] == step, "Filename and actual checkpoint step disagree") + from lightning.pytorch.callbacks import ModelCheckpoint + callback = cfg.callbacks.model_checkpoint + need(callback.filename == "epoch_{epoch:04d}-step_{step:09d}", "Production checkpoint template changed") + formatter = ModelCheckpoint(filename=callback.filename, + auto_insert_metric_name=callback.get("auto_insert_metric_name", True)) + need(Path(formatter.format_checkpoint_name({"epoch": checkpoint["epoch"], "step": step})).name == Path(path).name, + "Actual maintained checkpoint formatter/epoch does not match the retained filename") + saved = OmegaConf.create(checkpoint["hyper_parameters"]["config_tree"]) + checkpoint_science(saved, cfg) + need(saved.run_provenance.run_id == request["run_id"], "Checkpoint W&B identity differs") + need(checkpoint.get("loops") and checkpoint.get("callbacks"), "Lightning full-state loops/callbacks missing") + need(all(isinstance(v, torch.Tensor) and v.dtype == torch.float32 and torch.isfinite(v).all() + for v in checkpoint["state_dict"].values()), "Invalid/nonfinite FP32 model state") + optimizers, schedulers = checkpoint.get("optimizer_states", []), checkpoint.get("lr_schedulers", []) + need(len(optimizers) == 1 and optimizers[0]["state"], "Missing actual AdamW state") + policy = validate(request["lr_policy"]) + need(all(g["lr"] == update_lr(policy, step) and g["weight_decay"] == 0.01 + for g in optimizers[0]["param_groups"]), "Wrong checkpoint LR/weight decay") + verify_scheduler(policy, schedulers, step) + sampler = checkpoint["PixelProprioDataModule"] + need(sampler["schema"] == "mot-sampling-state-v1" and sampler["consumed_batches"] == step + and sampler["distributed_sampling"]["world_size"] == request["world_size"] + and sampler["distributed_sampling"]["global_batch_size"] == 32, + "Wrong production sampling cursor/topology") + check_rng(checkpoint["mot_rng_state"], request["world_size"]) + data = instantiate(cfg.data) + data.load_state_dict(sampler) + need(equal_state(data.state_dict(), sampler), "Complete sampling settings/state did not restore exactly") + context = data.prepare_context(mode="eval", normalization="prepared", normalizer=None, + restored_state=checkpoint["data_context"]) + need(context.fingerprint() == preflight["provenance"]["data_context_sha256"], + "Checkpoint context differs from actual CPU/smoke preflight") + validate_checkpoint_binding(checkpoint, saved, context) + with checkpoint_construction(): + wrapper = ModelWrapper(config_tree=saved, enable_grad_norm=False).cpu().eval() + context.bind(wrapper.model) + wrapper.data_context = context + need(sum(p.numel() for p in wrapper.model.parameters()) == request["parameters"], + "Strict reconstructed model size differs") + wrapper.load_state_dict(checkpoint["state_dict"], strict=True) + need(equal_state(wrapper.state_dict(), checkpoint["state_dict"]), "Strict weights reload differs") + optimizer = instantiate(cfg.model.optimizer, params=wrapper.parameters()) + if callable(optimizer): + optimizer = optimizer() + optimizer.load_state_dict(optimizers[0]) + scheduler = instantiate(cfg.model.scheduler, optimizer=optimizer) + if scheduler is not None: + scheduler.load_state_dict(schedulers[0]) + need(equal_state(optimizer.state_dict(), optimizers[0]) + and equal_state([] if scheduler is None else [scheduler.state_dict()], schedulers), + "Strict AdamW/scheduler reload differs") + usage = checkpoint["native_conditional_optimizer_usage"] + actual = conditional_record(optimizer, wrapper.named_parameters(), usage["parameter_update_counts"], step) + need(actual == usage, "Conditional parameter-order/count identity differs") + return dict(global_step=step, strict_full_wrapper_optimizer_scheduler_reload=True, + conditional_optimizer_usage_sha256=state_digest(usage), + model_state_sha256=weights_digest(checkpoint["state_dict"]), + loop_state_sha256=state_digest(checkpoint["loops"]), + data_context_sha256=context.fingerprint(), sampler_state_sha256=state_digest(sampler), + rank_rng_sha256=[state_digest(r) for r in checkpoint["mot_rng_state"]["ranks"]], + all_rank_rng_structure_verified=True, actual_gpu_restore_proven=False, + scheduled_validation_resume_proven=False) + + +def inspect_main(request_path): + """Allocated producer: native authentication/ownership before ML imports/hash loads.""" + from scripts.train.joint.native_interruption import command, native_identity, process_identity, destination_probe + + q = json.loads(Path(request_path).read_text()) + need(q["schema"] == "native-elmo-checkpoint-inspection-v1", "Wrong allocated inspection request") + need(os.environ.get("SLURM_JOB_ID", "").isdecimal() + and os.environ.get("SLURM_RESTART_COUNT", "0") == "0" + and not os.environ.get("SLURM_JOB_GPUS") and not os.environ.get("SLURM_STEP_GPUS"), + "Checkpoint inspection must run in a fresh owned CPU allocation") + expected = q["producer_native"] + process_identity(expected, cpu=True) + need(expected["JobId"] == os.environ["SLURM_JOB_ID"] + and expected["NodeList"] == os.environ["SLURMD_NODENAME"] + and expected["NumCPUs"] == "4" and expected["MinMemoryNode"] == "24G" + and expected["TimeLimit"] == "00:20:00" and expected["TresPerNode"] == "", + "Exact4CPU24G20minute bound inspection required") + native = command(["/usr/local/bin/scontrol", "show", "job", os.environ["SLURM_JOB_ID"], "-o", "-dd"]) + native_identity(native, expected, "RUNNING") + request = bound(q["parent_request"]) + need(request["kind"] == "native-subtask-world-model-v1" and request["phase"] == "full", + "Exact native production parent required") + from scripts.train.joint.interruption_admission import recovery_phase_admission + eligibility = recovery_phase_admission(q, request_path) + verify_source_binding(request, q) + # The allocated launcher owns current runtime/cache/source closure checks. + # This producer reopens its actual admitted receipt rather than accepts flags. + admission_path = Path(q["allocated_prerequisites_path"]) + need(admission_path.is_file() and not admission_path.is_symlink(), "Actual allocated prerequisite path missing") + admission = json.loads(admission_path.read_text()) + need(admission["job_id"] == os.environ["SLURM_JOB_ID"] + and admission["request_sha256"] == sha(request_path) + and admission["status"] == "ACTUAL_PREREQUISITES_AND_NATIVE_ALLOCATION_PASS" + , "Current allocated prerequisite/claim admission missing") + native_identity(admission["native"], expected, "RUNNING") + need(q["required_free_bytes"] >= request["reserved_checkpoint_bytes"], + "Inspection must reserve the actual complete production checkpoint budget") + destination = destination_probe(q["production_output"], os.environ["SLURM_JOB_ID"], q["required_free_bytes"]) + output = Path(q["output"]) + need(output.is_absolute() and not output.exists() + and shutil.disk_usage(output.parent).free >= q["required_free_bytes"] > 0, + "Fresh inspection output and actual free space required") + output.mkdir() + payload = ("native inspection " + os.environ["SLURM_JOB_ID"]).encode() + probe = output / "write-fsync-readback.probe" + with probe.open("xb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + need(probe.read_bytes() == payload, "Actual write/fsync/readback failed") + fd = os.open(output, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + need(sha(request["source_manifest"]) == request["source_manifest_sha256"] + and sha(request["data_manifest"]) == request["data_manifest_sha256"], "Source/data manifest changed") + for line in Path(request["source_manifest"]).read_text().splitlines(): + digest, relative = line.split(" ", 1) + need(sha(Path(request["source"]) / relative) == digest, "Checkpoint-bound source changed") + from omegaconf import OmegaConf + from scripts.train.joint.contract import verify_evidence + + _, preflight = verify_evidence(request) + bound(q["resolved"], False) + cfg = OmegaConf.load(q["resolved"]["path"]) + census = checkpoint_census(request) + write_once(output / "census-before.json", census) + results, selected = [], None + for row in census["candidates"]: + if selected is not None and row["global_step"] < selected["global_step"]: + break + path = Path(row["path"]) + digest = sha(path) + try: + proof = inspect_checkpoint(path, row["global_step"], request, cfg, preflight) + except (ValueError, KeyError, RuntimeError, EOFError) as error: + results.append(dict(**row, sha256=digest, status="REJECTED_REQUIRES_DIAGNOSIS", + error_type=type(error).__name__, error=str(error)[:4000])) + continue + valid = dict(**row, sha256=digest, status="CPU_STRICT_INSPECTION_PASS", proof=proof) + results.append(valid) + need(selected is None or selected["sha256"] == digest, "Conflicting healthy checkpoints at the same step") + selected = selected or valid + need(checkpoint_census(request) == census, "Checkpoint collection changed during strict inspection") + # Invalid newer milestones are retained and diagnosed; this tool never + # silently loops a corrupt/nonfinite scientific failure by falling backward. + ready = selected is not None and all(r["status"] == "CPU_STRICT_INSPECTION_PASS" for r in results) + write_once(output / "checkpoint-inspection.json", dict( + schema="native-elmo-checkpoint-inspection-result-v1", + status="INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION" if ready else "BLOCKED_CHECKPOINT_DIAGNOSIS", + job_id=os.environ["SLURM_JOB_ID"], node=os.environ["SLURMD_NODENAME"], + request_sha256=sha(request_path), parent_request=q["parent_request"], run_id=request["run_id"], + created_at=datetime.now(timezone.utc).isoformat(), census=census, results=results, selected=selected, + phase_admission=eligibility, + source_manifest_sha256=request["source_manifest_sha256"], data_manifest_sha256=request["data_manifest_sha256"], + highest_healthy_known=ready, actual_gpu_restore_proven=False, + source_config_context_strictly_checked=ready, write_fsync_readback_sha256=sha(probe), + production_destination=destination)) + + +if __name__ == "__main__": + import sys + inspect_main(sys.argv[1]) diff --git a/scripts/train/joint/continuation.py b/scripts/train/joint/continuation.py new file mode 100644 index 000000000..f8bea7602 --- /dev/null +++ b/scripts/train/joint/continuation.py @@ -0,0 +1,127 @@ +"""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", "model_execution", "action_contract", + "capacity_profile", "noise_comparison", "track_recipe_comparison", + "kind", "interfaces", "subtask_contract", "dataset_contract", +) + + +def allocation_steps(request): + """No continuation is inferred from a directory or W&B's status.""" + need(not (request.get("interruption") and request.get("continuation")), + "Cannot treat an interruption as scheduled boundary completion") + if request.get("interruption") is not None: + from scripts.train.joint.interruption_admission import admission + return admission(request) + 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") + from scripts.train.joint.native_subtasks import selected + if selected(request): + need(all(proof.get(key) is True for key in ( + "native_multimodal", "private_bank_optimizer_steps_verified", + "actual_all_rank_rng_restore_observed", "both_embodiments_optimized", + "both_embodiments_hidden_nan_20step_sample")) + and proof.get("diagnostic_scope") == "scheduled-smoke-checkpoint-only", + "Native scheduled continuation requires actual conditional recovery proof") + return start, stop + + +def configure_allocation(cfg, request, start, stop): + """Apply operational additions only AFTER the existing science comparison.""" + 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.get("interruption") or 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.""" + from scripts.train.joint.production_history import has_interruption, segments + if has_interruption(request): + return [row["csv"] for row in segments(request, request.get("allocation_stop_step", request["target_steps"]))] + 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..b9a2c126e 100644 --- a/scripts/train/joint/contract.py +++ b/scripts/train/joint/contract.py @@ -4,7 +4,10 @@ import hashlib import json import math +import uuid +from scripts.train.joint.lr_policy import validate, update_lr from pathlib import Path +from datetime import datetime def need(condition, message): @@ -12,6 +15,179 @@ def need(condition, message): raise ValueError(message) + +def verify_model_execution(request, cfg): + """Bind native historical image/state execution without changing legacy admission.""" + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + return native_subtasks.verify_model_execution(request, cfg) + 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" + track = isinstance(policy, dict) and policy.get("implementation") == "native_track_joint" + primary = "track" if track else "image" + need(policy == {"implementation": "native_track_joint" if track else "native_image_joint", "compile_mode": "none" if track else "default", + "compile_dynamic": True, "compile_backend": "inductor", + **({"activation_checkpointing": True} if track else {})}, + "Native image execution requires an explicit historical request binding") + image = model["image_model"] + need(image.get("_target_") == ("egomimic.models.track_proprio.TrackTokenCore" if track else "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") + independent = cfg["data"].get("training_conditioning", "mixed_context") == "independent_frame_noise" + need(image.get("require_observed_context", True) is (not independent), + "Independent noise and clean-context model guard disagree") + streams = request["active_streams"] + action = streams == [primary, "proprio", "action"] + need(streams in ((["track", "proprio", "action"],) if track else (["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 track: + need(cfg["data"].get("spatial_stream") == "track" + and image["image_dim"] == 3 and image["max_points"] == 1120 + and image["spatial_grid_size"] == [28, 40] + and image["use_camera_embedding"] is False + and image.get("activation_checkpointing") is True + and cfg["evaluator"].get("spatial_stream") == "track" + and cfg["evaluator"].get("pixel_diagnostics") is None + and cfg["evaluator"].get("prediction_key") == "track_proprio_action_prediction", + "XYZ grid/no-RGB contract changed") + 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") in ( + "native-joint-command-v1", "native-joint-command-v2"), + "Missing explicit native action contract") + native, encoded = contract.get("native_dimensions"), contract.get("encoded_dimensions") + mode = cfg["data"]["mode"] + if contract["schema"] == "native-joint-command-v1": + need(mode == "discrete", "Historical command v1 is discrete-only") + else: + need(mode in ("discrete", "continuous") + and contract.get("temporal_mode") == mode, + "Command temporal mode is not explicitly bound") + group = cfg["data"]["stride"] if mode == "discrete" else 1 + need(isinstance(native, dict) and native and all(type(v) is int and v > 0 for v in native.values()) + 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_capacity_profile(request, cfg): + """Admit the explicit wider image expert without weakening legacy counts.""" + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + return request["expert_parameters"]["image"] + profile = request.get("capacity_profile") + if (request.get("model_execution") or {}).get("implementation") == "native_track_joint": + need(isinstance(profile, dict) and profile.get("schema") == "native-track-capacity-v1", + "Track constructor count binding required") + model = cfg["model"]["pipeline"]["stages"][cfg["model"]["pipeline"]["stage_ids"]["joint_denoiser"]]["model"] + core = model["image_model"] + counts = profile.get("expert_parameters", {}) + need(set(counts) == {"track", "proprio", "action"} + and all(type(v) is int and v > 0 for v in counts.values()) + and counts == request["expert_parameters"] + and sum(counts.values()) == profile["parameters"] == request["parameters"] + and request["expert_widths"] == [768, 256, 256] + and (core["width"], core["heads"], core["depth"], core.get("mlp_ratio", 4)) == (768, 6, 8, 4) + and core["image_dim"] == 3 and core["max_points"] == 1120 + and core["rope_axis_dims"] == [64, 32, 32] + and model["proprio_width"] == model["action_width"] == 256, + "Track capacity differs from bound constructor") + return counts["track"] + if profile is None: + return 88178688 + need(isinstance(profile, dict) + and set(profile) in ({"schema", "image_width", "image_heads", "image_depth", + "parameters", "expert_parameters"}, + {"schema", "image_width", "image_heads", "image_depth", + "image_mlp_ratio", "parameters", "expert_parameters"}) + and profile["schema"] == "native-image-capacity-v1" + and all(type(profile[k]) is int and profile[k] > 0 for k in ( + "image_width", "image_heads", "image_depth", "parameters")) + and profile["image_width"] > 768 and profile["image_depth"] == 8 + and profile["image_width"] == 128 * profile["image_heads"] + and type(profile.get("image_mlp_ratio", 4)) is int + and profile.get("image_mlp_ratio", 4) >= 4 + and isinstance(profile["expert_parameters"], dict) + and set(profile["expert_parameters"]) == {"image", "proprio", "action"} + and all(type(v) is int and v > 0 for v in profile["expert_parameters"].values()) + and sum(profile["expert_parameters"].values()) == profile["parameters"], + "Unknown native capacity profile") + model = cfg["model"]["pipeline"]["stages"][ + cfg["model"]["pipeline"]["stage_ids"]["joint_denoiser"]]["model"] + image = model["image_model"] + need(cfg["data"]["mode"] == request["arm"] == "discrete" + and request["active_streams"] == model["active_streams"] == ["image", "proprio", "action"] + and request["expert_widths"] == [profile["image_width"], 256, 256] + and request["parameters"] == profile["parameters"] + and request["expert_parameters"] == profile["expert_parameters"] + and (image["width"], image["heads"], image["depth"]) == ( + profile["image_width"], profile["image_heads"], profile["image_depth"]) + and image["image_dim"] == 768 and image["rope_axis_dims"] == [64, 32, 32] + and image.get("mlp_ratio", 4) == profile.get("image_mlp_ratio", 4) + and model["proprio_width"] == model["action_width"] == 256, + "Capacity profile differs from actual resolved native model") + return profile["expert_parameters"]["image"] + + +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 +208,123 @@ 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 validate_submission_job_name(request, argv): + """Fail before a new claim when a scheduler label differs from its request.""" + need(isinstance(argv, list) and all(isinstance(v, str) for v in argv), + "Require structured submission argv") + names = [v.split("=", 1)[1] for v in argv if v.startswith("--job-name=")] + need(names == [request["job_name"]], "Submitted JobName differs from exact request") + return names[0] + + +def smoke_submission_job_name(evidence, smoke): + """Allow a recorded native-subtask retry name without changing consumed inputs.""" + attempt = evidence.get("submission_attempt") + if attempt is None: + return smoke["job_name"] + from scripts.train.joint import native_subtasks + need(native_subtasks.selected(smoke), "Retry-name evidence requires native subtask mode") + need(isinstance(attempt, dict) and set(attempt) == {"plan", "accepted", "allocation"}, + "Incomplete bound submission-attempt evidence") + plan, accepted, allocation = (bound(attempt[key]) for key in ("plan", "accepted", "allocation")) + path, digest = evidence["request"]["path"], evidence["request"]["sha256"] + need(plan.get("request_remote") == accepted.get("request") == path + and plan.get("request_sha256") == accepted.get("request_sha256") == digest + and str(accepted.get("job_id")) == str(evidence["job_id"]), + "Retry submission does not bind the accepted job and exact request") + argv = plan.get("argv") + need(isinstance(argv, list) and all(isinstance(v, str) for v in argv) + and argv and Path(argv[0]).name == "sbatch", "Missing exact submitted sbatch argv") + names = [v.split("=", 1)[1] for v in argv if v.startswith("--job-name=")] + need(len(names) == 1 and names[0] == Path(path).stem + and names[0].startswith(smoke["job_name"] + "-"), "Unbound retry job name") + name = names[0] + raw = allocation.get("stdout", "") + need(isinstance(raw, str), "Missing structured native allocation stdout") + native = dict(v.split("=", 1) for v in raw.split() if "=" in v) + owner = smoke.get("quota", {}) + need(allocation.get("rc") == 0 and native.get("JobId") == str(evidence["job_id"]) + and native.get("JobName") == name and native.get("JobState") == "RUNNING" + and native.get("UserId") == f"{owner.get('user')}({owner.get('uid')})" + and native.get("NodeList") == smoke.get("lambda_runtime", {}).get("node") + and native.get("NumCPUs") == "32" and native.get("MinMemoryNode") == "192G" + and f"GRES=gpu:h100:{smoke['world_size']}(IDX:" in raw, + "Retry name lacks exact owned allocated CPU/memory/GPU evidence") + return name + + +def native_smoke_state(evidence, smoke, request): + """Read the declared native format; never manufacture a scontrol receipt. + + Accounting is an alternative for a purged Lambda job. Its original bound + nine-field row remains the evidence; the returned fields are only an + internal admission view. PACE still requires its association attestation. + """ + raw = bound(evidence["native"], False) + kind = evidence.get("native_format", "scontrol-v1") + if kind == "scontrol-v1": + return dict(s.split("=", 1) for s in raw.split() if "=" in s) + need(kind == "sacct-pipe-v1", "Unsupported native completion format") + need(smoke.get("allocation") is None and request.get("allocation") is None, + "Accounting-only completion requires Lambda allocation evidence") + rows = [line for line in raw.splitlines() if line.strip()] + need(len(rows) == 1, "Exactly one native accounting row required") + fields = rows[0].split("|") + need(len(fields) == 9, "Native accounting requires nine fields") + job, name, user, uid, state, exit_code, node, start, end = fields + owner = smoke.get("quota", {}) + need(type(owner.get("uid")) is int and owner["uid"] > 0 + and isinstance(owner.get("user"), str) and bool(owner["user"]) + and request.get("quota", {}).get("uid") == owner["uid"] + and request.get("quota", {}).get("user") == owner["user"], + "Missing matching native accounting owner binding") + need(job == str(evidence["job_id"]) and name == smoke_submission_job_name(evidence, smoke) + and user == owner["user"] and uid == str(owner["uid"]) + and state == "COMPLETED" and exit_code == "0:0" + and node == smoke.get("lambda_runtime", {}).get("node") + == request.get("lambda_runtime", {}).get("node"), + "Wrong native accounting identity or completion") + try: + started, ended = datetime.fromisoformat(start), datetime.fromisoformat(end) + need(ended > started, "Native accounting has no allocated runtime") + except (ValueError, TypeError) as error: + raise ValueError("Invalid native accounting allocation times") from error + return dict(JobId=job, JobName=name, UserId=f"{user}({uid})", + JobState=state, ExitCode=exit_code, NodeList=node) + + def verify_evidence(request): need( request["schema"] == "joint-training-v2" and request["phase"] == "full", @@ -46,30 +339,27 @@ def verify_evidence(request): completion = bound(evidence["completion"]) preflight = bound(evidence["preflight"]) capacity = bound(evidence["capacity"]) - native = dict( - s.split("=", 1) for s in bound(evidence["native"], False).split() if "=" in s - ) + native = native_smoke_state(evidence, smoke, request) need(bound(evidence["host_exit"], False).strip() == "0", "Smoke host failed") need( native["JobId"] == str(evidence["job_id"]) and native["JobState"] == "COMPLETED" and native["ExitCode"] == "0:0" - and native["UserId"].startswith("paphiwetsa3(") - and native["JobName"] == smoke["job_name"], + and native["UserId"].startswith((smoke.get("allocation") or {}).get("user", "paphiwetsa3")+"(") + and native["JobName"] == smoke_submission_job_name(evidence, smoke), "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 +405,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"], @@ -130,6 +422,9 @@ def verify_evidence(request): ) for key in ("source_revision", "source_manifest_sha256", "data_manifest_sha256"): need(preflight["provenance"][key] == request[key], "Unbound provenance: " + key) + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + return native_subtasks.verify_completion(request, smoke, completion, preflight, capacity) need( completion["schema"] == "joint-mot-smoke-v1" and completion["status"] == "PASS" @@ -147,11 +442,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}" @@ -159,6 +450,7 @@ def verify_evidence(request): for metric in ( { "image": "image_normalized_mse", + "track": "track_normalized_mse", "proprio": "proprio_normalized_mse", "action": "action_command_mse", }[name] @@ -172,14 +464,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): @@ -200,7 +492,7 @@ def verify_evidence(request): need( record["active_streams"] == request["active_streams"] and set(record["metrics"]) == metric_names - and len(record["paired_rgb_batch_sha256"]) == 3, + and len(record["paired_track_batch_sha256" if request["active_streams"][0] == "track" else "paired_rgb_batch_sha256"]) == 3, "Wrong active metrics or missing paired RGB proof", ) need( @@ -216,7 +508,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,14 +527,14 @@ 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) return completion, preflight -def compare_configs(production, smoke, request): +def _compare_configs_legacy(production, smoke, request): """Allow only the declared budget/cadence, smoke instrumentation, output and logging changes.""" def relocate(value): @@ -271,6 +563,7 @@ def relocate(value): "mode", "ckpt_path", "modality_ablation", + "track_model", ): need(production.get(key) == smoke.get(key), "Scientific config changed: " + key) full_trainer, smoke_trainer = ( @@ -286,8 +579,9 @@ def relocate(value): copy.deepcopy(smoke["evaluator"]), ) for item in (full_eval, smoke_eval): - item["pixel_diagnostics"].pop("artifact_dir") - item["pixel_diagnostics"].pop("log_wandb") + if item.get("pixel_diagnostics") is not None: + item["pixel_diagnostics"].pop("artifact_dir") + item["pixel_diagnostics"].pop("log_wandb") need(full_eval == smoke_eval, "Evaluator behavior changed") full_cb, smoke_cb = ( copy.deepcopy(production["callbacks"]), @@ -339,3 +633,323 @@ def relocate(value): and production["data"]["batch_size"] * trainer["devices"] == 32, "Fresh step0/globalbatch required", ) + + +def verify_gap_extension(request): + """Bind a fresh completion-gap trial to its immutable ungapped parent. + + Parent admission is cached evidence, not a request to reload its model. + Only the two maintained data modules may change scientific source bytes. + Native smoke-to-full verification remains mandatory for the new recipe. + """ + independent = request.get("noise_comparison") is not None + comparison = request["noise_comparison" if independent else "gap_comparison"] + need(comparison.get("schema") == ("native-independent-frame-noise-extension-v1" if independent else "native-completion-gap-extension-v1") + and comparison.get("scope") == ("all_training_frames" if independent else "completion_only"), + "Sampling comparison requires explicit noise/gap scope") + parent = bound(comparison["parent_request"]) + proof = bound(comparison["parent_admission"]) + need(proof.get("status") == "FULL_ADMISSION_PASS" + and proof.get("request_sha256") == comparison["parent_request"]["sha256"] + and proof.get("provenance", {}).get("source_revision") == parent["source_revision"], + "Gap parent lacks its exact cached native admission") + need(request["phase"] == parent["phase"] == "full" + and request["arm"] == parent["arm"] == "continuous" + and request["run_id"] != parent["run_id"] + and request["source_parent"] == parent["source_revision"], + "Gap comparison must be a fresh direct source child") + topology_changed = any(request[k] != parent[k] for k in ("world_size", "per_rank_batch")) + if topology_changed: + need(comparison.get("topology_control") == "global_batch_matched_rank_rng_unpaired" + and (parent["world_size"], parent["per_rank_batch"]) == (4, 8) + and (request["world_size"], request["per_rank_batch"]) == (2, 16) + and request["global_batch"] == parent["global_batch"] == 32, + "Gap topology change requires explicit two-rank global32 control") + else: + need(comparison.get("topology_control") in (None, "global_batch_matched_rank_rng_unpaired"), + "Unknown gap topology control") + for key in ("data_manifest_sha256", "dataset_contract", "active_streams", "parameters", + "expert_parameters", "expert_widths", "global_batch", "per_rank_batch", + "world_size", "target_steps", "lr_policy", "model_execution", "action_contract"): + if topology_changed and key in ("world_size", "per_rank_batch"): + continue + need(request[key] == parent[key], "Gap comparison changed control: " + key) + old_cfg = bound(comparison["parent_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) + # Hydra derives this operational directory from each run's own output. + # Check both exact owners before normalizing it; all runtime semantics stay + # under the strict comparison below. + derived_outputs = (("runtime", "slurm_signal_checkpoint_dir", "checkpoints"), + ("inference_config", "output_path", "checkpoints/inference-config.yaml")) + for section, key, suffix in derived_outputs: + if key in old_cfg.get(section, {}) or key in new_cfg.get(section, {}): + smoke_request = bound(request["smoke"]["request"]) + for cfg, owner in ((old_cfg, parent), (new_cfg, smoke_request)): + need(cfg.get(section, {}).get(key) == str(Path(owner["output"])/suffix), + "Gap derived artifact path is outside its exact output") + cfg[section][key] = "/"+suffix + if independent: + def strip_opt_in_context_guard(value, *, new): + count = 0 + if isinstance(value, dict): + if "require_observed_context" in value: + need(value.get("_target_") == "egomimic.models.image_proprio.ImageProprioModel" + and value["require_observed_context"] is (not new), + "Unbound model context-guard change") + value.pop("require_observed_context") + count += 1 + for child in value.values(): + count += strip_opt_in_context_guard(child, new=new) + elif isinstance(value, list): + for child in value: + count += strip_opt_in_context_guard(child, new=new) + return count + strip_opt_in_context_guard(old_cfg, new=False) + need(strip_opt_in_context_guard(new_cfg, new=True) > 0, + "Independent-noise model did not explicitly allow all-noised context") + for key in ("model", "seed", "runtime", "world_model_attention", "joint", + "norm_stats", "normalizer", "inference_config", "modality_ablation"): + need(old_cfg.get(key) == new_cfg.get(key), "Gap changed scientific config: " + key) + old_data, new_data = copy.deepcopy(old_cfg["data"]), copy.deepcopy(new_cfg["data"]) + if topology_changed: + for reference, data in ((parent, old_data), (request, new_data)): + need(data.pop("batch_size") == reference["per_rank_batch"] + and data.pop("distributed_world_size") == reference["world_size"], + "Gap resolved topology differs from request") + previous_schedule = old_data.pop("completion_gap_schedule", None) + if independent: + need(previous_schedule == comparison.get("parent_schedule", comparison["schedule"]) + and old_data.pop("training_conditioning", "mixed_context") == "mixed_context" + and old_data.pop("temporal_gap_schedule", None) is None, + "Independent-noise parent must be the exact mixed-mask gap trial") + need(new_data.pop("training_conditioning", None) == "independent_frame_noise" + and new_data.pop("completion_gap_schedule", None) is None, + "Independent-noise trial retained context masks or completion-only gap") + schedule = new_data.pop("temporal_gap_schedule", None) + else: + need(previous_schedule is None, "Gap parent must be ungapped") + schedule = new_data.pop("completion_gap_schedule", None) + need((schedule is None and independent and comparison["schedule"] is None) + or (isinstance(schedule, list) and schedule and schedule == comparison["schedule"]), + "Gap schedule differs from authorized stages") + last_step, last_gap = -1, 0 + for stage in schedule or []: + need(isinstance(stage, dict) and set(stage) == {"start_step", "max_gap_seconds"} + and type(stage["start_step"]) is int and stage["start_step"] > last_step + and type(stage["max_gap_seconds"]) in (int, float) + and math.isfinite(stage["max_gap_seconds"]) + and stage["max_gap_seconds"] > 0 and stage["max_gap_seconds"] >= last_gap, + "Invalid gap curriculum stage") + last_step, last_gap = stage["start_step"], stage["max_gap_seconds"] + need(schedule is None or (schedule[0]["start_step"] == 0 and schedule[-1]["start_step"] < request["target_steps"]), + "Gap curriculum must start at zero and finish inside training budget") + old_data.pop("manifest_path"); new_data.pop("manifest_path") + need(old_data == new_data, "Gap changed data beyond completion schedule") + old_runtime, new_runtime = copy.deepcopy(parent["lambda_runtime"]), copy.deepcopy(request["lambda_runtime"]) + for value in (old_runtime, new_runtime): + value.pop("node"); value.pop("mounts") + need(old_runtime == new_runtime, "Gap changed native runtime semantics") + old_manifest = bound(comparison["parent_source_manifest"], False) + need(hashlib.sha256(old_manifest.encode()).hexdigest() == parent["source_manifest_sha256"], + "Gap parent source manifest identity changed") + def scientific(text): + return {name: digest for digest, name in (line.split(" ", 1) for line in text.splitlines()) + if name.startswith("egomimic/")} + old_files = scientific(old_manifest) + new_files = scientific(Path(request["source_manifest"]).read_text()) + changed = {name for name in set(old_files) | set(new_files) + if old_files.get(name) != new_files.get(name)} + allowed = {"egomimic/rldb/pixel_proprio.py", "egomimic/rldb/pixel_proprio_datamodule.py"} + if independent: + allowed.add("egomimic/models/image_proprio.py") + need(changed == allowed, "Sampling source delta includes unrelated scientific changes") + return dict(status="PASS", scope="all_training_frames" if independent else "completion_only", schedule=schedule, + fairness="FAIR WITH CAVEATS", + caveats=(["Independent timestamp strengths replace the context-mask mixture; the same temporal gap stages now apply to every training clip.", + "Training targets and contexts differ; compare common held-out metrics and actual rollouts, not raw training loss."] if independent else + ["Dedicated node differs; compare optimizer updates and metrics, not raw throughput.", + "Completion curriculum intentionally leaves forecast sampling unchanged."]) + + (["Two ranks with batch16 preserve global32; rank noise draws and throughput are not paired."] + if topology_changed else [])) + + +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_capacity_extension(request, completion): + """Compare the wider discrete model with its actual three-stream smoke. + + Existing source, labels, optimization and deployment semantics must match. + Initialization and stochastic draws are recorded, not claimed identical + between networks of different shapes. + """ + comparison = request["capacity_comparison"] + need(comparison.get("schema") == "native-capacity-extension-v1", "Wrong capacity comparison") + parent = bound(comparison["parent_request"]) + baseline, _ = verify_evidence(parent) + need(parent.get("capacity_profile") is None and request.get("capacity_profile") is not None, + "Capacity comparison requires original and wider profiles") + for key in ("source_revision", "source_manifest_sha256", "source_archive_sha256", + "data_manifest_sha256", "world_size", "per_rank_batch", "global_batch", + "target_steps", "arm", "model_execution", "lr_policy", "active_streams", + "experiment", "action_contract"): + need(request[key] == parent[key], "Capacity comparison changed control: " + key) + need(request["arm"] == "discrete" and request["experiment"] == "image_proprio_action" + and request["run_id"] != parent["run_id"], "Wrong capacity comparison identities") + roots = [Path(r[k]) for r in (parent, request) for k in ("output", "local_logs")] + need(all(p.is_absolute() for p in roots) + and all(a != b and a not in b.parents and b not in a.parents + for i, a in enumerate(roots) for b in roots[i+1:]), + "Overlapping capacity output/log roots") + need(parent["smoke"]["job_id"] != request["smoke"]["job_id"], "Shared capacity GPU smoke") + from omegaconf import OmegaConf + normalized = [] + profile = request["capacity_profile"] + for reference, width, heads, ratio in ((parent, 768, 6, 4), + (request, profile["image_width"], profile["image_heads"], profile.get("image_mlp_ratio", 4))): + cfg = OmegaConf.to_container(OmegaConf.create(bound(reference["smoke"]["resolved"], False)), resolve=True) + need(verify_model_execution(reference, cfg) == "visual_state_action", "Wrong capacity graph") + expected_image = verify_capacity_profile(reference, cfg) + need(reference["expert_parameters"]["image"] == expected_image + and cfg["ckpt_path"] is None and cfg["seed"] == cfg["data"]["seed"] == 42, + "Capacity comparison is not a fresh seed42 trial") + smoke_request = bound(reference["smoke"]["request"]) + roots = ((smoke_request["output"], ""), (smoke_request["local_logs"], "")) + image_count = 0 + def normalize(value): + nonlocal image_count + if isinstance(value, dict): + value = {k: normalize(v) for k, v in value.items()} + if value.get("_target_") == "egomimic.models.image_proprio.ImageProprioModel": + need((value["width"], value["heads"], value["depth"]) == (width, heads, 8), + "Unexpected image scaling") + need(value.get("mlp_ratio", 4) == ratio, "Unexpected image feed-forward scaling") + value["width"], value["heads"] = "", "" + if "mlp_ratio" in value: + value["mlp_ratio"] = "" + image_count += 1 + return value + if isinstance(value, list): + return [normalize(v) for v in value] + if isinstance(value, str): + for root, token in roots: + if value == root or value.startswith(root + "/"): + return token + value[len(root):] + return value + normalized.append(normalize(cfg)) + need(image_count > 0, "Missing actual image constructor") + need(normalized[0] == normalized[1], "Capacity comparison changed unrelated resolved configuration") + need(len(completion["ranks"]) == len(baseline["ranks"]) == request["world_size"], + "Wrong capacity smoke rank count") + 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 + + +def compare_configs(production, smoke, request): + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + return native_subtasks.compare_configs(production, smoke, request, _compare_configs_legacy) + return _compare_configs_legacy(production, smoke, request) diff --git a/scripts/train/joint/dataset_contract.py b/scripts/train/joint/dataset_contract.py new file mode 100644 index 000000000..77cb550ea --- /dev/null +++ b/scripts/train/joint/dataset_contract.py @@ -0,0 +1,38 @@ +"""Manifest-bound prepared pixel corpus admission; legacy requests remain exact.""" +from collections import Counter +import json +from pathlib import Path +from scripts.train.joint.contract import need, sha + + +def verify_dataset(request): + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + return native_subtasks.verify_dataset(request) + policy = request.get("dataset_contract") + if policy is None: + return {"train": 450, "val": 50}, 100 + need(policy.get("schema") == "prepared-pixel-dataset-v1", "Unknown dataset contract") + manifest_path = Path(request["data_manifest"]) + need(sha(manifest_path) == request["data_manifest_sha256"] == policy["manifest_sha256"], + "Prepared dataset manifest changed") + manifest = json.loads(manifest_path.read_text()) + need(manifest["schema"] == "pixel-proprio-prepared-v1", "Wrong prepared pixel schema") + episodes = manifest["episodes"] + counts = dict(Counter(ep["split"] for ep in episodes)) + need(set(counts) == {"train", "val"} and counts == policy["splits"] + and sum(counts.values()) == policy["episodes"] and min(counts.values()) > 0, + "Incomplete prepared split") + physical = [(ep["native_source"]["source_demo_sha256"], + ep["native_source"]["episode"]) for ep in episodes] + need(len(set(physical)) == len(physical) + and len({(ep["embodiment"], ep["episode_id"]) for ep in episodes}) == len(episodes), + "Duplicate physical episode or train/validation overlap") + norm = manifest["normalization"] + need(norm["scope"] == "train_only" and norm["sha256"] == policy["normalization_sha256"] + and sha(manifest_path.parent / norm["path"]) == norm["sha256"], + "Train-only normalization changed") + anchors = policy["validation_anchors_per_episode"] + need(type(anchors) is int and anchors > 0 and policy["validation_batches"] == counts["val"] * anchors, + "Unbound validation anchor count") + return counts, policy["validation_batches"] diff --git a/scripts/train/joint/interruption_admission.py b/scripts/train/joint/interruption_admission.py new file mode 100644 index 000000000..d3e0b8cf9 --- /dev/null +++ b/scripts/train/joint/interruption_admission.py @@ -0,0 +1,338 @@ +"""Finalized, current-intent-bound native Elmo production interruption admission. + +No submission/cancellation occurs here. Producer completion is checked from raw +native records and host exits. The live launch gate repeats intent, queues and +all named claim directories; a local test or a producer's own PASS is insufficient. +""" +import copy +from datetime import datetime, timezone +import hashlib +import json +import os +from pathlib import Path +import re + +from scripts.train.joint.contract import bound, need, sha, write_once +from scripts.train.joint.native_interruption import ( + NATIVE_KEYS, command, fields, native_identity, classify_failure, check_steps, process_identity, destination_probe, +) + + +CHANGING_ATTEMPT_FIELDS = {"job_name", "output", "local_logs", "overrides", "operational_service", + "interruption", "continuation", "allocation_stop_step", "allocation_recovery", "allocation_boundaries"} + + +def fingerprint(request): + """Avoid circular finalizer/request hashes by omitting only the admission bundle.""" + value = copy.deepcopy(request) + value["interruption"].pop("admission", None) + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()).hexdigest() + + +def compatible(parent, child): + need(parent["schema"] == child["schema"] == "joint-training-v2" + and parent["phase"] == child["phase"] == "full" + and parent["kind"] == child["kind"] == "native-subtask-world-model-v1" + and child["parameters"] == 314187064 and child["world_size"] == 4, + "Exact production native Elmo314M identity required") + need({k: v for k, v in parent.items() if k not in CHANGING_ATTEMPT_FIELDS} + == {k: v for k, v in child.items() if k not in CHANGING_ATTEMPT_FIELDS}, + "Interrupted allocation changed science/runtime/cache/total budget") + need(child.get("continuation") is None and child.get("allocation_stop_step", child["target_steps"]) == child["target_steps"], + "Interruption v1 resumes to the original total target; no fake scheduled boundary") + need(child["output"] != parent["output"] and child["local_logs"] != parent["local_logs"] + and child["job_name"] != parent["job_name"], "Fresh immutable output/log/job namespaces required") + + +def completed_producer(bundle, result_schema, status): + """Successful new recovery phases need complete service-verified observers.""" + from scripts.train.joint.phase_observer_proof import verify_observer, MODES, bound as proof_bound, read_json, text + admitted = { + "native-elmo-interruption-evidence-v1": ("capture", "CAPTURED_REQUIRES_NATIVE_PRODUCER_COMPLETION"), + "native-elmo-checkpoint-inspection-result-v1": ("inspect", "INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION"), + "native-elmo-production-recovery-diagnostic-v1": ("diagnostic", "DIAGNOSTIC_COMPLETE_REQUIRES_HOST_AND_NATIVE_SUCCESS"), + "native-elmo-interruption-finalized-v1": ("finalize", "FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS"), + } + need(result_schema in admitted and status == admitted[result_schema][1], "Unknown recovery producer result/status contract") + mode = admitted[result_schema][0] + q, result = read_json(proof_bound(bundle["request"])), read_json(proof_bound(bundle["result"])) + need(q["schema"] == MODES[mode][0], "Recovery producer request/result phase differs") + filenames = dict(capture="interruption-capture.json", inspect="checkpoint-inspection.json", + diagnostic="production-recovery-diagnostic.json", finalize="interruption-finalized.json") + need(bundle["result"]["path"] == str(Path(q["output"]) / filenames[mode]), + "Producer result escaped its canonical allocated output") + expected = q["producer_native"] + proof = verify_observer(bundle["observer"], producer_request=bundle["request"], + producer_native=expected, terminal_native=bundle["native"]) + need(proof["spec"]["mode"] == mode and proof["request"] == q + and proof["running_admission_status"] == "VALIDATED", "Successful phase lacks actual matching RUNNING provenance") + native = native_identity(text(proof_bound(bundle["native"])), expected, "COMPLETED") + host_path = proof["root"] / "status" / (mode + "-" + expected["JobId"] + ".exit") + need(bundle["host_exit"]["path"] == str(host_path) and native["ExitCode"] == "0:0" + and text(proof_bound(bundle["host_exit"])).strip() == "0", "Producer did not actually exit canonical host0/nativeCOMPLETED0:0") + jobs = [result[k] for k in ("job_id", "producer_job_id") if k in result] + need(jobs and all(job == expected["JobId"] for job in jobs) + and result["request_sha256"] == bundle["request"]["sha256"] + and result["schema"] == result_schema and result["status"] == status, + "Completed native producer does not bind the exact request/result") + return q, result + + +def current_intent(path, run_id, expected_sha=None): + path = Path(path) + need(path.is_absolute() and path.is_file() and not path.is_symlink(), "Current human intent file missing") + if expected_sha is not None: + need(sha(path) == expected_sha, "Human intent changed since finalization; re-evaluate before restart") + intent = json.loads(path.read_text()) + need(intent["schema"] == "native-elmo-current-intent-v1" and intent["authority"] == "human" + and intent["status"] == "ACTIVE" and run_id in intent["authorized_recovery_run_ids"] + and run_id not in intent["retired_run_ids"] and not intent["hold_recovery"], + "Recovery is paused, held, retired, or outside current human authorization") + need(intent.get("source_reference") and intent.get("updated_at"), "Traceable current human instruction required") + return sha(path) + + +def quiescence(spec, run_id, *, current_job=None, current_request_sha=None): + """Check every registered claim directory and the fresh complete owner queue.""" + need(spec["user"] and spec["claim_roots"] and spec["scientific_job_names"] + and spec["scientific_name_prefixes"], "Complete campaign claim/queue scope required") + need(all(len(p) >= 16 for p in spec["scientific_name_prefixes"]), "Scientific writer prefix too broad") + queue = command(["/usr/local/bin/squeue", "--noheader", "--user", spec["user"], "--format=%i|%j|%T|%k"]) + # Do not target old IDs: Slurm may have purged them after our genuine native + # observer retained their terminal records. Still query the complete live + # owner step inventory and reject every surviving step of each prior writer. + steps = command(["/usr/local/bin/squeue", "--steps", "--noheader", "--user", spec["user"], "--format=%i|%T"]) + from scripts.train.joint.original_observer import strict_step_rows + strict_step_rows(steps) + for line in queue.splitlines(): + if not line.strip(): + continue + row = [v.strip() for v in line.split("|")] + need(len(row) == 4, "Unparseable actual owner queue") + job, name, state, comment = row + owned = name in spec["scientific_job_names"] or any(name.startswith(p) for p in spec["scientific_name_prefixes"]) or comment == run_id + if owned: + need(job == current_job and state == "RUNNING", "Other pending/active scientific writer: " + job) + from scripts.train.joint.phase_observer_proof import read_json, text, physical, check_owned_steps + seen, evidence = set(), [] + for item in spec["claim_roots"]: + root = physical(item["path"]) + need(root.is_absolute() and root.is_dir() and not root.is_symlink(), "Registered claim directory unavailable") + for claim in sorted(root.iterdir()): + need(claim.is_dir() and not claim.is_symlink() and re.fullmatch(r"[0-9a-f]{64}", claim.name), + "Unrecognized claim entry requires exact inspection") + prepared = read_json(claim / "PREPARED.json") + need(prepared["request_sha256"] == claim.name, "Claim/request identity mismatch") + job = text(claim / "PARENT_JOB_ID").strip() + need(job.isdecimal() and job not in seen, "Unknown/duplicate native submission acceptance") + seen.add(job) + if job == current_job: + need(claim.name == current_request_sha, "Current native writer belongs to another request claim") + evidence.append(dict(path=str(claim), job_id=job, current=True)) + continue + from scripts.train.joint.phase_observer_proof import verify_claim_terminal + terminal_proof = verify_claim_terminal(claim, item["terminal_directory"], run_id) + need(terminal_proof["job_id"] == job, "Terminal observer belongs to another accepted claim") + check_owned_steps(steps, job) + evidence.append(dict(path=str(claim), **terminal_proof)) + if current_job is not None: + need(current_job in seen, "Current allocation is not registered in the complete claim census") + return dict(queue=queue, owned_steps=steps, claims=evidence, checked_at=datetime.now(timezone.utc).isoformat()) + + +def validate_failure(bundle, parent_binding): + q, result = completed_producer(bundle, "native-elmo-interruption-evidence-v1", + "CAPTURED_REQUIRES_NATIVE_PRODUCER_COMPLETION") + need(q["parent_request"] == result["parent_request"] == parent_binding, + "Failure evidence belongs to another production parent") + raw = bound(result["evidence"]["parent-native.txt"], False) + from scripts.train.joint.original_observer import resolve_terminal, validate_steps + resolved_raw, provenance = resolve_terminal(q, bound(result["evidence"]["parent-native-query.json"])) + need(resolved_raw == raw and provenance == bound(result["evidence"]["parent-terminal-provenance.json"]) + and provenance["source"] == result["parent_terminal_source"], "Terminal provenance changed after capture") + state = fields(raw).get("JobState") + terminal = native_identity(raw, q["parent_native"], state) + running = bound(q["parent_running_admission"]) + need(running["job_id"] == q["parent_job_id"] and running["request_sha256"] == parent_binding["sha256"], + "Parent actual RUNNING claim differs") + native_identity(running["native"], q["parent_native"], "RUNNING") + steps = validate_steps(bound(result["evidence"]["parent-step-query.json"]), q["parent_job_id"]) + need(steps == bound(result["evidence"]["parent-live-steps.txt"], False), "Preserved complete step evidence differs") + tails = "\n".join(bound(value, False) for name, value in result["evidence"].items() if name.endswith(".tail64.txt")) + need(tails and classify_failure(terminal, tails) == result["reason"], "Actual failure classification differs") + from scripts.train.joint.native_interruption import parent_csv_paths, newest_loss + parent = bound(parent_binding) + need({name: value["path"] for name, value in result["loss_sources"].items()} == parent_csv_paths(parent), + "Failure loss evidence came from another production log namespace") + snapshots = {name: {part: bound(result["evidence"]["loss-" + name + "-" + part + ".txt"], False) + for part in ("header", "tail")} + for name, value in result["loss_sources"].items() if value["exists"]} + need(newest_loss(snapshots) == result["newest_numeric_loss"], "Preserved newest numeric loss evidence differs") + return result + + +def resume_native_contract(expected, parent_native, child): + """Bind future immutable allocation shape before its job ID exists.""" + need(set(expected) == (set(NATIVE_KEYS) | {"NodeList", "TresPerNode"}) - {"JobId"}, + "Complete future production native contract required") + stable = set(expected) - {"JobName", "Command"} + need(all(expected[key] == parent_native[key] for key in stable) + and expected["JobName"] == child["job_name"] + and expected["Command"] != parent_native["Command"] and Path(expected["Command"]).is_absolute() + and expected["NumCPUs"] == "32" and expected["MinMemoryNode"] == "192G" + and expected["TresPerNode"] == "gres/gpu:h100:4", + "Resume native resources/owner/node differ or command/job namespace is reused") + return expected + + +def recovery_phase_admission(q, request_path, *, scientific_writer=False): + """Recheck real failure, human intent and all writers before costly phases.""" + parent = bound(q["parent_request"]) + intent = current_intent(q["current_intent_path"], parent["run_id"]) + failure = validate_failure(q["failure"], q["parent_request"]) + options = dict(current_job=os.environ["SLURM_JOB_ID"], current_request_sha=sha(request_path)) if scientific_writer else {} + live = quiescence(q["quiescence"], parent["run_id"], **options) + return dict(current_intent_sha256=intent, failure_reason=failure["reason"], quiescence=live) + + +def finalize(request_path): + """Separate allocated CPU finalizer; emits no job or production checkpoint.""" + q = json.loads(Path(request_path).read_text()) + need(q["schema"] == "native-elmo-interruption-finalizer-v1", "Wrong native interruption finalizer request") + producer = q["producer_native"] + process_identity(producer, cpu=True) + need(producer["JobId"] == os.environ.get("SLURM_JOB_ID") + and producer["NodeList"] == os.environ.get("SLURMD_NODENAME") + and producer["NumCPUs"] == "4" and producer["MinMemoryNode"] == "24G" + and producer["TimeLimit"] == "00:20:00" and producer["TresPerNode"] == "" + and not os.environ.get("SLURM_JOB_GPUS") and not os.environ.get("SLURM_STEP_GPUS"), + "Separate real4CPU24G20min finalizer required") + raw = command(["/usr/local/bin/scontrol", "show", "job", producer["JobId"], "-o", "-dd"]) + native_identity(raw, producer, "RUNNING") + parent, child = bound(q["parent_request"]), bound(q["resume_request_template"]) + from scripts.train.joint.checkpoint_inspection import verify_source_binding + verify_source_binding(parent, q) + need(child["interruption"]["parent_request"] == q["parent_request"] + and child["interruption"].get("admission") is None, "Finalizer requires an unadmitted immutable template") + compatible(parent, child) + intent_sha = current_intent(q["current_intent_path"], parent["run_id"]) + failure = validate_failure(q["failure"], q["parent_request"]) + expected_resume = resume_native_contract(q["resume_native_contract"], failure["parent_native"], child) + inspect_q, inspection = completed_producer(q["inspection"], "native-elmo-checkpoint-inspection-result-v1", + "INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION") + need(inspect_q["parent_request"] == inspection["parent_request"] == q["parent_request"] + and inspection["highest_healthy_known"] is True, "Highest healthy exact-parent production inspection missing") + selected = inspection["selected"] + need(selected["status"] == "CPU_STRICT_INSPECTION_PASS" + and selected["global_step"] == max(row["global_step"] for row in inspection["census"]["candidates"]) + and any(all(selected[k] == row[k] for k in ("path", "global_step", "size_bytes", "mtime_ns", "inode", "device")) + for row in inspection["census"]["candidates"]) + and all(row["status"] == "CPU_STRICT_INSPECTION_PASS" for row in inspection["results"]), + "Selected checkpoint is not the highest strictly inspected retained production milestone") + need(inspect_q["production_output"] == inspection["production_destination"]["output"] == child["output"] + and inspection["production_destination"]["actual_write_fsync_readback"] is True + and sha(inspection["production_destination"]["probe_path"]) == inspection["production_destination"]["probe_sha256"], + "Allocated inspection did not prove the actual production destination") + need(child["interruption"]["checkpoint"] == {k: selected[k] for k in ("path", "sha256", "global_step", "size_bytes")} + and sha(selected["path"]) == selected["sha256"], "Resume selected a different or changed production checkpoint") + from scripts.train.joint.checkpoint_inspection import checkpoint_census + need(checkpoint_census(parent) == inspection["census"], "Retained production census changed after inspection") + diagnostic_q, diagnostic = completed_producer(q["diagnostic"], "native-elmo-production-recovery-diagnostic-v1", + "DIAGNOSTIC_COMPLETE_REQUIRES_HOST_AND_NATIVE_SUCCESS") + need(diagnostic_q["parent_request"] == q["parent_request"] + and diagnostic_q["inspection"] == q["inspection"]["result"] + and diagnostic["production_checkpoint"] == child["interruption"]["checkpoint"], + "Actual GPU diagnostic used another parent/checkpoint") + expected_rank_paths = {str(Path(diagnostic_q["output"]) / "proof" / f"resume-{first}-{last}-rank{rank}.json") + for first, last in ((selected["global_step"], selected["global_step"] + 2), + (selected["global_step"] + 2, selected["global_step"] + 6)) + for rank in range(4)} + need(len(diagnostic["rank_proofs"]) == 8 + and {b["path"] for b in diagnostic["rank_proofs"]} == expected_rank_paths, + "Completed diagnostic does not bind the exact eight actual rank receipts") + for binding in diagnostic["rank_proofs"]: + bound(binding) + need(child["operational_service"] == q["operational_service"] == diagnostic_q["operational_service"] + == inspect_q["operational_service"], "Finalizer/inspection/GPU/resume operational implementations differ") + from scripts.train.joint.production_recovery import read_phase_proofs, validate_phase_proofs + from scripts.train.joint.lr_policy import update_lr + from scripts.train.joint.production_history import retained_milestones, scheduled_validation + + start = selected["global_step"] + rows = read_phase_proofs(diagnostic_q["output"], start, parent["target_steps"], 4, diagnostic["job_id"]) + flags = validate_phase_proofs(rows, selected, diagnostic.get("allocation_attestation")) + for phase in rows: + for row in phase: + need(row["optimizer_update_lrs"] == [update_lr(parent["lr_policy"], i) + for i in range(row["start_step"], row["global_step"])], "Actual diagnostic LR sequence differs") + metric_names = set(rows[-1][0]["validation_metrics"]) + _, checkpoints = retained_milestones(parent, start) + _, history = scheduled_validation(parent, start, metric_names) + live = quiescence(q["quiescence"], parent["run_id"]) + need(not Path(child["output"]).exists() and not Path(child["local_logs"]).exists(), + "Fresh resume output/log namespace already exists") + destination = destination_probe(child["output"], producer["JobId"], child["reserved_checkpoint_bytes"]) + output = Path(q["output"]) + output.mkdir() + write_once(output / "interruption-finalized.json", dict( + schema="native-elmo-interruption-finalized-v1", status="FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS", + job_id=producer["JobId"], request_sha256=sha(request_path), parent_request=q["parent_request"], + run_id=parent["run_id"], resume_request_fingerprint=fingerprint(child), + checkpoint=child["interruption"]["checkpoint"], failure_reason=failure["reason"], + current_intent_path=q["current_intent_path"], current_intent_sha256=intent_sha, + retained_checkpoint_history=checkpoints, committed_validation_history=history, + production_destination=destination, + resume_native_contract=expected_resume, + quiescence=live, diagnostic=q["diagnostic"], inspection=q["inspection"], failure=q["failure"], + actual_updates=6, original_total_target=parent["target_steps"], + last_verified_production_step=start, + lost_unsaved_updates=None, + lost_unsaved_updates_status="Requires actual final worker-update evidence; checkpoint step is not crash progress", + finalized_at=datetime.now(timezone.utc).isoformat(), **flags)) + + +def admission(request): + """Static immutable chain check, safe for later full-completion verification.""" + interruption = request["interruption"] + parent = bound(interruption["parent_request"]) + compatible(parent, request) + q, result = completed_producer(interruption["admission"], "native-elmo-interruption-finalized-v1", + "FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS") + need(q["parent_request"] == interruption["parent_request"] == result["parent_request"] + and result["resume_request_fingerprint"] == fingerprint(request) + and result["checkpoint"] == interruption["checkpoint"] + and result["run_id"] == request["run_id"] and result["original_total_target"] == request["target_steps"] + and result["actual_updates"] == 6 and result["diagnostic_scope"] == "production-checkpoint-only", + "Finalized recovery admission differs from this exact resume request") + start = interruption["checkpoint"]["global_step"] + need(type(start) is int and 0 < start < request["target_steps"] and start % 20000 == 0 + and sha(interruption["checkpoint"]["path"]) == interruption["checkpoint"]["sha256"], + "Highest inspected production checkpoint changed") + return start, request["target_steps"] + + +def live_admission(request, request_path): + """Called once in actual full RUNNING allocation before any training starts.""" + admission(request) + q, result = completed_producer(request["interruption"]["admission"], "native-elmo-interruption-finalized-v1", + "FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS") + need(q["resume_native_contract"] == result["resume_native_contract"], "Finalized production native contract differs") + expected = dict(result["resume_native_contract"], JobId=os.environ["SLURM_JOB_ID"]) + process_identity(expected) + native = command(["/usr/local/bin/scontrol", "show", "job", expected["JobId"], "-o", "-dd"]) + native_identity(native, expected, "RUNNING") + current_intent(result["current_intent_path"], request["run_id"], result["current_intent_sha256"]) + need(not Path(request["output"]).exists() and not Path(request["local_logs"]).exists(), + "Resume attempt would overwrite an existing output/log namespace") + need(os.environ.get("SLURM_RESTART_COUNT", "0") == "0", "Fresh immutable no-requeue attempt required") + live = quiescence(q["quiescence"], request["run_id"], current_job=os.environ["SLURM_JOB_ID"], + current_request_sha=sha(request_path)) + live["production_destination"] = destination_probe(request["output"], os.environ["SLURM_JOB_ID"], + request["reserved_checkpoint_bytes"]) + live["actual_production_native"] = native + return live + + +if __name__ == "__main__": + import sys + finalize(sys.argv[1]) diff --git a/scripts/train/joint/lambda_full.sbatch b/scripts/train/joint/lambda_full.sbatch new file mode 100644 index 000000000..5dcc4f736 --- /dev/null +++ b/scripts/train/joint/lambda_full.sbatch @@ -0,0 +1,75 @@ +#!/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" +# Bind broad writable parents before the verified read-only source child. +# Otherwise a parent node-scratch mount hides the virtual source workdir. +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,$JOINT_NODE_SCRATCH:$JOINT_NODE_SCRATCH:rw,$JOINT_JOB_TMP:$JOINT_JOB_TMP:rw,$mounts$credentials_mount,$source_cache/source:$PIXEL_SOURCE_ROOT:ro" \ + --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..b73be2ff7 --- /dev/null +++ b/scripts/train/joint/lambda_inside.sh @@ -0,0 +1,28 @@ +#!/bin/bash +set -euo pipefail +root=${1:?Operations root} +request=${2:?Exact request} +# Batch --export=NONE can suppress ordinary host variables in the srun step. +# Resolve the source from the same hash-bound request inside the container. +export PIXEL_SOURCE_ROOT +PIXEL_SOURCE_ROOT=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["source"])' "$request") +[[ "$PIXEL_SOURCE_ROOT" == /* && -d "$PIXEL_SOURCE_ROOT" ]] +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/native_elmo_ops/FILES.sha256 b/scripts/train/joint/native_elmo_ops/FILES.sha256 new file mode 100644 index 000000000..e10797e7e --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/FILES.sha256 @@ -0,0 +1,23 @@ +762f5e2cc0cc10d481f7bd873779df0178430fc10357b2cc93ee4a650dd353f2 allocated.sbatch +1707fdcd216eb5fbcc71a17752d83ba01912b426c9feec6709812b60ffedd382 controller.py +149a9e5f54a62d5df83e723f5fe46571f67feffd5783d2677836a154d8461388 controller.sbatch +6541745310153c8de0ef95d1c26f6d27c20fd7892e3b2236ad73244c0644ff35 dispatch_after_runtime.py +822010203099bdad23c6b92d57bb265ecdf326c27333130f3ead2f5f67dba293 execution_spec.py +06b3cab135f90f12291708bdbe6d43ee6fbf054888428da2a8a436077621a64f fa3-binding.json +33b8ccffb9d5253d3d18471a15f34a28263281f06cef6f3290b1bbd71f8b0b87 gpu_health_r3.py +e1f3701fb5dc6e605b759a7bbc3d655582c207c2be7f19a3aaf64201a0a71b02 gpu_health_set.py +7e778c22f15beaa90bf31f17f5d471e6cded0109969b81a6830f70bc93bc7801 host_guard.py +174871f1fee48e2e175474f532de46a6c901afae2341df6bb35a9388f78a984f inside.sh +4d5695db4e703a2ac1ea29418d34261ffe06c040fc367e5af960788715944d19 lambda_runtime.py +50d3f31641ebaf494542d63cf43380486ca9ed50bc315c846cee354a2e91ce32 lambda_stage_source.sh +2a9ef29e84e2d282f94538a6eab45bc778d1ef614ba7a29a940235a9a8659554 launcher_common.sh +41a3af46c6ce2acafc3ce71219815925903a6543e0559c6da1d32e0fe8c45d60 native_client.py +5b4604db7643907444cafaa68d4a5c93f8f2c88618af8d947c8bb6a0d6afbca7 native_scontrol.sh +7e8e24b21ac12c8d8eba73c90de5cbf01bc8e4d107aa6b1e13b00510ab9a161b native_squeue.sh +d7443dc3364f13f5ed36a83c3b74b5b9f9fb6bbb0eb67588ea95f9b5230f21fb observe-phase.sbatch +f3c5eba2d20d0f9d029c467a7bdcee7903dc93e0daf26532370be2f6e9e0c3cc phase_observer.py +fffd65ed5b85b5ac0178d02b13b8e746a69dc29d143d909083ca4f13fb3ea2ee runtime-binding.json +9be61a71a8181eb30c6e3ad2c30721b8ce42c92696c3119395f04d8f41b1a4b4 runtime_bootstrap.py +4bffab76a9d0d6231fd169f7bd70aca2317527f8a163f901841a20884abbb987 runtime_guard.py +029234245bca0244659df4aece5fcced0a71a2a55bc41dfd0426ad2a019529b6 verify_joint_data_cache.py +4dc45077e7f94cc1e896152b454936fb1f2651e495b88db8529e84ba5c3f701c wrapper_common.py diff --git a/scripts/train/joint/native_elmo_ops/README.md b/scripts/train/joint/native_elmo_ops/README.md new file mode 100644 index 000000000..60d25b8a4 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/README.md @@ -0,0 +1,97 @@ +# Native Elmo composite recovery operations — service C pinned candidate + +This is a new, local 23-file operations assembly. It preserves the sealed runtime, +host-adapter, execution-v2, controller and observer fragments. No scientific source, +remote files, canonical campaign receipts, jobs or actual recovery bindings changed. +Both `execution_spec.SERVICE_COMMIT` / `SERVICE_SHA` and the wrapper equivalents +bind verified published service commit `abbb00499096faabf7035fde23bf4b1eba984377` and 35-file manifest +`adf45524de3c2e6dba186572e7151d48bfca0b3adc8e6cf9defd686b47744c87`. This operations closure +is a separate local candidate awaiting final review and its own publication. +The prepublication candidate and every sealed input remain unchanged. + +The controller is CPU orchestration only. It never proves failure recoverability, +checkpoint validity, zero optimizer progress, scientific completion or child native +completion. Its gateway transport still reopens remote bytes and native state; +all scientific admission remains inside the actual allocated host guards. An +unknown submission/release/capture outcome remains single-use and requires exact +inspection, never an automatic retry or cancellation. + +## Dependency and proof protocol + +`FILES.sha256` covers all static operations bytes and excludes dynamic execution, +request, observer, controller, claim and status directories. The ID-free execution +spec binds this manifest and predeclares only `release_observer_path`. The observer +spec binds the execution spec SHA. Thus neither source module embeds the other's +SHA and no manifest/spec/observer hash cycle is introduced. Fixed-name operations +imports verify the whole immutable closure, physical module path and source bytes +captured at import; foreign or changed cached modules fail closed. + +Every producer release requires an accepted, currently `PENDING/Dependency` +afterany observer and an exact `PENDING/JobHeldUser` producer. The release records +exclusive `RELEASE_HELD_NATIVE.txt`, its hash, and complete observer readiness +before the single release call. All five phases use the same observer contract; +the original production parents use their separate maintained `original_observer` +provenance path. + +`phase_observer.verify_captured(observer_spec_path, observer_spec_sha, ops_sha)` +reopens the accepted binding and delegates to the active maintained service: + +``` +phase_observer_proof.verify_observer( + {path: canonical native-JOB.observer.json, sha256: final commit hash}, + producer_request={path: actual request JSON, sha256: actual request hash}, + producer_native=exact native contract including actual JobId, + terminal_native={path: canonical native-JOB.txt, sha256: raw native hash}, +) +``` + +The service returns context including `result`, the exact observer commit JSON. +It reads and semantically validates evidence without importing operations code. +The operations wrapper preserves the existing `verify_captured -> result` API. +Both execution-claim and scientific-request-claim censuses require that final +commit, exact request/job/spec/native-path correspondence, and a fresh complete +owned step inventory. A raw terminal file alone never passes. Original-parent +claims require separately validated original-observer/failure evidence and their +actual original RUNNING admission; this is not substituted with a phase observer. + +The canonical same-arm controller ledger is also reopened before new controller +submission/release. Earlier PREPARED/spec/manifest/accepted IDs and held bytes must +agree. Every previously registered execution root and terminal directory must be +covered by the new execution census, and prior GPU writer claim roots by the +scientific census. An old controller RESULT saying a child was released is never +terminal proof, even after that child disappears from the queue. + +Missing RUNNING admission for a genuinely observed later terminal phase is only +terminal/quiescence evidence. It never proves zero progress or successful phase +completion. Maintained completion producers separately require actual RUNNING, +canonical host exit 0, native COMPLETED 0:0 and the exact phase-specific result. + +Proof metadata is limited to 4 MiB per file, native stdout/stderr to 1 MiB/64 KiB, +and immutable manifests to 128 members. Checkpoint and data I/O remain in the +allocated maintained producer, outside these metadata-only guards. + +## Local verification and remaining gates + +`tests/native_elmo_ops/` uses portable synthetic requests and native records. Native subprocesses are intercepted. The real composed binder, +observer, release hook and maintained proof reader run together for all three +arms and all five modes. Additional cases exercise changed module bytes, cached +foreign modules, unknown claims, partial observer commits, conflicting release +proofs, surviving worker steps and current held/RUNNING interfaces. The temporary +local service snapshot is test input only; it is not publication or production +proof. The final 71-case suite uses the exact published service constants without overriding them. +Only synthetic original-parent fixture profiles and native subprocess results are replaced in memory. + +Pending: independent final review; final operations publication; immutable actual deployment/readback; a genuine +recoverable parent failure; actual allocated controller/native authentication; +CPU highest-checkpoint inspection; all-rank GPU restoration and scheduled +validation; finalizer; human-visible resolved production preflight. No actual +request, claim, terminal certificate, checkpoint inspection or GPU proof is +created by this local candidate. + +Portable publication tests live in `tests/native_elmo_ops/` and synthesize all +requests and native evidence. Local retained-request tests remain audit-only. +The final portable suite contains 71 cases, including two executable-wrapper packaging +regressions. It runs at the intended repository layout without path overrides. +Repeated runs exercise the same matrix and must not be summed as unique tests. +Native wrapper files must retain mode 0755 in Git/export/tar; direct executable +bind mounts otherwise fail before actual native authentication. diff --git a/scripts/train/joint/native_elmo_ops/allocated.sbatch b/scripts/train/joint/native_elmo_ops/allocated.sbatch new file mode 100644 index 000000000..54e9e7b43 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/allocated.sbatch @@ -0,0 +1,76 @@ +#!/bin/bash +# Immutable held submission/release is owned by execution_spec.py, never this script. +set -euo pipefail +root=${1:?Exact operations root} +mode=${2:?capture inspect diagnostic finalize full} +spec=${3:?Exact execution specification} +spec_sha=${4:?Exact specification SHA256} +files_sha=${5:?Exact operations manifest SHA256} +[[ "$root" == /workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/* ]] +[[ "$mode" == capture || "$mode" == inspect || "$mode" == diagnostic || "$mode" == finalize || "$mode" == full ]] +[[ "$spec" == "$root"/execution/*.json && "$spec_sha" =~ ^[0-9a-f]{64}$ && "$files_sha" =~ ^[0-9a-f]{64}$ ]] +[[ "${SLURM_JOB_ID:-}" =~ ^[0-9]+$ && "${SLURM_RESTART_COUNT:-0}" == 0 ]] +timeout --kill-after=10s 120s bash -c ' + set -euo pipefail + printf "%s %s\n" "$2" "$1/FILES.sha256" | sha256sum --check --status + cd "$1" && sha256sum --check --status FILES.sha256 +' elmo-wrapper-manifest "$root" "$files_sha" +printf '%s %s\n' "$spec_sha" "$spec" | timeout --kill-after=10s 120s sha256sum --check --status +export PYTHONDONTWRITEBYTECODE=1 +source "$root/launcher_common.sh" +child= mirror= +receipt="$root/status/$mode-$SLURM_JOB_ID.exit" +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 [[ "$mode" == full && -d "${ELMO_LOCAL_LOGS:-}" ]]; then + destination="$ELMO_ATTEMPT/startup-log-mirror" + [[ ! -f "$ELMO_OUTPUT/.training-admitted" ]] || destination="$ELMO_OUTPUT/log-mirror" + mirror_once "$ELMO_LOCAL_LOGS" "$destination" || code=1 + fi + (set -o noclobber; printf '%s\n' "$code" > "$receipt") + exit "$code" +} +trap finish EXIT +trap 'exit 143' TERM +# Allocated host must authenticate actual native ownership before source/cache I/O. +exports=$(timeout --kill-after=10s 300s python3 "$root/runtime_guard.py" host "$spec" "$spec_sha" "$mode" "$files_sha") +eval "$exports" +native_exports=$(timeout --kill-after=10s 600s python3 "$root/native_client.py" "$spec" "$spec_sha" "$mode" "$files_sha") +eval "$native_exports" +mounts="/workspace:/workspace:rw,$ELMO_SCRATCH:$ELMO_SCRATCH:rw,$ELMO_NATIVE_MOUNTS" +workdir="$root" +if [[ "$mode" != capture ]]; then + runtime_exports=$(timeout --kill-after=10s 300s python3 "$root/lambda_runtime.py" "$ELMO_PARENT_REQUEST") + eval "$runtime_exports" + source_cache="$JOINT_NODE_SCRATCH/source/$SLURM_JOB_ID" + timeout --kill-after=10s 600s bash "$root/lambda_stage_source.sh" "$ELMO_ARCHIVE" "$ELMO_ARCHIVE_SHA" "$ELMO_SOURCE_MANIFEST" "$ELMO_SOURCE_MANIFEST_SHA" "$source_cache" + cp -- "$source_cache/source-stage.json" "$ELMO_ATTEMPT/source-stage.json" + mounts+=",$JOINT_JOB_TMP:$JOINT_JOB_TMP:rw,$(cat "$JOINT_RUNTIME_MOUNTS"),$ELMO_CACHE_ROOT:$ELMO_DATA_ROOT:ro,$source_cache/source:$ELMO_SOURCE:ro" + workdir="$ELMO_SOURCE" +fi +if [[ "$mode" == full ]]; then + 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 'Production W&B credential unavailable' + mounts+="$credentials_mount" +fi +gpu_args=(--gres=none) +if [[ "$mode" == diagnostic || "$mode" == full ]]; then gpu_args=(--gres=gpu:h100:4 --gpu-bind=none); fi +/usr/local/bin/srun --jobid="$SLURM_JOB_ID" --exclusive --exact --unbuffered --nodes=1 --ntasks=1 \ + --cpus-per-task="$ELMO_CPUS" --mem="$ELMO_MEMORY" "${gpu_args[@]}" \ + --container-name="elmo-recovery-$mode-$SLURM_JOB_ID" --container-image="$ELMO_CONTAINER_IMAGE" --container-writable \ + --container-mounts="$mounts" --container-workdir="$workdir" \ + --output="$root/logs/$mode-$SLURM_JOB_ID.out" --error="$root/logs/$mode-$SLURM_JOB_ID.err" \ + bash "$root/inside.sh" "$root" "$mode" "$spec" "$spec_sha" "$files_sha" & +child=$! +if [[ "$mode" == full ]]; then + watch_log_mirror "$child" "$ELMO_LOCAL_LOGS" "$ELMO_OUTPUT/log-mirror" "$ELMO_OUTPUT/.training-admitted" "$ELMO_ATTEMPT/unused" & + mirror=$! +fi +status=0 +wait "$child" || status=$? +child= +exit "$status" diff --git a/scripts/train/joint/native_elmo_ops/controller.py b/scripts/train/joint/native_elmo_ops/controller.py new file mode 100644 index 000000000..2b41de122 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/controller.py @@ -0,0 +1,487 @@ +"""Bounded Python-free gateway bootstrap and actual allocated-host controller. + +No model/checkpoint loading, remote mirror semantics, cancellation, or retry. +The allocated controller calls maintained producer and observer implementations. +""" +import argparse +import base64 +import hashlib +import importlib.util +import json +import os +from pathlib import Path, PurePosixPath +import re +import shlex +import subprocess +import sys +from datetime import datetime, timezone + +PREFIX = '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/' +PINS = { + 'plain': ('412789', 'gpu11', '936e9acd30fa517f6f38311c0bcfd2b88544da8d476bd41125c4008de2845073'), + 'fused': ('412798', 'gpu11', 'e34190bfef155942fe216a1e5e81317cfa31b65d4c98ccb025a579c874fd3876'), + 'boundary': ('412800', 'gpu03', '210e6f73c8a1a396716fab2f1e7e17f3b9d0663b099971eb05807181f9f00684'), +} +MODES = {'capture', 'inspect', 'diagnostic', 'finalize', 'full'} +FAILURE_STATES = {'FAILED', 'TIMEOUT', 'PREEMPTED', 'NODE_FAIL', 'BOOT_FAIL', 'CANCELLED'} +MAX_FILE = 4 << 20 + +_LOADED_IMPLEMENTATION_SHA256 = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + + +def need(value, message): + if not value: + raise ValueError(message) + + +def digest(value): + return hashlib.sha256(value).hexdigest() + + +def encode(value): + return (json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + '\n').encode() + + +def parse(value): + def pairs(items): + need(len(items) == len(dict(items)), 'Duplicate JSON field') + return dict(items) + return json.loads(value, object_pairs_hook=pairs, + parse_constant=lambda x: (_ for _ in ()).throw(ValueError('Nonfinite JSON'))) + + +def path(value): + p = PurePosixPath(value) + need(isinstance(value, str) and value.startswith(PREFIX) and '..' not in p.parts + and str(p) == value and re.fullmatch(r'[A-Za-z0-9_./=-]+', value), 'Unsafe/noncanonical campaign path') + return p + + +def fields(raw): + items = re.findall(r'(?:^|\s)([A-Za-z][A-Za-z0-9:/_-]*)=(\S+)', raw) + need(len(items) == len(dict(items)), 'Ambiguous native fields') + return dict(items) + + +def binding(value): + need(set(value) == {'path', 'sha256'} and re.fullmatch('[0-9a-f]{64}', value['sha256']), 'Exact path/hash binding required') + path(value['path']) + return value + + +class OutcomeUnknown(RuntimeError): + pass + + +class SSH: + """Explicit remote transport. Never executes an unquoted caller command.""" + def __init__(self, host='lambda'): + need(host == 'lambda', 'Only the configured Lambda endpoint is supported') + self.host = host + + def run(self, script, *, mutation=False): + try: + result = subprocess.run(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=10', + '-o', 'ServerAliveInterval=10', '-o', 'ServerAliveCountMax=2', self.host, + 'bash -c ' + shlex.quote(script)], capture_output=True, timeout=55) + except (OSError, subprocess.TimeoutExpired) as exc: + raise (OutcomeUnknown if mutation else RuntimeError)( + 'Transport failed; inspect exact existing claim/native state, never retry a mutation') from exc + need(len(result.stdout) <= MAX_FILE + 65536 and len(result.stderr) <= 65536, 'Unbounded transport output') + if result.returncode: + raise (OutcomeUnknown if mutation else RuntimeError)( + 'Remote operation did not certify success; inspect immutable claim; no retry') + return result.stdout + + def read(self, filename): + p = shlex.quote(str(path(filename))) + return self.run('set -euo pipefail; test "$(readlink -f -- ' + p + ')" = ' + p + + '; test -f ' + p + '; test "$(stat -c %s -- ' + p + ')" -le ' + str(MAX_FILE) + + '; timeout --kill-after=2s 15s cat -- ' + p) + + def native(self, argv): + need(argv[0] in {'/usr/local/bin/squeue', '/usr/local/bin/scontrol'}, 'Read-only native client required') + need(argv[1:2] != ['release'], 'Mutation through native reader prohibited') + return self.run('set -euo pipefail; timeout --kill-after=2s 20s ' + shlex.join(argv)).decode() + + def census(self, root): + p = shlex.quote(str(path(root))) + return self.run('set -euo pipefail; test "$(readlink -f -- ' + p + ')" = ' + p + + '; test -d ' + p + '; timeout --kill-after=2s 15s find ' + p + + ' -mindepth 1 -maxdepth 1 -printf "%f|%y\\n" | LC_ALL=C sort').decode() + + +class Reader: + def __init__(self, read): + self.read = read + self.readset = {} + + def file(self, filename): + path(filename) + value = self.read(filename) + need(isinstance(value, bytes) and len(value) <= MAX_FILE, 'Bounded byte reader required') + old = self.readset.setdefault(filename, digest(value)) + need(old == digest(value), 'Remote input changed during admission') + return value + + def bound(self, b, *, parsed=True): + binding(b) + value = self.file(b['path']) + need(digest(value) == b['sha256'], 'Changed bound remote input') + return parse(value) if parsed else value + + def manifest(self, root, b): + need(b['path'] == root + '/FILES.sha256', 'Manifest root differs') + data = self.bound(b, parsed=False).decode() + names = set() + for line in data.splitlines(): + sha, name = line.split(' ', 1) + need(name not in names and not name.startswith('/') and '..' not in PurePosixPath(name).parts, + 'Unsafe or duplicate manifest member') + names.add(name) + self.bound({'path': root + '/' + name, 'sha256': sha}, parsed=False) + need(names and len(names) <= 128, 'Unbounded operations manifest') + return names + + +def validate(reader, spec_path, spec_sha, ops_sha): + spec = reader.bound({'path': spec_path, 'sha256': spec_sha}) + need(spec['schema'] == 'native-elmo-allocated-controller-v1' and spec['arm'] in PINS, 'Controller schema/arm differs') + root = str(path(spec['root'])) + need(str(path(spec_path).parent) == root + '/controller-specs', 'Controller spec escaped root') + need(spec['ops_manifest'] == {'path': root + '/FILES.sha256', 'sha256': ops_sha}, 'Operations manifest differs') + need(spec['controller_claim_root']==PREFIX+'recovery-controller-claims/'+spec['arm'], + 'Canonical same-arm controller claim root required across operations packets') + names = reader.manifest(root, spec['ops_manifest']) + need({'controller.py', 'controller.sbatch', 'execution_spec.py', 'host_guard.py', 'phase_observer.py', 'observe-phase.sbatch'} <= names, + 'Controller/producer/observer immutable closure incomplete') + need(not any(PurePosixPath(n).parts[0] in {'controller-specs','controller-claims','execution','templates','requests','observer-specs','observer-claims','status','logs','scientific-claims','producer-claims','execution-claims'} for n in names), + 'Dynamic records must not enter immutable executable manifest') + parent_id, node, parent_sha = PINS[spec['arm']] + parent = reader.bound(spec['parent_request']) + need(spec['parent_request']['sha256'] == parent_sha, 'Controller changed original parent identity') + run_id = 'mot-elmo-pickplace-continuous-large314m-' + spec['arm'] + '-robot-human-camera-subtask-s42-300k-20261026-v1' + need(parent['run_id'] == run_id and parent['parameters'] == 314187064 and parent['world_size'] == 4 + and parent['target_steps'] == 300000 and parent['phase'] == 'full', 'Wrong original campaign identity') + intent = reader.bound(spec['current_intent']) + need(intent['schema'] == 'native-elmo-current-intent-v1' and intent['authority'] == 'human' + and intent['status'] == 'ACTIVE' and intent['hold_recovery'] is False + and run_id in intent['authorized_recovery_run_ids'] and run_id not in intent['retired_run_ids'], 'Recovery no longer authorized') + terminal = fields(reader.bound(spec['original_failed_native'], parsed=False).decode()) + need(terminal.get('JobId') == parent_id and terminal.get('JobName') == parent['job_name'] + and terminal.get('UserId') == 'paphiwetsa3(41003)' and terminal.get('NodeList') == node + and terminal.get('JobState') in FAILURE_STATES + and terminal.get('EndTime') not in {None, 'Unknown', 'None'} + and re.fullmatch(r'\d+:\d+', terminal.get('ExitCode', '')), 'Actual failed original parent required for controller allocation') + execution = spec['producer_execution'] + need(set(execution) == {'path','sha256','mode','ops_manifest_sha256'} and execution['mode'] in MODES + and execution['ops_manifest_sha256'] == ops_sha, 'Producer execution binding differs') + producer = reader.bound({k: execution[k] for k in ('path','sha256')}) + need(producer['root'] == root and producer['arm'] == spec['arm'] and producer['mode'] == execution['mode'] + and producer['parent_request'] == spec['parent_request'] and producer['runtime']['admission']['current_intent_path'] == spec['current_intent']['path'], 'Producer and controller scopes differ') + queue_scope=producer['runtime']['admission']['quiescence'] + need(spec['queue_scope']=={k:queue_scope[k] for k in ('user','scientific_job_names','scientific_name_prefixes')} + and spec['queue_scope']['user']=='paphiwetsa3' and spec['queue_scope']['scientific_job_names'] + and spec['queue_scope']['scientific_name_prefixes'] + and all(len(p)>=16 for p in spec['queue_scope']['scientific_name_prefixes']), 'Complete actual producer queue scope required') + reader.bound(producer['request_template']) + service = producer['operational_service'] + reader.manifest(service['root'], {'path':service['manifest'],'sha256':service['manifest_sha256']}) + observer = reader.bound(spec['observer_spec']) + need(observer == {'schema':'native-elmo-phase-observer-spec-v1', 'root':root, + 'producer_execution':execution, 'terminal_directory':spec['terminal_directory']}, 'Observer ID-free scope differs') + need(str(path(spec['observer_spec']['path']).parent) == root + '/observer-specs' + and producer['release_observer_path'] == spec['observer_spec']['path'], 'Mandatory release observer differs') + name = spec['job_name'] + need(re.fullmatch('elmo-recovery-controller-' + spec['arm'] + '-[a-z0-9-]{1,64}', name), 'Controller must have non-scientific scoped name') + need(spec['interpreter']['path'] == '/usr/bin/python3' and spec['interpreter']['minimum_version'] == [3,10], 'Explicit allocated-host interpreter required') + evidence = reader.bound(spec['interpreter']['provenance']) + need(evidence['schema'] == 'native-elmo-host-python-provenance-v1' and evidence['node'] == node + and evidence['interpreter_path'] == '/usr/bin/python3', 'Interpreter provenance belongs to another node') + script = reader.bound(evidence['script'], parsed=False).decode() + native = fields(reader.bound(evidence['native'], parsed=False).decode()) + need('/usr/bin/python3 -B ' in script and native['JobId'] == evidence['job_id'] and native['NodeList'] == node + and native['JobState'] == 'COMPLETED' and native['ExitCode'] == '0:0' + and reader.bound(evidence['host_exit'], parsed=False).strip() == b'0', 'No successful allocated-host interpreter route') + # Prior script proves the route, not current binary identity or GPU/runtime health. + return spec, producer + + +def controller_argv(spec, spec_path, spec_sha, ops_sha): + root=spec['root']; node=PINS[spec['arm']][1] + return ['/usr/local/bin/sbatch','--parsable','--hold','--no-requeue','--export=NONE', + '--account=loaner','--qos=loaner-normal','--partition=rl2','--nodes=1','--ntasks=1', + '--cpus-per-task=2','--mem=1G','--gres=none','--time=00:20:00','--nodelist='+node, + '--job-name='+spec['job_name'],'--comment=elmo-controller:'+spec_sha,'--chdir='+root, + '--output='+root+'/logs/%x-%j.out','--error='+root+'/logs/%x-%j.err', + root+'/controller.sbatch',root,spec_path,spec_sha,ops_sha] + + +def check_controller(raw,spec,spec_sha,job,state): + f=fields(raw); root=spec['root']; node=PINS[spec['arm']][1] + expected={'JobId':job,'JobName':spec['job_name'],'UserId':'paphiwetsa3(41003)', + 'Account':'loaner','QOS':'loaner-normal','Partition':'rl2','Requeue':'0','Restarts':'0', + 'JobState':state,'ReqNodeList':node,'NumCPUs':'2','NumTasks':'1','MinMemoryNode':'1G', + 'TimeLimit':'00:20:00','Command':root+'/controller.sbatch','Comment':'elmo-controller:'+spec_sha} + need(all(f.get(k)==v for k,v in expected.items()),'Controller actual native identity/resources differ') + need(f.get('SubmitTime') not in {None,'Unknown','None'} and f.get('NumNodes') in ({'1','1-1'} if state=='PENDING' else {'1'}), 'Controller submit/node contract missing') + if state=='PENDING': + need(f.get('Reason')=='JobHeldUser' and f.get('NodeList','') in {'','(null)',node} + and f.get('AllocTRES','') in {'','(null)'} and f.get('StartTime') in {None,'Unknown','N/A'}, 'Controller is not genuinely held') + else: + need(f.get('NodeList')==node and f.get('StartTime') not in {None,'Unknown','None'},'Controller allocation not actual') + for key in ('ReqTRES', *(['AllocTRES'] if state=='RUNNING' else [])): + values=set(f.get(key,'').split(',')) + need({'cpu=2','node=1'} <= values and bool(values & {'mem=1G','mem=1024M'}) and not any('gpu' in v.lower() for v in values),'Controller CPU/memory TRES conflict') + need(f.get('JOB_GRES','') in {'','(null)'} and f.get('GRES','') in {'','(null)'} + and f.get('TresPerNode','') in {'','(null)'} and f.get('TresPerTask','cpu:2')=='cpu:2', 'Controller has GPU/auxiliary resource claims') + need(not any('gpu' in k.lower() or ('tresper' in k.lower() and 'gpu' in v.lower()) for k,v in f.items()), 'Controller owns unexpected GPU resources') + return f + + +QUEUE=['/usr/local/bin/squeue','--noheader','--user','paphiwetsa3','--format=%i|%j|%T|%k|%u'] + + +def queue_guard(raw,spec,current=None): + for line in raw.splitlines(): + if not line.strip(): continue + row=line.split('|') + need(len(row)==5 and row[0].isdecimal() and row[4]=='paphiwetsa3','Malformed full owner queue') + scope=spec['queue_scope'] + run_id='mot-elmo-pickplace-continuous-large314m-'+spec['arm']+'-robot-human-camera-subtask-s42-300k-20261026-v1' + scientific=(row[1] in scope['scientific_job_names'] or any(row[1].startswith(p) for p in scope['scientific_name_prefixes']) + or row[3]==run_id) + controller=row[1].startswith('elmo-recovery-controller-'+spec['arm']+'-') + need(not scientific and (not controller or row[0]==current and row[2]=='PENDING'), 'Other active/pending scientific writer or controller') + + +def ledger(spec): + return spec['controller_claim_root']+'/'+spec['producer_execution']['sha256'] + + +def write_shell(filename,payload): + # GNU dd exclusive create plus fsync; never overwrites or deletes a record. + data=base64.b64encode(payload).decode() + return 'printf %s '+shlex.quote(data)+' | base64 --decode | dd status=none conv=excl,fsync of='+shlex.quote(filename)+'; ' + + +def transaction(reader,remote,spec,queue,census,body,*,held=None): + root=spec['root']; claims=spec['controller_claim_root']; q=shlex.quote + lock=claims+'.lock' + script='set -euo pipefail; umask 077; test "$(readlink -f -- '+q(str(PurePosixPath(claims).parent))+')" = '+q(str(PurePosixPath(claims).parent))+'; test ! -L '+q(lock)+'; exec 9>>'+q(lock)+'; flock -n 9; ' + for filename,sha in sorted(reader.readset.items()): + script+='test "$(readlink -f -- '+q(filename)+')" = '+q(filename)+'; printf "%s %s\\n" '+q(sha)+' '+q(filename)+' | sha256sum --check --status; ' + script+='current=$(timeout --kill-after=2s 20s '+shlex.join(QUEUE)+'); test "$(printf %s "$current" | sha256sum | cut -d" " -f1)" = '+q(digest(queue.rstrip('\n').encode()))+'; ' + for directory,observed in sorted(census.items()): + script+='current=$(find '+q(directory)+' -mindepth 1 -maxdepth 1 -printf "%f|%y\\n" | LC_ALL=C sort); test "$(printf %s "$current" | sha256sum | cut -d" " -f1)" = '+q(digest(observed.rstrip('\n').encode()))+'; ' + if held: + job,raw=held + script+='current=$(timeout --kill-after=2s 20s /usr/local/bin/scontrol show job '+q(job)+' -o -dd); test "$(printf %s "$current" | sha256sum | cut -d" " -f1)" = '+q(digest(raw.rstrip('\n').encode()))+'; ' + return remote.run(script+body,mutation=True) + + +def prior_census(reader,remote,spec,*,own=False): + census={}; current=spec['producer_execution']['sha256']; seen=set() + for directory in [spec['controller_claim_root']]: + raw=remote.census(directory); census[directory]=raw + for line in raw.splitlines(): + name,kind=line.split('|') + need(re.fullmatch('[0-9a-f]{64}',name) and kind=='d' and name not in seen,'Unexpected/duplicate controller claim entry') + seen.add(name) + if own and name==current: continue + need(name!=current,'Controller claim already exists; inspect, never retry') + old=directory+'/'+name + prepared=parse(reader.file(old+'/PREPARED.json')) + need(prepared['schema']=='native-elmo-controller-prepared-v1' + and prepared['producer_spec_sha256']==name,'Earlier controller intent differs') + previous=reader.bound({'path':prepared['spec_path'],'sha256':prepared['spec_sha256']}) + need(previous['schema']=='native-elmo-allocated-controller-v1' + and previous['arm']==spec['arm'] and previous['parent_request']==spec['parent_request'] + and previous['controller_claim_root']==spec['controller_claim_root'] + and str(path(prepared['spec_path']).parent)==previous['root']+'/controller-specs' + and previous['ops_manifest']=={'path':previous['root']+'/FILES.sha256','sha256':prepared['ops_manifest_sha256']}, + 'Earlier controller belongs to another root/identity/manifest') + reader.manifest(previous['root'],previous['ops_manifest']) + execution=previous['producer_execution'] + need(execution['sha256']==name and execution['ops_manifest_sha256']==prepared['ops_manifest_sha256'] + and str(path(execution['path']).parent)==previous['root']+'/execution', + 'Earlier producer execution escaped immutable controller intent') + earlier=reader.bound({k:execution[k] for k in ('path','sha256')}) + need(earlier['root']==previous['root'] and earlier['arm']==spec['arm'] + and earlier['parent_request']==spec['parent_request'] and earlier['mode']==execution['mode'], + 'Earlier producer belongs to another original campaign') + current_producer=reader.bound({k:spec['producer_execution'][k] for k in ('path','sha256')}) + old_roots=earlier['execution_claim_roots'] + own_root={'path':earlier['root']+'/execution-claims','terminal_directory':previous['terminal_directory']} + need(own_root in old_roots and old_roots + and all(item in current_producer['execution_claim_roots'] for item in old_roots), + 'Current execution census omits a prior released child root or terminal directory') + if earlier['mode'] in {'diagnostic','full'}: + scientific={'path':earlier['root']+'/scientific-claims','terminal_directory':previous['terminal_directory']} + need(scientific in current_producer['runtime']['admission']['quiescence']['claim_roots'], + 'Current scientific census omits a prior released writer claim root') + accepted=parse(reader.file(old+'/acceptance.json')); result=parse(reader.file(old+'/RESULT.json')) + prior_argv=controller_argv(previous,prepared['spec_path'],prepared['spec_sha256'],prepared['ops_manifest_sha256']) + need(prepared['argv']==prior_argv and accepted['argv']==prior_argv + and accepted['schema']=='native-elmo-controller-acceptance-v1' + and accepted['spec_sha256']==prepared['spec_sha256'] and accepted['producer_spec_sha256']==name + and str(accepted['job_id']).isdecimal() + and reader.file(old+'/PARENT_JOB_ID').strip()==accepted['job_id'].encode() + and reader.file(old+'/sbatch.stdout').strip()==accepted['job_id'].encode() + and reader.file(old+'/sbatch.exit').strip()==b'0', + 'Earlier controller acceptance is incomplete or differs from immutable intent') + held=reader.file(old+'/HELD_NATIVE.txt').decode() + need(digest(held.encode())==accepted['held_native_sha256'],'Earlier controller held bytes changed') + check_controller(held,previous,prepared['spec_sha256'],accepted['job_id'],'PENDING') + need(result['schema']=='native-elmo-controller-result-v1' and result['controller_job_id']==accepted['job_id'] + and result['producer_spec_sha256']==name and result['status']=='CHILD_RELEASE_ACCEPTED_NOT_COMPLETION', + 'Earlier controller outcome unknown; inspect child claims/native, never bypass') + # This result proves no child terminal state. The actual allocated + # host must reopen every included execution/request claim and its + # maintained observer final commit before submitting a new child. + return census + + +def submit_controller(remote,spec_path,spec_sha,ops_sha): + reader=Reader(remote.read); spec,_=validate(reader,spec_path,spec_sha,ops_sha) + census=prior_census(reader,remote,spec); queue=remote.native(QUEUE); queue_guard(queue,spec) + claim=ledger(spec); argv=controller_argv(spec,spec_path,spec_sha,ops_sha) + prepared=dict(schema='native-elmo-controller-prepared-v1',spec_path=spec_path,spec_sha256=spec_sha, + ops_manifest_sha256=ops_sha,producer_spec_sha256=spec['producer_execution']['sha256'],argv=argv) + body='mkdir '+shlex.quote(claim)+'; sync -f '+shlex.quote(spec['controller_claim_root'])+'; ' + body+=write_shell(claim+'/PREPARED.json',encode(prepared)) + body+='sync -f '+shlex.quote(claim)+'; ' + body+='set +e; (set -o noclobber; timeout --kill-after=3s 25s '+shlex.join(argv)+' > '+shlex.quote(claim+'/sbatch.stdout')+' 2> '+shlex.quote(claim+'/sbatch.stderr')+'); code=$?; set -e; ' + body+='printf "%s\\n" "$code" | dd status=none conv=excl,fsync of='+shlex.quote(claim+'/sbatch.exit')+'; sync -f '+shlex.quote(claim)+'; test "$code" = 0; cat '+shlex.quote(claim+'/sbatch.stdout')+'; ' + job=transaction(reader,remote,spec,queue,census,body).decode().strip() + need(job.isdecimal(),'Controller acceptance unknown; inspect existing claim, never retry') + # A failed binding read/write leaves the durable submission claim unresolved. + raw=remote.native(['/usr/local/bin/scontrol','show','job',job,'-o','-dd']) + check_controller(raw,spec,spec_sha,job,'PENDING') + need(reader.file(claim+'/sbatch.stdout').strip()==job.encode() and reader.file(claim+'/sbatch.exit').strip()==b'0','Actual controller acceptance not preserved') + census=prior_census(reader,remote,spec,own=True); queue=remote.native(QUEUE); queue_guard(queue,spec,job) + accepted=dict(schema='native-elmo-controller-acceptance-v1',job_id=job,spec_sha256=spec_sha, + producer_spec_sha256=spec['producer_execution']['sha256'],argv=argv,held_native_sha256=digest(raw.encode())) + body=write_shell(claim+'/PARENT_JOB_ID',(job+'\n').encode())+write_shell(claim+'/HELD_NATIVE.txt',raw.encode())+write_shell(claim+'/acceptance.json',encode(accepted)) + body+='sync -f '+shlex.quote(claim)+'; ' + transaction(reader,remote,spec,queue,census,body,held=(job,raw)) + return accepted + + +def release_controller(remote,spec_path,spec_sha,ops_sha): + reader=Reader(remote.read); spec,_=validate(reader,spec_path,spec_sha,ops_sha); claim=ledger(spec) + accepted=parse(reader.file(claim+'/acceptance.json')); prepared=parse(reader.file(claim+'/PREPARED.json')); job=accepted['job_id'] + argv=controller_argv(spec,spec_path,spec_sha,ops_sha) + need(accepted['schema']=='native-elmo-controller-acceptance-v1' and accepted['spec_sha256']==spec_sha + and accepted['argv']==prepared['argv']==argv and prepared['spec_sha256']==spec_sha + and reader.file(claim+'/PARENT_JOB_ID').strip()==job.encode() + and reader.file(claim+'/sbatch.stdout').strip()==job.encode() and reader.file(claim+'/sbatch.exit').strip()==b'0', 'Controller acceptance changed') + held=reader.file(claim+'/HELD_NATIVE.txt').decode(); old=check_controller(held,spec,spec_sha,job,'PENDING') + need(digest(held.encode())==accepted['held_native_sha256'],'Held evidence changed') + raw=remote.native(['/usr/local/bin/scontrol','show','job',job,'-o','-dd']); current=check_controller(raw,spec,spec_sha,job,'PENDING') + need(current['SubmitTime']==old['SubmitTime'],'Controller allocation replaced') + census=prior_census(reader,remote,spec,own=True); queue=remote.native(QUEUE); queue_guard(queue,spec,job) + attempt=dict(job_id=job,spec_sha256=spec_sha,held_native_sha256=digest(raw.encode()),completion_proven=False) + body=write_shell(claim+'/RELEASE_ATTEMPT.json',encode(attempt)) + body+='sync -f '+shlex.quote(claim)+'; ' + body+='set +e; (set -o noclobber; timeout --kill-after=3s 25s /usr/local/bin/scontrol release '+shlex.quote(job)+' > '+shlex.quote(claim+'/release.stdout')+' 2> '+shlex.quote(claim+'/release.stderr')+'); code=$?; set -e; ' + body+='printf "%s\\n" "$code" | dd status=none conv=excl,fsync of='+shlex.quote(claim+'/release.exit')+'; sync -f '+shlex.quote(claim)+'; test "$code" = 0; ' + transaction(reader,remote,spec,queue,census,body,held=(job,raw)) + return {'job_id':job,'status':'CONTROLLER_RELEASE_ACCEPTED_NOT_CHILD_AUTHORIZATION_OR_COMPLETION'} + + +def local_read(filename): + p=Path(str(path(filename))) + need(p.resolve(strict=True)==p and p.is_file() and p.stat().st_size<=MAX_FILE,'Actual physical allocated-host input required') + return p.read_bytes() + + +def write_once(filename,value): + p=Path(filename); need(p.parent.resolve(strict=True)==p.parent,'Physical output parent required') + with p.open('xb') as stream: + stream.write(encode(value)); stream.flush(); os.fsync(stream.fileno()) + fd=os.open(p.parent,os.O_RDONLY) + try: os.fsync(fd) + finally: os.close(fd) + + +def import_bound(root,name,manifest_sha): + need(name in {'execution_spec','phase_observer'},'No arbitrary controller import') + names=Reader(local_read).manifest(root,{'path':str(Path(root)/'FILES.sha256'),'sha256':manifest_sha}) + need(name+'.py' in names,'Controller module absent from immutable closure') + location=Path(root)/(name+'.py') + module_spec=importlib.util.spec_from_file_location(name,location) + module=importlib.util.module_from_spec(module_spec) + module._verified_ops_manifest_sha256=manifest_sha + module._verified_ops_member_sha256=digest(location.read_bytes()) + sys.modules[name]=module; module_spec.loader.exec_module(module) + return module + + +def advance(execution,observer,execution_args,observer_args,*,before_each): + """One producer, one observer, one release; no recovery/retry branch.""" + need(getattr(execution,'RELEASE_OBSERVER_REQUIRED',False) is True,'Execution release lacks mandatory observer guard') + before_each(); accepted=execution.submit_held(*execution_args) + before_each(); observed=observer.submit_once(*observer_args) + before_each(); ready=observer.verify_ready_for_release(*observer_args) + # release_bound must also enforce its mandatory observer hook internally. + released=execution.release_bound(*execution_args) + return {'accepted':accepted,'observer':observed,'observer_readiness':ready,'released':released} + + +def allocated(spec_path,spec_sha,ops_sha): + reader=Reader(local_read); spec,producer=validate(reader,spec_path,spec_sha,ops_sha); root=spec['root']; claim=ledger(spec) + job=os.environ.get('SLURM_JOB_ID','') + need(job.isdecimal() and os.getuid()==41003 and os.environ.get('SLURM_RESTART_COUNT','0')=='0' + and os.environ.get('SLURMD_NODENAME')==PINS[spec['arm']][1] + and not os.environ.get('SLURM_JOB_GPUS') and not os.environ.get('SLURM_STEP_GPUS'),'Actual CPU controller process identity missing') + need(sys.version_info >= (3,10) and Path(sys.executable).resolve()==Path('/usr/bin/python3').resolve(),'Unbound current host Python interpreter') + native=subprocess.run(['/usr/local/bin/scontrol','show','job',job,'-o','-dd'],capture_output=True,text=True,timeout=20) + need(native.returncode==0 and not native.stderr,'Controller native authentication failed') + current=check_controller(native.stdout,spec,spec_sha,job,'RUNNING') + accepted=parse(reader.file(claim+'/acceptance.json')); held=reader.file(claim+'/HELD_NATIVE.txt').decode() + prepared=parse(reader.file(claim+'/PREPARED.json')); released=parse(reader.file(claim+'/RELEASE_ATTEMPT.json')) + previous=check_controller(held,spec,spec_sha,job,'PENDING') + need(accepted['job_id']==job and accepted['spec_sha256']==spec_sha and accepted['argv']==controller_argv(spec,spec_path,spec_sha,ops_sha) + and accepted['held_native_sha256']==digest(held.encode()) and current['SubmitTime']==previous['SubmitTime'] + and released['job_id']==job and released['spec_sha256']==spec_sha + and prepared['spec_sha256']==spec_sha and prepared['argv']==accepted['argv'] + and reader.file(claim+'/PARENT_JOB_ID').strip()==job.encode() + and reader.file(claim+'/sbatch.stdout').strip()==job.encode() + and reader.file(claim+'/sbatch.exit').strip()==b'0','Controller not the accepted/released actual allocation') + write_once(claim+'/RUNNING.json',dict(job_id=job,native=native.stdout,interpreter=sys.executable, + interpreter_realpath=str(Path(sys.executable).resolve()),interpreter_sha256=digest(Path(sys.executable).read_bytes()),python_version=sys.version)) + sys.path.insert(0,root) + execution=import_bound(root,'execution_spec',ops_sha); observer=import_bound(root,'phase_observer',ops_sha) + e=spec['producer_execution']; execution_args=(e['path'],e['sha256'],e['mode'],ops_sha) + observer_args=(spec['observer_spec']['path'],spec['observer_spec']['sha256'],ops_sha) + def refresh(): + # Real shared filesystem reread before every phase transition; mutation fails closed. + fresh=Reader(local_read); actual,_=validate(fresh,spec_path,spec_sha,ops_sha) + need(actual==spec,'Controller spec changed') + result=advance(execution,observer,execution_args,observer_args,before_each=refresh) + record=dict(schema='native-elmo-controller-result-v1',controller_job_id=job, + producer_spec_sha256=e['sha256'],status='CHILD_RELEASE_ACCEPTED_NOT_COMPLETION',details=result, + controller_native_completion='UNKNOWN_NOT_USED_TO_CERTIFY_CHILD',created_at=datetime.now(timezone.utc).isoformat()) + write_once(claim+'/RESULT.json',record) + return record + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('operation',choices=['plan','submit-controller','release-controller','allocated']) + parser.add_argument('spec_path'); parser.add_argument('spec_sha'); parser.add_argument('ops_sha') + args=parser.parse_args(); values=(args.spec_path,args.spec_sha,args.ops_sha) + if args.operation=='allocated': result=allocated(*values) + else: + remote=SSH() + if args.operation=='plan': + reader=Reader(remote.read); spec,_=validate(reader,*values) + result={'argv':controller_argv(spec,*values),'allocates':False,'remote_readset':reader.readset} + else: result={'submit-controller':submit_controller,'release-controller':release_controller}[args.operation](remote,*values) + print(json.dumps(result,indent=2)) + + +if __name__=='__main__': + main() diff --git a/scripts/train/joint/native_elmo_ops/controller.sbatch b/scripts/train/joint/native_elmo_ops/controller.sbatch new file mode 100644 index 000000000..959a179a8 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/controller.sbatch @@ -0,0 +1,26 @@ +#!/bin/bash +# CPU orchestration only. Host exit never certifies any child native terminal. +set -euo pipefail +root=${1:?Operations root} +spec=${2:?Controller specification} +spec_sha=${3:?Controller SHA256} +files_sha=${4:?Operations manifest SHA256} +[[ "$root" == /workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/* ]] +[[ "$spec" == "$root"/controller-specs/*.json ]] +[[ "$spec_sha" =~ ^[0-9a-f]{64}$ && "$files_sha" =~ ^[0-9a-f]{64}$ ]] +[[ "${SLURM_JOB_ID:-}" =~ ^[0-9]+$ && "${SLURM_RESTART_COUNT:-0}" == 0 ]] +[[ "$(id -u)" == 41003 && -z "${SLURM_JOB_GPUS:-}" && -z "${SLURM_STEP_GPUS:-}" ]] +finish() { + code=$? + trap - EXIT + (set -o noclobber; printf '%s\n' "$code" > "$root/status/controller-$SLURM_JOB_ID.exit") || code=74 + exit "$code" +} +trap finish EXIT +printf '%s %s\n' "$files_sha" "$root/FILES.sha256" | sha256sum --check --status +(cd "$root" && timeout --kill-after=5s 120s sha256sum --check --status FILES.sha256) +printf '%s %s\n' "$spec_sha" "$spec" | sha256sum --check --status +[[ -x /usr/bin/python3 ]] +export CUDA_VISIBLE_DEVICES=-1 PYTHONDONTWRITEBYTECODE=1 +unset PYTHONPATH PYTHONHOME WANDB_API_KEY PIXEL_WANDB_NETRC +timeout --kill-after=10s 900s /usr/bin/python3 -I -B "$root/controller.py" allocated "$spec" "$spec_sha" "$files_sha" diff --git a/scripts/train/joint/native_elmo_ops/dispatch_after_runtime.py b/scripts/train/joint/native_elmo_ops/dispatch_after_runtime.py new file mode 100644 index 000000000..311575c31 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/dispatch_after_runtime.py @@ -0,0 +1,53 @@ +"""Fixed phase adapter, reached after exact runtime_bootstrap/service_bootstrap.""" +import os +from pathlib import Path +import runpy +import sys +from wrapper_common import MODES, bound, need, read_json, sha, write_once, now +from runtime_guard import attempt_path, capsule, native_auth, prerequisite_path + + +def main(request_path): + from execution_spec import load_execution + args = (os.environ['ELMO_EXECUTION_SPEC'], os.environ['ELMO_EXECUTION_SPEC_SHA256'], + os.environ['ELMO_EXECUTION_MODE'], os.environ['ELMO_OPS_MANIFEST_SHA256']) + spec, path, q, job, claim = load_execution(*args) + mode = spec['mode'] + need(mode in {'inspect', 'diagnostic', 'finalize', 'full'} and str(path) == request_path, + 'Fixed non-capture phase adapter required') + from scripts.train.joint.service_bootstrap import activate + root = activate(q) + need(root == Path(spec['operational_service']['root']), 'Active operational service differs') + import torch + parent = bound(spec['parent_request']) + need(Path(torch.__file__).is_relative_to(Path(parent['lambda_runtime']['fa4_root']) / 'torch-runtime'), + 'Pinned runtime bootstrap did not select exact Torch') + attempt = attempt_path(spec, job) + actual = read_json(attempt / 'pre-torch-prerequisites.json') + need(actual['job_id'] == job and actual['request_sha256'] == sha(path), 'Pre-Torch actual evidence differs') + auth = native_auth(spec, path, q, job) + gpu = None + if mode in {'diagnostic', 'full'}: + from gpu_health_set import admit_gpu_set + gpu = admit_gpu_set(attempt / 'gpu-health.json', 4) + need(gpu['status'] == 'PASS', 'Actual four-H100 health admission failed') + else: + need(not os.environ.get('SLURM_JOB_GPUS') and not os.environ.get('SLURM_STEP_GPUS'), 'CPU phase owns GPUs') + receipt = dict(schema='native-elmo-allocated-prerequisites-v1', + status='ACTUAL_PREREQUISITES_AND_NATIVE_ALLOCATION_PASS', + request_sha256=sha(path), job_id=job, native=auth['native'], checked_at=now(), + parent_request=spec['parent_request'], execution_spec_sha256=args[1], + ops_manifest_sha256=args[3], operational_service=spec['operational_service'], + actual_prerequisites=actual, actual_gpu_health=gpu, + native_client_manifest_sha256=capsule(spec, job), + host_auth={'path': str(attempt / 'host-native.json'), 'sha256': sha(attempt / 'host-native.json')}, + container_auth={'path': str(attempt / 'container-native.json'), 'sha256': sha(attempt / 'container-native.json')}, + actual_runtime_bootstrap=True, actual_service_bootstrap=True) + write_once(prerequisite_path(spec, q), receipt) + module = 'scripts.train.joint.' + MODES[mode] + sys.argv = [module, str(path)] + runpy.run_module(module, run_name='__main__', alter_sys=True) + + +if __name__ == '__main__': + main(sys.argv[1]) diff --git a/scripts/train/joint/native_elmo_ops/execution_spec.py b/scripts/train/joint/native_elmo_ops/execution_spec.py new file mode 100644 index 000000000..eb13ae1f8 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/execution_spec.py @@ -0,0 +1,617 @@ +"""Immutable native Elmo phase submission and late actual-job request binding. + +Standard library only; no source/checkpoint loads, training, or implicit retries. +The hash-bound host guard owns current scientific/infrastructure admission. +""" +import argparse +import copy +import hashlib +import importlib +import importlib.util +import json +import os +from pathlib import Path +import re +import subprocess +import sys +from datetime import datetime, timezone + + +RELEASE_OBSERVER_REQUIRED = True + +SERVICE_COMMIT = "abbb00499096faabf7035fde23bf4b1eba984377" # Verified upstream service publication. +SERVICE_SHA = "adf45524de3c2e6dba186572e7151d48bfca0b3adc8e6cf9defd686b47744c87" +SERVICE_FILE_COUNT = 35 +PINS = { + "plain": ("412789", "gpu11", "936e9acd30fa517f6f38311c0bcfd2b88544da8d476bd41125c4008de2845073"), + "fused": ("412798", "gpu11", "e34190bfef155942fe216a1e5e81317cfa31b65d4c98ccb025a579c874fd3876"), + "boundary": ("412800", "gpu03", "210e6f73c8a1a396716fab2f1e7e17f3b9d0663b099971eb05807181f9f00684"), +} +MODES = { + "capture": ("native-elmo-interruption-capture-v1", "2", "1G", "00:10:00", "", "native_interruption"), + "inspect": ("native-elmo-checkpoint-inspection-v1", "4", "24G", "00:20:00", "", "checkpoint_inspection"), + "diagnostic": ("native-elmo-production-recovery-request-v1", "32", "192G", "01:00:00", "gres/gpu:h100:4", "production_recovery"), + "finalize": ("native-elmo-interruption-finalizer-v1", "4", "24G", "00:20:00", "", "interruption_admission"), + "full": ("joint-training-v2", "32", "192G", "7-00:00:00", "gres/gpu:h100:4", "run"), +} +NATIVE_KEYS = {"JobName", "UserId", "Account", "QOS", "Partition", "ReqNodeList", "NodeList", + "NumCPUs", "MinMemoryNode", "TimeLimit", "Requeue", "Restarts", "Command", "TresPerNode"} + +_LOADED_IMPLEMENTATION_SHA256 = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + + +def need(value, message): + if not value: + raise ValueError(message) + + +def sha(path): + p = Path(path) + need(p.stat().st_size <= 4 * 1024 * 1024, "Operations metadata exceeds 4 MiB") + return hashlib.sha256(p.read_bytes()).hexdigest() + + +def absolute(path): + p = Path(path) + need(p.is_absolute() and ".." not in p.parts and p.resolve() == p, "Noncanonical/symlink path") + return p + + +def bound(binding): + p = absolute(binding["path"]) + need(p.is_file() and not p.is_symlink() and sha(p) == binding["sha256"], "Bound file missing or changed: " + str(p)) + return p + + +def read_json(path): + def pairs(items): + need(len(items) == len(dict(items)), "Duplicate JSON keys") + return dict(items) + need(Path(path).stat().st_size <= 4 * 1024 * 1024, "Operations JSON exceeds 4 MiB") + return json.loads(Path(path).read_text(), object_pairs_hook=pairs, + parse_constant=lambda x: (_ for _ in ()).throw(ValueError("Nonfinite JSON: " + x))) + + +def encoded(value): + return (json.dumps(value, sort_keys=True, indent=2, allow_nan=False) + "\n").encode() + + +def write_once(path, value): + p = absolute(path) + payload = value if isinstance(value, bytes) else encoded(value) + with p.open("xb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + fd = os.open(p.parent, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + need(p.read_bytes() == payload, "Immutable write readback mismatch") + + +def mkdir_once(path): + p = absolute(path) + p.mkdir() + fd = os.open(p.parent, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def manifest(root, binding): + root = absolute(root) + path = bound(binding) + need(path == root / "FILES.sha256", "Manifest outside exact root") + names = set() + lines = path.read_text().splitlines() + need(0 < len(lines) <= 128, "Unbounded operations/service manifest") + for line in lines: + digest, name = line.split(" ", 1) + need(re.fullmatch(r"[a-f0-9]{64}", digest) and name not in names + and not Path(name).is_absolute() and ".." not in Path(name).parts, "Invalid manifest entry") + bound(dict(path=str(root / name), sha256=digest)) + names.add(name) + need(names, "Empty manifest") + return names + + +def fields(raw): + pairs = re.findall(r"(?:^|\s)([A-Za-z][A-Za-z0-9:/_-]*)=(\S+)", raw) + need(len(pairs) == len(dict(pairs)), "Ambiguous native record") + return dict(pairs) + + +def command(argv): + result = subprocess.run(argv, capture_output=True, text=True, timeout=30) + need(len(result.stdout.encode()) <= 1048576 and len(result.stderr.encode()) <= 65536, "Native query exceeded bounds") + need(result.returncode == 0 and not result.stderr.strip(), "Native query failed: " + result.stderr[-1500:]) + return result.stdout + + +def paths(spec, spec_sha, job=None, request_sha=None): + root = Path(spec["root"]) + return dict(spec_claim=root / "execution-claims" / spec_sha, + request=root / "requests" / (spec_sha + "-" + str(job) + ".json"), + request_root=root / ("scientific-claims" if spec["mode"] in {"diagnostic", "full"} else "producer-claims"), + request_claim=root / ("scientific-claims" if spec["mode"] in {"diagnostic", "full"} else "producer-claims") / str(request_sha)) + + +def read_spec(spec_path, spec_sha, mode, manifest_sha): + need(isinstance(SERVICE_COMMIT, str) and re.fullmatch(r"[0-9a-f]{40}", SERVICE_COMMIT), + "Adapter service is unpublished; explicit reviewed publication binding required") + need(re.fullmatch(r"[a-f0-9]{64}", spec_sha) and re.fullmatch(r"[a-f0-9]{64}", manifest_sha), "Invalid hashes") + p = bound(dict(path=str(spec_path), sha256=spec_sha)) + spec = read_json(p) + need(spec["schema"] == "native-elmo-execution-spec-v1" and spec["mode"] == mode and mode in MODES, + "Wrong execution schema/mode") + root = absolute(spec["root"]) + need(p.parent == root / "execution" and spec["arm"] in PINS, "Spec root/arm differs") + claim_roots = spec["execution_claim_roots"] + need(isinstance(claim_roots, list) and claim_roots + and len({r["path"] for r in claim_roots}) == len(claim_roots) + and str(root / "execution-claims") in {r["path"] for r in claim_roots}, + "Complete unique same-identity execution claim roots required") + for item in claim_roots: + need(absolute(item["path"]).is_dir() and absolute(item["terminal_directory"]).is_dir(), + "Execution claim/terminal directory unavailable") + need(spec["ops_manifest"] == dict(path=str(root / "FILES.sha256"), sha256=manifest_sha), "Ops manifest binding differs") + files = manifest(root, spec["ops_manifest"]) + need({"execution_spec.py", "allocated.sbatch", "host_guard.py", "wrapper_common.py", "phase_observer.py", "observe-phase.sbatch"} <= files + and all(Path(name).parts[0] not in {"execution", "templates", "execution-claims", "requests", "scientific-claims", "producer-claims", "status", "logs", "observer-specs", "observer-claims", "controller-specs"} for name in files), + "Ops manifest missing executable guards or includes circular runtime artifacts") + need(sha(root / "execution_spec.py") == sha(__file__), "Execution implementation drift") + observer_path = absolute(spec["release_observer_path"]) + need(observer_path.parent == root / "observer-specs" and observer_path.suffix == ".json", + "Every execution must predeclare its immutable observer path") + parent_id, node, parent_sha = PINS[spec["arm"]] + need(str(spec["original_parent_job_id"]) == parent_id and spec["parent_request"]["sha256"] == parent_sha, + "Only the exact three original production parents are supported") + parent = read_json(bound(spec["parent_request"])) + run_id = "mot-elmo-pickplace-continuous-large314m-" + spec["arm"] + "-robot-human-camera-subtask-s42-300k-20261026-v1" + need(parent["run_id"] == run_id and parent["kind"] == "native-subtask-world-model-v1" + and parent["phase"] == "full" and parent["target_steps"] == 300000 + and parent["parameters"] == 314187064 and parent["world_size"] == 4, "Parent scientific identity differs") + service = spec["operational_service"] + need(service["manifest_sha256"] == SERVICE_SHA, "Wrong published recovery service") + service_names = manifest(service["root"], dict(path=service["manifest"], sha256=SERVICE_SHA)) + need(len(service_names) == SERVICE_FILE_COUNT, "Incomplete adapter recovery service") + template_path = bound(spec["request_template"]) + need(template_path.parent == root / "templates", "Template outside immutable template directory") + template = read_json(template_path) + need(template["schema"] == MODES[mode][0] and template["operational_service"] == service, "Request phase/service differs") + parent_binding = template.get("parent_request") if mode != "full" else template["interruption"]["parent_request"] + need(parent_binding == spec["parent_request"], "Template belongs to another production parent") + native = spec["native_contract"] + _, cpus, memory, limit, gres, _ = MODES[mode] + need(set(native) == NATIVE_KEYS and all(isinstance(v, str) for v in native.values()), "ID-free complete native contract required") + need(native["UserId"] == "paphiwetsa3(41003)" and native["Account"] == "loaner" + and native["Partition"] == "rl2" and native["QOS"] == "loaner-normal" + and native["NodeList"] == native["ReqNodeList"] == node + and native["NumCPUs"] == cpus and native["MinMemoryNode"] == memory + and native["TimeLimit"] == limit and native["TresPerNode"] == gres + and native["Requeue"] == native["Restarts"] == "0" + and native["Command"] == str(root / "allocated.sbatch"), "Native phase resources/owner/node changed") + prefix = "elmo-native-recovery-" if mode not in {"diagnostic", "full"} else "elmo-large314m-continuous-" + expected_prefix = prefix + spec["arm"] + "-" + mode + "-" + need(native["JobName"].startswith(expected_prefix) + and re.fullmatch(r"[A-Za-z0-9_-]{1,128}", native["JobName"]), "Unsafe or out-of-scope job name") + if mode == "full": + need("producer_native" not in template and template["job_name"] == native["JobName"], "Full identity cannot gain producer fields") + service_module(spec, "interruption_admission").compatible(parent, template) + else: + need(template["producer_native"] == native, "Template must omit only the future producer JobId") + if mode != "capture": + for key in ("source", "source_revision", "source_manifest", "source_manifest_sha256", + "data_manifest", "data_manifest_sha256", "lambda_runtime", "data_cache"): + need(template.get(key) == parent.get(key), "Producer changed parent binding: " + key) + if mode == "capture": + need(str(template["parent_job_id"]) == parent_id and template["parent_native"]["JobId"] == parent_id, + "Capture parent ID differs") + return spec, template + + +def submission_argv(spec, spec_path, spec_sha, manifest_sha): + n, root = spec["native_contract"], Path(spec["root"]) + return ["/usr/local/bin/sbatch", "--parsable", "--hold", "--no-requeue", "--export=NONE", + "--account=" + n["Account"], "--qos=" + n["QOS"], "--partition=" + n["Partition"], + "--job-name=" + n["JobName"], "--nodelist=" + n["ReqNodeList"], "--nodes=1", "--ntasks=1", + "--cpus-per-task=" + n["NumCPUs"], "--mem=" + n["MinMemoryNode"], "--time=" + n["TimeLimit"], + "--gres=" + ("gpu:h100:4" if n["TresPerNode"] else "none"), + "--comment=elmo-execution:" + spec_sha, "--chdir=" + str(root), + "--output=" + str(root / "logs/%x-%j.out"), "--error=" + str(root / "logs/%x-%j.err"), + str(root / "allocated.sbatch"), str(root), spec["mode"], str(spec_path), spec_sha, manifest_sha] + + +def held_native(raw, spec, spec_sha, job): + f, n = fields(raw), spec["native_contract"] + need(f.get("JobId") == job and f.get("JobState") == "PENDING" and f.get("Reason") == "JobHeldUser", + "Exact held native acceptance required") + need(all(f.get(k, "") == v for k, v in n.items() if k != "NodeList") + and f.get("NodeList", "") in {"", "(null)", n["NodeList"]} + and f.get("NumNodes") in {"1", "1-1"} and f.get("NumTasks") == "1" + and f.get("Comment") == "elmo-execution:" + spec_sha + and f.get("SubmitTime") not in {None, "Unknown", "None"}, "Held native spec/argv/resources differ") + need(f.get("AllocTRES", "") in {"", "(null)"} and f.get("GRES", "") in {"", "(null)"} + and f.get("StartTime") in {None, "Unknown", "N/A"}, + "Held job already has an allocation/start") + requested = f.get("ReqTRES", "").split(",") + memory = n["MinMemoryNode"] + need("cpu=" + n["NumCPUs"] in requested and "node=1" in requested + and bool(set(requested) & {"mem=" + memory, "mem=" + str(int(memory[:-1]) * 1024) + "M"}), + "Held requested CPU/memory/node TRES differ") + gpu = {part for part in requested if "gpu" in part.lower()} + need(gpu <= ({"gres/gpu=4", "gres/gpu:h100=4"} if n["TresPerNode"] else set()), + "Held requested GPU TRES conflict") + # Pending allocations may not yet expose JOB_GRES. If present it must + # agree; absence never substitutes for eventual RUNNING GPU-index proof. + job_gres = f.get("JOB_GRES", "") + if n["TresPerNode"]: + need("JOB_GRES" not in f or job_gres == "gpu:h100:4", "Held JOB_GRES conflicts with four H100s") + else: + need(job_gres in {"", "(null)"}, "CPU held allocation has JOB_GRES") + for key, value in f.items(): + if key in {"ReqTRES", "AllocTRES", "TresPerNode", "JOB_GRES", "GRES"}: + continue + if key.lower().startswith("tresper") or "gpu" in key.lower(): + # One-task CPU allocation metadata is legitimate; per-task GPU + # requests or any other GPU-bearing auxiliary field are not part + # of the exact spec, even if the primary TRES fields look right. + need("gpu" not in value.lower() and "gpu" not in key.lower(), + "Unexpected held auxiliary GPU resource field") + if key == "TresPerTask": + need(value == "cpu:" + n["NumCPUs"], "Held per-task CPU resources differ") + return f + + +def materialized(spec, template, job): + if spec["mode"] == "full": + return bound(spec["request_template"]).read_bytes() + q = copy.deepcopy(template) + q["producer_native"]["JobId"] = job + return encoded(q) + + +def bind_accepted(spec_path, spec_sha, mode, manifest_sha, job, raw): + """Separate operation after exact accepted native ID; never submits/releases.""" + spec, template = read_spec(spec_path, spec_sha, mode, manifest_sha) + need(str(job).isdecimal(), "Actual accepted job ID required") + job = str(job) + held_native(raw, spec, spec_sha, job) + p = paths(spec, spec_sha, job) + prepared = read_json(p["spec_claim"] / "PREPARED.json") + argv = submission_argv(spec, spec_path, spec_sha, manifest_sha) + need(prepared["spec_sha256"] == spec_sha and prepared["argv"] == argv, "Submission argv/spec claim differs") + need((p["spec_claim"] / "PARENT_JOB_ID").read_text().strip() == job, "Accepted ID not preserved in original spec claim") + need((p["spec_claim"] / "sbatch.stdout").read_text().strip() == job + and (p["spec_claim"] / "sbatch.exit").read_text().strip() == "0", + "Binding requires the preserved unambiguous accepted sbatch outcome") + payload = materialized(spec, template, job) + request_sha = hashlib.sha256(payload).hexdigest() + p = paths(spec, spec_sha, job, request_sha) + write_once(p["request"], payload) + mkdir_once(p["request_claim"]) + write_once(p["request_claim"] / "PREPARED.json", dict(request_sha256=request_sha, spec_sha256=spec_sha, + spec_path=str(spec_path), actual_request_path=str(p["request"]), argv=argv, + original_parent_job_id=str(spec["original_parent_job_id"]))) + write_once(p["request_claim"] / "PARENT_JOB_ID", (job + "\n").encode()) + write_once(p["spec_claim"] / "HELD_NATIVE.txt", raw.encode()) + acceptance = dict(schema="native-elmo-held-execution-binding-v1", spec_sha256=spec_sha, + ops_manifest_sha256=manifest_sha, job_id=job, argv=argv, + request=dict(path=str(p["request"]), sha256=request_sha), + request_claim=str(p["request_claim"]), held_native=dict(path=str(p["spec_claim"] / "HELD_NATIVE.txt"), + sha256=sha(p["spec_claim"] / "HELD_NATIVE.txt"))) + write_once(p["spec_claim"] / "acceptance.json", acceptance) + return acceptance + + +def read_binding(spec_path, spec_sha, mode, manifest_sha): + spec, template = read_spec(spec_path, spec_sha, mode, manifest_sha) + sc = paths(spec, spec_sha)["spec_claim"] + a = read_json(sc / "acceptance.json") + job = a["job_id"] + need(a["schema"] == "native-elmo-held-execution-binding-v1" and str(job).isdecimal() + and a["spec_sha256"] == spec_sha and a["ops_manifest_sha256"] == manifest_sha + and a["argv"] == submission_argv(spec, spec_path, spec_sha, manifest_sha), "Acceptance spec/argv differs") + need(a["held_native"]["path"] == str(sc / "HELD_NATIVE.txt"), "Held source escaped spec claim") + held = held_native(bound(a["held_native"]).read_text(), spec, spec_sha, job) + request_path = bound(a["request"]) + p = paths(spec, spec_sha, job, a["request"]["sha256"]) + need(request_path == p["request"] and a["request_claim"] == str(p["request_claim"]) + and request_path.read_bytes() == materialized(spec, template, job), "Actual request differs from sole authorized late binding") + for directory in (sc, p["request_claim"]): + need(absolute(directory).is_dir() and (directory / "PARENT_JOB_ID").read_text().strip() == job, + "Spec/request accepted job mismatch") + prepared = read_json(p["request_claim"] / "PREPARED.json") + need(prepared["request_sha256"] == a["request"]["sha256"] and prepared["spec_sha256"] == spec_sha + and prepared["argv"] == a["argv"] and prepared["actual_request_path"] == str(request_path) + and prepared["spec_path"] == str(spec_path) + and prepared["original_parent_job_id"] == str(spec["original_parent_job_id"]), "Canonical request-SHA claim mismatch") + original = read_json(sc / "PREPARED.json") + need(original["spec_sha256"] == spec_sha and original["argv"] == a["argv"] + and original["spec_path"] == str(spec_path) and original["arm"] == spec["arm"] + and original["parent_request"] == spec["parent_request"], "Original submission intent changed") + need((sc / "sbatch.stdout").read_text().strip() == job and (sc / "sbatch.exit").read_text().strip() == "0", + "Original submission acceptance became unresolved") + return spec, request_path, read_json(request_path), job, p["request_claim"], held + + +def service_module(spec, name): + """Activate only the verified published service before importing its stdlib guards.""" + need(isinstance(SERVICE_COMMIT, str) and re.fullmatch(r"[0-9a-f]{40}", SERVICE_COMMIT), + "Adapter service is unpublished; explicit reviewed publication binding required") + need(name in {"native_interruption", "interruption_admission", "original_observer", "phase_observer_proof"}, + "Unapproved pre-Torch service import") + binding = spec["operational_service"] + need(binding["manifest_sha256"] == SERVICE_SHA, "Wrong published recovery service") + root = absolute(binding["root"]) + names = manifest(root, dict(path=binding["manifest"], sha256=SERVICE_SHA)) + need(len(names) == SERVICE_FILE_COUNT, "Incomplete adapter recovery service") + path = root / "scripts/train/joint/service_bootstrap.py" + import_spec = importlib.util.spec_from_file_location("_elmo_bound_bootstrap", path) + module = importlib.util.module_from_spec(import_spec) + import_spec.loader.exec_module(module) + module.activate(dict(operational_service=spec["operational_service"])) + return importlib.import_module("scripts.train.joint." + name) + + +def load_execution(spec_path, spec_sha, mode, manifest_sha): + """Read-only guard, safe before Torch; authenticates native in its own environment.""" + spec, request_path, q, job, claim, held = read_binding(spec_path, spec_sha, mode, manifest_sha) + expected = dict(spec["native_contract"], JobId=job) + native = service_module(spec, "native_interruption") + native.process_identity(expected, cpu=mode not in {"diagnostic", "full"}) + raw = command(["/usr/local/bin/scontrol", "show", "job", job, "-o", "-dd"]) + f = native.native_identity(raw, expected, "RUNNING") + need(f.get("SubmitTime") == held["SubmitTime"] and f.get("Comment") == "elmo-execution:" + spec_sha, + "RUNNING job is not the held execution") + release = read_json(paths(spec, spec_sha)["spec_claim"] / "RELEASE_ATTEMPT.json") + need(release["job_id"] == job and release["spec_sha256"] == spec_sha + and release["request_sha256"] == sha(request_path), "Preserved release attempt differs") + return spec, request_path, q, job, claim + + +def record_running_admission(spec_path, spec_sha, mode, manifest_sha): + spec, path, q, job, claim = load_execution(spec_path, spec_sha, mode, manifest_sha) + raw = command(["/usr/local/bin/scontrol", "show", "job", job, "-o", "-dd"]) + service_module(spec, "native_interruption").native_identity(raw, dict(spec["native_contract"], JobId=job), "RUNNING") + record = dict(job_id=job, request_sha256=sha(path), spec_sha256=spec_sha, native=raw) + write_once(claim / "RUNNING_ADMISSION.json", record) + return record + + +def execution_claim_census(spec, *, current_job=None, current_request_sha=None, allow_held=False): + """All registered spec ledgers, including claims without accepted/materialized IDs. + + Never equate an absent/partial acceptance with no job. Current held admission + is separate from the RUNNING-only scientific gate and preserves raw state. + """ + need((current_job is None) == (current_request_sha is None), "Incomplete current execution binding") + native = service_module(spec, "native_interruption") + observer = service_module(spec, "original_observer") + step_query = observer.query(observer.step_argv("paphiwetsa3")) + observer.query_record(step_query, observer.step_argv("paphiwetsa3")) + need(step_query["returncode"] == 0 and not step_query["stderr"].strip(), "Complete owned step query failed") + observer.strict_step_rows(step_query["stdout"]) + evidence, seen = [], set() + for entry in spec["execution_claim_roots"]: + directory = absolute(entry["path"]) + for claim in sorted(directory.iterdir()): + need(claim.is_dir() and not claim.is_symlink() and re.fullmatch(r"[0-9a-f]{64}", claim.name), + "Unknown execution ledger entry requires inspection") + prepared = read_json(claim / "PREPARED.json") + need(prepared["spec_sha256"] == claim.name and prepared["arm"] == spec["arm"] + and prepared["parent_request"] == spec["parent_request"], "Execution root includes another identity") + previous_spec_path = bound(dict(path=prepared["spec_path"], sha256=claim.name)) + previous = read_json(previous_spec_path) + need(previous["arm"] == spec["arm"] and previous["parent_request"] == spec["parent_request"] + and paths(previous, claim.name)["spec_claim"] == claim, + "Historical execution claim escaped its actual spec root") + need((claim / "acceptance.json").is_file(), "Unresolved submission or partial actual-ID binding; inspect, never retry") + args = (str(previous_spec_path), claim.name, previous["mode"], prepared["ops_manifest_sha256"]) + old, path, q, job, request_claim, held = read_binding(*args) + need(job not in seen, "Duplicate actual job across execution claims") + seen.add(job) + expected = dict(old["native_contract"], JobId=job) + if job == current_job: + need(Path(old["root"]) == Path(spec["root"]) and sha(path) == current_request_sha, + "Current job belongs to another execution/request") + raw = command(["/usr/local/bin/scontrol", "show", "job", job, "-o", "-dd"]) + if allow_held: + f = held_native(raw, old, claim.name, job) + service_module(spec, "phase_observer_proof").check_owned_steps(step_query["stdout"], job) + else: + f = native.native_identity(raw, expected, "RUNNING") + need(f.get("SubmitTime") == held["SubmitTime"] and f.get("Comment") == "elmo-execution:" + claim.name, + "Current native identity differs from accepted execution") + evidence.append(dict(job_id=job, spec_sha256=claim.name, current=True, actual_state=f["JobState"])) + continue + captured, raw = observed_terminal(*args, expected_job=job, expected_request_sha=sha(path), + terminal_directory=entry["terminal_directory"]) + service_module(spec, "phase_observer_proof").check_owned_steps(step_query["stdout"], job) + evidence.append(dict(job_id=job, spec_sha256=claim.name, observer_spec=captured["observer_spec"], + terminal_sha256=captured["producer_native"]["sha256"], terminal_state=captured["producer_terminal_state"])) + if current_job is not None: + need(current_job in seen, "Current job absent from complete execution-claim census") + return evidence + + +def operations_module(spec, name): + """Load fixed operations code only after the caller's entire immutable closure.""" + need(name in {"host_guard", "wrapper_common", "phase_observer"}, "Unapproved operations module") + root = absolute(spec["root"]) + names = manifest(root, spec["ops_manifest"]) + need("execution_spec.py" in names and sha(root / "execution_spec.py") == sha(__file__), + "Operations caller implementation drift") + need(name + ".py" in names, "Operations module absent from immutable manifest") + if name == "host_guard": + operations_module(spec, "wrapper_common") + location = root / (name + ".py") + existing = sys.modules.get(name) + if existing is not None: + need(getattr(existing, "__file__", None) == str(location) + and getattr(existing, "_LOADED_IMPLEMENTATION_SHA256", None) == sha(location) + and getattr(existing, "_verified_ops_manifest_sha256", None) in {None, spec["ops_manifest"]["sha256"]}, + "Foreign or unverified operations module already imported") + existing._verified_ops_manifest_sha256 = spec["ops_manifest"]["sha256"] + existing._verified_ops_member_sha256 = sha(location) + return existing + loader = importlib.util.spec_from_file_location(name, location) + module = importlib.util.module_from_spec(loader) + module._verified_ops_manifest_sha256 = spec["ops_manifest"]["sha256"] + module._verified_ops_member_sha256 = sha(location) + sys.modules[name] = module + try: + loader.loader.exec_module(module) + except BaseException: + del sys.modules[name] + raise + return module + + +def observed_terminal(spec_path, spec_sha, mode, manifest_sha, *, expected_job, + expected_request_sha, terminal_directory): + """Observer final commit is required even when an isolated raw terminal exists.""" + spec, path, q, job, claim, held = read_binding(spec_path, spec_sha, mode, manifest_sha) + need(job == expected_job and sha(path) == expected_request_sha, + "Observer terminal belongs to another accepted request") + observer_path = absolute(spec["release_observer_path"]) + need(observer_path.parent == Path(spec["root"]) / "observer-specs", "Observer escaped execution root") + result = operations_module(spec, "phase_observer").verify_captured(str(observer_path), sha(observer_path), manifest_sha) + need(result["schema"] == "native-elmo-phase-observer-result-v1" + and result["status"] == "GENUINE_PRODUCER_TERMINAL_CAPTURED" + and result["producer_execution"] == dict(path=str(spec_path), sha256=spec_sha, mode=mode, ops_manifest_sha256=manifest_sha) + and result["producer_request"] == dict(path=str(path), sha256=expected_request_sha) + and result["producer_job_id"] == job and result["producer_phase"] == mode + and result["observer_spec"] == dict(path=str(observer_path), sha256=sha(observer_path)), + "Observer final commit identity differs from accepted execution") + terminal = absolute(terminal_directory) / ("native-" + job + ".txt") + need(result["producer_native"]["path"] == str(terminal), "Terminal escaped registered claim evidence directory") + raw = bound(result["producer_native"]).read_text() + native = service_module(spec, "native_interruption") + state = native.fields(raw).get("JobState") + need(state in native.TERMINAL and state == result["producer_terminal_state"], "Observer terminal state differs") + f = native.native_identity(raw, dict(spec["native_contract"], JobId=job), state) + need(f.get("SubmitTime") == held["SubmitTime"] and f.get("Comment") == "elmo-execution:" + spec_sha, + "Terminal identity differs from actual held execution") + return result, raw + + +def request_claim_terminal(spec, claim, prepared, job, terminal_directory): + """Follow a canonical request-SHA claim back to its immutable execution spec.""" + old_path = bound(dict(path=prepared["spec_path"], sha256=prepared["spec_sha256"])) + old = read_json(old_path) + need(old["arm"] == spec["arm"] and old["parent_request"] == spec["parent_request"], + "Prior scientific claim belongs to another identity") + need(paths(old, prepared["spec_sha256"], job, claim.name)["request_claim"] == claim, + "Prior request claim escaped actual execution root") + return observed_terminal(str(old_path), prepared["spec_sha256"], old["mode"], old["ops_manifest"]["sha256"], + expected_job=job, expected_request_sha=claim.name, terminal_directory=terminal_directory) + + +def host_preflight(spec, template, *, current_job=None, current_request_sha=None): + operations_module(spec, "host_guard").admit_before_submission( + spec, template, current_job=current_job, current_request_sha=current_request_sha) + + +def submit_held(spec_path, spec_sha, mode, manifest_sha): + """One sbatch maximum. Any existing claim/unknown outcome stops this operation.""" + spec, template = read_spec(spec_path, spec_sha, mode, manifest_sha) + host_preflight(spec, template) + sc = paths(spec, spec_sha)["spec_claim"] + mkdir_once(sc) # exclusive durable claim before the sole native side effect + argv = submission_argv(spec, spec_path, spec_sha, manifest_sha) + write_once(sc / "PREPARED.json", dict(spec_sha256=spec_sha, ops_manifest_sha256=manifest_sha, + spec_path=str(spec_path), arm=spec["arm"], parent_request=spec["parent_request"], + argv=argv, original_parent_job_id=str(spec["original_parent_job_id"]), + created_at=datetime.now(timezone.utc).isoformat())) + try: + result = subprocess.run(argv, capture_output=True, text=True, timeout=30) + except (OSError, subprocess.TimeoutExpired) as error: + write_once(sc / "SUBMISSION_UNKNOWN.json", dict(error=type(error).__name__, inspect_required=True)) + raise RuntimeError("Submission outcome unknown; inspect exact existing spec claim/native; never retry") from error + write_once(sc / "sbatch.stdout", result.stdout.encode()) + write_once(sc / "sbatch.stderr", result.stderr.encode()) + write_once(sc / "sbatch.exit", (str(result.returncode) + "\n").encode()) + job = result.stdout.strip() + need(result.returncode == 0 and job.isdecimal(), "Submission outcome unresolved; inspect existing claim/native; never retry") + write_once(sc / "PARENT_JOB_ID", (job + "\n").encode()) + raw = command(["/usr/local/bin/scontrol", "show", "job", job, "-o", "-dd"]) + return bind_accepted(spec_path, spec_sha, mode, manifest_sha, job, raw) + + +def require_release_observer(spec, spec_sha, manifest_sha, job, request_sha): + """Fresh actual observer dependency is mandatory for every producer release.""" + observer_path = absolute(spec["release_observer_path"]) + binding = {"path": str(observer_path), "sha256": sha(observer_path)} + observer = read_json(bound(binding)) + need(observer["producer_execution"] == dict(path=read_json(paths(spec, spec_sha)["spec_claim"] / "PREPARED.json")["spec_path"], + sha256=spec_sha, mode=spec["mode"], ops_manifest_sha256=manifest_sha), + "Observer spec does not bind this exact producer execution") + need(observer_path.parent == Path(spec["root"]) / "observer-specs", + "Release observer spec escaped exact operations root") + module = operations_module(spec, "phase_observer") + ready = module.verify_ready_for_release(str(observer_path), binding["sha256"], manifest_sha) + need(ready["schema"] == "native-elmo-phase-observer-ready-v1" + and ready["observer_spec_sha256"] == binding["sha256"] + and ready["production_completion_proven"] is False + and ready["producer_job_id"] == job and ready["producer_request_sha256"] == request_sha + and str(ready["observer_job_id"]).isdecimal(), "Observer readiness belongs to another producer") + return ready + + +def release_bound(spec_path, spec_sha, mode, manifest_sha): + """Separate release only after immutable late binding; an unknown release is never repeated.""" + spec, path, q, job, claim, held = read_binding(spec_path, spec_sha, mode, manifest_sha) + host_preflight(spec, q, current_job=job, current_request_sha=sha(path)) + raw = command(["/usr/local/bin/scontrol", "show", "job", job, "-o", "-dd"]) + current = held_native(raw, spec, spec_sha, job) + need(current["SubmitTime"] == held["SubmitTime"], "Accepted held identity changed") + observer = require_release_observer(spec, spec_sha, manifest_sha, job, sha(path)) + sc = paths(spec, spec_sha)["spec_claim"] + write_once(sc / "RELEASE_HELD_NATIVE.txt", raw.encode()) + write_once(sc / "RELEASE_ATTEMPT.json", dict(job_id=job, spec_sha256=spec_sha, + request_sha256=sha(path), held_native_sha256=hashlib.sha256(raw.encode()).hexdigest(), + observer_readiness=observer)) + result = subprocess.run(["/usr/local/bin/scontrol", "release", job], capture_output=True, text=True, timeout=30) + write_once(sc / "RELEASE_RESULT.json", dict(job_id=job, returncode=result.returncode, + stdout=result.stdout, stderr=result.stderr, completion_proven=False)) + need(result.returncode == 0, "Release unresolved; inspect exact native state; never retry release") + return dict(job_id=job, request_sha256=sha(path), status="RELEASE_ACCEPTED_NOT_RUNNING_OR_COMPLETED") + + +def inspect_existing(spec_path, spec_sha, mode, manifest_sha): + """Read only. Report exact queue candidates; do not infer an ID or repair a claim.""" + spec, _ = read_spec(spec_path, spec_sha, mode, manifest_sha) + sc = paths(spec, spec_sha)["spec_claim"] + queue = command(["/usr/local/bin/squeue", "--noheader", "--user", "paphiwetsa3", "--format=%i|%j|%T|%k"]) + rows = [] + for line in queue.splitlines(): + row = line.split("|") + need(len(row) == 4, "Ambiguous owner queue") + if row[1] == spec["native_contract"]["JobName"] or row[3] == "elmo-execution:" + spec_sha: + rows.append(row) + return dict(spec_claim=str(sc), claim_exists=sc.exists(), native_candidates=rows, + automatic_retry_authorized=False, binding_exists=(sc / "acceptance.json").is_file()) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("operation", choices=["plan", "submit-held", "release-bound", "inspect-existing"]) + parser.add_argument("spec_path"); parser.add_argument("spec_sha"); parser.add_argument("mode", choices=MODES) + parser.add_argument("manifest_sha") + args = parser.parse_args() + values = (args.spec_path, args.spec_sha, args.mode, args.manifest_sha) + if args.operation == "plan": + spec, _ = read_spec(*values) + result = dict(argv=submission_argv(spec, args.spec_path, args.spec_sha, args.manifest_sha), submitted=False) + else: + result = {"submit-held": submit_held, "release-bound": release_bound, "inspect-existing": inspect_existing}[args.operation](*values) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/native_elmo_ops/fa3-binding.json b/scripts/train/joint/native_elmo_ops/fa3-binding.json new file mode 100644 index 000000000..415bbf248 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/fa3-binding.json @@ -0,0 +1,90 @@ +{ + "build_flags": { + "FLASHATTENTION_DISABLE_APPENDKV": true, + "FLASHATTENTION_DISABLE_BACKWARD": false, + "FLASHATTENTION_DISABLE_CLUSTER": false, + "FLASHATTENTION_DISABLE_FP16": true, + "FLASHATTENTION_DISABLE_FP8": true, + "FLASHATTENTION_DISABLE_HDIM128": false, + "FLASHATTENTION_DISABLE_HDIM192": true, + "FLASHATTENTION_DISABLE_HDIM256": true, + "FLASHATTENTION_DISABLE_HDIM64": true, + "FLASHATTENTION_DISABLE_HDIM96": true, + "FLASHATTENTION_DISABLE_LOCAL": true, + "FLASHATTENTION_DISABLE_PACKGQA": true, + "FLASHATTENTION_DISABLE_PAGEDKV": true, + "FLASHATTENTION_DISABLE_SM8x": true, + "FLASHATTENTION_DISABLE_SOFTCAP": true, + "FLASHATTENTION_DISABLE_SPLIT": true, + "FLASHATTENTION_DISABLE_VARLEN": false, + "FLASHATTENTION_ENABLE_VCOLMAJOR": false, + "FLASH_ATTENTION_DISABLE_HDIMDIFF192": true, + "FLASH_ATTENTION_DISABLE_HDIMDIFF64": true + }, + "evidence_sha256": { + "evidence/fa3-source-build/binary-build-paths.txt": "2256c6cdbd481ffda8494dd3225c39629778a40489fba245ff164d37165a9931", + "evidence/fa3-source-build/fa3-build-411351.err": "d78a09b06b4457c937a66373fc3817f8cd51b27560d780912019b48efadba161", + "evidence/fa3-source-build/fa3-build-411351.out": "f274035b51e41df0a91918734bdd908e41afbaab21e537dab080263b3defb480", + "evidence/fa3-source-build/installed-sha256.txt": "b070608e6e03ea0f39f92bbe281b26808723fd1d47577b2cadfd29d234c493d5" + }, + "file_sha256": { + "flash_attn_3-3.0.1.dist-info/METADATA": "a12ff5f3e6d15eff14fdaeb987e0398d3f93289e3552fd54bc7d9112921413b1", + "flash_attn_3-3.0.1.dist-info/RECORD": "38b95b864644497a000c9a50020e1e4a6938683548fd683294c2b38a9bdf8173", + "flash_attn_3-3.0.1.dist-info/direct_url.json": "c6a43d33aeaa0fb199d8bd290eed679b892e08f2553e21fc9137a7f35552001f", + "flash_attn_3/_C.abi3.so": "0f978aee4dde000b10a7fda072469253caa2244c30a27e7275e3dab4bf8497d2", + "flash_attn_3/__init__.py": "362bca1d0502eca70a2d051eea4ad27b5ea8ec94132df156d61ff4326e6aaa23", + "flash_attn_3/flash_attn_config.py": "255bbac5b1c027fc5862396753e81cf66ffbc7e3c40ee3b5d1a098efc842be00", + "flash_attn_3/flash_attn_interface.py": "39f173452b4906d0ba7371e4c21947be7c7156d06825ec058364fd3fa4b0de49", + "flash_attn_config.py": "bbd989bb61127bbc9cfa6cd401ec788efc3239559845e03e69b34cae442d556b", + "flash_attn_interface.py": "887eb85ea5380707a78cd15dc04169e654f4dee58f9bddab49a2ce71f1680e35" + }, + "root": "/workspace/users/paphiwetsa3/tasks/pixel-video-gripper-optimization-20261017-v1/runtime/fa3", + "schema": "fa3-runtime-binding-v1", + "source": { + "subdirectory": "hopper", + "url": "https://github.com/Dao-AILab/flash-attention.git", + "vcs_info": { + "commit_id": "965c9212e6be33c821ee9f2b461564c5c91a2503", + "requested_revision": "965c9212e6be33c821ee9f2b461564c5c91a2503", + "vcs": "git" + } + }, + "source_build_proof": { + "embedded_paths": [ + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/algorithm", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/arch", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/atom", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/container", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/numeric", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cute/util", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cutlass", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cutlass/arch", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cutlass/detail", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cutlass/pipeline", + "/tmp/pip-req-build-nfojukf_/csrc/cutlass/include/cutlass/platform", + "/tmp/pip-req-build-nfojukf_/hopper", + "/tmp/pip-req-build-nfojukf_/hopper/build/temp.linux-x86_64-cpython-312", + "/tmp/pip-req-build-nfojukf_/hopper/flash_api_stable.cpp", + "/tmp/pip-req-build-nfojukf_/hopper/flash_bwd_launch_template.h", + "/tmp/pip-req-build-nfojukf_/hopper/flash_fwd_launch_template.h", + "/tmp/pip-req-build-nfojukf_/hopper/flash_prepare_scheduler.cu", + "/tmp/pip-req-build-nfojukf_/hopper/instantiations/flash_bwd_hdim128_bf16_sm90.cu", + "/tmp/pip-req-build-nfojukf_/hopper/instantiations/flash_fwd_hdim128_bf16_sm90.cu" + ], + "method": "installed_binary_embeds_unique_build_checkout_paths", + "node": "gpu19", + "pip_build_directory": "/tmp/pip-req-build-nfojukf_", + "result": "SOURCE_COMPILED", + "slurm_job_id": 411351 + }, + "verification_scope": "Read-only installed artifact identity and source-build provenance; CUDA correctness and throughput are separate gates.", + "versions": { + "flash-attn-3": "3.0.1" + }, + "wheel": { + "filename": "flash_attn_3-3.0.1-cp310-abi3-linux_x86_64.whl", + "sha256": "2811ee540413366dcefadadda16369bee7c84f56fd51d4e167b74670cbbad0c9", + "size_bytes": 2082988 + } +} diff --git a/scripts/train/joint/native_elmo_ops/gpu_health_r3.py b/scripts/train/joint/native_elmo_ops/gpu_health_r3.py new file mode 100644 index 000000000..640f2cc18 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/gpu_health_r3.py @@ -0,0 +1,225 @@ +"""Read-only admission of the one CUDA-visible H100 assigned to this job. + +NVIDIA documents GPU Recovery Action as the replacement for deprecated reset +flags. An N/A reset flag is therefore acceptable with Recovery Action=None. +N/A ECC values are not zeros: require a usable legacy total or both SRAM/DRAM +coverage in each of the volatile and aggregate sections. Row-remap pending and +failure flags must both be explicitly clear on this H100 workload. + +Accept both row_remapper and the remapped_rows XML emitted by the allocated +580.105.08 driver. Every reported pending/failure alias must be explicitly +clear; one clear schema cannot hide an unsafe second schema. Row counts are +historical remapping totals, recorded separately from these status flags. + +Reference: https://docs.nvidia.com/deploy/nvidia-smi/index.html +Unit fixtures are synthetic schema examples. The allocation's actual XML is +saved alongside the JSON receipt; these checks never reset or modify hardware. +""" +from __future__ import annotations + +import argparse +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +import re +import subprocess +import uuid +import xml.etree.ElementTree as ET + + +class HealthError(ValueError): + pass + + +def _uuid(value): + if isinstance(value, bytes): + value = str(uuid.UUID(bytes=value)) if len(value) == 16 else value.decode("ascii") + value = str(value) + if value.startswith("GPU-"): + value = value[4:] + try: + return "GPU-" + str(uuid.UUID(value)) + except (ValueError, AttributeError) as error: + raise HealthError("CUDA did not expose a valid full GPU UUID") from error + + +def _text(node, path): + value = node.findtext(path) + return value.strip() if value is not None else None + + +def _clear(value, label): + if value is None or value.lower() not in {"no", "false", "0"}: + raise HealthError(f"{label} is not explicitly clear: {value!r}") + + +def _ecc(scope, label): + if scope is None: + raise HealthError(f"Missing {label} ECC counters") + numeric, unavailable = {}, [] + + def visit(node, prefix=""): + path = f"{prefix}/{node.tag}" if prefix else node.tag + if list(node): + for child in node: + visit(child, path) + return + if not any(word in path.lower() for word in ("uncorrectable", "uncorrected", "double_bit")): + return + value = (node.text or "").strip() + if value.upper() in {"N/A", "NOT SUPPORTED"}: + unavailable.append(path) + elif not re.fullmatch(r"\d+", value): + raise HealthError(f"Ambiguous {label} ECC counter {path}: {value!r}") + else: + numeric[path] = int(value) + if numeric[path] != 0: + raise HealthError(f"Nonzero {label} uncorrectable ECC counter {path}: {value}") + + visit(scope) + names = {path.rsplit("/", 1)[-1] for path in numeric} + legacy_total = any(path.endswith("/double_bit/total") for path in numeric) + dram = bool(names & {"dram_uncorrectable", "dram_uncorrected"}) + sram = bool(names & {"sram_uncorrectable", "sram_uncorrected"}) or { + "sram_uncorrectable_parity", "sram_uncorrectable_secded" + }.issubset(names) + if not (legacy_total or (dram and sram)): + raise HealthError(f"Incomplete {label} uncorrectable ECC coverage") + return {"zero_counters": numeric, "unavailable_redundant_counters": unavailable} + + +def _row_remapper(gpu): + """Normalize both XML schemas without interpreting missing flags as clear.""" + sections = [ + node for node in gpu if node.tag in {"row_remapper", "remapped_rows"} + ] + if not sections: + raise HealthError("Missing H100 row-remapper status") + result = {"pending": {}, "failure": {}, "counts": {}, "schemas": []} + for index, row in enumerate(sections): + # Preserve every value, including duplicate aliases, so contradictory + # status fields cannot be hidden by Element.find's first-match behavior. + prefix = f"{row.tag}[{index}]" + result["schemas"].append(row.tag) + for label, aliases in { + "pending": ("pending", "row_remapping_pending", "remapped_row_pending"), + "failure": ( + "remapping_failure_occurred", "row_remapping_failure", + "remapped_row_failure", + ), + }.items(): + fields = [child for child in row if child.tag in aliases] + if not fields: + raise HealthError(f"Missing row-remapper {label} flag in {row.tag}") + for field_index, field in enumerate(fields): + path = f"{prefix}/{field.tag}[{field_index}]" + value = (field.text or "").strip() + _clear(value, path) + result[label][path] = value + for field_index, field in enumerate(row): + if field.tag not in {"remapped_row_corr", "remapped_row_unc"}: + continue + path = f"{prefix}/{field.tag}[{field_index}]" + value = (field.text or "").strip() + if not re.fullmatch(r"\d+", value): + raise HealthError(f"Ambiguous row-remapper count {path}: {value!r}") + # A completed remapping is not a pending or failed remapping. + # The existing admission policy requires clear status and zero + # uncorrectable ECC, not zero historical correctable row remaps. + result["counts"][path] = int(value) + return result + + +def inspect_xml(xml, expected_uuid): + """Validate selected-device XML without consulting or modifying hardware.""" + expected_uuid = _uuid(expected_uuid) + try: + root = ET.fromstring(xml) + except ET.ParseError as error: + raise HealthError("nvidia-smi did not return valid XML") from error + devices = root.findall("gpu") + if len(devices) != 1: + raise HealthError("Selected-device nvidia-smi query must return exactly one GPU") + gpu = devices[0] + if _uuid(_text(gpu, "uuid")) != expected_uuid: + raise HealthError("nvidia-smi UUID differs from the CUDA-visible device") + name = _text(gpu, "product_name") or "" + if not re.search(r"\bH100\b", name): + raise HealthError(f"Expected an H100, got {name!r}") + checks = {"uuid": expected_uuid, "product_name": name} + checks["uncorrectable_ecc"] = { + scope: _ecc(gpu.find(f"ecc_errors/{scope}"), scope) + for scope in ("volatile", "aggregate") + } + recovery = _text(gpu, "gpu_recovery_action") + reset = { + field: _text(gpu, f"gpu_reset_status/{field}") + for field in ("reset_required", "drain_and_reset_recommended") + } + for field, value in reset.items(): + if value is not None and value.upper() not in {"N/A", "NOT SUPPORTED"}: + _clear(value, field) + if recovery is not None and recovery.upper() not in {"N/A", "NOT SUPPORTED"}: + if recovery.lower() != "none": + raise HealthError(f"GPU recovery action required: {recovery}") + recovery_basis = "gpu_recovery_action" + else: + for field, value in reset.items(): + _clear(value, field) + recovery_basis = "legacy_reset_flags" + checks["recovery"] = {"action": recovery, "reset_flags": reset, "basis": recovery_basis} + checks["row_remapper"] = _row_remapper(gpu) + # Newer drivers may expose these repair flags; absent/unsupported optional + # fields are covered by the required recovery status, not treated as zeros. + repair = {} + for leaf in gpu.iter(): + if leaf.tag in {"sram_threshold_exceeded", "channel_repair_pending", "tpc_repair_pending", "unrepairable_memory"}: + value = (leaf.text or "").strip() + if value.upper() not in {"N/A", "NOT SUPPORTED"}: + _clear(value, leaf.tag) + repair[leaf.tag] = value + checks["optional_repair_flags"] = repair + return checks + + +def admit_gpu(receipt_path): + """Save admission evidence and raise on an unhealthy/ambiguous allocation.""" + import torch + + receipt_path = Path(receipt_path) + receipt_path.parent.mkdir(parents=True, exist_ok=True) + xml_path = receipt_path.with_suffix(".xml") + receipt = { + "schema": "dino-proprio-gpu-admission-v1", + "checked_at_utc": datetime.now(timezone.utc).isoformat(), + "status": "FAIL", + "policy": "one-H100-zero-volatile-and-aggregate-uncorrectable-ecc-clear-recovery-and-row-remap", + } + try: + if not torch.cuda.is_available() or torch.cuda.device_count() != 1: + raise HealthError("Exactly one CUDA-visible GPU is required") + properties = torch.cuda.get_device_properties(0) + if not re.search(r"\bH100\b", properties.name): + raise HealthError(f"CUDA-visible GPU is not H100: {properties.name!r}") + assigned_uuid = _uuid(getattr(properties, "uuid", None)) + receipt["cuda_device"] = {"ordinal": 0, "name": properties.name, "uuid": assigned_uuid} + command = ["nvidia-smi", "-q", "-x", "-i", assigned_uuid] + result = subprocess.run(command, check=True, capture_output=True, text=True, timeout=30) + xml_path.write_text(result.stdout) + receipt["xml"] = {"path": str(xml_path), "sha256": hashlib.sha256(result.stdout.encode()).hexdigest()} + receipt["checks"] = inspect_xml(result.stdout, assigned_uuid) + receipt["status"] = "PASS" + except Exception as error: + receipt["error"] = f"{type(error).__name__}: {error}" + receipt_path.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n") + raise HealthError(f"GPU admission failed; see {receipt_path}: {error}") from error + receipt_path.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n") + return receipt + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--receipt", type=Path, required=True) + options = parser.parse_args() + print(json.dumps(admit_gpu(options.receipt), sort_keys=True)) diff --git a/scripts/train/joint/native_elmo_ops/gpu_health_set.py b/scripts/train/joint/native_elmo_ops/gpu_health_set.py new file mode 100644 index 000000000..1f4e68f6f --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/gpu_health_set.py @@ -0,0 +1,43 @@ +"""Read-only admission of the exact requested CUDA-visible H100s in one run step.""" +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +import subprocess +from gpu_health_r3 import HealthError, _uuid, inspect_xml + + +def admit_gpu_set(receipt_path, expected): + if expected not in (2, 4): + raise HealthError("Expected two/four dedicated H100s") + import torch + receipt_path = Path(receipt_path) + receipt_path.parent.mkdir(parents=True, exist_ok=True) + if receipt_path.exists(): + raise HealthError("Refusing to overwrite GPU-set admission") + record = {"schema": "joint-h100-admission-v1", "status": "FAIL", "devices": [], + "checked_at_utc": datetime.now(timezone.utc).isoformat()} + try: + if not torch.cuda.is_available() or torch.cuda.device_count() != expected: + raise HealthError("The exact requested CUDA-visible GPU count is required") + for ordinal in range(expected): + properties = torch.cuda.get_device_properties(ordinal) + assigned_uuid = _uuid(properties.uuid) + result = subprocess.run(["nvidia-smi", "-q", "-x", "-i", assigned_uuid], + check=True, capture_output=True, text=True, timeout=30) + xml_path = receipt_path.with_suffix(f".gpu{ordinal}.xml") + xml_path.write_text(result.stdout) + checks = inspect_xml(result.stdout, assigned_uuid) + record["devices"].append({"ordinal": ordinal, "uuid": assigned_uuid, + "name": properties.name, "checks": checks, + "xml": str(xml_path), + "xml_sha256": hashlib.sha256(result.stdout.encode()).hexdigest()}) + if len({device["uuid"] for device in record["devices"]}) != expected: + raise HealthError("GPU UUIDs must be distinct") + record["status"] = "PASS" + except Exception as error: + record["error"] = f"{type(error).__name__}: {error}" + receipt_path.write_text(json.dumps(record, indent=2) + "\n") + raise + receipt_path.write_text(json.dumps(record, indent=2) + "\n") + return record diff --git a/scripts/train/joint/native_elmo_ops/host_guard.py b/scripts/train/joint/native_elmo_ops/host_guard.py new file mode 100644 index 000000000..b1be5f840 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/host_guard.py @@ -0,0 +1,234 @@ +"""Actual read-only admission before claim, before release, and inside allocation. + +Metadata proofs are re-opened here. Large checkpoint hashing/loading remains in +the allocated maintained producers. No flag or mock receipt grants readiness. +""" +import json +import hashlib +import os +from pathlib import Path +import re +import shutil +from wrapper_common import bound, command, fields, need, physical, read_json, service, sha, now + +_LOADED_IMPLEMENTATION_SHA256 = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + + +def context(spec, q): + parent = bound(spec['parent_request']) + need(parent['parameters'] == 314187064 and parent['world_size'] == 4 + and parent['kind'] == 'native-subtask-world-model-v1' and parent['target_steps'] == 300000, + 'Exact frozen native campaign parent required') + runtime = spec['runtime'] + need(runtime['schema'] == 'native-elmo-wrapper-runtime-v1', 'Explicit wrapper runtime contract required') + need(spec['native_contract']['NodeList'] == parent['lambda_runtime']['node'] == parent['data_cache']['node'], + 'New node requires new actual matching runtime/cache/diagnostic proof') + admission = service(spec, 'interruption_admission') + need(q['operational_service'] == spec['operational_service'], 'Producer service differs from execution spec') + if spec['mode'] not in {'capture', 'full'}: + for key in ('source', 'source_revision', 'source_manifest', 'source_manifest_sha256', + 'data_manifest', 'data_manifest_sha256', 'lambda_runtime', 'data_cache'): + need(q.get(key) == parent.get(key), 'Producer changed parent science/runtime/cache: ' + key) + if spec['mode'] == 'full': + admission.compatible(parent, q) + gate, result = admission.completed_producer(q['interruption']['admission'], + 'native-elmo-interruption-finalized-v1', 'FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS') + need(result['resume_request_fingerprint'] == admission.fingerprint(q) + and result['checkpoint'] == q['interruption']['checkpoint'] + and result['run_id'] == parent['run_id'] and result['original_total_target'] == 300000 + and result['actual_updates'] == 6 and result['diagnostic_scope'] == 'production-checkpoint-only' + and result['resume_native_contract'] == spec['native_contract'], + 'Completed finalizer does not admit this exact fresh full request') + need(result['parent_request'] == gate['parent_request'] == spec['parent_request'], + 'Completed finalizer belongs to another parent') + intent_hash = admission.current_intent(result['current_intent_path'], parent['run_id'], result['current_intent_sha256']) + else: + gate = q + need(gate['parent_request'] == spec['parent_request'], 'Producer parent binding differs') + intent_hash = admission.current_intent(gate['current_intent_path'], parent['run_id']) + need(runtime['admission'] == {k: gate[k] for k in ('current_intent_path', 'quiescence')}, + 'Submission and allocated intent/claim scope differ') + return parent, gate, admission, intent_hash + + +def phase_evidence(spec, q, gate, admission): + native = service(spec, 'native_interruption') + if spec['mode'] == 'capture': + expected = q['parent_native']; job = str(spec['original_parent_job_id']) + need(str(q['parent_job_id']) == job and expected['JobId'] == job, + 'Capture parent native identity differs') + accepted = bound(q['parent_running_admission']) + need(accepted['job_id'] == job and accepted['request_sha256'] == spec['parent_request']['sha256'], + 'Actual parent RUNNING admission missing') + native.native_identity(accepted['native'], expected, 'RUNNING') + observer = service(spec, 'original_observer') + query = observer.query(['/usr/local/bin/scontrol', 'show', 'job', job, '-o', '-dd']) + raw, provenance = observer.resolve_terminal(q, query) + steps = observer.query(observer.step_argv('paphiwetsa3')) + observer.validate_steps(steps, job) + service(spec, 'phase_observer_proof').check_owned_steps(steps['stdout'], job) + # Classification and finite numeric loss evidence belong to the actual capture producer. + return {'parent_native': raw, 'parent_terminal_query': query, + 'parent_terminal_provenance': provenance, 'parent_step_query': steps, + 'capture_is_not_recovery_authorization': True} + failure = admission.validate_failure(gate['failure'], spec['parent_request']) + result = {'failure_reason': failure['reason'], 'parent_native': bound(failure['evidence']['parent-native.txt'], False), + 'validated_failure_parent': spec['parent_request']} + if spec['mode'] in {'diagnostic', 'finalize', 'full'}: + bundle = q['inspection_producer'] if spec['mode'] == 'diagnostic' else gate['inspection'] + iq, ir = admission.completed_producer(bundle, 'native-elmo-checkpoint-inspection-result-v1', + 'INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION') + need(iq['parent_request'] == ir['parent_request'] == spec['parent_request'] + and ir['highest_healthy_known'] is True and ir['selected']['status'] == 'CPU_STRICT_INSPECTION_PASS', + 'Actual highest production checkpoint inspection missing') + need(ir['selected']['global_step'] == max(x['global_step'] for x in ir['census']['candidates']) + and all(x['status'] == 'CPU_STRICT_INSPECTION_PASS' for x in ir['results']), + 'Inspection skipped a newer or unhealthy retained checkpoint') + if spec['mode'] == 'diagnostic': + need(q['inspection'] == bundle['result'], 'Diagnostic references another inspection') + policy = spec['runtime']['native_client'] + prior = bound(policy['prior_container_observation']) + prerequisite = bound(policy['prior_allocated_prerequisites']) + need(policy['prior_allocated_prerequisites']['path'] == iq['allocated_prerequisites_path'] + and prerequisite['status'] == 'ACTUAL_PREREQUISITES_AND_NATIVE_ALLOCATION_PASS' + and prerequisite['job_id'] == prior['job_id'] == iq['producer_native']['JobId'] + and prerequisite['request_sha256'] == prior['request_sha256'] == bundle['request']['sha256'] + and prerequisite['container_auth'] == policy['prior_container_observation'] + and prerequisite['native_client_manifest_sha256'] == prior['native_client_manifest_sha256'] == policy['expected_manifest_sha256'] + and prior['parent_request'] == spec['parent_request'], + 'Prior capsule proof is not the completed inspection producer actual prerequisite chain') + native.native_identity(prior['native'], iq['producer_native'], 'RUNNING') + result['inspection_result_sha256'] = bundle['result']['sha256'] + if spec['mode'] in {'finalize', 'full'}: + dq, dr = admission.completed_producer(gate['diagnostic'], 'native-elmo-production-recovery-diagnostic-v1', + 'DIAGNOSTIC_COMPLETE_REQUIRES_HOST_AND_NATIVE_SUCCESS') + need(dq['parent_request'] == spec['parent_request'] and len(dr['rank_proofs']) == 8, + 'Completed actual all-rank diagnostic missing') + for binding in dr['rank_proofs']: + bound(binding) + result['diagnostic_result_sha256'] = gate['diagnostic']['result']['sha256'] + return result + + +def own_current(spec, job, request_sha, *, running=False): + """Validate the sole allowed held/current exception through immutable acceptance.""" + from execution_spec import read_binding, held_native, load_execution + need(str(job).isdecimal() and re.fullmatch('[0-9a-f]{64}', request_sha or ''), 'Exact current job/request required') + root = physical(spec['root']) + claim = root / ('scientific-claims' if spec['mode'] in {'diagnostic', 'full'} else 'producer-claims') / request_sha + prepared = read_json(claim / 'PREPARED.json') + args = (prepared['spec_path'], prepared['spec_sha256'], spec['mode'], spec['ops_manifest']['sha256']) + observed = load_execution(*args) if running else read_binding(*args) + need(observed[0] == spec and observed[3] == job and sha(observed[1]) == request_sha and observed[4] == claim, + 'Own claim is not the accepted exact spec/request') + raw = command(['/usr/local/bin/scontrol', 'show', 'job', job, '-o', '-dd']) + if running: + service(spec, 'native_interruption').native_identity(raw, dict(spec['native_contract'], JobId=job), 'RUNNING') + else: + held_native(raw, spec, prepared['spec_sha256'], job) + return raw, claim + + +def quiescence(spec, scope, run_id, current_job=None, current_request_sha=None, *, running=False, original_terminal=None): + """All prior claims require terminal proof; only this exact held/RUNNING job may remain.""" + native = service(spec, 'native_interruption') + observer = service(spec, 'original_observer') + need(scope['claim_roots'] and scope['scientific_job_names'] and scope['scientific_name_prefixes'] + and all(len(p) >= 16 for p in scope['scientific_name_prefixes']), 'Complete identity claim/queue scope required') + own = None + if current_job is not None: + _, own = own_current(spec, current_job, current_request_sha, running=running) + queue = command(['/usr/local/bin/squeue', '--noheader', '--user', scope['user'], '--format=%i|%j|%T|%k']) + step_query = observer.query(observer.step_argv(scope['user'])) + # A purged exact job query cannot certify release. Read the full owned inventory. + observer.query_record(step_query, observer.step_argv(scope['user'])) + need(step_query['returncode'] == 0 and not step_query['stderr'].strip(), 'Complete owned step query failed') + observer.strict_step_rows(step_query['stdout']) + if current_job is not None and not running: + service(spec, 'phase_observer_proof').check_owned_steps(step_query['stdout'], current_job) + for line in queue.splitlines(): + if not line.strip(): + continue + row = [x.strip() for x in line.split('|')] + need(len(row) == 4, 'Unparseable actual owner queue') + job, name, state, comment = row + scoped = name in scope['scientific_job_names'] or any(name.startswith(p) for p in scope['scientific_name_prefixes']) or comment == run_id + if scoped: + need(job == current_job and state == ('RUNNING' if running else 'PENDING'), + 'Other pending/active scientific writer: ' + job) + seen = set(); evidence = []; own_seen = False + for item in scope['claim_roots']: + root = physical(item['path']); need(root.is_dir(), 'Claim root unavailable') + for claim in sorted(root.iterdir()): + physical(claim) + need(claim.is_dir() and re.fullmatch('[0-9a-f]{64}', claim.name), 'Unknown claim entry') + p = read_json(claim / 'PREPARED.json'); job = (claim / 'PARENT_JOB_ID').read_text().strip() + need(p['request_sha256'] == claim.name and job.isdecimal() and job not in seen, + 'Unknown/duplicate accepted scientific claim') + seen.add(job) + if claim == own: + need(job == current_job and claim.name == current_request_sha, 'Own claim identity changed') + own_seen = True; continue + if job == str(spec['original_parent_job_id']): + need(claim.name == spec['parent_request']['sha256'] and original_terminal is not None, + 'Original parent needs separately validated original observer/failure evidence') + actual = read_json(claim / 'RUNNING_ADMISSION.json') + need(actual['job_id'] == job and actual['request_sha256'] == claim.name, + 'Original parent lacks actual RUNNING provenance') + f = fields(actual['native']) + expected = {k: f.get(k, '') for k in (*native.NATIVE_KEYS, 'NodeList', 'TresPerNode')} + native.native_identity(actual['native'], expected, 'RUNNING') + raw = original_terminal; state = fields(raw).get('JobState') + need(state in native.TERMINAL, 'Original parent is not terminal') + terminal_fields = native.native_identity(raw, expected, state) + need(terminal_fields.get('StartTime') == f.get('StartTime') + and terminal_fields.get('SubmitTime') == f.get('SubmitTime'), + 'Original parent terminal differs from actual RUNNING allocation') + import hashlib + proof = {'route': 'validated-original-parent', 'terminal_native_sha256': hashlib.sha256(raw.encode()).hexdigest()} + else: + from execution_spec import request_claim_terminal + captured, raw = request_claim_terminal(spec, claim, p, job, item['terminal_directory']) + proof = {'route': 'maintained-phase-observer', 'observer_spec': captured['observer_spec'], + 'terminal_native_sha256': captured['producer_native']['sha256']} + service(spec, 'phase_observer_proof').check_owned_steps(step_query['stdout'], job) + evidence.append(dict(claim=str(claim), job_id=job, **proof)) + if current_job is not None and spec['mode'] in {'diagnostic', 'full'}: + need(own_seen, 'Current scientific claim omitted from complete scope') + return {'queue': queue, 'step_query': step_query, 'prior_claims': evidence, + 'own_current_claim': str(own) if own else None} + + +def storage_and_node(spec, q, parent): + output = Path(q['output']); physical(output.parent) + need(not output.exists() and not output.is_symlink(), 'Output exists or is a symlink; preserve prior attempts') + if spec['mode'] == 'full': + need(not Path(q['local_logs']).exists(), 'Full local log namespace already exists') + reserve = parent['reserved_checkpoint_bytes'] + need(type(reserve) is int and reserve >= 80_000_000_000, 'Complete production retention reserve required') + free = shutil.disk_usage(output.parent).free + need(free > reserve, 'Insufficient actual destination free space') + node = command(['/usr/local/bin/scontrol', 'show', 'node', spec['native_contract']['NodeList'], '-o']) + f = fields(node) + need(f.get('NodeName') == spec['native_contract']['NodeList'] and f.get('State') in {'IDLE', 'ALLOCATED', 'MIXED'}, + 'Node has an unavailable/drain/failure flag') + need(int(f['CPUTot']) >= int(spec['native_contract']['NumCPUs']) + and int(f['RealMemory']) >= int(spec['native_contract']['MinMemoryNode'][:-1]) * 1024, + 'Node cannot supply requested resources') + return {'node_native': node, 'available_bytes': free, 'reserved_bytes': reserve, + 'write_fsync_certified': False} + + +def admit_before_submission(spec, template, current_job=None, current_request_sha=None, *, running=False): + from execution_spec import execution_claim_census + execution_claims = execution_claim_census(spec, current_job=current_job, + current_request_sha=current_request_sha, allow_held=current_job is not None and not running) + parent, gate, admission, intent = context(spec, template) + evidence = phase_evidence(spec, template, gate, admission) + live = quiescence(spec, gate['quiescence'], parent['run_id'], current_job, current_request_sha, running=running, original_terminal=evidence['parent_native']) + health = storage_and_node(spec, template, parent) + return {'schema': 'native-elmo-readonly-host-admission-v1', 'checked_at': now(), + 'mode': spec['mode'], 'parent_request': spec['parent_request'], + 'current_intent_sha256': intent, 'evidence': evidence, 'quiescence': live, + 'execution_claims': execution_claims, 'node_storage': health, 'checkpoint_content_checked': False} diff --git a/scripts/train/joint/native_elmo_ops/inside.sh b/scripts/train/joint/native_elmo_ops/inside.sh new file mode 100644 index 000000000..211b65ff5 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/inside.sh @@ -0,0 +1,37 @@ +#!/bin/bash +set -euo pipefail +root=${1:?Operations root} +mode=${2:?Mode} +spec=${3:?Execution specification} +spec_sha=${4:?Specification hash} +files_sha=${5:?Operations manifest hash} +[[ "${SLURM_JOB_ID:-}" =~ ^[0-9]+$ && "${SLURM_STEP_ID:-}" =~ ^[0-9]+$ && "${SLURM_RESTART_COUNT:-0}" == 0 ]] +export PYTHONDONTWRITEBYTECODE=1 +# Verify private client bytes and actual authenticated scontrol+squeue before +# touching the data cache or entering runtime_bootstrap (which imports Torch). +exports=$(python3 "$root/runtime_guard.py" inside "$spec" "$spec_sha" "$mode" "$files_sha") +eval "$exports" +if [[ "$mode" != capture ]]; then + runtime_exports=$(python3 "$root/lambda_runtime.py" "$ELMO_PARENT_REQUEST" --inside) + eval "$runtime_exports" + export PIXEL_SOURCE_ROOT="$ELMO_SOURCE" PIXEL_FA4_ROOT="$JOINT_FA4_ROOT" PIXEL_FA3_ROOT="$JOINT_FA3_ROOT" 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 NVIDIA_TF32_OVERRIDE=0 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" 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" +fi +if [[ "$mode" != diagnostic && "$mode" != full ]]; then + [[ -z "${SLURM_JOB_GPUS:-}" && -z "${SLURM_STEP_GPUS:-}" ]] + export CUDA_VISIBLE_DEVICES=-1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 +fi +if [[ "$mode" == full ]]; then + export PIXEL_WANDB_NETRC=/run/pixel-wandb.netrc +else + unset WANDB_API_KEY PIXEL_WANDB_NETRC + export WANDB_MODE=disabled +fi +exec python3 "$root/runtime_guard.py" dispatch "$spec" "$spec_sha" "$mode" "$files_sha" diff --git a/scripts/train/joint/native_elmo_ops/lambda_runtime.py b/scripts/train/joint/native_elmo_ops/lambda_runtime.py new file mode 100644 index 000000000..1ae806093 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/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/native_elmo_ops/lambda_stage_source.sh b/scripts/train/joint/native_elmo_ops/lambda_stage_source.sh new file mode 100644 index 000000000..e4bae4d43 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/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/native_elmo_ops/launcher_common.sh b/scripts/train/joint/native_elmo_ops/launcher_common.sh new file mode 100644 index 000000000..09fc5af6f --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/launcher_common.sh @@ -0,0 +1,69 @@ +#!/bin/bash +# Resume-only operational helper. Original/live launchers remain unchanged. +fail() { printf 'PIXEL STARTUP ERROR: %s\n' "$*" >&2; exit 2; } +need_command() { command -v "$1" >/dev/null 2>&1 || fail "Required command unavailable: $1"; } + +_mirror_health() { + local source=${1:?} code=${2:?} previous="$1/.mirror-health.json" + local last_success=0 failures=0 now temporary + if [[ -f "$previous" ]]; then + last_success=$(sed -n 's/.*"last_success_unix":\([0-9][0-9]*\).*/\1/p' "$previous") + failures=$(sed -n 's/.*"consecutive_failures":\([0-9][0-9]*\).*/\1/p' "$previous") + fi + [[ "$last_success" =~ ^[0-9]+$ ]] || last_success=0 + [[ "$failures" =~ ^[0-9]+$ ]] || failures=0 + now=$(date +%s) + if [[ "$code" == 0 ]]; then last_success=$now; failures=0; else failures=$((failures + 1)); fi + temporary=$(mktemp "$source/.mirror-health.tmp.XXXXXX") || return + if ! printf '{"schema":"pixel-log-mirror-health-v1","checked_at_unix":%s,"last_success_unix":%s,"consecutive_failures":%s,"last_exit_code":%s}\n' \ + "$now" "$last_success" "$failures" "$code" > "$temporary"; then + rm -f "$temporary" + return 1 + fi + mv -f "$temporary" "$previous" +} + +# Default is the final flush: at most three bounded attempts, requiring exit0. +# The periodic watcher passes attempts=1 and retries on its next interval. +mirror_once() ( + local source=${1:?} destination=${2:?} attempts=${3:-3} + local attempt code=1 + [[ -d "$source" ]] || { printf 'PIXEL LOG MIRROR ERROR: local log directory missing\n' >&2; return 2; } + [[ "$attempts" =~ ^[123]$ ]] || { printf 'PIXEL LOG MIRROR ERROR: attempts must be1..3\n' >&2; return 2; } + # The lock belongs to this subshell for the entire copy/retry/health update. + # It is outside the transfer and cannot be copied into its own destination. + exec 9> "$source/.mirror.lock" || return + flock --wait 330 9 || { printf 'PIXEL LOG MIRROR WARNING: lock wait expired\n' >&2; return 75; } + for ((attempt=1; attempt<=attempts; attempt++)); do + code=0 + # Bound mkdir as well as rsync: either can stall on the shared filesystem. + timeout --kill-after=15s 300s bash -c ' + mkdir -p "$2" && exec rsync -a --no-owner --no-group --delay-updates \ + --exclude=/wandb-cache/ --exclude=/.mirror.lock --exclude=/.mirror-health.tmp.\* \ + -- "$1/" "$2/" + ' pixel-log-mirror "$source" "$destination" || code=$? + if ! _mirror_health "$source" "$code"; then + printf 'PIXEL LOG MIRROR ERROR: could not write local health receipt\n' >&2 + return 74 + fi + [[ "$code" == 0 ]] && return 0 + printf 'PIXEL LOG MIRROR WARNING: attempt%s/%s exited%s; local files retained\n' "$attempt" "$attempts" "$code" >&2 + [[ "$attempt" == "$attempts" ]] || sleep 2 + done + return "$code" +) + +watch_log_mirror() { + local child=${1:?} source=${2:?} destination=${3:?} admission=${4:?} + # Keep the old fifth argument for dispatch compatibility. Transient periodic + # errors must not create a permanent failure marker or signal the trainer. + local unused_failure_path=${5:?} interval=${6:-60} + while kill -0 "$child" 2>/dev/null; do + sleep "$interval" + kill -0 "$child" 2>/dev/null || break + [[ -f "$admission" ]] || continue + if ! mirror_once "$source" "$destination" 1; then + printf 'PIXEL LOG MIRROR WARNING: periodic copy incomplete; retrying next interval, training continues\n' >&2 + fi + done +} diff --git a/scripts/train/joint/native_elmo_ops/native_client.py b/scripts/train/joint/native_elmo_ops/native_client.py new file mode 100644 index 000000000..fd8934f24 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/native_client.py @@ -0,0 +1,270 @@ +"""Stage an isolated native Slurm client closure inside an actual allocation. + +The private host loader/libc are used only by child Slurm clients. They never +replace the pinned container's loader, libc or Torch dependencies. +""" +import argparse +import hashlib +import json +import os +import posixpath +from pathlib import Path +import re +import shlex +import shutil +import stat +import subprocess + +from wrapper_common import bound, need, physical, sha, write_once + +LOGICAL = Path('/run/elmo-slurm-client') +INPUT = {'schema': 'native-elmo-client-closure-input-v1', + 'scontrol': '/usr/local/bin/scontrol', 'squeue': '/usr/local/bin/squeue', + 'config': '/etc/slurm/slurm.conf', 'munge_socket': '/run/munge/munge.socket.2'} + + +def ldd(path): + result = subprocess.run(['ldd', str(path)], capture_output=True, text=True, timeout=20) + loader = Path(path).name.startswith('ld-linux-x86-64.so.') + need((result.returncode == 0 or loader and result.returncode == 1) + and 'not found' not in result.stdout, + 'Could not close native ELF dependencies: ' + str(path)) + files = [] + for line in result.stdout.splitlines(): + if 'linux-vdso' in line or not line.strip() or loader and line.strip() == 'statically linked': + continue + paths = re.findall(r'(? reserved >= 80_000_000_000, 'Actual persistent storage reserve unavailable') + probe = output.parent / ('.elmo-recovery-fsync-' + job + '-' + request_sha[:16] + '.json') + value = dict(job_id=job, request_sha256=request_sha, output=str(output)) + write_once(probe, value) + need(read_json(probe) == value, 'Exclusive persistent write/fsync/readback failed') + return dict(probe=str(probe), probe_sha256=sha(probe), output=str(output), + available_bytes=free, reserved_bytes=reserved, actual_write_fsync_readback=True) + + +def dispatch_argv(spec, request_path): + mode = spec['mode']; need(mode in MODES, 'Unknown phase') + root = Path(spec['root']); service_boot = Path(spec['operational_service']['root']) / 'scripts/train/joint/service_bootstrap.py' + if mode == 'capture': + # Evidence-only exception approved by root: avoid Torch in 2CPU/1GiB capture. + return [sys.executable, str(service_boot), str(request_path), '--', '-m', + 'scripts.train.joint.native_interruption', str(request_path)] + return [os.environ['JOINT_BOOTSTRAP_PYTHON'], str(root / 'runtime_bootstrap.py'), + str(service_boot), str(request_path), '--', str(root / 'dispatch_after_runtime.py'), + str(request_path)] + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('stage', choices=('host', 'inside', 'dispatch')) + parser.add_argument('spec'); parser.add_argument('spec_sha'); parser.add_argument('mode', choices=MODES) + parser.add_argument('manifest_sha'); args = parser.parse_args(argv) + from execution_spec import load_execution, record_running_admission + # Authenticate capsule bytes before calling any container native wrapper. + if args.stage != 'host': + spec_early = bound({'path': args.spec, 'sha256': args.spec_sha}) + capsule(spec_early, os.environ['SLURM_JOB_ID']) + spec, path, q, job, claim = load_execution(args.spec, args.spec_sha, args.mode, args.manifest_sha) + parent = bound(spec['parent_request']); attempt = attempt_path(spec, job) + auth = native_auth(spec, path, q, job) + gate = admit_before_submission(spec, q, job, sha(path), running=True) + prerequisite_path(spec, q) + if args.stage == 'host': + attempt.mkdir() + running = record_running_admission(args.spec, args.spec_sha, args.mode, args.manifest_sha) + write_once(attempt / 'host-native.json', dict(auth, running_admission=running, live_admission=gate)) + elif args.stage == 'inside': + host = read_json(attempt / 'host-native.json') + need(host['request_sha256'] == sha(path) and host['job_id'] == job, 'Actual host authentication missing') + record = dict(auth, schema='native-elmo-container-auth-v1', status='CONTAINER_NATIVE_AUTHENTICATED', + parent_request=spec['parent_request'], original_parent_job_id=str(spec['original_parent_job_id']), + native_client_manifest_sha256=capsule(spec, job), host_native_sha256=sha(attempt / 'host-native.json')) + write_once(attempt / 'container-native.json', record) + if args.mode != 'capture': + actual = verify_scientific_prerequisites(spec, q, parent, job) + actual['storage'] = storage_probe(spec, q, parent, job, sha(path)) + write_once(attempt / 'pre-torch-prerequisites.json', dict(actual, request_sha256=sha(path), job_id=job)) + else: + need((attempt / 'container-native.json').is_file(), 'Container native authentication not performed') + if args.mode != 'capture': + need((attempt / 'pre-torch-prerequisites.json').is_file(), 'Pre-Torch source/cache/storage gate missing') + os.environ.update(ELMO_EXECUTION_SPEC=args.spec, ELMO_EXECUTION_SPEC_SHA256=args.spec_sha, + ELMO_OPS_MANIFEST_SHA256=args.manifest_sha, ELMO_EXECUTION_MODE=args.mode) + os.execv(dispatch_argv(spec, path)[0], dispatch_argv(spec, path)) + print(env_exports(spec, q, job, path)) + + +if __name__ == '__main__': + main() diff --git a/scripts/train/joint/native_elmo_ops/verify_joint_data_cache.py b/scripts/train/joint/native_elmo_ops/verify_joint_data_cache.py new file mode 100644 index 000000000..c5309ca7d --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/verify_joint_data_cache.py @@ -0,0 +1,377 @@ +"""Allocated-host admission of the immutable complete joint-training data cache. + +No ML imports. The producer hashes every artifact; this admission binds that +receipt and native completion, rehashes the manifest, and checks all file sizes. +The DataModule still performs its full content/provenance validation afterward. +""" + +import argparse +import hashlib +import json +import os +import stat +from pathlib import Path + + +def sha(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def require(condition, message): + if not condition: + raise ValueError(message) + + +def verify_cache(request, node, cache_parent=Path("/tmp/pixel-joint-complete-data")): + binding = request["data_cache"] + if binding.get("kind") == "prepared-track-v1": + return verify_track_cache(request, node) + if binding.get("kind") == "prepared-pixel-v1": + return verify_pixel_cache(request, node) + receipt_path, native_path = ( + Path(binding["receipt"]), + Path(binding["native_receipt"]), + ) + require( + sha(receipt_path) == binding["receipt_sha256"], "Cache receipt hash mismatch" + ) + require( + sha(native_path) == binding["native_receipt_sha256"], + "Native completion hash mismatch", + ) + receipt = json.loads(receipt_path.read_text()) + job = str(binding["job_id"]) + native = native_path.read_text().split() + fields = dict(part.split("=", 1) for part in native if "=" in part) + require( + fields.get("JobId") == job + and fields.get("JobName") == "joint-complete-data-v1" + and fields.get("JobState") == "COMPLETED" + and fields.get("ExitCode") == "0:0" + and fields.get("UserId", "").startswith("paphiwetsa3(") + and fields.get("NodeList") == node, + "Cache producer lacks exact successful native completion", + ) + cache = Path(binding["root"]) + require( + job.isdigit() + and cache == cache_parent / job + and cache.resolve() == cache + and cache.is_dir() + and node == binding["node"] == receipt["node"], + "Wrong physical cache location or allocation node", + ) + logical = str(Path(request["data_manifest"]).parent) + require( + receipt["schema"] == "joint-complete-data-cache-v1" + and receipt["status"] == "PASS" + and str(receipt["job_id"]) == job + and receipt["cache_root"] == str(cache) + and receipt["logical_data_root"] == logical + and receipt["manifest_sha256"] == request["data_manifest_sha256"] + and receipt["files"] == 1502 + and receipt["rgb_files"] == 500 + and receipt["splits"] == {"train": 450, "val": 50}, + "Complete joint cache contract mismatch", + ) + require( + sha(cache / "manifest.json") == request["data_manifest_sha256"], + "Cached manifest changed", + ) + manifest = json.loads((cache / "manifest.json").read_text()) + expected = {"manifest.json": (request["data_manifest_sha256"], False)} + + def add(name, digest, rgb=False): + path = Path(name) + require( + not path.is_absolute() and ".." not in path.parts, "Unbounded data artifact" + ) + require(name not in expected, "Duplicate data artifact") + expected[name] = (digest, rgb) + + add(manifest["normalization"]["path"], manifest["normalization"]["sha256"]) + episodes = manifest["episodes"] + require( + len(episodes) == 500 + and {k: sum(e["split"] == k for e in episodes) for k in ("train", "val")} + == receipt["splits"], + "Cached manifest split changed", + ) + for episode in episodes: + add(episode["path"], episode["sha256"]) + add(episode["rgb_path"], episode["rgb_sha256"], True) + add( + episode["action_provenance"]["path"], episode["action_provenance"]["sha256"] + ) + records = receipt["artifacts"] + require( + len(records) == 1502 and len(expected) == 1502, "Incomplete artifact inventory" + ) + seen, total, rgb_count = set(), 0, 0 + for record in records: + name = record["path"] + require( + name in expected and name not in seen, + "Unexpected or duplicate cache artifact", + ) + seen.add(name) + require( + (record["sha256"], record["rgb_hardlink"]) == expected[name], + "Artifact digest/role mismatch", + ) + target = cache / name + require( + target.resolve(strict=True) == target + and stat.S_ISREG(target.lstat().st_mode), + "Cache artifact must be a regular non-symlink file", + ) + require( + type(record["bytes"]) is int + and record["bytes"] > 0 + and target.stat().st_size == record["bytes"], + "Missing, partial or resized cache artifact", + ) + total += record["bytes"] + rgb_count += record["rgb_hardlink"] + require( + seen == set(expected) + and total == receipt["total_bytes"] > 0 + and rgb_count == 500, + "Cache inventory totals mismatch", + ) + return dict( + status="DATA_CACHE_ADMISSION_PASS", + files=1502, + rgb_files=500, + total_bytes=total, + producer_job=job, + node=node, + root=str(cache), + receipt_sha256=binding["receipt_sha256"], + native_receipt_sha256=binding["native_receipt_sha256"], + ) + + +def track_inventory(root, manifest_sha): + """Exact prepared XYZ contract, sharing the existing immutable cache owner.""" + root = Path(root) + manifest = json.loads((root / "manifest.json").read_text()) + identity = manifest.get("track_identity", {}) + require(manifest.get("schema") == "track-proprio-prepared-v1" + and identity == dict(grid_shape=[28, 40], reference_camera="agentview", + selection="regular_reference_view_grid", point_identity="persistent_surface_point", + frame="fixed_reference", units="meters", interpolated_endpoints=False) + and manifest["normalization"]["scope"] == "train_only", "Unbound track geometry/normalization") + episodes = manifest["episodes"] + require(len(episodes) == 500 and {k: sum(ep["split"] == k for ep in episodes) for k in ("train", "val")} + == dict(train=450, val=50), "Track split changed") + physical = [(ep["native_source"]["source_demo_sha256"], ep["native_source"]["episode"]) for ep in episodes] + require(len(set(physical)) == len(physical), "Duplicate physical track episodes") + expected = {"manifest.json": manifest_sha, + manifest["normalization"]["path"]: manifest["normalization"]["sha256"]} + for ep in episodes: + require(not any(key in ep for key in ("rgb_path", "rgb_sha256")), "RGB cache in XYZ corpus") + for name, digest in ((ep["path"], ep["sha256"]), + (ep["action_provenance"]["path"], ep["action_provenance"]["sha256"])): + require(name not in expected, "Duplicate track artifact") + expected[name] = digest + expected["provenance/source-manifest.json"] = manifest["source_manifest_sha256"] + locations = root / "provenance/native-locations.json" + canonical = json.dumps(json.loads(locations.read_text()), sort_keys=True, separators=(",", ":"), allow_nan=False) + require(hashlib.sha256(canonical.encode()).hexdigest() == manifest["native_locations_sha256"], + "Native track locations changed") + expected["provenance/native-locations.json"] = sha(locations) + require(len(expected) == 1004, "Incomplete prepared track inventory") + for name in expected: + path = Path(name) + require(not path.is_absolute() and ".." not in path.parts, "Unbounded track artifact") + return expected + + +def verify_track_producer(binding): + path = Path(binding["native_receipt"]) + require(sha(path) == binding["native_receipt_sha256"], "Track producer accounting changed") + fields = dict(p.split("=", 1) for p in path.read_text().split() if "=" in p) + require(fields.get("JobId") == str(binding["job_id"]) + and fields.get("JobName") == binding["producer_name"] + and fields.get("UserId") == f"{binding['user']}({binding['uid']})" + and fields.get("NodeList") == binding["producer_node"] + and fields.get("JobState") == "COMPLETED" and fields.get("ExitCode") == "0:0", + "Track producer lacks exact native completion") + + +def verify_track_cache(request, node): + binding = request["data_cache"] + verify_track_producer(binding) + require(sha(binding["receipt"]) == binding["receipt_sha256"], "Track staging receipt changed") + receipt = json.loads(Path(binding["receipt"]).read_text()) + root = Path(binding["root"]) + require(root.is_absolute() and root.resolve(strict=True) == root and root.stat().st_uid == binding["uid"] + and receipt["schema"] == "prepared-track-cache-v1" and receipt["status"] == "PASS" + and receipt["cache_root"] == str(root) + and receipt["manifest_sha256"] == request["data_manifest_sha256"] + and receipt["producer_native_sha256"] == binding["native_receipt_sha256"], + "Track cache placement/receipt changed") + require(sha(root / "manifest.json") == request["data_manifest_sha256"], "Track manifest changed") + expected = track_inventory(root, request["data_manifest_sha256"]) + seen, total = set(), 0 + for item in receipt["artifacts"]: + name = item["path"] + require(name in expected and name not in seen and item["sha256"] == expected[name], "Track inventory changed") + seen.add(name) + path = root / name + require(path.resolve(strict=True) == path and stat.S_ISREG(path.lstat().st_mode) + and path.stat().st_size == item["bytes"] > 0, "Missing/partial track artifact") + total += item["bytes"] + require(seen == set(expected) and receipt["total_bytes"] == total, "Track inventory incomplete") + return dict(status="DATA_CACHE_ADMISSION_PASS", spatial_stream="track", files=len(seen), + rgb_files=0, total_bytes=total, producer_job=str(binding["job_id"]), node=node, + root=str(root), receipt_sha256=binding["receipt_sha256"], + native_receipt_sha256=binding["native_receipt_sha256"]) + + +def stage_track_cache(request, source, node): + return stage_prepared_cache(request, source, node, spatial="track") + + +def pixel_inventory(root, manifest_sha, request=None): + """Deduplicate combined RGB/state HDF5s while rejecting conflicting roles.""" + root = Path(root) + require(sha(root / "manifest.json") == manifest_sha, "Pixel manifest changed") + manifest = json.loads((root / "manifest.json").read_text()) + require(manifest["schema"] == "pixel-proprio-prepared-v1", "Wrong pixel corpus schema") + expected = {"manifest.json": manifest_sha} + def add(name, digest): + path = Path(name) + require(not path.is_absolute() and ".." not in path.parts, "Unbounded pixel artifact") + require(name not in expected or expected[name] == digest, "Conflicting shared pixel artifact") + expected[name] = digest + add(manifest["normalization"]["path"], manifest["normalization"]["sha256"]) + for episode in manifest["episodes"]: + add(episode["path"], episode["sha256"]) + add(episode["rgb_path"], episode["rgb_sha256"]) + native_subtask = request and request.get("kind") == "native-subtask-world-model-v1" + provenance = episode.get("action_provenance") + if provenance: + add(provenance["path"], provenance["sha256"]) + else: + require(native_subtask and episode["embodiment"] not in request["interfaces"]["action_dims"], + "Missing command provenance for an action-bearing embodiment") + if native_subtask: + native = episode.get("native_source", {}) + if native.get("provenance_path"): + add(native["provenance_path"], native["provenance_sha256"]) + if request and request.get("kind") == "native-subtask-world-model-v1": + logical = Path(request["data_manifest"]).parent + for binding in request["subtask_contract"].values(): + relative = Path(binding["path"]).relative_to(logical) + add(str(relative), binding["sha256"]) + require(sha(root / relative) == binding["sha256"], "Changed native subtask sidecar") + return expected + + +def verify_pixel_cache(request, node): + binding = request["data_cache"] + verify_track_producer(binding) # same exact native terminal/owner/resource binding + require(sha(binding["receipt"]) == binding["receipt_sha256"], "Pixel staging receipt changed") + receipt = json.loads(Path(binding["receipt"]).read_text()) + root = Path(binding["root"]) + require(root.is_absolute() and root.resolve(strict=True) == root and root.stat().st_uid == binding["uid"] + and node == binding["node"] == receipt["staging_node"] + and receipt["schema"] == "prepared-pixel-cache-v1" and receipt["status"] == "PASS" + and str(receipt["staging_job_id"]) == str(binding["job_id"]) + and receipt["cache_root"] == str(root) + and receipt["manifest_sha256"] == request["data_manifest_sha256"], "Pixel cache placement changed") + expected = pixel_inventory(root, request["data_manifest_sha256"], request) + seen, total = set(), 0 + for item in receipt["artifacts"]: + name = item["path"] + require(name in expected and name not in seen and item["sha256"] == expected[name], "Pixel inventory changed") + seen.add(name) + path = root / name + require(path.resolve(strict=True) == path and stat.S_ISREG(path.lstat().st_mode) + and path.stat().st_size == item["bytes"] > 0, "Missing/partial pixel artifact") + total += item["bytes"] + require(seen == set(expected) and total == receipt["total_bytes"], "Incomplete pixel inventory") + return dict(status="DATA_CACHE_ADMISSION_PASS", files=len(seen), total_bytes=total, + producer_job=str(binding["job_id"]), node=node, root=str(root), + receipt_sha256=binding["receipt_sha256"], native_receipt_sha256=binding["native_receipt_sha256"]) + + +def stage_pixel_cache(request, source, node): + return stage_prepared_cache(request, source, node, spatial="pixel") + + +def stage_prepared_cache(request, source, node, *, spatial): + """One immutable copy/hash implementation for prepared track or pixel corpora.""" + import shutil + import tempfile + require(spatial in ("pixel", "track"), "Unknown prepared corpus") + binding = request["data_cache"] + require(binding.get("kind") == "prepared-" + spatial + "-v1" + and os.getuid() == binding["uid"], "Wrong prepared staging owner") + if spatial == "track": + verify_track_producer(binding) + else: + require(node == binding["node"], "Wrong pixel staging node") + source, destination = Path(source), Path(binding["root"]) + require(source.resolve(strict=True) == source and not destination.exists() + and destination.parent.resolve(strict=True) == destination.parent, "Preserve existing/partial corpus") + require(sha(source / "manifest.json") == request["data_manifest_sha256"], "Preparation manifest changed") + inventory = pixel_inventory if spatial == "pixel" else track_inventory + expected = (inventory(source, request["data_manifest_sha256"], request) if spatial == "pixel" + else inventory(source, request["data_manifest_sha256"])) + needed = sum((source / name).stat().st_size for name in expected) + require(shutil.disk_usage(destination.parent).free > needed + request["reserved_checkpoint_bytes"], + "Insufficient prepared data plus checkpoint reserve") + with tempfile.TemporaryFile(dir=destination.parent) as f: + f.write(b"prepared-cache-route-check"); f.flush(); os.fsync(f.fileno()); f.seek(0) + require(f.read() == b"prepared-cache-route-check", "Prepared write/read/fsync failed") + destination.mkdir() + records = [] + for name in sorted(expected, key=lambda name: name == "manifest.json"): + src, dst = source / name, destination / name + require(src.resolve(strict=True) == src and src.is_file(), "Unsafe source artifact") + dst.parent.mkdir(parents=True, exist_ok=True) + digest, count = hashlib.sha256(), 0 + with src.open("rb") as inp, dst.open("xb") as out: + for block in iter(lambda: inp.read(8 << 20), b""): + out.write(block); digest.update(block); count += len(block) + out.flush(); os.fsync(out.fileno()) + require(digest.hexdigest() == expected[name] and dst.stat().st_size == count > 0, + "Prepared transfer bytes/hash mismatch") + records.append(dict(path=name, sha256=digest.hexdigest(), bytes=count)) + result = dict(schema="prepared-" + spatial + "-cache-v1", status="PASS", staging_job_id=os.environ["SLURM_JOB_ID"], + staging_node=node, cache_root=str(destination), manifest_sha256=request["data_manifest_sha256"], + total_bytes=sum(r["bytes"] for r in records), artifacts=records) + if spatial == "track": + result["producer_native_sha256"] = binding["native_receipt_sha256"] + with Path(binding["receipt"]).open("x") as out: + json.dump(result, out, indent=2); out.write("\n"); out.flush(); os.fsync(out.fileno()) + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("request", type=Path) + parser.add_argument("proof", type=Path) + parser.add_argument("--stage-from", type=Path) + args = parser.parse_args() + require( + os.environ.get("SLURM_RESTART_COUNT", "0") == "0", "Fresh allocation required" + ) + request, node = json.loads(args.request.read_text()), os.environ["SLURMD_NODENAME"] + stage = stage_pixel_cache if request["data_cache"].get("kind") == "prepared-pixel-v1" else stage_track_cache + proof = (stage(request, args.stage_from, node) if args.stage_from else verify_cache(request, node)) + with args.proof.open("x") as stream: + json.dump(proof, stream, indent=2, sort_keys=True) + stream.write("\n") + print(json.dumps(proof), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/joint/native_elmo_ops/wrapper_common.py b/scripts/train/joint/native_elmo_ops/wrapper_common.py new file mode 100644 index 000000000..8f1bc59b3 --- /dev/null +++ b/scripts/train/joint/native_elmo_ops/wrapper_common.py @@ -0,0 +1,113 @@ +"""Small stdlib-only helpers; never imports scientific/model modules.""" +from datetime import datetime, timezone +import hashlib +import importlib +import importlib.util +import json +import os +from pathlib import Path +import re +import subprocess + +SERVICE_COMMIT = "abbb00499096faabf7035fde23bf4b1eba984377" # Verified upstream service publication. +SERVICE_SHA = "adf45524de3c2e6dba186572e7151d48bfca0b3adc8e6cf9defd686b47744c87" +SERVICE_FILE_COUNT = 35 +MODES = {'capture': 'native_interruption', 'inspect': 'checkpoint_inspection', + 'diagnostic': 'production_recovery', 'finalize': 'interruption_admission', 'full': 'run'} + +_LOADED_IMPLEMENTATION_SHA256 = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + + +def need(value, message): + if not value: + raise ValueError(message) + + +def physical(path): + p = Path(path) + need(p.is_absolute() and '..' not in p.parts and p.resolve(strict=True) == p, + 'Absolute physical existing path required: ' + str(p)) + return p + + +def sha(path): + h = hashlib.sha256() + with Path(path).open('rb') as stream: + for block in iter(lambda: stream.read(8 << 20), b''): + h.update(block) + return h.hexdigest() + + +def read_json(path): + def pairs(items): + need(len(items) == len(dict(items)), 'Duplicate JSON field') + return dict(items) + need(physical(path).stat().st_size <= 4 * 1024 * 1024, 'Operations JSON exceeds 4 MiB') + return json.loads(physical(path).read_text(), object_pairs_hook=pairs, + parse_constant=lambda value: (_ for _ in ()).throw(ValueError('Nonfinite operations JSON'))) + + +def bound(binding, parsed=True): + need(set(binding) == {'path', 'sha256'} and re.fullmatch('[0-9a-f]{64}', binding['sha256']), + 'Exact path/hash binding required') + path = physical(binding['path']) + need(path.is_file() and path.stat().st_size <= 4 * 1024 * 1024 and sha(path) == binding['sha256'], 'Changed bound evidence: ' + str(path)) + return read_json(path) if parsed else path.read_text() + + +def write_once(path, value): + path = Path(path) + physical(path.parent) + payload = value if isinstance(value, bytes) else (json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + '\n').encode() + with path.open('xb') as stream: + stream.write(payload); stream.flush(); os.fsync(stream.fileno()) + fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def command(argv): + p = subprocess.run(argv, capture_output=True, text=True, timeout=30) + need(len(p.stdout.encode()) <= 1048576 and len(p.stderr.encode()) <= 65536, 'Native query exceeded bounds') + need(p.returncode == 0 and not p.stderr.strip(), 'Actual bounded command failed: ' + str(argv[0])) + return p.stdout + + +def fields(raw): + pairs = re.findall(r'(?:^|\s)([A-Za-z][A-Za-z0-9:/_-]*)=(\S+)', raw) + need(len(pairs) == len(dict(pairs)), 'Ambiguous native record') + return dict(pairs) + + +def service(spec, name): + need(isinstance(SERVICE_COMMIT, str) and re.fullmatch('[0-9a-f]{40}', SERVICE_COMMIT), + 'Adapter service is unpublished; explicit reviewed publication binding required') + need(name in {'native_interruption', 'interruption_admission', 'checkpoint_inspection', 'original_observer', 'phase_observer_proof'}, + 'Unapproved pre-Torch service import') + binding = spec['operational_service'] + need(binding['manifest_sha256'] == SERVICE_SHA and sha(binding['manifest']) == SERVICE_SHA, + 'Exact published recovery service required') + root = physical(binding['root']) + need(Path(binding['manifest']) == root / 'FILES.sha256', 'Service manifest escaped root') + names = set() + lines = Path(binding['manifest']).read_text().splitlines() + need(0 < len(lines) <= 128, 'Unbounded service manifest') + for line in lines: + digest, name_in_manifest = line.split(' ', 1) + relative = Path(name_in_manifest) + need(not relative.is_absolute() and '..' not in relative.parts and name_in_manifest not in names, + 'Unsafe or duplicate operational service member') + names.add(name_in_manifest) + need(sha(physical(root / relative)) == digest, 'Operational service member changed') + need(len(names) == SERVICE_FILE_COUNT, 'Incomplete exact adapter service') + loader = importlib.util.spec_from_file_location('_elmo_wrapper_service', root / 'scripts/train/joint/service_bootstrap.py') + module = importlib.util.module_from_spec(loader); loader.loader.exec_module(module) + module.activate({'operational_service': binding}) + return importlib.import_module('scripts.train.joint.' + name) + + +def now(): + return datetime.now(timezone.utc).isoformat() diff --git a/scripts/train/joint/native_interruption.py b/scripts/train/joint/native_interruption.py new file mode 100644 index 000000000..0db1b4c94 --- /dev/null +++ b/scripts/train/joint/native_interruption.py @@ -0,0 +1,357 @@ +"""Read-only native interruption evidence for the maintained Elmo service. + +This module never cancels, submits, requeues, or imports ML libraries. A capture +is evidence, not restart authorization. Terminal history must be captured while +Slurm still has it: an absent record is deliberately not accepted as success. +""" +from datetime import datetime, timezone +import csv +import io +import json +import math +import os +from pathlib import Path +import re +import subprocess + +from scripts.train.joint.contract import bound, need, sha, write_once + + +TERMINAL = {"COMPLETED", "FAILED", "TIMEOUT", "NODE_FAIL", "PREEMPTED", "CANCELLED", "OUT_OF_MEMORY"} +INFRASTRUCTURE = {"TIMEOUT", "NODE_FAIL", "PREEMPTED"} +NATIVE_KEYS = ("JobId", "JobName", "UserId", "Account", "QOS", "Partition", "ReqNodeList", + "NumCPUs", "MinMemoryNode", "TimeLimit", "Requeue", "Restarts", "Command") + + +def fields(raw): + pairs = re.findall(r"(?:^|\s)([A-Za-z][A-Za-z0-9:/_-]*)=(\S+)", raw) + need(len(pairs) == len(dict(pairs)), "Duplicate or multiple native job records") + return dict(pairs) + + +def process_identity(expected, *, cpu=False): + owner = re.fullmatch(r"[^()]+\((\d+)\)", expected["UserId"]) + need(owner is not None and int(owner[1]) == os.getuid() + and expected["JobId"] == os.environ.get("SLURM_JOB_ID") + and expected["NodeList"] == os.environ.get("SLURMD_NODENAME") + and os.environ.get("SLURM_RESTART_COUNT", "0") == "0", + "Actual process UID/node/fresh allocation differs from the bound native producer") + if cpu: + need(not os.environ.get("SLURM_JOB_GPUS") and not os.environ.get("SLURM_STEP_GPUS"), + "CPU producer unexpectedly owns GPUs") + + +def destination_probe(output, job, required_bytes): + """Preserve a fresh exclusive fsync/readback probe beside the real output.""" + import shutil + + output = Path(output) + need(output.is_absolute() and output.parent.resolve(strict=True) == output.parent + and not output.exists() and type(required_bytes) is int and required_bytes > 0 + and str(job).isdecimal(), "Exact fresh production destination/storage reserve required") + free = shutil.disk_usage(output.parent).free + need(free > required_bytes, "Insufficient actual production destination free space") + path = output.parent / (".native-recovery-fsync-" + str(job) + ".probe") + payload = (str(output) + "\n" + str(job) + "\n").encode() + with path.open("xb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + need(path.read_bytes() == payload, "Actual production destination readback failed") + fd = os.open(output.parent, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + return dict(output=str(output), probe_path=str(path), probe_sha256=sha(path), + free_bytes=free, reserved_bytes=required_bytes, actual_write_fsync_readback=True) + + +def native_identity(raw, expected, state): + """Bind exact owner, node, resources, walltime and fresh-allocation identity.""" + f = fields(raw) + need(set(expected) == set(NATIVE_KEYS) | {"NodeList", "TresPerNode"}, + "Complete exact parent native resource binding required") + need(all(f.get(k, "") == str(v) for k, v in expected.items()), + "Actual native identity/resources differ from the bound parent") + need(expected["Requeue"] == expected["Restarts"] == "0" + and f.get("JobState") == state and f.get("NumNodes") == "1" + and f.get("NumTasks") == "1", "Wrong fresh single-node native allocation") + need(f.get("ReqTRES") and f.get("AllocTRES"), "Actual requested and allocated TRES missing") + for key in ("ReqTRES", "AllocTRES"): + parts = f[key].split(",") + need("cpu=" + expected["NumCPUs"] in parts and "node=1" in parts, + "Native CPU/node TRES differ") + mem = expected["MinMemoryNode"] + allowed = {"mem=" + mem} + if mem.endswith("G"): + allowed.add("mem=" + str(int(mem[:-1]) * 1024) + "M") + need(bool(set(parts) & allowed), "Native allocated memory differs") + gpu = [p for p in parts if "gpu" in p.lower()] + if expected["TresPerNode"] == "gres/gpu:h100:4": + need(set(gpu) <= {"gres/gpu=4", "gres/gpu:h100=4"}, + "Conflicting optional GPU TRES") + else: + need(expected["TresPerNode"] == "" and not gpu, + "CPU producer unexpectedly owns GPUs") + if expected["TresPerNode"] == "gres/gpu:h100:4": + need(f.get("JOB_GRES") == "gpu:h100:4", "Actual four-H100 JOB_GRES missing") + allocated = f.get("GRES", "") + if state == "RUNNING" or allocated: + match = re.fullmatch(r"gpu:h100:4\(IDX:([0-9,-]+)\)", allocated) + need(match is not None, "Actual four-H100 allocation indices missing") + indices = [] + for part in match[1].split(","): + endpoints = part.split("-") + need(len(endpoints) <= 2, "Invalid GPU index range") + first, last = int(endpoints[0]), int(endpoints[-1]) + need(0 <= first <= last < 8, "Invalid H100 index range") + indices.extend(range(first, last + 1)) + need(len(indices) == len(set(indices)) == 4, "GPU indices are not four distinct devices") + else: + need(f.get("JOB_GRES", "(null)") == "(null)" and not f.get("GRES"), + "CPU producer has nonempty native GPU allocation") + if state in TERMINAL: + end = f.get("EndTime", "") + need(end not in {"", "Unknown", "None", "N/A"} + and re.fullmatch(r"\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d", end) + and re.fullmatch(r"\d+:\d+", f.get("ExitCode", "")), + "Authoritative terminal EndTime/ExitCode missing") + return f + + +def classify_failure(terminal, worker_tail): + """Only mechanical infrastructure evidence admits automatic recovery. + + Cancellation, OOM and arbitrary FAILED remain blocked for diagnosis. A + future operational-stop extension must bind its actual guarded cancellation + and prior crash; it must not reinterpret a human cancellation here. + """ + state = terminal["JobState"] + scientific = re.compile(r"(?:nan loss|nonfinite|non-finite|FloatingPointError|" + r"(?:loss|mse)[^\n]{0,80}?(?:[:=(]\s*|\s+)[+-]?(?:nan|inf(?:inity)?)(?![a-z])|" + r"CUDA error: device-side assert|AssertionError|ValueError|" + r"out of memory|OutOfMemoryError)", re.I) + need(not scientific.search(worker_tail), + "Scientific, deterministic, or capacity failure requires diagnosis") + if state in INFRASTRUCTURE: + return state + if state == "FAILED": + patterns = (r"(?:ProcessGroupNCCL|NCCL)[^\n]*(?:timed out|timeout|watchdog)", + r"Watchdog caught collective operation timeout") + need(terminal["ExitCode"] != "0:0" and any(re.search(p, worker_tail, re.I) for p in patterns), + "FAILED without preserved infrastructure worker-crash evidence") + return "VERIFIED_NCCL_WORKER_CRASH" + raise ValueError("Terminal state is not automatically recoverable: " + state) + + +def check_steps(raw, job): + """squeue --steps is live: any row for the failed parent blocks release.""" + rows = [line.split("|", 1)[0].strip() for line in raw.splitlines() if line.strip()] + need(not any(row == str(job) or row.startswith(str(job) + ".") for row in rows), + "Failed parent still has active steps") + return True + + +def check_queue(raw, names, allowed_job=None): + """Complete exact owned queue, using explicit scientific names from all attempts.""" + need(names and len(names) == len(set(names)), "Complete unique scientific writer names required") + for line in raw.splitlines(): + if not line.strip(): + continue + columns = line.split("|") + need(len(columns) == 3, "Unparseable native job queue") + job, name, state = (v.strip() for v in columns) + if name in names: + need(job == allowed_job, "Other active/pending scientific writer: " + job) + need(state == "RUNNING", "Current admitted allocation is not RUNNING") + return True + + +def read_tail(path): + """Bound bytes and lines without loading unbounded worker logs.""" + path = Path(path) + need(path.is_file() and not path.is_symlink(), "Actual bounded worker log missing") + with path.open("rb") as stream: + stream.seek(max(0, path.stat().st_size - 65536)) + raw = stream.read(65536) + return b"\n".join(raw.splitlines()[-64:]).decode("utf8", errors="replace") + + +def parent_csv_paths(parent): + from scripts.train.joint.production_history import start_step + + relative = Path("csv") / f"step{start_step(parent):09d}" / "metrics.csv" + return {"local": str(Path(parent["local_logs"]) / relative), + "mirror": str(Path(parent["output"]) / "log-mirror" / relative)} + + +def newest_loss(snapshots): + """Select newest numeric-step loss row before any finiteness filtering.""" + candidates = [] + for name, snapshot in snapshots.items(): + columns = next(csv.reader([snapshot["header"]])) + need(len(columns) == len(set(columns)) and "step" in columns, "Ambiguous loss CSV header") + for position, row in enumerate(csv.DictReader(io.StringIO(snapshot["tail"]), fieldnames=columns)): + if list(row.values()) == columns: + continue + validation = {key: float(value) for key, value in row.items() if isinstance(key, str) + and key.startswith("val/") and value not in (None, "")} + need(all(math.isfinite(value) for value in validation.values()), + "Observed validation is nonfinite; scientific diagnosis required") + metrics = {key: value for key, value in row.items() if isinstance(key, str) + and key.lower().startswith("train/") and ("loss" in key.lower() or "mse" in key.lower()) + and value not in (None, "")} + if not metrics: + continue + need(None not in row and all(value is not None for value in row.values()), "Truncated loss row") + step = float(row["step"]) + need(math.isfinite(step) and step.is_integer() and step >= 0, "Invalid numeric loss step") + # NaN and Inf deliberately stay eligible here; newest is selected first. + numeric = {key: float(value) for key, value in metrics.items()} + candidates.append(dict(source=name, step=int(step), position=position, metrics=numeric)) + need(candidates, "No actual numeric loss row; scientific state unresolved (absence is not zero-update proof)") + newest = max(row["step"] for row in candidates) + selected = {} + for row in candidates: + if row["step"] == newest: + selected[row["source"]] = row + rows = list(selected.values()) + required = {"train/loss", "train/joint_image_mse", "train/joint_proprio_mse", + "train/joint_camera_mse", "train/source_0_loss_action"} + need(all(required <= {key.lower() for key in row["metrics"]} for row in rows), + "Newest loss row lacks total/image/proprio/camera/action-support metrics") + need(all(all(math.isfinite(value) for value in row["metrics"].values()) for row in rows), + "Newest numeric loss row is nonfinite; scientific diagnosis required") + need(all(row["metrics"] == rows[0]["metrics"] for row in rows), "Local/mirror newest loss rows disagree") + # Lightning CSV retains the previous value of other-embodiment metrics. + # Presence alone cannot identify the current batch: require exact agreement + # of all three current generic metrics with one embodiment's metrics. + metrics = {key.lower(): value for key, value in rows[0]["metrics"].items()} + active = [embodiment for embodiment in ("eva_bimanual", "human_bimanual") + if all(metrics.get(f"train/joint_{stream}_mse/{embodiment}") == metrics[f"train/joint_{stream}_mse"] + for stream in ("image", "proprio", "camera"))] + need(len(active) == 1, "Newest row active-embodiment/stream support is unknown or ambiguous") + if active[0] == "human_bimanual": + need(metrics["train/source_0_loss_action"] == 0, "Human batch unexpectedly has action loss") + else: + need("train/joint_action_mse" in metrics and "train/joint_action_mse/eva_bimanual" in metrics + and metrics["train/joint_action_mse"] == metrics["train/joint_action_mse/eva_bimanual"], + "Robot batch lacks the actual supervised action metric") + return dict(step=newest, metrics=rows[0]["metrics"], finite=True, selected_rows=rows, + active_embodiment=active[0], supervised_action=active[0] == "eva_bimanual", + selection="highest_numeric_step_then_last_loss_row_per_source_before_finiteness") + + +def loss_snapshot(parent): + snapshots, sources = {}, {} + for name, value in parent_csv_paths(parent).items(): + path = Path(value) + if not path.exists(): + sources[name] = dict(path=value, exists=False) + continue + need(path.is_file() and not path.is_symlink(), "Actual loss CSV must be a regular file") + before = path.stat() + with path.open("rb") as stream: + header = stream.readline(65536) + need(header.endswith(b"\n"), "Loss CSV header exceeds bound/is incomplete") + offset = max(0, before.st_size - 1048576) + stream.seek(offset) + raw = stream.read(1048576) + after = path.stat() + need((before.st_size, before.st_mtime_ns, before.st_ino) == + (after.st_size, after.st_mtime_ns, after.st_ino), "Loss CSV changed during bounded observation") + if offset: + raw = raw.partition(b"\n")[2] # never parse a partial leading row + lines = raw.splitlines() + need(len(lines) >= 64 or offset == 0, "Loss CSV tail exceeds bounded 1MiB window") + snapshots[name] = dict(header=header.decode("utf8"), tail=b"\n".join(lines[-64:]).decode("utf8")) + sources[name] = dict(path=value, exists=True, size_bytes=before.st_size, + mtime_ns=before.st_mtime_ns, inode=before.st_ino, device=before.st_dev) + return sources, snapshots + + +def command(argv): + result = subprocess.run(argv, capture_output=True, text=True, timeout=25, check=False) + need(result.returncode == 0, "Native query failed; history unavailable: " + result.stderr[:1000]) + return result.stdout + + +def capture(request_path): + """Capture actual parent evidence. Caller must use an owned CPU capture job.""" + q = json.loads(Path(request_path).read_text()) + need(q["schema"] == "native-elmo-interruption-capture-v1", "Wrong capture request") + parent = bound(q["parent_request"]) + need(parent["kind"] == "native-subtask-world-model-v1" and parent["phase"] == "full", + "Only exact native Elmo production failure may enter this path") + job, expected = q["parent_job_id"], q["parent_native"] + need(str(job).isdecimal() and expected["JobId"] == str(job) + and expected["JobName"] == parent["job_name"], "Parent job/name/request binding differs") + need(os.environ.get("SLURM_JOB_ID", "").isdecimal() + and not os.environ.get("SLURM_JOB_GPUS") and not os.environ.get("SLURM_STEP_GPUS") + and os.environ.get("SLURM_RESTART_COUNT", "0") == "0", "Fresh allocated CPU capture required") + process_identity(q["producer_native"], cpu=True) + producer = command(["/usr/local/bin/scontrol", "show", "job", os.environ["SLURM_JOB_ID"], "-o", "-dd"]) + native_identity(producer, q["producer_native"], "RUNNING") + need(q["producer_native"]["JobId"] == os.environ["SLURM_JOB_ID"] + and q["producer_native"]["NodeList"] == os.environ["SLURMD_NODENAME"], + "CPU producer is not this actual job/node") + accepted = bound(q["parent_running_admission"]) + need(accepted["job_id"] == str(job) + and accepted["request_sha256"] == q["parent_request"]["sha256"], + "Actual launch claim does not identify parent request") + native_identity(accepted["native"], expected, "RUNNING") + from scripts.train.joint.original_observer import query, resolve_terminal, step_argv, validate_steps + output = Path(q["output"]) + output.mkdir() # preserve a consumed observation even if provenance rejects + terminal_query = query(["/usr/local/bin/scontrol", "show", "job", str(job), "-o", "-dd"]) + write_once(output / "parent-native-query.json", terminal_query) + raw, terminal_provenance = resolve_terminal(q, terminal_query) + terminal = fields(raw) + need(terminal.get("JobState") in TERMINAL, "Parent is not terminal") + native_identity(raw, expected, terminal["JobState"]) + # A complete owned inventory remains queryable after a parent ID is purged. + # Missing controller history is never itself resource-release evidence. + steps_query = query(step_argv("paphiwetsa3")) + write_once(output / "parent-step-query.json", steps_query) + steps = validate_steps(steps_query, job) + tails = {key: read_tail(value) for key, value in q["worker_logs"].items()} + need(tails, "Preserved bounded stdout/stderr evidence required") + sources, snapshots = loss_snapshot(parent) + csv_evidence = {"loss-" + name + "-" + part + ".txt": text for name, snap in snapshots.items() + for part, text in snap.items()} + for name, text in {"parent-native.txt": raw, "parent-live-steps.txt": steps, + "parent-terminal-provenance.json": json.dumps(terminal_provenance, sort_keys=True), + "producer-running-native.txt": producer, **csv_evidence, + **{k + ".tail64.txt": v for k, v in tails.items()}}.items(): + need(Path(name).name == name, "Capture filename escaped evidence directory") + with (output / name).open("x") as stream: + stream.write(text) + stream.flush() + os.fsync(stream.fileno()) + try: + loss = newest_loss(snapshots) + reason = classify_failure(terminal, "\n".join(tails.values())) + except (ValueError, KeyError) as error: + write_once(output / "interruption-blocked.json", dict( + schema="native-elmo-interruption-blocked-v1", status="SCIENTIFIC_OR_INFRASTRUCTURE_DIAGNOSIS_REQUIRED", + request_sha256=sha(request_path), parent_request=q["parent_request"], parent_job_id=str(job), + loss_sources=sources, error=str(error), automatically_restart_authorized=False, + evidence={p.name: dict(path=str(p), sha256=sha(p)) for p in output.iterdir()})) + raise + write_once(output / "interruption-capture.json", dict( + schema="native-elmo-interruption-evidence-v1", status="CAPTURED_REQUIRES_NATIVE_PRODUCER_COMPLETION", + request_sha256=sha(request_path), parent_request=q["parent_request"], parent_job_id=str(job), + run_id=parent["run_id"], terminal_state=terminal["JobState"], terminal_exit=terminal["ExitCode"], + terminal_end=terminal["EndTime"], reason=reason, producer_job_id=os.environ["SLURM_JOB_ID"], + captured_at=datetime.now(timezone.utc).isoformat(), + parent_terminal_source=terminal_provenance["source"], + newest_numeric_loss=loss, loss_sources=sources, + evidence={p.name: dict(path=str(p), sha256=sha(p)) for p in output.iterdir()}, + parent_running_admission=q["parent_running_admission"], parent_native=expected, + automatically_restart_authorized=False)) + + +if __name__ == "__main__": + import sys + capture(sys.argv[1]) diff --git a/scripts/train/joint/native_recovery.py b/scripts/train/joint/native_recovery.py new file mode 100644 index 000000000..bf33b0502 --- /dev/null +++ b/scripts/train/joint/native_recovery.py @@ -0,0 +1,259 @@ +"""Operational proof helpers for the native conditional Elmo model. + +These reuse scientific checkpoint and sampler implementations. They neither +train nor authorize an interrupted production allocation by themselves. +""" +import hashlib +import json +import math +from contextlib import contextmanager + +import numpy as np +import torch + +from scripts.train.joint.contract import need, sha + + +def checkpoint_science(saved_config, current_config): + """Loss and sampling policy must match even when tensor shapes are unchanged.""" + from omegaconf import OmegaConf + + saved = OmegaConf.to_container(OmegaConf.create(saved_config), resolve=True) + current = OmegaConf.to_container(OmegaConf.create(current_config), resolve=True) + need(all(saved[key] == current[key] for key in ("model", "data")), + "Recovery changed checkpoint-bound model/loss/data/sampling semantics") + + +def state_digest(value): + """Stable digest of RNG/container state without drawing random numbers.""" + h = hashlib.sha256() + + def add(item): + h.update(type(item).__name__.encode() + b"\0") + if isinstance(item, torch.Tensor): + tensor = item.detach().cpu().contiguous() + h.update(str(tensor.dtype).encode()) + h.update(str(tuple(tensor.shape)).encode()) + h.update(tensor.numpy().tobytes()) + elif isinstance(item, np.ndarray): + h.update(str(item.dtype).encode()) + h.update(str(item.shape).encode()) + h.update(item.tobytes()) + elif isinstance(item, dict): + for key in sorted(item): + add(key) + add(item[key]) + elif isinstance(item, (list, tuple)): + for entry in item: + add(entry) + else: + h.update(json.dumps(item, allow_nan=False).encode()) + h.update(b"\xff") + + add(value) + return h.hexdigest() + + +def conditional_record(optimizer, named_parameters, counts, step): + """Bind counters to the real optimizer parameter order and moment shapes.""" + from egomimic.pl_utils.joint_smoke import check_conditional_adam_state + + named = list(named_parameters) + names = [name for name, _ in named] + parameters = [parameter for _, parameter in named] + ordered = [p for group in optimizer.param_groups for p in group["params"]] + need(len(ordered) == len(parameters) == len({id(p) for p in ordered}) + and all(a is b for a, b in zip(ordered, parameters)), + "Conditional optimizer parameter identity/order changed") + record = dict(schema="native-conditional-optimizer-usage-v1", global_step=step, + parameter_names=names, parameter_update_counts=dict(counts)) + state = optimizer.state_dict() + ids = [key for group in state["param_groups"] for key in group["params"]] + need(set(state["state"]) <= set(ids), "Unexpected extra optimizer state") + check_conditional_adam_state(state, record, step) + for parameter in parameters: + moments = optimizer.state.get(parameter) + if moments is not None: + need(all(moments[key].shape == parameter.shape + for key in ("exp_avg", "exp_avg_sq")), + "Conditional Adam moment shape changed") + return record + + +def restored_conditional_counts(trainer, module, checkpoint): + from egomimic.pl_utils.joint_smoke import ConditionalOptimizerUsage + + record = checkpoint.get("native_conditional_optimizer_usage") + need(isinstance(record, dict), "Missing native conditional checkpoint counts") + actual = conditional_record(trainer.optimizers[0], module.named_parameters(), + record["parameter_update_counts"], trainer.global_step) + need(actual == record, "Saved conditional parameter names/order changed") + callbacks = [c for c in trainer.callbacks if isinstance(c, ConditionalOptimizerUsage)] + need(len(callbacks) == 1 and callbacks[0].counts == record["parameter_update_counts"] + and callbacks[0].restored == record, + "Maintained conditional usage callback was not strictly restored") + return dict(record["parameter_update_counts"]), callbacks[0] + + +def install_rng_observer(module, saved, equal_state): + """Observe the existing restore at the first default training-step boundary. + + WorldModelTrainingBehavior restores RNG and acknowledges the sampler before + this call. The observer draws no randomness and never restores state itself. + """ + from egomimic.pl_utils.world_model_behavior import _capture_rng_state, _rank_rng_state + + expected = _rank_rng_state(saved) + need(expected.get("cuda_scope") == "current_device" + and len(expected.get("torch_cuda", [])) == 1, + "Recovery requires exactly one saved CUDA RNG state per rank") + previous = module._default_training_step + record = dict(observed=False, expected_sha256=state_digest(expected)) + + def observe(*args, **kwargs): + need(module.training_behavior._pending_rng is None, + "Scientific RNG restore did not execute before the training step") + actual = _capture_rng_state(current_cuda_only=True) + need(equal_state(actual, expected), "Actual first-update rank RNG differs from checkpoint") + record.update(observed=True, actual_sha256=state_digest(actual)) + module._default_training_step = previous + return previous(*args, **kwargs) + + module._default_training_step = observe + + def cleanup(): + module._default_training_step = previous + + return record, cleanup + + +def expected_metrics(network): + return {f"val/{emb}/{kind}/{('action_command_mse' if stream == 'action' else stream + '_normalized_mse')}" + for emb in network.pose_dims for kind in ("forecast", "completion") + for stream in network.active_streams(emb)} + + +@contextmanager +def preserved_rng(): + from egomimic.pl_utils.world_model_behavior import _capture_rng_state, _restore_rng_state + + state = _capture_rng_state(current_cuda_only=True) + try: + yield + finally: + _restore_rng_state(state) + + +def validate_hidden_sample(sample, output, streams): + need(set(output) == set(streams), "Unexpected native sampled streams") + need("language_features" in sample and "pixel_metadata" in sample, + "Missing sampled subtask/language provenance") + for name in streams: + source, value = sample["world"][name], output[name] + need(torch.isnan(source["value"][~source["observed"]]).all(), + "Withheld labels were not poisoned before recovery sampling") + need(value.shape == source["value"].shape and value.dtype == torch.float32 + and torch.isfinite(value).all(), "Nonfinite or malformed recovery sample") + need(torch.equal(value[source["observed"]], source["value"][source["observed"]]), + "Recovery sampling changed clean observations") + need(not value[~source["valid"]].count_nonzero(), "Recovery sampling emitted padded values") + + +def sample_both_embodiments(module, network, output): + """Use the configured maintained 20-step sampler on harvested held-out data.""" + sampler = module.model.pipeline.stage_by_id("joint_sampler").sampler + need(sampler.state_steps == 20 and sampler.action_steps == 20, + "Recovery hidden-label sample must retain the exact20-step schedule") + was_training = module.model.training + rows = {} + try: + module.model.eval() + module.model.device = next(module.parameters()).device + need(module.model.device.type == "cuda", "Actual native recovery sampling requires CUDA") + with preserved_rng(), torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + for emb in sorted(network.pose_dims): + path = output / f"sample-input-{emb}.pt" + need(path.is_file() and not path.is_symlink(), "Missing actual validation sample") + sample = torch.load(path, map_location="cpu", weights_only=False) + groups = module.model.process_batch_for_training({emb: sample}) + # Check poisoning before invoking the maintained sampler too. + for stream in groups[emb]["world"].values(): + need(torch.isnan(stream["value"][~stream["observed"]]).all(), + "Recovery sample contains withheld clean targets") + result = module.model.forward_eval(groups)[emb]["image_proprio_prediction"] + validate_hidden_sample(groups[emb], result, network.active_streams(emb)) + rows[emb] = dict(sample_input_sha256=sha(path), state_steps=20, + action_steps=20 if emb in network.action_dims else 0, + streams=list(result), finite=True, hidden_targets_nan=True) + finally: + module.model.train(was_training) + return rows + + +def summarize_native_phases(phases, *, start_step=50, production=False): + """Validate all-rank conditional evidence for the actual2+4 diagnostic.""" + need(type(start_step) is int and ((production is False and start_step == 50) + or (production is True and start_step > 0 and start_step % 20000 == 0)), + "Wrong explicit smoke/production diagnostic origin") + expected_streams = {"eva_bimanual": ["image", "proprio", "action", "camera"], + "human_bimanual": ["image", "proprio", "camera"]} + expected = {f"val/{emb}/{kind}/{('action_command_mse' if s == 'action' else s + '_normalized_mse')}" + for emb, streams in expected_streams.items() + for kind in ("forecast", "completion") for s in streams} + need(len(phases) == 2 and len(phases[0]) == len(phases[1]) in (2, 4), + "Native recovery requires both diagnostic phases and all ranks") + for phase, (start, stop) in zip(phases, ((start_step, start_step + 2), (start_step + 2, start_step + 6))): + first = phase[0] + for row in phase: + need(row.get("multimodal") is True + and row.get("private_bank_optimizer_steps_verified") is True + and row.get("human_action_absent") is True + and row.get("active_streams_by_embodiment") == expected_streams, + "Missing native conditional recovery proof") + record = row["conditional_optimizer_record"] + initial, final = row["initial_parameter_usage"], record["parameter_update_counts"] + names = record["parameter_names"] + need(record["schema"] == "native-conditional-optimizer-usage-v1" + and record["global_step"] == stop + and set(initial) == set(final) == set(names) and len(names) == len(set(names)) + and all(type(initial[n]) is type(final[n]) is int + and 0 <= initial[n] <= start and initial[n] <= final[n] <= stop + and final[n] - initial[n] <= stop - start for n in names), + "Invalid restored conditional counter progression") + need(initial == first["initial_parameter_usage"] + and record == first["conditional_optimizer_record"], + "Conditional counter restoration differs across ranks") + batches = row["embodiment_optimizer_batches"] + need(set(batches) == set(expected_streams) + and all(type(v) is int and v >= 0 for v in batches.values()) + and sum(batches.values()) == stop - start + and batches == first["embodiment_optimizer_batches"], + "Missing actual native embodiment updates") + rng = row["actual_rank_rng_restore"] + need(rng.get("observed") is True and isinstance(rng.get("actual_sha256"), str) + and len(rng["actual_sha256"]) == 64 + and rng["actual_sha256"] == rng["expected_sha256"], + "Actual rank RNG restore is not proven") + metrics = row["validation_metrics"] + need(set(metrics) == expected and all(math.isfinite(v) for v in metrics.values()) + and metrics == first["validation_metrics"], + "Missing finite native validation metrics on all ranks") + samples = row["hidden_label_samples"] + need(set(samples) == set(expected_streams), "Missing hidden-label native samples") + for emb, sample in samples.items(): + need(sample.get("state_steps") == 20 + and sample.get("action_steps") == (20 if emb == "eva_bimanual" else 0) + and sample.get("streams") == expected_streams[emb] + and sample.get("finite") is True and sample.get("hidden_targets_nan") is True + and isinstance(sample.get("sample_input_sha256"), str) + and len(sample["sample_input_sha256"]) == 64, + "Missing finite20-step hidden-label recovery sample") + for a, b in zip(phases[0], phases[1]): + need(a["conditional_optimizer_record"]["parameter_update_counts"] == b["initial_parameter_usage"], + "Second process did not restore conditional optimizer counts") + need(all(sum(phase[0]["embodiment_optimizer_batches"][emb] for phase in phases) > 0 + for emb in expected_streams), "Actual2+4 recovery did not exercise both embodiments") + return dict(native_multimodal=True, private_bank_optimizer_steps_verified=True, + actual_all_rank_rng_restore_observed=True, both_embodiments_optimized=True, + both_embodiments_hidden_nan_20step_sample=True, + diagnostic_scope=("production-checkpoint-only" if production else "scheduled-smoke-checkpoint-only")) diff --git a/scripts/train/joint/native_subtasks.py b/scripts/train/joint/native_subtasks.py new file mode 100644 index 000000000..16922a06c --- /dev/null +++ b/scripts/train/joint/native_subtasks.py @@ -0,0 +1,256 @@ +"""Opt-in native subtask/world-model admission; historical requests stay exact. + +The generic training service still owns execution. This module validates the +additional immutable data, conditional expert and hardware contracts only. +""" +import copy +import math +from collections import Counter +from pathlib import Path + +KIND = "native-subtask-world-model-v1" +STREAMS = ["image", "proprio", "action", "camera"] +POLICY = dict(implementation="native_subtask_joint", attention_backend="flash4", + token_execution="packed_valid", compile_mode="default", + compile_dynamic=True, compile_backend="inductor", + recompile_limit=64, accumulated_recompile_limit=256) + + +def selected(request): + kind = request.get("kind") + if kind is not None and kind != KIND: + raise ValueError("Unknown optional joint training request kind") + return kind == KIND + + +def _need(value, message): + if not value: + raise ValueError(message) + + +def verify_model_execution(request, cfg): + from scripts.train.joint.lr_policy import from_config, validate + _need(selected(request), "Native subtask request kind required") + policy = request.get("model_execution") + _need(policy == POLICY, "Production requires exact Flash4/packed/compiled policy") + stages, ids = cfg["model"]["pipeline"]["stages"], cfg["model"]["pipeline"]["stage_ids"] + net = stages[ids["joint_denoiser"]]["model"] + core = net["image_model"] + interfaces = request["interfaces"] + _need(net["_target_"] == "egomimic.models.joint_image_proprio.JointImageProprioModel" + and core["_target_"] == "egomimic.models.image_proprio.ImageProprioModel", + "Native multimodal constructor required") + _need(request["active_streams"] == net["active_streams"] == STREAMS + and net.get("separate_proprio_experts") is True, + "Private embodiment proprio banks and camera stream required") + for name in ("pose_dims", "action_dims", "camera_dims"): + _need(net[name] == interfaces[name] == cfg["joint"][name], "Changed native interfaces: " + name) + embodiments = set(interfaces["pose_dims"]) + _need(len(embodiments) >= 2 and set(interfaces["camera_dims"]) == embodiments + and set(interfaces["action_dims"]) == set(cfg["data"]["action_embodiments"]) + and set(interfaces["action_dims"]) < embodiments, + "Human/robot interfaces must not fabricate missing human commands") + _need(net["camera_max_count"] == core["max_cameras"] == interfaces["max_cameras"] + and net["language_feature_dim"] == interfaces["language_feature_dim"] > 0 + and core.get("require_observed_context") is False, + "Camera identity/language/per-frame noise contract changed") + for key in ("attention_backend", "token_execution", "compile_mode", "compile_dynamic", "compile_backend"): + _need(core[key] == policy[key], "Changed effective execution: " + key) + for key in ("recompile_limit", "accumulated_recompile_limit"): + _need(core["compile_" + key] == policy[key], "Changed bounded compilation budget") + from scripts.train.joint import boundary_uncertainty + boundary_uncertainty.request_config(request,cfg) + data = cfg["data"] + _need(data["training_conditioning"] == "independent_frame_noise" + and data["camera_pose"] is True and data["task"] == "visual_state_action" + and data["image_only"] is False and data["couple_state_observations"] is True, + "Subtask multimodal training/noising path changed") + _need(data["validation_anchors_per_episode"] is None + and data["validation_batches_per_embodiment"] == 50 + and data["validation_batch_size"] == 1 + and cfg["trainer"]["limit_val_batches"] == (100 // request["world_size"] if request["phase"] == "smoke" else 1.0), + "Balanced100 held-out validation must include50 robot and50 human batches") + _need(cfg["trainer"]["strategy"] == "ddp_find_unused_parameters_true" + and cfg["trainer"]["devices"] == request["world_size"] in (2, 4) + and data["distributed_world_size"] == request["world_size"] + and data["batch_size"] == request["per_rank_batch"] + and data["batch_size"] * request["world_size"] == request["global_batch"] == 32, + "Conditional DDP/global32 contract changed") + if request["arm"] == "discrete": + _need(data["mode"] == "discrete" and data["discrete_frames"] == 16 and data["stride"] == 4, + "Discrete production must use native16/stride4") + else: + _need(request["arm"] == data["mode"] == "continuous" and data["min_frames"] == 2 + and data["max_frames"] == 64 and data["sampling_strategy"] == "full_episode_uniform", + "Continuous production must use native N2..64") + for name, field in (("subtask_manifest", "subtask_manifest"), ("language_features", "language_features")): + binding = request["subtask_contract"][name] + _need(data[field + "_path"] == binding["path"] + and data[field + "_sha256"] == binding["sha256"], "Unbound " + name) + _need(from_config(cfg["model"]) == validate(request["lr_policy"]), "Unbound optimizer policy") + _need(cfg["evaluator"].get("active_streams_from_data") is True, + "Native evaluator must respect actual per-embodiment streams") + if request["phase"] == "full": + _need(any(v.get("_target_") == "egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage" + for v in cfg["callbacks"].values() if isinstance(v, dict)), + "Production conditional optimizer usage callback required") + if request["phase"] == "smoke": + callback = cfg["callbacks"]["joint_smoke"] + _need(callback["_target_"] == "egomimic.pl_utils.joint_smoke.JointSmokeProof" + and callback.get("multimodal") is True + and type(cfg["trainer"]["limit_val_batches"]) is int + and cfg["trainer"]["limit_val_batches"] == 100 // request["world_size"], + "Missing actual multimodal smoke instrumentation/100 global validation batches") + return "visual_state_action" + + +def verify_dataset(request): + from scripts.train.joint.contract import bound, sha + policy = request["dataset_contract"] + _need(policy.get("schema") == "native-subtask-dataset-v1", "Explicit native subtask dataset contract required") + manifest = bound(dict(path=request["data_manifest"], sha256=request["data_manifest_sha256"])) + _need(manifest["schema"] == "pixel-proprio-prepared-v1", "Wrong native pixel manifest") + norm = manifest["normalization"] + _need(norm["scope"] == "train_only" and norm["sha256"] == policy["normalization_sha256"] + and sha(Path(request["data_manifest"]).parent / norm["path"]) == norm["sha256"], + "Train-only normalization changed") + contract = request["subtask_contract"] + selection, index = bound(contract["selection"]), bound(contract["subtask_manifest"]) + _need(manifest.get("preparation", {}).get("selection_sha256") == contract["selection"]["sha256"], + "Prepared corpus must bind the exact selected original Scale episodes") + language = bound(contract["language_features"]) + _need(language.get("schema") == "pixel-task-language-features-v1" + and language.get("prepared_manifest_sha256") == request["data_manifest_sha256"] + and language.get("subtask_index_sha256") == contract["subtask_manifest"]["sha256"], + "Only final prepared/index-bound language features admit data") + expected = {(e["embodiment"], e["episode_id"]): e for e in selection["episodes"]} + physical = {(e["embodiment"], e["episode_id"]): e for e in manifest["episodes"]} + _need(len(expected) == len(selection["episodes"]) == policy["episodes"] + and len(physical) == len(manifest["episodes"]) and set(expected) == set(physical), + "Physical corpus coverage/identity changed") + counts = dict(Counter(e["split"] for e in manifest["episodes"])) + _need(counts == policy["splits"] and set(counts) == {"train", "val"} + and min(counts.values()) > 0, "Physical train/validation split changed") + _need(index["schema"] == "pixel-subtask-index-v1" + and index["prepared_manifest_sha256"] == request["data_manifest_sha256"] + and index["interval_semantics"] == "half_open" + and index["concurrent_policy"] == "combine_active_hands_atomic_intervals", + "Subtask clock/concurrency contract changed") + seen = set() + for record in index["episodes"]: + key = record["embodiment"], record["episode_id"] + _need(key in expected and key not in seen, "Unknown/duplicate subtask episode") + seen.add(key) + source, ep = expected[key], physical[key] + _need(record["split"] == ep["split"] == source["split"] + and record["frames"] == ep["frames"] == source["frames"], + "Subtask expansion changed physical split or native length") + previous = 0 + _need(bool(record["subtasks"]), "Selected episode has no admitted subtasks") + for unit in record["subtasks"]: + lo, hi = unit["start_idx"], unit["end_idx"] + _need(type(lo) is int and type(hi) is int and previous <= lo < hi <= ep["frames"], + "Subtask must be disjoint, half-open and in-bounds") + _need(any(a <= lo < hi <= b for a, b in ep.get("continuity_segments", [[0, ep["frames"]]])), + "Subtask crosses a native recording discontinuity") + previous = hi + _need(unit.get("quality_and_names_verified") is True + and unit.get("source_clock_alignment_verified") is True + and unit.get("instruction", "").strip(), "Unreviewed names/source clock") + evidence = unit["annotation_provenance"]["source"] + _need(evidence.get("annotation_source_kind") == "scale_existing" + and evidence.get("task_id") == source["annotation_task_id"] + and evidence.get("raw_sha256") == source["annotation_raw_sha256"] + and unit["annotation_provenance"]["events_sha256"] == index["annotation_events_sha256"] + and unit["annotation_provenance"].get("event_ids"), + "Only exact original selected Scale annotations are admitted; no generated labels") + _need(seen == set(expected), "Subtask index does not cover physical corpus") + tasks = {u["task_id"]: u["instruction"] for e in index["episodes"] for u in e["subtasks"]} + _need(set(language.get("tasks", {})) == set(tasks) + and all(language["tasks"][k]["instruction"] == text for k, text in tasks.items()), + "Language vectors must exactly cover final subtask instructions") + batches = policy["validation_batches"] + _need(type(batches) is int and batches == 100 + and policy["validation_batches_per_embodiment"] == 50, + "Need exactly50 robot plus50 human held-out global batches") + from scripts.train.joint import boundary_uncertainty + boundary_uncertainty.eligibility(request,manifest,index,bound) + return counts, batches + + +def _rank_capability(row): + capability = row.get("gpu_capability", row.get("compute_capability")) + _need(isinstance(capability, (list, tuple)) and len(capability) == 2 + and all(type(value) is int and value >= 0 for value in capability), + "Missing typed native GPU capability") + if "gpu_capability" in row and "compute_capability" in row: + _need(row["gpu_capability"] == row["compute_capability"], "Conflicting GPU capability aliases") + return capability + + +def verify_completion(request, smoke, completion, preflight, capacity): + from scripts.train.joint.contract import bound, canonical_uuid + from scripts.train.joint.lr_policy import update_lr + for key in ("kind", "interfaces", "subtask_contract", "dataset_contract", "boundary_sampling_contract"): + _need(request.get(key) == smoke.get(key), "Production differs from multimodal smoke: " + key) + _need(completion.get("schema") == "joint-multimodal-smoke-v1" + and completion.get("status") == "PASS" + and completion["job_id"] == str(request["smoke"]["job_id"]) + and completion["global_step"] == 50 and completion["validation_batches"] == 100 + and completion["world_size"] == request["world_size"], "Missing actual production topology smoke") + for key in ("strict_full_wrapper_reload", "optimizer_sampler_all_rank_rng_verified", + "lr_scheduler_strictly_restored", "all_rank_weights_agree", "sample_finite", + "private_bank_optimizer_steps_verified", "human_action_absent", + "language_conditioning_verified", "camera_targets_verified"): + _need(completion.get(key) is True, "Missing actual multimodal proof: " + key) + from scripts.train.joint import boundary_uncertainty + boundary_uncertainty.completion(request,smoke,completion,preflight) + _need(completion["active_streams"] == request["active_streams"] + and completion["expert_parameters"] == request["expert_parameters"], "Smoke expert identity changed") + ranks = completion["ranks"] + _need(len(ranks) == request["world_size"] + and len({canonical_uuid(r["gpu_uuid"]) for r in ranks}) == len(ranks), "Distinct owned GPU ranks required") + for rank, row in enumerate(ranks): + _need(row["rank"] == rank and row["world_size"] == len(ranks) + and row["global_step"] == row["bf16_training_outputs"] == row["fp32_losses"] == 50 + and row.get("fp32_gradients_verified") is True + and "H100" in row["gpu_name"] and _rank_capability(row)[0] >= 9, + "Actual H100/BF16/FP32 smoke rank proof missing") + _need(row["metrics"] and all(math.isfinite(float(v)) for v in row["metrics"].values()), + "Missing finite native held-out metrics") + _need(len(row["optimizer_update_lrs"]) == 50 + and all(math.isclose(lr, update_lr(request["lr_policy"], i), rel_tol=1e-12) + for i, lr in enumerate(row["optimizer_update_lrs"])), + "Actual optimizer warmup differs from bound policy") + _need(capacity.get("schema") == "native-multimodal-capacity-v1" + and capacity.get("status") == "PASS" + and capacity["config_sha256"] == request["smoke"]["resolved"]["sha256"] + and capacity["source_manifest_sha256"] == request["source_manifest_sha256"] + and capacity["data_manifest_sha256"] == request["data_manifest_sha256"] + and capacity["per_rank_batch"] == request["per_rank_batch"] + and capacity["world_size"] == request["world_size"] + and capacity["frames"] == (16 if request["arm"] == "discrete" else 64) and capacity["cameras"] == request["interfaces"]["max_cameras"] + and capacity["image_size"] == [160, 160] + and capacity["attention_backend"] == "flash4" and capacity["token_execution"] == "packed_valid" + and capacity.get("compiled_bf16_forward_backward") is True + and capacity["peak_memory_bytes"] < capacity["device_memory_bytes"], + "Actual maximum-shape production memory/backend proof required") + verify_dataset(request) + return completion, preflight + + +def compare_configs(production, smoke, request, legacy): + from scripts.train.joint import boundary_uncertainty + boundary_uncertainty.request_config(request,production) + boundary_uncertainty.request_config(request,smoke) + _need(production["trainer"].get("limit_val_batches", 1.0) == 1.0 + and type(smoke["trainer"]["limit_val_batches"]) is int + and smoke["trainer"]["limit_val_batches"] == 100 // request["world_size"], + "Only smoke diagnostic validation may be capped") + adjusted = copy.deepcopy(smoke) + adjusted["trainer"]["limit_val_batches"] = production["trainer"].get("limit_val_batches", 1.0) + if "limit_val_batches" not in production["trainer"]: + adjusted["trainer"].pop("limit_val_batches") + legacy(production, adjusted, request) + for field in ("world_model_execution", "pixel_video"): + _need(production.get(field) == smoke.get(field), "Scientific execution config changed: " + field) diff --git a/scripts/train/joint/observer_proof_README.md b/scripts/train/joint/observer_proof_README.md new file mode 100644 index 000000000..e78e1acb9 --- /dev/null +++ b/scripts/train/joint/observer_proof_README.md @@ -0,0 +1,100 @@ +# Native phase-observer proof admission + +`phase_observer_proof.py` is the maintained read-only semantic verifier for +native Elmo recovery phase observers. It parses hash-bound evidence with the +standard library. It never imports external operations code, issues a native +command, loads a model/checkpoint, or authorizes a recovery allocation. + +## API + +```python +verified = verify_observer( + observer_binding, + producer_request=request_binding, + producer_native=expected_native, + terminal_native=native_binding, +) +``` + +`observer_binding` is `{path, sha256}` for the canonical **final commit** +`TERMINAL_DIRECTORY/native-JOB.observer.json`, not an observer specification. +`producer_request` is the actual materialized request `{path, sha256}`, not its +parsed JSON. `producer_native` is the complete expected native contract dict +including actual JobId. `terminal_native` is the raw canonical `native-JOB.txt` +`{path, sha256}`. These three comparison arguments are optional for generic +census and required as appropriate by stricter consumers. + +`verified['result']` preserves the original final commit JSON. Other context +includes parsed request/spec, physical request/spec paths, actual job, accepted +held native record, observer acceptance/ID, terminal native fields, +`expected_native`, and `running_admission_status`. Physical paths are Path +objects; the original result remains JSON-compatible. + +`verify_claim_terminal(claim, terminal_directory, run_id)` preserves the exact +reviewed original-parent afterany route and requires a complete spec-bound phase +observer for all new claims. `check_owned_steps(raw, job)` strictly parses a +complete owner inventory and rejects parent, parent+'.', and parent+'_' rows. +Callers must obtain fresh owner queue/steps separately before admission; the +proof reader itself performs no query. + +## Evidence contract + +The active service root, complete manifest and imported guard origins must match +the producer's exact service binding. Ops manifests are read and hashed as data; +no external code is imported. The reader reopens ID-free execution and observer +specifications, templates, actual materialized requests, accepted held-native +records, PREPARED/acceptance/submission outcomes, exact observer afterany +acceptance and the complete final terminal/step commit. + +The producer spec's `release_observer_path` must match the bound observer spec. +The producer execution claim preserves `RELEASE_HELD_NATIVE.txt` and +`RELEASE_ATTEMPT.json`, whose fields are `job_id`, `spec_sha256`, +`request_sha256`, `held_native_sha256`, and `observer_readiness`. The actual +held native record and exact approved readiness query/identity are revalidated. +No operator-supplied PASS boolean replaces this chain. + +Producer and observer owner/name/node/resources/Command, accepted SubmitTime and +spec comments must agree. Real StartTime/EndTime must put producer termination +before observer start. Native query success, producer host exit, observer +terminal state and scientific completion are distinct. Every metadata file is +bounded to 4 MiB, manifests to 128 members, and query stdout/stderr to +1 MiB/64 KiB. Oversized checkpoint-like inputs reject without loading. + +The final commit is canonical, exclusive and produced only after fsync/readback +and successful full owned-step parsing. A raw or partial native file alone is +insufficient. Observer own terminal remains `UNKNOWN_NOT_REQUIRED` after purge. +Missing producer RUNNING provenance may remain `NOT_RECORDED` in generic failed +terminal census; it establishes neither zero work nor successful execution. + +## Maintained completion consumers + +`completed_producer` requires a complete observer binding, exact known +result-schema/status → request-schema/phase mapping, validated actual producer +RUNNING provenance, matching actual request/native bindings, canonical host exit +0, native COMPLETED0:0 and consistent result job IDs. Unknown schemas/statuses +have no compatibility bypass. + +The canonical result files under `request['output']` are: + +- capture: `interruption-capture.json`; +- inspect: `checkpoint-inspection.json`; +- diagnostic: `production-recovery-diagnostic.json`; +- finalize: `interruption-finalized.json`. + +Producer host exits are `ROOT/status/MODE-JOB.exit`. Full terminal census uses +the shared proof API separately and receives no phase-completion authority. +All historical new-claim terminal consumers must use the complete proof reader; +original-parent terminal captures retain their separately pinned provenance. + +Build service first, then bind its manifest in reviewed ops. Do not embed ops or +own service-manifest hashes inside this verifier. Dynamic specs/claims/results +stay outside immutable ops FILES, avoiding a hash cycle. A source publication is +not deployment or actual allocated recovery proof: controller/observer/runtime +assembly, current intent, source/cache/storage/native authentication, highest +checkpoint inspection, all-rank restore, finalizer and full preflight remain +separate gates. + +Portable CPU fixtures live together in `tests/joint_phase_observer/`. They use +fresh temporary service copies and synthetic profiles/claims/native records, +block native subprocesses and ML imports, and exercise the real shared reader. +Retained campaign-specific original-profile audits are separate local evidence. diff --git a/scripts/train/joint/original_observer.py b/scripts/train/joint/original_observer.py new file mode 100644 index 000000000..dfa7f9e78 --- /dev/null +++ b/scripts/train/joint/original_observer.py @@ -0,0 +1,220 @@ +"""Exact original-afterany terminal provenance; no submission or ML imports. + +Only the three already accepted original full allocations are supported. A +replacement needs a new reviewed profile; an arbitrary path cannot be adopted. +""" +from datetime import datetime, timezone +import json +from pathlib import Path +import re +import subprocess + +from scripts.train.joint.contract import bound, need + +ORIGINAL_OBSERVERS = {'mot-elmo-pickplace-continuous-large314m-boundary-robot-human-camera-subtask-s42-300k-20261026-v1': {'capture_script_sha256': 'be8ca57b2e5c09503ec312c9d84c91451d932d4d9a168c484f22439ccf521b57', + 'files_sha256': 'fbce43e2cf1ad9c090563bc7e5492c654f00db7e84890ef83d039e66cc9aea5f', + 'has_smoke_identity': True, + 'node': 'gpu03', + 'observer_id': '412801', + 'observer_name': 'elmo-native-full-large314m-boundary-v2', + 'parent_id': '412800', + 'parent_name': 'elmo-large314m-continuous-boundary-full-v2', + 'request_path': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/boundary-full-preparation-v2/ops/requests/elmo-large314m-continuous-boundary-full-v2.json', + 'request_sha256': '210e6f73c8a1a396716fab2f1e7e17f3b9d0663b099971eb05807181f9f00684', + 'root': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/boundary-full-preparation-v2', + 'submitter_sha256': '9e1722e56a117f8b42297e7fa3c82ef973a3435dd25bff51b69b09557fd68ddb'}, + 'mot-elmo-pickplace-continuous-large314m-fused-robot-human-camera-subtask-s42-300k-20261026-v1': {'capture_script_sha256': '87b011c6bd1f5d71c0758a75fb6ad9e05a807b7e5af33eba28264240b6071930', + 'files_sha256': '3a82e42346138a97a78fbca4201cc89b9635fe8987132620ceace54965386b58', + 'has_smoke_identity': True, + 'node': 'gpu11', + 'observer_id': '412799', + 'observer_name': 'elmo-native-full-large314m-fused-v2', + 'parent_id': '412798', + 'parent_name': 'elmo-large314m-continuous-fused-full-v2', + 'request_path': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/fused-full-preparation-v2/ops/requests/elmo-large314m-continuous-fused-full-v2.json', + 'request_sha256': 'e34190bfef155942fe216a1e5e81317cfa31b65d4c98ccb025a579c874fd3876', + 'root': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/fused-full-preparation-v2', + 'submitter_sha256': 'bede26ca7e2678c874e8442c5e5f6de01fdc58f38ee97c61d0f9f2c4c2a49dea'}, + 'mot-elmo-pickplace-continuous-large314m-plain-robot-human-camera-subtask-s42-300k-20261026-v1': {'capture_script_sha256': '5eeb85de75821627054787e0046972913442f2b2e7a751b681f2e71f81da21d4', + 'files_sha256': 'e2af24ae9daa20defd6a75235cb191101f4bc1abfe1bcf374b117659b2b1328d', + 'has_smoke_identity': False, + 'node': 'gpu11', + 'observer_id': '412790', + 'observer_name': 'elmo-native-full-large314m-v1', + 'parent_id': '412789', + 'parent_name': 'elmo-large314m-continuous-plain-full-v1', + 'request_path': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/baseline-full-preparation-v1/ops/requests/elmo-large314m-continuous-plain-full-v1.json', + 'request_sha256': '936e9acd30fa517f6f38311c0bcfd2b88544da8d476bd41125c4008de2845073', + 'root': '/workspace/users/paphiwetsa3/tasks/elmo-large-continuous-20261026-v1/baseline-full-preparation-v1', + 'submitter_sha256': '94c077facd7836fb55708199dab988fb5902db5348ce449daa861a888d95901f'}} + +PURGED_STDERR = "slurm_load_jobs error: Invalid job id specified" + + +def read_binding(binding, json_value=True): + path = Path(binding["path"]) + need(path.is_absolute() and path.is_file() and not path.is_symlink() + and path.resolve(strict=True) == path, "Observer proof must be a physical regular file") + return bound(binding, json_value) + + +def query(argv): + result = subprocess.run(argv, capture_output=True, text=True, timeout=25, check=False) + need(len(result.stdout.encode()) <= 1048576 and len(result.stderr.encode()) <= 65536, + "Native observation exceeded bound") + return dict(schema="native-elmo-terminal-query-v1", argv=argv, returncode=result.returncode, + stdout=result.stdout, stderr=result.stderr, observed_at_utc=datetime.now(timezone.utc).isoformat()) + + +def query_record(record, argv): + need(set(record) == {"schema", "argv", "returncode", "stdout", "stderr", "observed_at_utc"} + and record["schema"] == "native-elmo-terminal-query-v1" and record["argv"] == argv + and type(record["returncode"]) is int and isinstance(record["stdout"], str) + and isinstance(record["stderr"], str), "Wrong actual native query observation") + observed = datetime.fromisoformat(record["observed_at_utc"]) + need(observed.tzinfo is not None and observed <= datetime.now(timezone.utc), "Native query time missing/future") + + +def validate_original(q, reader=read_binding): + from scripts.train.joint.native_interruption import fields, native_identity, TERMINAL + + parent = reader(q["parent_request"]) + profile = ORIGINAL_OBSERVERS.get(parent.get("run_id")) + need(profile is not None, "No reviewed original observer for this identity") + root = Path(profile["root"]); pid, oid = profile["parent_id"], profile["observer_id"] + request_sha = profile["request_sha256"] + need(q["parent_request"] == dict(path=profile["request_path"], sha256=request_sha) + and str(q["parent_job_id"]) == pid and parent["job_name"] == profile["parent_name"], + "Observer does not belong to the exact original production request") + proof = q["original_observer"] + claim = root / "claims" / request_sha + expected = dict(files_manifest=root/"FILES.sha256", capture_script=root/"capture-native.sbatch", + guard=root/"guard_full.py", submitter=root/"submit_full_once.py", + parent_script=root/"full.sbatch", + prepared=claim/"PREPARED.json", parent_id=claim/"PARENT_JOB_ID", + observer_id=claim/"CAPTURE_JOB_ID", pending_native=claim/"CAPTURE_PENDING_NATIVE.txt", + running_admission=claim/"RUNNING_ADMISSION.json", terminal_native=root/("native-"+pid+".txt"), + query_exit=root/("native-"+pid+".exit"), query_stderr=root/("native-"+pid+".stderr")) + if profile["has_smoke_identity"]: + expected["smoke_identity"] = root / "smoke-identity.json" + need(set(proof) == set(expected) | {"schema"} and proof["schema"] == "native-elmo-original-observer-v1", + "Complete original observer provenance required") + for key, path in expected.items(): + value = proof[key] + need(set(value) == {"path", "sha256"} and value["path"] == str(path) + and re.fullmatch(r"[0-9a-f]{64}", value["sha256"]), "Unreviewed original observer path/hash: "+key) + need(proof["files_manifest"]["sha256"] == profile["files_sha256"] + and proof["capture_script"]["sha256"] == profile["capture_script_sha256"] + and proof["submitter"]["sha256"] == profile["submitter_sha256"] + and proof["running_admission"] == q["parent_running_admission"], "Original immutable code/claim pins differ") + raw = {key: reader(value, False) for key, value in proof.items() if key != "schema"} + manifest = {} + for line in raw["files_manifest"].splitlines(): + m = re.fullmatch(r"([0-9a-f]{64}) (.+)", line) + need(m and m[2] not in manifest and not Path(m[2]).is_absolute() and ".." not in Path(m[2]).parts, + "Unsafe/ambiguous original FILES manifest") + manifest[m[2]] = m[1] + for key in ("capture_script", "guard", "submitter", "parent_script", *(["smoke_identity"] if profile["has_smoke_identity"] else [])): + need(manifest.get(str(expected[key].relative_to(root))) == proof[key]["sha256"], + "Observer dependency not bound by original manifest: "+key) + need(manifest.get(str(Path(profile["request_path"]).relative_to(root))) == request_sha, + "Original scientific request absent from immutable manifest") + prepared = json.loads(raw["prepared"]) + need(prepared["request_sha256"] == request_sha and prepared["manifest_sha256"] == profile["files_sha256"] + and raw["parent_id"].strip() == pid and raw["observer_id"].strip() == oid, + "Original accepted claim does not bind this parent/observer") + running = json.loads(raw["running_admission"]) + need(running["status"] == "ACTUAL_BOUND_FULL_NATIVE_ALLOCATION_PASS" + and running["job_id"] == pid and running["request_sha256"] == request_sha, + "Actual original RUNNING claim missing") + native = q["parent_native"] + need(native["JobId"] == pid and native["JobName"] == profile["parent_name"] + and native["UserId"] == "paphiwetsa3(41003)" and native["Account"] == "loaner" + and native["QOS"] == "loaner-normal" and native["Partition"] == "rl2" + and native["NodeList"] == native["ReqNodeList"] == profile["node"] + and native["NumCPUs"] == "32" and native["MinMemoryNode"] in {"192G", "196608M", "196608"} + and native["TimeLimit"] == "7-00:00:00" and native["Requeue"] == native["Restarts"] == "0" + and native["Command"] == str(root/"full.sbatch") and native["TresPerNode"] == "gres/gpu:h100:4", + "Original full native owner/resources/Command differ") + actual_running = native_identity(running["native"], native, "RUNNING") + pending = fields(raw["pending_native"]) + required = dict(JobId=oid, JobName=profile["observer_name"], UserId=native["UserId"], Account="loaner", + QOS="loaner-normal", Partition="rl2", ReqNodeList=profile["node"], NumTasks="1", + NumCPUs="2", TimeLimit="00:10:00", Requeue="0", Restarts="0", JobState="PENDING", + Command=str(root/"capture-native.sbatch"), WorkDir=str(root)) + need(all(pending.get(k) == v for k, v in required.items()) and pending.get("NumNodes") in {"1", "1-1"} + and pending.get("MinMemoryNode") in {"1G", "1024M", "1024"} + and re.fullmatch(r"afterany:"+pid+r"\((?:unfulfilled|unsatisfied)\)", pending.get("Dependency", "")), + "Original observer accepted dependency/native identity differs") + need(pending.get("StdOut") == str(root/"logs"/(profile["observer_name"]+"-"+oid+".out")) + and pending.get("StdErr") == str(root/"logs"/(profile["observer_name"]+"-"+oid+".err")), + "Original observer output namespace differs") + tres = pending.get("ReqTRES", "").split(",") + need("cpu=2" in tres and "node=1" in tres and any(x in tres for x in ("mem=1G", "mem=1024M")), + "Original observer CPU TRES differ") + for key, value in pending.items(): + if key in {"ReqTRES", "AllocTRES"}: + need(not any("gpu" in v.lower() for v in value.split(",")), "Observer has GPU allocation") + if key.lower().startswith(("gres", "job_gres", "tresper")): + need(value.lower() in {"", "(null)", "none", "n/a", "0"} or key == "TresPerTask" and value == "cpu:2", + "Observer has GPU/GRES resources") + need(raw["query_exit"].strip() == "0" and not raw["query_stderr"].strip(), + "Original observer parent-query did not succeed cleanly") + terminal_fields = fields(raw["terminal_native"]) + need(terminal_fields.get("JobState") in TERMINAL, "Original observer did not capture terminal parent") + terminal = native_identity(raw["terminal_native"], native, terminal_fields["JobState"]) + start = datetime.fromisoformat(terminal["StartTime"]); end = datetime.fromisoformat(terminal["EndTime"]) + need(start <= end and terminal["StartTime"] == actual_running["StartTime"], "Terminal interval/original allocation differs") + return raw["terminal_native"], dict(schema="native-elmo-original-observer-validation-v1", original_parent_id=pid, + original_observer_id=oid, observer_own_terminal="UNKNOWN_NOT_REQUIRED", capture_query_exit=0, + scientific_parent_state=terminal["JobState"], scientific_parent_exit=terminal["ExitCode"], proofs=proof) + + +def resolve_terminal(q, record, reader=read_binding): + from scripts.train.joint.native_interruption import fields, native_identity, TERMINAL + + argv = ["/usr/local/bin/scontrol", "show", "job", str(q["parent_job_id"]), "-o", "-dd"] + query_record(record, argv) + observer = validate_original(q, reader) if q.get("original_observer") is not None else None + if record["returncode"] == 0: + need(not record["stderr"].strip(), "Live parent query returned diagnostics") + raw = record["stdout"]; current = fields(raw) + need(current.get("JobState") in TERMINAL, "Live parent is not terminal; observer cannot override it") + native_identity(raw, q["parent_native"], current["JobState"]) + if observer: + old = fields(observer[0]) + need(all(current.get(k) == old.get(k) for k in ("JobState", "ExitCode", "StartTime", "EndTime")), + "Live and captured parent terminal records contradict") + return raw, dict(source="LIVE_NATIVE_QUERY", original_observer=None if observer is None else observer[1]) + need(record["returncode"] == 1 and not record["stdout"].strip() and record["stderr"].strip() == PURGED_STDERR, + "Native query failed for a reason other than the exact reviewed purged-ID response") + need(observer is not None, "Purged parent requires genuine original observer provenance") + return observer[0], dict(source="ORIGINAL_AFTERANY_CAPTURE_LIVE_ID_PURGED", original_observer=observer[1]) + + +def step_argv(user): + need(user == "paphiwetsa3", "Exact original owner step scope required") + return ["/usr/local/bin/squeue", "--steps", "--noheader", "--user", user, "--format=%i|%T"] + + +def strict_step_rows(raw): + """Parse complete owner inventory; malformed success is not release proof.""" + for line in raw.splitlines(): + if not line.strip(): + continue + row = [v.strip() for v in line.split("|")] + need(len(row) == 2 and re.fullmatch(r"[0-9]+(?:_[0-9]+)?(?:\.[A-Za-z0-9_-]+)?", row[0]) + and row[1] in {"PENDING", "RUNNING", "COMPLETING", "CONFIGURING", "SUSPENDED", "STOPPED", "RESIZING", + "COMPLETED", "FAILED", "CANCELLED", "TIMEOUT", "NODE_FAIL", "OUT_OF_MEMORY", "PREEMPTED"}, + "Ambiguous complete owned step inventory") + return raw + + +def validate_steps(record, parent): + from scripts.train.joint.native_interruption import check_steps + + query_record(record, step_argv("paphiwetsa3")) + need(record["returncode"] == 0 and not record["stderr"].strip(), "Complete owned step query failed") + strict_step_rows(record["stdout"]) + check_steps(record["stdout"], parent) + return record["stdout"] diff --git a/scripts/train/joint/phase_observer_proof.py b/scripts/train/joint/phase_observer_proof.py new file mode 100644 index 000000000..f9a3d1f28 --- /dev/null +++ b/scripts/train/joint/phase_observer_proof.py @@ -0,0 +1,423 @@ +"""Service-owned read-only semantic closure for new native phase observers. + +No ops module is imported or executed, no native query is issued, and no model +or checkpoint is loaded. Ops code and dynamic records are data with hashes. +""" +import copy +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +import re + +from scripts.train.joint.native_interruption import fields, native_identity, NATIVE_KEYS, TERMINAL +from scripts.train.joint.original_observer import ORIGINAL_OBSERVERS, strict_step_rows + +MODES = { + 'capture': ('native-elmo-interruption-capture-v1','2','1G','00:10:00',''), + 'inspect': ('native-elmo-checkpoint-inspection-v1','4','24G','00:20:00',''), + 'diagnostic': ('native-elmo-production-recovery-request-v1','32','192G','01:00:00','gres/gpu:h100:4'), + 'finalize': ('native-elmo-interruption-finalizer-v1','4','24G','00:20:00',''), + 'full': ('joint-training-v2','32','192G','7-00:00:00','gres/gpu:h100:4'), +} +RESULT_MODES = { + 'native-elmo-interruption-evidence-v1':'capture', + 'native-elmo-checkpoint-inspection-result-v1':'inspect', + 'native-elmo-production-recovery-diagnostic-v1':'diagnostic', + 'native-elmo-interruption-finalized-v1':'finalize', +} +OBSERVER_KEYS={'schema','root','producer_execution','terminal_directory'} +MAX_BYTES=4*1024*1024 +MAX_MEMBERS=128 + + +def need(value,message): + if not value: raise ValueError(message) + + +def physical(path): + p=Path(path) + need(p.is_absolute() and '..' not in p.parts and p.resolve(strict=True)==p,'Physical existing proof path required') + return p + + +def data(path): + p=physical(path);need(p.is_file() and not p.is_symlink() and p.stat().st_size<=MAX_BYTES,'Bounded regular proof metadata required') + with p.open('rb') as stream:payload=stream.read(MAX_BYTES+1) + need(len(payload)<=MAX_BYTES,'Proof input exceeds metadata bound') + return payload + + +def text(path):return data(path).decode('utf-8') + + +def sha(path):return hashlib.sha256(data(path)).hexdigest() + + +def read_json(path): + def pairs(items): + need(len(items)==len(dict(items)),'Duplicate proof JSON keys') + return dict(items) + return json.loads(text(path),object_pairs_hook=pairs, + parse_constant=lambda x: (_ for _ in ()).throw(ValueError('Nonfinite proof JSON'))) + + +def bound(binding): + need(set(binding)=={'path','sha256'} and re.fullmatch('[0-9a-f]{64}',binding['sha256']),'Exact proof path/hash binding required') + path=physical(binding['path']);need(sha(path)==binding['sha256'],'Proof bytes changed') + return path + + +def binding(path):return {'path':str(path),'sha256':sha(path)} + + +def encoded(value):return (json.dumps(value,sort_keys=True,indent=2,allow_nan=False)+'\n').encode() + + +def manifest(root,proof): + root=physical(root);path=bound(proof);need(path==root/'FILES.sha256','Manifest escaped exact root') + names=set() + for line in text(path).splitlines(): + match=re.fullmatch(r'([0-9a-f]{64}) (.+)',line) + need(match is not None,'Malformed proof manifest') + digest,name=match.groups();relative=Path(name) + need(name not in names and not relative.is_absolute() and '..' not in relative.parts,'Unsafe/duplicate manifest member') + need(sha(root/relative)==digest,'Manifest member changed') + names.add(name) + need(len(names)<=MAX_MEMBERS,'Proof manifest exceeds metadata member bound') + need(names,'Empty proof manifest');return names + + +def service(binding_value): + root=physical(binding_value['root']) + need(root==Path(__file__).resolve().parents[3] and Path(binding_value['manifest'])==root/'FILES.sha256', + 'Producer service is not the active proof implementation') + names=manifest(root,{'path':binding_value['manifest'],'sha256':binding_value['manifest_sha256']}) + from scripts.train.joint import native_interruption, original_observer, service_bootstrap + for module in (native_interruption,original_observer,service_bootstrap): + need(Path(module.__file__).resolve()==root/'scripts/train/joint'/(module.__name__.rsplit('.',1)[1]+'.py'), + 'Proof dependency imported outside active service') + need(names==set(service_bootstrap.service_files(root)) and 'scripts/train/joint/phase_observer_proof.py' in names, + 'Active service export closure differs') + return root + + +def native_argv(job):return ['/usr/local/bin/scontrol','show','job',str(job),'-o','-dd'] + + +def successful(record,argv): + need(set(record)=={'schema','argv','returncode','stdout','stderr','observed_at_utc'} + and record['schema']=='native-elmo-phase-observer-query-v1' and record['argv']==argv + and type(record['returncode']) is int and record['returncode']==0 + and isinstance(record['stdout'],str) and isinstance(record['stderr'],str) and not record['stderr'].strip() + and len(record['stdout'].encode())<=1048576 and len(record['stderr'].encode())<=65536, + 'Actual clean successful observer query required') + stamp=datetime.fromisoformat(record['observed_at_utc']) + need(stamp.tzinfo is not None and stamp<=datetime.now(timezone.utc),'Missing/future proof query time') + return record['stdout'] + + +def held_native(raw,spec,spec_sha,job): + f=fields(raw);n=spec['native_contract'] + need(f.get('JobId')==job and f.get('JobState')=='PENDING' and f.get('Reason')=='JobHeldUser' + and all(f.get(k,'')==v for k,v in n.items() if k!='NodeList') + and f.get('NodeList','') in {'','(null)',n['NodeList']} and f.get('NumNodes') in {'1','1-1'} + and f.get('NumTasks')=='1' and f.get('Comment')=='elmo-execution:'+spec_sha + and f.get('SubmitTime') not in {None,'Unknown','None'},'Wrong actual held execution identity') + need(f.get('AllocTRES','') in {'','(null)'} and f.get('GRES','') in {'','(null)'} + and f.get('StartTime') in {None,'Unknown','N/A'},'Held execution already allocated/started') + items=[p.split('=',1) for p in f.get('ReqTRES','').split(',')] + need(all(len(v)==2 for v in items) and len(items)==len(dict(items)),'Ambiguous held TRES') + tres=dict(items);memory=n['MinMemoryNode'] + need(tres.get('cpu')==n['NumCPUs'] and tres.get('node')=='1' + and tres.get('mem') in {memory,str(int(memory[:-1])*1024)+'M'},'Held CPU/memory/node TRES differ') + gpu={k+'='+v for k,v in tres.items() if 'gpu' in k.lower()} + need(gpu<=({'gres/gpu=4','gres/gpu:h100=4'} if n['TresPerNode'] else set()),'Held GPU TRES conflict') + if n['TresPerNode']:need('JOB_GRES' not in f or f['JOB_GRES']=='gpu:h100:4','Held JOB_GRES conflict') + else:need(f.get('JOB_GRES','') in {'','(null)'},'CPU held allocation has GPU') + for k,v in f.items(): + if k in {'ReqTRES','AllocTRES','TresPerNode','JOB_GRES','GRES'}:continue + if k.lower().startswith('tresper') or 'gpu' in k.lower(): + need('gpu' not in k.lower() and 'gpu' not in v.lower(),'Auxiliary held GPU resources') + if k=='TresPerTask':need(v=='cpu:'+n['NumCPUs'],'Held per-task CPU differs') + return f + + +def producer_argv(spec,spec_path,spec_sha,ops_sha): + n=spec['native_contract'];root=Path(spec['root']) + return ['/usr/local/bin/sbatch','--parsable','--hold','--no-requeue','--export=NONE','--account='+n['Account'], + '--qos='+n['QOS'],'--partition='+n['Partition'],'--job-name='+n['JobName'],'--nodelist='+n['ReqNodeList'], + '--nodes=1','--ntasks=1','--cpus-per-task='+n['NumCPUs'],'--mem='+n['MinMemoryNode'],'--time='+n['TimeLimit'], + '--gres='+('gpu:h100:4' if n['TresPerNode'] else 'none'),'--comment=elmo-execution:'+spec_sha,'--chdir='+str(root), + '--output='+str(root/'logs/%x-%j.out'),'--error='+str(root/'logs/%x-%j.err'),str(root/'allocated.sbatch'),str(root), + spec['mode'],str(spec_path),spec_sha,ops_sha] + + +def execution(ex): + need(set(ex)=={'path','sha256','mode','ops_manifest_sha256'} and ex['mode'] in MODES,'Wrong exact phase execution binding') + path=bound({'path':ex['path'],'sha256':ex['sha256']});spec=read_json(path);root=physical(spec['root']);mode=ex['mode'] + need(spec['schema']=='native-elmo-execution-spec-v1' and spec['mode']==mode and path.parent==root/'execution' + and spec['arm'] in {'plain','fused','boundary'},'Execution schema/root/arm differs') + ops={'path':str(root/'FILES.sha256'),'sha256':ex['ops_manifest_sha256']} + need(spec['ops_manifest']==ops,'Execution/observer operations manifest differs') + names=manifest(root,ops) + need({'execution_spec.py','allocated.sbatch','host_guard.py','phase_observer.py','observe-phase.sbatch'}<=names + and all(Path(n).parts[0] not in {'execution','execution-claims','requests','scientific-claims','producer-claims','observer-specs','observer-claims','status','logs'} for n in names), + 'Missing ops role closure or circular dynamic manifest') + parent=read_json(bound(spec['parent_request']));run_id='mot-elmo-pickplace-continuous-large314m-'+spec['arm']+'-robot-human-camera-subtask-s42-300k-20261026-v1' + profile=ORIGINAL_OBSERVERS.get(run_id) + need(profile is not None and spec['parent_request']=={'path':profile['request_path'],'sha256':profile['request_sha256']} + and str(spec['original_parent_job_id'])==profile['parent_id'] and parent['run_id']==run_id + and parent['schema']=='joint-training-v2' and parent['kind']=='native-subtask-world-model-v1' and parent['phase']=='full' + and parent['target_steps']==300000 and parent['parameters']==314187064 and parent['world_size']==4,'Exact original scientific parent differs') + service(spec['operational_service']) + template_path=bound(spec['request_template']);need(template_path.parent==root/'templates','Template escaped immutable namespace') + template=read_json(template_path);schema,cpus,memory,limit,gres=MODES[mode] + need(template['schema']==schema and template['operational_service']==spec['operational_service'],'Producer phase/schema/service differs') + parent_binding=template.get('parent_request') if mode!='full' else template['interruption']['parent_request'] + need(parent_binding==spec['parent_request'],'Template belongs to another original parent') + n=spec['native_contract'] + need(set(n)==(set(NATIVE_KEYS)|{'NodeList','TresPerNode'})-{'JobId'} and all(isinstance(v,str) for v in n.values()),'Complete ID-free native contract required') + need(n['UserId']=='paphiwetsa3(41003)' and n['Account']=='loaner' and n['QOS']=='loaner-normal' and n['Partition']=='rl2' + and n['ReqNodeList']==n['NodeList']==profile['node'] and n['NumCPUs']==cpus and n['MinMemoryNode']==memory + and n['TimeLimit']==limit and n['TresPerNode']==gres and n['Requeue']==n['Restarts']=='0' + and n['Command']==str(root/'allocated.sbatch'),'Exact native phase owner/resources/Command differ') + prefix=('elmo-large314m-continuous-' if mode in {'full','diagnostic'} else 'elmo-native-recovery-')+spec['arm']+'-'+mode+'-' + need(n['JobName'].startswith(prefix) and re.fullmatch('[A-Za-z0-9_-]{1,128}',n['JobName']),'Wrong native phase name') + if mode=='full': + from scripts.train.joint.interruption_admission import compatible + need('producer_native' not in template and template['job_name']==n['JobName'],'Full template gains producer fields') + compatible(parent,template) + else: + need(template['producer_native']==n,'Template may omit only producer JobId') + if mode!='capture': + for k in ('source','source_revision','source_manifest','source_manifest_sha256','data_manifest','data_manifest_sha256','lambda_runtime','data_cache'): + need(template.get(k)==parent.get(k),'Producer changed parent source/runtime/cache: '+k) + else:need(str(template['parent_job_id'])==profile['parent_id'] and template['parent_native']['JobId']==profile['parent_id'],'Capture original-parent ID differs') + entries=spec['execution_claim_roots'] + need(entries and len({v['path'] for v in entries})==len(entries),'Complete unique execution-claim roots required') + own=[v for v in entries if v['path']==str(root/'execution-claims')] + need(len(own)==1,'Unique own terminal namespace required') + for row in entries:need(physical(row['path']).is_dir() and physical(row['terminal_directory']).is_dir(),'Unavailable registered execution root') + sc=root/'execution-claims'/ex['sha256'];a=read_json(sc/'acceptance.json');job=a['job_id'] + need(type(job) is str and re.fullmatch('[0-9]+',job),'Actual accepted producer job ID missing') + argv=producer_argv(spec,path,ex['sha256'],ex['ops_manifest_sha256']) + request_path=root/'requests'/(ex['sha256']+'-'+job+'.json') + actual=bound(a['request']);request_sha=sha(actual) + request_claim=root/('scientific-claims' if mode in {'diagnostic','full'} else 'producer-claims')/request_sha + need(set(a)=={'schema','spec_sha256','ops_manifest_sha256','job_id','argv','request','request_claim','held_native'} + and a['schema']=='native-elmo-held-execution-binding-v1' and a['spec_sha256']==ex['sha256'] + and a['ops_manifest_sha256']==ex['ops_manifest_sha256'] and a['argv']==argv and actual==request_path + and a['request_claim']==str(request_claim) and a['held_native']['path']==str(sc/'HELD_NATIVE.txt'),'Producer acceptance/spec/argv/request differs') + held=held_native(text(bound(a['held_native'])),spec,ex['sha256'],job) + q=read_json(actual);expected=copy.deepcopy(template) + if mode!='full':expected['producer_native']['JobId']=job;expected_bytes=encoded(expected) + else:expected_bytes=data(template_path) + need(data(actual)==expected_bytes and q['operational_service']==spec['operational_service'],'Actual request changed beyond sole actual JobId binding') + for directory in (sc,request_claim):need(physical(directory).is_dir() and text(directory/'PARENT_JOB_ID').strip()==job,'Spec/request actual IDs differ') + prepared=read_json(sc/'PREPARED.json') + need(set(prepared)=={'spec_sha256','ops_manifest_sha256','spec_path','arm','parent_request','argv','original_parent_job_id','created_at'} + and prepared['spec_sha256']==ex['sha256'] and prepared['ops_manifest_sha256']==ex['ops_manifest_sha256'] + and prepared['spec_path']==str(path) and prepared['arm']==spec['arm'] and prepared['parent_request']==spec['parent_request'] + and prepared['argv']==argv and str(prepared['original_parent_job_id'])==profile['parent_id'],'Original producer submission intent differs') + canonical=read_json(request_claim/'PREPARED.json') + need(canonical==dict(request_sha256=request_sha,spec_sha256=ex['sha256'],spec_path=str(path),actual_request_path=str(actual),argv=argv, + original_parent_job_id=profile['parent_id']),'Canonical materialized request-SHA claim differs') + need(text(sc/'sbatch.stdout').strip()==job and text(sc/'sbatch.exit').strip()=='0','Producer submission acceptance unresolved') + text(sc/'sbatch.stderr') + return dict(spec=spec,spec_path=path,root=root,request_path=actual,request=q,request_claim=request_claim,held=held,job=job, + execution=ex,terminal_directory=physical(own[0]['terminal_directory']),run_id=run_id) + + +def observer_name(ctx): + spec=ctx['spec'];value='elmo-native-phase-'+spec['arm']+'-'+spec['mode']+'-'+ctx['execution']['sha256'][:16] + scope=spec.get('runtime',{}).get('admission',{}).get('quiescence',{}) + need(value!=spec['native_contract']['JobName'] and value not in scope.get('scientific_job_names',[]) + and not any(value.startswith(p) for p in scope.get('scientific_name_prefixes',[])) + and not value.startswith(('joint-','elmo-large314m-continuous-')),'Observer collides with scientific writer scope') + return value + + +def observer_native(ctx): + n=ctx['spec']['native_contract'] + return dict(JobName=observer_name(ctx),UserId=n['UserId'],Account=n['Account'],QOS=n['QOS'],Partition=n['Partition'], + ReqNodeList=n['ReqNodeList'],NodeList=n['NodeList'],NumCPUs='2',MinMemoryNode='1G',TimeLimit='00:10:00',Requeue='0',Restarts='0', + Command=str(ctx['root']/'observe-phase.sbatch'),TresPerNode='') + + +def observer_argv(ctx,path,spec_sha): + n=observer_native(ctx);root=ctx['root'];comment='elmo-phase-observer:'+spec_sha + need(comment!=ctx['request'].get('run_id'),'Observer comment is scientific identity') + return ['/usr/local/bin/sbatch','--parsable','--no-requeue','--export=NONE','--dependency=afterany:'+ctx['job'], + '--account='+n['Account'],'--qos='+n['QOS'],'--partition='+n['Partition'],'--job-name='+n['JobName'], + '--nodelist='+n['ReqNodeList'],'--nodes=1','--ntasks=1','--cpus-per-task=2','--mem=1G','--time=00:10:00','--gres=none', + '--comment='+comment,'--chdir='+str(root),'--output='+str(root/'logs/%x-%j.out'),'--error='+str(root/'logs/%x-%j.err'), + str(root/'observe-phase.sbatch'),str(root),str(path),spec_sha,ctx['execution']['ops_manifest_sha256']] + + +def pending_native(raw,ctx,observer,spec_sha): + f=fields(raw);n=observer_native(ctx) + need(f.get('JobId')==observer and f.get('JobState')=='PENDING' and f.get('Reason')=='Dependency' + and all(f.get(k,'')==v for k,v in n.items() if k not in {'NodeList','TresPerNode'}) + and f.get('NodeList','') in {'','(null)',n['NodeList']} and f.get('NumNodes') in {'1','1-1'} and f.get('NumTasks')=='1' + and re.fullmatch('afterany:'+ctx['job']+r'\((?:unfulfilled|unsatisfied)\)',f.get('Dependency','')) + and f.get('Comment')=='elmo-phase-observer:'+spec_sha and f.get('SubmitTime') not in {None,'Unknown','None'} + and f.get('WorkDir')==str(ctx['root']),'Observer accepted pending native differs') + for k,suffix in [('StdOut','.out'),('StdErr','.err')]: + need(f.get(k)==str(ctx['root']/'logs'/(n['JobName']+'-'+observer+suffix)),'Observer output namespace differs') + items=[p.split('=',1) for p in f.get('ReqTRES','').split(',')] + need(all(len(v)==2 for v in items) and len(items)==len(dict(items)),'Ambiguous observer TRES') + tres=dict(items);need(tres.get('cpu')=='2' and tres.get('node')=='1' and tres.get('mem') in {'1G','1024M'} + and f.get('AllocTRES','') in {'','(null)'},'Observer pending requested/allocated TRES differ') + for k,v in f.items(): + if k in {'ReqTRES','AllocTRES'}:need('gpu' not in v.lower(),'Observer has GPU TRES') + if k.lower().startswith(('gres','job_gres','tresper')): + need(v.lower() in {'','(null)','none','n/a','0'} or k=='TresPerTask' and v=='cpu:2','Observer GPU/GRES conflict') + return f + + +def observer_context(observer_binding): + result_path=bound(observer_binding);result=read_json(result_path) + os_binding=result['observer_spec'];os_path=bound(os_binding);ospec=read_json(os_path) + need(set(ospec)==OBSERVER_KEYS and ospec['schema']=='native-elmo-phase-observer-spec-v1','Wrong ID-free observer spec') + ctx=execution(ospec['producer_execution']);root=ctx['root'];ps=ctx['execution']['sha256'];osh=os_binding['sha256'] + need(ospec['root']==str(root) and os_path.parent==root/'observer-specs' and ctx['spec']['release_observer_path']==str(os_path) + and ospec['terminal_directory']==str(ctx['terminal_directory']),'Observer spec escaped actual producer namespace') + ledger=physical(root/'observer-claims'/ps);job=ctx['job'];observer=result['observer_job_id'] + need(type(observer) is str and re.fullmatch('[0-9]+',observer) and observer!=job,'Wrong actual observer ID') + expected_argv=observer_argv(ctx,os_path,osh);a=read_json(ledger/'acceptance.json') + need(a==dict(schema='native-elmo-phase-observer-acceptance-v1',observer_spec=os_binding, + ops_manifest_sha256=ctx['execution']['ops_manifest_sha256'],producer_execution=ctx['execution'], + producer_request=binding(ctx['request_path']),producer_job_id=job,observer_job_id=observer,argv=expected_argv, + pending_native=binding(ledger/'PENDING_NATIVE.txt')),'Observer acceptance belongs to another exact producer') + intent=read_json(ledger/'PREPARED.json') + need(intent==dict(schema='native-elmo-phase-observer-intent-v1',observer_spec=os_binding, + ops_manifest_sha256=ctx['execution']['ops_manifest_sha256'],producer_execution=ctx['execution'], + producer_request=binding(ctx['request_path']),producer_job_id=job,argv=expected_argv),'Original observer submission intent differs') + need(text(ledger/'sbatch.stdout').strip()==observer and text(ledger/'sbatch.exit').strip()=='0' + and text(ledger/'OBSERVER_JOB_ID').strip()==observer,'Observer actual submission acceptance unresolved') + text(ledger/'sbatch.stderr') + pending_raw=successful(read_json(ledger/'ACCEPTANCE_QUERY.json'),native_argv(observer)) + need(pending_raw==text(bound(a['pending_native'])),'Observer acceptance query/native bytes differ') + pending=pending_native(pending_raw,ctx,observer,osh) + original=held_native(successful(read_json(ledger/'PRODUCER_HELD_QUERY.json'),native_argv(job)),ctx['spec'],ps,job) + need(original['SubmitTime']==ctx['held']['SubmitTime'],'Observer submitted against another actual held producer') + release=read_json(root/'execution-claims'/ps/'RELEASE_ATTEMPT.json') + need(set(release)=={'job_id','spec_sha256','request_sha256','held_native_sha256','observer_readiness'} + and release['job_id']==job and release['spec_sha256']==ps and release['request_sha256']==sha(ctx['request_path']), + 'Producer release intent differs from accepted execution') + release_held_path=root/'execution-claims'/ps/'RELEASE_HELD_NATIVE.txt' + need(sha(release_held_path)==release['held_native_sha256'],'Actual held native release record changed') + released_held=held_native(text(release_held_path),ctx['spec'],ps,job) + need(released_held['SubmitTime']==ctx['held']['SubmitTime'],'Released held allocation differs from acceptance') + ready=release['observer_readiness'] + need(set(ready)=={'schema','producer_job_id','producer_request_sha256','observer_job_id','observer_spec_sha256', + 'observer_pending_query','observer_own_terminal','production_completion_proven'} + and ready['schema']=='native-elmo-phase-observer-ready-v1' and ready['producer_job_id']==job + and ready['producer_request_sha256']==sha(ctx['request_path']) and ready['observer_job_id']==observer + and ready['observer_spec_sha256']==osh and ready['observer_own_terminal']=='UNKNOWN_NOT_REQUIRED' + and ready['production_completion_proven'] is False,'Exact observer-before-release readiness missing') + current=pending_native(successful(ready['observer_pending_query'],native_argv(observer)),ctx,observer,osh) + need(current['SubmitTime']==pending['SubmitTime'],'Observer changed before producer release') + ctx.update(result=result,result_path=result_path,observer_spec=os_binding,observer_spec_path=os_path,observer_id=observer, + observer_ledger=ledger,observer_acceptance=a,pending=pending) + return ctx + + +def check_owned_steps(raw,job): + strict_step_rows(raw) + for line in raw.splitlines(): + if not line.strip():continue + step=line.split('|',1)[0].strip() + need(step!=job and not step.startswith(job+'.') and not step.startswith(job+'_'), + 'Producer still has native workers') + + +def verify_observer(observer_binding,*,producer_request=None,producer_native=None,terminal_native=None): + """Reopen all evidence, returning verified context; never issues a query.""" + ctx=observer_context(observer_binding);result=ctx['result'];job=ctx['job'];root=ctx['root'];terminal_dir=ctx['terminal_directory'] + canonical=dict(native=terminal_dir/('native-'+job+'.txt'),query=terminal_dir/('native-'+job+'.query.json'), + query_exit=terminal_dir/('native-'+job+'.exit'),query_stderr=terminal_dir/('native-'+job+'.stderr'), + steps=terminal_dir/('native-'+job+'.steps-query.json'),result=terminal_dir/('native-'+job+'.observer.json')) + need(ctx['result_path']==canonical['result'],'Observer final commit is not canonical producer destination') + own_query=read_json(ctx['observer_ledger']/'RUNNING_NATIVE_QUERY.json') + own=native_identity(successful(own_query,native_argv(ctx['observer_id'])),dict(observer_native(ctx),JobId=ctx['observer_id']),'RUNNING') + osh=ctx['observer_spec']['sha256'] + need(own.get('SubmitTime')==ctx['pending']['SubmitTime'] and own.get('Comment')=='elmo-phase-observer:'+osh,'Own observer differs from accepted afterany allocation') + query=read_json(canonical['query']);raw=successful(query,native_argv(job));state=fields(raw).get('JobState') + need(state in TERMINAL,'Producer is not genuinely terminal') + expected_native=dict(ctx['spec']['native_contract'],JobId=job) + terminal=native_identity(raw,expected_native,state) + need(terminal.get('SubmitTime')==ctx['held']['SubmitTime'] and terminal.get('Comment')=='elmo-execution:'+ctx['execution']['sha256'], + 'Producer terminal contradicts exact accepted held identity') + need(datetime.fromisoformat(terminal['StartTime'])<=datetime.fromisoformat(terminal['EndTime'])<=datetime.fromisoformat(own['StartTime']), + 'Producer terminal/afterany time order differs') + running_path=ctx['request_claim']/'RUNNING_ADMISSION.json';running_binding=None + if running_path.exists() or running_path.is_symlink(): + running=read_json(running_path) + need(running['job_id']==job and running['request_sha256']==sha(ctx['request_path']) and running['spec_sha256']==ctx['execution']['sha256'],'Producer RUNNING claim differs') + actual=native_identity(running['native'],expected_native,'RUNNING') + need(actual.get('SubmitTime')==ctx['held']['SubmitTime'] and actual.get('Comment')=='elmo-execution:'+ctx['execution']['sha256'] + and actual['StartTime']==terminal['StartTime'],'Producer RUNNING contradicts held/terminal identity') + running_binding=binding(running_path) + steps=read_json(canonical['steps']);step_argv=['/usr/local/bin/squeue','--steps','--noheader','--user','paphiwetsa3','--format=%i|%T'] + step_raw=successful(steps,step_argv);check_owned_steps(step_raw,job) + need(text(canonical['native'])==raw and text(canonical['query_exit']).strip()=='0' and not text(canonical['query_stderr']), + 'Canonical native/query artifacts disagree') + stamp=datetime.fromisoformat(result['captured_at_utc']);need(stamp.tzinfo is not None and stamp<=datetime.now(timezone.utc),'Missing/future final commit time') + expected=dict(schema='native-elmo-phase-observer-result-v1',status='GENUINE_PRODUCER_TERMINAL_CAPTURED',producer_execution=ctx['execution'], + producer_request=binding(ctx['request_path']),producer_job_id=job,producer_phase=ctx['spec']['mode'],producer_terminal_state=terminal['JobState'], + producer_terminal_exit=terminal['ExitCode'],producer_native=binding(canonical['native']),producer_query=binding(canonical['query']), + producer_query_exit=binding(canonical['query_exit']),producer_query_stderr=binding(canonical['query_stderr']),owned_steps=binding(canonical['steps']), + producer_running_admission=running_binding,producer_running_admission_status='VALIDATED' if running_binding else 'NOT_RECORDED', + observer_spec=ctx['observer_spec'],observer_acceptance=binding(ctx['observer_ledger']/'acceptance.json'), + observer_running_query=binding(ctx['observer_ledger']/'RUNNING_NATIVE_QUERY.json'),observer_job_id=ctx['observer_id'], + observer_own_terminal='UNKNOWN_NOT_REQUIRED',observer_host_exit_is_native_proof=False,phase_completion_proven=False,zero_work_proven=False, + captured_at_utc=result['captured_at_utc'],automatic_recovery_authorized=False) + need(result==expected,'Final observer commit/evidence semantic closure differs') + if producer_request is not None:need(producer_request==binding(ctx['request_path']),'Observer certifies another producer request') + if producer_native is not None:need(producer_native==expected_native,'Observer certifies another native contract') + if terminal_native is not None:need(terminal_native==binding(canonical['native']),'Observer certifies another terminal native binding') + ctx.update(terminal=terminal,expected_native=expected_native,running_admission=running_binding, + running_admission_status='VALIDATED' if running_binding else 'NOT_RECORDED') + return ctx + + +def verify_claim_terminal(claim,terminal_directory,run_id): + """Exact old-original route or complete new spec-bound phase observer.""" + claim=physical(claim);terminal_directory=physical(terminal_directory) + prepared=read_json(claim/'PREPARED.json');job=text(claim/'PARENT_JOB_ID').strip() + need(job.isdecimal() and prepared['request_sha256']==claim.name,'Unknown/incorrect accepted request claim') + profile=ORIGINAL_OBSERVERS.get(run_id) + need(profile is not None,'Claim outside exact supported original identity') + if job==profile['parent_id']: + root=Path(profile['root']) + need(claim==root/'claims'/profile['request_sha256'] and terminal_directory==root,'Unreviewed original parent/claim/terminal route') + running_binding=binding(claim/'RUNNING_ADMISSION.json');running=read_json(claim/'RUNNING_ADMISSION.json') + f=fields(running['native']);expected={k:f.get(k,'') for k in (*NATIVE_KEYS,'NodeList','TresPerNode')} + paths=dict(files_manifest=root/'FILES.sha256',capture_script=root/'capture-native.sbatch',guard=root/'guard_full.py', + submitter=root/'submit_full_once.py',parent_script=root/'full.sbatch',prepared=claim/'PREPARED.json', + parent_id=claim/'PARENT_JOB_ID',observer_id=claim/'CAPTURE_JOB_ID',pending_native=claim/'CAPTURE_PENDING_NATIVE.txt', + running_admission=claim/'RUNNING_ADMISSION.json',terminal_native=root/('native-'+job+'.txt'), + query_exit=root/('native-'+job+'.exit'),query_stderr=root/('native-'+job+'.stderr')) + if profile['has_smoke_identity']:paths['smoke_identity']=root/'smoke-identity.json' + proof=dict(schema='native-elmo-original-observer-v1',**{k:binding(v) for k,v in paths.items()}) + q=dict(parent_request={'path':profile['request_path'],'sha256':profile['request_sha256']},parent_job_id=job, + parent_native=expected,parent_running_admission=running_binding,original_observer=proof) + from scripts.train.joint.original_observer import validate_original + def reader(b,json_value=True):return read_json(bound(b)) if json_value else text(bound(b)) + raw,provenance=validate_original(q,reader) + return dict(job_id=job,terminal_native_sha256=sha(paths['terminal_native']),terminal_state=fields(raw)['JobState'], + provenance='EXACT_ORIGINAL_AFTERANY',observer_own_terminal=provenance['observer_own_terminal']) + need('spec_path' in prepared and 'spec_sha256' in prepared and 'actual_request_path' in prepared, + 'New accepted claim lacks exact execution/materialized request provenance') + commit=terminal_directory/('native-'+job+'.observer.json') + verified=verify_observer(binding(commit),producer_request={'path':prepared['actual_request_path'],'sha256':claim.name}, + terminal_native=binding(terminal_directory/('native-'+job+'.txt'))) + need(verified['request_claim']==claim and verified['job']==job and verified['run_id']==run_id + and verified['execution']['path']==prepared['spec_path'] and verified['execution']['sha256']==prepared['spec_sha256'], + 'New phase terminal observer belongs to another exact claim/spec/identity') + return dict(job_id=job,terminal_native_sha256=verified['result']['producer_native']['sha256'], + terminal_state=verified['terminal']['JobState'],provenance='COMPLETE_PHASE_OBSERVER_COMMIT', + observer_commit=binding(commit),observer_own_terminal='UNKNOWN_NOT_REQUIRED', + running_admission_status=verified['running_admission_status']) diff --git a/scripts/train/joint/production_history.py b/scripts/train/joint/production_history.py new file mode 100644 index 000000000..a38d38f0c --- /dev/null +++ b/scripts/train/joint/production_history.py @@ -0,0 +1,106 @@ +"""Retain immutable production checkpoints and validation across crash attempts. + +Every old file is preserved. The selected production checkpoint defines the +committed prefix of a failed attempt; later unsaved work is audit history only. +""" +import csv +import math +from pathlib import Path +import re + +from scripts.train.joint.contract import bound, need, sha +from scripts.train.joint.checkpoint_inspection import attempt_chain, scheduled_filename + + +def has_interruption(request): + if request.get("kind") != "native-subtask-world-model-v1": + return False + return any(row.get("interruption") is not None for row in attempt_chain(request)) + + +def start_step(request): + if request.get("interruption"): + return request["interruption"]["checkpoint"]["global_step"] + if request.get("continuation"): + return bound(request["continuation"]["completion"])["global_step"] + return 0 + + +def segments(request, stop): + """Chronological attempt ownership intervals: (start, end].""" + attempts = attempt_chain(request)[::-1] + result = [] + for index, attempt in enumerate(attempts): + start = start_step(attempt) + end = start_step(attempts[index + 1]) if index + 1 < len(attempts) else stop + need(type(start) is type(end) is int and 0 <= start <= end <= request["target_steps"] + and start % 20000 == end % 20000 == 0, "Invalid committed attempt history") + relative = Path("csv") / f"step{start:09d}" / "metrics.csv" + local, mirror = Path(attempt["local_logs"]) / relative, Path(attempt["output"]) / "log-mirror" / relative + need(local != mirror, "Local/mirror log namespace collision") + result.append(dict(request=attempt, start=start, end=end, csv=local if local.is_file() else mirror)) + need(len({str(row["csv"]) for row in result}) == len(result), "Repeated production log namespace") + return result + + +def retained_milestones(request, stop): + milestones, excluded = {}, [] + for row in segments(request, stop): + directory = Path(row["request"]["output"]) / "checkpoints" + need(directory.is_dir() and not directory.is_symlink(), "Retained attempt checkpoint directory missing") + for path in directory.glob("*.ckpt"): + if path.name == "last.ckpt": + continue + need(path.is_file() and not path.is_symlink(), "Invalid retained checkpoint name/type") + step = scheduled_filename(path.name) + need(step > 0 and step % 20000 == 0 and step <= request["target_steps"], + "Unexpected unscheduled production checkpoint") + if not row["start"] < step <= row["end"]: + excluded.append(dict(path=str(path), global_step=step, retained=True)) + continue + if step in milestones: + need(milestones[step].resolve() == path.resolve(), "Duplicate committed production milestone") + milestones[step] = path + need(set(milestones) == set(range(20000, stop + 1, 20000)), "Committed production history lost a scheduled milestone") + return milestones, dict(schema="native-elmo-retained-attempt-history-v1", retained_steps=sorted(milestones), + files=[dict(global_step=step, path=str(path), sha256=sha(path)) + for step, path in sorted(milestones.items())], + excluded_from_committed_prefix_but_preserved=excluded) + + +def scheduled_validation(request, stop, metric_names): + """Merge disjoint committed prefixes, rejecting duplicate/partial/nonfinite events.""" + validation, evidence = {}, [] + for item in segments(request, stop): + selected, discarded = {}, set() + path = item["csv"] + need(path.is_file() and not path.is_symlink(), "Actual retained validation CSV missing") + with path.open() as stream: + for row in csv.DictReader(stream): + metrics = {k: float(v) for k, v in row.items() if k.startswith("val/") and v} + if not metrics: + continue + need(all(math.isfinite(value) for value in metrics.values()), + "Observed validation is nonfinite, including uncommitted history") + raw = float(row["step"]) + need(math.isfinite(raw) and raw.is_integer(), "Invalid validation step") + step = int(raw) + step = step + 1 if step % 20000 == 19999 else step + need(step > 0 and step % 20000 == 0 and step <= request["target_steps"], + "Unexpected non-scheduled production validation") + need(step != item["start"], "Resumed attempt repeated checkpoint-boundary validation") + if not item["start"] < step <= item["end"]: + discarded.add(step) + continue + need(set(metrics) <= set(metric_names) and all(math.isfinite(v) for v in metrics.values()), + "Unexpected/nonfinite committed validation metric") + target = selected.setdefault(step, {}) + need(not target.keys() & metrics.keys(), "Repeated validation metric/event within attempt") + target.update(metrics) + need(not selected.keys() & validation.keys(), "Repeated validation across committed attempts") + need(all(set(v) == set(metric_names) for v in selected.values()), "Incomplete scheduled validation event") + validation.update(selected) + evidence.append(dict(path=str(path), sha256=sha(path), start=item["start"], end=item["end"], + included_steps=sorted(selected), uncommitted_steps_preserved=sorted(discarded))) + need(set(validation) == set(range(20000, stop + 1, 20000)), "Missing committed scheduled validation event") + return validation, evidence diff --git a/scripts/train/joint/production_recovery.py b/scripts/train/joint/production_recovery.py new file mode 100644 index 000000000..f404999b5 --- /dev/null +++ b/scripts/train/joint/production_recovery.py @@ -0,0 +1,238 @@ +"""Actual2+4-update recovery diagnostic from an inspected production milestone. + +The maintained training loader and ResumeOptimizerProof perform all restoration +and updates. This executable diagnostic does not admit a production relaunch. +Its extra checkpoints are diagnostic evidence only, never production inputs. +""" +import copy +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys + +from scripts.train.joint.contract import bound, need, sha, write_once, canonical_uuid + + +def phases(start, target): + need(type(start) is type(target) is int and 0 < start < target + and start % 20000 == target % 20000 == 0 and start + 6 < target, + "Production diagnostic needs a retained nonterminal20k milestone") + return ((start, start + 2), (start + 2, start + 6)) + + +def phase_config(original, *, output, checkpoint, origin, start, stop, target, world, proof): + from omegaconf import OmegaConf, open_dict + + need((start, stop) in phases(origin, target) and world == 4, + "Only exact4-rank2+4 production diagnostic allowed") + cfg = OmegaConf.create(copy.deepcopy(OmegaConf.to_container(original, resolve=True))) + science = {k: OmegaConf.to_container(cfg[k], resolve=True) for k in ("model", "data")} + need(any(c.get("_target_") == "egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage" + for c in cfg.callbacks.values()), "Maintained native conditional optimizer callback required") + with open_dict(cfg): + cfg.ckpt_path = checkpoint + cfg.paths.output_dir = str(output) + cfg.val_at_start = cfg.val_at_end = False + cfg.trainer.max_steps = origin + 6 + cfg.trainer.val_check_interval = 2 + cfg.trainer.limit_val_batches = 100 // world + cfg.trainer.num_sanity_val_steps = 0 + cfg.trainer.check_val_every_n_epoch = None + cfg.callbacks.pop("joint_smoke", None) + cfg.callbacks.pop("allocation_boundary", None) + cfg.callbacks.pop("wandb_profiler", None) + cfg.logger.pop("wandb", None) + cfg.logger.csv.save_dir = str(output / "logs") + cfg.logger.csv.version = "resume-" + str(start) + 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 = str(output / "checkpoints") + cb.filename = "diagnostic-epoch{epoch:04d}-step{step:09d}" + cfg.callbacks.resume_proof = dict( + _target_="scripts.train.joint.recovery_check.ResumeOptimizerProof", + output=str(proof), world_size=world, start_step=start, stop_step=stop, + checkpoint=str(checkpoint), multimodal=True) + if stop == origin + 2: + cfg.callbacks.allocation_boundary = dict( + _target_="scripts.train.joint.allocation_boundary.ScheduledAllocationBoundary", + start_step=start, stop_step=stop, target_step=origin + 6, cadence=2, + output=str(proof)) + if cfg.evaluator.get("pixel_diagnostics") is not None: + cfg.evaluator.pixel_diagnostics.log_wandb = False + cfg.evaluator.pixel_diagnostics.artifact_dir = str(proof / ("validation-" + str(stop))) + need(all(OmegaConf.to_container(cfg[k], resolve=True) == v for k, v in science.items()), + "Production diagnostic changed model/loss/data/sampling") + return cfg + + +def read_phase_proofs(output, start, target, world, job_id): + """Reopen exact callback files; reused by an independent finalizer.""" + output = Path(output) + proof = output / "proof" + rows = [] + for first, last in phases(start, target): + ranks = [json.loads((proof / f"resume-{first}-{last}-rank{rank}.json").read_text()) + for rank in range(world)] + need(all(row["status"] == "RESUME_PASS" and row["job_id"] == job_id + and row["rank"] == rank and row["world_size"] == world + and row["start_step"] == first and row["global_step"] == last + and row["optimizer_updates"] == row["fp32_losses"] == row["bf16_outputs"] == last - first + and row["fp32_gradients_verified"] is True + and row["strict_start_weights_optimizer_scheduler_rng_cursor"] is True + and row["validation_steps"] == list(range(first + 2, last + 1, 2)) + for rank, row in enumerate(ranks)), "Actual all-rank production restore proof incomplete") + rows.append(ranks) + return rows + + +def validate_phase_proofs(rows, selected, admission=None): + from scripts.train.joint.native_recovery import summarize_native_phases, state_digest + + start = selected["global_step"] + flags = summarize_native_phases(rows, start_step=start, production=True) + need(all(row["initial_weights_sha256"] == selected["proof"]["model_state_sha256"] + for row in rows[0]), "Diagnostic did not initialize from inspected production weights") + for rank, row in enumerate(rows[0]): + need(row["actual_rank_rng_restore"]["expected_sha256"] == selected["proof"]["rank_rng_sha256"][rank], + "Diagnostic restored RNG differs from the strictly inspected production checkpoint") + restored_usage = dict(schema="native-conditional-optimizer-usage-v1", global_step=start, + parameter_names=row["conditional_optimizer_record"]["parameter_names"], + parameter_update_counts=row["initial_parameter_usage"]) + need(state_digest(restored_usage) == selected["proof"]["conditional_optimizer_usage_sha256"], + "Diagnostic conditional counters differ from inspected production state") + # Both sides use the maintained weights_digest. The launch path separately + # checks the byte-identical complete input checkpoint SHA. + need(all(a["weights_sha256"] == b["initial_weights_sha256"] for a, b in zip(*rows)), + "Diagnostic second process did not restore first-process weights") + devices = [[canonical_uuid(row["gpu_uuid"]) for row in phase] for phase in rows] + need(len(set(devices[0])) == 4 and devices[0] == devices[1], + "Actual four distinct GPUs must be unchanged across both diagnostic processes") + 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 rows), "Actual restored ranks differ from admitted GPUs") + return flags + + +def summarize(request_path, q, original, admission, selected): + """Publish diagnostic results, leaving host/native completion to a producer.""" + output, start, world = Path(q["output"]), selected["global_step"], original["world_size"] + proof = output / "proof" + rows = read_phase_proofs(output, start, original["target_steps"], world, os.environ["SLURM_JOB_ID"]) + flags = validate_phase_proofs(rows, selected, admission) + boundary = json.loads((proof / f"allocation-step{start + 2:09d}.json").read_text()) + need(boundary["status"] == "PASS" and boundary["checkpoint_saved_after_validation"] is True + and boundary["start_step"] == start and boundary["global_step"] == start + 2 + and sha(boundary["checkpoint"]["path"]) == boundary["checkpoint"]["sha256"], + "Actual diagnostic post-validation process boundary missing") + write_once(output / "production-recovery-diagnostic.json", dict( + schema="native-elmo-production-recovery-diagnostic-v1", + status="DIAGNOSTIC_COMPLETE_REQUIRES_HOST_AND_NATIVE_SUCCESS", + request_sha256=sha(request_path), job_id=os.environ["SLURM_JOB_ID"], + parent_request=q["parent_request"], inspection=q["inspection"], + production_checkpoint={k: selected[k] for k in ("path", "sha256", "global_step", "size_bytes")}, + diagnostic_checkpoint=boundary["checkpoint"], actual_updates=6, + run_id=original["run_id"], world_size=world, resumed_from=start, resumed_to=start + 6, + no_extra_validation=True, model_data_unchanged=True, + source_manifest_sha256=original["source_manifest_sha256"], + data_manifest_sha256=original["data_manifest_sha256"], + rank_proofs=[dict(path=str(p), sha256=sha(p)) for p in sorted(proof.glob("resume-*-rank*.json"))], + allocation_attestation=admission, production_relaunch_authorized=False, **flags)) + + +def main(request_path): + # The hash-bound allocated launcher must perform current runtime/cache/GPU + # health/claim admission before dispatching this module. Check that actual + # RUNNING native record again before any Torch/checkpoint import. + from scripts.train.joint.native_interruption import command, native_identity, process_identity + + q = json.loads(Path(request_path).read_text()) + need(q["schema"] == "native-elmo-production-recovery-request-v1", "Wrong production diagnostic request") + native = q["producer_native"] + process_identity(native) + need(native["JobId"] == os.environ.get("SLURM_JOB_ID") + and native["NodeList"] == os.environ.get("SLURMD_NODENAME") + and native["NumCPUs"] == "32" and native["MinMemoryNode"] == "192G" + and native["TimeLimit"] == "01:00:00" and native["TresPerNode"] == "gres/gpu:h100:4" + and os.environ.get("SLURM_RESTART_COUNT", "0") == "0", "Exact fresh4H10032CPU192G1h diagnostic required") + raw = command(["/usr/local/bin/scontrol", "show", "job", native["JobId"], "-o", "-dd"]) + native_identity(raw, native, "RUNNING") + prerequisite_path = Path(q["allocated_prerequisites_path"]) + need(prerequisite_path.is_file() and not prerequisite_path.is_symlink(), "Actual allocated prerequisites missing") + prerequisite = json.loads(prerequisite_path.read_text()) + need(prerequisite["status"] == "ACTUAL_PREREQUISITES_AND_NATIVE_ALLOCATION_PASS" + and prerequisite["job_id"] == native["JobId"] and prerequisite["request_sha256"] == sha(request_path), + "Actual producer prerequisite/claim admission differs") + native_identity(prerequisite["native"], native, "RUNNING") + original = bound(q["parent_request"]) + need(original["phase"] == "full" and original["kind"] == "native-subtask-world-model-v1" + and original["world_size"] == 4 and original["parameters"] == 314187064, + "Exact native Elmo314M parent required") + from scripts.train.joint.checkpoint_inspection import verify_source_binding + from scripts.train.joint.interruption_admission import recovery_phase_admission, completed_producer + eligibility = recovery_phase_admission(q, request_path, scientific_writer=True) + verify_source_binding(original, q) + _, inspection = completed_producer(q["inspection_producer"], "native-elmo-checkpoint-inspection-result-v1", + "INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION") + need(q["inspection"] == q["inspection_producer"]["result"], "GPU diagnostic bound another inspection result") + need(inspection["parent_request"] == q["parent_request"] + and inspection["status"] == "INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION" + and inspection["highest_healthy_known"] is True, "Actual highest-production inspection missing") + selected = inspection["selected"] + need(sha(selected["path"]) == selected["sha256"], "Inspected production checkpoint changed") + phases(selected["global_step"], original["target_steps"]) + bound(q["resolved"], False) + for line in Path(original["source_manifest"]).read_text().splitlines(): + digest, relative = line.split(" ", 1) + need(sha(Path(original["source"]) / relative) == digest, "Checkpoint-bound scientific source changed") + need(sha(original["source_manifest"]) == original["source_manifest_sha256"] + and sha(original["data_manifest"]) == original["data_manifest_sha256"], "Source/data binding changed") + output, root = Path(q["output"]), Path(request_path).parent.parent + need(not output.exists() and q["reserved_checkpoint_bytes"] >= 4 * selected["size_bytes"] + 1000000000 + and shutil.disk_usage(output.parent).free > q["reserved_checkpoint_bytes"], + "Fresh diagnostic output and retained full-state reserve required") + output.mkdir() + write_once(output / "phase-admission.json", dict(request_sha256=sha(request_path), **eligibility)) + (output / "proof").mkdir() + (output / "configs").mkdir() + from omegaconf import OmegaConf + from scripts.train.joint.allocation import admit + from scripts.train.joint.service_bootstrap import command as service_command + + execution = copy.deepcopy(original) + execution["operational_service"] = q["operational_service"] + execution["allocation"] = q.get("allocation") + execution["lambda_runtime"] = q.get("lambda_runtime", original.get("lambda_runtime")) + admission = admit(execution, status_root=root / "status", bootstrap=root / "runtime_bootstrap.py") + cfg = OmegaConf.load(q["resolved"]["path"]) + checkpoint = selected["path"] + origin = selected["global_step"] + for start, stop in phases(origin, original["target_steps"]): + phase = phase_config(cfg, output=output, checkpoint=checkpoint, origin=origin, + start=start, stop=stop, target=original["target_steps"], world=4, + proof=output / "proof") + path = output / "configs" / f"resume-{start}-{stop}.yaml" + with path.open("x") as stream: + stream.write(OmegaConf.to_yaml(phase, resolve=True)) + argv = [sys.executable, str(root / "runtime_bootstrap.py"), "-m", "torch.distributed.run", + "--standalone", "--nnodes=1", "--nproc-per-node=4", "--max-restarts=0", + *service_command(q, Path(request_path), root / "runtime_bootstrap.py", + ["-m", "egomimic.trainHydra"])[1:], + "--config-path", str(path.parent), "--config-name", path.stem, + "hydra.run.dir=" + str(output / f"hydra-{stop}"), "hydra.job.chdir=false"] + subprocess.run(argv, check=True, cwd=original["source"], env=dict(os.environ, WANDB_MODE="disabled")) + if stop == origin + 2: + boundary = json.loads((output / "proof" / f"allocation-step{stop:09d}.json").read_text()) + need(boundary["status"] == "PASS" and boundary["checkpoint_saved_after_validation"], + "Actual first diagnostic process did not complete after validation") + checkpoint = boundary["checkpoint"]["path"] + need(sha(checkpoint) == boundary["checkpoint"]["sha256"], "Diagnostic boundary checkpoint changed") + summarize(request_path, q, original, admission, selected) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/scripts/train/joint/recovery_check.py b/scripts/train/joint/recovery_check.py new file mode 100644 index 000000000..9f87cccba --- /dev/null +++ b/scripts/train/joint/recovery_check.py @@ -0,0 +1,403 @@ +"""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, multimodal=False): + super().__init__(output, world_size, expected_steps=stop_step, warmup_steps=start_step, + multimodal=multimodal) + self.start_step, self.checkpoint = start_step, checkpoint + self.rng_cleanup = None + + 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") + from scripts.train.joint.native_recovery import checkpoint_science + checkpoint_science(saved["hyper_parameters"]["config_tree"], pl_module.hparams.config_tree) + 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") + need(equal_state(trainer.datamodule.state_dict(), saved["PixelProprioDataModule"]), + "Complete sampler continuation state differs") + need(equal_state(pl_module.data_context.snapshot(), saved["data_context"]), + "Checkpoint data context was not strictly 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 + if self.multimodal: + from egomimic.pl_utils.joint_smoke import _check_pilot_interfaces + from scripts.train.joint.native_recovery import restored_conditional_counts, install_rng_observer + + _check_pilot_interfaces(self.network, set(self.network.pose_dims), set(self.network.action_dims)) + need(set(self.network.pose_dims) == {"eva_bimanual", "human_bimanual"} + and set(self.network.action_dims) == {"eva_bimanual"} + and "camera" in self.network.streams, "Wrong native Elmo interfaces") + need(torch.distributed.get_backend() == "nccl", "Native recovery requires NCCL") + prop = torch.cuda.get_device_properties(pl_module.device) + need("H100" in prop.name and prop.major >= 9, "Native recovery requires actual H100") + image = self.network.image_model + need(image.token_execution == "packed_valid" and image.blocks[0].attention_backend == "flash4", + "Native recovery changed Flash4/packed execution") + need(trainer.limit_val_batches == 100 // self.world_size + and len(trainer.datamodule.val_dataloader()) >= 100 // self.world_size, + "Native recovery needs100 global validation anchors") + self.parameter_usage, self.usage_callback = restored_conditional_counts(trainer, pl_module, saved) + self.initial_parameter_usage = dict(self.parameter_usage) + self.embodiment_batches = {emb: 0 for emb in self.network.pose_dims} + self.current_embodiment = None + self.rng_observation, self.rng_cleanup = install_rng_observer( + pl_module, saved["mot_rng_state"], equal_state) + + def observe(module, args, outputs): + need(set(outputs) == set(self.network.active_streams(args[1]) if self.multimodal else 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 + if self.multimodal: + self.current_embodiment = args[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) + super().on_validation_start(trainer, pl_module) + + def on_validation_batch_end(self, *args, **kwargs): + if self.multimodal: + super().on_validation_batch_end(*args, **kwargs) + + def teardown(self, trainer, pl_module, stage): + if self.rng_cleanup is not None: + self.rng_cleanup() + + 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 + and pl_module.evaluator._validation_batches == 100 // self.world_size, + "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") + extra = {} + if self.multimodal: + from scripts.train.joint.native_recovery import conditional_record, expected_metrics, sample_both_embodiments + + need(set(metrics) == expected_metrics(self.network), "Missing native embodiment/stream validation metrics") + record = conditional_record(trainer.optimizers[0], pl_module.named_parameters(), + self.parameter_usage, self.expected_steps) + need(self.usage_callback.counts == self.parameter_usage, + "Maintained optimizer counters differ from observed resumed updates") + need(self.rng_observation["observed"], "Actual rank RNG restore was never observed") + sample_proof = sample_both_embodiments(pl_module, self.network, self.output) + extra = dict(multimodal=True, conditional_optimizer_record=record, + initial_parameter_usage=self.initial_parameter_usage, + embodiment_optimizer_batches=self.embodiment_batches, + active_streams_by_embodiment={e: list(self.network.active_streams(e)) for e in self.network.pose_dims}, + actual_rank_rng_restore=self.rng_observation, + hidden_label_samples=sample_proof, + private_bank_optimizer_steps_verified=True, human_action_absent=True) + else: + need(all(float(s["step"]) == self.expected_steps for s in trainer.optimizers[0].state.values()), + "Optimizer restarted instead of continuing") + verify_scheduler(self.lr_policy, [c.scheduler.state_dict() for c in trainer.lr_scheduler_configs], self.expected_steps) + 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)), + **extra, + )) + + +def phase_config(original, *, output, checkpoint, start_step, stop_step, service_output, world): + need((start_step, stop_step) in ((50, 52), (52, 56)) and world in (2, 4), + "This diagnostic proves exactly50->52->56; it is not production crash admission") + cfg = OmegaConf.create(copy.deepcopy(OmegaConf.to_container(original, resolve=True))) + original_model, original_data = copy.deepcopy(cfg.model), copy.deepcopy(cfg.data) + multimodal = bool(cfg.callbacks.joint_smoke.get("multimodal", False)) + 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, + multimodal=multimodal, + ) + 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" + if cfg.evaluator.get("pixel_diagnostics") is not None: + cfg.evaluator.pixel_diagnostics.artifact_dir = service_output + f"/validation-{stop_step}" + if multimodal: + need(any(c.get("_target_") == "egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage" + for c in cfg.callbacks.values()), "Missing maintained conditional optimizer callback") + 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") + from scripts.train.joint.native_subtasks import selected + extra = {} + if selected(original): + from scripts.train.joint.native_recovery import summarize_native_phases + extra = summarize_native_phases(phases) + 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", + **extra, + )) + + +def main(): + request_path = Path(sys.argv[1]) + q = json.loads(request_path.read_text()) + if q.get("schema") == "native-elmo-production-recovery-request-v1": + need(len(sys.argv) == 2, "Production recovery does not support legacy postflight repair") + from scripts.train.joint.production_recovery import main as production_recovery_main + return production_recovery_main(request_path) + 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") + from scripts.train.joint.native_subtasks import selected + if selected(original): + need(completion.get("schema") == "joint-multimodal-smoke-v1" + and completion.get("private_bank_optimizer_steps_verified") is True + and cfg.callbacks.joint_smoke.get("multimodal") is True, + "Native recovery requires the matching actual multimodal 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..627439aa7 100644 --- a/scripts/train/joint/run.py +++ b/scripts/train/joint/run.py @@ -11,12 +11,15 @@ 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, + verify_capacity_profile, write_once, ) @@ -36,6 +39,76 @@ def credentials(): os.environ["WANDB_API_KEY"] = values[2] +def validate_resolved_request(request, cfg): + """Check pure resolved science/mode guards before any allocation or construction. + + The runtime and metadata preflight must use this same gate. Configuration + group selection alone may leave experiment-inherited logger dictionaries. + """ + from omegaconf import OmegaConf + + need(request["schema"] == "joint-training-v2" and request["phase"] in ("smoke", "full"), + "Wrong training request") + smoke = request["phase"] == "smoke" + need(request["target_steps"] in ((50,) if smoke else (100000, 300000)), + "Wrong authorized training budget") + resolved = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + expected_data_task = verify_model_execution(request, resolved) + from scripts.train.joint import native_subtasks + if native_subtasks.selected(request): + need(cfg.data.get("_convert_") == "all", + "Native subtask data requires Hydra _convert_=all before construction") + world = request["world_size"] + need( + cfg.seed == 42 + and cfg.trainer.devices == cfg.data.distributed_world_size == world, + "Seed/topology mismatch", + ) + need( + cfg.data.batch_size == request["per_rank_batch"] + and cfg.data.batch_size * world == 32 + and cfg.model.optimizer.lr == validate(request["lr_policy"])["peak_lr"] + and cfg.model.optimizer.weight_decay == 0.01, + "Batch/optimizer mismatch", + ) + 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") + and not cfg.model.pipeline.get("trainability"), + "Fresh step0/trainable experts required", + ) + need( + ( + (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( + cfg.trainer.precision == "bf16-mixed" + and not cfg.model.enable_grad_norm + and cfg.runtime.disable_tf32, + "Precision/clipping policy changed", + ) + if smoke: + need( + cfg.trainer.max_steps == cfg.trainer.val_check_interval == 50 + and cfg.trainer.num_sanity_val_steps == 0, + "Scheduled optimization and validation required", + ) + need( + "wandb" not in cfg.logger + and (cfg.evaluator.pixel_diagnostics is None or cfg.evaluator.pixel_diagnostics.log_wandb is False), + "Smoke cannot write W&B", + ) + need("joint_smoke" in cfg.callbacks, "Missing all-rank smoke instrumentation") + return resolved + + def main(): path = Path(sys.argv[1]) request = json.loads(path.read_text()) @@ -45,14 +118,16 @@ def main(): and request["phase"] in ("smoke", "full"), "Wrong training request", ) + from scripts.train.joint import native_subtasks + native_subtask = native_subtasks.selected(request) smoke = request["phase"] == "smoke" 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", @@ -64,33 +139,73 @@ def main(): sha(request["data_manifest"]) == request["data_manifest_sha256"], "Changed data manifest", ) + if native_subtask: + native_subtasks.verify_dataset(request) preflight = None + start, stop = 0, request["target_steps"] 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 request.get("interruption") is not None: + from scripts.train.joint.interruption_admission import live_admission + current_recovery = live_admission(request, path) + write_once(path.with_suffix(".interruption-live-admission.json"), dict( + status="CURRENT_INTENT_CLAIMS_QUEUE_STORAGE_PASS", job_id=os.environ["SLURM_JOB_ID"], + request_sha256=sha(path), **current_recovery)) + check_storage(request, completion) + from scripts.train.joint.verify_temporal_pair import admission_mode + + mode = "native_subtask" if native_subtask else admission_mode(request) + if mode == "native_subtask": + native_subtasks.verify_dataset(request) + elif mode == "temporal_comparison": + from scripts.train.joint.verify_temporal_pair import admit as admit_temporal_pair + admit_temporal_pair(request) + elif mode == "pixel_recipe_comparison": + from scripts.train.joint.verify_temporal_pair import admit_track_recipe + admit_track_recipe(request, spatial="image") + elif mode == "track_recipe_comparison": + from scripts.train.joint.verify_temporal_pair import admit_track_recipe + admit_track_recipe(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) + elif mode == "capacity_comparison": + from scripts.train.joint.contract import verify_capacity_extension + verify_capacity_extension(request, completion) + elif mode in ("gap_comparison", "noise_comparison"): + from scripts.train.joint.contract import verify_gap_extension + verify_gap_extension(request) + 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 @@ -102,72 +217,16 @@ def main(): cfg = compose( config_name="train_zarr_cartesian", overrides=request["overrides"] ) - resolved = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + resolved = validate_resolved_request(request, cfg) world = request["world_size"] - need( - cfg.seed == 42 - and cfg.trainer.devices == cfg.data.distributed_world_size == world, - "Seed/topology mismatch", - ) - 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.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( - cfg.ckpt_path is None - and not cfg.model.pipeline.get("initialization") - and not cfg.model.pipeline.get("trainability"), - "Fresh step0/trainable experts required", - ) - need( - cfg.data.couple_state_observations - and not cfg.data.image_only - and cfg.data.task == "visual_state_action", - "Paired full data/noising contract changed", - ) - need( - cfg.trainer.precision == "bf16-mixed" - and not cfg.model.enable_grad_norm - and cfg.runtime.disable_tf32, - "Precision/clipping policy changed", - ) - if smoke: - need( - cfg.trainer.max_steps == cfg.trainer.val_check_interval == 50 - and cfg.trainer.num_sanity_val_steps == 0, - "Scheduled optimization and validation required", - ) - need( - "wandb" not in cfg.logger - and cfg.evaluator.pixel_diagnostics.log_wandb is False, - "Smoke cannot write W&B", - ) - need("joint_smoke" in cfg.callbacks, "Missing all-rank smoke instrumentation") - else: + if not smoke: smoke_cfg = OmegaConf.to_container( 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) with torch.device("meta"): network = instantiate( cfg.model.pipeline.stages[cfg.model.pipeline.stage_ids.joint_denoiser].model @@ -182,8 +241,9 @@ def main(): "Active expert count/width/trainability mismatch", ) need( - counts["image"] == 88178688 - and all(v <= 20000000 for k, v in counts.items() if k != "image"), + counts[request["active_streams"][0]] == verify_capacity_profile(request, resolved) + and all(v <= request["vector_expert_parameter_budget"] if native_subtask else v <= 20000000 + for k, v in counts.items() if k != request["active_streams"][0]), "Expert budget mismatch", ) del network @@ -191,24 +251,34 @@ def main(): context = data.prepare_context( mode="train", normalization="prepared", normalizer=None ) + from scripts.train.joint import boundary_uncertainty + boundary_uncertainty.context(resolved,context.snapshot(),data.state_dict()) + boundary_contract=request.get('boundary_sampling_contract') + if boundary_contract is not None: + census=bound(boundary_contract['eligibility']['report']) + need(census.get('data_context_sha256')==context.fingerprint(), + 'Actual training context differs from allocated boundary census') if preflight is not None: need( context.fingerprint() == preflight["provenance"]["data_context_sha256"], "Context changed after smoke", ) + from scripts.train.joint.dataset_contract import verify_dataset + expected_split, expected_validation = verify_dataset(request) need( - {k: len(v.episodes) for k, v in data.datasets.items()} - == {"train": 450, "val": 50} - and len(data.val_dataloader()) == 100, + {k: len(v.episodes) for k, v in data.datasets.items()} == expected_split + and len(data.val_dataloader()) == expected_validation, "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"]) + fresh_output = not start or request.get("interruption") is not None need( - str(cfg.paths.output_dir) == str(output) and not output.exists(), + str(cfg.paths.output_dir) == str(output) + and output.exists() == (not fresh_output) + and (fresh_output or (output / ".training-admitted").is_file()), "Output collision/mismatch", ) provenance = dict( @@ -217,6 +287,10 @@ def main(): data_manifest_sha256=request["data_manifest_sha256"], data_context_sha256=context.fingerprint(), ) + if boundary_contract is not None: + provenance.update(boundary_policy_sha256=boundary_uncertainty.digest(boundary_uncertainty.POLICY), + boundary_objective_sha256=boundary_uncertainty.digest(boundary_uncertainty.OBJECTIVE), + boundary_eligibility_report_sha256=boundary_contract['eligibility']['report']['sha256']) if not smoke: provenance["run_id"] = request["run_id"] if allocation is None: @@ -246,18 +320,36 @@ def main(): ) with path.with_suffix(".resolved.yaml").open("x") as stream: stream.write(OmegaConf.to_yaml(cfg, resolve=True)) + from scripts.train.joint.service_bootstrap import command if smoke: - 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 native_subtask: + # The same source-owned verifier owns actual H100/Flash4 maxshape. + # A40/pilot or historical two-camera evidence never substitutes. + bootstrap = root / "runtime_bootstrap.py" + subprocess.run([ + sys.executable, str(bootstrap), "-m", "torch.distributed.run", + "--standalone", "--nnodes=1", "--nproc-per-node=" + str(world), + "--max-restarts=0", + *command(request, path, bootstrap, + [str(source / "scripts/world_model/verify_joint_smoke.py"), + str(output), "--capacity-config", str(path.with_suffix(".resolved.yaml")), + "--capacity-output", str(path.with_suffix(".capacity.json")), + "--capacity-source-manifest-sha256", request["source_manifest_sha256"], + "--capacity-data-manifest-sha256", request["data_manifest_sha256"]])[1:]], + check=True, cwd=source) + else: + from scripts.world_model.check_joint_capacity import check_capacity + check_capacity(cfg, data, path.with_suffix(".capacity.json")) + if fresh_output: + 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 +360,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 +380,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" - ) - arguments = [str(output), "--world-size", str(world)] if smoke else [str(path)] + 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), *(["--multimodal"] if native_subtask else [])] 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..e0a6793d6 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,17 @@ 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") + from scripts.train.joint.native_subtasks import selected + native_subtask = selected(request) streams = request["active_streams"] count = target // 20000 import torch @@ -27,6 +34,9 @@ def main(): from egomimic.pipeline.construction import checkpoint_construction from egomimic.pl_utils.pl_model import ModelWrapper + from scripts.train.joint.production_history import has_interruption, retained_milestones, scheduled_validation + interrupted_history = has_interruption(request) + retained_history = validation_history = None milestones = {} directory = output / "checkpoints" for path in directory.glob("*.ckpt"): @@ -40,18 +50,32 @@ 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", - ) + if interrupted_history: + milestones, retained_history = retained_milestones(request, target) + 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 + + 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,22 +98,21 @@ 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) + if native_subtask: + from egomimic.pl_utils.joint_smoke import check_conditional_adam_state + check_conditional_adam_state(optimizers[0], checkpoint["native_conditional_optimizer_usage"], target) for state in optimizers[0]["state"].values(): need( - float(state["step"]) == target + (native_subtask or float(state["step"]) == target) and all( state[k].dtype == torch.float32 and torch.isfinite(state[k]).all() for k in ("exp_avg", "exp_avg_sq") @@ -141,55 +164,65 @@ 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())) - sample = batch["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}) + expected_embodiments = set(request["interfaces"]["pose_dims"]) if native_subtask else {"libero_gripper"} + sampled = set() 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" - ] - 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(), - "Invalid terminal sample", - ) - need( - torch.equal(value[stream["observed"]], stream["value"][stream["observed"]]), - "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 + for batch in data.val_dataloader(): + for embodiment in set(batch) & expected_embodiments - sampled: + sample = batch[embodiment] + for stream in sample["world"].values(): + stream["value"] = torch.where(stream["observed"][..., None], stream["value"], float("nan")) + groups = wrapper.model.process_batch_for_training({embodiment: sample}) + torch.manual_seed(42) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + predictions = wrapper.model.forward_eval(groups)[embodiment][ + getattr(cfg.evaluator, "prediction_key", "image_proprio_prediction")] + for name, value in predictions.items(): + stream = groups[embodiment]["world"][name] + need(value.dtype == torch.float32 and value.shape == stream["value"].shape + and torch.isfinite(value).all(), "Invalid terminal sample") + need(torch.equal(value[stream["observed"]], stream["value"][stream["observed"]]), + "Observed values changed") + sampled.add(embodiment) + if sampled == expected_embodiments: + break + need(sampled == expected_embodiments, "Missing terminal sample embodiment") + if interrupted_history: + metric_names = {f"val/{emb}/{kind}/{('action_command_mse' if stream == 'action' else stream + '_normalized_mse')}" + for emb in expected_embodiments for kind in ("forecast", "completion") + for stream in streams if stream != "action" or emb in request["interfaces"]["action_dims"]} + validation, validation_history = scheduled_validation(request, target, metric_names) + else: + validation = {} + from scripts.train.joint.continuation import log_segments + + for csv_path in 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 any(k.startswith("val/" + e + "/") for e in expected_embodiments) 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) need( len(validation) == count - and all(len(v) == 2 * len(streams) for v in validation.values()) + and all(len(v) == (sum(2 * (len(streams) - (0 if e in request["interfaces"]["action_dims"] else 1)) + for e in expected_embodiments) if native_subtask else 2 * len(streams)) + for v in validation.values()) and all( step in validation or step - 1 in validation for step in range(20000, target + 1, 20000) @@ -197,21 +230,26 @@ 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, + **(dict(retained_attempt_history=retained_history, committed_validation_history=validation_history) + if interrupted_history else {}), ), ) diff --git a/scripts/train/joint/verify_temporal_pair.py b/scripts/train/joint/verify_temporal_pair.py new file mode 100644 index 000000000..f2909d263 --- /dev/null +++ b/scripts/train/joint/verify_temporal_pair.py @@ -0,0 +1,426 @@ +"""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", + "capacity_comparison", "factorial_controls", "track_recipe_comparison", "pixel_recipe_comparison", "gap_comparison", "noise_comparison") +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() + + +def admit_track_recipe(request, *, spatial="track"): + """Bind independent D/C track recipes to both native smokes and their controls.""" + need(spatial in ("image", "track"), "Unknown spatial recipe") + mode = "pixel_recipe_comparison" if spatial == "image" else "track_recipe_comparison" + schema = "native-pixel-recipe-pair-v1" if spatial == "image" else "native-track-recipe-pair-v1" + need(admission_mode(request) == mode, "Wrong native recipe admission mode") + selection = request[mode] + need(selection.get("schema") == schema, "Unknown native recipe pair") + pairs = [bound(evidence_binding(item)) for item in selection["prototypes"]] + need(len(pairs) == 2 and {p["arm"] for p in pairs} == {"discrete", "continuous"}, + "Both independent track arms required") + current = copy.deepcopy(request) + current.pop(mode) + need(sum(current == p for p in pairs) == 1, "Full request differs from selected prototype") + outputs, normalized, initial_images = [], [], [] + for item in pairs: + need(item["phase"] == "full" and item["target_steps"] == 300000 + and item["world_size"] == 4 and item["per_rank_batch"] == 8 + and item["global_batch"] == 32 + and item["active_streams"] == [spatial, "proprio", "action"] + and item["lr_policy"] == dict(peak_lr=1e-4, schedule="constant", warmup_optimizer_steps=0), + "Track full recipe changed") + completion, _ = verify_evidence(item) + cfg = compose_request(item) + from omegaconf import OmegaConf + smoke_cfg = OmegaConf.to_container(OmegaConf.create(bound(item["smoke"]["resolved"], False)), resolve=True) + compare_configs(cfg, smoke_cfg, item) + verify_model_execution(item, cfg) + data, arm = copy.deepcopy(cfg["data"]), item["arm"] + gap = [dict(start_step=step, max_gap_seconds=seconds) + for step, seconds in ((0, .10), (50000, .15), (100000, .20), (200000, .30))] + need(data["mode"] == arm and data["min_frames"] == 2 and data["max_frames"] == 64 + and (arm != "discrete" or (data["discrete_frames"] == 16 and data["stride"] == 4)) + and data["sampling_strategy"] == ("discrete_window" if arm == "discrete" else "full_episode_uniform") + and data["validation_sampling_strategy"] == ("discrete_window" if spatial == "image" or arm == "discrete" else "full_episode_uniform") + and data.get("completion_gap_schedule") == (gap if spatial == "track" and arm == "continuous" else None) + and cfg["seed"] == data["seed"] == 42 and cfg["ckpt_path"] is None, + "Track timeline, seed or completion-only gap changed") + if spatial == "image": + initial_images.extend(r["initial_image_weights_sha256"] for r in completion["ranks"]) + batch_key = "paired_track_batch_sha256" if spatial == "track" else "paired_rgb_batch_sha256" + need(all(len(r[batch_key]) == 3 for r in completion["ranks"]), "Actual corruption hashes missing") + if spatial == "image": + from scripts.train.joint.dataset_contract import verify_dataset + counts, validation = verify_dataset(item) + need(item.get("dataset_contract") is not None + and data["validation_anchors_per_episode"] == item["dataset_contract"]["validation_anchors_per_episode"] + and validation == 100 and counts == item["dataset_contract"]["splits"], "Unbound pixel dataset") + normalized.append(track_recipe_controls(item, cfg, spatial=spatial)) + smoke_request = bound(item["smoke"]["request"]) + need(smoke_request["intended_full_run_id"] == item["run_id"] + and smoke_request["intended_full_output"] == item["output"], "Track smoke belongs to another run") + outputs.extend((item["output"], item["local_logs"], smoke_request["output"], smoke_request["local_logs"])) + for key in ("source_revision", "source_manifest_sha256", "source_archive_sha256", "data_manifest_sha256", + "model_execution", "operational_service"): + need(pairs[0].get(key) == pairs[1].get(key), "Native recipe shared control changed: " + key) + if spatial == "image": + for key in ("dataset_contract", "data_cache", "lambda_runtime"): + need(pairs[0].get(key) == pairs[1].get(key), "Pixel recipe shared control changed: " + key) + need(pairs[0]["smoke"]["job_id"] != pairs[1]["smoke"]["job_id"] + and pairs[0]["run_id"] != pairs[1]["run_id"], "Track pair shares experiment identity") + need(normalized[0] == normalized[1], "Undeclared track recipe difference") + disjoint(outputs) + if spatial == "image": + need(len(initial_images) == 8 and len(set(initial_images)) == 1, "Actual initial image weights differ") + return dict(schema=schema, status="PASS", fairness="FAIR WITH CAVEATS", + claim="Discrete versus continuous native image/state/action recipes on the same prepared corpus", + caveats=["Equal optimizer updates/global batch; temporal exposure, compute and vector initialization RNG differ.", + "Validation anchors do not increase the number of independent held-out episodes."]) + return dict(schema="native-track-recipe-pair-v1", status="PASS", fairness="FAIR WITH CAVEATS", + claim="Independent discrete versus continuous completion-gap track recipes", + caveats=["Equal updates/global batch; temporal exposure, sampling, compute and initialization RNG differ.", + "Completion-only continuous gap is intentional; this is not a gap-only causal ablation."]) + + +def track_recipe_controls(item, cfg, *, spatial="track"): + """Normalize only declared temporal fields and output-derived operational locations.""" + data, model = copy.deepcopy(cfg["data"]), copy.deepcopy(cfg["model"]) + for key in ("mode", "sampling_strategy", "validation_sampling_strategy", "validation_discrete_frames", + "validation_stride", "completion_gap_schedule", "discrete_frames", "stride"): + data.pop(key, None) + arm = item["arm"] + for stage in model["pipeline"]["stages"]: + if "model" in stage and stage["model"].get("_target_", "").endswith("JointImageProprioModel"): + core = stage["model"]["image_model"] + need(stage["model"]["pose_dims"] == {"libero_gripper": 36 if arm == "discrete" else 9} + and stage["model"]["action_dims"] == {"libero_gripper": 28 if arm == "discrete" else 7} + and core["time_mode"] == arm + and core["max_tokens"] == (16 if arm == "discrete" else 64) * (1122 if spatial == "track" else 202), + "Track model temporal dimensions changed") + stage["model"]["pose_dims"] = "" + stage["model"]["action_dims"] = "" + core["time_mode"] = "" + core["max_tokens"] = "" + def locations(value): + if isinstance(value, dict): return {k: locations(v) for k,v in value.items()} + if isinstance(value, list): return [locations(v) for v in value] + if isinstance(value, str): + for key in ("output", "local_logs"): + prefix = item[key] + if value == prefix or value.startswith(prefix + "/"): + return "<" + key + ">" + value[len(prefix):] + return value + return locations(dict(data=data, model=model, runtime=cfg["runtime"], evaluator=cfg["evaluator"], + normalizer=cfg["normalizer"], norm_stats=cfg["norm_stats"])) diff --git a/scripts/world_model/benchmark_pixel_execution.py b/scripts/world_model/benchmark_pixel_execution.py index faa25938f..d809df63a 100644 --- a/scripts/world_model/benchmark_pixel_execution.py +++ b/scripts/world_model/benchmark_pixel_execution.py @@ -53,9 +53,7 @@ def require(condition, message): def write_once(path, value): with Path(path).open("x") as stream: - stream.write( - json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n" - ) + stream.write(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") def digest(path): @@ -72,78 +70,38 @@ def frame_schedule(scenario, batch_size): if scenario != "mixed": raise ValueError(f"Unknown scenario: {scenario}") cycle = len(MIXED_FRAMES) // math.gcd(batch_size, len(MIXED_FRAMES)) - return [ - tuple( - MIXED_FRAMES[(step * batch_size + row) % len(MIXED_FRAMES)] - for row in range(batch_size) - ) - for step in range(cycle) - ] + return [tuple(MIXED_FRAMES[(step * batch_size + row) % len(MIXED_FRAMES)] + for row in range(batch_size)) for step in range(cycle)] def complete_cycles(steps, schedule): return math.ceil(steps / len(schedule)) * len(schedule) -def model_arguments( - backend, - scenario, - *, - compile_mode="none", - compile_dynamic=False, - token_execution="padded", -): - if ( - backend not in BACKENDS - or scenario not in SCENARIOS - or token_execution not in TOKEN_EXECUTIONS - ): +def model_arguments(backend, scenario, *, compile_mode="none", compile_dynamic=False, + token_execution="padded"): + if backend not in BACKENDS or scenario not in SCENARIOS or token_execution not in TOKEN_EXECUTIONS: raise ValueError("Unknown execution candidate") return dict( - objective="visual_only", - phase_a_condition_on_proprio=False, - pose_dims=None, - action_dims=None, - image_dim=768, - width=768, - depth=8, - heads=6, - mlp_ratio=4, - max_cameras=2, + objective="visual_only", phase_a_condition_on_proprio=False, + pose_dims=None, action_dims=None, image_dim=768, width=768, + depth=8, heads=6, mlp_ratio=4, max_cameras=2, max_tokens=3200 if scenario == "fixed16" else 25600, time_mode="discrete" if scenario == "fixed16" else "continuous", - use_camera_embedding=True, - use_modality_embedding=False, - use_embodiment_embedding=False, - allow_patch_compression=False, - use_phase_embedding=False, - noise_conditioning="adaln_zero", - position_encoding="spatiotemporal_rope", - rope_axis_dims=(64, 32, 32), + use_camera_embedding=True, use_modality_embedding=False, + use_embodiment_embedding=False, allow_patch_compression=False, + use_phase_embedding=False, noise_conditioning="adaln_zero", + position_encoding="spatiotemporal_rope", rope_axis_dims=(64, 32, 32), spatial_grid_size=(10, 10), - attention_backend=backend, - attention_dtype="bfloat16", - strict_fp32=False, - attention_query_chunk_size=0, - activation_checkpointing=False, - loss_reduction="mean_per_example", - compile_mode=compile_mode, - compile_dynamic=compile_dynamic, - compile_backend="inductor", - token_execution=token_execution, + attention_backend=backend, attention_dtype="bfloat16", strict_fp32=False, + attention_query_chunk_size=0, activation_checkpointing=False, + loss_reduction="mean_per_example", compile_mode=compile_mode, + compile_dynamic=compile_dynamic, compile_backend="inductor", token_execution=token_execution, ) -def make_batch( - frames, - *, - seed, - device, - image_size=160, - patch_size=16, - cameras=2, - frame_seconds=0.05, -): +def make_batch(frames, *, seed, device, image_size=160, patch_size=16, cameras=2, + frame_seconds=0.05): """True RGB patch geometry; only batch padding is invalid.""" if not frames or any(type(length) is not int or length < 2 for length in frames): raise ValueError("Every synthetic episode needs at least two frames") @@ -166,49 +124,32 @@ def make_batch( size = length * per_frame value[row, :size] = torch.rand(size, dimension, generator=generator) * 2 - 1 valid[row, :size] = True - observed[row, : max(1, length // 2) * per_frame] = True + observed[row, :max(1, length // 2) * per_frame] = True indices = torch.arange(length, dtype=torch.float32).repeat_interleave(per_frame) index[row, :size] = indices clock[row, :size] = indices * frame_seconds - return { - "image": { - key: tensor.to(device) - for key, tensor in dict( - value=value, - valid=valid, - observed=observed, - time=clock, - index=index, - camera=camera[None].expand(len(frames), -1).clone(), - xy=xy[None].expand(len(frames), -1, -1).clone(), - ).items() - } - } + return {"image": {key: tensor.to(device) for key, tensor in dict( + value=value, valid=valid, observed=observed, time=clock, index=index, + camera=camera[None].expand(len(frames), -1).clone(), + xy=xy[None].expand(len(frames), -1, -1).clone(), + ).items()}} def single_example(batch, row): indices = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] - return { - "image": { - key: value[row : row + 1].index_select(1, indices) - for key, value in batch["image"].items() - } - } + return {"image": {key: value[row:row + 1].index_select(1, indices) + for key, value in batch["image"].items()}} def fixed_corruption(batch, seed): image = batch["image"] generator = torch.Generator(device=image["value"].device).manual_seed(seed) - noise = torch.randn( - image["value"].shape, - dtype=torch.float32, - device=image["value"].device, - generator=generator, - ) + noise = torch.randn(image["value"].shape, dtype=torch.float32, + device=image["value"].device, generator=generator) levels = (0.15 + 0.7 * torch.remainder(image["index"], 7) / 6).masked_fill( image["observed"] | ~image["valid"], 0 ) - noisy = (1 - levels[..., None]) * image["value"] + levels[..., None] * noise + noisy = ((1 - levels[..., None]) * image["value"] + levels[..., None] * noise) return noisy.masked_fill(~image["valid"][..., None], 0), levels @@ -228,7 +169,6 @@ def execution_context(backend, device): with torch.autocast(device.type, dtype=torch.bfloat16): if backend == "sdpa": from torch.nn.attention import SDPBackend, sdpa_kernel - with sdpa_kernel(SDPBackend.FLASH_ATTENTION): yield else: @@ -237,11 +177,9 @@ def execution_context(backend, device): def relative_error(actual, expected): error = (actual.float() - expected.float()).double() - return dict( - squared_error=float(error.square().sum()), - squared_reference=float(expected.double().square().sum()), - max_abs=float(error.abs().max()), - ) + return dict(squared_error=float(error.square().sum()), + squared_reference=float(expected.double().square().sum()), + max_abs=float(error.abs().max())) @contextmanager @@ -255,18 +193,12 @@ def nontrivial_conditioning_probe(model, seed): if getattr(model, "noise_conditioning", "additive") != "adaln_zero": yield {"nonzero_adaln_probe": False} return - selected = { - name: parameter - for name, parameter in model.named_parameters() - if name.startswith(("image_out.", "final_adaln.")) or ".adaln_image." in name - } - require( - selected and any("adaln_image" in name for name in selected), - "AdaLN correctness probe could not find modulation parameters", - ) - original = { - name: parameter.detach().clone() for name, parameter in selected.items() - } + selected = {name: parameter for name, parameter in model.named_parameters() + if name.startswith(("image_out.", "final_adaln.")) + or ".adaln_image." in name} + require(selected and any("adaln_image" in name for name in selected), + "AdaLN correctness probe could not find modulation parameters") + original = {name: parameter.detach().clone() for name, parameter in selected.items()} device = next(model.parameters()).device devices = [device.index] if device.type == "cuda" else [] try: @@ -277,11 +209,8 @@ def nontrivial_conditioning_probe(model, seed): torch.cuda.default_generators[device.index].manual_seed(seed) for parameter in selected.values(): parameter.normal_(mean=0.0, std=0.02) - yield { - "nonzero_adaln_probe": True, - "seed": seed, - "restores_training_initialization": True, - } + yield {"nonzero_adaln_probe": True, "seed": seed, + "restores_training_initialization": True} finally: with torch.no_grad(): for name, parameter in selected.items(): @@ -297,21 +226,13 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): with execution_context(backend, device): prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] loss = example_mean_loss(prediction, batch) - require( - loss.dtype == torch.float32 and torch.isfinite(loss).item(), - "Nonfinite/non-FP32 probe loss", - ) + require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), "Nonfinite/non-FP32 probe loss") actual = prediction.detach() batched_gradients = None if gradients: loss.backward() - require( - all(p.grad is not None for p in model.parameters()), - "Unused trainable parameter", - ) - batched_gradients = { - name: p.grad.detach().cpu().clone() for name, p in model.named_parameters() - } + require(all(p.grad is not None for p in model.parameters()), "Unused trainable parameter") + batched_gradients = {name: p.grad.detach().cpu().clone() for name, p in model.named_parameters()} del prediction, loss model.zero_grad(set_to_none=True) errors = [] @@ -320,59 +241,36 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): selected = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] independent = single_example(batch, row) with execution_context(backend, device): - predicted = model( - independent, - "libero_gripper", - noisy[row : row + 1].index_select(1, selected), - levels[row : row + 1].index_select(1, selected), - phase="A", - )["image"] + predicted = model(independent, "libero_gripper", + noisy[row:row + 1].index_select(1, selected), + levels[row:row + 1].index_select(1, selected), phase="A")["image"] individual_loss = example_mean_loss(predicted, independent) / batch_size - errors.append( - relative_error( - actual[row : row + 1].index_select(1, selected), predicted.detach() - ) - ) + errors.append(relative_error(actual[row:row + 1].index_select(1, selected), predicted.detach())) if gradients: individual_loss.backward() del predicted, individual_loss - output_relative = math.sqrt( - sum(v["squared_error"] for v in errors) - / max(sum(v["squared_reference"] for v in errors), 1e-30) - ) - require( - math.isfinite(output_relative) and output_relative <= tolerance, - f"Batched output differs from independent B1: relative L2={output_relative}", - ) - report = { - "examples": batch_size, - "output_relative_l2": output_relative, - "output_max_abs": max(v["max_abs"] for v in errors), - } + output_relative = math.sqrt(sum(v["squared_error"] for v in errors) / + max(sum(v["squared_reference"] for v in errors), 1e-30)) + require(math.isfinite(output_relative) and output_relative <= tolerance, + f"Batched output differs from independent B1: relative L2={output_relative}") + report = {"examples": batch_size, "output_relative_l2": output_relative, + "output_max_abs": max(v["max_abs"] for v in errors)} if gradients: error_sum, reference_sum = 0.0, 0.0 for name, parameter in model.named_parameters(): - require( - parameter.grad is not None and parameter.grad.dtype == torch.float32, - "Missing/non-FP32 reference gradient", - ) - measured = relative_error( - batched_gradients[name], parameter.grad.detach().cpu() - ) + require(parameter.grad is not None and parameter.grad.dtype == torch.float32, + "Missing/non-FP32 reference gradient") + measured = relative_error(batched_gradients[name], parameter.grad.detach().cpu()) error_sum += measured["squared_error"] reference_sum += measured["squared_reference"] gradient_relative = math.sqrt(error_sum / max(reference_sum, 1e-30)) - require( - math.isfinite(gradient_relative) and gradient_relative <= tolerance, - f"Batched gradient differs from mean independent B1: relative L2={gradient_relative}", - ) + require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Batched gradient differs from mean independent B1: relative L2={gradient_relative}") report["gradient_relative_l2"] = gradient_relative model.zero_grad(set_to_none=True) # Perturb every other example, preserving row zero and its noise exactly. if batch_size > 1: - changed = { - "image": {key: value.clone() for key, value in batch["image"].items()} - } + changed = {"image": {key: value.clone() for key, value in batch["image"].items()}} changed["image"]["value"][1:] *= -0.7 perturbed = noisy.clone() perturbed[1:] *= -0.7 @@ -380,29 +278,17 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): # Inductor may select different kernels for grad-enabled and # no-grad execution. Compare intervention/control under the same # mode and graph; separately bound the mode's BF16 rounding drift. - control = model(batch, "libero_gripper", noisy, levels, phase="A")["image"][ - 0 - ].clone() - independent = model( - changed, "libero_gripper", perturbed, levels, phase="A" - )["image"][0] + control = model(batch, "libero_gripper", noisy, levels, phase="A")["image"][0].clone() + independent = model(changed, "libero_gripper", perturbed, levels, phase="A")["image"][0] mode_error = relative_error(control, actual[0]) - mode_relative = math.sqrt( - mode_error["squared_error"] / max(mode_error["squared_reference"], 1e-30) - ) - require( - math.isfinite(mode_relative) and mode_relative <= tolerance, - f"Grad/no-grad output differs beyond BF16 tolerance: relative L2={mode_relative}", - ) + mode_relative = math.sqrt(mode_error["squared_error"] / max(mode_error["squared_reference"], 1e-30)) + require(math.isfinite(mode_relative) and mode_relative <= tolerance, + f"Grad/no-grad output differs beyond BF16 tolerance: relative L2={mode_relative}") report["grad_vs_no_grad_relative_l2"] = mode_relative error = relative_error(independent, control) - leak = math.sqrt( - error["squared_error"] / max(error["squared_reference"], 1e-30) - ) - require( - math.isfinite(leak) and leak <= 1e-6, - f"Cross-example attention leakage detected: same-mode relative L2={leak}", - ) + leak = math.sqrt(error["squared_error"] / max(error["squared_reference"], 1e-30)) + require(math.isfinite(leak) and leak <= 1e-6, + f"Cross-example attention leakage detected: same-mode relative L2={leak}") report["cross_example_relative_l2"] = leak else: report["cross_example_check"] = "not_applicable_batch_one" @@ -411,102 +297,55 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): def check_eager_compiled(eager, compiled, batch, *, backend, seed, tolerance): """Require state-identical eager/compiled AMP forward and FP32 backward parity.""" - require( - eager.token_execution == compiled.token_execution, - "Eager/compiled token execution differs", - ) + require(eager.token_execution == compiled.token_execution, + "Eager/compiled token execution differs") reference_state, candidate_state = eager.state_dict(), compiled.state_dict() - require( - reference_state.keys() == candidate_state.keys(), - "Eager/compiled state keys differ", - ) - require( - all( - torch.equal(value, candidate_state[name]) - for name, value in reference_state.items() - ), - "Eager/compiled initial state differs", - ) - require( - all( - p.dtype == torch.float32 - for model in (eager, compiled) - for p in model.parameters() - ), - "Eager/compiled stored parameters must be FP32", - ) + require(reference_state.keys() == candidate_state.keys(), "Eager/compiled state keys differ") + require(all(torch.equal(value, candidate_state[name]) for name, value in reference_state.items()), + "Eager/compiled initial state differs") + require(all(p.dtype == torch.float32 for model in (eager, compiled) for p in model.parameters()), + "Eager/compiled stored parameters must be FP32") noisy, levels = fixed_corruption(batch, seed) device = batch["image"]["value"].device def forward_backward(model): model.zero_grad(set_to_none=True) with execution_context(backend, device): - prediction = model(batch, "libero_gripper", noisy, levels, phase="A")[ - "image" - ] + prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] loss = example_mean_loss(prediction, batch) - require( - loss.dtype == torch.float32 and torch.isfinite(loss).item(), - "Eager/compiled loss must be finite FP32", - ) + require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), + "Eager/compiled loss must be finite FP32") # Match standard mixed precision: backward is outside forward autocast. loss.backward() - require( - all( - p.grad is not None - and p.grad.dtype == torch.float32 - and torch.isfinite(p.grad).all() - for p in model.parameters() - ), - "Eager/compiled gradient must be present and finite FP32", - ) + require(all(p.grad is not None and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() for p in model.parameters()), + "Eager/compiled gradient must be present and finite FP32") return prediction.detach().cpu(), float(loss.detach()) try: reference_output, reference_loss = forward_backward(eager) - reference_gradients = { - name: p.grad.detach().cpu().clone() for name, p in eager.named_parameters() - } + reference_gradients = {name: p.grad.detach().cpu().clone() for name, p in eager.named_parameters()} eager.zero_grad(set_to_none=True) output, loss = forward_backward(compiled) measured = relative_error(output, reference_output) - output_relative = math.sqrt( - measured["squared_error"] / max(measured["squared_reference"], 1e-30) - ) + output_relative = math.sqrt(measured["squared_error"] / max(measured["squared_reference"], 1e-30)) loss_relative = abs(loss - reference_loss) / max(abs(reference_loss), 1e-30) - require( - math.isfinite(output_relative) and output_relative <= tolerance, - f"Compiled output differs from eager: relative L2={output_relative}", - ) - require( - math.isfinite(loss_relative) and loss_relative <= tolerance, - f"Compiled loss differs from eager: relative error={loss_relative}", - ) - errors = [ - relative_error(p.grad.detach().cpu(), reference_gradients[name]) - for name, p in compiled.named_parameters() - ] - gradient_relative = math.sqrt( - sum(value["squared_error"] for value in errors) - / max(sum(value["squared_reference"] for value in errors), 1e-30) - ) - require( - math.isfinite(gradient_relative) and gradient_relative <= tolerance, - f"Compiled gradient differs from eager: relative L2={gradient_relative}", - ) - return { - "output_relative_l2": output_relative, - "output_max_abs": measured["max_abs"], - "loss_relative_error": loss_relative, - "eager_loss": reference_loss, - "compiled_loss": loss, - "gradient_relative_l2": gradient_relative, - "parameters_and_gradients_dtype": "float32", - "backward_autocast": "off", - "same_state_verified": True, - "backend": backend, - "token_execution": compiled.token_execution, - } + require(math.isfinite(output_relative) and output_relative <= tolerance, + f"Compiled output differs from eager: relative L2={output_relative}") + require(math.isfinite(loss_relative) and loss_relative <= tolerance, + f"Compiled loss differs from eager: relative error={loss_relative}") + errors = [relative_error(p.grad.detach().cpu(), reference_gradients[name]) + for name, p in compiled.named_parameters()] + gradient_relative = math.sqrt(sum(value["squared_error"] for value in errors) / + max(sum(value["squared_reference"] for value in errors), 1e-30)) + require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Compiled gradient differs from eager: relative L2={gradient_relative}") + return {"output_relative_l2": output_relative, "output_max_abs": measured["max_abs"], + "loss_relative_error": loss_relative, "eager_loss": reference_loss, + "compiled_loss": loss, "gradient_relative_l2": gradient_relative, + "parameters_and_gradients_dtype": "float32", "backward_autocast": "off", + "same_state_verified": True, "backend": backend, + "token_execution": compiled.token_execution} finally: eager.zero_grad(set_to_none=True) compiled.zero_grad(set_to_none=True) @@ -514,11 +353,7 @@ def forward_backward(model): def compiler_counters(): from torch._dynamo.utils import counters - - return { - group: {str(k): int(v) for k, v in values.items()} - for group, values in counters.items() - } + return {group: {str(k): int(v) for k, v in values.items()} for group, values in counters.items()} def counter_total(counters, group, key=None): @@ -527,11 +362,9 @@ def counter_total(counters, group, key=None): def backend_inventory(backend): - module_names = { - "flash4": ("flash_attn.cute.interface",), - "flash3": ("flash_attn_3.flash_attn_interface", "flash_attn_interface"), - "sdpa": ("torch",), - }[backend] + module_names = {"flash4": ("flash_attn.cute.interface",), + "flash3": ("flash_attn_3.flash_attn_interface", "flash_attn_interface"), + "sdpa": ("torch",)}[backend] try: for index, module_name in enumerate(module_names): try: @@ -540,45 +373,27 @@ def backend_inventory(backend): except ModuleNotFoundError as error: # Match the model's official FA3 package/legacy alias support; # an installed backend with a missing dependency must fail. - if index == len(module_names) - 1 or error.name not in { - module_name, - module_name.split(".")[0], - }: + if (index == len(module_names) - 1 + or error.name not in {module_name, module_name.split(".")[0]}): raise except ImportError as error: - raise UnsupportedBackend( - f"Required backend import failed: {module_name}: {error}" - ) from error + raise UnsupportedBackend(f"Required backend import failed: {module_name}: {error}") from error module_path = Path(module.__file__) versions = {} - for name in ( - "torch", - "flash-attn-4", - "flash-attn-3", - "flash-attn", - "nvidia-cutlass-dsl", - ): + for name in ("torch", "flash-attn-4", "flash-attn-3", "flash-attn", "nvidia-cutlass-dsl"): try: versions[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: versions[name] = None - return { - "module": module_name, - "path": str(module_path), - "sha256": digest(module_path), - "versions": versions, - "sdpa_kernel": "FLASH_ATTENTION_ONLY" if backend == "sdpa" else None, - } + return {"module": module_name, "path": str(module_path), "sha256": digest(module_path), + "versions": versions, "sdpa_kernel": "FLASH_ATTENTION_ONLY" if backend == "sdpa" else None} def timed_step(model, batch, optimizer, backend, device): optimizer.zero_grad(set_to_none=True) with execution_context(backend, device): losses = model.training_losses(batch, "libero_gripper", pose_weight=0.0) - require( - losses["total"].dtype == losses["image"].dtype == torch.float32, - "MSE must remain FP32", - ) + require(losses["total"].dtype == losses["image"].dtype == torch.float32, "MSE must remain FP32") losses["total"].backward() optimizer.step() return losses["total"].detach() @@ -588,131 +403,73 @@ def run_candidate(request, destination): from egomimic.models.image_proprio import ImageProprioModel if not torch.cuda.is_available(): - raise UnsupportedBackend( - "A scheduled H100 CUDA device is required; no CPU benchmark fallback" - ) + raise UnsupportedBackend("A scheduled H100 CUDA device is required; no CPU benchmark fallback") device = torch.device("cuda", request["device"]) torch.cuda.set_device(device) properties = torch.cuda.get_device_properties(device) if "H100" not in properties.name: - raise UnsupportedBackend( - f"This benchmark binds H100 hardware, got {properties.name}" - ) + raise UnsupportedBackend(f"This benchmark binds H100 hardware, got {properties.name}") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision("highest") torch.manual_seed(request["seed"]) inventory = backend_inventory(request["backend"]) - arguments = model_arguments( - request["backend"], - request["scenario"], - compile_mode=request["compile_mode"], - compile_dynamic=request["compile_dynamic"], - token_execution=request.get("token_execution", "padded"), - ) + arguments = model_arguments(request["backend"], request["scenario"], + compile_mode=request["compile_mode"], compile_dynamic=request["compile_dynamic"], + token_execution=request.get("token_execution", "padded")) model = ImageProprioModel(**arguments).to(device).train() - require( - sum(p.numel() for p in model.parameters()) == PARAMETERS, - "Exact architecture parameter count changed", - ) - require( - all(p.dtype == torch.float32 for p in model.parameters()), - "Stored parameters must be FP32", - ) + require(sum(p.numel() for p in model.parameters()) == PARAMETERS, "Exact architecture parameter count changed") + require(all(p.dtype == torch.float32 for p in model.parameters()), "Stored parameters must be FP32") before_counters = compiler_counters() schedule = frame_schedule(request["scenario"], request["batch_size"]) - batches = [ - make_batch( - frames, - seed=request["seed"] + index, - device=device, - frame_seconds=0.2 if request["scenario"] == "fixed16" else 0.05, - ) - for index, frames in enumerate(schedule) - ] + batches = [make_batch(frames, seed=request["seed"] + index, device=device, + frame_seconds=0.2 if request["scenario"] == "fixed16" else 0.05) + for index, frames in enumerate(schedule)] with nontrivial_conditioning_probe(model, request["seed"] + 1000) as probe_record: - probe = make_batch( - (2, 3) if request["batch_size"] > 1 else (2,), - seed=request["seed"], - device=device, - ) - correctness = { - "bounded_gradient_probe": check_outputs( - model, - probe, - backend=request["backend"], - seed=request["seed"], - tolerance=request["tolerance"], - gradients=True, - ) - } + probe = make_batch((2, 3) if request["batch_size"] > 1 else (2,), seed=request["seed"], device=device) + correctness = {"bounded_gradient_probe": check_outputs( + model, probe, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"], gradients=True)} if request["compile_mode"] != "none": # This is a separate eager model, not another invocation of the same # compiled callable. Keep synthetic noise RNG unchanged by construction. with torch.random.fork_rng(devices=[device.index]): - eager = ( - ImageProprioModel(**{**arguments, "compile_mode": "none"}) - .to(device) - .train() - ) + eager = ImageProprioModel(**{**arguments, "compile_mode": "none"}).to(device).train() eager.load_state_dict(model.state_dict(), strict=True) correctness["compiled_vs_eager"] = check_eager_compiled( - eager, - model, - probe, - backend=request["backend"], - seed=request["seed"], - tolerance=request["tolerance"], - ) + eager, model, probe, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"]) del eager del probe correctness["conditioning_probe"] = probe_record correctness["full_shape_outputs"] = [] with torch.no_grad(): for batch in batches: - correctness["full_shape_outputs"].append( - check_outputs( - model, - batch, - backend=request["backend"], - seed=request["seed"], - tolerance=request["tolerance"], - gradients=False, - ) - ) + correctness["full_shape_outputs"].append(check_outputs( + model, batch, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"], gradients=False)) torch.cuda.empty_cache() # Release unused correctness-probe allocator blocks. optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01) warmup = complete_cycles(request["warmup"], schedule) steps = complete_cycles(request["steps"], schedule) warm_start = time.perf_counter() for step in range(warmup): - timed_step( - model, batches[step % len(batches)], optimizer, request["backend"], device - ) + timed_step(model, batches[step % len(batches)], optimizer, request["backend"], device) torch.cuda.synchronize(device) warmup_seconds = time.perf_counter() - warm_start warm_counters = compiler_counters() if request["compile_mode"] != "none": - require( - counter_total(warm_counters, "stats", "unique_graphs") - > counter_total(before_counters, "stats", "unique_graphs"), - "Requested compilation captured no tensor regions", - ) - require( - counter_total(warm_counters, "graph_break") - == counter_total(before_counters, "graph_break"), - "A compiled tensor region broke its fullgraph contract", - ) + require(counter_total(warm_counters, "stats", "unique_graphs") > counter_total(before_counters, "stats", "unique_graphs"), + "Requested compilation captured no tensor regions") + require(counter_total(warm_counters, "graph_break") == counter_total(before_counters, "graph_break"), + "A compiled tensor region broke its fullgraph contract") torch.cuda.reset_peak_memory_stats(device) walls, cuda_times, loss_values = [], [], [] valid_tokens = 0 for step in range(steps): batch = batches[step % len(batches)] torch.cuda.synchronize(device) - start_event, end_event = ( - torch.cuda.Event(enable_timing=True), - torch.cuda.Event(enable_timing=True), - ) + start_event, end_event = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) start = time.perf_counter() start_event.record() loss = timed_step(model, batch, optimizer, request["backend"], device) @@ -723,92 +480,41 @@ def run_candidate(request, destination): loss_values.append(float(loss)) valid_tokens += sum(schedule[step % len(schedule)]) * 200 after_counters = compiler_counters() - require( - all(math.isfinite(value) for value in loss_values), - "Nonfinite measured training loss", - ) - require( - counter_total(after_counters, "stats", "unique_graphs") - == counter_total(warm_counters, "stats", "unique_graphs"), - "New compiled graphs appeared during steady-state timing", - ) - require( - all( - p.grad is not None - and p.grad.dtype == torch.float32 - and torch.isfinite(p.grad).all() - for p in model.parameters() - ), - "Missing, nonfinite or non-FP32 measured gradient", - ) - require( - all( - state[key].dtype == torch.float32 and torch.isfinite(state[key]).all() - for state in optimizer.state.values() - for key in ("exp_avg", "exp_avg_sq") - ), - "Nonfinite or non-FP32 AdamW moments", - ) - require( - all( - p.dtype == torch.float32 and torch.isfinite(p).all() - for p in model.parameters() - ), - "Nonfinite or non-FP32 optimized parameters", - ) + require(all(math.isfinite(value) for value in loss_values), "Nonfinite measured training loss") + require(counter_total(after_counters, "stats", "unique_graphs") == counter_total(warm_counters, "stats", "unique_graphs"), + "New compiled graphs appeared during steady-state timing") + require(all(p.grad is not None and p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() + for p in model.parameters()), "Missing, nonfinite or non-FP32 measured gradient") + require(all(state[key].dtype == torch.float32 and torch.isfinite(state[key]).all() + for state in optimizer.state.values() for key in ("exp_avg", "exp_avg_sq")), + "Nonfinite or non-FP32 AdamW moments") + require(all(p.dtype == torch.float32 and torch.isfinite(p).all() for p in model.parameters()), + "Nonfinite or non-FP32 optimized parameters") result = { - "status": "PASS", - "candidate": request, - "parameters_total_trainable": PARAMETERS, - "model_arguments": arguments, - "backend_inventory": inventory, - "gpu": { - "name": properties.name, - "uuid": str(properties.uuid), - "total_memory_bytes": properties.total_memory, - }, - "correctness": correctness, - "frame_schedule": schedule, - "warmup_steps": warmup, - "timed_steps": steps, - "warmup_seconds": warmup_seconds, - "wall_seconds": walls, - "cuda_event_seconds": cuda_times, - "median_wall_seconds": statistics.median(walls), - "median_cuda_seconds": statistics.median(cuda_times), + "status": "PASS", "candidate": request, "parameters_total_trainable": PARAMETERS, + "model_arguments": arguments, "backend_inventory": inventory, + "gpu": {"name": properties.name, "uuid": str(properties.uuid), "total_memory_bytes": properties.total_memory}, + "correctness": correctness, "frame_schedule": schedule, + "warmup_steps": warmup, "timed_steps": steps, "warmup_seconds": warmup_seconds, + "wall_seconds": walls, "cuda_event_seconds": cuda_times, + "median_wall_seconds": statistics.median(walls), "median_cuda_seconds": statistics.median(cuda_times), "examples_per_second": steps * request["batch_size"] / sum(walls), - "valid_tokens_per_second": valid_tokens / sum(walls), - "valid_tokens_measured": valid_tokens, + "valid_tokens_per_second": valid_tokens / sum(walls), "valid_tokens_measured": valid_tokens, "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), - "resident_input_bytes": sum( - value.numel() * value.element_size() - for batch in batches - for value in batch["image"].values() - ), - "losses": loss_values, - "compiler_counters_before": before_counters, - "compiler_counters_after_warmup": warm_counters, - "compiler_counters_after_timing": after_counters, + "resident_input_bytes": sum(value.numel() * value.element_size() + for batch in batches for value in batch["image"].values()), + "losses": loss_values, "compiler_counters_before": before_counters, + "compiler_counters_after_warmup": warm_counters, "compiler_counters_after_timing": after_counters, "compile_scope": "declared_fullgraph_tensor_regions; attention_adapter_and_validation_eager", - "compiled_backward_autocast": "off" - if request["compile_mode"] != "none" - else None, + "compiled_backward_autocast": "off" if request["compile_mode"] != "none" else None, "measurement_scope": "single_gpu_device_resident_synthetic_model; no_loader_H2D_NCCL_or_checkpoint_IO", } if request["profile"]: trace = destination / "trace.json" - with torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - record_shapes=True, - profile_memory=True, - ) as profiler: - with torch.profiler.record_function( - "synthetic_device_resident_optimizer_step" - ): + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, profile_memory=True) as profiler: + with torch.profiler.record_function("synthetic_device_resident_optimizer_step"): timed_step(model, batches[0], optimizer, request["backend"], device) torch.cuda.synchronize(device) profiler.export_chrome_trace(str(trace)) @@ -826,29 +532,15 @@ def worker(request_path): except UnsupportedBackend as error: result = {"status": "UNSUPPORTED", "error": str(error), "candidate": request} except CorrectnessError as error: - result = { - "status": "CORRECTNESS_FAILED", - "error": str(error), - "candidate": request, - } + result = {"status": "CORRECTNESS_FAILED", "error": str(error), "candidate": request} except Exception as error: # Backend/API/compiler failures are never retried with another kernel. - unsupported = ( - isinstance(error, RuntimeError) - and any( - phrase in str(error) - for phrase in ("No available kernel", "No viable backend") - ) - ) or ( - isinstance(error, ValueError) - and str(error) == "Ragged packed_valid execution requires flash3 or flash4" - ) - result = { - "status": "UNSUPPORTED" if unsupported else "ERROR", - "error": f"{type(error).__name__}: {error}", - "traceback": traceback.format_exc(), - "candidate": request, - } + unsupported = (isinstance(error, RuntimeError) and any( + phrase in str(error) for phrase in ("No available kernel", "No viable backend") + )) or (isinstance(error, ValueError) and + str(error) == "Ragged packed_valid execution requires flash3 or flash4") + result = {"status": "UNSUPPORTED" if unsupported else "ERROR", "error": f"{type(error).__name__}: {error}", + "traceback": traceback.format_exc(), "candidate": request} write_once(request_path.parent / "result.json", result) @@ -861,13 +553,8 @@ def csv_choices(value, choices): def run_worker_process(command, folder, timeout): """Bound compiler subprocesses as well as the GPU-owning worker on timeout.""" - with ( - (folder / "stdout.log").open("x") as out, - (folder / "stderr.log").open("x") as err, - ): - process = subprocess.Popen( - command, stdout=out, stderr=err, start_new_session=True - ) + with (folder / "stdout.log").open("x") as out, (folder / "stderr.log").open("x") as err: + process = subprocess.Popen(command, stdout=out, stderr=err, start_new_session=True) try: return process.wait(timeout=timeout), False except subprocess.TimeoutExpired: @@ -882,57 +569,30 @@ def run_worker_process(command, folder, timeout): def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path) - parser.add_argument( - "--backends", type=lambda s: csv_choices(s, BACKENDS), default=list(BACKENDS) - ) - parser.add_argument( - "--scenarios", type=lambda s: csv_choices(s, SCENARIOS), default=list(SCENARIOS) - ) - parser.add_argument( - "--batch-sizes", - type=lambda s: [int(v) for v in csv_choices(s, ("1", "2", "4", "8"))], - default=[1, 2, 4, 8], - ) + parser.add_argument("--backends", type=lambda s: csv_choices(s, BACKENDS), default=list(BACKENDS)) + parser.add_argument("--scenarios", type=lambda s: csv_choices(s, SCENARIOS), default=list(SCENARIOS)) + parser.add_argument("--batch-sizes", type=lambda s: [int(v) for v in csv_choices(s, ("1", "2", "4", "8"))], default=[1, 2, 4, 8]) parser.add_argument("--compile-mode", choices=COMPILE_MODES, default="none") parser.add_argument("--compile-dynamic", action="store_true") - parser.add_argument( - "--token-execution", - choices=TOKEN_EXECUTIONS, - default="padded", - help="Keep padded tensors, or pack valid image tokens across the image-only expert", - ) + parser.add_argument("--token-execution", choices=TOKEN_EXECUTIONS, default="padded", + help="Keep padded tensors, or pack valid image tokens across the image-only expert") parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--steps", type=int, default=10) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--device", type=int, default=0) parser.add_argument("--tolerance", type=float, default=0.03) parser.add_argument("--worker-timeout", type=int, default=1800) - parser.add_argument( - "--python-bootstrap", - type=Path, - help="Run every child through this reviewed runtime bootstrap (.pth hooks included)", - ) - parser.add_argument( - "--profile-candidate", help="Exact candidate ID such as flash4-mixed-b4-none" - ) - parser.add_argument( - "--dry-run", - action="store_true", - help="Print the immutable candidate plan without CUDA work", - ) + parser.add_argument("--python-bootstrap", type=Path, + help="Run every child through this reviewed runtime bootstrap (.pth hooks included)") + parser.add_argument("--profile-candidate", help="Exact candidate ID such as flash4-mixed-b4-none") + parser.add_argument("--dry-run", action="store_true", help="Print the immutable candidate plan without CUDA work") parser.add_argument("--worker-request", type=Path, help=argparse.SUPPRESS) args = parser.parse_args() if args.worker_request: worker(args.worker_request) return - if ( - min(args.warmup, args.steps, args.worker_timeout) < 1 - or args.device < 0 - or not 0 < args.tolerance < 1 - ): - parser.error( - "Positive bounded warmup/steps/timeout/tolerance and nonnegative device required" - ) + if min(args.warmup, args.steps, args.worker_timeout) < 1 or args.device < 0 or not 0 < args.tolerance < 1: + parser.error("Positive bounded warmup/steps/timeout/tolerance and nonnegative device required") candidates = [] for backend in args.backends: for scenario in args.scenarios: @@ -940,32 +600,16 @@ def main(): identity = f"{backend}-{scenario}-b{batch_size}-{args.compile_mode}" if args.token_execution != "padded": identity += f"-{args.token_execution}" - candidates.append( - dict( - id=identity, - backend=backend, - scenario=scenario, - batch_size=batch_size, - compile_mode=args.compile_mode, - compile_dynamic=args.compile_dynamic, - token_execution=args.token_execution, - seed=args.seed, - device=args.device, - warmup=args.warmup, - steps=args.steps, - tolerance=args.tolerance, - profile=args.profile_candidate == identity, - ) - ) + candidates.append(dict(id=identity, backend=backend, scenario=scenario, + batch_size=batch_size, compile_mode=args.compile_mode, + compile_dynamic=args.compile_dynamic, token_execution=args.token_execution, + seed=args.seed, + device=args.device, warmup=args.warmup, steps=args.steps, + tolerance=args.tolerance, profile=args.profile_candidate == identity)) if args.profile_candidate and not any(c["profile"] for c in candidates): parser.error("--profile-candidate must identify an enabled candidate") if args.dry_run: - print( - json.dumps( - {"parameters_total_trainable": PARAMETERS, "candidates": candidates}, - indent=2, - ) - ) + print(json.dumps({"parameters_total_trainable": PARAMETERS, "candidates": candidates}, indent=2)) return if args.output is None: parser.error("--output is required") @@ -974,31 +618,15 @@ def main(): bootstrap_path = args.python_bootstrap.resolve(strict=True) bootstrap = {"path": str(bootstrap_path), "sha256": digest(bootstrap_path)} args.output.mkdir(parents=True, exist_ok=False) - source_files = [ - Path(__file__), - ROOT / "egomimic/models/image_proprio.py", - ROOT / "egomimic/models/pixel_diffusion.py", - ROOT / "egomimic/models/world_model.py", - ROOT / "egomimic/models/attention_execution.py", - ROOT / "egomimic/models/compilation.py", - ] - source = { - str(path.relative_to(ROOT)): digest(path) - for path in source_files - if path.exists() - } - write_once( - args.output / "plan.json", - { - "schema": "pixel-execution-benchmark-v1", - "candidates": candidates, - "source_sha256": source, - "slurm_job_id": os.environ.get("SLURM_JOB_ID"), - "python_bootstrap": bootstrap, - "python_executable": sys.executable, - "scope": "synthetic_single_gpu_execution_only; not_training_readiness_or_policy_quality", - }, - ) + source_files = [Path(__file__), ROOT / "egomimic/models/image_proprio.py", + ROOT / "egomimic/models/pixel_diffusion.py", + ROOT / "egomimic/models/world_model.py", ROOT / "egomimic/models/attention_execution.py", + ROOT / "egomimic/models/compilation.py"] + source = {str(path.relative_to(ROOT)): digest(path) for path in source_files if path.exists()} + write_once(args.output / "plan.json", {"schema": "pixel-execution-benchmark-v1", "candidates": candidates, + "source_sha256": source, "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "python_bootstrap": bootstrap, "python_executable": sys.executable, + "scope": "synthetic_single_gpu_execution_only; not_training_readiness_or_policy_quality"}) results = [] for candidate in candidates: folder = args.output / candidate["id"] @@ -1008,41 +636,19 @@ def main(): command = [sys.executable] if args.python_bootstrap: command.append(str(args.python_bootstrap.resolve())) - command += [ - str(Path(__file__).resolve()), - "--worker-request", - str(request_path.resolve()), - ] + command += [str(Path(__file__).resolve()), "--worker-request", str(request_path.resolve())] returncode, timed_out = run_worker_process(command, folder, args.worker_timeout) if timed_out: - result = { - "status": "TIMEOUT", - "seconds": args.worker_timeout, - "candidate": candidate, - } + result = {"status": "TIMEOUT", "seconds": args.worker_timeout, "candidate": candidate} elif returncode != 0 or not (folder / "result.json").exists(): - result = { - "status": "PROCESS_ERROR", - "returncode": returncode, - "candidate": candidate, - } + result = {"status": "PROCESS_ERROR", "returncode": returncode, "candidate": candidate} else: result = json.loads((folder / "result.json").read_text()) results.append(result) - print( - json.dumps({"candidate": candidate["id"], "status": result["status"]}), - flush=True, - ) - write_once( - args.output / "report.json", - { - "schema": "pixel-execution-benchmark-v1", - "source_sha256": source, - "status": "COMPLETED", - "results": results, - "scope": "Only PASS candidates have measured and checked synthetic throughput; real-data DDP smoke still required", - }, - ) + print(json.dumps({"candidate": candidate["id"], "status": result["status"]}), flush=True) + write_once(args.output / "report.json", {"schema": "pixel-execution-benchmark-v1", "source_sha256": source, + "status": "COMPLETED", "results": results, + "scope": "Only PASS candidates have measured and checked synthetic throughput; real-data DDP smoke still required"}) if __name__ == "__main__": diff --git a/scripts/world_model/check_joint_capacity.py b/scripts/world_model/check_joint_capacity.py index ada6a0202..231d70451 100644 --- a/scripts/world_model/check_joint_capacity.py +++ b/scripts/world_model/check_joint_capacity.py @@ -70,7 +70,8 @@ def check_capacity(cfg, data, destination): ) optimization_lr = optimizer.param_groups[0]["lr"] optimizer.step() - scheduler.step() + if scheduler is not None: + scheduler.step() torch.cuda.synchronize() elapsed = time.perf_counter() - start peak = torch.cuda.max_memory_allocated() diff --git a/scripts/world_model/verify_joint_smoke.py b/scripts/world_model/verify_joint_smoke.py index 1b8169383..5ff13946f 100644 --- a/scripts/world_model/verify_joint_smoke.py +++ b/scripts/world_model/verify_joint_smoke.py @@ -22,26 +22,64 @@ from egomimic.pipeline.construction import checkpoint_construction from egomimic.pl_utils.joint_smoke import ( load_smoke_configs, + _check_pilot_adam_usage, require, weights_digest, write_once, ) from egomimic.pl_utils.pl_model import ModelWrapper +from scripts.train.joint.lr_policy import from_config, update_lr, verify_scheduler +from scripts.train.joint import boundary_uncertainty def main(): parser = argparse.ArgumentParser() parser.add_argument("output", type=Path) - parser.add_argument("--world-size", type=int, choices=(2, 4), required=True) + parser.add_argument("--world-size", type=int, choices=(2, 4)) + parser.add_argument("--pilot-config", type=Path, help="One-GPU four-update TRAIN-pilot code validation; distinct from production smoke") + parser.add_argument("--pilot-postflight-only", action="store_true", help="Resume only verification of existing immutable step4 checkpoint; zero optimizer updates") + parser.add_argument("--multimodal", action="store_true", help="Native robot/human camera/language DDP smoke contract") + parser.add_argument("--capacity-config",type=Path,help="Distributed native C64/3camera/rankbatch maximum-shape backward; zero optimizer updates") + parser.add_argument("--capacity-output",type=Path,help="Immutable sibling capacity receipt; smoke output remains absent") + parser.add_argument("--capacity-source-manifest-sha256",help="Exact staged source manifest SHA256") + parser.add_argument("--capacity-data-manifest-sha256",help="Exact request-bound prepared data manifest SHA256") args = parser.parse_args() + if args.capacity_config is not None: + require(not args.multimodal and args.pilot_config is None and args.world_size is None, + 'Capacity mode is separate from optimizer smoke/pilot') + from egomimic.pl_utils.joint_smoke import validate_multimodal_capacity + require(args.capacity_output is not None and args.capacity_source_manifest_sha256 and args.capacity_data_manifest_sha256, + "Capacity requires exact independent output/source binding") + validate_multimodal_capacity(args.capacity_config,args.capacity_output, + source_manifest_sha256=args.capacity_source_manifest_sha256, + data_manifest_sha256=args.capacity_data_manifest_sha256) + return + require(args.capacity_output is None and args.capacity_source_manifest_sha256 is None + and args.capacity_data_manifest_sha256 is None, + 'Capacity-only bindings require --capacity-config') + if args.pilot_config is not None: + require(not args.multimodal, "Pilot cannot certify distributed multimodal smoke") + require(args.world_size is None, "Pilot mode owns its explicit single-GPU topology") + from egomimic.pl_utils.joint_smoke import validate_multimodal_pilot + if args.pilot_postflight_only: + validate_multimodal_pilot(args.pilot_config, args.output, postflight_only=True) + else: + validate_multimodal_pilot(args.pilot_config, args.output) + return + require(not args.pilot_postflight_only, "Postflight-only requires --pilot-config") + require(args.world_size in (2,4), "Historical smoke requires --world-size") root, world = args.output, args.world_size job = os.environ["SLURM_JOB_ID"] ranks = [ json.loads((root / f"fit-rank{rank}.json").read_text()) for rank in range(world) ] + cfg, model_config = load_smoke_configs(root) + policy = from_config(cfg.model) + peak, warmup = policy["peak_lr"], policy["warmup_optimizer_steps"] for rank, record in enumerate(ranks): require( - record["status"] == "FIT_PASS" + record["schema"] == ("joint-multimodal-fit-v1" if args.multimodal else "joint-mot-fit-v1") + and record["status"] == "FIT_PASS" and record["job_id"] == job and record["world_size"] == world and record["rank"] == rank, @@ -58,10 +96,11 @@ def main(): "Incomplete real optimization/validation proof", ) require( - record["lr_warmup_steps"] == 3000 + record["lr_warmup_steps"] == warmup + and record["lr_policy"] == policy and len(record["optimizer_update_lrs"]) == 50 and all( - math.isclose(lr, 5e-6 * (step + 1) / 3000, rel_tol=1e-12) + math.isclose(lr, update_lr(policy, step), rel_tol=1e-12) for step, lr in enumerate(record["optimizer_update_lrs"]) ), "Actual optimizer updates did not use the configured LR warmup", @@ -91,13 +130,7 @@ def main(): cfg, model_config = load_smoke_configs(root) require( checkpoint["global_step"] == 50 - and cfg.model.optimizer.lr == 5e-6 - and 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 - and cfg.data.couple_state_observations + and (cfg.data.image_only or cfg.data.couple_state_observations) and cfg.trainer.devices == world and cfg.data.batch_size * world == 32, "Checkpoint/config scientific contract mismatch", @@ -119,22 +152,17 @@ def main(): ) require( all( - math.isclose(g["lr"], 5e-6 * 51 / 3000, rel_tol=1e-12) - and g["initial_lr"] == 5e-6 + math.isclose(g["lr"], update_lr(policy, 50), rel_tol=1e-12) + and (not warmup or g["initial_lr"] == peak) for g in optimizers[0]["param_groups"] ), "Wrong next-update LR after50 completed updates", ) schedulers = checkpoint.get("lr_schedulers", []) - require( - len(schedulers) == 1 - and schedulers[0]["last_epoch"] == 50 - and schedulers[0]["base_lrs"] == [5e-6] - and schedulers[0]["lr_lambdas"] == [{"warmup_steps": 3000}], - "Missing exact LR scheduler continuation state", - ) + verify_scheduler(policy, schedulers, 50) for state in optimizers[0]["state"].values(): - require(float(state["step"]) == 50, "AdamW update count mismatch") + if not args.multimodal: + require(float(state["step"]) == 50, "AdamW update count mismatch") require( all( state[k].dtype == torch.float32 and torch.isfinite(state[k]).all() @@ -180,6 +208,15 @@ def main(): context.fingerprint() == ranks[0]["data_context_sha256"], "Restored data context changed", ) + boundary_uncertainty.context(cfg,checkpoint['data_context'],sampler) + boundary_proof=None + if args.multimodal: + boundary_proof=boundary_uncertainty.fit_rows(ranks, + selected=boundary_uncertainty.sampling_config(cfg),world=world, + per_rank_batch=cfg.data.batch_size,embodiments=cfg.joint.pose_dims, + context_sha256=context.fingerprint()) + else: + require(cfg.data.get('boundary_sampling') is None,'Boundary uncertainty requires actual multimodal smoke') validate_checkpoint_binding(checkpoint, model_config, context) with checkpoint_construction(): wrapper = ( @@ -193,15 +230,27 @@ def main(): restored_optimizer = restored_optimizer() restored_scheduler = instantiate(cfg.model.scheduler, optimizer=restored_optimizer) restored_optimizer.load_state_dict(optimizers[0]) - restored_scheduler.load_state_dict(schedulers[0]) + if restored_scheduler is not None: + restored_scheduler.load_state_dict(schedulers[0]) require( - restored_scheduler.state_dict() == schedulers[0] + (restored_scheduler.state_dict() == schedulers[0] if warmup else restored_scheduler is None) and all( - math.isclose(g["lr"], 5e-6 * 51 / 3000, rel_tol=1e-12) + math.isclose(g["lr"], update_lr(policy, 50), rel_tol=1e-12) for g in restored_optimizer.param_groups ), "Strict optimizer/scheduler restoration changed warmup state", ) + if args.multimodal: + usage=ranks[0]['parameter_update_counts'] + require(all(r['parameter_update_counts']==usage for r in ranks), + 'Conditional optimizer usage differs across ranks') + require(set(usage)==set(dict(wrapper.named_parameters())),'Optimizer parameter identity differs') + _check_pilot_adam_usage(restored_optimizer,wrapper.named_parameters(),usage) + require(all(r.get('private_bank_optimizer_steps_verified') and r.get('human_action_absent') + and all(r.get(key) for key in ('language_conditioning_verified','camera_targets_verified')) + and r.get('attention_backend')=='flash4' and r.get('token_execution')=='packed_valid' + and 'H100' in r.get('gpu_name','') and r.get('gpu_capability',[0])[0]>=9 for r in ranks), + 'Missing actual H100/multimodal optimizer proof') require(all(p.requires_grad for p in wrapper.parameters()), "An expert was frozen") network = wrapper.model.pipeline.stage_by_id("joint_denoiser").model require( @@ -221,32 +270,50 @@ def main(): sum(expert_parameters.values()) == sum(p.numel() for p in wrapper.parameters()), "Per-expert counts do not account for every trained parameter", ) - sample = torch.load(root / "sample-input.pt", map_location="cpu", weights_only=True) - groups = wrapper.model.process_batch_for_training({"libero_gripper": sample}) + if args.multimodal: + samples={emb:torch.load(root/f'sample-input-{emb}.pt',map_location='cpu',weights_only=False) + for emb in network.pose_dims} + if boundary_proof is not None:boundary_uncertainty.strict_validation_samples(samples) + require(all(r['active_streams_by_embodiment']=={e:list(network.active_streams(e)) for e in network.pose_dims} + for r in ranks),'Actual embodiment streams differ') + else: + samples={'libero_gripper':torch.load(root/'sample-input.pt',map_location='cpu',weights_only=True)} wrapper.model.device = next(wrapper.parameters()).device - torch.manual_seed(42) - torch.cuda.synchronize() - start = time.perf_counter() - with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): - output = wrapper.model.forward_eval(groups)["libero_gripper"][ - "image_proprio_prediction" - ] - torch.cuda.synchronize() - sample_seconds = time.perf_counter() - start - for name, value in output.items(): - source = groups["libero_gripper"]["world"][name] - require( - value.shape == source["value"].shape - and value.dtype == torch.float32 - and torch.isfinite(value).all(), - "Strictly reloaded sample is invalid", - ) - require( - torch.equal(value[source["observed"]], source["value"][source["observed"]]), - "Sampling changed observed state/commands", - ) + sample_seconds=0.0 + for embodiment,sample in samples.items(): + groups=wrapper.model.process_batch_for_training({embodiment:sample}) + torch.manual_seed(42) + torch.cuda.synchronize() + start=time.perf_counter() + with torch.no_grad(),torch.autocast('cuda',dtype=torch.bfloat16): + output=wrapper.model.forward_eval(groups)[embodiment]['image_proprio_prediction'] + torch.cuda.synchronize() + sample_seconds+=time.perf_counter()-start + if args.multimodal: + require(set(output)==set(network.active_streams(embodiment)), 'Unexpected native sampled streams') + require('language_features' in sample and 'pixel_metadata' in sample, + 'Missing actual subtask/language sample provenance') + for name,value in output.items(): + source=groups[embodiment]['world'][name] + require(value.shape==source['value'].shape and value.dtype==torch.float32 + and torch.isfinite(value).all(),'Strictly reloaded sample is invalid') + require(torch.equal(value[source['observed']],source['value'][source['observed']]), + 'Sampling changed observed state/commands') + sampler = wrapper.model.pipeline.stage_by_id("joint_sampler").sampler + sample_state_steps = sampler.state_steps + sample_action_steps = getattr(sampler, "action_steps", 0) if "action" in network.streams else 0 + require(sample_state_steps == 20 and ("action" not in network.streams or sample_action_steps == 20), + "Actual configured sampler steps differ from the smoke contract") + execution = None + if hasattr(network, "image_model"): + require(network.image_model.compile_mode == "default" and network.image_model.compile_dynamic is True + and network.image_model.compile_backend == "inductor", + "Native image compile policy differs from historical training") + execution = {"implementation": "native_image_joint", "compile_mode": "default", + "compile_dynamic": True, "compile_backend": "inductor"} receipt = dict( schema="joint-mot-smoke-v1", + model_execution=execution, status="PASS", job_id=job, world_size=world, @@ -258,15 +325,39 @@ def main(): strict_full_wrapper_reload=True, optimizer_sampler_all_rank_rng_verified=True, lr_scheduler_strictly_restored=True, - lr_warmup_steps=3000, + lr_warmup_steps=warmup, + lr_policy=policy, checkpoint_sha256=checkpoint_sha, all_rank_weights_agree=True, sample_finite=True, - sample_state_steps=20, - sample_action_steps=20, + sample_state_steps=sample_state_steps, + sample_action_steps=sample_action_steps, sample_seconds=sample_seconds, ranks=ranks, ) + if args.multimodal: + require(network.image_model.token_execution=='packed_valid' + and network.image_model.blocks[0].attention_backend=='flash4' + and cfg.evaluator.active_streams_from_data, + 'Production multimodal execution/evaluator contract differs') + require(len(network.pose_dims)==2 and len(network.action_dims)==1 + and network.separate_proprio_experts and network.camera_dims, + 'Missing native private human/robot expert contract') + receipt.update(schema='joint-multimodal-smoke-v1', + private_bank_optimizer_steps_verified=True,human_action_absent=True, + subtask_boundaries_verified=bool(cfg.data.subtask_manifest_sha256), + language_conditioning_verified=network.language_feature_dim==1024, + camera_targets_verified=bool(cfg.data.camera_pose), + active_streams_by_embodiment={e:list(network.active_streams(e)) for e in network.pose_dims}, + attention_backend='flash4',token_execution='packed_valid', + data_context_sha256=context.fingerprint()) + if boundary_proof is not None: + receipt.update(subtask_boundaries_verified=False, + original_subtask_annotation_identity_verified=True, + subtask_boundary_verification_policy=boundary_uncertainty.POLICY['policy'], + boundary_uncertainty=boundary_proof, + boundary_policy_sha256=boundary_uncertainty.digest(boundary_uncertainty.POLICY), + boundary_objective_sha256=boundary_uncertainty.digest(boundary_uncertainty.OBJECTIVE)) write_once(root / "smoke-completion.json", receipt) print( json.dumps({"status": "PASS", "receipt": str(root / "smoke-completion.json")}) diff --git a/tests/fixtures/actual-smoke-csv-provenance.json b/tests/fixtures/actual-smoke-csv-provenance.json new file mode 100644 index 000000000..2bef7bf03 --- /dev/null +++ b/tests/fixtures/actual-smoke-csv-provenance.json @@ -0,0 +1,7 @@ +{ + "source_path": "/Users/elmopa/Documents/Codex/2026-10-17/the-exact-code-used-by-all-2/work/elmo-large-continuous-20261026-v1/smoke-preparation-v5/verified-output-20261026T085105Z/files/output/elmo-large314m-continuous-plain-smoke-v5/log-mirror/csv/version_0/metrics.csv", + "source_sha256": "87ce76d6a64358d7c3602fea432fc2a104baa4946b09550f528683d4fc85834a", + "selection": "Header and actual rows whose step is 0,1,2,3; no value changes", + "fixture_sha256": "ef17d625a1dac31f6f81de930eb2b81b614ee6a46d53947456e97ffec02b28f5", + "actual_recovery_proof": false +} diff --git a/tests/fixtures/actual-smoke-first-four-updates.csv b/tests/fixtures/actual-smoke-first-four-updates.csv new file mode 100644 index 000000000..0cbc70b92 --- /dev/null +++ b/tests/fixtures/actual-smoke-first-four-updates.csv @@ -0,0 +1,9 @@ +Optimizer/param_group_0_lr,Timing/Compute_Losses_Sec,Timing/Forward_Pass_Sec,Timing/Process_Batch_Sec,Train/Loss,Train/joint_action_mse,Train/joint_action_mse/eva_bimanual,Train/joint_camera_mse,Train/joint_camera_mse/eva_bimanual,Train/joint_camera_mse/human_bimanual,Train/joint_image_mse,Train/joint_image_mse/eva_bimanual,Train/joint_image_mse/human_bimanual,Train/joint_proprio_mse,Train/joint_proprio_mse/eva_bimanual,Train/joint_proprio_mse/human_bimanual,Train/loss,Train/source_0_log_joint_action_mse,Train/source_0_log_joint_camera_mse,Train/source_0_log_joint_image_mse,Train/source_0_log_joint_proprio_mse,Train/source_0_loss,Train/source_0_loss_action,Train/source_0_loss_camera,Train/source_0_loss_image,Train/source_0_loss_proprio,epoch,lr-AdamW,step,val/eva_bimanual/completion/action_command_mse,val/eva_bimanual/completion/camera_normalized_mse,val/eva_bimanual/completion/image_normalized_mse,val/eva_bimanual/completion/proprio_normalized_mse,val/eva_bimanual/forecast/action_command_mse,val/eva_bimanual/forecast/camera_normalized_mse,val/eva_bimanual/forecast/image_normalized_mse,val/eva_bimanual/forecast/proprio_normalized_mse,val/human_bimanual/completion/camera_normalized_mse,val/human_bimanual/completion/image_normalized_mse,val/human_bimanual/completion/proprio_normalized_mse,val/human_bimanual/forecast/camera_normalized_mse,val/human_bimanual/forecast/image_normalized_mse,val/human_bimanual/forecast/proprio_normalized_mse +,,,,,,,,,,,,,,,,,,,,,,,,,,,1.6666666666666667e-09,0,,,,,,,,,,,,,, +,0.00021004676818847656,26.692527770996094,0.0013590455055236816,1.2624458074569702,,,0.2016090750694275,,0.2016090750694275,0.2683567404747009,,0.2683567404747009,0.7924798727035522,,0.7924798727035522,1.2624458074569702,,0.2016090750694275,0.2683567404747009,0.7924798727035522,1.2624458074569702,0.0,0.2016090750694275,0.2683567404747009,0.7924798727035522,0,,0,,,,,,,,,,,,,, +,,,,,,,,,,,,,,,,,,,,,,,,,,,3.3333333333333334e-09,1,,,,,,,,,,,,,, +,0.00016188621520996094,0.5318794250488281,0.0014432072639465332,1.4678915739059448,,,0.20105388760566711,,0.20105388760566711,0.2697792649269104,,0.2697792649269104,0.9970583915710449,,0.9970583915710449,1.4678915739059448,,0.20105388760566711,0.2697792649269104,0.9970583915710449,1.4678915739059448,-0.0,0.20105388760566711,0.2697792649269104,0.9970583915710449,0,,1,,,,,,,,,,,,,, +,,,,,,,,,,,,,,,,,,,,,,,,,,,5e-09,2,,,,,,,,,,,,,, +,0.00015032291412353516,10.984643936157227,0.0009689927101135254,1.6105961799621582,0.29422447085380554,0.29422447085380554,3.4238759138251113e-32,3.4238759138251113e-32,0.20105388760566711,0.20437857508659363,0.20437857508659363,0.2697792649269104,1.1119931936264038,1.1119931936264038,0.9970583915710449,1.6105961799621582,0.29422447085380554,3.4238759138251113e-32,0.20437857508659363,1.1119931936264038,1.6105961799621582,0.29422447085380554,3.4238759138251113e-32,0.20437857508659363,1.1119931936264038,0,,2,,,,,,,,,,,,,, +,,,,,,,,,,,,,,,,,,,,,,,,,,,6.666666666666667e-09,3,,,,,,,,,,,,,, +,0.00016325712203979492,1.674769639968872,0.0010260939598083496,1.4463703632354736,0.29750120639801025,0.29750120639801025,4.8066246123767534e-29,4.8066246123767534e-29,0.20105388760566711,0.20404383540153503,0.20404383540153503,0.2697792649269104,0.9448253512382507,0.9448253512382507,0.9970583915710449,1.4463703632354736,0.29750120639801025,4.8066246123767534e-29,0.20404383540153503,0.9448253512382507,1.4463703632354736,0.29750120639801025,4.8066246123767534e-29,0.20404383540153503,0.9448253512382507,0,,3,,,,,,,,,,,,,, diff --git a/tests/joint_phase_observer/proof_fixture.py b/tests/joint_phase_observer/proof_fixture.py new file mode 100644 index 000000000..c89c14d96 --- /dev/null +++ b/tests/joint_phase_observer/proof_fixture.py @@ -0,0 +1,138 @@ +"""Portable synthetic observer evidence builder. Never invokes native/ML code.""" +import copy +from datetime import datetime,timezone +import json +from pathlib import Path +import sys +from scripts.train.joint import phase_observer_proof as p + + +def put(path,value): + path.parent.mkdir(parents=True,exist_ok=True) + path.write_bytes(value if isinstance(value,bytes) else value.encode() if isinstance(value,str) else p.encoded(value)) + return p.binding(path) + + +def raw(value):return ' '.join(k+'='+str(v) for k,v in value.items())+'\n' + + +def query(argv,stdout):return dict(schema='native-elmo-phase-observer-query-v1',argv=argv,returncode=0,stdout=stdout,stderr='',observed_at_utc=datetime.now(timezone.utc).isoformat()) + + +def build(root,mode='inspect',arm='plain',running=True,terminal_state='COMPLETED'): + root=Path(root);root.mkdir() + for name in ['execution','execution-claims','templates','requests','scientific-claims','producer-claims','status','logs','observer-specs','observer-claims']:(root/name).mkdir() + runtime=['execution_spec.py','allocated.sbatch','host_guard.py','phase_observer.py','observe-phase.sbatch'] + for name in runtime:put(root/name,"raise AssertionError('EXTERNAL OPS CODE MUST NEVER EXECUTE')\n") + manifest=''.join(p.sha(root/name)+' '+name+'\n' for name in runtime);ops=put(root/'FILES.sha256',manifest) + service=Path(p.__file__).resolve().parents[3] + sb=dict(root=str(service),manifest=str(service/'FILES.sha256'),manifest_sha256=p.sha(service/'FILES.sha256')) + node='gpu03' if arm=='boundary' else 'gpu11';old='111';job='900001';observer='900002' + run='mot-elmo-pickplace-continuous-large314m-'+arm+'-robot-human-camera-subtask-s42-300k-20261026-v1' + parent=dict(schema='joint-training-v2',phase='full',kind='native-subtask-world-model-v1',run_id=run,job_name='synthetic-original', + target_steps=300000,parameters=314187064,world_size=4,output=str(root/'original-output'),local_logs=str(root/'original-logs'), + operational_service=sb) + for key in ['source','source_revision','source_manifest','source_manifest_sha256','data_manifest','data_manifest_sha256','lambda_runtime','data_cache']:parent[key]='synthetic-'+key + pb=put(root/'parent.json',parent) + profile=dict(parent_id=old,node=node,request_path=pb['path'],request_sha256=pb['sha256']) + schema,cpu,memory,limit,gres=p.MODES[mode] + name=('elmo-large314m-continuous-' if mode in {'full','diagnostic'} else 'elmo-native-recovery-')+arm+'-'+mode+'-synthetic' + native=dict(JobName=name,UserId='paphiwetsa3(41003)',Account='loaner',QOS='loaner-normal',Partition='rl2',ReqNodeList=node,NodeList=node, + NumCPUs=cpu,MinMemoryNode=memory,TimeLimit=limit,Requeue='0',Restarts='0',Command=str(root/'allocated.sbatch'),TresPerNode=gres) + if mode=='full': + template=copy.deepcopy(parent);template.update(job_name=name,output=str(root/'phase-output'),local_logs=str(root/'phase-logs'),interruption=dict(parent_request=pb,admission=None)) + else: + template=dict(schema=schema,parent_request=pb,operational_service=sb,producer_native=native,output=str(root/'phase-output')) + for key in ['source','source_revision','source_manifest','source_manifest_sha256','data_manifest','data_manifest_sha256','lambda_runtime','data_cache']:template[key]=parent[key] + if mode=='capture':template.update(parent_job_id=old,parent_native=dict(JobId=old)) + tb=put(root/'templates/phase.json',template) + spec=dict(schema='native-elmo-execution-spec-v1',root=str(root),mode=mode,arm=arm,original_parent_job_id=old,parent_request=pb, + request_template=tb,native_contract=native,operational_service=sb,ops_manifest=ops, + execution_claim_roots=[dict(path=str(root/'execution-claims'),terminal_directory=str(root/'status'))],release_observer_path=str(root/'observer-specs/observer.json')) + sp=put(root/'execution/phase.json',spec);ex=dict(path=sp['path'],sha256=sp['sha256'],mode=mode,ops_manifest_sha256=ops['sha256']) + sc=root/'execution-claims'/sp['sha256'];sc.mkdir();argv=p.producer_argv(spec,Path(sp['path']),sp['sha256'],ops['sha256']) + held=dict(native,JobId=job,JobState='PENDING',Reason='JobHeldUser',NodeList='',NumNodes='1-1',NumTasks='1', + SubmitTime='2026-01-01T01:00:00',StartTime='Unknown',AllocTRES='(null)',ReqTRES='cpu='+cpu+',mem='+memory+',node=1,billing='+cpu, + Comment='elmo-execution:'+sp['sha256']) + hb=put(sc/'HELD_NATIVE.txt',raw(held)) + actual=copy.deepcopy(template) + if mode!='full':actual['producer_native']['JobId']=job + ab=put(root/'requests'/(sp['sha256']+'-'+job+'.json'),actual) + claim=root/('scientific-claims' if mode in {'diagnostic','full'} else 'producer-claims')/ab['sha256'];claim.mkdir() + put(sc/'PREPARED.json',dict(spec_sha256=sp['sha256'],ops_manifest_sha256=ops['sha256'],spec_path=sp['path'],arm=arm,parent_request=pb,argv=argv,original_parent_job_id=old,created_at=datetime.now(timezone.utc).isoformat())) + put(claim/'PREPARED.json',dict(request_sha256=ab['sha256'],spec_sha256=sp['sha256'],spec_path=sp['path'],actual_request_path=ab['path'],argv=argv,original_parent_job_id=old)) + for directory in [sc,claim]:put(directory/'PARENT_JOB_ID',job+'\n') + put(sc/'sbatch.stdout',job+'\n');put(sc/'sbatch.exit','0\n');put(sc/'sbatch.stderr','') + acceptance=dict(schema='native-elmo-held-execution-binding-v1',spec_sha256=sp['sha256'],ops_manifest_sha256=ops['sha256'],job_id=job,argv=argv,request=ab,request_claim=str(claim),held_native=hb) + put(sc/'acceptance.json',acceptance) + ospec=dict(schema='native-elmo-phase-observer-spec-v1',root=str(root),producer_execution=ex,terminal_directory=str(root/'status')) + ob=put(root/'observer-specs/observer.json',ospec);ledger=root/'observer-claims'/sp['sha256'];ledger.mkdir() + ctx=dict(root=root,spec=spec,execution=ex,request=actual,job=job) + on=p.observer_native(ctx);oa=p.observer_argv(ctx,Path(ob['path']),ob['sha256']) + pending=dict(on,JobId=observer,JobState='PENDING',Reason='Dependency',NodeList='',NumNodes='1-1',NumTasks='1', + SubmitTime='2026-01-01T01:00:01',StartTime='Unknown',ReqTRES='cpu=2,mem=1G,node=1,billing=2',AllocTRES='(null)', + Dependency='afterany:'+job+'(unfulfilled)',Comment='elmo-phase-observer:'+ob['sha256'],WorkDir=str(root), + StdOut=str(root/'logs'/(on['JobName']+'-'+observer+'.out')),StdErr=str(root/'logs'/(on['JobName']+'-'+observer+'.err'))) + pending_binding=put(ledger/'PENDING_NATIVE.txt',raw(pending));pq=query(p.native_argv(observer),raw(pending)) + put(ledger/'ACCEPTANCE_QUERY.json',pq);put(ledger/'PRODUCER_HELD_QUERY.json',query(p.native_argv(job),raw(held))) + intent=dict(schema='native-elmo-phase-observer-intent-v1',observer_spec=ob,ops_manifest_sha256=ops['sha256'],producer_execution=ex,producer_request=ab,producer_job_id=job,argv=oa) + put(ledger/'PREPARED.json',intent) + for name,value in [('sbatch.stdout',observer+'\n'),('sbatch.exit','0\n'),('sbatch.stderr',''),('OBSERVER_JOB_ID',observer+'\n')]:put(ledger/name,value) + oaccept=dict(schema='native-elmo-phase-observer-acceptance-v1',observer_spec=ob,ops_manifest_sha256=ops['sha256'],producer_execution=ex,producer_request=ab,producer_job_id=job,observer_job_id=observer,argv=oa,pending_native=pending_binding) + oab=put(ledger/'acceptance.json',oaccept) + ready=dict(schema='native-elmo-phase-observer-ready-v1',producer_job_id=job,producer_request_sha256=ab['sha256'],observer_job_id=observer, + observer_spec_sha256=ob['sha256'],observer_pending_query=pq,observer_own_terminal='UNKNOWN_NOT_REQUIRED',production_completion_proven=False) + rh=put(sc/'RELEASE_HELD_NATIVE.txt',raw(held));put(sc/'RELEASE_ATTEMPT.json',dict(job_id=job,spec_sha256=sp['sha256'],request_sha256=ab['sha256'],held_native_sha256=rh['sha256'],observer_readiness=ready)) + own=dict(on,JobId=observer,JobState='RUNNING',Reason='None',NumNodes='1',NumTasks='1',SubmitTime=pending['SubmitTime'],StartTime='2026-01-01T02:00:00', + ReqTRES='cpu=2,mem=1G,node=1,billing=2',AllocTRES='cpu=2,mem=1G,node=1,billing=2',Comment=pending['Comment']) + ownb=put(ledger/'RUNNING_NATIVE_QUERY.json',query(p.native_argv(observer),raw(own))) + rn=dict(native,JobId=job,JobState='RUNNING',Reason='None',NumNodes='1',NumTasks='1',SubmitTime=held['SubmitTime'],StartTime='2026-01-01T01:00:02', + ReqTRES=held['ReqTRES'],AllocTRES=held['ReqTRES'],Comment=held['Comment']) + if gres:rn.update(JOB_GRES='gpu:h100:4',GRES='gpu:h100:4(IDX:0-3)') + rb=put(claim/'RUNNING_ADMISSION.json',dict(job_id=job,request_sha256=ab['sha256'],spec_sha256=sp['sha256'],native=raw(rn))) if running else None + term=dict(rn,JobState=terminal_state,EndTime='2026-01-01T02:00:00',ExitCode='0:0' if terminal_state=='COMPLETED' else '1:0') + if gres:term['GRES']='' + nb=put(root/'status'/('native-'+job+'.txt'),raw(term));qb=put(root/'status'/('native-'+job+'.query.json'),query(p.native_argv(job),raw(term))) + eb=put(root/'status'/('native-'+job+'.exit'),'0\n');errb=put(root/'status'/('native-'+job+'.stderr'),'') + stepb=put(root/'status'/('native-'+job+'.steps-query.json'),query(['/usr/local/bin/squeue','--steps','--noheader','--user','paphiwetsa3','--format=%i|%T'],'900999.batch|RUNNING\n')) + result=dict(schema='native-elmo-phase-observer-result-v1',status='GENUINE_PRODUCER_TERMINAL_CAPTURED',producer_execution=ex,producer_request=ab, + producer_job_id=job,producer_phase=mode,producer_terminal_state=terminal_state,producer_terminal_exit=term['ExitCode'],producer_native=nb,producer_query=qb, + producer_query_exit=eb,producer_query_stderr=errb,owned_steps=stepb,producer_running_admission=rb,producer_running_admission_status='VALIDATED' if running else 'NOT_RECORDED', + observer_spec=ob,observer_acceptance=oab,observer_running_query=ownb,observer_job_id=observer,observer_own_terminal='UNKNOWN_NOT_REQUIRED',observer_host_exit_is_native_proof=False, + phase_completion_proven=False,zero_work_proven=False,captured_at_utc=datetime.now(timezone.utc).isoformat(),automatic_recovery_authorized=False) + orb=put(root/'status'/('native-'+job+'.observer.json'),result) + result_mode={v:k for k,v in p.RESULT_MODES.items()};statuses=dict(capture='CAPTURED_REQUIRES_NATIVE_PRODUCER_COMPLETION',inspect='INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION',diagnostic='DIAGNOSTIC_COMPLETE_REQUIRES_HOST_AND_NATIVE_SUCCESS',finalize='FINALIZED_REQUIRES_HOST_AND_NATIVE_SUCCESS') + result_schema=result_mode.get(mode,'synthetic-full-census-only');status=statuses.get(mode,'NOT_PHASE_AUTHORITY') + filenames=dict(capture='interruption-capture.json',inspect='checkpoint-inspection.json',diagnostic='production-recovery-diagnostic.json',finalize='interruption-finalized.json',full='full-census-not-phase-result.json') + phase_result=put(Path(actual['output'])/filenames[mode],dict(schema=result_schema,status=status,job_id=job,request_sha256=ab['sha256'])) + host=put(root/'status'/(mode+'-'+job+'.exit'),'0\n') + bundle=dict(request=ab,result=phase_result,native=nb,host_exit=host,observer=orb) + return dict(root=str(root),profiles={run:profile},bundle=bundle,result_schema=result_schema,status=status,mode=mode,arm=arm, + spec=sp,observer_spec=ob,execution_claim=str(sc),request_claim=str(claim),observer_claim=str(ledger),ops_manifest=ops) + + +def build_original(root,arm='plain'): + root=Path(root);root.mkdir();job='111';observer='112';run='mot-elmo-pickplace-continuous-large314m-'+arm+'-robot-human-camera-subtask-s42-300k-20261026-v1' + parent=dict(run_id=run,job_name='synthetic-original-parent') + pb=put(root/'ops/requests/parent.json',parent) + members={} + for name in ['capture-native.sbatch','guard_full.py','submit_full_once.py','full.sbatch']: + members[name]=put(root/name,'synthetic immutable original '+name+'\n')['sha256'] + members['ops/requests/parent.json']=pb['sha256'];mb=put(root/'FILES.sha256',''.join(v+' '+k+'\n' for k,v in sorted(members.items()))) + claim=root/'claims'/pb['sha256'];claim.mkdir(parents=True) + put(claim/'PREPARED.json',dict(request_sha256=pb['sha256'],manifest_sha256=mb['sha256'])) + put(claim/'PARENT_JOB_ID',job+'\n');put(claim/'CAPTURE_JOB_ID',observer+'\n') + n=dict(JobId=job,JobName=parent['job_name'],UserId='paphiwetsa3(41003)',Account='loaner',QOS='loaner-normal',Partition='rl2',ReqNodeList='gpu11',NodeList='gpu11',NumCPUs='32',MinMemoryNode='192G',TimeLimit='7-00:00:00',Requeue='0',Restarts='0',Command=str(root/'full.sbatch'),TresPerNode='gres/gpu:h100:4') + rn=dict(n,JobState='RUNNING',NumNodes='1',NumTasks='1',StartTime='2026-01-01T01:00:00',EndTime='2026-01-08T01:00:00',ExitCode='0:0',ReqTRES='cpu=32,mem=192G,node=1',AllocTRES='cpu=32,mem=192G,node=1',JOB_GRES='gpu:h100:4',GRES='gpu:h100:4(IDX:0-3)') + put(claim/'RUNNING_ADMISSION.json',dict(status='ACTUAL_BOUND_FULL_NATIVE_ALLOCATION_PASS',job_id=job,request_sha256=pb['sha256'],native=raw(rn))) + term=dict(rn,JobState='TIMEOUT',EndTime='2026-01-01T02:00:00',ExitCode='0:15',GRES='') + put(root/('native-'+job+'.txt'),raw(term));put(root/('native-'+job+'.exit'),'0\n');put(root/('native-'+job+'.stderr'),'') + oname='synthetic-original-observer' + pending=dict(JobId=observer,JobName=oname,UserId=n['UserId'],Account='loaner',QOS='loaner-normal',Partition='rl2',ReqNodeList='gpu11',NumTasks='1',NumCPUs='2',TimeLimit='00:10:00',Requeue='0',Restarts='0',JobState='PENDING',Command=str(root/'capture-native.sbatch'),WorkDir=str(root),NumNodes='1-1',MinMemoryNode='1G',Dependency='afterany:'+job+'(unfulfilled)',StdOut=str(root/'logs'/(oname+'-'+observer+'.out')),StdErr=str(root/'logs'/(oname+'-'+observer+'.err')),ReqTRES='cpu=2,mem=1G,node=1',AllocTRES='(null)',TresPerTask='cpu:2') + put(claim/'CAPTURE_PENDING_NATIVE.txt',raw(pending)) + profile=dict(root=str(root),parent_id=job,observer_id=observer,parent_name=parent['job_name'],observer_name=oname,node='gpu11',request_path=pb['path'],request_sha256=pb['sha256'],files_sha256=mb['sha256'],capture_script_sha256=members['capture-native.sbatch'],submitter_sha256=members['submit_full_once.py'],has_smoke_identity=False) + return dict(root=str(root),terminal_directory=str(root),profiles={run:profile},request_claim=str(claim),bundle={}) + + +if __name__=='__main__': + print(json.dumps(build_original(sys.argv[1],sys.argv[3]) if sys.argv[2]=='original' else build(sys.argv[1],sys.argv[2],sys.argv[3],sys.argv[4]=='yes',sys.argv[5]))) diff --git a/tests/joint_phase_observer/test_phase_observer_proof.py b/tests/joint_phase_observer/test_phase_observer_proof.py new file mode 100644 index 000000000..1c8c3e6b1 --- /dev/null +++ b/tests/joint_phase_observer/test_phase_observer_proof.py @@ -0,0 +1,266 @@ +"""Portable synthetic CPU tests: no campaign-relative data, native or ML calls.""" +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys + +import pytest + +HERE=Path(__file__).resolve().parent +SOURCE=Path(os.environ.get('ELMO_TEST_SERVICE_ROOT',HERE.parent/'service')).resolve() + + +def sha(path):return hashlib.sha256(Path(path).read_bytes()).hexdigest() +def put(path,value): + path=Path(path);path.write_text(value if isinstance(value,str) else json.dumps(value,sort_keys=True,indent=2)+'\n') + return dict(path=str(path),sha256=sha(path)) +def read(path):return json.loads(Path(path).read_text()) + + +@pytest.fixture(scope='session') +def service(tmp_path_factory): + source=SOURCE + if not (source/'scripts/train/joint/phase_observer_proof.py').is_file(): + spec=importlib.util.find_spec('scripts.train.joint.phase_observer_proof') + assert spec and spec.origin;source=Path(spec.origin).resolve().parents[3] + loader=importlib.util.spec_from_file_location('portable_service_export',source/'scripts/train/joint/service_bootstrap.py') + module=importlib.util.module_from_spec(loader);loader.loader.exec_module(module) + root=tmp_path_factory.mktemp('portable-service') + for rel in module.service_files(source): + dst=root/rel;dst.parent.mkdir(parents=True,exist_ok=True);shutil.copyfile(source/rel,dst) + names=module.service_files(root) + (root/'FILES.sha256').write_text(''.join(sha(root/n)+' '+n+'\n' for n in sorted(names))) + return root + + +def environment(service):return dict(os.environ,PYTHONPATH=str(service),PYTHONDONTWRITEBYTECODE='1') + + +def fixture(tmp_path,service,mode='inspect',arm='plain',running=True,state='COMPLETED'): + run=subprocess.run([sys.executable,'-B',str(HERE/'proof_fixture.py'),str(tmp_path/'ops'),mode,arm,'yes' if running else 'no',state], + env=environment(service),cwd=service,capture_output=True,text=True,check=False) + assert run.returncode==0,run.stderr + return json.loads(run.stdout) + + +def execute(f,service,operation='verify',**extra): + payload=dict(f,operation=operation,**extra) + program='''import builtins,json,sys,subprocess +from scripts.train.joint import phase_observer_proof as p +f=json.loads(sys.argv[1]);p.ORIGINAL_OBSERVERS=f['profiles'] +from scripts.train.joint import original_observer +original_observer.ORIGINAL_OBSERVERS=f['profiles'] +original=builtins.__import__ +def guarded(name,*a,**k): + if name.split('.')[0] in {'torch','lightning','egomimic','numpy'}: raise AssertionError('ML import forbidden') + return original(name,*a,**k) +builtins.__import__=guarded +subprocess.run=lambda *a,**k:(_ for _ in ()).throw(AssertionError('Native side effect forbidden')) +b=f['bundle'] +if f['operation']=='manifest': + print(json.dumps(sorted(p.manifest(f['root'],f['ops_manifest'])))) +elif f['operation']=='claim': + v=p.verify_claim_terminal(f['request_claim'],f.get('terminal_directory',str(__import__('pathlib').Path(f['root'])/'status')),next(iter(f['profiles']))) + print(json.dumps(v)) +elif f['operation']=='quiescence': + from scripts.train.joint import interruption_admission as gate + root=__import__('pathlib').Path(f['root']);gate.command=lambda argv:'900999.batch|RUNNING'+chr(10) if '--steps' in argv else '' + spec=dict(user='paphiwetsa3',claim_roots=[dict(path=str(__import__('pathlib').Path(f['request_claim']).parent),terminal_directory=f.get('terminal_directory',str(root/'status')))],scientific_job_names=['synthetic'],scientific_name_prefixes=['elmo-large314m-continuous-']) + print(json.dumps(gate.quiescence(spec,next(iter(f['profiles']))))) +elif f['operation']=='verify': + v=p.verify_observer(b['observer'],producer_request=b['request'],terminal_native=b['native']) + print(json.dumps({'result':v['result'],'running':v['running_admission_status'],'job':v['job'],'mode':v['spec']['mode']})) +else: + from scripts.train.joint.interruption_admission import completed_producer + q,r=completed_producer(b,f.get('call_schema',f['result_schema']),f.get('call_status',f['status'])) + print(json.dumps({'request_schema':q['schema'],'result_schema':r['schema']})) +''' + return subprocess.run([sys.executable,'-B','-c',program,json.dumps(payload)],env=environment(service),cwd=service,capture_output=True,text=True,check=False) + + +def refresh_observer(f): + path=Path(f['bundle']['observer']['path']);f['bundle']['observer']['sha256']=sha(path) + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize','full']) +@pytest.mark.parametrize('arm',['plain','fused','boundary']) +def test_all_phases_full_semantic_closure_without_external_code(tmp_path,service,mode,arm): + f=fixture(tmp_path,service,mode,arm);r=execute(f,service);assert r.returncode==0,r.stderr + out=json.loads(r.stdout);assert out['mode']==mode and out['result']['observer_own_terminal']=='UNKNOWN_NOT_REQUIRED' + assert out['result']['phase_completion_proven'] is False + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize']) +def test_four_successful_recovery_producers_require_full_observer(tmp_path,service,mode): + f=fixture(tmp_path,service,mode);r=execute(f,service,'complete');assert r.returncode==0,r.stderr + del f['bundle']['observer'];r=execute(f,service,'complete');assert r.returncode!=0 + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize']) +def test_unknown_result_status_or_mixed_phase_cannot_bypass(tmp_path,service,mode): + f=fixture(tmp_path,service,mode) + assert execute(f,service,'complete',call_schema='unknown-result').returncode!=0 + assert execute(f,service,'complete',call_status='PASS').returncode!=0 + other='inspect' if mode!='inspect' else 'capture' + schemas=dict(capture='native-elmo-interruption-evidence-v1',inspect='native-elmo-checkpoint-inspection-result-v1') + statuses=dict(capture='CAPTURED_REQUIRES_NATIVE_PRODUCER_COMPLETION',inspect='INSPECTED_REQUIRES_NATIVE_PRODUCER_COMPLETION') + assert execute(f,service,'complete',call_schema=schemas[other],call_status=statuses[other]).returncode!=0 + + +def test_missing_running_is_terminal_only_not_success_or_zero_progress(tmp_path,service): + f=fixture(tmp_path,service,running=False);r=execute(f,service);assert r.returncode==0,r.stderr + assert json.loads(r.stdout)['running']=='NOT_RECORDED' + assert execute(f,service,'complete').returncode!=0 + + +def test_failed_full_terminal_remains_failed_not_completion(tmp_path,service): + f=fixture(tmp_path,service,'full',running=False,state='FAILED');r=execute(f,service);assert r.returncode==0,r.stderr + result=json.loads(r.stdout)['result'];assert result['producer_terminal_exit']=='1:0' and result['phase_completion_proven'] is False + assert execute(f,service,'complete').returncode!=0 + + +@pytest.mark.parametrize('field',['job_id','producer_job_id']) +def test_result_job_keys_cannot_contradict_observed_producer(tmp_path,service,field): + f=fixture(tmp_path,service);path=Path(f['bundle']['result']['path']);d=read(path);d[field]='999999';f['bundle']['result']=put(path,d) + assert execute(f,service,'complete').returncode!=0 + + +@pytest.mark.parametrize('field',['request','native','host_exit','result']) +def test_bundle_cannot_relocate_observed_or_canonical_evidence(tmp_path,service,field): + f=fixture(tmp_path,service);src=Path(f['bundle'][field]['path']);target=tmp_path/('elsewhere-'+src.name);target.write_bytes(src.read_bytes()) + f['bundle'][field]=dict(path=str(target),sha256=sha(target)) + assert execute(f,service,'complete').returncode!=0 + + +@pytest.mark.parametrize('key',['producer_query','owned_steps','observer_running_query']) +def test_rehashed_query_errors_do_not_become_complete_commits(tmp_path,service,key): + f=fixture(tmp_path,service);orb=Path(f['bundle']['observer']['path']);result=read(orb);path=Path(result[key]['path']);d=read(path);d['stderr']='authentication failure' + result[key]=put(path,d);put(orb,result);refresh_observer(f) + assert execute(f,service).returncode!=0 + + +@pytest.mark.parametrize('replacement',['900001.batch|RUNNING\n','900001.extern|COMPLETING\n','900001.0|COMPLETED\n','900001_7.batch|RUNNING\n','malformed']) +def test_rehashed_steps_reject_surviving_parent_or_unknown_inventory(tmp_path,service,replacement): + f=fixture(tmp_path,service);orb=Path(f['bundle']['observer']['path']);result=read(orb);path=Path(result['owned_steps']['path']);d=read(path);d['stdout']=replacement + result['owned_steps']=put(path,d);put(orb,result);refresh_observer(f) + assert execute(f,service).returncode!=0 + + +@pytest.mark.parametrize('target',['result','query','steps']) +def test_partial_missing_capture_never_admits_raw_native(tmp_path,service,target): + f=fixture(tmp_path,service);result=read(f['bundle']['observer']['path']) + path=f['bundle']['observer']['path'] if target=='result' else result['producer_query' if target=='query' else 'owned_steps']['path'] + Path(path).unlink();assert execute(f,service,'complete').returncode!=0 + + +@pytest.mark.parametrize('file,field,value',[ + ('PREPARED.json','spec_sha256','f'*64),('acceptance.json','job_id','999999'), + ('RELEASE_ATTEMPT.json','observer_readiness',{}),('RELEASE_ATTEMPT.json','held_native_sha256','a'*64)]) +def test_exact_producer_claim_and_release_provenance_is_required(tmp_path,service,file,field,value): + f=fixture(tmp_path,service);path=Path(f['execution_claim'])/file;d=read(path);d[field]=value;put(path,d) + assert execute(f,service).returncode!=0 + + +@pytest.mark.parametrize('file', ['PRODUCER_HELD_QUERY.json','ACCEPTANCE_QUERY.json']) +def test_original_observer_queries_must_really_have_succeeded(tmp_path,service,file): + f=fixture(tmp_path,service);path=Path(f['observer_claim'])/file;d=read(path);d['returncode']=1;put(path,d) + assert execute(f,service).returncode!=0 + + +@pytest.mark.parametrize('change',['SubmitTime','Comment']) +def test_present_running_claim_cannot_contradict_exact_acceptance(tmp_path,service,change): + f=fixture(tmp_path,service);path=Path(f['request_claim'])/'RUNNING_ADMISSION.json';d=read(path) + d['native']=d['native'].replace('SubmitTime=2026-01-01T01:00:00','SubmitTime=2026-01-02T01:00:00') if change=='SubmitTime' else d['native'].replace('Comment=elmo-execution:','Comment=contradictory:') + put(path,d);assert execute(f,service).returncode!=0 + + +def test_host_exit_is_not_native_query_exit(tmp_path,service): + f=fixture(tmp_path,service);path=Path(f['bundle']['host_exit']['path']);f['bundle']['host_exit']=put(path,'1\n') + assert execute(f,service,'complete').returncode!=0 + + +def test_oversized_proof_payload_rejects_before_hashing_or_loading(tmp_path,service): + f=fixture(tmp_path,service);path=tmp_path/'oversized.ckpt';path.write_bytes(b'x'*(4*1024*1024+1)) + f['bundle']['observer']=dict(path=str(path),sha256=sha(path));r=execute(f,service);assert r.returncode!=0 + assert 'Bounded regular proof metadata' in r.stderr + + +def test_manifest_member_limit_and_unsafe_paths_reject(tmp_path,service): + f=fixture(tmp_path,service);path=Path(f['ops_manifest']['path']);path.write_text(path.read_text()+'a'*64+' ../outside.py\n') + assert execute(f,service).returncode!=0 + + +def test_source_contains_no_external_import_or_submission_entrypoint(service): + import ast + source=service/'scripts/train/joint/phase_observer_proof.py' + tree=ast.parse(source.read_text());names={n.names[0].name for n in ast.walk(tree) if isinstance(n,ast.Import)} + assert not names & {'subprocess','importlib','torch','numpy','egomimic'} + assert not any(isinstance(n,ast.Attribute) and n.attr in {'exec_module','run_module','load','loads_checkpoint'} for n in ast.walk(tree)) + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize','full']) +def test_new_claim_census_requires_complete_commit_not_native_only(tmp_path,service,mode): + f=fixture(tmp_path,service,mode);r=execute(f,service,'claim');assert r.returncode==0,r.stderr + assert json.loads(r.stdout)['provenance']=='COMPLETE_PHASE_OBSERVER_COMMIT' + Path(f['bundle']['observer']['path']).unlink() + assert execute(f,service,'claim').returncode!=0 + + +def test_maintained_quiescence_consumes_new_commit_and_fresh_owned_steps(tmp_path,service): + f=fixture(tmp_path,service,'full');r=execute(f,service,'quiescence');assert r.returncode==0,r.stderr + assert json.loads(r.stdout)['claims'][0]['provenance']=='COMPLETE_PHASE_OBSERVER_COMMIT' + Path(f['bundle']['observer']['path']).unlink() + assert execute(f,service,'quiescence').returncode!=0 + + +def test_generic_claim_census_does_not_hide_failed_early_wrapper(tmp_path,service): + f=fixture(tmp_path,service,'full',running=False,state='FAILED');r=execute(f,service,'claim');assert r.returncode==0,r.stderr + d=json.loads(r.stdout);assert d['terminal_state']=='FAILED' and d['running_admission_status']=='NOT_RECORDED' + + +def test_external_active_service_root_mismatch_rejects(tmp_path,service): + f=fixture(tmp_path,service);other=tmp_path/'another-service';shutil.copytree(service,other) + r=execute(f,other);assert r.returncode!=0 and 'active proof implementation' in r.stderr + + +def test_tampered_active_service_member_rejects_before_external_code(tmp_path,service): + other=tmp_path/'tampered-service';shutil.copytree(service,other);f=fixture(tmp_path,other) + member=other/'scripts/train/joint/phase_observer_proof.py';member.write_text(member.read_text()+'\n# unauthorized drift\n') + r=execute(f,other);assert r.returncode!=0 and 'Manifest member changed' in r.stderr + + +def test_exact_original_claim_route_is_preserved_without_new_phase_commit(tmp_path,service): + f=fixture(tmp_path,service,'original');r=execute(f,service,'claim');assert r.returncode==0,r.stderr + d=json.loads(r.stdout);assert d['provenance']=='EXACT_ORIGINAL_AFTERANY' and d['terminal_state']=='TIMEOUT' + assert not (Path(f['root'])/'native-111.observer.json').exists() + (Path(f['root'])/'native-111.exit').write_text('1\n') + assert execute(f,service,'claim').returncode!=0 + + +def test_original_job_id_does_not_license_an_arbitrary_claim_path(tmp_path,service): + f=fixture(tmp_path,service,'original');other=tmp_path/'relocated-claim';shutil.copytree(f['request_claim'],other) + f['request_claim']=str(other);assert execute(f,service,'claim').returncode!=0 + + +def test_manifest_member_bound_is_exercised_directly(tmp_path,service): + f=fixture(tmp_path,service);root=Path(f['root']);manifest=Path(f['ops_manifest']['path']);rows=manifest.read_text() + for i in range(124): + member=root/('metadata-'+str(i)+'.txt');member.write_text('small fixture\n');rows+=sha(member)+' '+member.name+'\n' + f['ops_manifest']=put(manifest,rows) + result=execute(f,service,'manifest');assert result.returncode!=0 and 'member bound' in result.stderr + + +def test_duplicate_final_commit_keys_reject_even_with_new_hash(tmp_path,service): + f=fixture(tmp_path,service);path=Path(f['bundle']['observer']['path']);raw=path.read_text() + path.write_text('{"producer_job_id":"999999",'+raw.lstrip()[1:]);refresh_observer(f) + r=execute(f,service);assert r.returncode!=0 and 'Duplicate proof JSON keys' in r.stderr + + +def test_symlink_commit_cannot_relocate_a_valid_proof(tmp_path,service): + f=fixture(tmp_path,service);path=Path(f['bundle']['observer']['path']);outside=tmp_path/'external-commit.json';outside.write_bytes(path.read_bytes()) + path.unlink();path.symlink_to(outside) + r=execute(f,service);assert r.returncode!=0 and 'Physical existing proof path' in r.stderr diff --git a/tests/native_elmo_ops/controller_fixture.py b/tests/native_elmo_ops/controller_fixture.py new file mode 100644 index 000000000..53eb86f52 --- /dev/null +++ b/tests/native_elmo_ops/controller_fixture.py @@ -0,0 +1,55 @@ +"""Portable in-memory remote filesystem fixture; no actual SSH or files.""" +import controller as c + +ROOT=c.PREFIX+"synthetic-controller-fixture" + +def fixture(): + files={} + def save(name,value,raw=False): + location=ROOT+'/'+name + data=value if raw else c.encode(value) + files[location]=data + return {'path':location,'sha256':c.digest(data)} + original=c.encode(dict(run_id='mot-elmo-pickplace-continuous-large314m-plain-robot-human-camera-subtask-s42-300k-20261026-v1',job_name='synthetic-original',parameters=314187064,world_size=4,target_steps=300000,phase='full')) + c.PINS=dict(c.PINS,plain=('412789','gpu11',c.digest(original))) + parent=save('inputs/parent.json',original,True); q=c.parse(original) + intent=save('inputs/current-intent.json',dict(schema='native-elmo-current-intent-v1',authority='human',status='ACTIVE',hold_recovery=False, + authorized_recovery_run_ids=[q['run_id']],retired_run_ids=[])) + terminal=save('inputs/parent-native.txt',('JobId=412789 JobName='+q['job_name']+' UserId=paphiwetsa3(41003) NodeList=gpu11 JobState=TIMEOUT EndTime=2026-10-27T00:00:00 ExitCode=0:1\n').encode(),True) + py_script=save('inputs/python-producer.sbatch',b'#!/bin/bash\n/usr/bin/python3 -B "$root/cache.py"\n',True) + py_native=save('inputs/python-native.txt',b'JobId=412768 NodeList=gpu11 JobState=COMPLETED ExitCode=0:0\n',True) + py_exit=save('inputs/python-exit.txt',b'0\n',True) + provenance=save('inputs/python-provenance.json',dict(schema='native-elmo-host-python-provenance-v1',node='gpu11',job_id='412768', + interpreter_path='/usr/bin/python3',script=py_script,native=py_native,host_exit=py_exit)) + service_member=save('service/member.py',b'# synthetic test fixture\n',True) + sm=save('service/FILES.sha256',(service_member['sha256']+' member.py\n').encode(),True) + ops={name:save(name,('# test '+name+'\n').encode(),True)['sha256'] for name in + ('controller.py','controller.sbatch','execution_spec.py','host_guard.py','phase_observer.py','observe-phase.sbatch')} + om=save('FILES.sha256',''.join(sha+' '+name+'\n' for name,sha in sorted(ops.items())).encode(),True) + scope=dict(user='paphiwetsa3',scientific_job_names=[q['job_name']],scientific_name_prefixes=['elmo-large314m-continuous-plain-']) + template=save('templates/producer.json',{'fixture':True}) + producer=dict(root=ROOT,arm='plain',mode='capture',parent_request=parent,request_template=template, + runtime={'admission':dict(current_intent_path=intent['path'],quiescence=scope)}, + operational_service=dict(root=ROOT+'/service',manifest=sm['path'],manifest_sha256=sm['sha256']), + release_observer_path=ROOT+'/observer-specs/observe.json') + pb=save('execution/producer.json',producer) + execution=dict(pb,mode='capture',ops_manifest_sha256=om['sha256']) + observer=save('observer-specs/observe.json',dict(schema='native-elmo-phase-observer-spec-v1',root=ROOT, + producer_execution=execution,terminal_directory=ROOT+'/status')) + spec=dict(schema='native-elmo-allocated-controller-v1',root=ROOT,arm='plain',parent_request=parent,ops_manifest=om, + current_intent=intent,original_failed_native=terminal,producer_execution=execution,observer_spec=observer, + terminal_directory=ROOT+'/status',job_name='elmo-recovery-controller-plain-fixture', + controller_claim_root=c.PREFIX+'recovery-controller-claims/plain',queue_scope=scope, + interpreter=dict(path='/usr/bin/python3',minimum_version=[3,10],provenance=provenance)) + sb=save('controller-specs/controller.json',spec) + return files,spec,(sb['path'],sb['sha256'],om['sha256']) + +def native(spec,state='PENDING',**changes): + f=dict(JobId='900001',JobName=spec['job_name'],UserId='paphiwetsa3(41003)',Account='loaner',QOS='loaner-normal', + Partition='rl2',Requeue='0',Restarts='0',JobState=state,ReqNodeList='gpu11',NumCPUs='2',NumTasks='1', + MinMemoryNode='1G',TimeLimit='00:20:00',Command=ROOT+'/controller.sbatch',Comment='elmo-controller:'+'a'*64, + SubmitTime='2026-10-27T00:01:00',NumNodes='1',Reason='JobHeldUser',NodeList='(null)',ReqTRES='cpu=2,mem=1G,node=1,billing=2', + AllocTRES='(null)',JOB_GRES='(null)',StartTime='Unknown',TresPerTask='cpu:2') + if state=='RUNNING': f.update(NodeList='gpu11',StartTime='2026-10-27T00:02:00',AllocTRES='cpu=2,mem=1G,node=1,billing=2') + f.update(changes) + return ' '.join(k+'='+v for k,v in f.items())+'\n' diff --git a/tests/native_elmo_ops/execution_fixture.py b/tests/native_elmo_ops/execution_fixture.py new file mode 100644 index 000000000..f0b65172b --- /dev/null +++ b/tests/native_elmo_ops/execution_fixture.py @@ -0,0 +1,75 @@ +"""Portable synthetic execution fixture; no deployment request is read.""" +import copy, hashlib, json, shutil, subprocess +from pathlib import Path + +e=HERE=SERVICE=None + +def put(p, obj): + p.write_bytes(e.encoded(obj)) + +def fixture(tmp_path, mode="inspect", arm="plain"): + root = tmp_path / (mode + "-" + arm) + root.mkdir() + for name in ("execution", "templates", "requests", "execution-claims", "scientific-claims", "producer-claims", "status", "logs"): + (root / name).mkdir() + shutil.copyfile(HERE / "execution_spec.py", root / "execution_spec.py") + (root / "allocated.sbatch").write_text("#!/bin/bash\nexit 99\n") + (root / "host_guard.py").write_text("def admit_before_submission(spec, template, **kwargs):\n raise RuntimeError('test guard never executes')\n") + (root / "FILES.sha256").write_text("".join(e.sha(root / name) + " " + name + "\n" for name in ("allocated.sbatch", "execution_spec.py", "host_guard.py"))) + parent_id, node, parent_sha = e.PINS[arm] + run_id="mot-elmo-pickplace-continuous-large314m-"+arm+"-robot-human-camera-subtask-s42-300k-20261026-v1" + parent=dict(schema="joint-training-v2",phase="full",kind="native-subtask-world-model-v1",run_id=run_id, + job_name="synthetic-original-"+arm,target_steps=300000,parameters=314187064,world_size=4, + output=str(root/"original-output"),local_logs=str(root/"original-logs")) + for key in ("source","source_revision","source_manifest","source_manifest_sha256","data_manifest","data_manifest_sha256","lambda_runtime","data_cache"): + parent[key]="synthetic-"+key + parent_path=root/"synthetic-parent.json";put(parent_path,parent) + parent_sha=e.sha(parent_path) + e.PINS=dict(e.PINS, **{arm:(parent_id,node,parent_sha)}) + _, cpus, mem, limit, gres, _ = e.MODES[mode] + n = dict(JobName=("elmo-native-recovery-" if mode not in {"diagnostic", "full"} else "elmo-large314m-continuous-") + arm + "-" + mode + "-test", + UserId="paphiwetsa3(41003)", Account="loaner", QOS="loaner-normal", Partition="rl2", ReqNodeList=node, + NodeList=node, NumCPUs=cpus, MinMemoryNode=mem, TimeLimit=limit, Requeue="0", Restarts="0", + Command=str(root / "allocated.sbatch"), TresPerNode=gres) + service = dict(root=str(SERVICE), manifest=str(SERVICE / "FILES.sha256"), manifest_sha256=e.SERVICE_SHA) + pb = dict(path=str(parent_path), sha256=parent_sha) + if mode == "full": + q = copy.deepcopy(parent) + q.update(job_name=n["JobName"], output=str(root / "production-output"), local_logs=str(root / "production-logs"), + operational_service=service, interruption=dict(parent_request=pb)) + else: + q = dict(schema=e.MODES[mode][0], parent_request=pb, operational_service=service, producer_native=n, + output=str(root / "phase-output"), allocated_prerequisites_path=str(root / "status/prerequisites.json")) + for key in ("source", "source_revision", "source_manifest", "source_manifest_sha256", "data_manifest", "data_manifest_sha256", "lambda_runtime", "data_cache"): + q[key] = parent.get(key) + if mode == "capture": + q.update(parent_job_id=parent_id, parent_native=dict(JobId=parent_id)) + template = root / "templates/phase.json" + put(template, q) + spec = dict(schema="native-elmo-execution-spec-v1", root=str(root), mode=mode, arm=arm, + original_parent_job_id=parent_id, parent_request=pb, request_template=dict(path=str(template), sha256=e.sha(template)), + native_contract=n, operational_service=service, + execution_claim_roots=[dict(path=str(root / "execution-claims"),terminal_directory=str(root / "status"))], + ops_manifest=dict(path=str(root / "FILES.sha256"), sha256=e.sha(root / "FILES.sha256"))) + sp = root / "execution/phase.json" + put(sp, spec) + return dict(spec=spec, template=q, path=sp, args=(str(sp), e.sha(sp), mode, e.sha(root / "FILES.sha256")), root=root) + +def native(f, state="PENDING", job="900001", **changes): + n = dict(f["spec"]["native_contract"], JobId=job, JobState=state, Reason="JobHeldUser", NumNodes="1-1", + NumTasks="1", SubmitTime="2026-10-27T01:00:00", StartTime="Unknown", NodeList="", AllocTRES="(null)", + ReqTRES="cpu=" + f["spec"]["native_contract"]["NumCPUs"] + ",mem=" + f["spec"]["native_contract"]["MinMemoryNode"] + ",node=1,billing=32", + Comment="elmo-execution:" + f["args"][1]) + if state == "RUNNING": + n.update(Reason="None", NumNodes="1", NodeList=f["spec"]["native_contract"]["NodeList"], + AllocTRES=n["ReqTRES"], StartTime="2026-10-27T01:00:02") + if n["TresPerNode"]: + n.update(JOB_GRES="gpu:h100:4", GRES="gpu:h100:4(IDX:0-3)") + n.update(changes) + return " ".join(k + "=" + v for k, v in n.items()) + "\n" + +def accepted(f, monkeypatch, job="900001"): + monkeypatch.setattr(e, "host_preflight", lambda *a, **kw: None) + monkeypatch.setattr(e.subprocess, "run", lambda *a, **kw: subprocess.CompletedProcess(a[0], 0, job+"\n", "")) + monkeypatch.setattr(e, "command", lambda argv: native(f,job=job)) + return e.submit_held(*f["args"]) diff --git a/tests/native_elmo_ops/observer_fixture.py b/tests/native_elmo_ops/observer_fixture.py new file mode 100644 index 000000000..2e138ce12 --- /dev/null +++ b/tests/native_elmo_ops/observer_fixture.py @@ -0,0 +1,39 @@ +"""Portable synthetic native observer fixtures, never native queries.""" +from datetime import datetime,timezone +import os, subprocess +from pathlib import Path + +o=support=None + +def observer_native(f,state='PENDING',**changes): + s,e,p,*_=o.read_spec(*f['observer_args']);n=o.native_contract(s,p) + row=dict(n,JobId=f['observer_id'],JobState=state,Reason='Dependency',NumNodes='1-1',NumTasks='1',NodeList='', + ReqTRES='cpu=2,mem=1G,node=1,billing=2',AllocTRES='(null)',SubmitTime='2026-10-27T01:00:01',StartTime='Unknown', + Dependency='afterany:900001(unfulfilled)',Comment='elmo-phase-observer:'+f['observer_args'][1],WorkDir=str(f['root']), + StdOut=str(f['root']/'logs'/(n['JobName']+'-'+f['observer_id']+'.out')), + StdErr=str(f['root']/'logs'/(n['JobName']+'-'+f['observer_id']+'.err'))) + if state=='RUNNING':row.update(NodeList=n['NodeList'],NumNodes='1',AllocTRES='cpu=2,mem=1G,node=1,billing=2',StartTime='2026-10-27T02:00:00',Reason='None',Dependency='(null)') + row.update(changes);return ' '.join(k+'='+v for k,v in row.items())+'\n' + +def record(argv,code=0,stdout='',stderr=''): + return dict(schema='native-elmo-phase-observer-query-v1',argv=argv,returncode=code,stdout=stdout,stderr=stderr, + observed_at_utc=datetime.now(timezone.utc).isoformat()) + +def accept(f,monkeypatch): + monkeypatch.setattr(o,'query',lambda argv:record(argv,stdout=support.native(f) if argv[3]=='900001' else observer_native(f))) + monkeypatch.setattr(o.subprocess,'run',lambda argv,**kwargs:subprocess.CompletedProcess(argv,0,f['observer_id']+'\n','')) + return o.submit_once(*f['observer_args']) + +def terminal(f,**changes): + row=support.e.fields(support.native(f,'RUNNING')) + row.update(JobState='TIMEOUT',Reason='TimeLimit',EndTime='2026-10-27T02:00:00',ExitCode='0:15') + if row.get('GRES'):row['GRES']='' + row.update(changes);return ' '.join(k+'='+v for k,v in row.items())+'\n' + +def ready_capture_after_accept(f,monkeypatch): + _,path,q,job,claim,held=support.e.read_binding(*f['args']) + support.e.write_once(claim/'RUNNING_ADMISSION.json',dict(job_id=job,request_sha256=o.sha(path),spec_sha256=f['args'][1],native=support.native(f,'RUNNING'))) + monkeypatch.setattr(os,'getuid',lambda:41003) + monkeypatch.setenv('SLURM_JOB_ID',f['observer_id']);monkeypatch.setenv('SLURMD_NODENAME',f['spec']['native_contract']['NodeList']) + monkeypatch.delenv('SLURM_JOB_GPUS',raising=False);monkeypatch.delenv('SLURM_STEP_GPUS',raising=False) + monkeypatch.setattr(o,'query',lambda argv:record(argv,stdout=('900999.batch|RUNNING\n' if argv[0].endswith('squeue') else observer_native(f,'RUNNING') if argv[3]==f['observer_id'] else terminal(f)))) diff --git a/tests/native_elmo_ops/test_native_elmo_ops.py b/tests/native_elmo_ops/test_native_elmo_ops.py new file mode 100644 index 000000000..c94d50d5c --- /dev/null +++ b/tests/native_elmo_ops/test_native_elmo_ops.py @@ -0,0 +1,400 @@ +"""Synthetic native effects; real published service pins and runtime bytes are tested.""" +import copy +import os +import tempfile +from datetime import datetime, timezone +import hashlib +import importlib.util +import json +from pathlib import Path +import shutil +import subprocess +import sys +from types import SimpleNamespace +import pytest + +HERE=Path(__file__).resolve().parent +REPO=HERE.parents[1] +OPS=Path(os.environ.get('ELMO_TEST_OPS_ROOT',REPO/'scripts/train/joint/native_elmo_ops')).resolve() +SOURCE=Path(os.environ.get('ELMO_TEST_SERVICE_ROOT',REPO)).resolve() +NAMES=('execution_spec','phase_observer','wrapper_common','host_guard') + + +def load(name,path): + loader=importlib.util.spec_from_file_location(name,path) + module=importlib.util.module_from_spec(loader);loader.loader.exec_module(module);return module + +# Export only the maintained service closure. No campaign receipts or models. +service_temp=tempfile.TemporaryDirectory(prefix='native-elmo-ops-service-') +SERVICE=Path(service_temp.name).resolve() +bootstrap=load('_portable_ops_service_export',SOURCE/'scripts/train/joint/service_bootstrap.py') +for name in bootstrap.service_files(SOURCE): + target=SERVICE/name;target.parent.mkdir(parents=True,exist_ok=True);shutil.copyfile(SOURCE/name,target) +(SERVICE/'FILES.sha256').write_text(''.join(hashlib.sha256((SERVICE/name).read_bytes()).hexdigest()+' '+name+'\n' for name in bootstrap.service_files(SERVICE))) +support=load('_portable_execution_fixture',HERE/'execution_fixture.py') +observer_support=load('_portable_observer_fixture',HERE/'observer_fixture.py') + +@pytest.fixture(autouse=True) +def isolate(monkeypatch): + for name in NAMES:monkeypatch.delitem(sys.modules,name,raising=False) + monkeypatch.setattr(subprocess,'run',lambda *a,**k:pytest.fail('Unintercepted subprocess forbidden')) + yield + for name in NAMES:sys.modules.pop(name,None) + + +def fixture(tmp_path,monkeypatch,mode='inspect',arm='plain',accept=True): + src=load('_composite_fixture_source',OPS/'execution_spec.py') + assert src.SERVICE_SHA==src.sha(SERVICE/'FILES.sha256') + support.e=src;support.HERE=OPS;support.SERVICE=SERVICE + f=support.fixture(tmp_path,mode,arm);root=f['root'] + for source in OPS.iterdir(): + if source.is_file():shutil.copy2(source,root/source.name) + elif not (root/source.name).exists():(root/source.name).mkdir() + f['spec']['release_observer_path']=str(root/'observer-specs/phase.json') + f['spec']['ops_manifest']['sha256']=src.sha(root/'FILES.sha256') + support.put(f['path'],f['spec']) + f['args']=(str(f['path']),src.sha(f['path']),mode,src.sha(root/'FILES.sha256')) + e=load('execution_spec',root/'execution_spec.py');monkeypatch.setitem(sys.modules,'execution_spec',e) + monkeypatch.setattr(e,'PINS',src.PINS) + assert e.SERVICE_SHA==e.sha(SERVICE/'FILES.sha256') + support.e=e + if accept:support.accepted(f,monkeypatch) + o=e.operations_module(f['spec'],'phase_observer') + wc=e.operations_module(f['spec'],'wrapper_common') + assert (wc.SERVICE_COMMIT,wc.SERVICE_SHA)==(e.SERVICE_COMMIT,e.SERVICE_SHA) + h=e.operations_module(f['spec'],'host_guard') + execution=dict(path=str(f['path']),sha256=f['args'][1],mode=mode,ops_manifest_sha256=f['args'][3]) + observer_spec=o.build_spec(execution);sp=Path(f['spec']['release_observer_path']);support.put(sp,observer_spec) + f.update(e=e,o=o,h=h,wc=wc,observer_spec=observer_spec,observer_path=sp, + observer_args=(str(sp),e.sha(sp),f['args'][3]),observer_id='900002') + observer_support.o=o;observer_support.support=support + return f + + +def accept_observer(f,monkeypatch):return observer_support.accept(f,monkeypatch) + +def capture(f,monkeypatch): + observer_support.ready_capture_after_accept(f,monkeypatch) + return f['o'].capture(*f['observer_args']) + + +def bind(p):return {'path':str(p),'sha256':hashlib.sha256(p.read_bytes()).hexdigest()} + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize','full']) +@pytest.mark.parametrize('arm',['plain','fused','boundary']) +def test_manifest_cycle_free_held_observer_release(tmp_path,monkeypatch,mode,arm): + f=fixture(tmp_path,monkeypatch,mode,arm);e=f['e'];o=f['o'] + accept_observer(f,monkeypatch) + # Real composed release guard invokes real observer dependency checks. + calls=[] + monkeypatch.setattr(e.subprocess,'run',lambda argv,**kw:(calls.append(argv) or subprocess.CompletedProcess(argv,0,'',''))) + e.release_bound(*f['args']) + assert calls==[['/usr/local/bin/scontrol','release','900001']] + release=e.read_json(e.paths(f['spec'],f['args'][1])['spec_claim']/'RELEASE_ATTEMPT.json') + release_held=e.paths(f['spec'],f['args'][1])['spec_claim']/'RELEASE_HELD_NATIVE.txt' + assert release_held.read_text()==support.native(f) + assert e.sha(release_held)==release['held_native_sha256'] + assert release['observer_readiness']['producer_job_id']=='900001' + assert release['observer_readiness']['observer_job_id']=='900002' + assert release['observer_readiness']['observer_pending_query']['stdout'] + assert e._verified_ops_member_sha256==e.sha(f['root']/'execution_spec.py') + assert o._verified_ops_member_sha256==e.sha(f['root']/'phase_observer.py') + assert 'EXECUTION_SHA' not in (f['root']/'phase_observer.py').read_text() + + +@pytest.mark.parametrize('conflict',['missing','unknown','not-dependency','wrong-job']) +def test_release_requires_actual_matching_observer(tmp_path,monkeypatch,conflict): + f=fixture(tmp_path,monkeypatch);e=f['e'];o=f['o'] + if conflict!='missing':accept_observer(f,monkeypatch) + if conflict=='unknown': + a=o.paths(dict(f['observer_spec'],producer_job_id='900001'),f['observer_args'][1])['claim']/'acceptance.json' + a.unlink() # Disposable synthetic fixture only. + if conflict=='not-dependency': + monkeypatch.setattr(o,'query',lambda argv:observer_support.record(argv,stdout=support.native(f) if argv[3]=='900001' else observer_support.observer_native(f,Reason='Resources'))) + if conflict=='wrong-job': + monkeypatch.setattr(o,'query',lambda argv:observer_support.record(argv,stdout=support.native(f) if argv[3]=='900001' else observer_support.observer_native(f,JobId='900099'))) + monkeypatch.setattr(subprocess,'run',lambda *a,**k:pytest.fail('release reached without observer proof')) + with pytest.raises((ValueError,FileNotFoundError)):e.release_bound(*f['args']) + assert not (e.paths(f['spec'],f['args'][1])['spec_claim']/'RELEASE_ATTEMPT.json').exists() + + +@pytest.mark.parametrize('member',['phase_observer.py','execution_spec.py','observe-phase.sbatch','wrapper_common.py']) +def test_manifest_changed_member_rejected_before_import(tmp_path,monkeypatch,member): + f=fixture(tmp_path,monkeypatch);p=f['root']/member;p.write_text(p.read_text()+'\n# drift\n') + with pytest.raises(ValueError):f['o'].engine(f['root'],f['args'][3]) + with pytest.raises(ValueError):f['e'].operations_module(f['spec'],'phase_observer') + + +def test_foreign_cached_execution_not_adopted(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch) + monkeypatch.setitem(sys.modules,'execution_spec',SimpleNamespace(__file__='/wrong/path.py')) + with pytest.raises(ValueError,match='Foreign'):f['o'].engine(f['root'],f['args'][3]) + + +def test_phase_verify_is_maintained_delegate_with_exact_identity(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch);accept_observer(f,monkeypatch);result=capture(f,monkeypatch) + original=f['e'].service_module;calls=[] + def verify(binding,**kwargs):calls.append((binding,kwargs));return {'result':result} + monkeypatch.setattr(f['e'],'service_module',lambda spec,name:SimpleNamespace(verify_observer=verify) if name=='phase_observer_proof' else original(spec,name)) + assert f['o'].verify_captured(*f['observer_args'])==result + assert calls[0][0]==bind(f['root']/'status/native-900001.observer.json') + assert calls[0][1]['producer_native']['JobId']=='900001' + assert calls[0][1]['producer_request']==result['producer_request'] + assert calls[0][1]['terminal_native']==result['producer_native'] + result['producer_job_id']='999' + with pytest.raises(ValueError,match='another producer'):f['o'].verify_captured(*f['observer_args']) + + +def terminal_route(f,monkeypatch): + accept_observer(f,monkeypatch);result=capture(f,monkeypatch) + original=f['e'].service_module + monkeypatch.setattr(f['e'],'service_module',lambda spec,name:SimpleNamespace(verify_observer=lambda *a,**k:{'result':copy.deepcopy(result)},check_owned_steps=original(spec,name).check_owned_steps) if name=='phase_observer_proof' else original(spec,name)) + adapter=original(f['spec'],'original_observer') + record={'schema':'native-elmo-native-query-v1','argv':adapter.step_argv('paphiwetsa3'),'returncode':0,'stdout':'','stderr':'','observed_at_utc':datetime.now(timezone.utc).isoformat()} + # Use adapter's own query builder with intercepted subprocess for its exact schema. + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + return result,adapter + + +def test_execution_census_requires_observer_final_commit_and_fresh_steps(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch);result,adapter=terminal_route(f,monkeypatch) + assert f['e'].execution_claim_census(f['spec'])[0]['observer_spec']==result['observer_spec'] + monkeypatch.setattr(f['o'],'verify_captured',lambda *a:(_ for _ in ()).throw(ValueError('missing final commit'))) + with pytest.raises(ValueError,match='final commit'):f['e'].execution_claim_census(f['spec']) + monkeypatch.setattr(f['o'],'verify_captured',lambda *a:result) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'900001.batch|RUNNING\n','')) + with pytest.raises(ValueError):f['e'].execution_claim_census(f['spec']) + + +def test_host_claim_census_uses_same_observer_consumer(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch,'diagnostic');result,adapter=terminal_route(f,monkeypatch) + scope={'user':'paphiwetsa3','scientific_job_names':[f['spec']['native_contract']['JobName']], + 'scientific_name_prefixes':['elmo-large314m-continuous-'], + 'claim_roots':[{'path':str(f['root']/'scientific-claims'),'terminal_directory':str(f['root']/'status')}]} + proof=f['h'].quiescence(f['spec'],scope,'run-id') + assert proof['prior_claims'][0]['route']=='maintained-phase-observer' + monkeypatch.setattr(f['o'],'verify_captured',lambda *a:(_ for _ in ()).throw(ValueError('corrupt observer ledger'))) + with pytest.raises(ValueError,match='observer ledger'):f['h'].quiescence(f['spec'],scope,'run-id') + + +def test_unknown_execution_claim_blocks_before_other_raw_terminals(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch,accept=False) + c=f['root']/'execution-claims'/f['args'][1];c.mkdir() + support.put(c/'PREPARED.json',dict(spec_sha256=f['args'][1],arm='plain',parent_request=f['spec']['parent_request'],spec_path=str(f['path']))) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + with pytest.raises(ValueError,match='Unresolved'):f['e'].execution_claim_census(f['spec']) + + +def test_source_publication_pins_match_exact_published_service(): + e=load('_publication_source',OPS/'execution_spec.py');wc=load('_publication_wrapper',OPS/'wrapper_common.py') + expected=('abbb00499096faabf7035fde23bf4b1eba984377', + 'adf45524de3c2e6dba186572e7151d48bfca0b3adc8e6cf9defd686b47744c87') + assert (e.SERVICE_COMMIT,e.SERVICE_SHA)==expected==(wc.SERVICE_COMMIT,wc.SERVICE_SHA) + assert e.sha(SERVICE/'FILES.sha256')==expected[1] + binding=dict(root=str(SERVICE),manifest=str(SERVICE/'FILES.sha256'),manifest_sha256=expected[1]) + module=e.service_module(dict(operational_service=binding),'phase_observer_proof') + assert Path(module.__file__).resolve()==SERVICE/'scripts/train/joint/phase_observer_proof.py' + + +@pytest.mark.parametrize('mode',['inspect','diagnostic']) +@pytest.mark.parametrize('running',[False,True]) +def test_actual_current_claim_interface_held_and_running(tmp_path,monkeypatch,mode,running): + import os + f=fixture(tmp_path,monkeypatch,mode);e=f['e'];h=f['h'] + if running: + accept_observer(f,monkeypatch) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + e.release_bound(*f['args']) + monkeypatch.setattr(os,'getuid',lambda:41003) + monkeypatch.setenv('SLURM_JOB_ID','900001');monkeypatch.setenv('SLURMD_NODENAME','gpu11') + monkeypatch.delenv('SLURM_JOB_GPUS',raising=False);monkeypatch.delenv('SLURM_STEP_GPUS',raising=False) + state='RUNNING' if running else 'PENDING' + raw=support.native(f,state) + monkeypatch.setattr(e,'command',lambda argv:raw) + def query(argv,**kw): + out=raw if argv[0].endswith('scontrol') else '' if '--steps' in argv else '900001|'+f['spec']['native_contract']['JobName']+'|'+state+'|\n' + return subprocess.CompletedProcess(argv,0,out,'') + monkeypatch.setattr(subprocess,'run',query) + _,qpath,q,job,claim,held=e.read_binding(*f['args']) + scope={'user':'paphiwetsa3','scientific_job_names':[f['spec']['native_contract']['JobName']], + 'scientific_name_prefixes':['elmo-large314m-continuous-'], + 'claim_roots':[{'path':str(claim.parent),'terminal_directory':str(f['root']/'status')}]} + # No stubs for own_current, read_binding, load_execution or native identity. + proof=h.quiescence(f['spec'],scope,'run-id',job,e.sha(qpath),running=running) + assert proof['own_current_claim']==str(claim) + assert proof['prior_claims']==[] + + +@pytest.mark.parametrize('area',['execution-json','observer-json','preserved-query']) +def test_metadata_and_query_bounds(tmp_path,monkeypatch,area): + f=fixture(tmp_path,monkeypatch) + if area=='preserved-query': + row=observer_support.record(['/native'],stdout='x'*(1048576+1)) + with pytest.raises(ValueError,match='bounds'):f['o'].successful(row,['/native']) + else: + p=tmp_path/'oversize.json';p.write_bytes(b' '*(4*1024*1024+1)) + module=f['e'] if area=='execution-json' else f['o'] + with pytest.raises(ValueError,match='4 MiB'):module.read_json(p) + + +def actual_service_capture(f,monkeypatch): + accept_observer(f,monkeypatch) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + f['e'].release_bound(*f['args']) + result=capture(f,monkeypatch) + proof=f['e'].service_module(f['spec'],'phase_observer_proof') + # Synthetic original-parent profiles are adjusted only in test memory. + # No deployed request or production profile is altered. + profiles=copy.deepcopy(proof.ORIGINAL_OBSERVERS) + run_id=f['e'].read_json(Path(f['spec']['parent_request']['path']))['run_id'] + profiles[run_id]['request_path']=f['spec']['parent_request']['path'] + profiles[run_id]['request_sha256']=f['spec']['parent_request']['sha256'] + monkeypatch.setattr(proof,'ORIGINAL_OBSERVERS',profiles) + return result,proof + + +@pytest.mark.parametrize('mode',['capture','inspect','diagnostic','finalize','full']) +@pytest.mark.parametrize('arm',['plain','fused','boundary']) +def test_actual_composite_capture_through_actual_maintained_service(tmp_path,monkeypatch,mode,arm): + f=fixture(tmp_path,monkeypatch,mode,arm);result,proof=actual_service_capture(f,monkeypatch) + # No stub for observer verify_captured, maintained verify_observer, binding, + # claim readers, native semantics, final commit or release readiness. + monkeypatch.setattr(subprocess,'run',lambda *a,**k:pytest.fail('Verifier made native mutation/query')) + assert f['o'].verify_captured(*f['observer_args'])==result + ctx=proof.verify_observer(bind(f['root']/'status/native-900001.observer.json'), + producer_request=result['producer_request'],producer_native=dict(f['spec']['native_contract'],JobId='900001'), + terminal_native=result['producer_native']) + assert ctx['running_admission_status']=='VALIDATED' + assert ctx['job']=='900001' + + +@pytest.mark.parametrize('broken',['release-held','final-commit','release-readiness','owned-step']) +def test_actual_service_rejects_incomplete_or_conflicting_composite_proof(tmp_path,monkeypatch,broken): + f=fixture(tmp_path,monkeypatch);result,proof=actual_service_capture(f,monkeypatch) + sc=f['e'].paths(f['spec'],f['args'][1])['spec_claim'] + if broken=='release-held':(sc/'RELEASE_HELD_NATIVE.txt').write_text(support.native(f,SubmitTime='2026-01-01T00:00:00')) + elif broken=='final-commit':(f['root']/'status/native-900001.observer.json').unlink() + elif broken=='release-readiness': + release=f['e'].read_json(sc/'RELEASE_ATTEMPT.json');release['observer_readiness']['observer_pending_query']['stdout']=observer_support.observer_native(f,Reason='Resources');support.put(sc/'RELEASE_ATTEMPT.json',release) + else: + path=Path(result['owned_steps']['path']);r=f['e'].read_json(path);r['stdout']='900001.batch|RUNNING\n';support.put(path,r) + result['owned_steps']=bind(path);support.put(f['root']/'status/native-900001.observer.json',result) + with pytest.raises((ValueError,FileNotFoundError)):f['o'].verify_captured(*f['observer_args']) + +@pytest.mark.parametrize('consumer',['execution','scientific']) +def test_both_censuses_consume_actual_service_proof(tmp_path,monkeypatch,consumer): + f=fixture(tmp_path,monkeypatch,'diagnostic');result,proof=actual_service_capture(f,monkeypatch) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + scope={'user':'paphiwetsa3','scientific_job_names':[f['spec']['native_contract']['JobName']], + 'scientific_name_prefixes':['elmo-large314m-continuous-'], + 'claim_roots':[{'path':str(f['root']/'scientific-claims'),'terminal_directory':str(f['root']/'status')}]} + def census(): + return f['e'].execution_claim_census(f['spec']) if consumer=='execution' else f['h'].quiescence(f['spec'],scope,'run-id') + assert census() + final=f['root']/'status/native-900001.observer.json';final.unlink() + assert (f['root']/'status/native-900001.txt').is_file() + with pytest.raises((ValueError,FileNotFoundError)):census() + + +def test_original_parent_never_falls_back_to_raw_phase_terminal(tmp_path,monkeypatch): + f=fixture(tmp_path,monkeypatch,'full');e=f['e'];h=f['h'];parent_job=f['spec']['original_parent_job_id'] + root=tmp_path/'original-scientific-claims';root.mkdir();claim=root/f['spec']['parent_request']['sha256'];claim.mkdir() + support.put(claim/'PREPARED.json',{'request_sha256':claim.name});(claim/'PARENT_JOB_ID').write_text(parent_job+'\n') + running=support.native(f,'RUNNING',job=parent_job) + support.put(claim/'RUNNING_ADMISSION.json',{'job_id':parent_job,'request_sha256':claim.name,'native':running}) + terminal=running.replace('JobState=RUNNING','JobState=TIMEOUT')+' EndTime=2026-10-27T02:00:00 ExitCode=0:15\n' + (f['root']/'status'/('native-'+parent_job+'.txt')).write_text(terminal) + scope={'user':'paphiwetsa3','scientific_job_names':['original'], 'scientific_name_prefixes':['elmo-large314m-continuous-'], + 'claim_roots':[{'path':str(root),'terminal_directory':str(f['root']/'status')}]} + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'','')) + with pytest.raises(ValueError,match='separately validated'):h.quiescence(f['spec'],scope,'run-id') + out=h.quiescence(f['spec'],scope,'run-id',original_terminal=terminal) + assert out['prior_claims'][0]['route']=='validated-original-parent' + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,parent_job+'.batch|RUNNING\n' if '--steps' in argv else '','')) + with pytest.raises(ValueError):h.quiescence(f['spec'],scope,'run-id',original_terminal=terminal) + + +def controller_history_fixture(monkeypatch): + c=load('_composite_controller',OPS/'controller.py');monkeypatch.setitem(sys.modules,'controller',c) + helper=load('_controller_virtual_fixture',HERE/'controller_fixture.py') + helper.c=c;helper.ROOT=c.PREFIX+'synthetic-prior-controller' + oldfiles,old,args=helper.fixture() + def save(files,loc,value): + files[loc]=c.encode(value);return {'path':loc,'sha256':c.digest(files[loc])} + def producer(files,spec,changes): + data=c.parse(files[spec['producer_execution']['path']]);data.update(changes) + spec['producer_execution'].update(save(files,spec['producer_execution']['path'],data)) + return data + oldroot={'path':old['root']+'/execution-claims','terminal_directory':old['terminal_directory']} + oldproducer=producer(oldfiles,old,dict(mode='diagnostic',execution_claim_roots=[oldroot])) + old['producer_execution']['mode']='diagnostic' + oldsb=save(oldfiles,args[0],old);oldargs=(oldsb['path'],oldsb['sha256'],args[2]) + previous=c.ledger(old);argv=c.controller_argv(old,*oldargs) + held=helper.native(old,Comment='elmo-controller:'+oldargs[1]);oldfiles[previous+'/HELD_NATIVE.txt']=held.encode() + prepared=dict(schema='native-elmo-controller-prepared-v1',spec_path=oldargs[0],spec_sha256=oldargs[1],ops_manifest_sha256=oldargs[2], + producer_spec_sha256=old['producer_execution']['sha256'],argv=argv) + save(oldfiles,previous+'/PREPARED.json',prepared) + save(oldfiles,previous+'/acceptance.json',dict(schema='native-elmo-controller-acceptance-v1',job_id='900001',spec_sha256=oldargs[1], + producer_spec_sha256=old['producer_execution']['sha256'],argv=argv,held_native_sha256=c.digest(held.encode()))) + for name,value in [('PARENT_JOB_ID',b'900001\n'),('sbatch.stdout',b'900001\n'),('sbatch.exit',b'0\n')]:oldfiles[previous+'/'+name]=value + save(oldfiles,previous+'/RESULT.json',dict(schema='native-elmo-controller-result-v1',controller_job_id='900001', + producer_spec_sha256=old['producer_execution']['sha256'],status='CHILD_RELEASE_ACCEPTED_NOT_COMPLETION')) + helper.ROOT=c.PREFIX+'synthetic-current-controller';newfiles,current,newargs=helper.fixture() + current['parent_request']=old['parent_request'] + newproducer=c.parse(newfiles[current['producer_execution']['path']]);newproducer['parent_request']=old['parent_request'] + newproducer['execution_claim_roots']=[oldroot,{'path':current['root']+'/execution-claims','terminal_directory':current['terminal_directory']}] + newproducer['runtime']['admission']['quiescence']['claim_roots']=[{'path':old['root']+'/scientific-claims','terminal_directory':old['terminal_directory']}] + current['producer_execution'].update(save(newfiles,current['producer_execution']['path'],newproducer)) + allfiles={**oldfiles,**newfiles} + remote=SimpleNamespace(census=lambda _:old['producer_execution']['sha256']+'|d\n') + return c,allfiles,current,old,previous,newproducer,remote + + +@pytest.mark.parametrize('omit',['execution-root','terminal-directory','scientific-root']) +def test_gateway_new_root_cannot_omit_purged_prior_child(monkeypatch,omit): + c,files,current,old,ledger,newproducer,remote=controller_history_fixture(monkeypatch) + assert c.prior_census(c.Reader(files.__getitem__),remote,current) + if omit=='execution-root':newproducer['execution_claim_roots']=newproducer['execution_claim_roots'][1:] + elif omit=='terminal-directory':newproducer['execution_claim_roots'][0]['terminal_directory']=current['terminal_directory'] + else:newproducer['runtime']['admission']['quiescence']['claim_roots']=[] + files[current['producer_execution']['path']]=c.encode(newproducer) + current['producer_execution']['sha256']=c.digest(files[current['producer_execution']['path']]) + # Empty native queue / disappeared child is immaterial: history cannot be omitted. + with pytest.raises(ValueError,match='omits'):c.prior_census(c.Reader(files.__getitem__),remote,current) + + +@pytest.mark.parametrize('broken',['prepared','spec','parent','acceptance','held','unknown-result']) +def test_gateway_prior_intent_and_acceptance_are_reopened(monkeypatch,broken): + c,files,current,old,ledger,newproducer,remote=controller_history_fixture(monkeypatch) + if broken=='prepared':files[ledger+'/PREPARED.json']=b'{}\n' + elif broken=='spec': + p=c.parse(files[ledger+'/PREPARED.json'])['spec_path'];files[p]+=b' ' + elif broken=='parent': + p=c.parse(files[ledger+'/PREPARED.json']);s=c.parse(files[p['spec_path']]);s['parent_request']['sha256']='f'*64 + files[p['spec_path']]=c.encode(s);p['spec_sha256']=c.digest(files[p['spec_path']]);files[ledger+'/PREPARED.json']=c.encode(p) + elif broken=='acceptance':files[ledger+'/PARENT_JOB_ID']=b'900099\n' + elif broken=='held':files[ledger+'/HELD_NATIVE.txt']+=b' conflict\n' + else:files.pop(ledger+'/RESULT.json') + with pytest.raises((ValueError,KeyError)):c.prior_census(c.Reader(files.__getitem__),remote,current) + +@pytest.mark.parametrize('consumer',['execution','scientific']) +def test_fresh_array_alias_cannot_escape_complete_step_census(tmp_path,monkeypatch,consumer): + f=fixture(tmp_path,monkeypatch,'diagnostic');result,proof=actual_service_capture(f,monkeypatch) + monkeypatch.setattr(subprocess,'run',lambda argv,**kw:subprocess.CompletedProcess(argv,0,'900001_7.batch|RUNNING\n' if '--steps' in argv else '','')) + scope={'user':'paphiwetsa3','scientific_job_names':[f['spec']['native_contract']['JobName']], + 'scientific_name_prefixes':['elmo-large314m-continuous-'], + 'claim_roots':[{'path':str(f['root']/'scientific-claims'),'terminal_directory':str(f['root']/'status')}]} + with pytest.raises(ValueError,match='native workers'): + if consumer=='execution':f['e'].execution_claim_census(f['spec']) + else:f['h'].quiescence(f['spec'],scope,'run-id') + + +@pytest.mark.parametrize('member',['native_scontrol.sh','native_squeue.sh']) +def test_native_file_mount_wrappers_preserve_executable_mode(tmp_path,monkeypatch,member): + f=fixture(tmp_path,monkeypatch) + assert (OPS/member).stat().st_mode & 0o111 + client=load('_composite_native_client_mode_check',f['root']/'native_client.py') + client.executable_wrappers(f['root']) + (f['root']/member).chmod(0o644) + with pytest.raises(ValueError,match='not executable'):client.executable_wrappers(f['root']) diff --git a/tests/test_admission_history.py b/tests/test_admission_history.py new file mode 100644 index 000000000..3b1efe219 --- /dev/null +++ b/tests/test_admission_history.py @@ -0,0 +1,254 @@ +"""Synthetic local evidence tests; never real native/GPU certification.""" +import copy +import csv +import hashlib +import json +from pathlib import Path +import uuid + +import pytest + +from scripts.train.joint.interruption_admission import ( + fingerprint, compatible, completed_producer, current_intent, quiescence, + resume_native_contract, +) +from scripts.train.joint.production_history import retained_milestones, scheduled_validation, segments +from scripts.train.joint.production_recovery import validate_phase_proofs +from scripts.train.joint.native_recovery import state_digest +from test_interruption import bind, request, milestone, expected, native + + +def intent(tmp_path): + value = dict(schema="native-elmo-current-intent-v1", authority="human", status="ACTIVE", + authorized_recovery_run_ids=["run"], retired_run_ids=[], hold_recovery=False, + source_reference="synthetic-test-fixture", updated_at="2026-10-26T00:00:00Z") + path = tmp_path / "intent.json"; path.write_text(json.dumps(value)) + return path, value + + +def test_later_human_intent_precedes_finalized_admission(tmp_path): + path, value = intent(tmp_path) + digest = current_intent(path, "run") + value["retired_run_ids"] = ["run"] + path.write_text(json.dumps(value)) + with pytest.raises(ValueError): current_intent(path, "run", digest) + with pytest.raises(ValueError): current_intent(path, "run") + + +@pytest.mark.parametrize("change", [dict(status="PAUSED"), dict(hold_recovery=True), + dict(authority="agent"), dict(authorized_recovery_run_ids=[])]) +def test_intent_cannot_be_inferred_or_overridden(tmp_path, change): + path, value = intent(tmp_path); value.update(change); path.write_text(json.dumps(value)) + with pytest.raises(ValueError): current_intent(path, "run") + + +def pair(tmp_path): + p = request(tmp_path) + p.update(world_size=4, job_name="test-parent", local_logs=str(tmp_path / "logs1"), + operational_service=dict(manifest_sha256="old"), lr_policy=dict(peak_lr=.000005, + warmup_optimizer_steps=3000, schedule="linear_warmup_then_constant"), + overrides=["old-output"]) + c = copy.deepcopy(p) + c.update(job_name="test-child", output=str(tmp_path / "new-output"), local_logs=str(tmp_path / "logs2"), + operational_service=dict(manifest_sha256="new"), overrides=["new-output"], + interruption=dict(parent_request=bind(tmp_path / "parent.json", p), checkpoint=dict(global_step=20000), admission=None)) + return p, c + + +def test_fingerprint_excludes_only_admission_and_science_remains_bound(tmp_path): + p, c = pair(tmp_path) + compatible(p, c) + digest = fingerprint(c) + c["interruption"]["admission"] = dict(real_completed_producer="future-test-only") + assert fingerprint(c) == digest + c["lr_policy"]["peak_lr"] = .0001 + assert fingerprint(c) != digest + with pytest.raises(ValueError): compatible(p, c) + + +def test_future_production_native_contract_has_fresh_command_and_exact_resources(): + parent = expected() + child = dict(job_name="new-scientific-attempt") + future = {key: value for key, value in parent.items() if key != "JobId"} + future.update(JobName=child["job_name"], Command="/synthetic-new-attempt/full.sbatch") + resume_native_contract(future, parent, child) + for changes in (dict(NumCPUs="16"), dict(NodeList="gpu12"), dict(UserId="other(1235)"), + dict(TimeLimit="00:20:00"), dict(Command=parent["Command"]), dict(JobId="123")): + with pytest.raises(ValueError): resume_native_contract(dict(future, **changes), parent, child) + + +@pytest.mark.parametrize("field,value", [("parameters", 1), ("world_size", 2), ("target_steps", 600000), + ("source_revision", "different"), ("data_manifest_sha256", "different"), ("phase", "smoke"), + ("allocation_stop_step", 200000)]) +def test_request_compatibility_rejects_changed_science_or_budget(tmp_path, field, value): + p, c = pair(tmp_path); c[field] = value + with pytest.raises(ValueError): compatible(p, c) + + +def bundle(tmp_path): + q = dict(producer_native=expected()) + request_binding = bind(tmp_path / "producer.json", q) + result_binding = bind(tmp_path / "result.json", dict(job_id="123", request_sha256=request_binding["sha256"], + schema="test-only", status="DONE_NOT_NATIVE_PROOF")) + n = tmp_path / "native.txt"; n.write_text(native(state="COMPLETED", ExitCode="0:0")) + h = tmp_path / "host.txt"; h.write_text("0\n") + return dict(request=request_binding, result=result_binding, + native=dict(path=str(n), sha256=hashlib.sha256(n.read_bytes()).hexdigest()), + host_exit=dict(path=str(h), sha256=hashlib.sha256(h.read_bytes()).hexdigest())) + + +def test_self_authored_status_is_not_native_completion(tmp_path): + b = bundle(tmp_path) + completed_producer(b, "test-only", "DONE_NOT_NATIVE_PROOF") + n = Path(b["native"]["path"]); n.write_text(native(state="FAILED", ExitCode="1:0")) + b["native"]["sha256"] = hashlib.sha256(n.read_bytes()).hexdigest() + with pytest.raises(ValueError): completed_producer(b, "test-only", "DONE_NOT_NATIVE_PROOF") + + +def test_host_exit_also_required(tmp_path): + b = bundle(tmp_path) + h = Path(b["host_exit"]["path"]); h.write_text("1\n") + b["host_exit"]["sha256"] = hashlib.sha256(h.read_bytes()).hexdigest() + with pytest.raises(ValueError): completed_producer(b, "test-only", "DONE_NOT_NATIVE_PROOF") + + +def test_quiescence_rechecks_unknown_claims_and_current_writer(tmp_path, monkeypatch): + import scripts.train.joint.interruption_admission as module + root = tmp_path / "claims"; root.mkdir() + spec = dict(user="test", claim_roots=[dict(path=str(root), terminal_directory=str(tmp_path))], + scientific_job_names=["test-parent"], scientific_name_prefixes=["test-elmo-scientific-"]) + monkeypatch.setattr(module, "command", lambda argv: "") + assert quiescence(spec, "run")["claims"] == [] + claim = root / ("a" * 64); claim.mkdir() + (claim / "PREPARED.json").write_text(json.dumps(dict(request_sha256=claim.name))) + with pytest.raises(FileNotFoundError): quiescence(spec, "run") + (claim / "PARENT_JOB_ID").write_text("123\n") + with pytest.raises(ValueError): quiescence(spec, "run") + monkeypatch.setattr(module, "command", lambda argv: "123|test-parent|RUNNING|run\n") + assert quiescence(spec, "run", current_job="123", current_request_sha=claim.name)["claims"][0]["current"] + with pytest.raises(ValueError): quiescence(spec, "run", current_job="123", current_request_sha="b" * 64) + with pytest.raises(ValueError): quiescence(spec, "run") + + +def csv_at(q, start, rows): + path = Path(q["local_logs"]) / "csv" / f"step{start:09d}" / "metrics.csv" + path.parent.mkdir(parents=True) + with path.open("w") as stream: + writer = csv.DictWriter(stream, fieldnames=["step", "val/a", "val/b"]) + writer.writeheader(); writer.writerows(rows) + return path + + +def history(tmp_path): + p, c = pair(tmp_path) + (Path(c["output"]) / "checkpoints").mkdir(parents=True) + milestone(p, 20000); milestone(c, 40000) + csv_at(p, 0, [dict(step=19999, **{"val/a": .1, "val/b": .2})]) + csv_at(c, 20000, [dict(step=40000, **{"val/a": .05, "val/b": .1})]) + return p, c + + +def test_history_retains_immutable_prefix_and_new_attempt(tmp_path): + p, c = history(tmp_path) + checkpoints, audit = retained_milestones(c, 40000) + assert set(checkpoints) == {20000, 40000} + assert checkpoints[20000].parent == Path(p["output"]) / "checkpoints" + assert checkpoints[40000].parent == Path(c["output"]) / "checkpoints" + values, logs = scheduled_validation(c, 40000, {"val/a", "val/b"}) + assert set(values) == {20000, 40000} and len(logs) == 2 + + +@pytest.mark.parametrize("damage", ["repeat_boundary", "repeat_metric", "nonfinite", "partial", "off_schedule", "missing"]) +def test_validation_merge_rejects_invalid_history(tmp_path, damage): + p, c = history(tmp_path) + path = segments(c, 40000)[-1]["csv"] + if damage == "repeat_boundary": path.write_text("step,val/a,val/b\n20000,0.1,0.1\n40000,0.1,0.1\n") + elif damage == "repeat_metric": path.write_text("step,val/a,val/b\n40000,0.1,0.1\n40000,0.1,\n") + elif damage == "nonfinite": path.write_text("step,val/a,val/b\n40000,nan,0.1\n") + elif damage == "partial": path.write_text("step,val/a,val/b\n40000,0.1,\n") + elif damage == "off_schedule": path.write_text("step,val/a,val/b\n40001,0.1,0.1\n") + else: path.unlink() + with pytest.raises(ValueError): scheduled_validation(c, 40000, {"val/a", "val/b"}) + + +def test_uncommitted_validation_nonfinite_is_not_hidden_by_checkpoint_selection(tmp_path): + p, c = history(tmp_path) + path = segments(c, 40000)[0]["csv"] + with path.open("a") as stream: + stream.write("40000,nan,0.1\n") + with pytest.raises(ValueError, match="including uncommitted"): + scheduled_validation(c, 40000, {"val/a", "val/b"}) + + +def production_rows(): + from test_native_conditional_recovery import phase_rows + rows = phase_rows() + names = rows[0][0]["conditional_optimizer_record"]["parameter_names"] + start = 20000 + for phase, (first, last) in zip(rows, ((start, start + 2), (start + 2, start + 6))): + for rank, row in enumerate(phase): + row["initial_parameter_usage"] = dict(shared=first, private=first // 2, unused=0) + row["conditional_optimizer_record"].update(global_step=last, + parameter_update_counts=dict(shared=last, private=last // 2, unused=0)) + row["initial_weights_sha256"] = "w" if first == start else "v" + row["weights_sha256"] = "v" if first == start else "z" + row["gpu_uuid"] = str(uuid.UUID(int=rank + 1)) + selected = dict(global_step=start, proof=dict(model_state_sha256="w", rank_rng_sha256=["a" * 64] * 4, + conditional_optimizer_usage_sha256=state_digest(dict(schema="native-conditional-optimizer-usage-v1", + global_step=start, parameter_names=names, parameter_update_counts=rows[0][0]["initial_parameter_usage"])))) + return rows, selected + + +def test_production_origin_has_separate_scope_and_inspected_state_binding(): + rows, selected = production_rows() + assert validate_phase_proofs(rows, selected)["diagnostic_scope"] == "production-checkpoint-only" + selected["proof"]["rank_rng_sha256"][0] = "b" * 64 + with pytest.raises(ValueError): validate_phase_proofs(rows, selected) + + +@pytest.mark.parametrize("damage", ["counts", "weights", "duplicate_gpu", "changed_gpu", "counter_chain"]) +def test_actual_diagnostic_must_bind_inspected_state_and_all_ranks(damage): + rows, selected = production_rows() + if damage == "counts": selected["proof"]["conditional_optimizer_usage_sha256"] = "wrong" + elif damage == "weights": rows[0][0]["initial_weights_sha256"] = "wrong" + elif damage == "duplicate_gpu": rows[0][1]["gpu_uuid"] = rows[0][0]["gpu_uuid"] + elif damage == "changed_gpu": rows[1][0]["gpu_uuid"] = str(uuid.UUID(int=5)) + else: rows[1][0]["initial_parameter_usage"]["private"] -= 1 + with pytest.raises(ValueError): validate_phase_proofs(rows, selected) + + +def test_interruption_does_not_require_or_fake_boundary_completion(monkeypatch): + import scripts.train.joint.interruption_admission as gate + from scripts.train.joint.continuation import allocation_steps + seen = [] + monkeypatch.setattr(gate, "admission", lambda request: seen.append(request) or (20000,300000)) + request = dict(target_steps=300000, interruption=dict(checkpoint=dict(global_step=20000))) + assert allocation_steps(request) == (20000,300000) and seen == [request] + request["continuation"] = dict(fake="boundary") + with pytest.raises(ValueError): allocation_steps(request) + + +def test_configure_resume_uses_selected_production_checkpoint_and_same_history(): + from omegaconf import OmegaConf + from scripts.train.joint.continuation import configure_allocation + cfg = OmegaConf.create(dict(logger=dict(csv=dict(save_dir="old",version="old"),wandb=dict(resume="never")), + ckpt_path=None, callbacks={})) + q = dict(local_logs="fresh", target_steps=300000, interruption=dict(checkpoint=dict(path="retained-production.ckpt"))) + configure_allocation(cfg, q, 20000, 300000) + assert cfg.ckpt_path == "retained-production.ckpt" and cfg.logger.wandb.resume == "must" + assert cfg.logger.csv.version == "step000020000" and not cfg.callbacks + + +def test_existing_dispatcher_routes_production_mode_without_legacy_repair(monkeypatch, tmp_path): + import sys + import scripts.train.joint.recovery_check as dispatcher + import scripts.train.joint.production_recovery as production + path = tmp_path / "request.json" + path.write_text(json.dumps(dict(schema="native-elmo-production-recovery-request-v1"))) + calls = [] + monkeypatch.setattr(production, "main", lambda p: calls.append(p)) + monkeypatch.setattr(sys, "argv", ["recovery_check", str(path)]) + dispatcher.main() + assert calls == [path] + monkeypatch.setattr(sys, "argv", ["recovery_check", str(path), "--postflight", "123"]) + with pytest.raises(ValueError): dispatcher.main() diff --git a/tests/test_interruption.py b/tests/test_interruption.py new file mode 100644 index 000000000..d9cc1252e --- /dev/null +++ b/tests/test_interruption.py @@ -0,0 +1,316 @@ +"""CPU contract tests; fixtures are deliberately never runtime proof artifacts.""" +import copy +import csv +import hashlib +import io +import json +import os +from pathlib import Path +import random + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +from scripts.train.joint.native_interruption import ( + fields, native_identity, classify_failure, check_steps, check_queue, read_tail, + newest_loss, loss_snapshot, destination_probe, +) +from scripts.train.joint.checkpoint_inspection import attempt_chain, checkpoint_census, check_rng, scheduled_filename +from scripts.train.joint.production_recovery import phases, phase_config +from scripts.train.joint.native_recovery import state_digest, summarize_native_phases + + +def expected(): + return dict(JobId="123", JobName="test-only-native-elmo", UserId="test(1234)", Account="loaner", + QOS="loaner-normal", Partition="rl2", ReqNodeList="gpu11", NodeList="gpu11", NumCPUs="32", + MinMemoryNode="192G", TimeLimit="7-00:00:00", Requeue="0", Restarts="0", + TresPerNode="gres/gpu:h100:4", Command="/synthetic-test-only/full.sbatch") + + +def native(state="TIMEOUT", **overrides): + row = dict(expected(), JobState=state, NumNodes="1", NumTasks="1", ExitCode="0:15", + EndTime="2026-10-26T08:00:00", ReqTRES="cpu=32,mem=192G,node=1,gres/gpu=4,gres/gpu:h100=4", + AllocTRES="cpu=32,mem=192G,node=1,gres/gpu=4,gres/gpu:h100=4") + row.update(overrides) + row.setdefault("JOB_GRES", "gpu:h100:4" if row["TresPerNode"] else "(null)") + row.setdefault("GRES", "gpu:h100:4(IDX:0-3)" if state == "RUNNING" and row["TresPerNode"] else "") + return " ".join(k + "=" + v for k, v in row.items()) + + +def test_terminal_raw_not_status_bool(): + record = native_identity(native(), expected(), "TIMEOUT") + assert classify_failure(record, "stepd terminated job due to time limit") == "TIMEOUT" + with pytest.raises(ValueError): fields(native() + " JobId=999") + + +def test_actual_lambda_cpu_record_has_no_tres_per_node_key(): + e = dict(expected(), NumCPUs="4", MinMemoryNode="24G", TimeLimit="00:20:00", TresPerNode="") + text = native(state="COMPLETED", NumCPUs="4", MinMemoryNode="24G", TimeLimit="00:20:00", + TresPerNode="", ReqTRES="cpu=4,mem=24G,node=1,billing=4", AllocTRES="cpu=4,mem=24G,node=1,billing=4", + ExitCode="0:0") + assert native_identity(text, e, "COMPLETED")["ExitCode"] == "0:0" + + +@pytest.mark.parametrize("change", [dict(JobId="999"), dict(UserId="other(1235)"), dict(Restarts="1"), + dict(Requeue="1"), dict(TimeLimit="01:00:00"), dict(NodeList="gpu12"), dict(NumCPUs="16"), + dict(MinMemoryNode="96G"), dict(EndTime="Unknown"), dict(ExitCode=""), dict(NumNodes="2"), + dict(AllocTRES="cpu=32,mem=192G,node=1,gres/gpu=8,gres/gpu:h100=8"), dict(JobState="COMPLETING")]) +def test_actual_native_shape_and_terminal_required(change): + with pytest.raises(ValueError): native_identity(native(**change), expected(), "TIMEOUT") + + +@pytest.mark.parametrize("state", ["CANCELLED", "OUT_OF_MEMORY", "COMPLETED", "FAILED"]) +def test_no_arbitrary_automatic_restart(state): + with pytest.raises(ValueError): classify_failure(dict(JobState=state, ExitCode="1:0"), "training stopped") + + +def test_verified_nccl_crash_and_scientific_veto(): + row = dict(JobState="FAILED", ExitCode="0:6") + assert classify_failure(row, "Watchdog caught collective operation timeout") == "VERIFIED_NCCL_WORKER_CRASH" + with pytest.raises(ValueError): classify_failure(row, "Watchdog caught collective operation timeout\nValueError: nonfinite loss") + with pytest.raises(ValueError): classify_failure(dict(JobState="TIMEOUT", ExitCode="0:15"), "OutOfMemoryError") + + +@pytest.mark.parametrize("tail", ["SIGABRT rank0 pid123", "exitcode : -6 rank0", "ProcessGroupNCCL terminated"]) +def test_generic_abort_is_not_nccl_timeout_evidence(tail): + with pytest.raises(ValueError): classify_failure(dict(JobState="FAILED", ExitCode="0:6"), tail) + + +@pytest.mark.parametrize("tail", ["train/loss=nan", "Train/Loss: +Inf", "joint_proprio_mse=-infinity", + "loss tensor(nan)", "train/loss = -nan"]) +def test_infrastructure_terminal_does_not_override_nonfinite_loss(tail): + with pytest.raises(ValueError): classify_failure(dict(JobState="TIMEOUT", ExitCode="0:15"), tail) + + +def test_lambda_gpu_tres_omission_is_bound_by_actual_gres(): + for state in ("RUNNING", "COMPLETED"): + raw = native(state=state, ReqTRES="cpu=32,mem=192G,node=1,billing=32", + AllocTRES="cpu=32,mem=192G,node=1,billing=32") + native_identity(raw, expected(), state) + with pytest.raises(ValueError): + native_identity(native(state="RUNNING", GRES=""), expected(), "RUNNING") + for gres in ("gpu:h100:4(IDX:0-2)", "gpu:h100:4(IDX:0-2,2)", "gpu:a100:4(IDX:0-3)"): + with pytest.raises(ValueError): native_identity(native(state="RUNNING", GRES=gres), expected(), "RUNNING") + with pytest.raises(ValueError): native_identity(native(JOB_GRES="gpu:h100:8"), expected(), "TIMEOUT") + + +LOSS_HEADER = "step,Train/loss,Train/joint_image_mse,Train/joint_proprio_mse,Train/joint_action_mse,Train/joint_camera_mse,Train/source_0_loss_action,Train/joint_image_mse/eva_bimanual,Train/joint_proprio_mse/eva_bimanual,Train/joint_camera_mse/eva_bimanual,Train/joint_action_mse/eva_bimanual,lr-AdamW\n" + + +def loss_row(step, value="1"): + return str(step) + "," + ",".join([value] * 10) + ",\n" + + +@pytest.mark.parametrize("bad", ["nan", "Inf", "-inf", "infinity"]) +def test_newest_loss_selected_before_finiteness_and_not_lr_row(bad): + with pytest.raises(ValueError, match="nonfinite"): + newest_loss(dict(local=dict(header=LOSS_HEADER, tail=loss_row(20) + loss_row(21, bad) + "22,,,,,,,,,,,0.0001\n"))) + result = newest_loss(dict(local=dict(header=LOSS_HEADER, tail=loss_row(20, bad) + loss_row(21)))) + assert result["step"] == 21 and result["finite"] is True + + +def test_loss_snapshot_uses_exact_parent_namespace_and_bounded_tail(tmp_path): + parent = dict(local_logs=str(tmp_path / "logs"), output=str(tmp_path / "output")) + local = Path(parent["local_logs"]) / "csv/step000000000/metrics.csv" + local.parent.mkdir(parents=True) + local.write_text(LOSS_HEADER + "".join(loss_row(step) for step in range(100))) + sources, snapshots = loss_snapshot(parent) + assert len(snapshots["local"]["tail"].splitlines()) == 64 + assert newest_loss(snapshots)["step"] == 99 and sources["mirror"]["exists"] is False + with pytest.raises(ValueError, match="absence is not zero-update"): + newest_loss({}) + + +def test_loss_sources_latest_conflict_and_missing_modality_fail_closed(): + with pytest.raises(ValueError, match="disagree"): + newest_loss(dict(local=dict(header=LOSS_HEADER, tail=loss_row(1)), mirror=dict(header=LOSS_HEADER, tail=loss_row(1, "2")))) + with pytest.raises(ValueError, match="lacks"): + newest_loss(dict(local=dict(header=LOSS_HEADER, tail="1,1,,,,,,,,,,\n"))) + result = newest_loss(dict(local=dict(header=LOSS_HEADER, tail=loss_row(2)), mirror=dict(header=LOSS_HEADER, tail=loss_row(1)))) + assert result["selected_rows"][0]["source"] == "local" + + +def test_actual_native_csv_human_rows_and_stale_other_embodiment_values(): + lines = (Path(__file__).parent / "fixtures/actual-smoke-first-four-updates.csv").read_text().splitlines() + columns = next(csv.reader([lines[0]])); index = columns.index("step") + for step, embodiment in ((0, "human_bimanual"), (1, "human_bimanual"), (2, "eva_bimanual"), (3, "eva_bimanual")): + rows = [line for line in lines[1:] if next(csv.reader([line]))[index] == str(step)] + result = newest_loss(dict(actual=dict(header=lines[0], tail="\n".join(rows)))) + assert result["active_embodiment"] == embodiment + assert result["supervised_action"] == (embodiment == "eva_bimanual") + + +@pytest.mark.parametrize("value", ["nan", "inf", "-inf", "0.5"]) +def test_observed_validation_veto_is_independent_of_latest_finite_training(value): + lines = (Path(__file__).parent / "fixtures/actual-smoke-first-four-updates.csv").read_text().splitlines() + columns = next(csv.reader([lines[0]])) + buffer = io.StringIO(); writer = csv.DictWriter(buffer, fieldnames=columns) + writer.writerow({"step": "19999", "val/human_bimanual/forecast/proprio_normalized_mse": value}) + snapshots = dict(actual=dict(header=lines[0], tail="\n".join(lines[1:]) + "\n" + buffer.getvalue())) + if value == "0.5": + assert newest_loss(snapshots)["finite"] + assert classify_failure(dict(JobState="TIMEOUT", ExitCode="0:15"), "time limit") == "TIMEOUT" + else: + with pytest.raises(ValueError, match="validation is nonfinite"): + newest_loss(snapshots) + + +@pytest.mark.parametrize("name", ["epoch_epoch=0000-step_step=000020000.ckpt", "epoch_0000-step_000020000.ckpt"]) +def test_maintained_template_metric_insertion_variants(name): + assert scheduled_filename(name) == 20000 + with pytest.raises(ValueError): scheduled_filename(name.replace(".ckpt", "-v1.ckpt")) + + +def test_destination_probe_is_preserved_exclusive_and_requires_space(tmp_path): + evidence = destination_probe(str(tmp_path / "new-run"), "123", 1) + assert evidence["actual_write_fsync_readback"] and Path(evidence["probe_path"]).is_file() + with pytest.raises(FileExistsError): destination_probe(str(tmp_path / "new-run"), "123", 1) + with pytest.raises(ValueError): destination_probe(str(tmp_path / "new-run"), "124", 10**30) + + +def test_live_steps_and_exact_writer_queue(): + assert check_steps("", "123") + with pytest.raises(ValueError): check_steps("123.batch|RUNNING\n", "123") + assert check_queue("124|other-project|RUNNING\n", ["this-writer"]) + with pytest.raises(ValueError): check_queue("125|this-writer|PENDING\n", ["this-writer"]) + assert check_queue("125|this-writer|RUNNING\n", ["this-writer"], "125") + with pytest.raises(ValueError): check_queue("125|this-writer|PENDING\n", ["this-writer"], "125") + + +def test_tail_bounded_and_symlinks_rejected(tmp_path): + path = tmp_path / "worker.log" + path.write_text("start\n" * 100000 + "\n".join(str(i) for i in range(100))) + assert read_tail(path).splitlines() == [str(i) for i in range(36, 100)] + linked = tmp_path / "link" + linked.symlink_to(path) + with pytest.raises(ValueError): read_tail(linked) + + +def request(tmp_path, name="attempt1"): + output = tmp_path / name + (output / "checkpoints").mkdir(parents=True) + return dict(schema="joint-training-v2", phase="full", kind="native-subtask-world-model-v1", + run_id="synthetic-test-only", source_revision="s", source_manifest_sha256="m", + data_manifest_sha256="d", parameters=314187064, target_steps=300000, output=str(output)) + + +def bind(path, value): + path.write_text(json.dumps(value)) + return dict(path=str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def milestone(q, step, data=b"test-only"): + p = Path(q["output"]) / "checkpoints" / f"epoch_epoch=0000-step_step={step:09d}.ckpt" + p.write_bytes(data) + return p + + +def test_census_preserves_every_attempt_and_selects_order_not_last_link(tmp_path): + parent = request(tmp_path) + old = milestone(parent, 20000) + newer = milestone(parent, 40000) + (old.parent / "last.ckpt").symlink_to(old.name) + (old.parent / "incomplete.tmp").write_bytes(b"") + child = request(tmp_path, "attempt2") + child["interruption"] = dict(parent_request=bind(tmp_path / "parent.json", parent)) + milestone(child, 60000) + census = checkpoint_census(child) + assert [r["global_step"] for r in census["candidates"]] == [60000, 40000, 20000] + assert census["preserved_noncheckpoint_entries"][0]["size_bytes"] == 0 + assert newer.exists() and old.exists() + + +@pytest.mark.parametrize("damage", ["missing_collection", "off_cadence", "symlink", "archive", "identity", "ambiguous_parent"]) +def test_census_unknowns_fail_closed(tmp_path, damage): + q = request(tmp_path) + path = milestone(q, 20000) + if damage == "missing_collection": q["output"] = str(tmp_path / "absent") + elif damage == "off_cadence": milestone(q, 20001) + elif damage == "symlink": (path.parent / "epoch_epoch=0000-step_step=000040000.ckpt").symlink_to(path) + elif damage == "archive": q["archive_policy"] = dict(status="PASS") + elif damage == "identity": + parent = dict(q, run_id="another") + q["interruption"] = dict(parent_request=bind(tmp_path / "parent.json", parent)) + else: + q["interruption"] = q["continuation"] = dict(parent_request={}) + with pytest.raises((ValueError, FileNotFoundError)): checkpoint_census(q) + + +def rng_fixture(): + row = dict(torch_cpu=torch.get_rng_state(), torch_cuda=[torch.zeros(16, dtype=torch.uint8)], + numpy=np.random.get_state(), python=random.getstate(), cuda_scope="current_device") + return dict(schema="mot_rng_state.distributed.v1", world_size=4, + ranks=[copy.deepcopy(row) for _ in range(4)]) + + +def test_rng_cpu_structure_check_preserves_all_global_rng(): + saved = rng_fixture() + before = state_digest(saved) + check_rng(saved, 4) + assert state_digest(rng_fixture()) == before + + +@pytest.mark.parametrize("damage", ["missing_rank", "empty_cuda", "float_cuda", "bad_python", "bad_numpy", "bad_cpu", "scope"]) +def test_rng_structure_corruption_rejected(damage): + saved = rng_fixture() + if damage == "missing_rank": saved["ranks"].pop() + elif damage == "empty_cuda": saved["ranks"][0]["torch_cuda"] = [] + elif damage == "float_cuda": saved["ranks"][0]["torch_cuda"][0] = torch.zeros(16) + elif damage == "bad_python": saved["ranks"][0]["python"] = (0, (), None) + elif damage == "bad_numpy": saved["ranks"][0]["numpy"] = ("invalid",) + elif damage == "bad_cpu": saved["ranks"][0]["torch_cpu"] = torch.zeros(1, dtype=torch.uint8) + else: saved["ranks"][0]["cuda_scope"] = "all_devices" + with pytest.raises((ValueError, RuntimeError, IndexError)): check_rng(saved, 4) + + +def original_config(): + return OmegaConf.create(dict(model=dict(loss=dict(kind="fused", cap=5, tau=1, scale=.73)), + data=dict(sampling=dict(boundary_probability=.25)), paths=dict(output_dir="old"), + ckpt_path=None, val_at_start=False, val_at_end=False, + trainer=dict(max_steps=300000, val_check_interval=20000, limit_val_batches=25, + num_sanity_val_steps=0, check_val_every_n_epoch=None), + callbacks=dict(conditional=dict(_target_="egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage"), + model_checkpoint=dict(every_n_train_steps=20000, every_n_epochs=0, save_top_k=-1, + save_last="link", dirpath="old", filename="old"), + wandb_profiler=dict(_target_="do-not-invoke")), + logger=dict(wandb=dict(id="must-not-write-history"), csv=dict(save_dir="old", version="old")), + evaluator=dict(pixel_diagnostics=None))) + + +@pytest.mark.parametrize("origin", [20000, 40000, 280000]) +def test_production_phases_preserve_model_data_and_disable_history(tmp_path, origin): + original = original_config() + before = OmegaConf.to_container(original, resolve=True) + for start, stop in phases(origin, 300000): + cfg = phase_config(original, output=tmp_path, checkpoint="production.ckpt", origin=origin, + start=start, stop=stop, target=300000, world=4, proof=tmp_path / "proof") + assert cfg.model == original.model and cfg.data == original.data + assert "wandb" not in cfg.logger and "wandb_profiler" not in cfg.callbacks + assert cfg.trainer.max_steps == origin + 6 and cfg.trainer.val_check_interval == 2 + assert cfg.callbacks.resume_proof.multimodal + assert ("allocation_boundary" in cfg.callbacks) == (stop == origin + 2) + assert OmegaConf.to_container(original, resolve=True) == before + + +@pytest.mark.parametrize("origin,target", [(0,300000), (50,300000), (300000,300000), (20001,300000), (20000,300001)]) +def test_smoke_zero_terminal_and_off_schedule_rejected(origin, target): + with pytest.raises(ValueError): phases(origin, target) + + +def test_production_diagnostic_rejects_wrong_rank_or_counts(tmp_path): + for world, start, stop in [(2,20000,20002), (4,20000,20004), (4,50,52)]: + with pytest.raises(ValueError): phase_config(original_config(), output=tmp_path, checkpoint="p", + origin=20000, start=start, stop=stop, target=300000, world=world, proof=tmp_path / "proof") + + +def test_import_does_not_execute_remote_or_load_checkpoint(monkeypatch, tmp_path): + import scripts.train.joint.production_recovery as module + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + called = [] + monkeypatch.setattr(torch, "load", lambda *a, **k: called.append("unsafe-load")) + q = dict(schema="native-elmo-production-recovery-request-v1", producer_native=expected()) + path = tmp_path / "unallocated.json"; path.write_text(json.dumps(q)) + with pytest.raises(ValueError): module.main(path) + assert not called diff --git a/tests/test_native_conditional_recovery.py b/tests/test_native_conditional_recovery.py new file mode 100644 index 000000000..ae6ca5104 --- /dev/null +++ b/tests/test_native_conditional_recovery.py @@ -0,0 +1,297 @@ +import copy +import random +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +from egomimic.pl_utils.joint_smoke import ConditionalOptimizerUsage +from egomimic.pl_utils import world_model_behavior as wb +from scripts.train.joint.native_recovery import ( + conditional_record, restored_conditional_counts, install_rng_observer, + state_digest, preserved_rng, expected_metrics, validate_hidden_sample, + summarize_native_phases, checkpoint_science, +) +from scripts.train.joint.recovery_check import ResumeOptimizerProof, equal_state, phase_config + + +class TinyConditional(torch.nn.Module): + def __init__(self): + super().__init__() + self.shared = torch.nn.Parameter(torch.ones(2)) + self.proprio_banks = torch.nn.ModuleDict({name: torch.nn.Linear(2, 2) + for name in ("eva_bimanual", "human_bimanual")}) + self.action = torch.nn.Linear(2, 2) + self.camera = torch.nn.Linear(2, 2) + self.never = torch.nn.Parameter(torch.ones(2)) + self.pose_dims = {"eva_bimanual": 2, "human_bimanual": 2} + self.action_dims = {"eva_bimanual": 2} + + def active_streams(self, emb): + return ("image", "proprio", "action", "camera") if emb == "eva_bimanual" else ("image", "proprio", "camera") + + def loss(self, emb): + value = self.shared.sum() + self.proprio_banks[emb](self.shared).sum() + self.camera(self.shared).sum() + if emb == "eva_bimanual": + value = value + self.action(self.shared).sum() + return value.square() + + +def trained(): + torch.manual_seed(7) + net = TinyConditional() + optimizer = torch.optim.AdamW(net.parameters(), lr=5e-6) + counts = {name: 0 for name, _ in net.named_parameters()} + for emb in ("eva_bimanual", "human_bimanual") * 2: + optimizer.zero_grad(set_to_none=True) + net.loss(emb).backward() + for name, p in net.named_parameters(): + counts[name] += int(p.grad is not None) + optimizer.step() + return net, optimizer, counts + + +def test_actual_conditional_adam_accepts_shared_private_and_never_used(): + net, optimizer, counts = trained() + record = conditional_record(optimizer, net.named_parameters(), counts, 4) + assert counts["shared"] == 4 and counts["never"] == 0 + assert counts["action.weight"] == 2 + assert counts["proprio_banks.human_bimanual.weight"] == 2 + assert optimizer.state.get(net.never) is None + clone = TinyConditional() + clone.load_state_dict(net.state_dict(), strict=True) + restored = torch.optim.AdamW(clone.parameters(), lr=5e-6) + restored.load_state_dict(optimizer.state_dict()) + cb = ConditionalOptimizerUsage() + trainer = SimpleNamespace(global_step=4, optimizers=[restored], callbacks=[cb]) + checkpoint = dict(global_step=4, optimizer_states=[optimizer.state_dict()], + native_conditional_optimizer_usage=record) + cb.on_load_checkpoint(trainer, clone, checkpoint) + cb.on_fit_start(trainer, clone) + actual, found = restored_conditional_counts(trainer, clone, checkpoint) + assert actual == counts and found is cb + restored.zero_grad(set_to_none=True) + clone.loss("human_bimanual").backward() + cb.on_before_optimizer_step(trainer, clone, restored) + restored.step() + trainer.global_step = 5 + after = conditional_record(restored, clone.named_parameters(), cb.counts, 5) + assert after["parameter_update_counts"]["action.weight"] == 2 + assert after["parameter_update_counts"]["shared"] == 5 + + +@pytest.mark.parametrize("damage", ["step", "missing_count", "bool_count", "moment_nan", "moment_dtype", "moment_shape", "unused_state", "order", "extra_state"]) +def test_conditional_state_rejects_corruption(damage): + net, optimizer, counts = trained() + if damage == "step": optimizer.state[net.shared]["step"] += 1 + elif damage == "missing_count": counts.pop("shared") + elif damage == "bool_count": counts["shared"] = True + elif damage == "moment_nan": optimizer.state[net.shared]["exp_avg"][0] = float("nan") + elif damage == "moment_dtype": optimizer.state[net.shared]["exp_avg"] = optimizer.state[net.shared]["exp_avg"].half() + elif damage == "moment_shape": optimizer.state[net.shared]["exp_avg"] = torch.zeros(3) + elif damage == "unused_state": optimizer.state[net.never] = copy.deepcopy(optimizer.state[net.shared]) + elif damage == "order": optimizer.param_groups[0]["params"].reverse() + elif damage == "extra_state": optimizer.state[987654] = copy.deepcopy(optimizer.state[net.shared]) + with pytest.raises((ValueError, KeyError)): + conditional_record(optimizer, net.named_parameters(), counts, 4) + + +@pytest.mark.parametrize("emb", ["eva_bimanual", "human_bimanual"]) +def test_resume_uses_multimodal_gradient_checks_and_counts(tmp_path, emb): + net = TinyConditional() + optimizer = torch.optim.AdamW(net.parameters(), lr=5e-6) + callback = ResumeOptimizerProof(output=tmp_path, world_size=4, start_step=50, + stop_step=52, checkpoint="unbound", multimodal=True) + callback.network = net + callback.lr_policy = dict(peak_lr=5e-6, warmup_optimizer_steps=0, schedule="constant") + callback.optimizer_lrs = [] + callback.parameter_usage = {n: 0 for n, _ in net.named_parameters()} + callback.embodiment_batches = {e: 0 for e in net.pose_dims} + callback.current_embodiment = emb + net.training_behavior = SimpleNamespace(_pending_rng=None) + net.loss(emb).backward() + callback.on_before_optimizer_step(SimpleNamespace(global_step=50), net, optimizer) + assert callback.gradients_checked and callback.parameter_usage["never"] == 0 + assert callback.parameter_usage["action.weight"] == int(emb == "eva_bimanual") + assert callback.embodiment_batches[emb] == 1 + assert callback.multimodal + + +def test_legacy_constructor_remains_opt_out(tmp_path): + assert not ResumeOptimizerProof(output=tmp_path, world_size=4, start_step=50, + stop_step=52, checkpoint="unbound").multimodal + + +@pytest.mark.parametrize("rank", range(4)) +def test_actual_scientific_cpu_rng_restore_observed_before_draw(monkeypatch, rank): + original = wb._capture_rng_state(current_cuda_only=True) + cuda = [torch.arange(8, dtype=torch.uint8)] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_rng_state", lambda: cuda[0].clone()) + monkeypatch.setattr(torch.cuda, "set_rng_state", lambda value: cuda.__setitem__(0, value.clone())) + monkeypatch.setattr(wb, "_distributed_world", lambda: (4, rank)) + try: + states = [] + for seed in range(4): + torch.manual_seed(130 + seed); np.random.seed(130 + seed); random.seed(130 + seed) + states.append(wb._capture_rng_state(current_cuda_only=True)) + saved = dict(schema="mot_rng_state.distributed.v1", world_size=4, ranks=states) + wb._restore_rng_state(states[rank]) + expected_draw = (torch.rand(2), np.random.rand(2), random.random()) + torch.manual_seed(999); np.random.seed(999); random.seed(999) + + class Context: + global_step = 50 + trainer = SimpleNamespace(datamodule=SimpleNamespace(checkpointable_sampling=False)) + def _default_training_step(self, batch, index): + return torch.rand(2), np.random.rand(2), random.random() + + context = Context() + context.training_behavior = wb.WorldModelTrainingBehavior(42) + context.training_behavior.bind(context) + context.training_behavior.on_load_checkpoint({"mot_rng_state": saved}) + record, cleanup = install_rng_observer(context, saved, equal_state) + actual = context.training_behavior.training_step({"human_bimanual": {}}, 0) + assert equal_state(actual, expected_draw) + assert record["observed"] and record["actual_sha256"] == record["expected_sha256"] + cleanup() + finally: + wb._restore_rng_state({**original, "torch_cuda": []}) + + +def test_rng_observer_fails_if_restore_was_not_executed(monkeypatch): + saved = wb._capture_rng_state(current_cuda_only=True) + saved["torch_cuda"] = [torch.zeros(4, dtype=torch.uint8)] + monkeypatch.setattr(wb, "_rank_rng_state", lambda value: saved) + module = SimpleNamespace(training_behavior=SimpleNamespace(_pending_rng={}), + _default_training_step=lambda: pytest.fail("must not run")) + record, cleanup = install_rng_observer(module, {}, equal_state) + with pytest.raises(ValueError, match="did not execute"): + module._default_training_step() + assert not record["observed"] + cleanup() + + +def test_rng_preserved_even_on_sampler_error(): + before = wb._capture_rng_state(current_cuda_only=True) + with pytest.raises(RuntimeError): + with preserved_rng(): + torch.rand(3); np.random.rand(3); random.random() + raise RuntimeError("sample failure") + assert equal_state(wb._capture_rng_state(current_cuda_only=True), before) + assert state_digest(before) == state_digest(copy.deepcopy(before)) + + +def original_config(mode): + return OmegaConf.create(dict( + model={"optimizer": {"lr": 5e-6}, "loss": {"mode": mode, "scale": 1.7}}, + data={"sampling_strategy": "full_episode_uniform", "boundary_policy": None if mode != "boundary" else {"jitter": 4}}, + ckpt_path=None, paths={"output_dir": "/original"}, val_at_start=False, val_at_end=False, + trainer={"max_steps": 50, "val_check_interval": 50}, + callbacks={"joint_smoke": {"multimodal": True}, "model_checkpoint": {}, + "usage": {"_target_": "egomimic.pl_utils.joint_smoke.ConditionalOptimizerUsage"}}, + logger={"csv": {"save_dir": "/original"}}, evaluator={"pixel_diagnostics": None})) + + +@pytest.mark.parametrize("mode", ["plain", "fused", "boundary"]) +@pytest.mark.parametrize("steps", [(50, 52), (52, 56)]) +def test_phase_changes_only_diagnostic_lifetime_preserving_science(mode, steps): + original = original_config(mode) + before = OmegaConf.to_container(original, resolve=True) + cfg = phase_config(original, output="/diag", checkpoint="/ckpt", start_step=steps[0], + stop_step=steps[1], service_output="/proof", world=4) + assert OmegaConf.to_container(original, resolve=True) == before + assert cfg.model == original.model and cfg.data == original.data + assert cfg.callbacks.resume_proof.multimodal + assert cfg.trainer.max_steps == 56 and cfg.trainer.val_check_interval == 2 + assert ("allocation_boundary" in cfg.callbacks) == (steps[1] == 52) + assert cfg.callbacks.model_checkpoint.save_top_k == -1 + + +def test_phase_rejects_attempt_to_use_smoke_path_for_production_crash(): + with pytest.raises(ValueError, match="not production crash"): + phase_config(original_config("plain"), output="/diag", checkpoint="/ckpt", start_step=20000, + stop_step=20002, service_output="/proof", world=4) + + +@pytest.mark.parametrize("mode", ["plain", "fused", "boundary"]) +def test_checkpoint_science_pins_objective_and_sampling_even_if_weights_fit(mode): + saved = original_config(mode) + resumed = copy.deepcopy(saved) + resumed.trainer.max_steps = 56 + checkpoint_science(saved, resumed) + resumed.model.loss.scale = 2.3 + with pytest.raises(ValueError, match="model/loss/data/sampling"): + checkpoint_science(saved, resumed) + resumed = copy.deepcopy(saved) + resumed.data.sampling_strategy = "different_policy" + with pytest.raises(ValueError, match="model/loss/data/sampling"): + checkpoint_science(saved, resumed) + + +def hidden_fixture(): + stream = {"value": torch.tensor([[[3.], [float("nan")], [float("nan")]]]), + "observed": torch.tensor([[True, False, False]]), + "valid": torch.tensor([[True, True, False]])} + return {"world": {"image": stream}, "language_features": torch.ones(1, 1024), "pixel_metadata": [{}]}, {"image": torch.tensor([[[3.], [2.], [0.]]])} + + +@pytest.mark.parametrize("damage", [None, "leak", "nonfinite", "observed", "padding", "extra"]) +def test_hidden_sample_requires_finite_outputs_clamped_context_and_no_labels(damage): + sample, output = hidden_fixture() + if damage == "leak": sample["world"]["image"]["value"][0, 1] = 8 + elif damage == "nonfinite": output["image"][0, 1] = float("nan") + elif damage == "observed": output["image"][0, 0] = 4 + elif damage == "padding": output["image"][0, 2] = 7 + elif damage == "extra": output["action"] = output["image"] + if damage is None: validate_hidden_sample(sample, output, ["image"]) + else: + with pytest.raises(ValueError): validate_hidden_sample(sample, output, ["image"]) + + +def phase_rows(): + net = TinyConditional() + streams = {e: list(net.active_streams(e)) for e in net.pose_dims} + result = [] + for start, stop in ((50, 52), (52, 56)): + row = dict(multimodal=True, private_bank_optimizer_steps_verified=True, human_action_absent=True, + active_streams_by_embodiment=streams, + initial_parameter_usage={"shared": start, "private": start // 2, "unused": 0}, + conditional_optimizer_record=dict(schema="native-conditional-optimizer-usage-v1", global_step=stop, + parameter_names=["shared", "private", "unused"], + parameter_update_counts={"shared": stop, "private": stop // 2, "unused": 0}), + embodiment_optimizer_batches={e: (stop-start)//2 for e in streams}, + actual_rank_rng_restore=dict(observed=True, actual_sha256="a"*64, expected_sha256="a"*64), + validation_metrics={k: 0.01 for k in expected_metrics(net)}, + hidden_label_samples={e: dict(sample_input_sha256="b"*64, state_steps=20, + action_steps=20 if e in net.action_dims else 0, streams=v, + finite=True, hidden_targets_nan=True) for e, v in streams.items()}) + result.append([copy.deepcopy(row) for _ in range(4)]) + return result + + +def test_summary_accepts_complete_synthetic_fixture_only(): + summary = summarize_native_phases(phase_rows()) + assert summary["diagnostic_scope"] == "scheduled-smoke-checkpoint-only" + assert summary["private_bank_optimizer_steps_verified"] + + +@pytest.mark.parametrize("damage", ["legacy", "human_action", "rng", "nan_metric", "missing_metric", "missing_rank", "counter", "sample", "second_restore", "absent_human"]) +def test_summary_rejects_incomplete_or_corrupt_all_rank_evidence(damage): + phases = phase_rows(); row = phases[1][0] + if damage == "legacy": row.pop("multimodal") + elif damage == "human_action": row["active_streams_by_embodiment"]["human_bimanual"].append("action") + elif damage == "rng": row["actual_rank_rng_restore"]["observed"] = False + elif damage == "nan_metric": row["validation_metrics"][next(iter(row["validation_metrics"]))] = float("nan") + elif damage == "missing_metric": row["validation_metrics"].popitem() + elif damage == "missing_rank": phases[1].pop() + elif damage == "counter": row["conditional_optimizer_record"]["parameter_update_counts"]["private"] = 56 + elif damage == "sample": row["hidden_label_samples"]["human_bimanual"]["hidden_targets_nan"] = False + elif damage == "second_restore": + for r in phases[1]: r["initial_parameter_usage"]["private"] = 25 + elif damage == "absent_human": + for phase, count in zip(phases, (2, 4)): + for r in phase: r["embodiment_optimizer_batches"] = {"eva_bimanual": count, "human_bimanual": 0} + with pytest.raises(ValueError): summarize_native_phases(phases)