diff --git a/tests/fixtures/legacy_alignment_reference.py b/tests/fixtures/legacy_alignment_reference.py
deleted file mode 100644
index 1358433a33..0000000000
--- a/tests/fixtures/legacy_alignment_reference.py
+++ /dev/null
@@ -1,75 +0,0 @@
-"""Verbatim methods from EgoVerse ec5c903c HPTModel, for parity only."""
-
-import torch
-from geomloss import SamplesLoss
-
-
-class LegacyAlignment:
- def resume_from_depth(self, block_outputs, depth):
- """
- Detach at trunk depth and resume trunk forward pass.
- Gradients will only flow from depth upward.
- """
- cut_tokens = block_outputs[depth - 1].detach()
-
- blocks = self.trunk["trunk"].blocks
- for blk in list(blocks)[depth:]:
- cut_tokens = blk(cut_tokens, attn_mask=None)
-
- if self.trunk["trunk"].post_transformer_layer is not None:
- cut_tokens = self.trunk["trunk"].post_transformer_layer(cut_tokens)
-
- return self.postprocess_tokens(cut_tokens)
-
- def make_custom_cost(self, scaling_mask):
- def custom_cost(x, y):
- cost = 0.5 * (((x.unsqueeze(1) - y.unsqueeze(0)) ** 2).sum(dim=-1))
- return cost * scaling_mask
-
- return custom_cost
-
- def compute_ot(
- self, tokens1, tokens2, emb1_actions, emb2_actions, supervised, lambd
- ):
- tokens1 = tokens1.reshape(tokens1.shape[0], -1)
- tokens2 = tokens2.reshape(tokens1.shape[0], -1)
-
- if not supervised:
- ot_loss_fn = SamplesLoss("sinkhorn", p=2, blur=0.05, truncate=18)
- ot_loss = ot_loss_fn(tokens2, tokens1)
- avg_feature_dist = torch.norm(tokens2 - tokens1, dim=-1).mean()
- return ot_loss, avg_feature_dist
- else:
- B = tokens1.shape[0]
- if not self.ot_6dof:
- emb1_actions = emb1_actions[..., :3]
- emb2_actions = emb2_actions[..., :3]
- if self.use_dtw:
- emb2_delta = emb2_actions
- emb1_delta = emb1_actions
- emb2_expand = emb2_delta.unsqueeze(1).expand(B, B, -1, -1)
- emb1_expand = emb1_delta.unsqueeze(0).expand(B, B, -1, -1)
- pairwise_dist = self.dtw(
- emb2_expand.reshape(B * B, *emb2_actions.shape[1:]),
- emb1_expand.reshape(B * B, *emb1_actions.shape[1:]),
- ).view(B, B)
- else:
- emb2_expand = emb2_actions.unsqueeze(1) # (B, 1, T, D)
- emb1_expand = emb1_actions.unsqueeze(0) # (1, B, T, D)
- pairwise_dist = ((emb2_expand - emb1_expand) ** 2).mean(
- dim=(2, 3)
- ) # (B, B) #changed
-
- labels = torch.argmin(pairwise_dist, dim=1)
- W = torch.ones(B, B).to(self.device)
- W[torch.arange(B), labels] = lambd
-
- custom_cost_fn = self.make_custom_cost(W)
-
- ot_loss_fn = SamplesLoss(
- loss="sinkhorn", p=2, blur=0.05, cost=custom_cost_fn, truncate=18
- )
-
- ot_loss = ot_loss_fn(tokens2, tokens1)
- avg_feature_dist = torch.norm(tokens2 - tokens1, dim=-1).mean()
- return ot_loss, avg_feature_dist
diff --git a/tests/fixtures/legacy_hpt_trunk/metadata.json b/tests/fixtures/legacy_hpt_trunk/metadata.json
deleted file mode 100644
index bb81e01d3e..0000000000
--- a/tests/fixtures/legacy_hpt_trunk/metadata.json
+++ /dev/null
@@ -1,46 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "constructor": {
- "embed_dim": 16,
- "num_blocks": 2,
- "num_heads": 2,
- "observation_horizon": 1,
- "action_horizon": 3
- },
- "seed": 9226,
- "torch_version": "2.7.1",
- "source_namespace": "trunk.",
- "weights_sha256": "a2b5a6e44028d34195be7de5de62caa1a2bb8a3f6eceb34eec5e9201234e1786",
- "reference_sha256": "bb550af3d1ce85a546dd5597c17acf8633663ea0707bd71f662e280829dd8bbd",
- "keys": [
- "trunk.blocks.0.attn.in_proj_weight",
- "trunk.blocks.0.attn.in_proj_bias",
- "trunk.blocks.0.attn.bias_k",
- "trunk.blocks.0.attn.bias_v",
- "trunk.blocks.0.attn.out_proj.weight",
- "trunk.blocks.0.attn.out_proj.bias",
- "trunk.blocks.0.norm_1.weight",
- "trunk.blocks.0.norm_1.bias",
- "trunk.blocks.0.mlp.fc1.weight",
- "trunk.blocks.0.mlp.fc1.bias",
- "trunk.blocks.0.mlp.fc2.weight",
- "trunk.blocks.0.mlp.fc2.bias",
- "trunk.blocks.0.norm_2.weight",
- "trunk.blocks.0.norm_2.bias",
- "trunk.blocks.1.attn.in_proj_weight",
- "trunk.blocks.1.attn.in_proj_bias",
- "trunk.blocks.1.attn.bias_k",
- "trunk.blocks.1.attn.bias_v",
- "trunk.blocks.1.attn.out_proj.weight",
- "trunk.blocks.1.attn.out_proj.bias",
- "trunk.blocks.1.norm_1.weight",
- "trunk.blocks.1.norm_1.bias",
- "trunk.blocks.1.mlp.fc1.weight",
- "trunk.blocks.1.mlp.fc1.bias",
- "trunk.blocks.1.mlp.fc2.weight",
- "trunk.blocks.1.mlp.fc2.bias",
- "trunk.blocks.1.norm_2.weight",
- "trunk.blocks.1.norm_2.bias"
- ],
- "scope": "actual legacy HPTModel trunk factory/loader, masked forward, loss and gradients; generated CPU parameters, not pretrained weights or a policy score"
-}
diff --git a/tests/fixtures/legacy_hpt_trunk/reference.pt b/tests/fixtures/legacy_hpt_trunk/reference.pt
deleted file mode 100644
index 5da2a0afac..0000000000
Binary files a/tests/fixtures/legacy_hpt_trunk/reference.pt and /dev/null differ
diff --git a/tests/fixtures/legacy_hpt_trunk/trunk.pth b/tests/fixtures/legacy_hpt_trunk/trunk.pth
deleted file mode 100644
index a3a1353366..0000000000
Binary files a/tests/fixtures/legacy_hpt_trunk/trunk.pth and /dev/null differ
diff --git a/tests/fixtures/legacy_language_schedule.json b/tests/fixtures/legacy_language_schedule.json
deleted file mode 100644
index 6b0a695279..0000000000
--- a/tests/fixtures/legacy_language_schedule.json
+++ /dev/null
@@ -1,2011 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "source_path": "egomimic/utils/scheduler_utils.py",
- "warmup_epochs": 100,
- "total_epochs": 2000,
- "eta_min": 1e-05,
- "peak_lr": 0.0002,
- "values": [
- 2.0000000000000002e-07,
- 2.1980000000000004e-06,
- 4.196e-06,
- 6.194e-06,
- 8.192e-06,
- 1.019e-05,
- 1.2188000000000001e-05,
- 1.4186e-05,
- 1.6184e-05,
- 1.8181999999999996e-05,
- 2.0179999999999997e-05,
- 2.2177999999999997e-05,
- 2.4175999999999994e-05,
- 2.617399999999999e-05,
- 2.8171999999999987e-05,
- 3.0169999999999987e-05,
- 3.216799999999999e-05,
- 3.4165999999999984e-05,
- 3.616399999999998e-05,
- 3.8161999999999985e-05,
- 4.015999999999998e-05,
- 4.2157999999999985e-05,
- 4.415599999999998e-05,
- 4.615399999999998e-05,
- 4.815199999999998e-05,
- 5.014999999999998e-05,
- 5.214799999999998e-05,
- 5.4145999999999986e-05,
- 5.614399999999999e-05,
- 5.814199999999999e-05,
- 6.014e-05,
- 6.213800000000001e-05,
- 6.413600000000001e-05,
- 6.613400000000001e-05,
- 6.813200000000002e-05,
- 7.013000000000002e-05,
- 7.212800000000002e-05,
- 7.412600000000003e-05,
- 7.612400000000002e-05,
- 7.812200000000002e-05,
- 8.012000000000003e-05,
- 8.211800000000003e-05,
- 8.411600000000003e-05,
- 8.611400000000004e-05,
- 8.811200000000004e-05,
- 9.011000000000003e-05,
- 9.210800000000003e-05,
- 9.410600000000004e-05,
- 9.610400000000004e-05,
- 9.810200000000003e-05,
- 0.00010010000000000003,
- 0.00010209800000000004,
- 0.00010409600000000004,
- 0.00010609400000000005,
- 0.00010809200000000004,
- 0.00011009000000000004,
- 0.00011208800000000003,
- 0.00011408600000000002,
- 0.00011608400000000004,
- 0.00011808200000000005,
- 0.00012008000000000006,
- 0.00012207800000000005,
- 0.00012407600000000004,
- 0.00012607400000000003,
- 0.00012807200000000002,
- 0.00013007000000000003,
- 0.00013206800000000002,
- 0.00013406600000000001,
- 0.000136064,
- 0.00013806200000000002,
- 0.00014006000000000004,
- 0.00014205800000000003,
- 0.00014405600000000005,
- 0.00014605400000000004,
- 0.00014805200000000003,
- 0.00015005000000000004,
- 0.00015204800000000006,
- 0.00015404600000000005,
- 0.00015604400000000007,
- 0.00015804200000000008,
- 0.00016004000000000007,
- 0.0001620380000000001,
- 0.00016403600000000008,
- 0.00016603400000000007,
- 0.0001680320000000001,
- 0.00017003000000000008,
- 0.0001720280000000001,
- 0.0001740260000000001,
- 0.00017602400000000013,
- 0.00017802200000000012,
- 0.00018002000000000014,
- 0.00018201800000000013,
- 0.00018401600000000014,
- 0.00018601400000000013,
- 0.00018801200000000012,
- 0.00019001000000000014,
- 0.00019200800000000013,
- 0.00019400600000000015,
- 0.00019600400000000017,
- 0.00019800200000000018,
- 0.0002,
- 0.00019999987013681378,
- 0.0001999994805476102,
- 0.00019999883123345428,
- 0.00019999792219612135,
- 0.00019999675343809658,
- 0.00019999532496257536,
- 0.0001999936367734631,
- 0.0001999916888753752,
- 0.0001999894812736372,
- 0.00019998701397428453,
- 0.00019998428698406274,
- 0.0001999813003104273,
- 0.00019997805396154366,
- 0.00019997454794628723,
- 0.0001999707822742433,
- 0.0001999667569557071,
- 0.00019996247200168362,
- 0.00019995792742388786,
- 0.00019995312323474449,
- 0.00019994805944738796,
- 0.0001999427360756625,
- 0.00019993715313412194,
- 0.0001999313106380299,
- 0.00019992520860335957,
- 0.00019991884704679357,
- 0.00019991222598572423,
- 0.00019990534543825328,
- 0.00019989820542319183,
- 0.00019989080596006042,
- 0.0001998831470690889,
- 0.00019987522877121643,
- 0.00019986705108809128,
- 0.00019985861404207096,
- 0.000199849917656222,
- 0.00019984096195432,
- 0.00019983174696084954,
- 0.00019982227270100403,
- 0.00019981253920068573,
- 0.00019980254648650566,
- 0.00019979229458578352,
- 0.00019978178352654763,
- 0.00019977101333753478,
- 0.0001997599840481903,
- 0.00019974869568866783,
- 0.00019973714828982935,
- 0.00019972534188324493,
- 0.0001997132765011929,
- 0.00019970095217665952,
- 0.00019968836894333904,
- 0.00019967552683563356,
- 0.00019966242588865288,
- 0.0001996490661382145,
- 0.00019963544762084345,
- 0.00019962157037377225,
- 0.00019960743443494074,
- 0.00019959303984299612,
- 0.00019957838663729255,
- 0.00019956347485789138,
- 0.00019954830454556083,
- 0.000199532875741776,
- 0.00019951718848871863,
- 0.00019950124282927707,
- 0.0001994850388070461,
- 0.00019946857646632694,
- 0.00019945185585212703,
- 0.00019943487701015982,
- 0.00019941763998684484,
- 0.0001994001448293075,
- 0.00019938239158537885,
- 0.0001993643803035956,
- 0.00019934611103319987,
- 0.0001993275838241392,
- 0.00019930879872706621,
- 0.0001992897557933387,
- 0.00019927045507501926,
- 0.0001992508966248754,
- 0.00019923108049637908,
- 0.0001992110067437069,
- 0.00019919067542173975,
- 0.00019917008658606263,
- 0.00019914924029296467,
- 0.00019912813659943884,
- 0.00019910677556318182,
- 0.00019908515724259386,
- 0.00019906328169677862,
- 0.00019904114898554304,
- 0.00019901875916939708,
- 0.00019899611230955361,
- 0.00019897320846792835,
- 0.0001989500477071395,
- 0.0001989266300905077,
- 0.00019890295568205578,
- 0.0001988790245465087,
- 0.00019885483674929332,
- 0.00019883039235653807,
- 0.00019880569143507302,
- 0.00019878073405242956,
- 0.0001987555202768402,
- 0.00019873005017723844,
- 0.00019870432382325859,
- 0.00019867834128523546,
- 0.0001986521026342044,
- 0.00019862560794190083,
- 0.0001985988572807602,
- 0.00019857185072391786,
- 0.00019854458834520867,
- 0.00019851707021916693,
- 0.00019848929642102617,
- 0.00019846126702671883,
- 0.00019843298211287628,
- 0.00019840444175682834,
- 0.0001983756460366033,
- 0.0001983465950309275,
- 0.00019831728881922536,
- 0.0001982877274816189,
- 0.0001982579110989277,
- 0.00019822783975266854,
- 0.00019819751352505543,
- 0.00019816693249899908,
- 0.00019813609675810686,
- 0.00019810500638668248,
- 0.00019807366146972583,
- 0.00019804206209293276,
- 0.00019801020834269472,
- 0.00019797810030609865,
- 0.00019794573807092669,
- 0.00019791312172565598,
- 0.00019788025135945836,
- 0.0001978471270622001,
- 0.00019781374892444188,
- 0.00019778011703743817,
- 0.00019774623149313724,
- 0.00019771209238418093,
- 0.00019767769980390425,
- 0.0001976430538463352,
- 0.0001976081546061945,
- 0.00019757300217889532,
- 0.00019753759666054312,
- 0.00019750193814793517,
- 0.00019746602673856049,
- 0.0001974298625305995,
- 0.0001973934456229238,
- 0.00019735677611509576,
- 0.0001973198541073685,
- 0.00019728267970068532,
- 0.00019724525299667968,
- 0.0001972075740976747,
- 0.00019716964310668306,
- 0.00019713146012740666,
- 0.00019709302526423628,
- 0.00019705433862225142,
- 0.00019701540030721986,
- 0.00019697621042559746,
- 0.00019693676908452785,
- 0.00019689707639184215,
- 0.00019685713245605867,
- 0.00019681693738638265,
- 0.0001967764912927058,
- 0.00019673579428560627,
- 0.00019669484647634805,
- 0.00019665364797688092,
- 0.00019661219889984002,
- 0.00019657049935854553,
- 0.00019652854946700243,
- 0.00019648634933990005,
- 0.00019644389909261197,
- 0.0001964011988411956,
- 0.00019635824870239168,
- 0.00019631504879362435,
- 0.00019627159923300046,
- 0.00019622790013930946,
- 0.000196183951632023,
- 0.00019613975383129466,
- 0.0001960953068579595,
- 0.00019605061083353387,
- 0.00019600566588021494,
- 0.00019596047212088057,
- 0.00019591502967908873,
- 0.00019586933867907734,
- 0.00019582339924576383,
- 0.00019577721150474488,
- 0.00019573077558229598,
- 0.0001956840916053712,
- 0.0001956371597016027,
- 0.00019558997999930054,
- 0.00019554255262745226,
- 0.00019549487771572245,
- 0.0001954469553944525,
- 0.0001953987857946602,
- 0.00019535036904803942,
- 0.00019530170528695969,
- 0.00019525279464446585,
- 0.00019520363725427777,
- 0.00019515423325078984,
- 0.00019510458276907072,
- 0.00019505468594486298,
- 0.00019500454291458257,
- 0.00019495415381531864,
- 0.00019490351878483309,
- 0.0001948526379615601,
- 0.00019480151148460598,
- 0.00019475013949374852,
- 0.00019469852212943678,
- 0.00019464665953279066,
- 0.00019459455184560052,
- 0.0001945421992103268,
- 0.00019448960177009963,
- 0.00019443675966871835,
- 0.00019438367305065128,
- 0.00019433034206103517,
- 0.00019427676684567494,
- 0.0001942229475510432,
- 0.0001941688843242798,
- 0.0001941145773131916,
- 0.0001940600266662518,
- 0.00019400523253259986,
- 0.00019395019506204087,
- 0.00019389491440504513,
- 0.0001938393907127478,
- 0.00019378362413694864,
- 0.00019372761483011123,
- 0.00019367136294536297,
- 0.00019361486863649432,
- 0.00019355813205795857,
- 0.00019350115336487137,
- 0.00019344393271301027,
- 0.00019338647025881435,
- 0.00019332876615938376,
- 0.00019327082057247932,
- 0.00019321263365652203,
- 0.0001931542055705927,
- 0.00019309553647443154,
- 0.00019303662652843756,
- 0.00019297747589366832,
- 0.0001929180847318394,
- 0.00019285845320532398,
- 0.00019279858147715236,
- 0.00019273846971101156,
- 0.00019267811807124487,
- 0.00019261752672285136,
- 0.00019255669583148545,
- 0.0001924956255634565,
- 0.0001924343160857282,
- 0.00019237276756591847,
- 0.00019231098017229848,
- 0.0001922489540737927,
- 0.00019218668943997801,
- 0.00019212418644108366,
- 0.0001920614452479904,
- 0.00019199846603223025,
- 0.00019193524896598602,
- 0.0001918717942220907,
- 0.0001918081019740272,
- 0.00019174417239592765,
- 0.0001916800056625731,
- 0.00019161560194939293,
- 0.00019155096143246445,
- 0.00019148608428851233,
- 0.00019142097069490822,
- 0.00019135562082967024,
- 0.00019129003487146237,
- 0.00019122421299959416,
- 0.00019115815539402003,
- 0.00019109186223533904,
- 0.0001910253337047941,
- 0.0001909585699842717,
- 0.00019089157125630124,
- 0.00019082433770405474,
- 0.00019075686951134613,
- 0.00019068916686263096,
- 0.00019062122994300556,
- 0.00019055305893820693,
- 0.00019048465403461202,
- 0.0001904160154192372,
- 0.00019034714327973782,
- 0.00019027803780440766,
- 0.00019020869918217852,
- 0.00019013912760261947,
- 0.0001900693232559366,
- 0.0001899992863329723,
- 0.00018992901702520487,
- 0.00018985851552474791,
- 0.00018978778202434984,
- 0.00018971681671739327,
- 0.00018964561979789473,
- 0.00018957419146050383,
- 0.00018950253190050295,
- 0.0001894306413138066,
- 0.00018935851989696086,
- 0.0001892861678471429,
- 0.00018921358536216055,
- 0.00018914077264045148,
- 0.00018906772988108294,
- 0.00018899445728375104,
- 0.00018892095504878024,
- 0.00018884722337712283,
- 0.0001887732624703584,
- 0.00018869907253069327,
- 0.00018862465376095987,
- 0.0001885500063646163,
- 0.00018847513054574567,
- 0.0001884000265090556,
- 0.00018832469445987773,
- 0.0001882491346041669,
- 0.00018817334714850095,
- 0.00018809733230007984,
- 0.0001880210902667253,
- 0.00018794462125688012,
- 0.00018786792547960767,
- 0.00018779100314459123,
- 0.00018771385446213359,
- 0.00018763647964315625,
- 0.00018755887889919905,
- 0.00018748105244241941,
- 0.00018740300048559195,
- 0.0001873247232421077,
- 0.00018724622092597366,
- 0.0001871674937518122,
- 0.00018708854193486037,
- 0.00018700936569096942,
- 0.0001869299652366042,
- 0.00018685034078884247,
- 0.00018677049256537453,
- 0.00018669042078450228,
- 0.00018661012566513897,
- 0.00018652960742680837,
- 0.00018644886628964424,
- 0.00018636790247438981,
- 0.00018628671620239703,
- 0.00018620530769562608,
- 0.0001861236771766447,
- 0.00018604182486862763,
- 0.00018595975099535596,
- 0.0001858774557812165,
- 0.0001857949394512012,
- 0.00018571220223090662,
- 0.0001856292443465331,
- 0.00018554606602488437,
- 0.0001854626674933668,
- 0.0001853790489799887,
- 0.00018529521071336002,
- 0.0001852111529226913,
- 0.00018512687583779332,
- 0.00018504237968907644,
- 0.0001849576647075499,
- 0.00018487273112482124,
- 0.0001847875791730956,
- 0.00018470220908517522,
- 0.00018461662109445857,
- 0.00018453081543494005,
- 0.000184444792341209,
- 0.00018435855204844927,
- 0.00018427209479243852,
- 0.00018418542080954762,
- 0.0001840985303367399,
- 0.0001840114236115706,
- 0.00018392410087218622,
- 0.00018383656235732376,
- 0.0001837488083063102,
- 0.00018366083895906171,
- 0.00018357265455608317,
- 0.00018348425533846736,
- 0.0001833956415478944,
- 0.00018330681342663095,
- 0.00018321777121752973,
- 0.00018312851516402874,
- 0.00018303904551015063,
- 0.00018294936250050201,
- 0.00018285946638027279,
- 0.00018276935739523558,
- 0.0001826790357917448,
- 0.0001825885018167364,
- 0.0001824977557177267,
- 0.00018240679774281213,
- 0.0001823156281406683,
- 0.00018222424716054941,
- 0.00018213265505228758,
- 0.00018204085206629216,
- 0.00018194883845354897,
- 0.00018185661446561976,
- 0.00018176418035464135,
- 0.00018167153637332507,
- 0.00018157868277495607,
- 0.0001814856198133925,
- 0.00018139234774306495,
- 0.00018129886681897569,
- 0.00018120517729669803,
- 0.0001811112794323755,
- 0.00018101717348272136,
- 0.00018092285970501757,
- 0.00018082833835711447,
- 0.00018073360969742978,
- 0.0001806386739849481,
- 0.00018054353147921998,
- 0.00018044818244036145,
- 0.00018035262712905309,
- 0.00018025686580653948,
- 0.00018016089873462842,
- 0.0001800647261756903,
- 0.0001799683483926571,
- 0.00017987176564902206,
- 0.00017977497820883872,
- 0.00017967798633672022,
- 0.00017958079029783863,
- 0.00017948339035792424,
- 0.00017938578678326474,
- 0.0001792879798407046,
- 0.00017918996979764427,
- 0.00017909175692203943,
- 0.0001789933414824004,
- 0.00017889472374779113,
- 0.00017879590398782886,
- 0.0001786968824726829,
- 0.00017859765947307441,
- 0.00017849823526027523,
- 0.00017839861010610737,
- 0.0001782987842829422,
- 0.0001781987580636997,
- 0.00017809853172184777,
- 0.00017799810553140137,
- 0.00017789747976692182,
- 0.0001777966547035162,
- 0.00017769563061683635,
- 0.0001775944077830783,
- 0.00017749298647898136,
- 0.00017739136698182752,
- 0.00017728954956944062,
- 0.00017718753452018553,
- 0.00017708532211296757,
- 0.00017698291262723152,
- 0.00017688030634296098,
- 0.0001767775035406777,
- 0.0001766745045014405,
- 0.00017657130950684492,
- 0.00017646791883902207,
- 0.00017636433278063812,
- 0.00017626055161489338,
- 0.00017615657562552163,
- 0.00017605240509678918,
- 0.00017594804031349434,
- 0.0001758434815609664,
- 0.00017573872912506508,
- 0.00017563378329217948,
- 0.00017552864434922748,
- 0.00017542331258365497,
- 0.00017531778828343492,
- 0.00017521207173706684,
- 0.00017510616323357565,
- 0.00017500006306251116,
- 0.00017489377151394723,
- 0.00017478728887848082,
- 0.0001746806154472315,
- 0.00017457375151184028,
- 0.00017446669736446906,
- 0.00017435945329779987,
- 0.0001742520196050338,
- 0.00017414439657989055,
- 0.00017403658451660732,
- 0.00017392858370993816,
- 0.00017382039445515318,
- 0.00017371201704803764,
- 0.00017360345178489127,
- 0.00017349469896252732,
- 0.00017338575887827188,
- 0.00017327663182996298,
- 0.0001731673181159498,
- 0.00017305781803509185,
- 0.00017294813188675813,
- 0.00017283825997082647,
- 0.00017272820258768245,
- 0.00017261796003821875,
- 0.00017250753262383427,
- 0.0001723969206464334,
- 0.00017228612440842498,
- 0.00017217514421272176,
- 0.00017206398036273934,
- 0.00017195263316239543,
- 0.00017184110291610904,
- 0.00017172938992879958,
- 0.00017161749450588616,
- 0.0001715054169532865,
- 0.00017139315757741638,
- 0.00017128071668518868,
- 0.0001711680945840125,
- 0.00017105529158179236,
- 0.00017094230798692732,
- 0.00017082914410831025,
- 0.00017071580025532687,
- 0.00017060227673785497,
- 0.00017048857386626345,
- 0.0001703746919514117,
- 0.00017026063130464844,
- 0.00017014639223781117,
- 0.00017003197506322513,
- 0.0001699173800937025,
- 0.0001698026076425416,
- 0.00016968765802352585,
- 0.00016957253155092317,
- 0.0001694572285394849,
- 0.00016934174930444514,
- 0.00016922609416151961,
- 0.00016911026342690514,
- 0.0001689942574172785,
- 0.0001688780764497957,
- 0.0001687617208420911,
- 0.0001686451909122764,
- 0.00016852848697894004,
- 0.0001684116093611461,
- 0.00016829455837843354,
- 0.00016817733435081528,
- 0.00016805993759877734,
- 0.00016794236844327792,
- 0.00016782462720574663,
- 0.0001677067142080835,
- 0.00016758862977265819,
- 0.00016747037422230903,
- 0.00016735194788034216,
- 0.00016723335107053067,
- 0.00016711458411711373,
- 0.00016699564734479566,
- 0.00016687654107874506,
- 0.00016675726564459388,
- 0.00016663782136843664,
- 0.0001665182085768294,
- 0.00016639842759678903,
- 0.00016627847875579205,
- 0.0001661583623817741,
- 0.0001660380788031287,
- 0.00016591762834870658,
- 0.0001657970113478147,
- 0.00016567622813021534,
- 0.00016555527902612516,
- 0.00016543416436621446,
- 0.00016531288448160602,
- 0.00016519143970387446,
- 0.00016506983036504514,
- 0.00016494805679759337,
- 0.00016482611933444345,
- 0.00016470401830896768,
- 0.00016458175405498568,
- 0.00016445932690676318,
- 0.00016433673719901142,
- 0.00016421398526688592,
- 0.00016409107144598574,
- 0.00016396799607235263,
- 0.00016384475948246995,
- 0.00016372136201326184,
- 0.00016359780400209225,
- 0.00016347408578676407,
- 0.00016335020770551816,
- 0.00016322617009703247,
- 0.0001631019733004211,
- 0.00016297761765523335,
- 0.00016285310350145282,
- 0.00016272843117949644,
- 0.00016260360103021354,
- 0.00016247861339488501,
- 0.00016235346861522226,
- 0.00016222816703336638,
- 0.00016210270899188708,
- 0.00016197709483378186,
- 0.00016185132490247503,
- 0.00016172539954181677,
- 0.00016159931909608225,
- 0.00016147308390997053,
- 0.00016134669432860382,
- 0.0001612201506975264,
- 0.00016109345336270375,
- 0.00016096660267052155,
- 0.00016083959896778468,
- 0.00016071244260171648,
- 0.0001605851339199576,
- 0.00016045767327056508,
- 0.00016033006100201147,
- 0.00016020229746318388,
- 0.00016007438300338293,
- 0.00015994631797232184,
- 0.00015981810272012553,
- 0.00015968973759732962,
- 0.0001595612229548794,
- 0.00015943255914412907,
- 0.00015930374651684055,
- 0.00015917478542518262,
- 0.00015904567622173,
- 0.00015891641925946232,
- 0.0001587870148917632,
- 0.0001586574634724192,
- 0.00015852776535561894,
- 0.0001583979208959522,
- 0.0001582679304484087,
- 0.00015813779436837737,
- 0.00015800751301164534,
- 0.0001578770867343968,
- 0.00015774651589321221,
- 0.00015761580084506732,
- 0.00015748494194733204,
- 0.00015735393955776958,
- 0.00015722279403453554,
- 0.00015709150573617672,
- 0.00015696007502163033,
- 0.00015682850225022293,
- 0.00015669678778166944,
- 0.00015656493197607218,
- 0.0001564329351939199,
- 0.0001563007977960868,
- 0.0001561685201438314,
- 0.0001560361025987958,
- 0.0001559035455230045,
- 0.00015577084927886347,
- 0.0001556380142291592,
- 0.00015550504073705764,
- 0.00015537192916610326,
- 0.000155238679880218,
- 0.00015510529324370024,
- 0.00015497176962122404,
- 0.00015483810937783787,
- 0.00015470431287896376,
- 0.00015457038049039614,
- 0.00015443631257830112,
- 0.00015430210950921524,
- 0.00015416777165004454,
- 0.00015403329936806366,
- 0.0001538986930309146,
- 0.000153763953006606,
- 0.00015362907966351191,
- 0.00015349407337037095,
- 0.0001533589344962851,
- 0.00015322366341071893,
- 0.00015308826048349841,
- 0.00015295272608481,
- 0.0001528170605851995,
- 0.00015268126435557128,
- 0.000152545337767187,
- 0.00015240928119166478,
- 0.0001522730950009781,
- 0.00015213677956745481,
- 0.00015200033526377611,
- 0.00015186376246297548,
- 0.00015172706153843775,
- 0.00015159023286389806,
- 0.00015145327681344076,
- 0.0001513161937614985,
- 0.00015117898408285104,
- 0.00015104164815262443,
- 0.0001509041863462899,
- 0.00015076659903966272,
- 0.00015062888660890138,
- 0.00015049104943050637,
- 0.0001503530878813193,
- 0.00015021500233852178,
- 0.00015007679317963438,
- 0.00014993846078251567,
- 0.00014980000552536116,
- 0.0001496614277867022,
- 0.00014952272794540512,
- 0.00014938390638066988,
- 0.0001492449634720294,
- 0.00014910589959934832,
- 0.00014896671514282188,
- 0.00014882741048297512,
- 0.00014868798600066162,
- 0.00014854844207706262,
- 0.00014840877909368588,
- 0.00014826899743236463,
- 0.00014812909747525667,
- 0.00014798907960484307,
- 0.00014784894420392738,
- 0.0001477086916556344,
- 0.0001475683223434093,
- 0.0001474278366510164,
- 0.00014728723496253816,
- 0.00014714651766237422,
- 0.00014700568513524042,
- 0.00014686473776616738,
- 0.0001467236759404999,
- 0.00014658250004389567,
- 0.0001464412104623241,
- 0.00014629980758206558,
- 0.00014615829178971014,
- 0.00014601666347215656,
- 0.0001458749230166113,
- 0.0001457330708105873,
- 0.00014559110724190306,
- 0.00014544903269868162,
- 0.00014530684756934933,
- 0.00014516455224263488,
- 0.0001450221471075683,
- 0.00014487963255347978,
- 0.00014473700896999863,
- 0.0001445942767470523,
- 0.00014445143627486528,
- 0.00014430848794395792,
- 0.00014416543214514551,
- 0.00014402226926953714,
- 0.00014387899970853466,
- 0.00014373562385383156,
- 0.00014359214209741194,
- 0.00014344855483154948,
- 0.00014330486244880627,
- 0.0001431610653420318,
- 0.00014301716390436185,
- 0.0001428731585292175,
- 0.0001427290496103039,
- 0.00014258483754160935,
- 0.0001424405227174041,
- 0.00014229610553223942,
- 0.00014215158638094638,
- 0.00014200696565863478,
- 0.0001418622437606922,
- 0.0001417174210827827,
- 0.0001415724980208461,
- 0.00014142747497109638,
- 0.00014128235233002107,
- 0.00014113713049437992,
- 0.0001409918098612039,
- 0.00014084639082779406,
- 0.00014070087379172052,
- 0.00014055525915082124,
- 0.0001404095473032011,
- 0.00014026373864723076,
- 0.00014011783358154552,
- 0.00013997183250504422,
- 0.0001398257358168882,
- 0.00013967954391650033,
- 0.00013953325720356363,
- 0.00013938687607802035,
- 0.00013924040094007094,
- 0.00013909383219017283,
- 0.0001389471702290394,
- 0.00013880041545763882,
- 0.00013865356827719305,
- 0.00013850662908917668,
- 0.0001383595982953158,
- 0.00013821247629758703,
- 0.00013806526349821626,
- 0.00013791796029967774,
- 0.00013777056710469266,
- 0.0001376230843162285,
- 0.00013747551233749748,
- 0.00013732785157195585,
- 0.0001371801024233024,
- 0.0001370322652954778,
- 0.000136884340592663,
- 0.0001367363287192785,
- 0.00013658823007998313,
- 0.00013644004507967292,
- 0.00013629177412348,
- 0.0001361434176167715,
- 0.00013599497596514845,
- 0.00013584644957444466,
- 0.0001356978388507256,
- 0.0001355491442002873,
- 0.00013540036602965527,
- 0.0001352515047455834,
- 0.00013510256075505264,
- 0.0001349535344652703,
- 0.00013480442628366844,
- 0.00013465523661790324,
- 0.0001345059658758535,
- 0.00013435661446561974,
- 0.000134207182795523,
- 0.0001340576712741038,
- 0.00013390808031012085,
- 0.00013375841031255026,
- 0.00013360866169058398,
- 0.00013345883485362905,
- 0.0001333089302113063,
- 0.00013315894817344934,
- 0.00013300888915010334,
- 0.0001328587535515239,
- 0.00013270854178817603,
- 0.000132558254270733,
- 0.00013240789141007506,
- 0.00013225745361728858,
- 0.00013210694130366473,
- 0.00013195635488069844,
- 0.00013180569476008725,
- 0.00013165496135373014,
- 0.00013150415507372654,
- 0.00013135327633237502,
- 0.00013120232554217226,
- 0.00013105130311581205,
- 0.00013090020946618386,
- 0.00013074904500637202,
- 0.00013059781014965432,
- 0.00013044650530950116,
- 0.0001302951308995741,
- 0.0001301436873337251,
- 0.00012999217502599502,
- 0.00012984059439061276,
- 0.00012968894584199396,
- 0.00012953722979473998,
- 0.00012938544666363667,
- 0.00012923359686365334,
- 0.00012908168080994156,
- 0.000128929698917834,
- 0.00012877765160284333,
- 0.00012862553928066117,
- 0.0001284733623671567,
- 0.00012832112127837586,
- 0.00012816881643053996,
- 0.00012801644824004467,
- 0.00012786401712345873,
- 0.00012771152349752308,
- 0.0001275589677791494,
- 0.00012740635038541924,
- 0.00012725367173358268,
- 0.00012710093224105738,
- 0.00012694813232542725,
- 0.00012679527240444145,
- 0.00012664235289601308,
- 0.00012648937421821837,
- 0.0001263363367892951,
- 0.00012618324102764176,
- 0.00012603008735181636,
- 0.00012587687618053516,
- 0.00012572360793267165,
- 0.00012557028302725537,
- 0.00012541690188347078,
- 0.00012526346492065606,
- 0.000125109972558302,
- 0.00012495642521605086,
- 0.00012480282331369522,
- 0.00012464916727117681,
- 0.0001244954575085854,
- 0.0001243416944461576,
- 0.00012418787850427572,
- 0.0001240340101034667,
- 0.00012388008966440086,
- 0.00012372611760789082,
- 0.00012357209435489032,
- 0.00012341802032649296,
- 0.00012326389594393138,
- 0.00012310972162857568,
- 0.00012295549780193255,
- 0.00012280122488564406,
- 0.00012264690330148647,
- 0.0001224925334713691,
- 0.00012233811581733323,
- 0.00012218365076155068,
- 0.00012202913872632315,
- 0.00012187458013408058,
- 0.00012171997540738028,
- 0.00012156532496890566,
- 0.00012141062924146514,
- 0.00012125588864799089,
- 0.00012110110361153782,
- 0.00012094627455528228,
- 0.00012079140190252103,
- 0.00012063648607666995,
- 0.00012048152750126302,
- 0.00012032652659995108,
- 0.00012017148379650068,
- 0.00012001639951479283,
- 0.00011986127417882218,
- 0.00011970610821269536,
- 0.00011955090204063027,
- 0.00011939565608695457,
- 0.00011924037077610488,
- 0.00011908504653262522,
- 0.00011892968378116621,
- 0.0001187742829464836,
- 0.00011861884445343741,
- 0.00011846336872699054,
- 0.00011830785619220763,
- 0.00011815230727425404,
- 0.00011799672239839464,
- 0.00011784110198999245,
- 0.00011768544647450775,
- 0.00011752975627749677,
- 0.00011737403182461062,
- 0.00011721827354159392,
- 0.00011706248185428392,
- 0.00011690665718860912,
- 0.00011675079997058826,
- 0.00011659491062632893,
- 0.00011643898958202677,
- 0.00011628303726396383,
- 0.0001161270540985079,
- 0.0001159710405121109,
- 0.0001158149969313081,
- 0.00011565892378271665,
- 0.0001155028214930346,
- 0.0001153466904890396,
- 0.00011519053119758795,
- 0.00011503434404561308,
- 0.00011487812946012477,
- 0.00011472188786820771,
- 0.00011456561969702045,
- 0.0001144093253737942,
- 0.00011425300532583165,
- 0.00011409665998050589,
- 0.0001139402897652591,
- 0.00011378389510760146,
- 0.00011362747643511001,
- 0.00011347103417542745,
- 0.00011331456875626091,
- 0.00011315808060538084,
- 0.00011300157015061996,
- 0.00011284503781987182,
- 0.00011268848404108985,
- 0.00011253190924228605,
- 0.00011237531385153003,
- 0.00011221869829694753,
- 0.00011206206300671958,
- 0.00011190540840908098,
- 0.0001117487349323195,
- 0.00011159204300477441,
- 0.00011143533305483544,
- 0.00011127860551094167,
- 0.00011112186080158016,
- 0.000110965099355285,
- 0.00011080832160063595,
- 0.00011065152796625746,
- 0.0001104947188808173,
- 0.00011033789477302554,
- 0.00011018105607163328,
- 0.0001100242032054316,
- 0.00010986733660325015,
- 0.00010971045669395633,
- 0.00010955356390645376,
- 0.00010939665866968136,
- 0.00010923974141261204,
- 0.00010908281256425159,
- 0.00010892587255363746,
- 0.00010876892180983773,
- 0.00010861196076194966,
- 0.00010845498983909879,
- 0.00010829800947043762,
- 0.0001081410200851445,
- 0.00010798402211242238,
- 0.0001078270159814978,
- 0.00010767000212161947,
- 0.00010751298096205728,
- 0.00010735595293210107,
- 0.00010719891846105954,
- 0.00010704187797825888,
- 0.00010688483191304183,
- 0.00010672778069476626,
- 0.00010657072475280428,
- 0.00010641366451654079,
- 0.00010625660041537248,
- 0.00010609953287870663,
- 0.00010594246233595989,
- 0.00010578538921655709,
- 0.00010562831394993017,
- 0.0001054712369655169,
- 0.00010531415869275978,
- 0.00010515707956110476,
- 0.00010500000000000023,
- 0.0001048429204388957,
- 0.00010468584130724072,
- 0.00010452876303448355,
- 0.0001043716860500703,
- 0.00010421461078344338,
- 0.0001040575376640406,
- 0.00010390046712129383,
- 0.00010374339958462799,
- 0.00010358633548345969,
- 0.00010342927524719621,
- 0.00010327221930523419,
- 0.00010311516808695866,
- 0.00010295812202174159,
- 0.00010280108153894093,
- 0.00010264404706789938,
- 0.0001024870190379432,
- 0.000102329997878381,
- 0.00010217298401850264,
- 0.00010201597788757803,
- 0.00010185897991485595,
- 0.00010170199052956281,
- 0.00010154501016090167,
- 0.0001013880392380508,
- 0.00010123107819016275,
- 0.00010107412744636295,
- 0.00010091718743574887,
- 0.0001007602585873884,
- 0.0001006033413303191,
- 0.00010044643609354666,
- 0.0001002895433060441,
- 0.00010013266339675026,
- 9.997579679456885e-05,
- 9.981894392836712e-05,
- 9.96621052269749e-05,
- 9.950528111918314e-05,
- 9.934847203374297e-05,
- 9.919167839936447e-05,
- 9.903490064471545e-05,
- 9.887813919842027e-05,
- 9.872139448905876e-05,
- 9.856466694516498e-05,
- 9.840795699522605e-05,
- 9.825126506768094e-05,
- 9.809459159091948e-05,
- 9.79379369932809e-05,
- 9.778130170305294e-05,
- 9.762468614847044e-05,
- 9.74680907577144e-05,
- 9.731151595891065e-05,
- 9.715496218012867e-05,
- 9.699842984938049e-05,
- 9.684191939461959e-05,
- 9.668543124373957e-05,
- 9.6528965824573e-05,
- 9.637252356489044e-05,
- 9.6216104892399e-05,
- 9.605971023474136e-05,
- 9.590334001949454e-05,
- 9.57469946741688e-05,
- 9.559067462620626e-05,
- 9.543438030298e-05,
- 9.527811213179273e-05,
- 9.512187053987569e-05,
- 9.496565595438737e-05,
- 9.480946880241252e-05,
- 9.465330951096081e-05,
- 9.449717850696585e-05,
- 9.434107621728378e-05,
- 9.418500306869233e-05,
- 9.402895948788953e-05,
- 9.387294590149257e-05,
- 9.371696273603661e-05,
- 9.356101041797369e-05,
- 9.340508937367149e-05,
- 9.324920002941221e-05,
- 9.30933428113913e-05,
- 9.293751814571654e-05,
- 9.278172645840654e-05,
- 9.262596817538988e-05,
- 9.247024372250367e-05,
- 9.231455352549273e-05,
- 9.215889801000804e-05,
- 9.200327760160585e-05,
- 9.184769272574638e-05,
- 9.169214380779284e-05,
- 9.153663127300994e-05,
- 9.138115554656304e-05,
- 9.122571705351685e-05,
- 9.107031621883429e-05,
- 9.091495346737525e-05,
- 9.075962922389559e-05,
- 9.060434391304587e-05,
- 9.044909795937023e-05,
- 9.029389178730512e-05,
- 9.01387258211783e-05,
- 8.998360048520761e-05,
- 8.982851620349982e-05,
- 8.967347340004938e-05,
- 8.951847249873743e-05,
- 8.936351392333052e-05,
- 8.920859809747946e-05,
- 8.905372544471818e-05,
- 8.889889638846267e-05,
- 8.874411135200959e-05,
- 8.858937075853538e-05,
- 8.843467503109482e-05,
- 8.828002459262023e-05,
- 8.812541986591992e-05,
- 8.797086127367735e-05,
- 8.781634923844979e-05,
- 8.766188418266728e-05,
- 8.750746652863135e-05,
- 8.735309669851397e-05,
- 8.719877511435638e-05,
- 8.70445021980679e-05,
- 8.689027837142476e-05,
- 8.673610405606905e-05,
- 8.658197967350744e-05,
- 8.642790564511014e-05,
- 8.627388239210958e-05,
- 8.611991033559955e-05,
- 8.59659898965337e-05,
- 8.581212149572469e-05,
- 8.565830555384278e-05,
- 8.550454249141496e-05,
- 8.535083272882353e-05,
- 8.519717668630514e-05,
- 8.504357478394947e-05,
- 8.489002744169834e-05,
- 8.473653507934427e-05,
- 8.458309811652954e-05,
- 8.442971697274494e-05,
- 8.42763920673287e-05,
- 8.412312381946516e-05,
- 8.396991264818392e-05,
- 8.38167589723585e-05,
- 8.366366321070519e-05,
- 8.351062578178188e-05,
- 8.335764710398713e-05,
- 8.32047275955588e-05,
- 8.305186767457298e-05,
- 8.289906775894282e-05,
- 8.274632826641749e-05,
- 8.259364961458095e-05,
- 8.244103222085076e-05,
- 8.228847650247707e-05,
- 8.213598287654138e-05,
- 8.198355175995546e-05,
- 8.183118356946013e-05,
- 8.167887872162425e-05,
- 8.15266376328434e-05,
- 8.137446071933897e-05,
- 8.122234839715674e-05,
- 8.10703010821661e-05,
- 8.091831919005853e-05,
- 8.076640313634674e-05,
- 8.061455333636338e-05,
- 8.04627702052601e-05,
- 8.03110541580061e-05,
- 8.015940560938728e-05,
- 8.0007824974005e-05,
- 7.985631266627493e-05,
- 7.970486910042588e-05,
- 7.955349469049884e-05,
- 7.940218985034565e-05,
- 7.925095499362798e-05,
- 7.90997905338161e-05,
- 7.89486968841879e-05,
- 7.879767445782768e-05,
- 7.864672366762496e-05,
- 7.849584492627338e-05,
- 7.834503864626977e-05,
- 7.819430523991267e-05,
- 7.804364511930147e-05,
- 7.789305869633517e-05,
- 7.774254638271133e-05,
- 7.759210858992487e-05,
- 7.744174572926695e-05,
- 7.729145821182386e-05,
- 7.714124644847602e-05,
- 7.699111084989659e-05,
- 7.684105182655057e-05,
- 7.669106978869361e-05,
- 7.65411651463709e-05,
- 7.639133830941596e-05,
- 7.624158968744966e-05,
- 7.609191968987902e-05,
- 7.594232872589613e-05,
- 7.57928172044769e-05,
- 7.564338553438016e-05,
- 7.54940341241464e-05,
- 7.534476338209666e-05,
- 7.519557371633143e-05,
- 7.504646553472962e-05,
- 7.489743924494724e-05,
- 7.474849525441653e-05,
- 7.45996339703446e-05,
- 7.44508557997126e-05,
- 7.43021611492743e-05,
- 7.415355042555527e-05,
- 7.400502403485144e-05,
- 7.38565823832284e-05,
- 7.37082258765199e-05,
- 7.355995492032698e-05,
- 7.341176992001677e-05,
- 7.326367128072145e-05,
- 7.311565940733694e-05,
- 7.296773470452212e-05,
- 7.281989757669748e-05,
- 7.267214842804409e-05,
- 7.252448766250243e-05,
- 7.237691568377144e-05,
- 7.222943289530726e-05,
- 7.208203970032223e-05,
- 7.193473650178363e-05,
- 7.178752370241289e-05,
- 7.164040170468411e-05,
- 7.149337091082326e-05,
- 7.134643172280685e-05,
- 7.119958454236108e-05,
- 7.10528297709605e-05,
- 7.090616780982707e-05,
- 7.075959905992896e-05,
- 7.061312392197957e-05,
- 7.046674279643631e-05,
- 7.032045608349958e-05,
- 7.017426418311172e-05,
- 7.002816749495576e-05,
- 6.988216641845445e-05,
- 6.973626135276916e-05,
- 6.959045269679885e-05,
- 6.944474084917872e-05,
- 6.929912620827944e-05,
- 6.915360917220585e-05,
- 6.900819013879603e-05,
- 6.886286950562e-05,
- 6.871764766997887e-05,
- 6.857252502890357e-05,
- 6.842750197915387e-05,
- 6.828257891721722e-05,
- 6.813775623930776e-05,
- 6.799303434136516e-05,
- 6.784841361905358e-05,
- 6.770389446776053e-05,
- 6.755947728259585e-05,
- 6.741516245839062e-05,
- 6.727095038969606e-05,
- 6.712684147078242e-05,
- 6.698283609563805e-05,
- 6.683893465796811e-05,
- 6.669513755119365e-05,
- 6.655144516845041e-05,
- 6.640785790258798e-05,
- 6.626437614616837e-05,
- 6.612100029146528e-05,
- 6.597773073046279e-05,
- 6.583456785485445e-05,
- 6.569151205604204e-05,
- 6.554856372513468e-05,
- 6.540572325294766e-05,
- 6.526299103000132e-05,
- 6.51203674465202e-05,
- 6.497785289243165e-05,
- 6.483544775736506e-05,
- 6.469315243065065e-05,
- 6.455096730131837e-05,
- 6.440889275809691e-05,
- 6.426692918941272e-05,
- 6.41250769833887e-05,
- 6.398333652784345e-05,
- 6.384170821028987e-05,
- 6.370019241793445e-05,
- 6.355878953767594e-05,
- 6.341749995610435e-05,
- 6.327632405950011e-05,
- 6.313526223383267e-05,
- 6.299431486475965e-05,
- 6.285348233762582e-05,
- 6.271276503746191e-05,
- 6.257216334898368e-05,
- 6.243167765659073e-05,
- 6.229130834436562e-05,
- 6.215105579607266e-05,
- 6.201092039515698e-05,
- 6.187090252474336e-05,
- 6.17310025676354e-05,
- 6.159122090631417e-05,
- 6.145155792293747e-05,
- 6.131201399933846e-05,
- 6.117258951702499e-05,
- 6.1033284857178234e-05,
- 6.089410040065178e-05,
- 6.0755036527970675e-05,
- 6.061609361933023e-05,
- 6.047727205459498e-05,
- 6.033857221329789e-05,
- 6.0199994474638934e-05,
- 6.0061539217484445e-05,
- 5.992320682036572e-05,
- 5.978499766147833e-05,
- 5.96469121186808e-05,
- 5.9508950569493744e-05,
- 5.9371113391098714e-05,
- 5.923340096033739e-05,
- 5.909581365371019e-05,
- 5.8958351847375664e-05,
- 5.8821015917149046e-05,
- 5.868380623850161e-05,
- 5.854672318655934e-05,
- 5.8409767136102034e-05,
- 5.827293846156237e-05,
- 5.813623753702468e-05,
- 5.799966473622403e-05,
- 5.786322043254533e-05,
- 5.772690499902202e-05,
- 5.759071880833536e-05,
- 5.745466223281315e-05,
- 5.731873564442885e-05,
- 5.718293941480063e-05,
- 5.704727391519016e-05,
- 5.691173951650171e-05,
- 5.677633658928117e-05,
- 5.6641065503715017e-05,
- 5.650592662962919e-05,
- 5.637092033648818e-05,
- 5.6236046993394116e-05,
- 5.610130696908551e-05,
- 5.596670063193644e-05,
- 5.583222834995554e-05,
- 5.569789049078483e-05,
- 5.556368742169895e-05,
- 5.542961950960391e-05,
- 5.5295687121036317e-05,
- 5.5161890622162184e-05,
- 5.5028230378776036e-05,
- 5.4894706756299826e-05,
- 5.4761320119782116e-05,
- 5.46280708338968e-05,
- 5.449495926294241e-05,
- 5.436198577084084e-05,
- 5.4229150721136586e-05,
- 5.40964544769956e-05,
- 5.396389740120427e-05,
- 5.383147985616869e-05,
- 5.3699202203913304e-05,
- 5.356706480608015e-05,
- 5.343506802392789e-05,
- 5.330321221833061e-05,
- 5.317149774977712e-05,
- 5.303992497836973e-05,
- 5.290849426382332e-05,
- 5.277720596546451e-05,
- 5.264606044223047e-05,
- 5.2515058052668014e-05,
- 5.2384199154932744e-05,
- 5.225348410678782e-05,
- 5.212291326560327e-05,
- 5.199248698835471e-05,
- 5.1862205631622704e-05,
- 5.1732069551591366e-05,
- 5.160207910404785e-05,
- 5.14722346443811e-05,
- 5.134253652758087e-05,
- 5.121298510823684e-05,
- 5.1083580740537716e-05,
- 5.0954323778270003e-05,
- 5.08252145748174e-05,
- 5.069625348315948e-05,
- 5.056744085587093e-05,
- 5.0438777045120604e-05,
- 5.0310262402670425e-05,
- 5.0181897279874485e-05,
- 5.005368202767819e-05,
- 4.9925616996617074e-05,
- 4.9797702536816114e-05,
- 4.966993899798853e-05,
- 4.954232672943491e-05,
- 4.941486608004237e-05,
- 4.928755739828352e-05,
- 4.9160401032215306e-05,
- 4.903339732947847e-05,
- 4.890654663729622e-05,
- 4.877984930247357e-05,
- 4.865330567139617e-05,
- 4.852691609002944e-05,
- 4.8400680903917706e-05,
- 4.82746004581832e-05,
- 4.814867509752493e-05,
- 4.802290516621807e-05,
- 4.789729100811286e-05,
- 4.7771832966633576e-05,
- 4.764653138477771e-05,
- 4.752138660511496e-05,
- 4.739639896978642e-05,
- 4.727156882050357e-05,
- 4.714689649854715e-05,
- 4.702238234476658e-05,
- 4.689802669957885e-05,
- 4.677382990296751e-05,
- 4.664979229448181e-05,
- 4.652591421323592e-05,
- 4.640219599790771e-05,
- 4.627863798673815e-05,
- 4.6155240517530006e-05,
- 4.603200392764729e-05,
- 4.59089285540142e-05,
- 4.578601473311404e-05,
- 4.566326280098849e-05,
- 4.554067309323671e-05,
- 4.541824594501422e-05,
- 4.529598169103218e-05,
- 4.517388066555648e-05,
- 4.505194320240652e-05,
- 4.4930169634954756e-05,
- 4.4808560296125436e-05,
- 4.468711551839388e-05,
- 4.456583563378545e-05,
- 4.444472097387471e-05,
- 4.432377186978452e-05,
- 4.4202988652185185e-05,
- 4.4082371651293287e-05,
- 4.3961921196871166e-05,
- 4.384163761822578e-05,
- 4.372152124420782e-05,
- 4.3601572403210874e-05,
- 4.3481791423170454e-05,
- 4.3362178631563205e-05,
- 4.324273435540601e-05,
- 4.312345892125483e-05,
- 4.300435265520421e-05,
- 4.288541588288617e-05,
- 4.276664892946925e-05,
- 4.264805211965776e-05,
- 4.252962577769091e-05,
- 4.241137022734172e-05,
- 4.229328579191645e-05,
- 4.217537279425332e-05,
- 4.2057631556722013e-05,
- 4.1940062401222605e-05,
- 4.182266564918466e-05,
- 4.1705441621566384e-05,
- 4.158839063885384e-05,
- 4.147151302105991e-05,
- 4.135480908772357e-05,
- 4.12382791579089e-05,
- 4.1121923550204244e-05,
- 4.100574258272146e-05,
- 4.088973657309483e-05,
- 4.077390583848033e-05,
- 4.065825069555484e-05,
- 4.054277146051504e-05,
- 4.042746844907679e-05,
- 4.031234197647412e-05,
- 4.019739235745835e-05,
- 4.008261990629741e-05,
- 3.996802493677484e-05,
- 3.985360776218878e-05,
- 3.973936869535153e-05,
- 3.962530804858827e-05,
- 3.95114261337365e-05,
- 3.9397723262145015e-05,
- 3.928419974467307e-05,
- 3.917085589168968e-05,
- 3.9057692013072654e-05,
- 3.894470841820761e-05,
- 3.883190541598744e-05,
- 3.871928331481127e-05,
- 3.860684242258358e-05,
- 3.8494583046713477e-05,
- 3.838250549411384e-05,
- 3.8270610071200376e-05,
- 3.815889708389098e-05,
- 3.8047366837604574e-05,
- 3.793601963726066e-05,
- 3.782485578727825e-05,
- 3.771387559157504e-05,
- 3.7603079353566624e-05,
- 3.7492467376165736e-05,
- 3.7382039961781235e-05,
- 3.727179741231754e-05,
- 3.716174002917352e-05,
- 3.7051868113241836e-05,
- 3.6942181964908154e-05,
- 3.68326818840502e-05,
- 3.672336817003699e-05,
- 3.66142411217281e-05,
- 3.650530103747265e-05,
- 3.639654821510872e-05,
- 3.628798295196235e-05,
- 3.617960554484681e-05,
- 3.60714162900618e-05,
- 3.5963415483392694e-05,
- 3.585560342010944e-05,
- 3.574798039496619e-05,
- 3.5640546702200116e-05,
- 3.55333026355309e-05,
- 3.542624848815973e-05,
- 3.5319384552768486e-05,
- 3.521271112151912e-05,
- 3.510622848605278e-05,
- 3.4999936937488824e-05,
- 3.489383676642433e-05,
- 3.4787928262933145e-05,
- 3.468221171656503e-05,
- 3.4576687416345054e-05,
- 3.447135565077253e-05,
- 3.43662167078205e-05,
- 3.426127087493493e-05,
- 3.415651843903357e-05,
- 3.405195968650564e-05,
- 3.3947594903210826e-05,
- 3.3843424374478376e-05,
- 3.373944838510661e-05,
- 3.363566721936188e-05,
- 3.353208116097791e-05,
- 3.342869049315507e-05,
- 3.332549549855949e-05,
- 3.322249645932229e-05,
- 3.311969365703899e-05,
- 3.301708737276847e-05,
- 3.291467788703243e-05,
- 3.281246547981447e-05,
- 3.27104504305594e-05,
- 3.2608633018172484e-05,
- 3.2507013521018675e-05,
- 3.240559221692172e-05,
- 3.230436938316366e-05,
- 3.22033452964838e-05,
- 3.210252023307819e-05,
- 3.200189446859869e-05,
- 3.190146827815225e-05,
- 3.180124193630029e-05,
- 3.170121571705783e-05,
- 3.1601389893892635e-05,
- 3.150176473972476e-05,
- 3.1402340526925595e-05,
- 3.1303117527317104e-05,
- 3.120409601217118e-05,
- 3.110527625220884e-05,
- 3.1006658517599575e-05,
- 3.090824307796054e-05,
- 3.081003020235569e-05,
- 3.0712020159295335e-05,
- 3.0614213216735204e-05,
- 3.051660964207573e-05,
- 3.0419209702161324e-05,
- 3.0322013663279762e-05,
- 3.022502179116125e-05,
- 3.0128234350977934e-05,
- 3.0031651607342883e-05,
- 2.9935273824309673e-05,
- 2.9839101265371512e-05,
- 2.9743134193460485e-05,
- 2.964737287094687e-05,
- 2.955181755963852e-05,
- 2.9456468520779957e-05,
- 2.936132601505186e-05,
- 2.9266390302570178e-05,
- 2.917166164288551e-05,
- 2.907714029498243e-05,
- 2.8982826517278678e-05,
- 2.8888720567624494e-05,
- 2.879482270330201e-05,
- 2.8701133181024336e-05,
- 2.86076522569351e-05,
- 2.8514380186607563e-05,
- 2.8421317225043974e-05,
- 2.8328463626674938e-05,
- 2.823581964535871e-05,
- 2.8143385534380284e-05,
- 2.805116154645106e-05,
- 2.795914793370788e-05,
- 2.786734494771246e-05,
- 2.7775752839450655e-05,
- 2.7684371859331767e-05,
- 2.7593202257187918e-05,
- 2.750224428227338e-05,
- 2.7411498183263677e-05,
- 2.732096420825524e-05,
- 2.7230642604764514e-05,
- 2.7140533619727284e-05,
- 2.7050637499498083e-05,
- 2.696095448984945e-05,
- 2.687148483597132e-05,
- 2.6782228782470374e-05,
- 2.669318657336915e-05,
- 2.660435845210569e-05,
- 2.651574466153272e-05,
- 2.642734544391693e-05,
- 2.633916104093839e-05,
- 2.625119169368993e-05,
- 2.616343764267633e-05,
- 2.6075899127813907e-05,
- 2.59885763884295e-05,
- 2.590146966326021e-05,
- 2.5814579190452517e-05,
- 2.5727905207561597e-05,
- 2.5641447951550874e-05,
- 2.555520765879116e-05,
- 2.5469184565060083e-05,
- 2.5383378905541523e-05,
- 2.529779091482493e-05,
- 2.521242082690451e-05,
- 2.5127268875178898e-05,
- 2.5042335292450212e-05,
- 2.4957620310923695e-05,
- 2.4873124162206826e-05,
- 2.4788847077308847e-05,
- 2.4704789286640097e-05,
- 2.4620951020011424e-05,
- 2.453733250663335e-05,
- 2.445393397511575e-05,
- 2.437075565346703e-05,
- 2.428779776909353e-05,
- 2.4205060548798956e-05,
- 2.4122544218783662e-05,
- 2.404024900464418e-05,
- 2.3958175131372522e-05,
- 2.3876322823355436e-05,
- 2.3794692304374045e-05,
- 2.371328379760311e-05,
- 2.363209752561035e-05,
- 2.3551133710355896e-05,
- 2.347039257319179e-05,
- 2.3389874334861157e-05,
- 2.330957921549787e-05,
- 2.322950743462562e-05,
- 2.314965921115764e-05,
- 2.307003476339596e-05,
- 2.299063430903075e-05,
- 2.2911458065139792e-05,
- 2.2832506248187967e-05,
- 2.2753779074026473e-05,
- 2.2675276757892457e-05,
- 2.2596999514408216e-05,
- 2.2518947557580736e-05,
- 2.2441121100801117e-05,
- 2.236352035684391e-05,
- 2.2286145537866558e-05,
- 2.2208996855408906e-05,
- 2.213207452039247e-05,
- 2.205537874312e-05,
- 2.1978909733274825e-05,
- 2.190266769992026e-05,
- 2.1826652851499143e-05,
- 2.1750865395833207e-05,
- 2.1675305540122366e-05,
- 2.1599973490944466e-05,
- 2.1524869454254394e-05,
- 2.144999363538377e-05,
- 2.1375346239040193e-05,
- 2.130092746930678e-05,
- 2.1226737529641624e-05,
- 2.1152776622877234e-05,
- 2.1079044951219805e-05,
- 2.1005542716248975e-05,
- 2.0932270118917064e-05,
- 2.085922735954852e-05,
- 2.0786414637839468e-05,
- 2.0713832152857093e-05,
- 2.0641480103039153e-05,
- 2.0569358686193427e-05,
- 2.049746809949704e-05,
- 2.0425808539496133e-05,
- 2.035438020210526e-05,
- 2.028318328260672e-05,
- 2.0212217975650166e-05,
- 2.0141484475252074e-05,
- 2.0070982974795096e-05,
- 2.0000713667027687e-05,
- 1.993067674406339e-05,
- 1.9860872397380515e-05,
- 1.9791300817821486e-05,
- 1.9721962195592336e-05,
- 1.9652856720262194e-05,
- 1.9583984580762825e-05,
- 1.951534596538799e-05,
- 1.9446941061793063e-05,
- 1.9378770056994478e-05,
- 1.9310833137369088e-05,
- 1.9243130488653877e-05,
- 1.917566229594529e-05,
- 1.9108428743698775e-05,
- 1.904143001572834e-05,
- 1.8974666295205897e-05,
- 1.8908137764660935e-05,
- 1.884184460597994e-05,
- 1.8775787000405825e-05,
- 1.8709965128537588e-05,
- 1.864437917032973e-05,
- 1.8579029305091746e-05,
- 1.8513915711487652e-05,
- 1.8449038567535537e-05,
- 1.838439805060702e-05,
- 1.8319994337426875e-05,
- 1.8255827604072303e-05,
- 1.8191898025972746e-05,
- 1.812820577790924e-05,
- 1.8064751034013946e-05,
- 1.800153396776969e-05,
- 1.793855475200957e-05,
- 1.7875813558916298e-05,
- 1.7813310560021967e-05,
- 1.7751045926207302e-05,
- 1.76890198277015e-05,
- 1.7627232434081526e-05,
- 1.756568391427177e-05,
- 1.75043744365435e-05,
- 1.7443304168514532e-05,
- 1.73824732771486e-05,
- 1.7321881928755075e-05,
- 1.7261530288988383e-05,
- 1.720141852284757e-05,
- 1.7141546794675948e-05,
- 1.7081915268160542e-05,
- 1.7022524106331624e-05,
- 1.6963373471562405e-05,
- 1.6904463525568425e-05,
- 1.684579442940726e-05,
- 1.6787366343477974e-05,
- 1.6729179427520695e-05,
- 1.6671233840616247e-05,
- 1.6613529741185683e-05,
- 1.6556067286989747e-05,
- 1.649884663512863e-05,
- 1.644186794204141e-05,
- 1.638513136350566e-05,
- 1.6328637054637013e-05,
- 1.6272385169888727e-05,
- 1.6216375863051337e-05,
- 1.6160609287252176e-05,
- 1.6105085594954873e-05,
- 1.60498049379591e-05,
- 1.5994767467400096e-05,
- 1.593997333374817e-05,
- 1.5885422686808397e-05,
- 1.5831115675720174e-05,
- 1.5777052448956777e-05,
- 1.572323315432503e-05,
- 1.5669657938964807e-05,
- 1.56163269493487e-05,
- 1.5563240331281616e-05,
- 1.551039822990034e-05,
- 1.545780078967313e-05,
- 1.5405448154399413e-05,
- 1.535334046720928e-05,
- 1.5301477870563166e-05,
- 1.5249860506251423e-05,
- 1.5198488515393948e-05,
- 1.5147362038439817e-05,
- 1.5096481215166858e-05,
- 1.5045846184681296e-05,
- 1.4995457085417375e-05,
- 1.4945314055136963e-05,
- 1.489541723092921e-05,
- 1.4845766749210105e-05,
- 1.4796362745722175e-05,
- 1.4747205355534067e-05,
- 1.4698294713040261e-05,
- 1.4649630951960514e-05,
- 1.4601214205339733e-05,
- 1.4553044605547445e-05,
- 1.4505122284277506e-05,
- 1.445744737254769e-05,
- 1.4410020000699399e-05,
- 1.4362840298397263e-05,
- 1.431590839462879e-05,
- 1.4269224417703989e-05,
- 1.4222788495255078e-05,
- 1.4176600754236111e-05,
- 1.4130661320922593e-05,
- 1.4084970320911206e-05,
- 1.4039527879119367e-05,
- 1.3994334119784972e-05,
- 1.3949389166466085e-05,
- 1.3904693142040434e-05,
- 1.3860246168705272e-05,
- 1.3816048367976926e-05,
- 1.3772099860690487e-05,
- 1.3728400766999503e-05,
- 1.3684951206375628e-05,
- 1.3641751297608287e-05,
- 1.3598801158804419e-05,
- 1.3556100907388011e-05,
- 1.3513650660099939e-05,
- 1.3471450532997587e-05,
- 1.3429500641454454e-05,
- 1.338780110015997e-05,
- 1.3346352023119073e-05,
- 1.3305153523651944e-05,
- 1.3264205714393723e-05,
- 1.3223508707294186e-05,
- 1.3183062613617335e-05,
- 1.3142867543941295e-05,
- 1.3102923608157822e-05,
- 1.3063230915472133e-05,
- 1.302378957440252e-05,
- 1.2984599692780111e-05,
- 1.2945661377748526e-05,
- 1.2906974735763676e-05,
- 1.2868539872593322e-05,
- 1.283035689331692e-05,
- 1.2792425902325284e-05,
- 1.2754747003320302e-05,
- 1.2717320299314639e-05,
- 1.2680145892631462e-05,
- 1.2643223884904177e-05,
- 1.2606554377076191e-05,
- 1.2570137469400482e-05,
- 1.2533973261439498e-05,
- 1.249806185206482e-05,
- 1.2462403339456867e-05,
- 1.242699782110464e-05,
- 1.239184539380548e-05,
- 1.2356946153664776e-05,
- 1.232230019609573e-05,
- 1.2287907615819052e-05,
- 1.2253768506862737e-05,
- 1.2219882962561836e-05,
- 1.2186251075558115e-05,
- 1.215287293779987e-05,
- 1.2119748640541646e-05,
- 1.208687827434402e-05,
- 1.2054261929073317e-05,
- 1.2021899693901367e-05,
- 1.19897916573053e-05,
- 1.1957937907067259e-05,
- 1.1926338530274176e-05,
- 1.1894993613317536e-05,
- 1.1863903241893152e-05,
- 1.1833067501000918e-05,
- 1.1802486474944566e-05,
- 1.1772160247331445e-05,
- 1.1742088901072316e-05,
- 1.1712272518381085e-05,
- 1.168271118077462e-05,
- 1.1653404969072462e-05,
- 1.1624353963396685e-05,
- 1.159555824317164e-05,
- 1.1567017887123709e-05,
- 1.1538732973281154e-05,
- 1.1510703578973834e-05,
- 1.1482929780833054e-05,
- 1.1455411654791319e-05,
- 1.1428149276082123e-05,
- 1.1401142719239768e-05,
- 1.1374392058099152e-05,
- 1.1347897365795578e-05,
- 1.1321658714764514e-05,
- 1.1295676176741409e-05,
- 1.126994982276155e-05,
- 1.1244479723159811e-05,
- 1.1219265947570456e-05,
- 1.1194308564926982e-05,
- 1.116960764346195e-05,
- 1.1145163250706704e-05,
- 1.1120975453491292e-05,
- 1.1097044317944236e-05,
- 1.1073369909492333e-05,
- 1.104995229286053e-05,
- 1.102679153207167e-05,
- 1.1003887690446406e-05,
- 1.0981240830602968e-05,
- 1.0958851014457006e-05,
- 1.0936718303221404e-05,
- 1.0914842757406183e-05,
- 1.089322443681821e-05,
- 1.0871863400561175e-05,
- 1.085075970703534e-05,
- 1.082991341393738e-05,
- 1.0809324578260273e-05,
- 1.0788993256293097e-05,
- 1.0768919503620928e-05,
- 1.0749103375124634e-05,
- 1.0729544924980744e-05,
- 1.0710244206661316e-05,
- 1.0691201272933795e-05,
- 1.0672416175860808e-05,
- 1.0653888966800111e-05,
- 1.0635619696404393e-05,
- 1.0617608414621115e-05,
- 1.0599855170692465e-05,
- 1.0582360013155102e-05,
- 1.0565122989840136e-05,
- 1.0548144147872937e-05,
- 1.0531423533672998e-05,
- 1.0514961192953843e-05,
- 1.0498757170722896e-05,
- 1.0482811511281315e-05,
- 1.0467124258223927e-05,
- 1.0451695454439093e-05,
- 1.0436525142108568e-05,
- 1.0421613362707412e-05,
- 1.0406960157003853e-05,
- 1.0392565565059195e-05,
- 1.037842962622771e-05,
- 1.036455237915651e-05,
- 1.0350933861785459e-05,
- 1.0337574111347066e-05,
- 1.0324473164366373e-05,
- 1.0311631056660878e-05,
- 1.02990478233404e-05,
- 1.0286723498807036e-05,
- 1.0274658116755e-05,
- 1.0262851710170606e-05,
- 1.02513043113321e-05,
- 1.0240015951809633e-05,
- 1.0228986662465152e-05,
- 1.021821647345233e-05,
- 1.0207705414216424e-05,
- 1.0197453513494281e-05,
- 1.0187460799314207e-05,
- 1.0177727298995905e-05,
- 1.016825303915038e-05,
- 1.015903804567991e-05,
- 1.0150082343777922e-05,
- 1.0141385957928984e-05,
- 1.0132948911908665e-05,
- 1.0124771228783522e-05,
- 1.0116852930911042e-05,
- 1.0109194039939557e-05,
- 1.0101794576808155e-05,
- 1.0094654561746707e-05,
- 1.0087774014275745e-05,
- 1.0081152953206404e-05,
- 1.0074791396640419e-05,
- 1.0068689361970052e-05,
- 1.006284686587803e-05,
- 1.0057263924337509e-05,
- 1.0051940552612027e-05,
- 1.0046876765255499e-05,
- 1.0042072576112127e-05,
- 1.0037527998316364e-05,
- 1.0033243044292918e-05,
- 1.0029217725756705e-05,
- 1.0025452053712766e-05,
- 1.0021946038456325e-05,
- 1.0018699689572687e-05,
- 1.0015713015937242e-05,
- 1.0012986025715443e-05,
- 1.0010518726362785e-05,
- 1.0008311124624773e-05,
- 1.0006363226536888e-05,
- 1.0004675037424615e-05,
- 1.0003246561903414e-05,
- 1.0002077803878662e-05,
- 1.0001168766545706e-05,
- 1.0000519452389823e-05,
- 1.0000129863186222e-05,
- 1e-05
- ]
-}
diff --git a/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.json b/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.json
deleted file mode 100644
index 75990205f8..0000000000
--- a/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "aria",
- "representation": "cartesian",
- "frame": "camframe",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian",
- "stride": 3
- },
- "normalization_frames": 54,
- "archive_sha256": "7aaf6dc3c82a95a2c844e76918334ad2762dd0e0ced038fe117baf928b5cfde6"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.npz b/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.npz
deleted file mode 100644
index 0529be4357..0000000000
Binary files a/tests/fixtures/legacy_pipeline/aria-cartesian-camframe.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.json b/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.json
deleted file mode 100644
index 6cac58d2a1..0000000000
--- a/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "aria",
- "representation": "cartesian",
- "frame": "eef_frame",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian_wristframe_ypr",
- "stride": 3
- },
- "normalization_frames": 54,
- "archive_sha256": "d456f93c3dec0d8eac559f46a5f56cabc0c42f813d2e4feac78fbf34d92fe65c"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.npz b/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.npz
deleted file mode 100644
index 0e34574bf1..0000000000
Binary files a/tests/fixtures/legacy_pipeline/aria-cartesian-eef_frame.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.json b/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.json
deleted file mode 100644
index 5f43a2dfa3..0000000000
--- a/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "aria",
- "representation": "keypoints",
- "frame": "camframe",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_keypoints",
- "observations.images.front_img_1",
- "observations.state.keypoints"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "keypoints_headframe_ypr",
- "stride": 3
- },
- "normalization_frames": 54,
- "archive_sha256": "4df87e74e7b8e88aa69fa4f41ca3325ca65c0c2a4a0e8b6dc0ee987c2160088f"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.npz b/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.npz
deleted file mode 100644
index 08607bb708..0000000000
Binary files a/tests/fixtures/legacy_pipeline/aria-keypoints-camframe.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.json b/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.json
deleted file mode 100644
index 5ce0ccb599..0000000000
--- a/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "aria",
- "representation": "keypoints",
- "frame": "eef_frame",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_keypoints",
- "observations.images.front_img_1",
- "observations.state.keypoints"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "keypoints_wristframe_ypr",
- "stride": 3
- },
- "normalization_frames": 54,
- "archive_sha256": "fac7145706aa87b6e20b008833488d659c82b0ba70c9d9eed9aee29ff91f2894"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.npz b/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.npz
deleted file mode 100644
index b834ed1c87..0000000000
Binary files a/tests/fixtures/legacy_pipeline/aria-keypoints-eef_frame.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.json b/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.json
deleted file mode 100644
index 3b11468aa6..0000000000
--- a/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.json
+++ /dev/null
@@ -1,27 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "eva",
- "representation": "cartesian",
- "frame": "camframe",
- "identity": 6,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.images.left_wrist_img",
- "observations.images.right_wrist_img",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian"
- },
- "normalization_frames": 54,
- "archive_sha256": "a4908f0de8cae703a2b5281f9da4e8ae61f30a0817c874000c8c7c056899ecac"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.npz b/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.npz
deleted file mode 100644
index a1e4d5b079..0000000000
Binary files a/tests/fixtures/legacy_pipeline/eva-cartesian-camframe.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.json b/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.json
deleted file mode 100644
index f357143472..0000000000
--- a/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.json
+++ /dev/null
@@ -1,27 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "eva",
- "representation": "cartesian",
- "frame": "eef_frame",
- "identity": 6,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.images.left_wrist_img",
- "observations.images.right_wrist_img",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian_wristframe_ypr"
- },
- "normalization_frames": 54,
- "archive_sha256": "687736aeb8cd1d280c79a368b34f4db612459d9aa4d98b5aa20be3d5b601cba7"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.npz b/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.npz
deleted file mode 100644
index f57778b1f8..0000000000
Binary files a/tests/fixtures/legacy_pipeline/eva-cartesian-eef_frame.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.json b/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.json
deleted file mode 100644
index ce6a59f921..0000000000
--- a/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "mecka",
- "representation": "cartesian",
- "frame": "camframe",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian",
- "stride": 1
- },
- "normalization_frames": 54,
- "archive_sha256": "7aaf6dc3c82a95a2c844e76918334ad2762dd0e0ced038fe117baf928b5cfde6"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.npz b/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.npz
deleted file mode 100644
index 0529be4357..0000000000
Binary files a/tests/fixtures/legacy_pipeline/mecka-cartesian-camframe.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.json b/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.json
deleted file mode 100644
index 07a13e6bda..0000000000
--- a/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "mecka",
- "representation": "cartesian",
- "frame": "eef_frame",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian_wristframe_ypr",
- "stride": 1
- },
- "normalization_frames": 54,
- "archive_sha256": "d456f93c3dec0d8eac559f46a5f56cabc0c42f813d2e4feac78fbf34d92fe65c"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.npz b/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.npz
deleted file mode 100644
index 0e34574bf1..0000000000
Binary files a/tests/fixtures/legacy_pipeline/mecka-cartesian-eef_frame.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.json b/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.json
deleted file mode 100644
index b10df9e19c..0000000000
--- a/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "scale",
- "representation": "cartesian",
- "frame": "camframe",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian",
- "stride": 1
- },
- "normalization_frames": 54,
- "archive_sha256": "7aaf6dc3c82a95a2c844e76918334ad2762dd0e0ced038fe117baf928b5cfde6"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.npz b/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.npz
deleted file mode 100644
index 0529be4357..0000000000
Binary files a/tests/fixtures/legacy_pipeline/scale-cartesian-camframe.npz and /dev/null differ
diff --git a/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.json b/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.json
deleted file mode 100644
index 5f33ed50f7..0000000000
--- a/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
- "source_commit": "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2",
- "vendor": "scale",
- "representation": "cartesian",
- "frame": "eef_frame",
- "identity": 3,
- "indices": [
- 0,
- 7,
- 29,
- 47,
- 50
- ],
- "model_keys": [
- "actions_cartesian",
- "observations.images.front_img_1",
- "observations.state.ee_pose"
- ],
- "api": "legacy",
- "transform_arguments": {
- "mode": "cartesian_wristframe_ypr",
- "stride": 1
- },
- "normalization_frames": 54,
- "archive_sha256": "d456f93c3dec0d8eac559f46a5f56cabc0c42f813d2e4feac78fbf34d92fe65c"
-}
\ No newline at end of file
diff --git a/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.npz b/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.npz
deleted file mode 100644
index 0e34574bf1..0000000000
Binary files a/tests/fixtures/legacy_pipeline/scale-cartesian-eef_frame.npz and /dev/null differ
diff --git a/tests/fixtures/libero_native_scontrol_408787.txt b/tests/fixtures/libero_native_scontrol_408787.txt
deleted file mode 100644
index 9b80c1cc7d..0000000000
--- a/tests/fixtures/libero_native_scontrol_408787.txt
+++ /dev/null
@@ -1 +0,0 @@
-JobId=408787 JobName=af-noaug-exact-gpu-v2 UserId=paphiwetsa3(41003) GroupId=loaner(41000) MCS_label=N/A Priority=10000 Nice=0 Account=loaner QOS=loaner-normal JobState=RUNNING Reason=None Dependency=(null) Requeue=0 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0 DerivedExitCode=0:0 RunTime=00:39:27 TimeLimit=02:00:00 TimeMin=N/A SubmitTime=2026-10-07T23:17:33 EligibleTime=2026-10-07T23:17:33 AccrueTime=2026-10-07T23:17:33 StartTime=2026-10-07T23:17:33 EndTime=2026-10-08T01:17:33 Deadline=N/A PreemptEligibleTime=2026-10-07T23:22:33 PreemptTime=None SuspendTime=None SecsPreSuspend=0 LastSchedEval=2026-10-07T23:17:33 Scheduler=Main Partition=rl2 AllocNode:Sid=loaner-gateway:61717 ReqNodeList=gpu19 ExcNodeList=(null) NodeList=gpu19 BatchHost=gpu19 NumNodes=1 NumCPUs=8 NumTasks=1 CPUs/Task=8 ReqB:S:C:T=0:0:*:* ReqTRES=cpu=8,mem=128G,node=1,billing=8 AllocTRES=cpu=8,mem=128G,node=1,billing=8 Socks/Node=* NtasksPerN:B:S:C=0:0:*:* CoreSpec=* JOB_GRES=gpu:h100:1 Nodes=gpu19 CPU_IDs=32-35,144-147 Mem=131072 GRES=gpu:h100:1(IDX:3) MinCPUsNode=8 MinMemoryNode=128G MinTmpDiskNode=0 Features=(null) DelayBoot=00:00:00 OverSubscribe=OK Contiguous=0 Licenses=(null) Network=(null) Command=/workspace/users/paphiwetsa3/tasks/cotrain-velocity-first-wave-20261007/rows/usocket-noaug/noaug_exact_gpu_v2.sbatch WorkDir=/workspace/users/paphiwetsa3/tasks/cotrain-velocity-first-wave-20261007/rows/usocket-noaug StdErr=/workspace/users/paphiwetsa3/tasks/cotrain-velocity-first-wave-20261007/rows/usocket-noaug/logs/noaug-exact-gpu-v2-408787.err StdIn=/dev/null StdOut=/workspace/users/paphiwetsa3/tasks/cotrain-velocity-first-wave-20261007/rows/usocket-noaug/logs/noaug-exact-gpu-v2-408787.out Power= MemPerTres=gpu:256000 TresPerNode=gres/gpu:h100:1 TresPerTask=cpu:8
diff --git a/tests/fixtures/libero_native_suite_corpus_authority.json b/tests/fixtures/libero_native_suite_corpus_authority.json
deleted file mode 100644
index 814ee77b94..0000000000
--- a/tests/fixtures/libero_native_suite_corpus_authority.json
+++ /dev/null
@@ -1,36 +0,0 @@
-{
- "recorded_authority_sha256": "559b6b3884eed7e39bdbdef12af67f4b033030af355a408db7c62354de23c31f",
- "schema": "libero-native-suite-corpus-authority/v1",
- "suites": {
- "libero10": {
- "dataset_logical_sha256": "3ed2b86cf97eaa27045af81e6d6c4834c8659d15eda69a2394d835e602014e1d",
- "historical_split_sha256": "c588c831ef0deab61a59a21840ea9e11dafe92a170d9023a87e0acee831df236",
- "training_seeds": [
- 42
- ]
- },
- "libero_goal": {
- "dataset_logical_sha256": "c190385dfc8a7d05b966a4b5f94651826b8a906b1e63d8b348edc70df52c1f60",
- "historical_split_sha256": "207c24115104bf0545e7f165abb91cf052eb99f6568446f11e835f9a3e4c6af8",
- "training_seeds": [
- 42,
- 43
- ]
- },
- "libero_object": {
- "dataset_logical_sha256": "40b5e3ccf4febcbae470690755adc8aa6a39471a4fcb7fa585a9bdb8e59c325d",
- "historical_split_sha256": "5ade2a86fad7552457e093944fd60efec8b51fbe9dfd20831f627d1c56e9b5d7",
- "training_seeds": [
- 42,
- 43
- ]
- },
- "libero_spatial": {
- "dataset_logical_sha256": "ee28d02e8c6cd28fb0a235fab2ccc649b662e6f3e10eb096d89be082d2871860",
- "historical_split_sha256": "15640eada8102c7071dcfcd153ef438cb98c61835474a2fb2c3c9415cf6d714e",
- "training_seeds": [
- 42
- ]
- }
- }
-}
diff --git a/tests/fixtures/pusht_before_real_sim_cleanup.json b/tests/fixtures/pusht_before_real_sim_cleanup.json
deleted file mode 100644
index 2cfe82172c..0000000000
--- a/tests/fixtures/pusht_before_real_sim_cleanup.json
+++ /dev/null
@@ -1,40 +0,0 @@
-{
- "contract": "SHA256 sorted compact JSON of unresolved composed config, excluding Hydra metadata; verified identical before/after directory move.",
- "recipes": {
- "action_flow_bc_usocket_bridge_likelihood_s42": "93a9a46c5c50dd063b4dc10ae457c6c503dc3f68f924b2bea94cc45ad6c6c4c5",
- "action_flow_bc_usocket_graph_section_s42": "f413c4d5adfbc84c9c2fb6145156b72171ac923fdd4e369c22fa5c96083a481a",
- "action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42": "9f8b97c9ab2a95a4be95d8bb125c774e1137e0c6d2ad9dfe0fd47db008249b87",
- "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42": "6223637fc6f6d47e180bcaf3fe7689e4e3e100ebb3bedb400a044dbe55ebcabf",
- "action_flow_bc_usocket_latent_fm_sg_recon1_s42": "9f0eea8acca4715cdcc174c4916edb54f4efcb024141bd2d78f28edc934d6f95",
- "action_flow_bc_usocket_recon100_s42": "33d779a7c1105d49f14807b7e9514f2cd8f9c3b6a1f4f01718eba898675756b4",
- "action_flow_bc_usocket_recon10_s42": "baf0e2745493a2c4c0f3dbe1b28375716fbb7dd1349245a3efd6e970cd6d4f1b",
- "action_flow_bc_usocket_recon10_warmup10k_s42": "b74cda0fa33c73b45eb2b3f29f0503d9bd12a48dae3f2e663bd5a6c0f63dda3f",
- "action_flow_bc_usocket_recon1_s42": "03fd25ca55fb0328904e41e1d00bd178d37b47555db1ff388130a54e7f432e69",
- "action_flow_chain_latent_fm_sg_unite_h384_sum14_cfg4_val8_s42": "e6401b0b233a11b7f577e536800b3b686979257f5f809a4f5e72a7f49e09ea8d",
- "action_flow_cotrain_uc_multiplier_interpolation": "1484a8fe2bafa7c297e58c5277254298c3331b48f1f88ed688ae265b75831270",
- "action_flow_usocket_candidate_common": "156337979165efd521a5aa752532b91c527e4a6cac487cda99345132f0c6dbd6",
- "action_flow_usocket_latent_fm_sg_unite_denoiser90m_sum14_cfg4_finalval1_s42": "9e35db31e65a8ffb4691c005ae8c955dda788add1f0d8aac1415f85b4d865708",
- "action_latent_vfm_usocket_val01_h16": "a7d96acf9a9afcc8b21a63dbbea43443075287315a145a4c7252d2f89c1f5182",
- "planar_uc_manual4919_dp_261m_af_obs_multiplier": "10799eebb8f927c1f265f4c50bacddec890525b270cdf25d2a35d0f98137a186",
- "planar_usocket_chain_manual4919_standard_dp_retimed_masked": "f9f7072e36d1a3f05b4ee669a6aa6644fb6fd69ba6263d8332975bfd675eb5de",
- "planar_v2_base": "2b81ed6636b484dba91f8f6c9924d65e25406bc564d392539a8e23f28ca066af",
- "planar_v2_chain_arc_bc": "4781bc8155faafecde4ce4f67d4bfa935db29ec6eee0aa04cb78c1fb2e8c3560",
- "planar_v2_chain_base": "58fe6f3c85eb69c6adf0e80b5567da445117ae9f033eba882da7721862fd0b78",
- "planar_v2_chain_direct_bc": "1c2ca35042bfecb658ea58eec0c1e58e2e9e9a778e302ce9dc0449b5461c870e",
- "planar_v2_chain_dp_paper": "4618c22f2295831805762865c2b08fc79461152824c1166ae6fd724a002efb21",
- "planar_v2_chain_dp_standard": "e52dd80834a0371d11d5614be73ba7b9329ec970376f638c21f14fbf34b1aaab",
- "planar_v2_cotrain_clean_dp_standard": "fafe963772ca0fb243b0d552d7c8f7f9e499d99f97f26c589214eb0f7180dfe5",
- "planar_v2_usocket_arc_bc": "9c0c7d575365f007220806fb9c226eb07a264641162f8aab7ee6d9860d351a97",
- "planar_v2_usocket_arc_graph_tok": "e52d3e962a3d4246e6637b2358477fdb35c7e64247aafa221b38bb1c04abfcab",
- "planar_v2_usocket_arc_hybrid_D40_M16_R24deg_bc": "88127379eaa9098b72dccda0bceb82fe145ddd633d809226186f8c9853e3ac34",
- "planar_v2_usocket_arc_paper_curvature_D40_M16_R24deg": "924e9d5aab3a586eb6c9b7170f80c7e49796240e23ab519ef5133140b22c0bb8",
- "planar_v2_usocket_arc_paper_uniform_D40_M16_R24deg": "609852fa193f2bc54655074c9ab90055117f3d698dad02d29de374e670ed1ae4",
- "planar_v2_usocket_base": "4eebb72701c219b503b96ddf8765630cc0a7160ce848620a26c77ab83c2d4abb",
- "planar_v2_usocket_direct_bc": "c4a46e4074481e2d82f03eb15078f884f3f10d22ffd0f3ed35bc2fb20a11f1b6",
- "planar_v2_usocket_dp_paper": "84cad67cd83cdb4324588dcbd2010e94245b5a29bf1cac7ebad1d1544d5bb38f",
- "planar_v2_usocket_dp_standard": "7e958a8973f08fd0b6afde306d3fc539a77f381215d5ded773a739368246d98f",
- "unite_cotrain_usocket_chain_val01_h16": "060912f7b91c8c9a2b1bfe4c432ffdcc03e1bb554498806fe8db4bc36b99f8fa",
- "unite_usocket_register_sweep_val01_h16": "2232b805b702d255cf97739d3cd26e5b5c45ea8679b2286fefe4aa4df740c34a"
- },
- "source_commit": "b6686ec9ad4fc8bbc89c97263fb63ec1891e1ae2"
-}
diff --git a/tests/fixtures/pusht_config_before_parent_cleanup.json b/tests/fixtures/pusht_config_before_parent_cleanup.json
deleted file mode 100644
index ff2c1db81e..0000000000
--- a/tests/fixtures/pusht_config_before_parent_cleanup.json
+++ /dev/null
@@ -1,88 +0,0 @@
-{
- "hash_contract": "SHA256 sorted compact JSON of unresolved composed config excluding hydra metadata; integer keys normalize through JSON",
- "recipes": {
- "action_flow_bc_usocket_bridge_likelihood_s42": "ff73c1ecfe60c29af17c6ae6806fb797984f22064af6656b7c40f745915caf56",
- "action_flow_bc_usocket_graph_section_s42": "4c7284b1f6cf021e6baccc8b3846b1155d73587c2162fe5606e798abe6851d74",
- "action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42": "4f586952aa5862aadd5fc5fc55387e1e33aeca0255b4cbd7082634fc6e2fdfd5",
- "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42": "3e2c27eb94568c5adb00c02b05e2a13ee700a598c53ed4224deecfdf1c370a09",
- "action_flow_bc_usocket_latent_fm_sg_recon1_s42": "bb976a267641d6db9d6d1bfcff5341dbec9e140463987e6517e53f73d040a493",
- "action_flow_bc_usocket_recon100_s42": "12617b5c6a477a08b1e35f5c0e079fe6dcca9c17b032439afc0f800dde4a420b",
- "action_flow_bc_usocket_recon10_s42": "d31d062bb28d81066265b958eb8da7a97a48fbffaf3797a29e7157f4cd2e72ab",
- "action_flow_bc_usocket_recon10_warmup10k_s42": "99e22eb4a39538e0fc05a4fab9fd123d2cf8fcb2ecc1dd57b18586baa01f2e59",
- "action_flow_bc_usocket_recon1_s42": "dd4f13fc095b59ed10fa7f29cd3b2e3e25cbd4e8ef4048c3c564438cd2f248fb",
- "action_flow_chain_latent_fm_sg_unite_h384_sum14_cfg4_val8_s42": "cc754f3f742641463584355992a5ab9a436da9bde736f961802a0ffd2c6176b2",
- "action_flow_cotrain_uc_multiplier_interpolation": "9755ee842621fc3e5aa483d373ff24348159045453d5ac2b258e82c2f0707af7",
- "action_flow_usocket_candidate_common": "5e4faf536d9e6d33d6c2e82b0a80c37bf4beb08954e4cdaf637afd7df7318f9d",
- "action_flow_usocket_latent_fm_sg_unite_denoiser90m_sum14_cfg4_finalval1_s42": "62150360dbb2e8daa02af64cf1405835e03e012c3ab30c45fe759a3fe66f0a7c",
- "action_latent_vfm_usocket_val01_h16": "c47bd51273b01fa5b516ec781ebe119e5d6e9774697990b6d3172705aca0b839",
- "planar_uc_manual4919_dp_261m_af_obs_multiplier": "e25f5f9f69b205cb10206e41a896922d4eeb5700a04597d6b086a4d3b44914a2",
- "planar_usocket_chain_manual4919_standard_dp_retimed_masked": "699dc867a101146e0fc6c6162b3c30514037d802dcf39f195bb9d4c29b0263ba",
- "planar_v2_base": "c9e7771e7f023a9742fff0e662693be6877372c736b3c1e8da75f32a3e2ee060",
- "planar_v2_chain_arc_bc": "da9b94f023d8bd098fcee001c8011a3af099c19982de203c62b6388a160fd068",
- "planar_v2_chain_base": "3ce62f0b980eabc2ce0b3228f9a11b2fbb73ca8ee1fd12eaac777b9ec3bc18da",
- "planar_v2_chain_direct_bc": "4d49cdd26c3066cd04f4a4632008369e7229dcf4562862a1356616fe966c21f1",
- "planar_v2_chain_dp_paper": "811fb66e1e7c25138c693251de0c698d437f579f26393e87a4fb89f86c942fd1",
- "planar_v2_chain_dp_standard": "bdf14917b32f5d1440417ffaca4321ff5531c284d60e46db29081236c8c8ba49",
- "planar_v2_cotrain_clean_dp_standard": "868969c07806c29e50846bbf8de0022437d54e4882ef6b0c3228f90b94737a2d",
- "planar_v2_usocket_arc_bc": "6085104c0e08534e8540282a5d33c5d01b247c2596f5f4281b10418e44a306a5",
- "planar_v2_usocket_arc_graph_tok": "8f019000f4a449dcbc890519682fb7156add1be19396a1d85857c9a81d4bfdda",
- "planar_v2_usocket_arc_hybrid_D40_M16_R24deg_bc": "752d6dda74daaebf2932ae57d544a7ed3c788fb83f1501993a8be51147f665db",
- "planar_v2_usocket_arc_paper_curvature_D40_M16_R24deg": "d464be794e7c2e700362aa3c9dbd448bdf53127057d37187d2874c99a77e7d53",
- "planar_v2_usocket_arc_paper_uniform_D40_M16_R24deg": "1874d622ce05170c9cc9456a3742a849bea7b86b24886d230d3455ae07af93ea",
- "planar_v2_usocket_base": "72885fa5bff3ef2b48c9a0a1a2fdee2fc75b7f87af3ba6fe834b4c191deb7f44",
- "planar_v2_usocket_direct_bc": "64b297d8bd5103de726f69c3c666f75fa6ce014b715448db092cb4edfd887923",
- "planar_v2_usocket_dp_paper": "79f87ad8c8fa6cc0835c82ef297222b4310b00d705bd895d350b41174523cc42",
- "planar_v2_usocket_dp_standard": "c7f54de8e746a66027e92bf9263e839fa1ce2d5db2a3825485f8b161c6a3d990",
- "unite_cotrain_usocket_chain_val01_h16": "56b9e9e6fdc1c8ad38d81148d7f7ba7713b2c817e54ba1dd14d4a15707a52d8b",
- "unite_usocket_register_sweep_val01_h16": "9823d8dd3107639ebe9a4755bba27b5f38757baa2c8b10e6e9395ef05082c84f"
- },
- "retired_data": [
- "action_flow_chain_points6_frozen4920_val01_h16",
- "action_flow_chain_points6_manual3000_val01_h16",
- "action_flow_chain_points6_manual4919_val01_h16",
- "action_flow_cotrain_usocket_chain_manual3000_val01_h16",
- "action_flow_cotrain_usocket_chain_manual4919_val01_h16",
- "action_flow_cotrain_usocket_chain_points6_frozen4920_val01_h16"
- ],
- "retired_recipes": {
- "action_flow_bc_usocket_latent_fm_sg_recon1_codec98k_s42": "7e1bba2ed089e54b2d74c2ddb34cb3472b77c8a2e79c9cdf9f560402edc8792d",
- "action_flow_bc_usocket_recon10_warmup10k_flow001_s42": "4f5bf734d0d8d2d72282fef6da53ffac86bbb76bb76a233793ca8923a2dca21e",
- "action_flow_chain_manual3000_avln_80k_s42": "07c2daf65519e44e0ae96d84343bcc68b16b384eb600035d68d8bfc8fb4eae34",
- "action_flow_chain_manual4919_avln_80k_s42": "3332315e2d95057076f5df303bfa5b8252c9f3fc5d185200ba12fb07cc128c42",
- "action_flow_chain_points6_latent_fm_sg_unite_h512d14h16_sum14_cfg4_val10k_s42": "e31546f683f40d60ffc719fe80dda8dee375826d4295a23ce488c648d57d97dd",
- "action_flow_cotrain_uc_latent_fm_sg_unite_h384d12h12_sum14_cfg4_val10k_s42": "79cbdbdd4c0221940f0fbd3574aab33ce5efc9f22b937fdca31126b1cf1a68f9",
- "action_flow_cotrain_uc_latent_fm_sg_unite_h512d14h16_sum14_cfg4_val10k_s42": "88040dd95a75a085058fd48adbc081ea6f00ec8485d6496b25b83ddc181f43b8",
- "action_flow_cotrain_uc_manual3000_avln_80k_s42": "ea67ff06292e4d080af6cd514ce642dc08830865f0ed25837a9979f534247dce",
- "action_flow_cotrain_uc_manual4919_avln_80k_s42": "978b5d06e0db777a8e3367cee1bd8b744ed2812f0749c35187f393ce4e018916",
- "action_flow_usocket_latent_fm_sg_unite_h384_s42": "3bbbb4f247acb88e2e74e96e20ce1617a18c2ee076d55e714f0d93a938a756d2",
- "action_flow_usocket_latent_fm_sg_unite_h384_sum14_cfg4_val8_deterministic_s42": "5f23aed86062a0827250927bba4b7433c450135688d58eeabf1da8b3acb819df",
- "action_flow_usocket_latent_fm_sg_unite_h384_sum14_cfg4_val8_s42": "51d39070058c361b37a0235eb72d7dea0f27163d30b16ea254a44c6c861a293e",
- "action_flow_usocket_latent_fm_sg_unite_h384_sum14_cfg4_val8_scale1_s42": "cd4eef7c8927a649e6c5fde911836275b3f1c76c64951265ced3df566d0d4796",
- "action_flow_usocket_latent_fm_sg_unite_h512d14h16_sum14_cfg4_val10k_s42": "afd8576ab5448604096f961b8a284d4a21f6fd8b05a83cbf59b355d5befe777a",
- "action_flow_usocket_refactored_c12_compat_s42": "a62734de32fedf721e3d1eda16e67e3deba5c9d03ec6451a164720ac2b7dfa22",
- "planar_chain_manual4919_standard_dp_retimed": "eb2ae443c790984c01a5159ca1f880169982a30e6bfe1c97b761449a44a4ab7e",
- "action_flow_cotrain_uc_speed_interpolation": "643f68448e85aff1d57273a2d2211689a3f2bb580f75c3d4c987bce7e981c045"
- },
- "retired_data": [
- "action_flow_chain_points6_frozen4920_val01_h16",
- "action_flow_chain_points6_manual3000_val01_h16",
- "action_flow_chain_points6_manual4919_val01_h16",
- "action_flow_cotrain_usocket_chain_manual3000_val01_h16",
- "action_flow_cotrain_usocket_chain_manual4919_val01_h16",
- "action_flow_cotrain_usocket_chain_points6_frozen4920_val01_h16"
- ],
- "parent_identity_metadata": {
- "parents/action_flow_usocket_h384_sum14": {
- "name": "action_flow_usocket_latent_fm_sg_unite_h384_sum14_cfg4_val8_s42",
- "description": "Action Flow H384 with released-UNITE flow aggregation, CFG, and validation coverage"
- },
- "parents/action_flow_usocket_h384": {
- "name": "action_flow_usocket_latent_fm_sg_unite_h384_s42",
- "description": "Action Flow stopgrad/JVP objective on separate released-UNITE H384 modules"
- },
- "parents/action_flow_cotrain_uc_manual4919": {
- "name": "action_flow_cotrain_uc_manual4919_avln_80k_s42",
- "description": "Manual U-Socket 2999 + ChainGripper 3000 standard/1919 obstacle recollections Action Flow AV-LN"
- }
- }
-}
diff --git a/tests/fixtures/stattempo_contracts.json b/tests/fixtures/stattempo_contracts.json
deleted file mode 100644
index 952c0035f8..0000000000
--- a/tests/fixtures/stattempo_contracts.json
+++ /dev/null
@@ -1,76 +0,0 @@
-{
- "source": "8e80f090c2463cdcc8007ad6be6856a178870187",
- "main": "7043c0251a350578c6da5a316a79c5ed7d21033d",
- "configs": {
- "bc_rl2_stattempo_slowpace366_arcdur_m200_d033": "a1eb967a2452aaa6d14ab601013ce42195a4eeb92f38211902c99a0b48f81095",
- "bc_rl2_stattempo_slowpace366_arcvel_m200_d033": "6b406ba86dfbdc90bed17aa4547773f3478abedfaba3aafb0cd523143c14854a",
- "bc_rl2_stattempo_slowpace366_arcdur_abstime_m200_d033": "b64f61df2c800bb109ac20ed2d969bb116aacb21d77c406d6e0020b4c26a1539",
- "bc_rl2_stattempo_slowpace366_arcvel_nosmooth_m200_d033": "50169b927f70cda897bc81f1750bf7b5bde73c3786d07eaecd940ba71fe196d2",
- "bc_rl2_stattempo_slowpace366_arc_projection_dur_m200_d033": "27db2736ea99844b2e237ab707f69e30f4bef775f06734d3506e612ff078ec4f",
- "bc_rl2_stattempo_slowpace366_arc_projection_vel_m200_d033": "207ea662c66645a669595f8d4a1276f0fa434b92bbaebb22832766291c2df3d8"
- },
- "codecs": {
- "motion/mean": "294796c37e92e78aa5bc3f5adc01e19bfd2fec768d4d4cb1e2a35f39019ebb67",
- "motion/profile": "55278af8ebf0c6af7c8d1dafe547ce440f5e4d98092abe57a0d440dafa2d4a84",
- "motion/logdur": "b551c891443c3adb741ee09c019edfec11112ed33ef4eb6bdf95a14cf0efd7f7",
- "motion/dur": "a24ab8b8a9ff57728eec12729a4a3623f276b9bb5d283d1c5eda89935a9d04bc",
- "motion/durhyb": "4c4d9b6eaa12ef474d5147d5fcc230341fe74285f67255687bd889c9ca95ac52",
- "motion/profhyb": "3eb3004d4bf2e926a85b9d1a44e39e5fd7e2e46059851a6db521fc17a53a6e51",
- "motion/durtri": "9f68f2c3ebc373e790a9166ead1f32f51cccd8e3c1d661ef5dab3d7a16214b68",
- "motion/proftri": "1a53518e0443895936bd6522387a985af24ffc19c442984712cfd8af2b3dc625",
- "motion/durevt": "44d9487249797e387d09ae9f11402336c63301c7fa541fc58bd4495cbd4ba671",
- "motion/stkevt": "b411a8edf5a2f7bfea162ae4aa4eeaa882097787954b7dc6459a9009f1380832",
- "motion/durref": "bf6273d564daa76e7b9a846a0724c1f4dbe3b279f575382584ca2ebd8735ac83",
- "motion/stkref": "02e34d462586e6d2969c00ccf0ad5eeba15d9275bc9be3f8ca6341584cfc4fc1",
- "hold/mean": "68f29b544b65f626c3259c7a70d99526939f19cdf8eddf031dfc0b035e7a572f",
- "hold/profile": "9c0db6929876854086ddaba0fe67763ea1304ecd5c5d8075d707b7c4ae2816d7",
- "hold/logdur": "0ab689937df851f021052424f4aa09eedd04a7f06d06b159235307062e4c1c37",
- "hold/dur": "9c0db6929876854086ddaba0fe67763ea1304ecd5c5d8075d707b7c4ae2816d7",
- "hold/durhyb": "519fc38d68309b6e2cb5c08ca8ea8c4037687c6536a387e5d49dc3047749fff9",
- "hold/profhyb": "4dfdebd0ced4b15e10874b09564eeb3aa64931e27a4ed27fc9d725e20e681053",
- "hold/durtri": "2c315840ab1fba1572920815c3b3f63ae3b530ff64f9c135b03c2c04761d4779",
- "hold/proftri": "a038fc0e34de75f7da5646b5c967afa05b0236337d9bff70a7733635b0d2832f",
- "hold/durevt": "a3e4f5317fb649b144281428e889e378ad1f7ff202d2401cafcec7498974c403",
- "hold/stkevt": "8ee78c048b8e808fe76641c67ab6a30274f0ad3acfc8fa2b9b8f3aa3051cd75e",
- "hold/durref": "e1074cc4893eb4bcfd9495a3d73ed438fbaf3f048e1d44139a0f8a25123c348f",
- "hold/stkref": "7cdbb56c003684d51679863bc161f3016ff8960d45b84880eb38763263ef77c9",
- "delayed/mean": "23599335b5b07fd33c7fa1ba5b93a9884cc8a79d5753f3592a7c1792ec8550a1",
- "delayed/profile": "8d98fc4d4a81d572983fb6e432b8d78845cffdb516d0fa737122d23d6c07266c",
- "delayed/logdur": "0a3fe9da66a43441e70978587adecd9fc137d86b228334235514ceaebba546e2",
- "delayed/dur": "2b0a368851cfbebaacaecf0c9fc392da57e2b15d0867afd48042df8d2253f04d",
- "delayed/durhyb": "9956527df703396df6352897c681502ff686153bb9f0ad06ce61731fc69de78e",
- "delayed/profhyb": "a66e43eef4f630e2ef5c120c820059c4e6f25fb56da30b6a33af00e9e79fcd50",
- "delayed/durtri": "2db07d70918af05ad829cf4281337092f1025f271020604bad5c3c47a455b205",
- "delayed/proftri": "521e27873e1bddb7d0f6fc520cb68272c433c5331c95626382d978d41b768799",
- "delayed/durevt": "2a84c28f9607b5f6d3c20fd6351a610e337a9c0de4c14ecb254e2a6323148581",
- "delayed/stkevt": "fca3113de8cf6aec42eac460490a44b3e8f4aa25ccbb102e9a8c83113d526549",
- "delayed/durref": "f97333f55bf13ac2cb1da76b5c0085a381a37e69331efc6d15e52ff531a32526",
- "delayed/stkref": "cb6a901566a5552a0aa62862f591e170ea7a282793617f346b07b8f3d976ed50"
- },
- "models": {
- "bc_rl2_stattempo_slowpace366_arcdur_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- },
- "bc_rl2_stattempo_slowpace366_arcvel_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- },
- "bc_rl2_stattempo_slowpace366_arcdur_abstime_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- },
- "bc_rl2_stattempo_slowpace366_arcvel_nosmooth_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- },
- "bc_rl2_stattempo_slowpace366_arc_projection_dur_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- },
- "bc_rl2_stattempo_slowpace366_arc_projection_vel_m200_d033": {
- "parameters": 180424144,
- "state_schema_sha256": "3bd8f28767616d8acb1357e0ca544a42418c5727411d272842f2754caad38694"
- }
- }
-}
diff --git a/tests/libero_ddp_resume_worker.py b/tests/libero_ddp_resume_worker.py
deleted file mode 100644
index 278453ecab..0000000000
--- a/tests/libero_ddp_resume_worker.py
+++ /dev/null
@@ -1,64 +0,0 @@
-"""Subprocess fixture for a real one-process to two-process Lightning resume."""
-
-import argparse
-import json
-from pathlib import Path
-
-import torch
-from omegaconf import OmegaConf, open_dict
-
-from egomimic.trainHydra import train
-from tests.test_libero_benchmark import make_replay
-from tests.test_oat_training import config_for
-
-
-def main():
- parser = argparse.ArgumentParser()
- parser.add_argument("--root", type=Path, required=True)
- parser.add_argument("--resume", action="store_true")
- args = parser.parse_args()
- torch.set_num_threads(1)
- args.root.mkdir(parents=True, exist_ok=True)
- replay = args.root / "replay.zarr"
- if not args.resume:
- make_replay(replay)
- cfg = config_for("libero_oattok", replay, args.root / "train")
- spec = cfg.model.pipeline.stages[0].tokenizer
- spec.emb_dim, spec.head_dim = 32, 8
- spec.encoder_depth, spec.decoder_depth, spec.num_registers = 1, 1, 4
- cfg.trainer.limit_train_batches = 1.0
- cfg.trainer.accumulate_grad_batches = 2 if args.resume else 4
- cfg.trainer.devices = 2 if args.resume else 1
- cfg.trainer.max_epochs = 2 if args.resume else 1
- cfg.callbacks.batch_budget = OmegaConf.create(
- {
- "_target_": "egomimic.pl_utils.oat_training.OATBatchBudgetCallback",
- "global_batch_size": 8,
- }
- )
- with open_dict(cfg.trainer):
- cfg.trainer.enable_progress_bar = False
- if args.resume:
- cfg.trainer.strategy = "ddp_find_unused_parameters_true"
- if args.resume:
- cfg.ckpt_path = str(args.root / "train/checkpoints/last.ckpt")
- _, objects = train(cfg)
- if objects["trainer"].is_global_zero:
- payload = torch.load(
- args.root / "train/checkpoints/last.ckpt",
- map_location="cpu",
- weights_only=False,
- )
- report = {
- "global_step": payload["global_step"],
- "ema_num_updates": payload["ema_num_updates"],
- "training_budget": payload["training_budget"],
- "optimizer_states": bool(payload["optimizer_states"]),
- }
- (args.root / ("resumed.json" if args.resume else "initial.json")).write_text(
- json.dumps(report)
- )
-
-
-if __name__ == "__main__":
- main()
diff --git a/tests/test_action_flow_affine_configs.py b/tests/test_action_flow_affine_configs.py
deleted file mode 100644
index 5c192ab8b2..0000000000
--- a/tests/test_action_flow_affine_configs.py
+++ /dev/null
@@ -1,138 +0,0 @@
-import hashlib
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-
-def test_builder_emits_frozen_eleven_variant_three_seed_matrix(tmp_path):
- source = Path(__file__).parents[1]
- output = tmp_path / "config"
- experiment = tmp_path / "experiment"
- training_dataset = tmp_path / "train.npz"
- evaluation_dataset = tmp_path / "eval.npz"
- training_dataset.write_bytes(b"frozen-training-data")
- evaluation_dataset.write_bytes(b"frozen-evaluation-data")
- subprocess.run(
- [
- sys.executable,
- str(source / "scripts/experiments/build_action_flow_affine_configs.py"),
- "--output-dir",
- str(output),
- "--experiment-root",
- str(experiment),
- "--training-dataset",
- str(training_dataset),
- "--evaluation-dataset",
- str(evaluation_dataset),
- ],
- check=True,
- )
- configs = sorted(output.glob("gaussian-torus-*.json"))
- assert len(configs) == 33
- resolved = [json.loads(path.read_text()) for path in configs]
- assert {config["seed"] for config in resolved} == {42, 43, 44}
- assert {config["checkpoint_every"] for config in resolved} == {50_000}
- assert {config["evaluation_particles"] for config in resolved} == {2048}
- assert {config["evaluation_dataset"] for config in resolved} == {
- str(evaluation_dataset)
- }
- action_configs = [
- config for config in resolved if config["architecture"] == "action_adapter_flow"
- ]
- assert len(action_configs) == 30
- assert {config["model"]["latent_dim"] for config in action_configs} == {8}
- manifest = json.loads((output / "manifest.json").read_text())
- assert manifest["primary_metric"] == "validation_generation_symmetric_nn_mse"
- assert manifest["pass_threshold"] is None
- assert len(manifest["variants"]) == 11
- assert manifest["datasets"]["training"] == {
- "path": str(training_dataset.resolve()),
- "sha256": hashlib.sha256(training_dataset.read_bytes()).hexdigest(),
- }
- assert manifest["datasets"]["evaluation"] == {
- "path": str(evaluation_dataset.resolve()),
- "sha256": hashlib.sha256(evaluation_dataset.read_bytes()).hexdigest(),
- }
-
- for index, config in enumerate(resolved):
- run = Path(config["output_dir"])
- run.mkdir(parents=True)
- (run / "summary.json").write_text(
- json.dumps(
- {
- "validation_generation_symmetric_nn_mse": float(index + 1),
- "validation_generation_energy_distance": float(index + 2),
- }
- )
- )
- summary_json = tmp_path / "summary.json"
- summary_markdown = tmp_path / "summary.md"
- subprocess.run(
- [
- sys.executable,
- str(source / "scripts/experiments/summarize_action_flow_affine.py"),
- "--manifest",
- str(output / "manifest.json"),
- "--output-json",
- str(summary_json),
- "--output-markdown",
- str(summary_markdown),
- ],
- check=True,
- )
- aggregate = json.loads(summary_json.read_text())
- assert len(aggregate["rows"]) == 11
- assert "mean" in aggregate["rows"][0][manifest["primary_metric"]]
- assert "Symmetric NN MSE" in summary_markdown.read_text()
-
- e_configs = [config for config in resolved if config["variant"].startswith("e-")]
- g_configs = [config for config in resolved if config["variant"].startswith("g-")]
- assert {config["lambda_reconstruction"] for config in e_configs} == {
- 1.0,
- 10.0,
- 100.0,
- }
- assert {config["lambda_reconstruction"] for config in g_configs} == {
- 1.0,
- 10.0,
- 100.0,
- }
- assert {config["lambda_action_velocity"] for config in e_configs} == {0.0}
- assert {config["lambda_action_velocity"] for config in g_configs} == {1.0}
-
-
-def test_builder_run_suffix_produces_distinct_retry_identities(tmp_path):
- source = Path(__file__).parents[1]
- training_dataset = tmp_path / "train.npz"
- evaluation_dataset = tmp_path / "eval.npz"
- training_dataset.write_bytes(b"train")
- evaluation_dataset.write_bytes(b"eval")
- output = tmp_path / "config"
- subprocess.run(
- [
- sys.executable,
- str(source / "scripts/experiments/build_action_flow_affine_configs.py"),
- "--output-dir",
- str(output),
- "--experiment-root",
- str(tmp_path / "experiment"),
- "--training-dataset",
- str(training_dataset),
- "--evaluation-dataset",
- str(evaluation_dataset),
- "--seeds",
- "42",
- "--run-suffix",
- "smoke-retry1",
- ],
- check=True,
- )
- config = json.loads(
- (
- output / "gaussian-torus-f-nonlinear-path-seed42-smoke-retry1.json"
- ).read_text()
- )
- assert config["wandb"]["id"].endswith("-smoke-retry1")
- assert config["wandb"]["dir"] == str(tmp_path / "experiment" / "wandb")
- assert config["output_dir"].endswith("-smoke-retry1")
diff --git a/tests/test_action_flow_candidate_codec.py b/tests/test_action_flow_candidate_codec.py
deleted file mode 100644
index ef6a7a8e7d..0000000000
--- a/tests/test_action_flow_candidate_codec.py
+++ /dev/null
@@ -1,151 +0,0 @@
-"""Exact-section and JVP behavior of the restricted graph diagnostic."""
-
-import pytest
-import torch
-import torch.nn as nn
-from torch.func import jvp
-
-from egomimic.models.action_flow_codec import GraphSectionSequenceCodec
-from egomimic.pipeline.core import Pipeline
-from egomimic.pipeline.stages_action_flow import (
- ActionFlowObjectiveStage,
- ConditionalVelocityStage,
- ContentDecoderStage,
- ContentEncoderStage,
- LatentBridgeStage,
-)
-
-
-class _EncoderView(nn.Module):
- def __init__(self, codec):
- super().__init__()
- self.graph = codec.graph
-
- def forward(self, content):
- return torch.cat((content, self.graph(content)), dim=-1)
-
-
-class _Field(nn.Module):
- def __init__(self):
- super().__init__()
- self.linear = nn.Linear(8, 8)
- self.condition = nn.Linear(3, 8)
-
- def forward(self, value, time, condition, *, condition_drop_mask):
- condition = condition.masked_fill(condition_drop_mask[:, None], 0)
- return self.linear(value).tanh() + self.condition(condition)[:, None]
-
-
-def _codec():
- return GraphSectionSequenceCodec(action_dim=4, latent_dim=8, horizon=4)
-
-
-def test_graph_section_identity_survives_joint_optimizer_updates_and_shared_ownership():
- torch.manual_seed(46)
- codec = _codec().double()
- encoder = ContentEncoderStage(_EncoderView(codec))
- decoder = ContentDecoderStage(codec)
- field = ConditionalVelocityStage(_Field().double())
- objective = ActionFlowObjectiveStage(
- reconstruction_weight=0.0, residual_key="action_flow/fm_velocity_residual"
- )
- pipeline = Pipeline(
- [encoder, LatentBridgeStage(samples_per_content=2), field, decoder, objective]
- )
- parameters = list(pipeline.parameters())
- assert len({id(parameter) for parameter in parameters}) == len(parameters)
- assert encoder.encoder.graph is decoder.decoder.graph
- optimizer = torch.optim.AdamW(parameters, lr=1e-3)
- assert len(optimizer.param_groups[0]["params"]) == len(parameters)
- snapshots = {
- name: [parameter.detach().clone() for parameter in module.parameters()]
- for name, module in (("graph", codec.graph), ("residual", codec.residual))
- }
- target = torch.randn(2, 4, 4, dtype=torch.float64)
- for _ in range(2):
- latent = codec.encode(target)
- torch.testing.assert_close(latent[..., :4], target, rtol=0, atol=0)
- torch.testing.assert_close(codec(latent), target, rtol=0, atol=1e-12)
- batch = pipeline(
- {
- "target": target,
- "sampler/noise": torch.randn(2, 4, 8, dtype=torch.float64),
- "condition": torch.randn(2, 3, dtype=torch.float64),
- }
- )
- torch.testing.assert_close(
- batch["loss/action_flow"],
- batch["log/action_flow_fm"] + batch["log/action_flow_action_velocity"],
- rtol=0,
- atol=0,
- )
- optimizer.zero_grad()
- batch["loss/action_flow"].backward()
- for module in (codec.graph, codec.residual, field.field):
- gradients = [p.grad for p in module.parameters() if p.grad is not None]
- assert gradients and all(torch.isfinite(g).all() for g in gradients)
- assert sum(g.abs().sum() for g in gradients) > 0
- optimizer.step()
- torch.testing.assert_close(codec(codec.encode(target)), target, rtol=0, atol=1e-12)
- for name, module in (("graph", codec.graph), ("residual", codec.residual)):
- assert any(
- not torch.equal(old, new)
- for old, new in zip(snapshots[name], module.parameters())
- )
-
-
-def test_graph_section_jvp_matches_finite_difference_with_trainable_graph():
- torch.manual_seed(7)
- codec = _codec().double()
- latent = torch.randn(2, 4, 8, dtype=torch.float64, requires_grad=True)
- tangent = torch.randn_like(latent, requires_grad=True)
- _, predicted = jvp(codec, (latent,), (tangent,))
- delta = 1e-6
- expected = (codec(latent + delta * tangent) - codec(latent - delta * tangent)) / (
- 2 * delta
- )
- torch.testing.assert_close(predicted, expected, rtol=2e-5, atol=2e-6)
- gradients = torch.autograd.grad(
- predicted.square().mean(), (latent, tangent, *codec.graph.parameters())
- )
- assert all(torch.isfinite(gradient).all() for gradient in gradients)
- assert sum(gradient.abs().sum() for gradient in gradients[2:]) > 0
-
-
-def test_graph_section_context_free_shape_and_small_private_networks():
- codec = GraphSectionSequenceCodec(action_dim=4, latent_dim=8, horizon=16)
- for module in (codec.graph, codec.residual):
- assert module.depth == 2
- assert sum(parameter.numel() for parameter in module.parameters()) < 12_000
- with pytest.raises(ValueError, match="expected sequence shape"):
- codec(torch.randn(2, 15, 8))
- with pytest.raises(ValueError, match="expected sequence shape"):
- codec.encode(torch.randn(2, 16, 3))
-
-
-def test_graph_section_mixed_precision_identity_and_jvp_backward_are_finite():
- torch.manual_seed(19)
- codec = _codec()
- content = torch.randn(2, 4, 4)
- latent = torch.randn(2, 4, 8, requires_grad=True)
- tangent = torch.randn_like(latent, requires_grad=True)
- with torch.autocast("cpu", dtype=torch.bfloat16):
- reconstructed = codec(codec.encode(content))
- _, decoded_tangent = jvp(codec, (latent,), (tangent,))
- loss = decoded_tangent.square().mean()
- torch.testing.assert_close(reconstructed, content, rtol=0, atol=0)
- assert torch.isfinite(decoded_tangent).all()
- loss.backward()
- for module in (codec.graph, codec.residual):
- gradients = [p.grad for p in module.parameters() if p.grad is not None]
- assert gradients and all(torch.isfinite(g).all() for g in gradients)
- assert sum(g.abs().sum() for g in gradients) > 0
-
-
-@pytest.mark.parametrize(
- "kwargs", [{"action_dim": 8}, {"latent_dim": 3}, {"dropout": 0.1}]
-)
-def test_graph_section_rejects_invalid_or_stochastic_section(kwargs):
- options = {"action_dim": 4, "latent_dim": 8, "horizon": 4, **kwargs}
- with pytest.raises(ValueError):
- GraphSectionSequenceCodec(**options)
diff --git a/tests/test_action_flow_candidate_configs.py b/tests/test_action_flow_candidate_configs.py
deleted file mode 100644
index 72d6231066..0000000000
--- a/tests/test_action_flow_candidate_configs.py
+++ /dev/null
@@ -1,181 +0,0 @@
-"""Exercise the maintained Hydra rows and their candidate-specific telemetry."""
-
-from pathlib import Path
-
-import hydra
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-from omegaconf import OmegaConf
-
-from egomimic.eval.action_flow_diagnostics import ActionFlowDiagnostics
-from egomimic.pl_utils.pl_model_action_flow import ActionFlowModelWrapper
-
-CONFIG_DIR = Path(__file__).parents[1] / "egomimic/hydra_configs"
-ROWS = [
- ("latent_fm_sg_recon1", "latent_fm_stopgrad", 1.0),
- ("graph_section", "graph_section_diagnostic", 0.0),
-]
-
-
-def _compose(row):
- with initialize_config_dir(version_base=None, config_dir=str(CONFIG_DIR.resolve())):
- return compose(
- config_name="train_zarr_cartesian",
- overrides=[
- f"+experiment=sim/pusht/action_flow_bc_usocket_{row}_s42",
- "++paths.root_dir=.",
- ],
- )
-
-
-@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS)
-def test_candidate_row_instantiates_exact_dimensions_and_shared_codec(
- row, method, reconstruction_weight
-):
- cfg = _compose(row)
- assert cfg.model.action_flow_method == cfg.train.action_flow_method == method
- assert cfg.run_provenance.objective.method == method
- assert cfg.model.action_horizon == 16
- assert cfg.model.latent_dim == 8
- assert cfg.model.action_dim == 4
- assert cfg.model.reconstruction_weight == reconstruction_weight
- assert cfg.run_provenance.objective.decoded_noise_scale_weight == 0
- assert cfg.model.flow_samples_per_content == 14
- assert cfg.model.condition_dropout_probability == 0.3
- assert cfg.model.optimizer.lr == 3e-5
-
- algo = hydra.utils.instantiate(cfg.model.pipeline, device="cpu")
- stages = algo.pipeline.stages
- encoder, field, decoder, objective = stages[3], stages[5], stages[6], stages[7]
- assert field.field.horizon == 16
- assert field.field.input_dim == field.field.output_dim == 8
- assert field.field.condition_dim == 67
- assert field.field.depth == 12
- assert 39_000_000 <= sum(p.numel() for p in field.parameters()) <= 41_000_000
- assert objective.reconstruction_weight == reconstruction_weight
- assert stages[1].num_tokens == 16 and stages[1].latent_dim == 8
- parameters = list(algo.nets.parameters())
- assert len(parameters) == len({id(parameter) for parameter in parameters})
- if method == "latent_fm_stopgrad":
- assert field.flow_clean_gradient_mode == "all_stopgrad"
- assert objective.residual_key == field.flow_residual_key
- assert decoder.residual_key == field.residual_key
- assert field.flow_residual_key != field.residual_key
- assert encoder.encoder.depth == decoder.decoder.depth == 2
- else:
- assert field.flow_clean_gradient_mode == "full"
- assert encoder.encoder.graph is decoder.decoder.graph
- encoder_parameters = {id(parameter) for parameter in encoder.parameters()}
- residual_parameters = {
- id(parameter) for parameter in decoder.decoder.residual.parameters()
- }
- assert encoder_parameters.isdisjoint(residual_parameters)
- assert decoder.decoder.latent_dim - decoder.decoder.action_dim == 4
- assert not cfg.run_provenance.objective.reconstruction_is_optimizer_objective
- assert cfg.run_provenance.requirements.r6 == "FAIL_restricted_diagnostic"
- assert not cfg.evaluator.action_flow_diagnostics.capture_activations
- target = torch.randn(2, 16, 4)
- torch.testing.assert_close(
- decoder.decoder(encoder.encoder(target)), target, rtol=0, atol=1e-6
- )
-
-
-@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS)
-def test_candidate_wrapper_gradient_routes_and_compute_contract(
- monkeypatch, row, method, reconstruction_weight
-):
- cfg = _compose(row)
- # Preserve the real Hydra codec/factory/stage wiring and all sequence
- # dimensions; replace only the expensive shared field width/depth. Inputs
- # are already encoded observations, so image/data nodes are not exercised.
- pipeline_cfg = OmegaConf.create(
- OmegaConf.to_container(cfg.model.pipeline, resolve=True)
- )
- pipeline_cfg.stages = list(pipeline_cfg.stages)[3:]
- field_config = pipeline_cfg.stages[2].field
- field_config.hidden_dim = 32
- field_config.depth = 2
- field_config.num_heads = 4
- field_config.feedforward_dim = 64
- field_config.time_embedding_dim = 32
- algo = hydra.utils.instantiate(pipeline_cfg, device="cpu")
- wrapper = ActionFlowModelWrapper(pipeline=algo, gradient_telemetry_cadence=1)
- logged = {}
- monkeypatch.setattr(
- wrapper, "log", lambda key, value, **kwargs: logged.update({key: value})
- )
- torch.manual_seed(51)
- loss = wrapper.training_step(
- {
- "fixture": {
- "target": torch.randn(2, 16, 4),
- "condition": torch.randn(2, 67),
- "sampler/noise": torch.randn(2, 16, 8),
- }
- },
- 0,
- )
- assert torch.isfinite(loss)
- loss.backward()
- assert wrapper._gradient_route_manifest is not None
- routes = wrapper._gradient_route_manifest["routes"]
- intersections = wrapper._gradient_route_manifest["intersections"]
- assert routes["FM"] and routes["ActionVelocity"]
- assert intersections["FM__ActionVelocity"]
- assert all(
- torch.isfinite(p.grad).all() for p in wrapper.parameters() if p.grad is not None
- )
- calls = 2 if method == "latent_fm_stopgrad" else 1
- assert logged["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] == calls
- assert (
- logged["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] == calls * 14
- )
- if method == "latent_fm_stopgrad":
- assert intersections["FM__Reconstruction"] == []
- assert logged["Train/ActionFlow/GradientCosineDefined/FM__Reconstruction"] == 0
- assert (
- logged[
- "Train/ActionFlow/GradientIntersectionParameterCount/FM__Reconstruction"
- ]
- == 0
- )
- assert routes["Reconstruction"]
- else:
- assert "Reconstruction" not in routes
- assert "Train/ActionFlow/GradientNorm/Reconstruction" not in logged
- torch.testing.assert_close(
- loss,
- logged["Train/ActionFlow/FlowMatchingLoss"]
- + logged["Train/ActionFlow/ActionVelocityLoss"],
- )
-
-
-def test_graph_factory_rejects_independent_encoder_configuration():
- cfg = _compose("graph_section")
- pipeline_cfg = OmegaConf.create(
- OmegaConf.to_container(cfg.model.pipeline, resolve=True)
- )
- pipeline_cfg.stages = [pipeline_cfg.stages[3], pipeline_cfg.stages[6]]
- pipeline_cfg.stages[0].encoder = {"_target_": "torch.nn.Identity"}
- with pytest.raises(hydra.errors.InstantiationException, match="second codec"):
- hydra.utils.instantiate(pipeline_cfg, device="cpu")
-
-
-@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS)
-def test_candidate_diagnostic_config_constructs_without_stale_activation_settings(
- tmp_path, row, method, reconstruction_weight
-):
- cfg = _compose(row)
- diagnostics = cfg.evaluator.action_flow_diagnostics
- diagnostics.noise_seed_bank_path = str(
- CONFIG_DIR / "evaluator/energy_score_seed_bank_k32_v1.json"
- )
- diagnostics.artifact_root = str(tmp_path / "diagnostics")
- config = OmegaConf.to_container(diagnostics, resolve=True)
- instance = ActionFlowDiagnostics(config)
- assert instance.jacobian_samples == 2
- assert instance.capture_activations == (method == "latent_fm_stopgrad")
- if method == "graph_section_diagnostic":
- assert instance.activation_layer_map == ()
- assert instance.cknna_k == 0
diff --git a/tests/test_action_flow_candidate_gates.py b/tests/test_action_flow_candidate_gates.py
deleted file mode 100644
index 915d02a014..0000000000
--- a/tests/test_action_flow_candidate_gates.py
+++ /dev/null
@@ -1,295 +0,0 @@
-"""Method-specific gates: real configs plus tiny, honest telemetry receipts."""
-
-import hashlib
-import re
-import subprocess
-
-import pytest
-import torch
-from hydra.core.hydra_config import HydraConfig
-from omegaconf import OmegaConf, open_dict
-
-from tests.test_verify_action_flow_training_smoke import HEAD, MODULE, _history_row
-from tools.validate_action_flow_config import (
- CHECKPOINT_METHOD_CONTRACTS,
- DEFAULT_CONFIG_ROOT,
- GRAPH_METHOD,
- LIKELIHOOD_METHOD,
- STOPGRAD_METHOD,
- STOPGRAD_UNITE_METHOD,
- PreflightError,
- compose_experiment,
- validate_experiment,
- validate_method_contract,
-)
-
-# Historical method identities are not selectable recipes. Exercise every current
-# recipe with the full gates, and require explicit rejection for removed ones.
-CURRENT_RECIPE_METHODS = {
- recipe: method
- for recipe, method in CHECKPOINT_METHOD_CONTRACTS.items()
- if (DEFAULT_CONFIG_ROOT / "experiment" / f"{recipe}.yaml").is_file()
-}
-RETIRED_METHOD_RECIPES = set(CHECKPOINT_METHOD_CONTRACTS) - set(CURRENT_RECIPE_METHODS)
-
-
-@pytest.mark.parametrize("experiment", sorted(RETIRED_METHOD_RECIPES))
-def test_retired_method_identity_cannot_select_a_training_recipe(experiment):
- with pytest.raises(PreflightError, match="absent or retired"):
- compose_experiment(experiment)
-
-
-@pytest.fixture(autouse=True)
-def offline_campaign_roots(monkeypatch, tmp_path):
- # Config composition only: no dataset reads or launch receipts are fabricated.
- monkeypatch.setenv("PUSHSHAPES_USOCKET_ROOT", str(tmp_path / "usocket"))
- monkeypatch.setenv("PUSHSHAPES_CHAIN_GRIPPER_ROOT", str(tmp_path / "chain"))
-
-
-@pytest.mark.parametrize("experiment", CURRENT_RECIPE_METHODS)
-def test_real_candidate_config_and_parameter_manifest(experiment):
- if (
- "chain_points6_latent_fm_sg_unite_h512" in experiment
- or "cotrain_uc_latent_fm_sg_unite_h512" in experiment
- ):
- # Campaign-owned manifests are deliberately absent from the shipped YAML;
- # this unit test must verify the gate, not invent production receipts.
- with pytest.raises(PreflightError, match="manifest.*missing"):
- validate_experiment(experiment)
- return
- report, _ = validate_experiment(experiment)
- assert report["status"] == "PASS"
- assert report["action_flow_method"] == CURRENT_RECIPE_METHODS[experiment]
- assert report["topology"]["shared_field_instance"]
- topology = report["topology"]
- if "private_codec_routes" in topology:
- assert set(topology["private_codec_routes"]) == set(
- report["dimensions"]["actions"]
- )
- assert "shared_decoder_instance" not in topology
- else:
- assert topology["shared_decoder_instance"]
- expected_stages = (
- 9 if CURRENT_RECIPE_METHODS[experiment] == STOPGRAD_UNITE_METHOD else 8
- )
- assert len(topology["train_order"]) == expected_stages
- assert (
- sum(
- v["total"] for k, v in report["parameters"].items() if k != "pipeline_total"
- )
- == report["parameters"]["pipeline_total"]["total"]
- )
-
-
-@pytest.mark.parametrize("experiment", CURRENT_RECIPE_METHODS)
-def test_candidate_two_optimizer_update_config_gate(experiment, tmp_path, monkeypatch):
- if experiment not in MODULE.APPROVED_EXPERIMENTS:
- # A shipped candidate config is not automatically an approved smoke recipe.
- with pytest.raises(
- MODULE.SmokeVerificationError, match="unapproved experiment"
- ):
- MODULE._validate_config(
- config_path=tmp_path / "unused.yaml",
- experiment=experiment,
- run_dir=tmp_path,
- expected_head=HEAD,
- expected_config_sha256=None,
- expected_split_sha256=None,
- expected_normalization_sha256=None,
- )
- return
- cfg = compose_experiment(experiment)
- normalization = tmp_path / "norm_stats.json"
- normalization.write_text("{}")
- norm_hash = hashlib.sha256(normalization.read_bytes()).hexdigest()
- with open_dict(cfg):
- cfg.trainer.max_steps = 2
- cfg.trainer.val_check_interval = (
- 2 if CURRENT_RECIPE_METHODS[experiment] == STOPGRAD_UNITE_METHOD else 1
- )
- cfg.trainer.limit_val_batches = 1
- cfg.trainer.precision = "bf16"
- cfg.callbacks.model_checkpoint.every_n_train_steps = 1
- cfg.model.gradient_telemetry_cadence = 2
- cfg.norm_stats.precomputed_norm_path = str(normalization)
- cfg.run_provenance.source_commit = HEAD
- cfg.run_provenance.normalization_sha256 = norm_hash
- cfg.evaluator.artifact_root = str(tmp_path / "energy")
- if CURRENT_RECIPE_METHODS[experiment] != LIKELIHOOD_METHOD:
- cfg.evaluator.action_flow_diagnostics.artifact_root = str(
- tmp_path / "diagnostics"
- )
- cfg.hydra.runtime.cwd = str(MODULE.REPOSITORY_ROOT)
- cfg.hydra.runtime.output_dir = str(tmp_path)
- HydraConfig.instance().set_config(cfg)
- selected = OmegaConf.masked_copy(cfg, [key for key in cfg if key != "hydra"])
- path = tmp_path / "config.yaml"
- OmegaConf.save(
- OmegaConf.create(OmegaConf.to_container(selected, resolve=True)), path
- )
- monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD)
- MODULE._validate_config(
- config_path=path,
- experiment=experiment,
- run_dir=tmp_path,
- expected_head=HEAD,
- expected_config_sha256=None,
- expected_split_sha256=None,
- expected_normalization_sha256=norm_hash,
- )
-
-
-@pytest.mark.parametrize(
- "experiment,path,value",
- [
- (
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42",
- "model.pipeline.stages.5.flow_clean_gradient_mode",
- "attached",
- ),
- (
- "sim/pusht/action_flow_bc_usocket_graph_section_s42",
- "model.reconstruction_weight",
- 1.0,
- ),
- (
- "sim/pusht/action_flow_bc_usocket_bridge_likelihood_s42",
- "run_provenance.inference.sampler",
- "reverse_euler",
- ),
- (
- "sim/pusht/action_flow_bc_usocket_bridge_likelihood_s42",
- "model.pipeline.stages.7.tau",
- 0.01,
- ),
- ],
-)
-def test_method_contract_rejects_scientific_identity_drift(experiment, path, value):
- cfg = compose_experiment(experiment)
- OmegaConf.update(cfg, path, value)
- with pytest.raises(PreflightError):
- validate_method_contract(cfg, experiment)
-
-
-def test_stopgrad_telemetry_requires_empty_fm_reconstruction_intersection():
- row = _history_row()
- for key in list(row):
- if key.endswith("/FM__Reconstruction"):
- row[key] = 0.0
- row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 2.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 28.0
- MODULE._validate_history({1: row}, method=STOPGRAD_METHOD)
- row["Train/ActionFlow/GradientIntersectionParameterCount/FM__Reconstruction"] = 1.0
- with pytest.raises(MODULE.SmokeVerificationError, match="reachability"):
- MODULE._validate_history({1: row}, method=STOPGRAD_METHOD)
-
-
-def test_graph_telemetry_does_not_invent_reconstruction_gradients():
- row = {
- key: value
- for key, value in _history_row().items()
- if not ("/Gradient" in key and "Reconstruction" in key)
- and "/Alignment/CK" not in key
- }
- row["Train/ActionFlow/TotalLoss_step"] = 2.5
- row[f"Train/ActionFlow/TotalLoss/{MODULE.SOURCE_LABEL}_step"] = 2.5
- row["Valid/ActionFlow/TotalLoss"] = 3.0
- MODULE._validate_history({1: row}, method=GRAPH_METHOD, reconstruction_weight=0.0)
-
-
-def test_likelihood_history_requires_real_nll_and_common_validation_only():
- row = {
- key: value
- for key, value in _history_row().items()
- if "/ActionFlow/" not in key or "/Compute/" in key
- }
- for component in ("InteriorBridgeNLL", "BoundaryNLL", "TotalLoss"):
- value = 2.5 if component == "TotalLoss" else 1.25
- row[f"Train/ActionFlow/{component}_step"] = value
- row[f"Train/ActionFlow/{component}/{MODULE.SOURCE_LABEL}_step"] = value
- row[f"Valid/ActionFlow/{component}"] = value
- for label in ("InteriorBridgeNLL", "BoundaryNLL"):
- row[f"Train/ActionFlow/GradientNorm/{label}"] = 0.5
- row[f"Train/ActionFlow/GradientParameterCount/{label}"] = 100.0
- pair = "InteriorBridgeNLL__BoundaryNLL"
- row[f"Train/ActionFlow/GradientCosine/{pair}"] = 0.25
- row[f"Train/ActionFlow/GradientCosineDefined/{pair}"] = 1.0
- row[f"Train/ActionFlow/GradientIntersectionParameterCount/{pair}"] = 50.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 15.0
- row["Train/ActionFlow/Compute/DecoderJVPCallsPerStep"] = 0.0
- MODULE._validate_history({1: row}, method=LIKELIHOOD_METHOD)
- del row["Valid/EnergyScore@32"]
- with pytest.raises(MODULE.SmokeVerificationError):
- MODULE._validate_history({1: row}, method=LIKELIHOOD_METHOD)
-
-
-@pytest.mark.parametrize("method", [STOPGRAD_METHOD, GRAPH_METHOD, LIKELIHOOD_METHOD])
-def test_checkpoint_gradient_receipt_uses_exact_method_routes(method):
- named = [
- (f"nets.pipeline.stages.{stage}.weight", torch.nn.Parameter(torch.ones(2)))
- for stage in (0, 3, 5, 6)
- ]
- groups = {
- "FM": [0, 2] if method == STOPGRAD_METHOD else [0, 1, 2],
- "Reconstruction": [1, 3],
- "ActionVelocity": [0, 1, 2, 3],
- }
- if method == GRAPH_METHOD:
- del groups["Reconstruction"]
- elif method == LIKELIHOOD_METHOD:
- groups = {"InteriorBridgeNLL": [0, 1, 2], "BoundaryNLL": [0, 1, 2, 3]}
- routes = {
- label: [
- {"name": named[i][0], "shape": [2], "numel": 2, "dtype": "torch.float32"}
- for i in indices
- ]
- for label, indices in groups.items()
- }
- labels = list(groups)
- core = {
- "schema_version": 1,
- "routes": routes,
- "route_sha256": {
- k: MODULE._canonical_json_sha256(v) for k, v in routes.items()
- },
- "intersections": {
- f"{left}__{right}": [
- named[i][0] for i in groups[left] if i in groups[right]
- ]
- for index, left in enumerate(labels)
- for right in labels[index + 1 :]
- },
- }
- manifest = {**core, "manifest_sha256": MODULE._canonical_json_sha256(core)}
- MODULE._validate_gradient_route_manifest(manifest, named, method=method)
-
-
-@pytest.mark.parametrize(
- "requested,selected,accepted",
- [
- ("gpu-h100,gpu-h200", "gpu-h100", True),
- ("gpu-h100,gpu-h200", "gpu-h200", True),
- ("ice-gpu", "ice-gpu", True),
- ("gpu-h100,gpu-h200", "gpu", False),
- ("gpu-h100", "gpu-h200", False),
- ("gpu-h100,", "gpu-h100", False),
- ],
-)
-def test_portable_launcher_accepts_only_explicit_partition_members(
- requested, selected, accepted
-):
- source = (
- MODULE.REPOSITORY_ROOT / "scripts/train/launch_action_flow_usocket.sbatch"
- ).read_text()
- function = re.search(r"partition_allowed\(\) \{.*?\n\}", source, re.S).group()
- result = subprocess.run(
- [
- "bash",
- "-c",
- function + '\npartition_allowed "$1" "$2"',
- "test",
- requested,
- selected,
- ]
- )
- assert (result.returncode == 0) is accepted
diff --git a/tests/test_action_flow_candidate_stages.py b/tests/test_action_flow_candidate_stages.py
deleted file mode 100644
index 7bf20cf6d9..0000000000
--- a/tests/test_action_flow_candidate_stages.py
+++ /dev/null
@@ -1,246 +0,0 @@
-"""Gradient routes for the separately approved latent-FM-only SG candidate."""
-
-import copy
-
-import pytest
-import torch
-import torch.nn as nn
-
-from egomimic.pipeline.stages_action_flow import (
- ActionFlowObjectiveStage,
- ConditionalVelocityStage,
- ContentDecoderStage,
- ContentEncoderStage,
- LatentBridgeStage,
-)
-
-
-class _Field(nn.Module):
- def __init__(self):
- super().__init__()
- self.weight = nn.Parameter(torch.tensor(0.4, dtype=torch.float64))
- self.condition = nn.Linear(3, 8).double()
- self.calls = []
-
- def forward(self, state, time, condition, *, condition_drop_mask):
- self.calls.append((state, time, condition, condition_drop_mask))
- masked = condition.masked_fill(condition_drop_mask[:, None], 0)
- return (
- self.weight * state.tanh()
- + self.condition(masked)[:, None]
- + time[:, None, None]
- )
-
-
-def _forward(
- mode, encoder, decoder, field, target, noise, condition, execution="separate"
-):
- stages = [
- ContentEncoderStage(encoder),
- LatentBridgeStage(samples_per_content=3, condition_dropout_probability=0.3),
- ConditionalVelocityStage(
- field, flow_clean_gradient_mode=mode, fm_field_execution=execution
- ),
- ContentDecoderStage(decoder),
- ActionFlowObjectiveStage(residual_key="action_flow/fm_velocity_residual"),
- ]
- batch = {"target": target, "sampler/noise": noise, "condition": condition}
- torch.manual_seed(72)
- for stage in stages:
- batch = stage(batch)
- return batch
-
-
-def _models_and_inputs():
- torch.manual_seed(41)
- return (
- nn.Sequential(nn.Linear(4, 8), nn.Tanh()).double(),
- nn.Sequential(nn.Linear(8, 6), nn.SiLU(), nn.Linear(6, 4)).double(),
- _Field(),
- torch.randn(3, 4, 4, dtype=torch.float64),
- torch.randn(3, 4, 8, dtype=torch.float64),
- torch.randn(3, 3, dtype=torch.float64, requires_grad=True),
- )
-
-
-def test_latent_only_stopgrad_blocks_both_clean_fm_routes_not_shared_condition():
- encoder, decoder, field, target, noise, condition = _models_and_inputs()
- batch = _forward("all_stopgrad", encoder, decoder, field, target, noise, condition)
- fm = batch["log/action_flow_fm"]
- gradients = torch.autograd.grad(
- fm, tuple(encoder.parameters()), allow_unused=True, retain_graph=True
- )
- assert all(
- gradient is None or torch.count_nonzero(gradient) == 0 for gradient in gradients
- )
- field_gradients = torch.autograd.grad(
- fm, tuple(field.parameters()), retain_graph=True
- )
- assert all(torch.isfinite(gradient).all() for gradient in field_gradients)
- assert sum(gradient.abs().sum() for gradient in field_gradients) > 0
- condition_gradient = torch.autograd.grad(fm, condition)[0]
- assert torch.isfinite(condition_gradient).all()
- assert condition_gradient.abs().sum() > 0
-
- attached, detached = field.calls
- assert attached[0].requires_grad and not detached[0].requires_grad
- torch.testing.assert_close(attached[0], detached[0], rtol=0, atol=0)
- # No new bridge noise, time, observation encoding, or dropout mask draw.
- for index in (1, 2, 3):
- assert attached[index] is detached[index]
- torch.testing.assert_close(
- batch["action_flow/fm_velocity_residual"],
- batch["action_flow/velocity_residual"],
- rtol=0,
- atol=0,
- )
-
-
-def test_shared_forward_matches_two_forward_reference_gradients():
- encoder, decoder, field, target, noise, condition = _models_and_inputs()
- reference_encoder, reference_decoder, reference_field = (
- copy.deepcopy(encoder),
- copy.deepcopy(decoder),
- copy.deepcopy(field),
- )
- reference_condition = condition.detach().clone().requires_grad_()
-
- output = _forward(
- "all_stopgrad",
- encoder,
- decoder,
- field,
- target,
- noise,
- condition,
- execution="shared",
- )
- torch.manual_seed(72)
- clean = reference_encoder(target)
- bridge = LatentBridgeStage(
- samples_per_content=3,
- condition_dropout_probability=0.3,
- )(
- {
- "action_flow/clean_latent": clean,
- "sampler/noise": noise,
- "condition": reference_condition,
- }
- )
- state = bridge["action_flow/state"]
- time = bridge["action_flow/time"]
- repeated_condition = bridge["action_flow/condition"]
- drop_mask = bridge["action_flow/condition_drop_mask"]
- target_velocity = bridge["action_flow/target_velocity"]
- action_prediction = reference_field(
- state,
- time,
- repeated_condition,
- condition_drop_mask=drop_mask,
- )
- flow_prediction = reference_field(
- state.detach(),
- time,
- repeated_condition,
- condition_drop_mask=drop_mask,
- )
- reference_batch = ContentDecoderStage(reference_decoder)(
- {
- **bridge,
- "target": target,
- "action_flow/velocity_residual": action_prediction - target_velocity,
- }
- )
- reference_batch["action_flow/fm_velocity_residual"] = (
- flow_prediction - target_velocity.detach()
- )
- reference_output = ActionFlowObjectiveStage(
- residual_key="action_flow/fm_velocity_residual"
- )(reference_batch)
-
- parameters = (
- *encoder.parameters(),
- *decoder.parameters(),
- *field.parameters(),
- condition,
- )
- reference_parameters = (
- *reference_encoder.parameters(),
- *reference_decoder.parameters(),
- *reference_field.parameters(),
- reference_condition,
- )
- gradients = torch.autograd.grad(output["loss/action_flow"], parameters)
- reference_gradients = torch.autograd.grad(
- reference_output["loss/action_flow"], reference_parameters
- )
- for actual, expected in zip(gradients, reference_gradients):
- torch.testing.assert_close(actual, expected, rtol=1.0e-10, atol=1.0e-12)
- assert len(field.calls) == 1
- assert len(reference_field.calls) == 2
-
-
-def test_latent_only_sg_preserves_all_action_flow_gradients_and_diagnostics():
- originals = _models_and_inputs()
- results = []
- for mode in ("full", "all_stopgrad"):
- encoder, decoder, field = [copy.deepcopy(module) for module in originals[:3]]
- target, noise = originals[3:5]
- condition = originals[5].detach().clone().requires_grad_()
- batch = _forward(mode, encoder, decoder, field, target, noise, condition)
- params = (
- *encoder.parameters(),
- *decoder.parameters(),
- *field.parameters(),
- condition,
- )
- gradients = torch.autograd.grad(
- batch["log/action_flow_action_velocity"],
- params,
- allow_unused=True,
- retain_graph=True,
- )
- # The encoder must still receive generative supervision through the
- # bridge state, target tangent, and decoder Jacobian evaluation point.
- encoder_gradients = gradients[: len(tuple(encoder.parameters()))]
- assert all(gradient is not None for gradient in encoder_gradients)
- assert sum(gradient.abs().sum() for gradient in encoder_gradients) > 0
- if mode == "full":
- fm_gradients = torch.autograd.grad(
- batch["log/action_flow_fm"], tuple(encoder.parameters())
- )
- assert sum(gradient.abs().sum() for gradient in fm_gradients) > 0
- assert len(field.calls) == 1
- assert (
- batch["action_flow/fm_velocity_residual"]
- is batch["action_flow/velocity_residual"]
- )
- results.append((batch, gradients))
- for key in (
- "action_flow/state",
- "action_flow/target_velocity",
- "action_flow/predicted_velocity",
- "action_flow/velocity_residual",
- "action_flow/decoded_velocity_residual",
- ):
- torch.testing.assert_close(
- results[0][0][key], results[1][0][key], rtol=0, atol=0
- )
- for full, isolated in zip(results[0][1], results[1][1]):
- if full is None:
- assert isolated is None
- else:
- torch.testing.assert_close(full, isolated, rtol=0, atol=0)
-
-
-@pytest.mark.parametrize("mode", ["target_stopgrad", "all", None])
-def test_candidate_rejects_unsupported_gradient_modes(mode):
- with pytest.raises(ValueError, match="flow_clean_gradient_mode"):
- ConditionalVelocityStage(_Field(), flow_clean_gradient_mode=mode)
-
-
-def test_fm_key_cannot_clobber_action_flow_residual():
- with pytest.raises(ValueError, match="must be distinct"):
- ConditionalVelocityStage(
- _Field(), flow_residual_key="action_flow/velocity_residual"
- )
diff --git a/tests/test_action_flow_diagnostics.py b/tests/test_action_flow_diagnostics.py
deleted file mode 100644
index 7f1d8c55e1..0000000000
--- a/tests/test_action_flow_diagnostics.py
+++ /dev/null
@@ -1,929 +0,0 @@
-from __future__ import annotations
-
-import hashlib
-import json
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-from hydra import compose, initialize_config_dir
-from omegaconf import OmegaConf
-
-from egomimic.eval.action_flow_diagnostic_forward import (
- _decoder_singular_values,
- _fixed_decoder_singular_values,
- collect_action_flow_diagnostics,
-)
-from egomimic.eval.action_flow_diagnostics import ActionFlowDiagnostics
-from egomimic.eval.pipeline_diagnostics import ActionFlowDiagnosticProvider
-from egomimic.eval.planar_action_eval import (
- USOCKET_NATIVE_ERROR_CONFIG,
- PlanarActionEval,
-)
-from egomimic.models.unite_action_decoder import UniteActionDecoder
-from egomimic.pipeline.algo import PipelineAlgo
-from egomimic.pipeline.pushshapes import USocketRotVecNativeDecoder
-from egomimic.pipeline.stages_action_flow import (
- ConditionalVelocityStage,
- ContentDecoderStage,
- ContentEncoderStage,
- RoutedContentDecoderStage,
- RoutedContentEncoderStage,
-)
-from egomimic.pl_utils.pl_model import ModelWrapper
-from egomimic.pl_utils.training_behavior_action_flow import ActionFlowTrainingBehavior
-
-_CONFIG_DIR = Path(__file__).parents[1] / "egomimic/hydra_configs"
-
-
-def test_chunked_decoder_jacobian_singular_values_match_full_jacrev():
- torch.manual_seed(29)
- decoder = nn.Sequential(nn.Linear(3, 12), nn.Tanh(), nn.Linear(12, 10))
- values = torch.randn(2, 5, 3)
- actual = _decoder_singular_values(decoder, values, sample_count=2)
- expected = []
- for value in values:
- jacobian = torch.func.jacrev(
- lambda item: decoder(item.unsqueeze(0)).squeeze(0)
- )(value)
- matrix = jacobian.float().reshape(jacobian.numel() // value.numel(), -1)
- expected.append(torch.linalg.svdvals(matrix))
- torch.testing.assert_close(actual, torch.stack(expected))
-
-
-def test_forward_math_decoder_jacobian_matches_reverse_for_transformer():
- torch.manual_seed(31)
- decoder = UniteActionDecoder(
- latent_dim=4,
- action_dim=2,
- num_latent_tokens=2,
- action_horizon=3,
- hidden_dim=24,
- depth=2,
- num_heads=4,
- dropout=0.0,
- gradient_checkpointing=False,
- ).eval()
- values = torch.randn(2, 2, 4)
- forward = _decoder_singular_values(
- decoder, values, sample_count=2, jacobian_method="forward_math_chunk4"
- )
- reverse = _decoder_singular_values(
- decoder, values, sample_count=2, jacobian_method="reverse_chunked"
- )
- torch.testing.assert_close(forward, reverse, rtol=1e-4, atol=1e-5)
-
-
-def _endpoint_decoder():
- return UniteActionDecoder(
- latent_dim=4,
- action_dim=2,
- num_latent_tokens=2,
- action_horizon=3,
- hidden_dim=24,
- depth=2,
- num_heads=4,
- dropout=0.0,
- gradient_checkpointing=False,
- ).eval()
-
-
-def test_exact_endpoint_reuse_preserves_spectra_state_and_rng(monkeypatch):
- import egomimic.eval.action_flow_diagnostic_forward as owner
-
- torch.manual_seed(51)
- decoder = _endpoint_decoder()
- clean, noise = torch.randn(2, 2, 4), torch.randn(2, 2, 4)
- levels = (0.0, 0.25, 0.5, 0.75, 1.0)
- fixed = torch.stack([(1 - t) * clean + t * noise for t in levels])
- method = "forward_math_chunk4"
- clean_s = _decoder_singular_values(decoder, clean, 2, jacobian_method=method)
- noise_s = _decoder_singular_values(decoder, noise, 2, jacobian_method=method)
- expected = torch.stack(
- [
- _decoder_singular_values(decoder, value, 2, jacobian_method=method)
- for value in fixed
- ]
- )
- weights = {k: v.clone() for k, v in decoder.state_dict().items()}
- flags = [p.requires_grad for p in decoder.parameters()]
- rng = torch.get_rng_state().clone()
- calls = []
-
- def counted(*args, **kwargs):
- calls.append(True)
- return _decoder_singular_values(*args, **kwargs)
-
- monkeypatch.setattr(owner, "_decoder_singular_values", counted)
- actual = _fixed_decoder_singular_values(
- decoder,
- fixed,
- clean,
- noise,
- clean_s,
- noise_s,
- 2,
- jacobian_method=method,
- )
- assert len(calls) == 3
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
- assert torch.equal(rng, torch.get_rng_state())
- assert flags == [p.requires_grad for p in decoder.parameters()]
- assert all(torch.equal(v, weights[k]) for k, v in decoder.state_dict().items())
- assert all(p.grad is None for p in decoder.parameters())
-
-
-@pytest.mark.parametrize(
- "unsupported", ["training", "child_training", "hook", "subclass", "other"]
-)
-def test_endpoint_reuse_keeps_legacy_calls_for_unproven_decoders(
- monkeypatch, unsupported
-):
- import egomimic.eval.action_flow_diagnostic_forward as owner
-
- decoder = _endpoint_decoder()
- if unsupported == "training":
- decoder.train()
- elif unsupported == "child_training":
- next(iter(decoder.children())).train()
- elif unsupported == "hook":
- decoder.register_forward_hook(lambda module, inputs, output: output)
- elif unsupported == "subclass":
-
- class CustomDecoder(UniteActionDecoder):
- pass
-
- decoder.__class__ = CustomDecoder
- else:
- decoder = nn.Identity().eval()
- clean, noise = torch.ones(2, 2, 4), torch.zeros(2, 2, 4)
- fixed = torch.stack([clean, 0.5 * clean, noise])
- calls = []
-
- def counted(_decoder, _value, count, **kwargs):
- calls.append(True)
- return torch.ones(count, 2)
-
- monkeypatch.setattr(owner, "_decoder_singular_values", counted)
- _fixed_decoder_singular_values(
- decoder,
- fixed,
- clean,
- noise,
- torch.ones(2, 2),
- torch.ones(2, 2),
- 2,
- jacobian_method="forward_math_chunk4",
- )
- assert len(calls) == 3
-
-
-def test_endpoint_reuse_requires_exact_selected_input_bytes(monkeypatch):
- import egomimic.eval.action_flow_diagnostic_forward as owner
-
- decoder = _endpoint_decoder()
- clean, noise = torch.ones(2, 2, 4), torch.zeros(2, 2, 4)
- changed = clean.clone()
- changed[0, 0, 0] = torch.nextafter(changed[0, 0, 0], torch.tensor(2.0))
- negative_zero = noise.clone()
- negative_zero[0, 0, 0] = -0.0
- unsampled_change = clean.clone()
- unsampled_change[1] += 1
- calls = []
-
- def counted(_decoder, _value, count, **kwargs):
- calls.append(True)
- return torch.ones(count, 2)
-
- monkeypatch.setattr(owner, "_decoder_singular_values", counted)
- _fixed_decoder_singular_values(
- decoder,
- torch.stack([changed, negative_zero, unsampled_change]),
- clean,
- noise,
- torch.ones(1, 2),
- torch.ones(1, 2),
- 1,
- jacobian_method="forward_math_chunk4",
- )
- assert len(calls) == 2 # Only the unchanged sampled row can reuse its spectrum.
-
-
-@pytest.fixture(autouse=True)
-def _isolate_scheduler_execution_identity(monkeypatch):
- """Keep artifact-path tests independent of the scheduler running pytest."""
- monkeypatch.delenv("SLURM_JOB_ID", raising=False)
- monkeypatch.delenv("SLURM_RESTART_COUNT", raising=False)
-
-
-class _IdentityNormalizer:
- @staticmethod
- def unnormalize(values, _selector):
- return values
-
-
-def _seed_bank(tmp_path, seeds=(101, 202, 303, 404)):
- path = tmp_path / "diagnostic-seeds.json"
- payload = json.dumps({"seeds": list(seeds)}, sort_keys=True).encode()
- path.write_bytes(payload)
- return path, hashlib.sha256(payload).hexdigest()
-
-
-def _config(tmp_path, **overrides):
- path, digest = _seed_bank(tmp_path)
- config = {
- "noise_seed_bank_path": str(path),
- "noise_seed_bank_sha256": digest,
- "raw_noise_levels": [0.0, 0.5, 1.0],
- "max_batches_per_rank": 1,
- "max_samples": None,
- "jacobian_samples": 2,
- "capture_activations": True,
- "activation_layer_map": {0: 0, 1: 2},
- "cknna_k": 2,
- "artifact_root": str(tmp_path / "action-flow-artifacts"),
- "validation_view": {
- "definition": "first_fixed_validation_batch_per_rank",
- "per_rank_batch_size": 6,
- "world_size": 1,
- "split_manifest_sha256": "split-sha",
- },
- "provenance": {
- "source_commit": "0123456789abcdef",
- "resolved_config_sha256": "config-sha",
- "normalization_sha256": "normalization-sha",
- "sampler": "reverse_euler",
- "sampler_steps": 2,
- },
- }
- config.update(overrides)
- return config
-
-
-def _decode(latent, action_dim=2):
- if action_dim == 4:
- return torch.cat(
- (
- latent[..., :2],
- torch.cos(latent[..., 2:3]),
- torch.sin(latent[..., 2:3]),
- ),
- dim=-1,
- )
- return latent[..., :action_dim]
-
-
-def _diagnostic(
- seed=101,
- *,
- nonfinite=False,
- include_activations=True,
- action_dim=2,
- latent_tokens=2,
- action_horizon=None,
- zero_clean=False,
-):
- batch_size, horizon, latent_dim = 6, latent_tokens, 3
- action_horizon = horizon if action_horizon is None else int(action_horizon)
- if action_horizon % horizon:
- raise ValueError("test action_horizon must be a multiple of latent_tokens")
-
- def decode_action(value):
- decoded = _decode(value, action_dim)
- return decoded.repeat_interleave(action_horizon // horizon, dim=-2)
-
- levels = torch.tensor([0.0, 0.5, 1.0])
- clean = (
- torch.arange(batch_size * horizon * latent_dim, dtype=torch.float32)
- .reshape(batch_size, horizon, latent_dim)
- .div(10.0)
- .add(0.25)
- )
- if zero_clean:
- clean = torch.zeros_like(clean)
- noise = torch.flip(clean, dims=(0,)).neg().sub(0.5)
- fixed_states = torch.stack(
- [(1.0 - level) * clean + level * noise for level in levels]
- )
- predicted_clean = torch.stack([clean + float(level) * 0.03 for level in levels])
- fixed_final = torch.stack([clean + float(level) * 0.01 for level in levels])
- generated = fixed_final[-1]
- middle = 0.5 * (noise + generated)
- trajectory = torch.stack((noise, middle, generated))
- target = decode_action(clean)
- diagnostic = {
- "schema": "action-flow-validation-diagnostics/v1",
- "source": "validation/usocket",
- "noise_seed": seed,
- "returned_samples": batch_size,
- "jacobian_samples": 2,
- "integration_semantics": (
- "reverse_euler_partial_to_lower_canonical_grid_then_uniform_to_zero"
- ),
- "noise_levels": levels,
- "num_inference_steps": 2,
- "target": target,
- "condition": torch.arange(batch_size * 4, dtype=torch.float32).reshape(
- batch_size, 4
- ),
- "latent/clean": clean,
- "latent/noise": noise,
- "latent/generated": generated,
- "latent/fixed_states": fixed_states,
- "latent/predicted_clean": predicted_clean,
- "latent/fixed_final": fixed_final,
- "field/predicted_velocity": noise.unsqueeze(0).repeat(3, 1, 1, 1),
- "field/velocity_residual": torch.full_like(fixed_states, 0.125),
- "latent/trajectory": trajectory,
- "decoded/reconstruction": target,
- "decoded/noise": decode_action(noise),
- "decoded/generated": decode_action(generated),
- "decoded/fixed_states": decode_action(fixed_states),
- "decoded/predicted_clean": decode_action(predicted_clean),
- "decoded/fixed_final": decode_action(fixed_final),
- "decoded/trajectory": decode_action(trajectory),
- "fixed_level_field_evaluations": torch.tensor([0, 1, 2]),
- "decoder_jacobian/clean_singular_values": torch.tensor(
- [[2.0, 1.0, 0.5, 0.25], [1.8, 0.9, 0.4, 0.2]]
- ),
- "decoder_jacobian/noise_singular_values": torch.tensor(
- [[1.5, 0.75, 0.3, 0.1], [1.4, 0.7, 0.25, 0.08]]
- ),
- "decoder_jacobian/fixed_singular_values": torch.tensor(
- [
- [[2.0, 1.0, 0.5, 0.25], [1.8, 0.9, 0.4, 0.2]],
- [[1.7, 0.8, 0.4, 0.2], [1.6, 0.7, 0.3, 0.1]],
- [[1.5, 0.7, 0.3, 0.1], [1.4, 0.6, 0.2, 0.08]],
- ]
- ),
- }
- if include_activations:
- row = torch.arange(batch_size, dtype=torch.float32)
- base = torch.stack((row, row.square(), row.sin(), row.cos()), dim=-1)
- base = base[:, None, :].repeat(1, horizon, 1)
- encoder = torch.stack((base, 2.0 * base + 0.1))
- field = torch.empty(3, 3, batch_size, horizon, 4)
- for level in range(3):
- field[level, 0] = encoder[0] * (1.0 + level)
- field[level, 1] = base.roll(1, dims=0)
- field[level, 2] = encoder[1] * (1.0 + 0.5 * level)
- diagnostic.update(
- {
- "activation/encoder_blocks": encoder,
- "activation/encoder_block_indices": torch.tensor([0, 1]),
- "activation/field_blocks": field,
- "activation/field_block_indices": torch.tensor([0, 1, 2]),
- }
- )
- if nonfinite:
- diagnostic["latent/generated"] = generated.clone()
- diagnostic["latent/generated"][0, 0, 0] = float("nan")
- return diagnostic
-
-
-class _DiagnosticModel:
- def __init__(self, diagnostic):
- self.diagnostic = diagnostic
- self.calls = []
-
- @staticmethod
- def forward_eval(batch):
- return {
- source: {"pred_action": source_batch["actions"]}
- for source, source_batch in batch.items()
- }
-
- def run_diagnostic(self, capability, batch, **kwargs):
- assert capability == "action_flow"
- assert not torch.is_inference_mode_enabled()
- self.calls.append(kwargs)
- return {source: self.diagnostic for source in batch}
-
-
-class _TwoBlockCodec(nn.Module):
- def __init__(self):
- super().__init__()
- self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
-
- def forward(self, value):
- for block in self.blocks:
- value = torch.tanh(block(value))
- return value
-
-
-class _TwoBlockField(nn.Module):
- def __init__(self):
- super().__init__()
- self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
-
- def forward(self, value, time, condition, *, condition_drop_mask=None):
- if condition_drop_mask is None:
- raise AssertionError("diagnostics must supply an explicit drop mask")
- effective = torch.where(
- condition_drop_mask[:, None], torch.zeros_like(condition), condition
- )
- value = value + time[:, None, None] + effective[:, None, :]
- for block in self.blocks:
- value = torch.tanh(block(value))
- return value
-
-
-def _batch(target):
- return {
- "validation/usocket": {
- "actions": target,
- "embodiment": torch.full((target.shape[0],), 19),
- }
- }
-
-
-def test_planar_evaluator_logs_and_hashes_action_flow_diagnostics(tmp_path):
- diagnostic = _diagnostic()
- config = _config(tmp_path)
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- action_flow_diagnostics={"enabled": True, **config},
- )
- evaluator.bind_data_context(normalizer=_IdentityNormalizer())
- evaluator.model = _DiagnosticModel(diagnostic)
- logged = {}
- evaluator.trainer = SimpleNamespace(
- current_epoch=3,
- global_step=41,
- global_rank=0,
- precision="bf16-mixed",
- lightning_module=SimpleNamespace(
- log_dict=lambda metrics, **_kwargs: logged.update(metrics)
- ),
- )
- evaluator.on_validation_start()
- torch.manual_seed(8675309)
- rng_state = torch.random.get_rng_state().clone()
-
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=7)
-
- assert torch.equal(torch.random.get_rng_state(), rng_state)
- assert evaluator.model.calls == [
- {
- "raw_noise_levels": [0.0, 0.5, 1.0],
- "noise_seed": 101,
- "max_samples": None,
- "jacobian_samples": 2,
- "capture_activations": True,
- }
- ]
- required_metrics = (
- "Valid/ActionFlow/Latent/clean/RMS",
- "Valid/ActionFlow/Latent/generated/EffectiveRank",
- "Valid/ActionFlow/Latent/clean/CovarianceEigenvalue/eig_00",
- "Valid/ActionFlow/DecoderJacobian/clean/spectral_norm_mean",
- "Valid/ActionFlow/DecodedFullNoise/token_radius_mean",
- "Valid/ActionFlow/Compute/InferenceFieldEvaluations",
- "Valid/ActionFlow/DenoisingTrajectory/LatentMSE/t1000",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedMSE/t0000",
- "Valid/ActionFlow/FixedLevel/final_latent_mse/t0500",
- "Valid/ActionFlow/Alignment/FinalLatentCosine/t1000",
- "Valid/ActionFlow/Alignment/CKA/encoder_00__field_00/t0500",
- "Valid/ActionFlow/Alignment/CKNNA/encoder_01__field_02/t1000",
- )
- for name in required_metrics:
- assert name in logged
- assert logged[name].ndim == 0
- assert bool(torch.isfinite(logged[name]))
- assert f"{name}/pushshapes_sim_u_socket" in logged
- assert float(
- logged["Valid/ActionFlow/Alignment/CKA/encoder_00__field_00/t0500"]
- ) == pytest.approx(1.0, abs=1.0e-5)
-
- artifact = tmp_path / "action-flow-artifacts/epoch-3-step-41/rank-0-batch-7.pt"
- sidecar = artifact.with_name(f"{artifact.name}.sha256")
- assert artifact.is_file() and sidecar.is_file()
- sidecar_payload = json.loads(sidecar.read_text())
- assert (
- sidecar_payload["sha256"] == hashlib.sha256(artifact.read_bytes()).hexdigest()
- )
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- assert payload["metric"] == "ActionFlowValidationDiagnostics"
- assert (
- payload["identity_sha256"] == evaluator._action_flow_diagnostics.identity_sha256
- )
- assert payload["noise_seed"] == 101
- assert payload["sampler"]["name"] == "reverse_euler"
- assert payload["statistics"]["trajectory_error"] == (
- "paired_diagnostic_not_distributional_score"
- )
- source = payload["sources"]["pushshapes_sim_u_socket"]
- assert source["computed"]["final_latent_cosine_by_condition"].shape == (3, 6)
- assert source["computed"]["latent_statistics"]["clean"][
- "covariance_eigenvalues"
- ].shape == (3,)
-
- evaluator.on_validation_start()
- with pytest.raises(FileExistsError, match="refusing to overwrite"):
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=7)
-
-
-def test_action_flow_diagnostics_allow_latent_tokens_to_differ_from_action_horizon(
- tmp_path,
-):
- diagnostic = _diagnostic(latent_tokens=2, action_horizon=4)
- config = _config(tmp_path)
- config["provenance"]["latent_shape"] = [2, 3]
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- action_flow_diagnostics={"enabled": True, **config},
- )
- evaluator.bind_data_context(normalizer=_IdentityNormalizer())
- evaluator.model = _DiagnosticModel(diagnostic)
- logged = {}
- evaluator.trainer = SimpleNamespace(
- current_epoch=0,
- global_step=1,
- global_rank=0,
- precision="bf16-mixed",
- lightning_module=SimpleNamespace(
- log_dict=lambda metrics, **_kwargs: logged.update(metrics)
- ),
- )
- evaluator.on_validation_start()
-
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=0)
-
- assert "Valid/ActionFlow/CleanReconstructionMSE" in logged
- artifact = tmp_path / "action-flow-artifacts/epoch-0-step-1/rank-0-batch-0.pt"
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- saved = payload["sources"]["pushshapes_sim_u_socket"]["diagnostic"]
- assert saved["latent/clean"].shape == (6, 2, 3)
- assert saved["target"].shape == (6, 4, 2)
-
-
-def test_action_flow_diagnostics_reject_undefined_zero_target_cosine(tmp_path):
- diagnostic = _diagnostic(zero_clean=True)
- runner = ActionFlowDiagnostics(_config(tmp_path))
-
- with pytest.raises(ValueError, match="zero clean-latent norm"):
- runner.run(
- model=_DiagnosticModel(diagnostic),
- batch=_batch(diagnostic["target"]),
- batch_idx=0,
- rank=0,
- epoch=0,
- global_step=1,
- precision="bf16-mixed",
- source_labels={"validation/usocket": "usocket"},
- )
-
-
-def test_action_flow_artifacts_preserve_slurm_attempts_and_same_attempt_refusal(
- tmp_path, monkeypatch
-):
- diagnostic = _diagnostic()
- legacy = tmp_path / "action-flow-artifacts/epoch-3-step-41/rank-0-batch-7.pt"
- legacy.parent.mkdir(parents=True)
- legacy.write_bytes(b"historical diagnostic")
- for job_id, restart in (("5714540", 0), ("5714540", 1), ("5715500", 0)):
- monkeypatch.setenv("SLURM_JOB_ID", job_id)
- monkeypatch.setenv("SLURM_RESTART_COUNT", str(restart))
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- action_flow_diagnostics={"enabled": True, **_config(tmp_path)},
- )
- evaluator.bind_data_context(normalizer=_IdentityNormalizer())
- evaluator.model = _DiagnosticModel(diagnostic)
- evaluator.trainer = SimpleNamespace(
- current_epoch=3,
- global_step=41,
- global_rank=0,
- precision="32-true",
- lightning_module=SimpleNamespace(log_dict=lambda *_args, **_kwargs: None),
- )
- evaluator.on_validation_start()
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=7)
- artifact = (
- legacy.parent.parent
- / f"job-{job_id}-restart-{restart}"
- / legacy.parent.name
- / legacy.name
- )
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- assert payload["execution"] == {
- "slurm_job_id": job_id,
- "slurm_restart_count": restart,
- }
- assert (
- payload["identity_sha256"]
- == evaluator._action_flow_diagnostics.identity_sha256
- )
- sidecar = json.loads(Path(f"{artifact}.sha256").read_text())
- assert sidecar["sha256"] == hashlib.sha256(artifact.read_bytes()).hexdigest()
- evaluator.on_validation_start()
- with pytest.raises(FileExistsError, match="refusing to overwrite"):
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=7)
- assert legacy.read_bytes() == b"historical diagnostic"
-
-
-def test_planar_action_flow_diagnostics_emit_native_circular_errors(tmp_path):
- diagnostic = _diagnostic(action_dim=4)
- config = _config(tmp_path, native_error=USOCKET_NATIVE_ERROR_CONFIG)
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- native_decoder=USocketRotVecNativeDecoder(),
- action_flow_diagnostics={"enabled": True, **config},
- )
- evaluator.bind_data_context(normalizer=_IdentityNormalizer())
- evaluator.model = _DiagnosticModel(diagnostic)
- logged = {}
- evaluator.trainer = SimpleNamespace(
- current_epoch=0,
- global_step=2,
- global_rank=0,
- precision="bf16-mixed",
- lightning_module=SimpleNamespace(
- log_dict=lambda metrics, **_kwargs: logged.update(metrics)
- ),
- )
- evaluator.on_validation_start()
-
- evaluator.on_validation_step(_batch(diagnostic["target"]), batch_idx=0)
-
- required = (
- "Valid/ActionFlow/CleanReconstructionNativeMSE",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedNativeMSE/t1000",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedNativeMSE/t0500",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedNativeMSE/t0000",
- "Valid/ActionFlow/FixedLevel/final_decoded_native_mse/t0500",
- )
- for name in required:
- assert name in logged
- assert f"{name}/pushshapes_sim_u_socket" in logged
- assert bool(torch.isfinite(logged[name]))
- assert float(logged["Valid/ActionFlow/CleanReconstructionNativeMSE"]) == 0.0
-
- artifact = tmp_path / "action-flow-artifacts/epoch-0-step-2/rank-0-batch-0.pt"
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- assert payload["identity"]["native_error"] == USOCKET_NATIVE_ERROR_CONFIG
- assert payload["statistics"]["native_action_error"] == (USOCKET_NATIVE_ERROR_CONFIG)
- computed = payload["sources"]["pushshapes_sim_u_socket"]["computed"]
- assert computed["clean_reconstruction_native_mse_by_condition"].shape == (6,)
- assert computed["trajectory_decoded_native_mse_by_condition"].shape == (3, 6)
-
-
-def test_native_action_flow_error_wraps_theta_at_pi_boundary():
- target_theta = torch.full((2, 3), torch.pi - 0.01)
- predicted_theta = torch.full((2, 3), -torch.pi + 0.01)
-
- def rotvec(theta):
- return torch.stack(
- (
- torch.zeros_like(theta),
- torch.zeros_like(theta),
- torch.cos(theta),
- torch.sin(theta),
- ),
- dim=-1,
- )
-
- decoder = USocketRotVecNativeDecoder()
- values = PlanarActionEval._native_mse_by_condition(
- decoder.decode(rotvec(predicted_theta)),
- decoder.decode(rotvec(target_theta)),
- decoder,
- )
-
- expected = torch.full((2,), (0.02**2) / 3.0)
- torch.testing.assert_close(values, expected, atol=1.0e-7, rtol=0.0)
-
-
-def test_action_flow_consumer_matches_real_wrapper_schema(tmp_path):
- torch.manual_seed(91)
- wrapper = ModelWrapper(
- pipeline=PipelineAlgo(
- stages=[
- ContentEncoderStage(_TwoBlockCodec()),
- ConditionalVelocityStage(_TwoBlockField(), num_inference_steps=2),
- ContentDecoderStage(_TwoBlockCodec()),
- ],
- device="cpu",
- ),
- training_behavior=ActionFlowTrainingBehavior(gradient_telemetry_cadence=0),
- diagnostic_provider=ActionFlowDiagnosticProvider(),
- )
- wrapper.eval()
- runner = ActionFlowDiagnostics(_config(tmp_path, activation_layer_map={0: 0, 1: 1}))
- batch = {
- "opaque_source": {
- "target": torch.randn(6, 2, 2) + 0.25,
- "condition": torch.randn(6, 2),
- }
- }
-
- metrics = runner.run(
- model=wrapper,
- batch=batch,
- batch_idx=0,
- rank=0,
- epoch=0,
- global_step=0,
- precision="32-true",
- source_labels={"opaque_source": "opaque"},
- )
-
- assert "Valid/ActionFlow/Latent/clean/RMS" in metrics
- assert "Valid/ActionFlow/Alignment/CKA/encoder_01__field_01/t1000" in metrics
- assert all(bool(torch.isfinite(value)) for value in metrics.values())
-
-
-def test_action_flow_diagnostic_forward_selects_private_routed_codecs():
- torch.manual_seed(92)
- model = PipelineAlgo(
- stages=[
- RoutedContentEncoderStage(
- encoders={"route_a": _TwoBlockCodec(), "route_b": _TwoBlockCodec()},
- route_key="route",
- ),
- ConditionalVelocityStage(_TwoBlockField(), num_inference_steps=2),
- RoutedContentDecoderStage(
- decoders={"route_a": _TwoBlockCodec(), "route_b": _TwoBlockCodec()},
- route_key="route",
- ),
- ],
- device="cpu",
- )
- batch = {
- route: {
- "target": torch.randn(3, 2, 2),
- "condition": torch.randn(3, 2),
- "route": route,
- }
- for route in ("route_a", "route_b")
- }
-
- diagnostics = collect_action_flow_diagnostics(
- model,
- batch,
- raw_noise_levels=(0.0, 1.0),
- noise_seed=101,
- max_samples=3,
- jacobian_samples=1,
- capture_activations=True,
- already_processed=True,
- )
-
- assert tuple(diagnostics) == ("route_a", "route_b")
- for diagnostic in diagnostics.values():
- assert diagnostic["provenance/encoder_class"] == "_TwoBlockCodec"
- assert diagnostic["provenance/decoder_class"] == "_TwoBlockCodec"
- assert diagnostic["activation/encoder_blocks"].shape[:2] == (2, 3)
- assert diagnostic["decoded/generated"].shape == (3, 2, 2)
-
-
-@pytest.mark.parametrize(
- "experiment",
- (
- "action_flow_bc_usocket_recon1_s42",
- "action_flow_bc_usocket_recon10_s42",
- ),
-)
-def test_real_action_flow_config_constructs_strict_diagnostics(tmp_path, experiment):
- with initialize_config_dir(
- version_base=None, config_dir=str(_CONFIG_DIR.resolve())
- ):
- config = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/pusht/{experiment}", "++paths.root_dir=."],
- )
- diagnostics = config.evaluator.action_flow_diagnostics
- diagnostics.noise_seed_bank_path = str(
- _CONFIG_DIR / "evaluator/energy_score_seed_bank_k32_v1.json"
- )
- diagnostics.artifact_root = str(tmp_path / experiment)
- diagnostics = OmegaConf.to_container(diagnostics, resolve=True)
-
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- action_flow_diagnostics=diagnostics,
- )
-
- runner = evaluator._action_flow_diagnostics
- assert evaluator.action_flow_diagnostics_enabled
- assert runner.noise_levels == [0.0, 0.25, 0.5, 0.75, 1.0]
- assert runner.max_samples == 16
- assert runner.jacobian_samples == 2
- assert runner.activation_layer_map == ((0, 0), (1, 11))
- assert runner.cknna_k == 10
-
-
-def test_action_flow_diagnostics_fail_closed_without_requested_activations(tmp_path):
- runner = ActionFlowDiagnostics(_config(tmp_path))
- diagnostic = _diagnostic(include_activations=False)
- model = _DiagnosticModel(diagnostic)
- batch = _batch(diagnostic["target"])
-
- with pytest.raises(TypeError, match="activation/encoder_blocks"):
- runner.run(
- model=model,
- batch=batch,
- batch_idx=0,
- rank=0,
- epoch=0,
- global_step=0,
- precision="32-true",
- source_labels={"validation/usocket": "usocket"},
- )
- assert not list((tmp_path / "action-flow-artifacts").rglob("*.pt"))
-
-
-def test_action_flow_diagnostics_reject_nonfinite_model_outputs(tmp_path):
- runner = ActionFlowDiagnostics(_config(tmp_path))
- diagnostic = _diagnostic(nonfinite=True)
-
- with pytest.raises(ValueError, match="non-finite.*latent/generated"):
- runner.run(
- model=_DiagnosticModel(diagnostic),
- batch=_batch(diagnostic["target"]),
- batch_idx=0,
- rank=0,
- epoch=0,
- global_step=0,
- precision=None,
- source_labels={"validation/usocket": "usocket"},
- )
-
-
-@pytest.mark.parametrize(
- ("overrides", "match"),
- (
- ({"jacobian_samples": 0}, "jacobian_samples must be positive"),
- ({"max_samples": 2, "jacobian_samples": 3}, "cannot exceed max_samples"),
- ({"capture_activations": True, "activation_layer_map": {}}, "non-empty"),
- ({"cknna_k": 6}, "k < diagnostic batch size"),
- ({"raw_noise_levels": [0.5, 0.5]}, "strictly increase"),
- ),
-)
-def test_action_flow_diagnostic_config_fails_closed(tmp_path, overrides, match):
- with pytest.raises(ValueError, match=match):
- ActionFlowDiagnostics(_config(tmp_path, **overrides))
-
-
-@pytest.mark.parametrize("restart", [None, 0, 1])
-def test_existing_latent_diagnostic_native_conversion_uses_decoder_argument(
- tmp_path,
- monkeypatch,
- restart,
-):
- monkeypatch.delenv("SLURM_JOB_ID", raising=False)
- if restart is not None:
- monkeypatch.setenv("SLURM_JOB_ID", "5714540")
- monkeypatch.setenv("SLURM_RESTART_COUNT", str(restart))
- evaluator = PlanarActionEval(energy_score_enabled=False)
- evaluator.bind_data_context(normalizer=_IdentityNormalizer())
- evaluator.unite_diagnostics = {"artifact_root": str(tmp_path / "legacy-diagnostic")}
- evaluator._unite_noise_seeds = [101]
- evaluator._unite_noise_seed_sha256 = "seed-sha"
- evaluator._unite_noise_levels = [0.0, 1.0]
- evaluator._unite_noise_labels = ["t000", "t1000"]
- evaluator._unite_cknna_k = 2
- target = _diagnostic()["target"]
- clean = torch.cat((target, torch.ones_like(target[..., :1])), dim=-1)
- feature = torch.arange(1, 7, dtype=torch.float32)[:, None, None].repeat(1, 2, 3)
- evaluator.model = SimpleNamespace(
- run_diagnostic=lambda capability, _batch, raw_noise_levels: {
- "validation/usocket": {
- "clean_latent": clean,
- "sampler_latents": torch.stack((clean + 1.0, clean)),
- "decoded_actions_normalized": torch.stack((target + 1.0, target)),
- "noise_level_final_predictions": torch.stack((clean, clean + 0.1)),
- "tokenization_activations": {"layer": feature},
- "denoising_activations": {
- "layer": torch.stack((feature, feature * 2.0))
- },
- }
- }
- if capability == "unite"
- else None
- )
- evaluator.trainer = SimpleNamespace(
- current_epoch=1,
- global_step=2,
- global_rank=0,
- )
- batch = _batch(target)
-
- metrics = evaluator._unite_metrics_and_artifact(batch, batch_idx=0)
-
- assert "Valid/DenoisingTrajectory/DecodedNativeMSE/step_0" in metrics
- root = tmp_path / "legacy-diagnostic"
- if restart is not None:
- root = root / f"job-5714540-restart-{restart}"
- artifact = root / "epoch-1-step-2/rank-0-batch-0.pt"
- assert artifact.is_file()
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- assert payload["execution"] == (
- None
- if restart is None
- else {"slurm_job_id": "5714540", "slurm_restart_count": restart}
- )
- with pytest.raises(FileExistsError, match="refusing to overwrite"):
- evaluator._unite_metrics_and_artifact(batch, batch_idx=0)
diff --git a/tests/test_action_flow_jacobian_checkpoint_performance.py b/tests/test_action_flow_jacobian_checkpoint_performance.py
deleted file mode 100644
index 5408c6ad14..0000000000
--- a/tests/test_action_flow_jacobian_checkpoint_performance.py
+++ /dev/null
@@ -1,555 +0,0 @@
-"""Opt-in, read-only checkpoint equivalence/timing test for exact endpoint reuse.
-
-Set ACTION_FLOW_JACOBIAN_PROFILE_REQUEST to a JSON request outside the checkout.
-Ordinary CPU suites skip the GPU benchmark. No optimizer or policy sampler runs.
-"""
-
-from __future__ import annotations
-
-import copy
-import gc
-import hashlib
-import json
-import os
-import statistics
-import time
-from pathlib import Path
-from types import SimpleNamespace
-from unittest.mock import patch
-
-import numpy as np
-import pytest
-import torch
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-
-import egomimic.eval.action_flow_diagnostic_forward as owner
-
-
-def prepare_training_profile(request):
- """Restore native data/model state on CPU; capture one real training batch."""
- from torch.utils.data import default_collate
-
- from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset
-
- assert _digest(request["split"]) == request["split_sha256"]
- split = json.loads(Path(request["split"]).read_text())
- for domain in ("yam", "human"):
- for part in ("train", "valid"):
- os.environ[f"{domain.upper()}_{part.upper()}_IDS_SHA256"] = split[domain][
- "summary"
- ][f"{part}_ids_sha256"]
- assert _digest(request["checkpoint"]) == request["checkpoint_sha256"]
- payload = torch.load(
- request["checkpoint"], map_location="cpu", mmap=True, weights_only=False
- )
- assert payload["global_step"] == request["checkpoint_step"]
- cfg = OmegaConf.load(request["config"])
- assert payload["hyper_parameters"]["config_tree"][
- "model"
- ] == OmegaConf.to_container(cfg.model, resolve=True)
- assert [int(cfg.model.num_latent_tokens), int(cfg.model.latent_dim)] == request[
- "latent_shape"
- ]
- assert (
- int(cfg.data.train_dataloader_params[request["source_key"]].batch_size)
- == request["batch_size"]
- )
- state = copy.deepcopy(payload["data_context"]["normalizer_state"])
- for field in ("key_types", "zarr_keys", "shapes", "norm_stats"):
- state[field] = {int(k): v for k, v in state[field].items()}
- normalizer = MultiDataset.from_state(state)
- fixture_path = Path(request["fixture"])
- if fixture_path.exists():
- if request.get("fixture_sha256"):
- assert _digest(fixture_path) == request["fixture_sha256"]
- fixture = torch.load(fixture_path, weights_only=False)
- assert fixture["checkpoint"] == request.get(
- "fixture_checkpoint", request["checkpoint"]
- )
- assert fixture["source_key"] == request["source_key"]
- else:
- dataset = instantiate(cfg.data.train_datasets[request["source_key"]])
- dataset.set_norm_stats_from(normalizer)
- indices = np.linspace(
- 0, len(dataset) - 1, request["batch_size"], dtype=int
- ).tolist()
- flat = default_collate([dataset[index] for index in indices])
- fixture = {
- "checkpoint": request["checkpoint"],
- "source_key": request["source_key"],
- "indices": indices,
- "batch": flat,
- }
- torch.save(fixture, fixture_path)
- algo = instantiate(cfg.model.pipeline, device="cpu")
- algo.bind_data_context(normalizer=normalizer)
- weights = {
- k.removeprefix("nets."): v
- for k, v in payload["state_dict"].items()
- if k.startswith("nets.")
- }
- algo.nets.load_state_dict(weights, strict=True)
- return cfg, payload, algo, fixture
-
-
-def _cpu_tree(value):
- if torch.is_tensor(value):
- return value.detach().cpu().clone()
- if isinstance(value, dict):
- return {k: _cpu_tree(v) for k, v in value.items()}
- if isinstance(value, (tuple, list)):
- return type(value)(_cpu_tree(v) for v in value)
- return copy.deepcopy(value)
-
-
-def _compare_training_tree(actual, expected):
- """Compare every tensor; retain exact mismatch counts and largest errors."""
- mismatches, tensor_count = [], 0
-
- def visit(a, b, path):
- nonlocal tensor_count
- if torch.is_tensor(a):
- tensor_count += 1
- if not (torch.is_tensor(b) and a.shape == b.shape and a.dtype == b.dtype):
- mismatches.append(
- {"path": path, "max_abs": 0.0, "contract_changed": True}
- )
- return
- if not torch.equal(a, b):
- diff = (a.double() - b.double()).abs()
- mismatches.append(
- {
- "path": path,
- "max_abs": float(diff.max()),
- "different_elements": int(torch.count_nonzero(diff)),
- "relative_rms": float(
- torch.linalg.vector_norm(diff)
- / torch.linalg.vector_norm(b.double()).clamp_min(1e-30)
- ),
- }
- )
- elif isinstance(a, dict):
- assert a.keys() == b.keys(), path
- for k in a:
- visit(a[k], b[k], path + "/" + str(k))
- elif isinstance(a, (tuple, list)):
- assert len(a) == len(b), path
- for k, (x, y) in enumerate(zip(a, b)):
- visit(x, y, path + "/" + str(k))
- elif torch.is_tensor(b) or a != b:
- mismatches.append(
- {"path": path, "actual": str(a), "expected": str(b), "max_abs": 0.0}
- )
-
- visit(actual, expected, "")
- return {
- "bit_exact": not mismatches,
- "tensor_count": tensor_count,
- "mismatch_count": len(mismatches),
- "largest_errors": sorted(mismatches, key=lambda x: x["max_abs"], reverse=True)[
- :12
- ],
- }
-
-
-@pytest.mark.skipif(
- not os.environ.get("ACTION_FLOW_TRAINING_PROFILE_REQUEST"),
- reason="Opt-in native checkpoint training test",
-)
-def test_checkpoint_training_execution_equivalence_and_timing():
- """Test existing execution options without changing/adopting model weights."""
- from egomimic.utils.libero_dit_half import LiberoDiTHalf
-
- assert torch.cuda.is_available() and torch.cuda.device_count() == 1
- request = json.loads(
- Path(os.environ["ACTION_FLOW_TRAINING_PROFILE_REQUEST"]).read_text()
- )
- out = Path(request["output_root"])
- out.mkdir(parents=True, exist_ok=True)
- cfg, payload, initial, fixture = prepare_training_profile(request)
- assert str(cfg.trainer.precision) == "bf16-mixed"
- baseline, records = [], []
- initial_state = _cpu_tree(initial.nets.state_dict())
- del initial
- for case_index, variant in enumerate(request["variants"]):
- algo = instantiate(cfg.model.pipeline, device="cpu")
- algo.nets.load_state_dict(initial_state, strict=True)
- holder = torch.nn.Module()
- holder.nets = algo.nets
- holder.cuda().train()
- algo.device = torch.device("cuda")
- optimizer = instantiate(
- cfg.model.optimizer, named_params=tuple(holder.named_parameters())
- )()
- scheduler = instantiate(cfg.model.scheduler, optimizer=optimizer)()
- optimizer.load_state_dict(copy.deepcopy(payload["optimizer_states"][0]))
- scheduler.load_state_dict(copy.deepcopy(payload["lr_schedulers"][0]))
- ema = instantiate(cfg.callbacks.ema)
- ema.on_load_checkpoint(None, holder, payload)
- ema.on_fit_start(SimpleNamespace(global_step=payload["global_step"]), holder)
- policy = LiberoDiTHalf(
- source_key=request["source_key"],
- checkpoint_periods={
- "encoder": 2,
- "velocity": {
- "third_flow_recomputation": 3,
- "quarter_flow_recomputation": 4,
- }[variant],
- }
- if variant in {"third_flow_recomputation", "quarter_flow_recomputation"}
- else None,
- )
- policy.on_fit_start(None, SimpleNamespace(model=algo))
- encoder, field, decoder = (
- algo.pipeline.stages[3].encoder,
- algo.pipeline.stages[5],
- algo.pipeline.stages[6].decoder,
- )
- if variant == "no_dit_recomputation":
- encoder.backbone.gradient_checkpointing = False
- field.field.backbone.gradient_checkpointing = False
- elif variant == "no_reconstruction_recomputation":
- decoder.gradient_checkpointing = False
- elif variant == "shared_field_prediction":
- field.fm_field_execution = "shared"
- else:
- assert variant in {
- "baseline",
- "third_flow_recomputation",
- "quarter_flow_recomputation",
- }
- batch = algo.process_batch_for_training(
- {request["source_key"]: fixture["batch"]}
- )
- timings, comparisons = [], []
- failure = None
- try:
- for step in range(request["steps"]):
- torch.manual_seed(request["seed"] + step)
- torch.cuda.synchronize()
- torch.cuda.reset_peak_memory_stats()
- started = time.perf_counter()
- optimizer.zero_grad(set_to_none=True)
- with torch.autocast("cuda", dtype=torch.bfloat16):
- prediction = algo.forward_training(batch)
- loss = algo.compute_losses(prediction, batch)["loss"]
- assert bool(torch.isfinite(loss))
- loss.backward()
- torch.cuda.synchronize()
- backward_done = time.perf_counter()
- gradients = _cpu_tree(
- {name: p.grad for name, p in holder.named_parameters()}
- )
- assert all(
- value is None or bool(torch.isfinite(value).all())
- for value in gradients.values()
- )
- components = _cpu_tree(
- {
- k: v
- for k, v in prediction[request["source_key"]].items()
- if k.startswith("log/action_flow_") or k == "loss/action_flow"
- }
- )
- # Copying proof tensors is deliberately outside measured GPU work.
- optimizer_start = time.perf_counter()
- clip = float(cfg.trainer.get("gradient_clip_val", 0) or 0)
- if clip:
- assert cfg.trainer.get("gradient_clip_algorithm", "norm") == "norm"
- torch.nn.utils.clip_grad_norm_(holder.parameters(), clip)
- optimizer.step()
- scheduler.step()
- ema._update(holder)
- torch.cuda.synchronize()
- ended = time.perf_counter()
- timings.append(
- {
- "step": step,
- "forward_backward_seconds": backward_done - started,
- "optimizer_scheduler_ema_seconds": ended - optimizer_start,
- "seconds": backward_done - started + ended - optimizer_start,
- "peak_cuda_bytes": torch.cuda.max_memory_allocated(),
- }
- )
- actual = {
- "gradients": gradients,
- "components": components,
- "state": _cpu_tree(holder.state_dict()),
- "optimizer": _cpu_tree(optimizer.state_dict()),
- "scheduler": _cpu_tree(scheduler.state_dict()),
- "ema": _cpu_tree(ema._shadow),
- "rng_cpu": torch.get_rng_state(),
- "rng_cuda": torch.cuda.get_rng_state(),
- }
- if case_index == 0:
- baseline.append(actual)
- else:
- comparisons.append(
- {
- section: _compare_training_tree(
- value, baseline[step][section]
- )
- for section, value in actual.items()
- }
- )
- print(
- "TRAINING_PROFILE_STEP",
- variant,
- json.dumps(timings[-1]),
- flush=True,
- )
- del prediction, loss, actual, gradients, components
- except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
- failure = {"type": type(exc).__name__, "message": str(exc)}
- if case_index == 0:
- raise
- finally:
- policy.on_fit_end(None, None)
- record = {
- "variant": variant,
- "timings": timings,
- "comparisons": comparisons,
- "failure": failure,
- "bit_exact": failure is None
- and all(all(s["bit_exact"] for s in step.values()) for step in comparisons),
- "median_seconds": statistics.median(t["seconds"] for t in timings[1:])
- if len(timings) > 1
- else None,
- }
- records.append(record)
- (out / "TRAINING_PROGRESS.json").write_text(json.dumps(records, indent=2))
- prediction = loss = actual = gradients = components = None
- del (
- prediction,
- loss,
- actual,
- gradients,
- components,
- algo,
- holder,
- optimizer,
- scheduler,
- ema,
- batch,
- policy,
- encoder,
- field,
- decoder,
- )
- gc.collect()
- torch.cuda.empty_cache()
- result = {
- "status": "COMPLETE",
- "checkpoint_step": request["checkpoint_step"],
- "batch_size": request["batch_size"],
- "latent_shape": request["latent_shape"],
- "precision": "bf16-mixed_native_FP32_decoder_JVP",
- "torch": torch.__version__,
- "gpu": torch.cuda.get_device_name(),
- "records": records,
- "scope": "native full pipeline forward/backward/optimizer/scheduler/EMA; frozen real training batch; excludes data-loader, telemetry and proof-copy overhead; no weights adopted",
- }
- (out / "TRAINING_RESULT.json").write_text(json.dumps(result, indent=2))
-
-
-def _digest(path):
- result = hashlib.sha256()
- with Path(path).open("rb") as handle:
- for chunk in iter(lambda: handle.read(1024 * 1024), b""):
- result.update(chunk)
- return result.hexdigest()
-
-
-def load_profile_components(request, device="cpu"):
- """Strict component restoration; also used by the scheduled CPU preflight."""
- assert _digest(request["checkpoint"]) == request["checkpoint_sha256"]
- cfg = OmegaConf.load(request["config"])
- state = torch.load(
- request["checkpoint"], map_location="cpu", mmap=True, weights_only=False
- )
- assert state["global_step"] == request["checkpoint_step"]
- assert state["hyper_parameters"]["config_tree"]["model"] == OmegaConf.to_container(
- cfg.model, resolve=True
- )
- result = {}
- for name, spec in request["components"].items():
- module = instantiate(OmegaConf.select(cfg, spec["config_key"]))
- prefix = spec["state_prefix"]
- weights = {
- k[len(prefix) :]: v
- for k, v in state["state_dict"].items()
- if k.startswith(prefix)
- }
- assert weights, f"Checkpoint component prefix absent: {prefix}"
- module.load_state_dict(weights, strict=True)
- module.eval().float().to(device)
- for parameter in module.parameters():
- parameter.requires_grad_(False)
- result[name] = module
- assert int(cfg.model.num_latent_tokens) == request["latent_shape"][0]
- assert int(cfg.model.latent_dim) == request["latent_shape"][1]
- assert int(cfg.model.action_horizon) == request["action_shape"][0]
- assert int(cfg.model.action_dim) == request["action_shape"][1]
- return result
-
-
-@pytest.mark.skipif(
- not os.environ.get("ACTION_FLOW_JACOBIAN_PROFILE_REQUEST"),
- reason="Opt-in checkpoint GPU benchmark",
-)
-def test_checkpoint_endpoint_reuse_equivalence_and_timing():
- assert torch.cuda.is_available() and torch.cuda.device_count() == 1
- request_path = Path(os.environ["ACTION_FLOW_JACOBIAN_PROFILE_REQUEST"])
- request = json.loads(request_path.read_text())
- request_sha = _digest(request_path)
- out = Path(request["output_root"])
- out.mkdir(parents=True, exist_ok=True)
- torch.backends.cuda.matmul.allow_tf32 = False
- torch.backends.cudnn.allow_tf32 = False
- modules = load_profile_components(request, "cuda")
- encoder, decoder = modules["encoder"], modules["decoder"]
- assert _digest(request["targets"]) == request["targets_sha256"]
- with np.load(request["targets"], allow_pickle=False) as artifact:
- target = (
- torch.from_numpy(artifact["target"][request["target_indices"]].copy())
- .float()
- .cuda()
- )
- assert target.shape[1:] == tuple(request["action_shape"])
- torch.manual_seed(request["seed"])
- with torch.no_grad():
- clean = encoder(target)
- assert clean.shape[1:] == tuple(request["latent_shape"])
- generator = torch.Generator(device=clean.device).manual_seed(request["seed"])
- noise = torch.randn(
- clean.shape, device=clean.device, dtype=clean.dtype, generator=generator
- )
- fixed = torch.stack(
- [(1.0 - level) * clean + level * noise for level in request["noise_levels"]]
- )
- count = request["jacobian_samples"]
- assert count <= len(target)
- method = "forward_math_chunk4"
- native = owner._decoder_singular_values
- parameters = {k: v.detach().cpu().clone() for k, v in decoder.state_dict().items()}
- rng_cpu, rng_cuda = torch.get_rng_state(), torch.cuda.get_rng_state()
- flags = [p.requires_grad for p in decoder.parameters()]
-
- def measure(kind):
- with patch.object(owner, "_decoder_singular_values", wraps=native) as calls:
- clean_s = owner._decoder_singular_values(
- decoder, clean, count, jacobian_method=method
- )
- noise_s = owner._decoder_singular_values(
- decoder, noise, count, jacobian_method=method
- )
- if kind == "baseline":
- fixed_s = torch.stack(
- [
- owner._decoder_singular_values(
- decoder, value, count, jacobian_method=method
- )
- for value in fixed
- ]
- )
- else:
- fixed_s = owner._fixed_decoder_singular_values(
- decoder,
- fixed,
- clean,
- noise,
- clean_s,
- noise_s,
- count,
- jacobian_method=method,
- )
- assert calls.call_count == (
- 2 + len(fixed) if kind == "baseline" else len(fixed)
- )
- return {
- "clean": clean_s.cpu(),
- "noise": noise_s.cpu(),
- "fixed": fixed_s.cpu(),
- }, calls.call_count
-
- records, reference = [], None
- # Alternate order to reduce first-case/allocator bias; completed cases are reusable.
- order = [
- kind
- for repeat in range(request["repeats"])
- for kind in (
- ("baseline", "optimized") if repeat % 2 == 0 else ("optimized", "baseline")
- )
- ]
- for index, kind in enumerate(order):
- artifact = out / f"case-{index:02d}-{kind}.pt"
- if artifact.exists():
- payload = torch.load(artifact, map_location="cpu", weights_only=False)
- assert payload["request_sha256"] == request_sha
- else:
- torch.cuda.synchronize()
- torch.cuda.reset_peak_memory_stats()
- started = time.perf_counter()
- spectra, calls = measure(kind)
- torch.cuda.synchronize()
- payload = {
- "request_sha256": request_sha,
- "kind": kind,
- "seconds": time.perf_counter() - started,
- "full_jacobians": calls * count,
- "peak_cuda_bytes": torch.cuda.max_memory_allocated(),
- "spectra": spectra,
- }
- temporary = artifact.with_suffix(".partial")
- torch.save(payload, temporary)
- temporary.replace(artifact)
- if reference is None:
- reference = payload["spectra"]
- maximum = 0.0
- bit_exact = True
- for key, actual in payload["spectra"].items():
- torch.testing.assert_close(actual, reference[key], rtol=1e-5, atol=1e-6)
- maximum = max(maximum, float((actual - reference[key]).abs().max()))
- bit_exact = bit_exact and torch.equal(actual, reference[key])
- record = {k: v for k, v in payload.items() if k != "spectra"}
- record.update(
- case=index, max_spectrum_absolute_difference=maximum, bit_exact=bit_exact
- )
- records.append(record)
- print("JACOBIAN_PROFILE_CASE", json.dumps(record), flush=True)
- (out / "progress.json").write_text(
- json.dumps({"status": "RUNNING", "cases": records}, indent=2)
- )
-
- assert torch.equal(rng_cpu, torch.get_rng_state())
- assert torch.equal(rng_cuda, torch.cuda.get_rng_state())
- assert flags == [p.requires_grad for p in decoder.parameters()]
- assert all(
- torch.equal(v.cpu(), parameters[k]) for k, v in decoder.state_dict().items()
- )
- assert all(p.grad is None for p in decoder.parameters())
- baseline = statistics.median(
- r["seconds"] for r in records if r["kind"] == "baseline"
- )
- optimized = statistics.median(
- r["seconds"] for r in records if r["kind"] == "optimized"
- )
- result = {
- "status": "PASS",
- "request_sha256": request_sha,
- "checkpoint_step": request["checkpoint_step"],
- "weights": "raw_strict_components",
- "precision": "FP32_math_SDPA_noTF32",
- "torch": torch.__version__,
- "gpu": torch.cuda.get_device_name(),
- "cases": records,
- "baseline_median_seconds": baseline,
- "optimized_median_seconds": optimized,
- "speedup": baseline / optimized,
- "parameter_state_unchanged": True,
- "rng_unchanged": True,
- "all_singular_values_checked": True,
- }
- (out / "RESULT.json").write_text(json.dumps(result, indent=2))
- print("JACOBIAN_PROFILE_PASS", json.dumps(result), flush=True)
diff --git a/tests/test_action_flow_launcher.py b/tests/test_action_flow_launcher.py
deleted file mode 100644
index 3a622c52a2..0000000000
--- a/tests/test_action_flow_launcher.py
+++ /dev/null
@@ -1,315 +0,0 @@
-from __future__ import annotations
-
-import os
-import re
-import subprocess
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).resolve().parents[1]
-LAUNCHER = ROOT / "scripts" / "train" / "launch_action_flow_usocket.sbatch"
-DATASET_VALIDATOR = ROOT / "scripts" / "ice" / "validate_planar_dataset.py"
-
-
-def _source() -> str:
- return LAUNCHER.read_text()
-
-
-def test_launcher_has_valid_shell_syntax():
- completed = subprocess.run(
- ["bash", "-n", str(LAUNCHER)],
- check=False,
- capture_output=True,
- text=True,
- )
- assert completed.returncode == 0, completed.stderr
- assert "#SBATCH --kill-on-invalid-dep=yes" in _source()
-
-
-def test_launcher_uses_the_maintained_dataset_validator_success_token():
- validator_source = DATASET_VALIDATOR.read_text()
- match = re.search(r'"status":\s*"([A-Z_]+)"', validator_source)
- assert match is not None
- assert f'dataset_report["status"] == "{match.group(1)}"' in _source()
-
-
-def test_launcher_is_one_portable_fail_closed_contract():
- source = _source()
- assert "AF_EXPECTED_HEAD" in source
- assert "AF_EXPECTED_LAUNCHER_SHA256" in source
- assert "AF_EXPECTED_SPLIT_MANIFEST_SHA256" in source
- assert "AF_EXPECTED_CONTENT_MANIFEST_SHA256" in source
- assert "AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256" in source
- assert "AF_EXPECTED_NORM_SHA256" in source
- assert "AF_PREFLIGHT_RESULT" in source
- assert "AF_EXPECTED_PREFLIGHT_SHA256" in source
- assert "validate_action_flow_config.py" in source
- assert "capture_runtime_lock.py" in source
- assert "dataset_loader_probe.json" in source
- assert "dataset = instantiate(dataset_cfg)" in source
- assert "validate_slurm_job_contract.py" in source
- assert "check_checkpoint_storage.py" in source
- assert "--gres=gpu:1 --constraint='H100|H200'" in source
- assert '--allowed-gpu-name "NVIDIA H100 80GB HBM3"' in source
- assert '--allowed-gpu-name "NVIDIA H200"' in source
- assert "AF_EXPECTED_GPU_CONSTRAINT" in source
- assert '--expected-constraint "$AF_EXPECTED_GPU_CONSTRAINT"' in source
- assert "trainer.devices=1" in source
- assert "trainer.strategy=auto" in source
- assert "#SBATCH --requeue" in source
- assert "#SBATCH --signal=B:USR1@600" in source
- assert "/coc/" not in source
- assert "/storage/ice" not in source
- assert "/storage/project" not in source
- assert 'absolute_path AF_EXTRA_PYTHONPATH "$AF_EXTRA_PYTHONPATH"' in source
- assert (
- 'AF_EXTRA_PYTHONPATH=${AF_EXTRA_PYTHONPATH:-$(dirname "$AF_CONFIG_VALIDATOR")}'
- in source
- )
- assert (
- 'PYTHONPATH="$AF_REPO:${AF_RESUME_HELPER_DIR:+$AF_RESUME_HELPER_DIR:}'
- '$(dirname "$AF_CONFIG_VALIDATOR")${AF_EXTRA_PYTHONPATH:+:$AF_EXTRA_PYTHONPATH}"'
- in source
- )
-
-
-def test_launcher_accepts_only_the_approved_sweep_and_pins_training_semantics():
- source = _source()
- assert "sim/pusht/action_flow_bc_usocket_recon1_s42" in source
- assert "sim/pusht/action_flow_bc_usocket_recon10_s42" in source
- assert "sim/pusht/action_flow_bc_usocket_recon100_s42" in source
- assert (
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"
- in source
- )
- assert (
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42"
- in source
- )
- assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon1_s42" in source
- assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon10_s42" in source
- assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon100_s42" in source
- assert "expected_name = {" not in source
- assert "AF_FULL_MAX_STEPS=240000" in source
- assert "AF_FULL_VALIDATE_EVERY=10000" in source
- assert "AF_FULL_CHECKPOINT_EVERY=40000" in source
- assert "AF_FULL_MAX_STEPS=150000" in source
- assert "TELEMETRY_EVERY=100" in source
- assert "data.train_dataloader_params.$AF_SOURCE.batch_size=32" in source
- assert (
- "data.valid_dataloader_params.$AF_SOURCE.batch_size=$AF_VALID_BATCH_SIZE"
- in source
- )
- assert "AF_FULL_LIMIT_VAL_BATCHES=8" in source
- assert "AF_VALID_BATCH_SIZE=32" in source
- assert (
- "cfg.data.valid_dataloader_params[source].batch_size == valid_batch_size"
- in source
- )
- assert (
- "diagnostics.validation_view.per_rank_batch_size == valid_batch_size" in source
- )
- assert "cfg.trainer.val_check_interval == full_validate_every" in source
- assert (
- "cfg.callbacks.model_checkpoint.every_n_train_steps == full_checkpoint_every"
- in source
- )
- assert "ckpt_path=null" in source
- assert "norm_stats.precomputed_norm_path=$AF_NORM_STATS_PATH" in source
- assert "++run_provenance.source_commit=$AF_EXPECTED_HEAD" in source
- assert "++run_provenance.content_manifest_sha256=" in source
- assert "++run_provenance.dataset_content_aggregate_sha256=" in source
- assert "++run_provenance.normalization_sha256=$AF_EXPECTED_NORM_SHA256" in source
- assert (
- "++run_provenance.preflight_result_sha256=$AF_EXPECTED_PREFLIGHT_SHA256"
- in source
- )
- assert '"runtime_lock_sha256": digest(runtime_lock_path)' in source
- assert 'current["installed_distributions"]["canonical_sha256"]' in source
- assert "[run-preflight] PASS kind={run_kind} parameters={count}" in source
- assert "expected_count = 50_725_221" in source
- assert "50_801_685" in source
- assert "expected_count = 199_754_837" in source
- assert "AF_EXPECTED_PARAMETER_COUNT=97956100" in source
- assert "ReleasedUniteCompositeOptimizer" in source
- assert 'optimizer.muon_adjust_lr_fn == "match_rms_adamw"' in source
- assert "AF_SECOND_DATASET_DIR" in source
- assert "AF_SECOND_SPLIT_MANIFEST" in source
- assert "AF_EXPECTED_SECOND_DATASET_CONTENT_AGGREGATE_SHA256" in source
- assert "RoutedContentEncoderStage" in source
- assert "RoutedContentDecoderStage" in source
- assert "cfg.model.flow_mini_batch == 14" in source
- assert "activation_layer_map.items()" in source
-
-
-def test_smoke_runs_optimizer_validation_checkpoint_and_verifier():
- source = _source()
- assert "MAX_STEPS=2" in source
- assert "VALIDATE_EVERY=2" in source
- assert "LIMIT_VAL_BATCHES=1" in source
- assert "CHECKPOINT_EVERY=1" in source
- assert "TELEMETRY_EVERY=2" in source
- assert "verify_action_flow_training_smoke.py" in source
- assert "--planned-checkpoint-count 14" in source
- assert '--expected-constraint "$AF_EXPECTED_GPU_CONSTRAINT"' in source
- assert "--expected-reconstruction-weight" in source
- assert "--expected-config-sha256" in source
- assert (
- '--expected-config-sha256 "$(sha256 '
- '"$AF_OUTPUT_DIR/.hydra/config.yaml")"' in source
- )
- assert "--expected-normalization-sha256" in source
- assert "--expected-content-manifest-sha256" in source
- assert "--expected-dataset-content-aggregate-sha256" in source
- assert "--expected-preflight-sha256" in source
- assert "callbacks.model_checkpoint.save_last=link" in source
- assert "single-gpu,$AF_RUN_KIND" in source
-
-
-def test_full_and_smoke_share_the_requeue_and_strict_checkpoint_path():
- source = _source()
- assert 'source "$AF_REPO/scripts/train/action_flow_run_child.sh"' in source
- # The canonical runner arguments live in the shared child entry point.
- child = (ROOT / "scripts/train/action_flow_run_child.sh").read_text()
- assert "ice_requeue_runner.py" in source
- assert "validate_lightning_checkpoint.py" in source
- assert '--checkpoint-validator "$AF_CHECKPOINT_VALIDATOR"' in child
- assert "--requeue-owner runner" in child
- assert "--confirm-child-requeue-disabled" in child
- assert "runtime.slurm_requeue_owner=runner" in source
- assert "runtime.slurm_save_signal=SIGUSR2" in source
-
-
-def test_full_mode_requires_the_exact_passed_smoke_gate():
- source = _source()
- assert "required AF_SMOKE_RESULT" in source
- assert "required AF_EXPECTED_SMOKE_SHA256" in source
- assert 'payload["status"] == "PASS"' in source
- assert 'payload["experiment"] == experiment' in source
- assert 'identities["repo_head"] == head' in source
- assert 'identities["split_manifest_sha256"] == split_sha' in source
- assert 'identities["normalization_sha256"] == norm_sha' in source
- assert 'identities["content_manifest_sha256"] == content_sha' in source
- assert 'payload["preflight"]["sha256"] == preflight_sha' in source
- assert "++run_provenance.smoke_result_sha256=$AF_EXPECTED_SMOKE_SHA256" in source
- assert 'cp -- "$AF_SMOKE_RESULT" "$ATTEMPT/SMOKE_RESULT.json"' in source
-
-
-def test_second_preflight_can_reuse_the_first_train_only_normalization():
- source = _source()
- assert 'if test -z "${AF_NORM_STATS_PATH:-}"; then' in source
- assert 'test -s "$EFFECTIVE_NORM_FILE"' in source
- assert "normalization SHA-256 mismatch" in source
-
-
-def test_preflight_reuses_hashed_dataset_evidence_and_removes_logger_group():
- source = _source()
- assert "AF_CACHED_DATASET_VALIDATION" in source
- assert "AF_EXPECTED_CACHED_DATASET_VALIDATION_SHA256" in source
- assert 'payload["status"] == "DATASET_VALIDATED"' in source
- assert "'~logger'" in source
- assert "logger=null" not in source
- assert "++logger.wandb.offline=false" in source
- assert "++logger.wandb.name=$AF_WANDB_NAME" in source
- assert "++logger.wandb.resume=allow" in source
- for key in ("entity", "project", "group", "id", "tags"):
- assert f"++logger.wandb.{key}=" in source
-
-
-@pytest.mark.parametrize(
- "selector",
- [
- "pusht/action_flow_chain_manual4919_avln_80k_s42",
- "sim/pusht/parents/action_flow_usocket_h384",
- "pusht_historical/action_flow_cotrain_uc_speed_interpolation",
- ],
-)
-def test_retired_and_parent_selectors_fail_before_runtime(selector):
- source = _source()
- required = source.split("for variable in ", 1)[1].split("; do", 1)[0]
- variables = required.replace("\\", " ").split()
- env = {**os.environ, **dict.fromkeys(variables, "fixture")}
- env.update(AF_REPO=str(ROOT), AF_EXPERIMENT=selector, AF_PYTHON="/must-not-execute")
- completed = subprocess.run(
- ["bash", str(LAUNCHER)], env=env, capture_output=True, text=True
- )
- assert completed.returncode != 0
- assert "retired" in completed.stderr or "historical config" in completed.stderr
- assert "must-not-execute" not in completed.stderr
-
-
-@pytest.mark.parametrize(
- "run_kind,expected",
- [("smoke", ["evaluator.limit_val_batches=1"]), ("full", [])],
-)
-def test_raw_multiplier_smoke_bounds_actual_evaluator_override(run_kind, expected):
- from omegaconf import OmegaConf
-
- from egomimic.eval.planar_action_eval import PlanarActionEval
-
- source = _source()
- start = source.index("# Bound the evaluator too:")
- end = source.index("fi\n\n", start) + len("fi\n\n")
- block = source[start:end]
- command = (
- "BASE_OVERRIDES=(); "
- + block
- + '\nfor item in "${BASE_OVERRIDES[@]}"; do printf "%s\\n" "$item"; done'
- )
- completed = subprocess.run(
- ["bash", "-c", command],
- env={
- **os.environ,
- "AF_RUN_KIND": run_kind,
- "AF_SPEED_DIAGNOSTIC": "true",
- "LIMIT_VAL_BATCHES": "1" if run_kind == "smoke" else "8",
- },
- capture_output=True,
- text=True,
- check=True,
- )
- overrides = completed.stdout.splitlines()
- assert overrides == expected
- limit = (
- OmegaConf.from_dotlist(overrides).evaluator.limit_val_batches
- if overrides
- else 1.0
- )
- evaluator = PlanarActionEval(limit_val_batches=limit, energy_score_enabled=False)
- actual = evaluator.trainer_overrides()["limit_val_batches"]
- assert type(actual) is (int if run_kind == "smoke" else float)
- assert actual == (1 if run_kind == "smoke" else 1.0)
-
-
-@pytest.mark.parametrize(
- "mode,step,message",
- [
- ("schema", "0", "typed preflight"),
- ("run", "extern", "numeric scheduled step"),
- ("run", "0", "AF_JOINT_PHASE_INPUTS"),
- ],
-)
-def test_joint_canonical_route_rejects_missing_operational_identity_before_runtime(
- mode, step, message
-):
- required = _source().split("for variable in ", 1)[1].split("; do", 1)[0]
- variables = required.replace("\\", " ").split()
- env = {key: value for key, value in os.environ.items() if not key.startswith("AF_")}
- env.update(dict.fromkeys(variables, "fixture"))
- env.update(
- AF_REPO=str(ROOT),
- AF_EXPERIMENT="sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- AF_PYTHON="/must-not-execute",
- AF_LAUNCH_MODE=mode,
- AF_RUN_KIND="smoke",
- AF_EXPECTED_MEMORY="128G",
- AF_EXPECTED_TIME_LIMIT="1:00:00",
- SLURM_STEP_ID=step,
- )
- result = subprocess.run(
- ["bash", str(LAUNCHER)], env=env, capture_output=True, text=True
- )
- assert result.returncode != 0
- assert message in result.stderr
- assert "must-not-execute" not in result.stderr
diff --git a/tests/test_action_flow_manual4919_config.py b/tests/test_action_flow_manual4919_config.py
deleted file mode 100644
index 1a17a54196..0000000000
--- a/tests/test_action_flow_manual4919_config.py
+++ /dev/null
@@ -1,134 +0,0 @@
-"""Guard the corrected manual standard + obstacle-gen ChainGripper recipes."""
-
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-
-from egomimic.eval.energy_score import USOCKET_ENERGY_DISTANCE_CONFIG
-from egomimic.eval.planar_action_eval import PlanarActionEval
-from egomimic.pipeline.pushshapes import ChainGripperPointsNativeDecoder
-
-CONFIG_DIR = Path(__file__).parents[1] / "egomimic" / "hydra_configs"
-ROWS = ("parents/action_flow_cotrain_uc_manual4919",)
-
-
-@pytest.mark.parametrize("row", ROWS)
-def test_manual4919_recipe_uses_effective_split_and_cadence(monkeypatch, row):
- monkeypatch.setenv("PUSHSHAPES_USOCKET_ROOT", "/verified/usocket-clean2999")
- monkeypatch.setenv("PUSHSHAPES_CHAIN_GRIPPER_ROOT", "/verified/manual4919-view")
- with initialize_config_dir(
- version_base="1.3", config_dir=str(CONFIG_DIR.resolve())
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/pusht/{row}"],
- )
- chain = cfg.data.train_datasets.pushshapes_sim_chain_gripper
- valid = cfg.data.valid_datasets.pushshapes_sim_chain_gripper
- assert "data" not in cfg.data
- assert set(cfg.data.train_datasets) == (
- {"pushshapes_sim_chain_gripper"}
- if row.startswith("action_flow_chain_")
- else {"pushshapes_sim_u_socket", "pushshapes_sim_chain_gripper"}
- )
- for dataset in (chain, valid):
- assert dataset.resolver.folder_path == "/verified/manual4919-view"
- assert dataset.resolver.expected_episode_count == 4919
- assert dataset.expected_train_episode_count == 4870
- assert dataset.expected_valid_episode_count == 49
- assert dataset.valid_ratio == 0.01
- assert dataset.split_seed == 42
- assert cfg.run_provenance.chain_collection.standard_manual_episodes == 3000
- assert cfg.run_provenance.chain_collection.obstacle_gen_manual_episodes == 1919
- assert cfg.run_provenance.chain_collection.excluded_frozen_mixed_view
- assert cfg.trainer.max_steps == 80000
- assert cfg.trainer.val_check_interval == 15000
- assert cfg.val_at_end
- assert cfg.callbacks.model_checkpoint.every_n_train_steps == 5000
- assert cfg.callbacks.model_checkpoint.save_top_k == -1
- assert cfg.evaluator.action_flow_diagnostics.native_error is None
-
-
-@pytest.mark.parametrize("row", ROWS)
-def test_manual4919_native_evaluator_decoder_instantiates(monkeypatch, row):
- """Catch inherited Hydra kwargs before an allocated GPU smoke starts."""
- monkeypatch.setenv("PUSHSHAPES_USOCKET_ROOT", "/verified/usocket-clean2999")
- monkeypatch.setenv("PUSHSHAPES_CHAIN_GRIPPER_ROOT", "/verified/manual4919-view")
- with initialize_config_dir(
- version_base="1.3", config_dir=str(CONFIG_DIR.resolve())
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/pusht/{row}"],
- )
- decoder_cfg = (
- cfg.evaluator.native_decoder
- if row.startswith("action_flow_chain_")
- else cfg.evaluator.native_decoders.pushshapes_sim_chain_gripper
- )
- decoder = instantiate(decoder_cfg)
- assert isinstance(decoder, ChainGripperPointsNativeDecoder)
- if row.startswith("action_flow_chain_"):
- assert decoder_cfg.action_horizon == 16
- assert decoder_cfg.native_action_dim == 4
-
-
-def test_chain_native_decoder_rejects_wrong_native_shape():
- with pytest.raises(ValueError, match="native_action_dim=4"):
- ChainGripperPointsNativeDecoder(native_action_dim=6)
-
-
-def test_manual4919_cotrain_resolves_per_domain_energy_distance(monkeypatch):
- monkeypatch.setenv("PUSHSHAPES_USOCKET_ROOT", "/verified/usocket-clean2999")
- monkeypatch.setenv("PUSHSHAPES_CHAIN_GRIPPER_ROOT", "/verified/manual4919-view")
- with initialize_config_dir(
- version_base="1.3", config_dir=str(CONFIG_DIR.resolve())
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "+experiment=sim/pusht/parents/action_flow_cotrain_uc_manual4919"
- ],
- )
- assert cfg.evaluator.energy_score_distance is None
- distances = OmegaConf.to_container(
- cfg.evaluator.energy_score_distances_by_embodiment, resolve=True
- )
- assert distances == {
- "pushshapes_sim_u_socket": {
- **USOCKET_ENERGY_DISTANCE_CONFIG,
- "complete_normalized_chunk_shape": [16, 4],
- "normalized_translation_indices": [0, 1],
- },
- "pushshapes_sim_chain_gripper": None,
- }
- assert (
- OmegaConf.to_container(
- cfg.evaluator.energy_score_provenance.distance_contract, resolve=True
- )
- == distances
- )
- assert (
- OmegaConf.to_container(
- cfg.run_provenance.energy_score_contract.distance, resolve=True
- )
- == distances
- )
- assert cfg.evaluator.action_flow_diagnostics.native_error is None
- evaluator = PlanarActionEval(
- energy_score_enabled=False,
- native_decoders={
- label: instantiate(decoder)
- for label, decoder in cfg.evaluator.native_decoders.items()
- },
- semantic_blocks_by_embodiment=OmegaConf.to_container(
- cfg.evaluator.semantic_blocks_by_embodiment, resolve=True
- ),
- energy_score_distances_by_embodiment=distances,
- )
- assert set(evaluator.energy_score_distance_metadata["by_embodiment"]) == set(
- distances
- )
diff --git a/tests/test_action_flow_smoke_schedule.py b/tests/test_action_flow_smoke_schedule.py
deleted file mode 100644
index dd9d53dac6..0000000000
--- a/tests/test_action_flow_smoke_schedule.py
+++ /dev/null
@@ -1,336 +0,0 @@
-"""Regressions for the actual Phoenix 12935484/12935488 verifier failures."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from omegaconf import OmegaConf
-
-from tests.test_verify_action_flow_training_smoke import MODULE, _checkpoint_payload
-from tools.validate_action_flow_config import (
- GRAPH_METHOD,
- STOPGRAD_METHOD,
- compose_experiment,
-)
-
-# Exact six-key schedules printed by the failed final verifiers. The jobs'
-# training steps exited 0; both experiments intentionally train jointly at 0.
-OBSERVED_SCHEDULES = (
- (
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42",
- STOPGRAD_METHOD,
- {
- "joint_objective_begins_at_global_step": 0,
- "reconstruction_only_optimizer_steps": 0,
- "joint_flow_weight": 1.0,
- "joint_reconstruction_weight": 1.0,
- "joint_action_velocity_weight": 1.0,
- "schema_version": 1,
- },
- ),
- (
- "sim/pusht/action_flow_bc_usocket_graph_section_s42",
- GRAPH_METHOD,
- {
- "joint_objective_begins_at_global_step": 0,
- "reconstruction_only_optimizer_steps": 0,
- "joint_flow_weight": 1.0,
- "joint_reconstruction_weight": 0.0,
- "joint_action_velocity_weight": 1.0,
- "schema_version": 1,
- },
- ),
-)
-
-
-@pytest.mark.parametrize("experiment,method,schedule", OBSERVED_SCHEDULES)
-@pytest.mark.parametrize("restored_warmup", (0, 1))
-def test_observed_joint_schedule_and_strict_reload(
- tmp_path, monkeypatch, experiment, method, schedule, restored_warmup
-):
- config = compose_experiment(experiment, overrides=["hydra/launcher=basic"])
- checkpoint_dir = tmp_path / "checkpoints"
- checkpoint_dir.mkdir()
- payload = _checkpoint_payload()
- payload["action_flow_loss_schedule"] = schedule
- immutable = checkpoint_dir / "epoch-0-step-2.ckpt"
- torch.save(payload, immutable)
- (checkpoint_dir / "last.ckpt").symlink_to(immutable.name)
-
- class Wrapper:
- reconstruction_only_warmup_steps = restored_warmup
- encoder_e, field_v, decoder_g = object(), object(), object()
- model = SimpleNamespace(pipeline=SimpleNamespace(stages=[]))
- nets = SimpleNamespace(named_parameters=lambda **kwargs: [])
-
- @classmethod
- def load_from_checkpoint(cls, *args, **kwargs):
- assert kwargs["strict"] is True
- return cls()
-
- def parameters(self):
- return [SimpleNamespace(numel=lambda: MODULE.EXPECTED_PARAMETER_COUNT)]
-
- monkeypatch.setattr(MODULE, "ActionFlowModelWrapper", Wrapper)
- monkeypatch.setattr(MODULE, "_validate_dimensions_and_modules", lambda *args: None)
- monkeypatch.setattr(MODULE, "_validate_gradient_route_manifest", lambda *args: {})
- kwargs = dict(
- reconstruction_weight=schedule["joint_reconstruction_weight"],
- flow_weight=1.0,
- method=method,
- config=config,
- )
- if restored_warmup:
- with pytest.raises(MODULE.SmokeVerificationError, match="strict reload lost"):
- MODULE._validate_checkpoint(tmp_path, **kwargs)
- else:
- result = MODULE._validate_checkpoint(tmp_path, **kwargs)
- assert result["loss_schedule"] == schedule
- assert result["strict_checkpoint_reload"] == "passed"
-
-
-@pytest.mark.parametrize("key", tuple(OBSERVED_SCHEDULES[0][2]))
-def test_every_schedule_field_remains_checked(key):
- experiment, _, schedule = OBSERVED_SCHEDULES[0]
- changed = {**schedule, key: schedule[key] + 1}
- with pytest.raises(MODULE.SmokeVerificationError, match="unexpected.*schedule"):
- MODULE._validate_checkpoint_loss_schedule(
- changed,
- compose_experiment(experiment, overrides=["hydra/launcher=basic"]),
- reconstruction_weight=1,
- flow_weight=1,
- )
-
-
-@pytest.mark.parametrize("flow_weight", (1.0, 0.01))
-def test_actual_warmup_smoke_still_requires_configured_step_one(flow_weight):
- config = compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon10_warmup10k_s42",
- overrides=["hydra/launcher=basic"],
- )
- # Exercise both verifier weights without reviving the retired flow001 leaf.
- OmegaConf.update(config, "model.flow_weight", flow_weight)
- # The maintained launcher reduces this experiment's full 10k warmup to one
- # update for the real two-update warmup smoke.
- OmegaConf.update(config, "model.reconstruction_only_warmup_steps", 1)
- schedule = {
- "joint_objective_begins_at_global_step": 1,
- "reconstruction_only_optimizer_steps": 1,
- "joint_flow_weight": flow_weight,
- "joint_reconstruction_weight": 10.0,
- "joint_action_velocity_weight": 1.0,
- "schema_version": 1,
- }
- assert (
- MODULE._validate_checkpoint_loss_schedule(
- schedule, config, reconstruction_weight=10, flow_weight=flow_weight
- )
- == 1
- )
- schedule["joint_objective_begins_at_global_step"] = 0
- schedule["reconstruction_only_optimizer_steps"] = 0
- with pytest.raises(MODULE.SmokeVerificationError, match="unexpected.*schedule"):
- MODULE._validate_checkpoint_loss_schedule(
- schedule, config, reconstruction_weight=10, flow_weight=flow_weight
- )
-
-
-def test_schedule_cannot_be_validated_without_resolved_config():
- with pytest.raises(MODULE.SmokeVerificationError, match="exact config"):
- MODULE._validate_checkpoint_loss_schedule(
- OBSERVED_SCHEDULES[0][2], None, reconstruction_weight=1, flow_weight=1
- )
-
-
-def test_joint_checkpoint_owners_reject_codec_and_parameter_aliases():
- from egomimic.pipeline.stages_action_flow import (
- RoutedContentDecoderStage,
- RoutedContentEncoderStage,
- )
-
- def stages_with(encoders, decoders, field):
- stages = [None] * 8
- stages[4] = (
- RoutedContentEncoderStage(encoders, route_key="embodiment")
- if encoders
- else SimpleNamespace(encoder={})
- )
- stages[6] = SimpleNamespace(field=field)
- stages[7] = RoutedContentDecoderStage(decoders, route_key="embodiment")
- return stages
-
- names = ("libero_panda", "libero_panda_sharpa")
- encoders = {name: torch.nn.Linear(2, 2) for name in names}
- decoders = {name: torch.nn.Linear(2, 2) for name in names}
- field = torch.nn.Linear(2, 2)
- assert (
- len(MODULE._joint_checkpoint_owners(stages_with(encoders, decoders, field)))
- == 5
- )
- decoders[names[0]].weight = encoders[names[0]].weight
- with pytest.raises(MODULE.SmokeVerificationError, match="share parameters"):
- MODULE._joint_checkpoint_owners(stages_with(encoders, decoders, field))
- decoders[names[0]] = encoders[names[0]]
- with pytest.raises(MODULE.SmokeVerificationError, match="aliased"):
- MODULE._joint_checkpoint_owners(stages_with(encoders, decoders, field))
- with pytest.raises(MODULE.SmokeVerificationError, match="coverage"):
- MODULE._joint_checkpoint_owners(stages_with({}, decoders, field))
-
-
-def test_joint_smoke_uses_unique_immutable_step_eight_without_rolling_last(tmp_path):
- checkpoint = tmp_path / "epoch-0000-step-000000008.ckpt"
- with pytest.raises(MODULE.SmokeVerificationError, match="one immutable"):
- MODULE._joint_smoke_checkpoint_path(tmp_path)
- checkpoint.write_bytes(b"path-contract-only")
- assert MODULE._joint_smoke_checkpoint_path(tmp_path) == checkpoint
- duplicate = tmp_path / "epoch-0001-step-000000008.ckpt"
- duplicate.write_bytes(b"path-contract-only")
- with pytest.raises(MODULE.SmokeVerificationError, match="one immutable"):
- MODULE._joint_smoke_checkpoint_path(tmp_path)
- duplicate.unlink()
- (tmp_path / "last.ckpt").symlink_to(checkpoint.name)
- with pytest.raises(MODULE.SmokeVerificationError, match="rolling"):
- MODULE._joint_smoke_checkpoint_path(tmp_path)
-
-
-def test_joint_history_requires_both_source_rows_at_step_seven():
- from tests.test_verify_action_flow_training_smoke import _history_row
-
- row = _history_row()
- for key, value in list(row.items()):
- if MODULE.SOURCE_LABEL in key:
- for label in ("libero_panda", "libero_panda_sharpa"):
- row[key.replace(MODULE.SOURCE_LABEL, label)] = value
- row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 2.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 28.0
- for prefix in (
- "GradientCosine",
- "GradientCosineDefined",
- "GradientIntersectionParameterCount",
- ):
- row[f"Train/ActionFlow/{prefix}/FM__Reconstruction"] = 0.0
- row["Optimizer/LR/AdamW"] = row["Optimizer/LR/Muon"] = 2.5e-8
- options = dict(
- method=MODULE.STOPGRAD_UNITE_METHOD,
- source_labels=("libero_panda", "libero_panda_sharpa"),
- minimum_step=7,
- )
- assert MODULE._validate_history({7: row}, **options)["valid_step"] == 7
- with pytest.raises(MODULE.SmokeVerificationError, match="gradient telemetry"):
- MODULE._validate_history({1: row}, **options)
- del row["Train/MSE/libero_panda_sharpa"]
- with pytest.raises(MODULE.SmokeVerificationError, match="gradient telemetry"):
- MODULE._validate_history({7: row}, **options)
-
-
-def test_joint_ema_reload_updates_every_alias_and_rejects_wrong_updates():
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_ema_reload,
- )
-
- model = torch.nn.Module()
- model.model = torch.nn.Sequential(torch.nn.Linear(3, 2))
- model.nets = model.model
- # The native normalizer rebuilds its ParameterDict during strict reload.
- from egomimic.models.oat.model.common.dict_of_tensor_mixin import DictOfTensorMixin
-
- model.model.add_module(
- "normalizer",
- DictOfTensorMixin(
- torch.nn.ParameterDict(
- {"scale": torch.nn.Parameter(torch.ones(2), requires_grad=False)}
- )
- ),
- )
- settings = dict(
- decay=0.9978,
- use_warmup=False,
- update_after_step=0,
- inv_gamma=1.0,
- power=0.75,
- min_decay=0.0,
- )
- cfg = dict(callbacks=dict(ema=settings))
- ema = {
- name: torch.full_like(parameter, 7)
- for name, parameter in model.named_parameters()
- }
- payload = dict(
- global_step=8, ema_num_updates=8, ema_config=settings, ema_state_dict=ema
- )
- result = validate_joint_ema_reload(model, payload, cfg)
- assert result["state_parameter_aliases"] > result["canonical_parameters"]
- assert all(
- torch.equal(value, torch.full_like(value, 7))
- for value in model.state_dict().values()
- )
- payload["ema_num_updates"] = 7
- with pytest.raises(ValueError, match="update count"):
- validate_joint_ema_reload(model, payload, cfg)
- payload["ema_num_updates"] = 8
- ema[next(iter(ema))].fill_(float("nan"))
- with pytest.raises(ValueError, match="EMA tensor differs"):
- validate_joint_ema_reload(model, payload, cfg)
-
-
-def test_joint_smoke_cli_dispatch_is_explicit_and_requires_all_hashes(
- tmp_path, monkeypatch
-):
- from scripts.train import verify_libero_native_action_flow_smoke as native
-
- experiment = "sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42"
- args = MODULE._parser().parse_args(
- [str(tmp_path), "--experiment", experiment, "--expected-head", "a" * 40]
- )
- assert args.experiment == experiment
- # Missing identities must fail before source inspection, checkpoint loads or API reads.
- with pytest.raises(ValueError, match="every explicit expected SHA"):
- MODULE.verify_smoke(tmp_path, experiment=experiment, expected_head="a" * 40)
- captured = {}
-
- def routed(**kwargs):
- captured.update(kwargs)
- return {"dispatch_only": True}
-
- monkeypatch.setattr(native, "verify_joint_smoke", routed)
- assert MODULE.verify_smoke(
- tmp_path,
- experiment=experiment,
- expected_head="a" * 40,
- expected_config_sha256="b" * 64,
- expected_split_sha256="c" * 64,
- expected_normalization_sha256="d" * 64,
- expected_preflight_sha256="e" * 64,
- ) == {"dispatch_only": True}
- assert captured["expected_preflight_sha256"] == "e" * 64
- assert captured["run_dir"] == tmp_path
-
-
-@pytest.mark.parametrize("transport", [None, "slurm-steps", "trainer-pidfd"])
-def test_joint_shared_child_rejects_unbound_signal_transport(transport):
- import os
- import subprocess
- from pathlib import Path
-
- env = {
- key: value
- for key, value in os.environ.items()
- if not key.startswith(("AF_", "ICE_"))
- }
- env["AF_EXPERIMENT"] = (
- "sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42"
- )
- if transport is not None:
- env["AF_CHECKPOINT_TRANSPORT"] = transport
- child = (
- Path(__file__).resolve().parents[1] / "scripts/train/action_flow_run_child.sh"
- )
- result = subprocess.run(
- ["bash", str(child)], env=env, capture_output=True, text=True
- )
- assert result.returncode == 2
- assert (
- "requires audited trainer-pidfd"
- if transport != "trainer-pidfd"
- else "PIDfd binding missing"
- ) in result.stderr
diff --git a/tests/test_alignment_parity.py b/tests/test_alignment_parity.py
deleted file mode 100644
index b76a4be95c..0000000000
--- a/tests/test_alignment_parity.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""Numerical parity with immutable source methods and configured graph training."""
-
-from copy import deepcopy
-from functools import partial
-
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.utils import instantiate
-
-from egomimic.models.cores.hpt_transformer import MultiheadAttention, SimpleTransformer
-from egomimic.pipeline.core import Pipeline
-from egomimic.pipeline.stages_alignment import (
- RepresentationAlignment,
- ScheduledLossReduction,
-)
-from egomimic.pipeline.stages_hpt import HPTTrunkStage
-from tests.test_retained_recipe_steps import CONFIGS, small_cpu_model
-
-pytest.importorskip("geomloss")
-SoftDTWLossPyTorch = pytest.importorskip("tslearn.metrics").SoftDTWLossPyTorch
-LegacyAlignment = pytest.importorskip(
- "tests.fixtures.legacy_alignment_reference"
-).LegacyAlignment
-
-
-@pytest.mark.parametrize("supervision", [None, "mse", "soft_dtw"])
-def test_sinkhorn_loss_and_gradients_match_source(supervision):
- torch.manual_seed(121)
- left = torch.randn(3, 4, 8, requires_grad=True)
- right = torch.randn(3, 4, 8, requires_grad=True)
- actions = torch.randn(3, 5, 6)
- other_actions = actions.flip(0) + 0.01
- reference = LegacyAlignment()
- reference.device = torch.device("cpu")
- reference.ot_6dof = False
- reference.use_dtw = supervision == "soft_dtw"
- reference.dtw = SoftDTWLossPyTorch(gamma=0.1)
- expected, expected_distance = reference.compute_ot(
- left, right, actions, other_actions, supervision is not None, 0.5
- )
- expected_grads = torch.autograd.grad(expected, (left, right))
- stage = RepresentationAlignment(
- "left",
- "right",
- supervision=supervision,
- cost_layout="legacy_broadcast",
- left_actions="actions",
- right_actions="other_actions",
- left_action_indices=[0, 1, 2],
- right_action_indices=[0, 1, 2],
- )
- result = stage(
- {
- "left": left,
- "right": right,
- "actions": actions,
- "other_actions": other_actions,
- }
- )
- torch.testing.assert_close(result["alignment/loss"], expected)
- torch.testing.assert_close(result["log/feature_distance"], expected_distance)
- for actual, expected_gradient in zip(
- torch.autograd.grad(result["alignment/loss"], (left, right)), expected_grads
- ):
- torch.testing.assert_close(actual, expected_gradient)
-
-
-def test_freeze_boundary_preserves_alignment_gradients():
- torch.manual_seed(42)
- trunk = SimpleTransformer(
- partial(MultiheadAttention, embed_dim=8, num_heads=2, batch_first=True),
- embed_dim=8,
- num_blocks=3,
- )
- stage = HPTTrunkStage(
- trunk,
- 8,
- action_token_count=2,
- squeeze_action_token=False,
- representation_block=1,
- detach_condition_before_block=2,
- )
- result = stage({"hpt/tokens": torch.randn(2, 3, 8)})
- result["condition"].square().sum().backward(retain_graph=True)
- assert trunk.blocks[0].attn.in_proj_weight.grad is None
- assert trunk.blocks[1].attn.in_proj_weight.grad is None
- assert trunk.blocks[2].attn.in_proj_weight.grad.abs().sum() > 0
- stage.zero_grad()
- result["hpt/representation"].square().sum().backward()
- assert trunk.blocks[0].attn.in_proj_weight.grad.abs().sum() > 0
- assert trunk.blocks[1].attn.in_proj_weight.grad.abs().sum() > 0
- assert trunk.blocks[2].attn.in_proj_weight.grad is None
- # The resumed path is the source implementation, with unchanged weights.
- _, blocks = trunk(torch.randn(2, 5, 8))
- reference = LegacyAlignment()
- reference.trunk = {"trunk": trunk}
- reference.postprocess_tokens = lambda value: value
- torch.testing.assert_close(
- trunk.resume_from_block(blocks, 2), reference.resume_from_depth(blocks, 2)
- )
- assert "hpt/representation" not in stage.execute(
- {"hpt/tokens": torch.randn(2, 3, 8)}, mode="inference"
- )
-
-
-def test_loss_schedule_and_resume_use_training_batches():
- def graph():
- return Pipeline(
- [
- ScheduledLossReduction(
- {
- "source_losses": {"start_batch": 2},
- "alignment": {"weight": 0.7, "start_batch": 3},
- },
- denominator_key="source_count",
- )
- ]
- )
-
- reducer = graph()
- batch = {
- "source_losses": torch.tensor([2.0, 4.0]),
- "alignment": torch.tensor(10.0),
- "source_count": 2,
- }
- assert reducer(batch)["loss/total"] == 0
- assert reducer(batch)["loss/total"] == 3
- restored = graph()
- restored.load_state_dict(deepcopy(reducer.state_dict()), strict=True)
- assert restored(batch)["loss/total"] == pytest.approx(6.5)
- assert restored(batch)["log/training_batches"] == 4
-
-
-def test_egobridge_two_source_optimizer_and_independent_feature_passes():
- with initialize_config_dir(version_base=None, config_dir=str(CONFIGS)):
- config = compose(
- config_name="train_zarr_cartesian", overrides=["model=egobridge"]
- )
- cpu = small_cpu_model(config, "hpt")
- cpu.model.pipeline.stages[1].representation_block = 0
- cpu.model.pipeline.loss_pipeline.stages[0].supervision = "mse"
- model = instantiate(cpu.model.pipeline, device="cpu")
- calls = []
- hook = model.pipeline.stage_by_id("stems").register_forward_hook(
- lambda *args: calls.append(1)
- )
- batch = {}
- for name, identity, width in [("eva_bimanual", 6, 14), ("human_bimanual", 3, 12)]:
- batch[name] = {
- "embodiment": torch.tensor([identity, identity]),
- "actions_cartesian": torch.randn(2, 100, width),
- "observations.state.ee_pose": torch.randn(2, width),
- **{
- f"observations.images.{key}": torch.rand(2, 3, 32, 32)
- for key in ("front_img_1", "left_wrist_img", "right_wrist_img")
- },
- }
- optimizer = torch.optim.Adam(model.nets.parameters(), lr=1e-4)
- for _ in range(2):
- optimizer.zero_grad()
- predictions = model.forward_training(batch)
- losses = model.compute_losses(predictions, batch)
- expected = (
- losses["source_0_loss"]
- + losses["source_1_loss"]
- + 0.7 * losses["log_alignment_loss"]
- ) / 2
- torch.testing.assert_close(losses["loss"], expected)
- losses["loss"].backward()
- assert model.pipeline.stage_by_id("trunk").action_token.grad.abs().sum() > 0
- optimizer.step()
- hook.remove()
- assert (
- len(calls) == 8
- ), "Two sources, two steps, independent BC and OT feature passes"
- clone = instantiate(cpu.model.pipeline, device="cpu")
- clone.nets.load_state_dict(deepcopy(model.nets.state_dict()), strict=True)
- assert clone.nets["loss_pipeline"].stages[1].training_batches == 2
- clone.nets.eval()
- assert clone.forward_eval({"human_bimanual": batch["human_bimanual"]})[
- "human_bimanual"
- ]["pred_action"].shape == (2, 100, 12)
-
-
-def test_alignment_rejects_unmatched_batches():
- stage = RepresentationAlignment("left", "right")
- with pytest.raises(ValueError, match="matching"):
- stage({"left": torch.zeros(2, 3), "right": torch.zeros(3, 3)})
diff --git a/tests/test_arc_tsne.py b/tests/test_arc_tsne.py
deleted file mode 100644
index 89092da7b9..0000000000
--- a/tests/test_arc_tsne.py
+++ /dev/null
@@ -1,295 +0,0 @@
-import json
-
-import numpy as np
-import pytest
-
-from egomimic.scripts.data_visualization.arc_tsne import (
- TokenizeBimanualArcLengthCartesian,
- arc_features,
- baseline_features,
- baseline_leaf,
- collect,
- load_data,
- project,
- write_report,
-)
-
-
-def raw_chunk(frame=0):
- t = np.arange(201) / 30
- raw = np.zeros((len(t), 14))
- raw[:, 0] = t * (0.1 + frame * 0.003)
- raw[:, 7] = t * (0.07 + frame * 0.001)
- raw[:, 1] = np.sin(t + frame) * 0.04
- raw[:, 10] = t * (0.03 + frame * 0.001)
- return raw
-
-
-def codec(**kwargs):
- return TokenizeBimanualArcLengthCartesian(
- min_distance_unit=0.4,
- rotation_distance_unit=0.3,
- resampled_vector_length=28,
- **kwargs,
- )
-
-
-@pytest.mark.parametrize("mode", ["velocity", "duration"])
-@pytest.mark.parametrize("streams", ["full_separation", "translation_rotation"])
-def test_clock_removal_and_layout_parity(mode, streams):
- c = codec(clock_mode=mode, streams=streams)
- token = c.transform({"actions_cartesian": raw_chunk()})["actions_cartesian"]
- geometry = arc_features(token, c, False)
- assert geometry.shape == (28 * 14,)
- mutated = token.copy()
- mutated[:, 14:] += 55
- np.testing.assert_array_equal(geometry, arc_features(mutated, c, False))
- assert not np.array_equal(arc_features(token, c), arc_features(mutated, c))
-
-
-class Leaf:
- def __init__(self):
- self.transform = [codec(clock_mode="velocity")]
- self.metadata = {"task": "fold "}
-
- def __len__(self):
- return 18
-
- def __getitem__(self, frame):
- data = {"actions_cartesian": raw_chunk(frame), "frame_index": frame}
- for transform in self.transform:
- data = transform.transform(data)
- return data
-
-
-def test_paired_sampling_and_report(tmp_path):
- leaf = Leaf()
- original = leaf.transform.copy()
- first, records, codecs = collect({"one": leaf}, samples_per_episode=8)
- second, records2, _ = collect({"one": leaf}, samples_per_episode=8)
- assert records == records2
- assert leaf.transform == original
- assert codecs[0]["tokenizer_config"]["min_distance_unit"] == 0.4
- for k in first:
- np.testing.assert_array_equal(first[k], second[k])
- for i, record in enumerate(records):
- np.testing.assert_array_equal(
- first["baseline"][i], raw_chunk(record["frame"])[:100].ravel()
- )
- xy, diag = project(first)
- xy2, _ = project(first)
- assert xy == xy2
- assert all(np.asarray(v).shape == (8, 2) for v in xy.values())
- output = tmp_path / "report"
- write_report(output, first, records, xy, {"projections": diag})
- html = (output / "index.html").read_text()
- assert "fold " not in html
- assert "__REPORT_DATA__" not in html
- assert "https://" not in html
- assert json.loads((output / "report.json").read_text())["records"] == records
- assert set(np.load(output / "features.npz").files) == set(first)
- with pytest.raises(FileExistsError):
- write_report(output, first, records, xy, {})
-
-
-def test_baseline_padding_and_bad_values():
- raw = raw_chunk()[:2]
- padded = baseline_features(raw, 4).reshape(4, 14)
- np.testing.assert_array_equal(padded[2:], np.repeat(raw[-1:], 2, axis=0))
- with pytest.raises(ValueError, match="identical"):
- project({"bad": np.ones((4, 12))})
- with pytest.raises(ValueError, match="Nonfinite"):
- baseline_features(np.full((3, 14), np.nan))
-
-
-def test_baseline_window_does_not_inherit_distance_cutoff():
- leaf = Leaf()
- leaf.key_map = {
- "actions": {"horizon": {"distance": 0.4}, "key_type": "action_keys"},
- "obs": {"horizon": 1, "key_type": "proprio_keys"},
- }
- time = baseline_leaf(leaf, leaf.transform, 0, 100)
- assert time.key_map["actions"]["horizon"] == 100
- assert time.key_map["obs"]["horizon"] == 1
- assert leaf.key_map["actions"]["horizon"] == {"distance": 0.4}
- assert time.transform == []
- assert len(leaf.transform) == 1
-
-
-@pytest.mark.parametrize("embodiment", ["human_bimanual", "yam_bimanual"])
-def test_rl2_inventory_configs_preserve_native_source(embodiment):
- from pathlib import Path
-
- from hydra.utils import instantiate
- from omegaconf import OmegaConf
-
- from egomimic.rldb.zarr.action_chunk_transforms import (
- InterpolateLinear,
- InterpolatePose,
- )
-
- path = (
- Path(__file__).parents[1]
- / "egomimic/hydra_configs/visualization"
- / f"rl2_organize_{embodiment}.yaml"
- )
- cfg = OmegaConf.load(path)
- assert cfg.data.valid_ratio == 0.05
- key_map = instantiate(cfg.data.key_map)
- assert not any(v.get("key_type") == "camera_keys" for v in key_map.values())
- transforms = instantiate(cfg.data.transform_list)
- for transform in transforms:
- if isinstance(transform, (InterpolatePose, InterpolateLinear)):
- assert transform.new_chunk_length is None
- tokenizer = next(
- t for t in transforms if isinstance(t, TokenizeBimanualArcLengthCartesian)
- )
- assert tokenizer.M == 100
- assert tokenizer.clock_mode == "duration"
- assert tokenizer.streams == "translation_rotation"
- assert tokenizer.tokenizer.config.min_distance_unit == 0.42
- assert "elmo" in str(cfg.data.filters.filter_lambdas[0])
-
-
-def test_clock_layout_features_separate_waypoints_from_clocks():
- from egomimic.rldb.zarr.arc_length_tokenizer import CLOCK_COLUMNS
- from egomimic.scripts.data_visualization.arc_tsne import arc_features
-
- class Codec:
- M, clock_mode, streams = 4, "duration", "translation_rotation"
-
- token = np.arange(4 * 18, dtype=float).reshape(4, 18) + 1
- with_clock = arc_features(token, Codec, True).reshape(-1, 14)
- without = arc_features(token, Codec, False).reshape(-1, 14)
- assert np.array_equal(without, token[:, :14])
- assert np.array_equal(with_clock[4:, list(CLOCK_COLUMNS)], token[:, 14:])
-
-
-def test_joint_fit_drops_gripper_and_scores_mixing():
- from egomimic.scripts.data_visualization.arc_tsne import (
- drop_gripper,
- mixing,
- project,
- )
-
- rng = np.random.default_rng(0)
- human = rng.normal(size=(40, 2 * 14))
- robot = rng.normal(size=(40, 2 * 14))
- human.reshape(40, 2, 14)[:, :, [6, 13]] = 0.0
- robot.reshape(40, 2, 14)[:, :, [6, 13]] = 5.0 + rng.normal(size=(40, 2, 2))
- dropped = drop_gripper({"baseline": np.concatenate([human, robot])})["baseline"]
- assert dropped.shape == (80, 2 * 12)
- groups = ["human"] * 40 + ["robot"] * 40
- _, mixed = mixing(dropped, groups)
- _, split = mixing(np.concatenate([human, robot]), groups)
- assert mixed > 0.7 and split < 0.2
- xy, diagnostics, neighbors = project({"baseline": dropped}, groups=groups)
- assert len(xy["baseline"]) == 80 and len(neighbors["baseline"][0]) == 8
- assert diagnostics["baseline"]["mixing"] > 0.7
-
-
-def test_inventory_sampling_uses_full_split():
- from egomimic.scripts.data_visualization.arc_tsne_inventory import (
- select_episode_ids,
- )
-
- names = [f"episode_{i:03}" for i in range(100)]
- sampled, train, valid = select_episode_ids(names, 0.05, 42, 42, 12)
- assert len(valid) == 5
- assert len(train) == 95
- assert len(sampled) == 12
- assert set(sampled) <= set(train)
- assert not set(sampled) & set(valid)
- assert (sampled, train, valid) == select_episode_ids(names[::-1], 0.05, 42, 42, 12)
-
-
-def test_missing_codec_and_baseline_only():
- leaf = Leaf()
- leaf.transform = []
- with pytest.raises(ValueError, match="requires a data config"):
- collect({"one": leaf})
- values, _, _ = collect({"one": leaf}, representations="baseline")
- assert set(values) == {"baseline"}
-
-
-def test_fallback_is_rejected_and_transforms_restored():
- class BadLeaf(Leaf):
- def __getitem__(self, frame):
- return super().__getitem__(frame + 1)
-
- leaf = BadLeaf()
- original = leaf.transform.copy()
- with pytest.raises(ValueError, match="substituted"):
- collect({"one": leaf})
- assert leaf.transform == original
-
-
-def test_real_zarr_config_pipeline(tmp_path):
- import zarr
- from omegaconf import OmegaConf
-
- root = tmp_path / "episodes"
- for i in range(4):
- store = zarr.open_group(str(root / f"episode_{i}.zarr"), mode="w")
- raw = raw_chunk(i)
- store.create_array("actions", data=raw)
- store.attrs.update(
- total_frames=len(raw),
- embodiment="yam_bimanual",
- task="test",
- features={"actions": {"dtype": "float64"}},
- )
- target = "egomimic.rldb.zarr."
- config = {
- "train_datasets": {
- "yam": {
- "mode": "train",
- "valid_ratio": 0.05,
- "split_seed": 42,
- "resolver": {
- "_target_": target + "zarr_dataset_multi.LocalEpisodeResolver",
- "folder_path": str(root),
- "key_map": {
- "actions_cartesian": {"zarr_key": "actions", "horizon": 150}
- },
- "transform_list": [
- {
- "_target_": target
- + "arc_length_tokenizer.TokenizeBimanualArcLengthCartesian",
- "rotation_distance_unit": 0.3,
- "clock_mode": "velocity",
- "resampled_vector_length": 28,
- }
- ],
- },
- }
- }
- }
- config_path = tmp_path / "data.yaml"
- OmegaConf.save(OmegaConf.create(config), config_path)
- leaves, provenance = load_data(config_path, [], None, "train", False)
- assert len(leaves) == 3 # 5% split retains one validation episode.
- assert provenance["config"]["valid_ratio"] == 0.05
- values, records, codecs = collect(leaves, samples_per_episode=3)
- assert len(records) == 9
- assert values["arc_clock"].shape == (9, 28 * 28)
- assert values["arc_no_clock"].shape == (9, 28 * 14)
- assert codecs[0]["streams"] == "full_separation"
- projections, diagnostics = project(values)
- write_report(
- tmp_path / "report",
- values,
- records,
- projections,
- {**provenance, "projections": diagnostics},
- )
- # The same file is valid under a full resolved training-config data section.
- OmegaConf.save(OmegaConf.create({"data": config}), config_path)
- full_leaves, _ = load_data(config_path, [], None, "train", False)
- assert set(full_leaves) == set(leaves)
- config["train_datasets"]["yam"]["resolver"]["_target_"] = (
- target + "zarr_dataset_multi.S3EpisodeResolver"
- )
- OmegaConf.save(OmegaConf.create(config), config_path)
- with pytest.raises(ValueError, match="Non-local resolver refused"):
- load_data(config_path, [], None, "train", False)
diff --git a/tests/test_batch_signal_relay.py b/tests/test_batch_signal_relay.py
deleted file mode 100644
index 35870aed7f..0000000000
--- a/tests/test_batch_signal_relay.py
+++ /dev/null
@@ -1,161 +0,0 @@
-"""Real bash/child regression for Slurm batch-shell-only signal delivery."""
-
-from __future__ import annotations
-
-import json
-import os
-import shlex
-import signal
-import subprocess
-import sys
-import time
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).parents[1]
-HELPER = ROOT / "scripts/ice/relay_batch_signals.sh"
-CHILD = r"""
-import json, os, pathlib, signal, sys, time
-root = pathlib.Path(sys.argv[1])
-root.joinpath("started").write_text(str(os.getpid()))
-if sys.argv[2] == "early":
- raise SystemExit(int(sys.argv[3]))
-time.sleep(float(sys.argv[4]))
-def handle(received, frame):
- root.joinpath("received").write_text(signal.Signals(received).name)
- if received == signal.SIGUSR1 and sys.argv[2] != "boundary":
- root.joinpath("COMPLETE.json").write_text("complete")
- time.sleep(.05)
- raise SystemExit(0 if received == signal.SIGUSR1 else 128 + received)
-for received in (signal.SIGUSR1, signal.SIGTERM, signal.SIGINT):
- signal.signal(received, handle)
-pathlib.Path(os.environ["ICE_RUNNER_SIGNAL_READY_FILE"]).write_text(str(os.getpid()))
-if sys.argv[2] == "normal":
- raise SystemExit(int(sys.argv[3]))
-time.sleep(5)
-"""
-
-
-def wait_file(path):
- deadline = time.monotonic() + 3
- while time.monotonic() < deadline:
- if path.exists() and path.stat().st_size:
- return
- time.sleep(0.01)
- raise AssertionError(f"child readiness timeout: {path}")
-
-
-def launch(root, mode="signal", exit_code=0, delay=0):
- argv = [sys.executable, "-c", CHILD, str(root), mode, str(exit_code), str(delay)]
- command = (
- f"set -euo pipefail; source {shlex.quote(str(HELPER))}; "
- f"ice_run_with_signal_relay {shlex.quote(str(root / 'ready'))} "
- f"{shlex.join(argv)}; "
- f'if test "$ICE_BATCH_BOUNDARY_FORWARDED" = 1 && test ! -s {shlex.quote(str(root / "COMPLETE.json"))}; '
- "then printf 'REQUEUE_ACCEPTED\\n'; exit 0; fi; "
- "printf 'POST_RUNNER_VERIFIER\\n'"
- )
- return subprocess.Popen(
- ["bash", "-c", command],
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- )
-
-
-def assert_reaped(root):
- pid = int((root / "started").read_text())
- with pytest.raises(ProcessLookupError):
- os.kill(pid, 0)
-
-
-@pytest.mark.parametrize("received", [signal.SIGUSR1, signal.SIGTERM, signal.SIGINT])
-@pytest.mark.parametrize("during_startup", [False, True])
-def test_exact_signal_reaches_child_and_wait_preserves_status(
- tmp_path, received, during_startup
-):
- child = launch(tmp_path, delay=0.2 if during_startup else 0)
- try:
- wait_file(tmp_path / ("started" if during_startup else "ready"))
- child.send_signal(received)
- output, errors = child.communicate(timeout=4)
- expected = 0 if received == signal.SIGUSR1 else 128 + received
- assert child.returncode == expected, (output, errors)
- assert (tmp_path / "received").read_text() == received.name
- assert ("POST_RUNNER_VERIFIER" in output) == (expected == 0)
- assert_reaped(tmp_path)
- finally:
- if child.poll() is None:
- child.kill()
- child.wait(timeout=3)
-
-
-@pytest.mark.parametrize("mode", ["normal", "early"])
-@pytest.mark.parametrize("exit_code", [0, 7])
-def test_early_and_normal_exit_status_is_not_replaced_by_wait(
- tmp_path, mode, exit_code
-):
- child = launch(tmp_path, mode=mode, exit_code=exit_code)
- output, errors = child.communicate(timeout=4)
- assert child.returncode == exit_code, (output, errors)
- assert ("POST_RUNNER_VERIFIER" in output) == (exit_code == 0)
- assert_reaped(tmp_path)
-
-
-def test_existing_readiness_marker_refuses_before_launch(tmp_path):
- (tmp_path / "ready").write_text("stale")
- child = launch(tmp_path)
- output, errors = child.communicate(timeout=3)
- assert child.returncode == 73, (output, errors)
- assert not (tmp_path / "started").exists()
- assert "POST_RUNNER_VERIFIER" not in output
-
-
-def test_successful_requeue_without_completion_does_not_run_final_verifier(tmp_path):
- child = launch(tmp_path, mode="boundary")
- wait_file(tmp_path / "ready")
- child.send_signal(signal.SIGUSR1)
- output, errors = child.communicate(timeout=4)
- assert child.returncode == 0, (output, errors)
- assert "REQUEUE_ACCEPTED" in output
- assert "POST_RUNNER_VERIFIER" not in output
- assert_reaped(tmp_path)
-
-
-def test_actual_runner_publishes_readiness_before_command_validation(tmp_path):
- ready = tmp_path / "ready.json"
- environment = {**os.environ, "ICE_RUNNER_SIGNAL_READY_FILE": str(ready)}
- result = subprocess.run(
- [
- sys.executable,
- str(ROOT / "scripts/ice/ice_requeue_runner.py"),
- "--state-dir",
- str(tmp_path / "state"),
- "--checkpoint-glob",
- str(tmp_path / "*.ckpt"),
- "--checkpoint-signal",
- "USR2",
- "--checkpoint-forwarding",
- "slurm-steps",
- "--requeue-owner",
- "runner",
- ],
- env=environment,
- capture_output=True,
- text=True,
- timeout=3,
- )
- assert result.returncode != 0
- payload = json.loads(ready.read_text())
- assert payload["signal_handlers_ready"] is True
- assert payload["pid"] > 0
- launcher = (ROOT / "scripts/train/launch_action_flow_usocket.sbatch").read_text()
- assert (
- "AF_SIGNAL_RELAY=${AF_SIGNAL_RELAY:-$AF_TOOL_DIR/relay_batch_signals.sh}"
- in launcher
- )
- assert 'source "$AF_REPO/scripts/train/action_flow_run_child.sh"' in launcher
- child = (ROOT / "scripts/train/action_flow_run_child.sh").read_text()
- assert 'source "$AF_SIGNAL_RELAY"' in child
- assert 'ice_run_with_signal_relay "$ATTEMPT/runner-signal-ready.json"' in child
diff --git a/tests/test_c12_compatibility.py b/tests/test_c12_compatibility.py
deleted file mode 100644
index 6f1b0e9cfb..0000000000
--- a/tests/test_c12_compatibility.py
+++ /dev/null
@@ -1,224 +0,0 @@
-"""Historical behavior contracts, not a claim of full training equivalence."""
-
-import ast
-import os
-import random
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from egomimic.pipeline.algo import PipelineAlgo
-from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper
-from egomimic.rldb.zarr import zarr_dataset_multi as dm
-from tests.test_bounds_check_gate import _dataset
-
-
-def module(mode):
- return MultiDataModuleWrapper({}, {}, {}, {}, compatibility_mode=mode)
-
-
-def test_invalid_mode_fails_closed():
- for make in (
- module,
- lambda mode: PipelineAlgo([], device="cpu", compatibility_mode=mode),
- lambda mode: dm.MultiDataset(state={}, compatibility_mode=mode),
- ):
- with pytest.raises(ValueError, match="compatibility_mode"):
- make("typo")
-
-
-def test_bounds_false_is_explicitly_overridden_only_in_legacy_mode():
- ds = dm.MultiDataset(state={}, bounds_check=False, compatibility_mode="legacy_c12")
- assert ds.bounds_check is True
- assert ds.bounds_semantics == "legacy_full_vector"
- assert ds.fallback_policy == "legacy_random"
- assert dm.MultiDataset(state={}, bounds_check=False).bounds_check is False
-
-
-@pytest.mark.parametrize("bad_value", [1e6, float("nan"), float("inf"), -1e-7])
-def test_legacy_filter_rejects_nonfinite_and_strict_outliers(bad_value):
- ds, leaf = _dataset(bounds_check=False, bad_value=bad_value)
- # Constructor-derived settings, transplanted onto a tiny normalized fixture.
- settings = dm.MultiDataset(
- state={}, bounds_check=False, compatibility_mode="legacy_c12"
- )
- for key in (
- "bounds_check",
- "bounds_semantics",
- "fallback_policy",
- "compatibility_mode",
- ):
- setattr(ds, key, getattr(settings, key))
- ds[1]
- assert leaf.served == [1, 0]
-
-
-def test_loader_matches_default_shuffle_worker_seed_and_global_rng():
- legacy = module("legacy_c12")
- for workers in (0, 2):
- results = []
- for generator in (None, legacy._loader_generator("train", "source")):
- torch.manual_seed(42)
- loader = torch.utils.data.DataLoader(
- torch.arange(64),
- batch_size=8,
- shuffle=True,
- num_workers=workers,
- generator=generator,
- )
- iterator = iter(loader)
- worker_seed = iterator._base_seed
- batches = torch.cat(list(iterator))
- results.append((worker_seed, batches, torch.get_rng_state().clone()))
- assert results[0][0] == results[1][0]
- assert torch.equal(results[0][1], results[1][1])
- assert torch.equal(results[0][2], results[1][2])
- assert not legacy._loader_generators
-
-
-def test_checkpoint_cannot_silently_change_loader_semantics():
- for left, right in (("current", "legacy_c12"), ("legacy_c12", "current")):
- with pytest.raises(ValueError, match="compatibility mode"):
- module(right).load_state_dict(module(left).state_dict())
- legacy = module("legacy_c12")
- legacy.load_state_dict(legacy.state_dict())
-
-
-def test_input_dtype_preserved_recursively_only_in_legacy_mode():
- value = torch.tensor([1.0000000001], dtype=torch.float64)
- for mode, expected in (("current", torch.float32), ("legacy_c12", torch.float64)):
- algo = PipelineAlgo([], device="cpu", compatibility_mode=mode)
- batch = algo._move_value({"a": [value, (value,)], "i": torch.tensor([1])})
- assert batch["a"][0].dtype == expected
- assert batch["a"][1][0].dtype == expected
- assert batch["i"].dtype == torch.int64
- if mode == "legacy_c12":
- assert torch.equal(batch["a"][0], value)
-
-
-def test_retry_against_pinned_historical_source():
- path = os.environ.get("C12_DATASET_SOURCE")
- if not path:
- pytest.skip("set C12_DATASET_SOURCE to the pinned historical dataset module")
- tree = ast.parse(Path(path).read_text())
- helper = next(
- n
- for n in tree.body
- if isinstance(n, ast.FunctionDef) and n.name == "get_fallback_idx"
- )
- cls = next(
- n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "MultiDataset"
- )
- method = next(
- n
- for n in cls.body
- if isinstance(n, ast.FunctionDef) and n.name == "_next_after_failure"
- )
- ns = dict(vars(dm))
- exec(
- compile(ast.Module(body=[helper, method], type_ignores=[]), str(path), "exec"),
- ns,
- )
- ds = dm.MultiDataset(state={}, compatibility_mode="legacy_c12")
- ds.index_map = [("a", i) for i in range(64)] + [("b", i) for i in range(4)]
- ds._global_indices_by_dataset = {"a": list(range(64)), "b": list(range(64, 68))}
- original_rng = random.getstate()
- try:
- for seed in (0, 42, 999):
- for attempt in (None, 0, 24, 25, 26, 62, 63, 999):
- results = []
- for method in (
- ns["_next_after_failure"],
- dm.MultiDataset._next_after_failure,
- ):
- random.seed(seed)
- try:
- result = method(ds, 0, "a", attempt, reason="fixture")
- except RuntimeError as exc:
- result = (type(exc).__name__, str(exc))
- results.append((result, random.getstate()))
- assert results[0] == results[1], (seed, attempt)
- finally:
- random.setstate(original_rng)
-
-
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.bfloat16])
-def test_metric_dtype_matches_historical_reducer_without_changing_gradient(dtype):
- from collections import OrderedDict
-
- from egomimic.pl_utils.training_behavior_action_flow import (
- ActionFlowTrainingBehavior,
- )
- from egomimic.pl_utils.training_metrics import reduce_component_means
-
- path = os.environ.get("C12_DATASET_SOURCE")
- if not path:
- pytest.skip("set C12_DATASET_SOURCE to the pinned historical dataset module")
- source = Path(path).parents[2] / "pl_utils/pl_model_action_flow.py"
- tree = ast.parse(source.read_text())
- cls = next(
- n
- for n in tree.body
- if isinstance(n, ast.ClassDef) and n.name == "ActionFlowModelWrapper"
- )
- method = next(
- n
- for n in cls.body
- if isinstance(n, ast.FunctionDef) and n.name == "_reduce_component_means"
- )
- method.decorator_list = []
- ns = {"torch": torch, "OrderedDict": OrderedDict, "Mapping": dict}
- exec(compile(ast.Module(body=[method], type_ignores=[]), str(source), "exec"), ns)
- value = torch.tensor(0.123456789, dtype=dtype, requires_grad=True)
- state = torch.get_rng_state().clone()
- expected, count = ns["_reduce_component_means"]({"loss": value}, 17)
- actual, actual_count = reduce_component_means(
- {"loss": value}, 17, label="fixture", preserve_input_dtype=True
- )
- assert count == actual_count
- assert actual["loss"].dtype == expected["loss"].dtype
- assert torch.equal(actual["loss"], expected["loss"])
- assert value.grad is None and not actual["loss"].requires_grad
- assert torch.equal(state, torch.get_rng_state())
- assert (
- ActionFlowTrainingBehavior(compatibility_mode="legacy_c12").compatibility_mode
- == "legacy_c12"
- )
- with pytest.raises(ValueError, match="compatibility_mode"):
- ActionFlowTrainingBehavior(compatibility_mode="invalid")
-
-
-def test_normalization_collection_matches_historical_precision():
- path = os.environ.get("C12_DATASET_SOURCE")
- if not path:
- pytest.skip("set C12_DATASET_SOURCE to the pinned historical dataset module")
- tree = ast.parse(Path(path).read_text())
- cls = next(
- n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "MultiDataset"
- )
- method = next(
- n
- for n in cls.body
- if isinstance(n, ast.FunctionDef) and n.name == "_collect_norm_samples"
- )
- ns = dict(vars(dm))
- exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), "exec"), ns)
- for dtype in (torch.float32, torch.float64):
- samples = torch.tensor(
- [[1.0000000001], [1.0000000003], [1.0000000009]], dtype=dtype
- )
- ds = dm.MultiDataset(state={}, compatibility_mode="legacy_c12")
- ds.keyname_to_zarr_key = lambda key, embodiment: key
- args = ([{"action": samples}], ["action"], "fixture", 2, 3, 0)
- expected = ns["_collect_norm_samples"](ds, *args)["action"][0]
- actual = ds._collect_norm_samples(*args)["action"][0]
- assert actual.dtype == expected.dtype
- np.testing.assert_array_equal(actual, expected)
- for key, value in ds._compute_stats_for_array(expected).items():
- np.testing.assert_array_equal(
- ds._compute_stats_for_array(actual)[key], value
- )
- ds.compatibility_mode = "current"
- assert ds._collect_norm_samples(*args)["action"][0].dtype == np.float32
diff --git a/tests/test_capture_runtime_lock.py b/tests/test_capture_runtime_lock.py
deleted file mode 100644
index 9dc5eec1b7..0000000000
--- a/tests/test_capture_runtime_lock.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import importlib.util
-import json
-from pathlib import Path
-
-import pytest
-
-MODULE_PATH = Path(__file__).parents[1] / "scripts" / "ice" / "capture_runtime_lock.py"
-SPEC = importlib.util.spec_from_file_location("capture_runtime_lock", MODULE_PATH)
-assert SPEC is not None and SPEC.loader is not None
-MODULE = importlib.util.module_from_spec(SPEC)
-SPEC.loader.exec_module(MODULE)
-
-
-def test_canonical_bytes_are_stable_and_newline_terminated():
- first = MODULE.canonical_bytes({"z": 1, "a": [2, 3]})
- second = MODULE.canonical_bytes({"a": [2, 3], "z": 1})
-
- assert first == second
- assert first.endswith(b"\n")
- assert json.loads(first) == {"a": [2, 3], "z": 1}
-
-
-def test_atomic_create_refuses_overwrite(tmp_path):
- output = tmp_path / "locks" / "runtime.json"
- MODULE.atomic_create(output, b"first\n")
-
- with pytest.raises(FileExistsError, match="refusing to overwrite"):
- MODULE.atomic_create(output, b"second\n")
-
- assert output.read_bytes() == b"first\n"
-
-
-def test_runtime_lock_source_has_no_deployment_specific_paths():
- text = MODULE_PATH.read_text()
-
- assert "/coc/" not in text
- assert "/storage/" not in text
- assert "paphiwetsa3" not in text
diff --git a/tests/test_check_checkpoint_storage.py b/tests/test_check_checkpoint_storage.py
deleted file mode 100644
index 8aebbea1b1..0000000000
--- a/tests/test_check_checkpoint_storage.py
+++ /dev/null
@@ -1,150 +0,0 @@
-import importlib.util
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-
-ROOT = Path(__file__).parents[1]
-SCRIPT = ROOT / "scripts" / "train" / "check_checkpoint_storage.py"
-
-
-def _load_module():
- spec = importlib.util.spec_from_file_location("check_checkpoint_storage", SCRIPT)
- assert spec is not None and spec.loader is not None
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- return module
-
-
-def _filesystem(*, available_blocks=1000, free_blocks=1100, read_only=False):
- return SimpleNamespace(
- f_frsize=4096,
- f_bavail=available_blocks,
- f_bfree=free_blocks,
- f_blocks=2000,
- f_flag=int(getattr(os, "ST_RDONLY", 1)) if read_only else 0,
- )
-
-
-def test_storage_budget_uses_larger_measured_or_expected_size(tmp_path, monkeypatch):
- module = _load_module()
- monkeypatch.setattr(module.os, "statvfs", lambda _: _filesystem())
- report, failures = module.storage_evidence(
- tmp_path,
- planned_checkpoint_count=3,
- measured_checkpoint_bytes=100,
- expected_checkpoint_bytes=120,
- safety_reserve_bytes=500,
- )
- assert failures == []
- assert report["bytes_per_checkpoint_budget"] == 120
- assert report["checkpoint_budget_bytes"] == 360
- assert report["required_available_bytes"] == 860
- assert report["filesystem_available_bytes"] == 1000 * 4096
-
-
-def test_storage_budget_reports_shortfall_and_read_only(tmp_path, monkeypatch):
- module = _load_module()
- monkeypatch.setattr(
- module.os,
- "statvfs",
- lambda _: _filesystem(available_blocks=1, free_blocks=2, read_only=True),
- )
- report, failures = module.storage_evidence(
- tmp_path,
- planned_checkpoint_count=3,
- measured_checkpoint_bytes=2000,
- expected_checkpoint_bytes=None,
- safety_reserve_bytes=1000,
- )
- assert report["required_available_bytes"] == 7000
- assert {failure["field"] for failure in failures} == {
- "filesystem_read_only",
- "filesystem_available_bytes",
- }
-
-
-def test_storage_cli_writes_pass_evidence(tmp_path):
- output = tmp_path / "storage.json"
- subprocess.run(
- [
- sys.executable,
- str(SCRIPT),
- "--target",
- str(tmp_path),
- "--planned-checkpoint-count",
- "1",
- "--measured-checkpoint-bytes",
- "1",
- "--expected-checkpoint-bytes",
- "2",
- "--safety-reserve-bytes",
- "1",
- "--output",
- str(output),
- ],
- check=True,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- payload = json.loads(output.read_text())
- assert payload["status"] == "CHECKPOINT_STORAGE_VALIDATED"
- assert payload["failures"] == []
- assert payload["storage"]["target"] == str(tmp_path.resolve())
-
-
-def test_storage_cli_writes_failure_evidence_for_insufficient_space(tmp_path):
- available = os.statvfs(tmp_path).f_bavail * os.statvfs(tmp_path).f_frsize
- output = tmp_path / "storage.json"
- completed = subprocess.run(
- [
- sys.executable,
- str(SCRIPT),
- "--target",
- str(tmp_path),
- "--planned-checkpoint-count",
- "2",
- "--expected-checkpoint-bytes",
- str(available + 1),
- "--safety-reserve-bytes",
- "1",
- "--output",
- str(output),
- ],
- check=False,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- assert completed.returncode == 1
- payload = json.loads(output.read_text())
- assert payload["status"] == "CHECKPOINT_STORAGE_FAILED"
- failure = payload["failures"][0]
- assert failure["field"] == "filesystem_available_bytes"
- assert failure["shortfall_bytes"] > 0
-
-
-def test_storage_cli_requires_checkpoint_size(tmp_path):
- completed = subprocess.run(
- [
- sys.executable,
- str(SCRIPT),
- "--target",
- str(tmp_path),
- "--planned-checkpoint-count",
- "1",
- "--safety-reserve-bytes",
- "1",
- "--output",
- str(tmp_path / "storage.json"),
- ],
- check=False,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- assert completed.returncode != 0
- assert "one or both" in completed.stderr
diff --git a/tests/test_config_graph.py b/tests/test_config_graph.py
deleted file mode 100644
index 43d7c028d9..0000000000
--- a/tests/test_config_graph.py
+++ /dev/null
@@ -1,249 +0,0 @@
-from __future__ import annotations
-
-import importlib.util
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-from egomimic.pipeline.stages_sampler import GaussianLatentNoise
-
-_TOOL = Path(__file__).resolve().parents[1] / "tools/config_graph.py"
-_SPEC = importlib.util.spec_from_file_location("config_graph_under_test", _TOOL)
-assert _SPEC is not None and _SPEC.loader is not None
-config_graph = importlib.util.module_from_spec(_SPEC)
-_SPEC.loader.exec_module(config_graph)
-
-
-def _model_yaml(note: str = "") -> str:
- return f"""
-_target_: egomimic.pl_utils.pl_model.ModelWrapper
-pipeline:
- _target_: egomimic.pipeline.algo.PipelineAlgo
- stages:
- - _target_: egomimic.pipeline.stages_sampler.KeyedFeatureProjection
- selector_key: route
- input_key: features
- output_key: condition
- output_dim: 3
- selector_aliases: {{1: image}}
- projections:
- image:
- input_dim: 2
- hidden_dim: 4
- semantic: {json.dumps(note)}
- - _target_: egomimic.pipeline.stages_sampler.GaussianLatentNoise
- num_tokens: 4
- latent_dim: 3
- - _target_: egomimic.pipeline.stages_io.ActionTargetBuilder
-""".lstrip()
-
-
-def _data_yaml(*, second_source_has_route: bool = True) -> str:
- second_batch_keys = "[route]" if second_source_has_route else "[]"
- return f"""
-train_datasets:
- opaque/source-A:
- batch_keys: [route]
- key_map: &keymap
- features: {{zarr_key: observations.features, key_type: proprio_keys}}
- actions: {{zarr_key: actions, key_type: action_keys}}
- "42":
- batch_keys: {second_batch_keys}
- key_map:
- <<: *keymap
-""".lstrip()
-
-
-def _write_full_config(
- path: Path, *, note: str = "", second_source_has_route: bool = True
-) -> None:
- model = "\n".join(f" {line}" for line in _model_yaml(note).splitlines())
- data = "\n".join(
- f" {line}"
- for line in _data_yaml(
- second_source_has_route=second_source_has_route
- ).splitlines()
- )
- path.write_text(f"model:\n{model}\ndata:\n{data}\n")
-
-
-def test_clean_pipeline_contracts_and_nested_params_are_preserved(
- tmp_path: Path, monkeypatch
-) -> None:
- selected = tmp_path / "selected.yaml"
- note = "nested-stage-parameter-" * 30
- _write_full_config(selected, note=note)
- calls = []
- original = GaussianLatentNoise.contract
-
- def record_contract(self, mode="train"):
- calls.append(mode)
- return original(self, mode)
-
- monkeypatch.setattr(GaussianLatentNoise, "contract", record_contract)
- train = config_graph.build_graph(selected, mode="train")
- inference = config_graph.build_graph(selected, mode="inference")
-
- assert calls == ["train", "inference"]
- assert train["lint"] == inference["lint"] == []
- assert train["seed_keys_by_source"] == {
- "opaque/source-A": ["actions", "features", "route"],
- "42": ["actions", "features", "route"],
- }
- assert inference["seed_keys"] == ["features", "route"]
- assert [node["t"] for node in inference["nodes"]] == [
- "KeyedFeatureProjection",
- "GaussianLatentNoise",
- ]
- assert inference["skipped_stages"] == [
- {
- "source_i": 2,
- "t": "ActionTargetBuilder",
- "target": "egomimic.pipeline.stages_io.ActionTargetBuilder",
- "reason": "train-only",
- }
- ]
- assert train["edges"] == [{"a": 0, "b": 1, "k": "condition", "s": "shared"}]
- params = train["nodes"][0]["p"]
- assert params["projections"]["image"]["semantic"] == note
- assert len(params["projections"]["image"]["semantic"]) > 400
-
-
-def test_source_names_are_opaque_and_missing_keymap_inputs_are_linted(
- tmp_path: Path,
-) -> None:
- selected = tmp_path / "missing.yaml"
- _write_full_config(selected, second_source_has_route=False)
-
- graph = config_graph.build_graph(selected)
-
- assert set(graph["seed_keys_by_source"]) == {"opaque/source-A", "42"}
- assert any(
- "source '42'" in problem and "reads 'route'" in problem
- for problem in graph["lint"]
- )
- assert not any(
- "source 'opaque/source-A'" in problem and "reads 'route'" in problem
- for problem in graph["lint"]
- )
-
-
-def test_hydra_experiment_accepts_exact_model_override(tmp_path: Path) -> None:
- root = tmp_path / "hydra_configs"
- (root / "model").mkdir(parents=True)
- (root / "data").mkdir()
- (root / "experiment").mkdir()
- (root / "train_zarr_cartesian.yaml").write_text(
- "defaults:\n - model: null\n - data: null\n - _self_\n"
- )
- (root / "model/selected.yaml").write_text(_model_yaml())
- (root / "data/selected.yaml").write_text(_data_yaml())
- experiment = root / "experiment/selected.yaml"
- experiment.write_text(
- "# @package _global_\n"
- "defaults:\n"
- " - override /data: selected\n"
- " - _self_\n"
- )
-
- overrides = ["model=selected", "++run_tag=exact-row"]
- graph = config_graph.build_graph(experiment, overrides=overrides)
-
- assert graph["lint"] == []
- assert graph["seed_key_source"] == "selected-data-keymap"
- assert graph["component_sources"]["config_root"] == str(root)
- assert graph["component_sources"]["overrides"] == overrides
-
-
-def test_legacy_model_key_is_not_considered(tmp_path: Path) -> None:
- selected = tmp_path / "legacy.yaml"
- selected.write_text(
- "robomimic_model:\n"
- " stages:\n"
- " - _target_: egomimic.pipeline.stages_io.ActionTargetBuilder\n"
- )
-
- with pytest.raises(ValueError, match=r"model\.pipeline\.stages"):
- config_graph.build_graph(selected)
-
-
-def test_lint_reports_duplicate_concrete_writers_and_cycles() -> None:
- graph = {
- "nodes": [
- {"i": 0, "t": "First", "in": ["b"], "out": ["a"]},
- {"i": 1, "t": "Second", "in": ["a"], "out": ["b", "a"]},
- ],
- "edges": [
- {"a": 1, "b": 0, "k": "b", "s": "shared"},
- {"a": 0, "b": 1, "k": "a", "s": "shared"},
- ],
- "seed_keys_by_source": {"anything": []},
- }
-
- problems = config_graph.lint(graph)
-
- assert any("reads 'b' before" in problem for problem in problems)
- assert any("duplicate writer for 'a'" in problem for problem in problems)
- assert any("dependency cycle" in problem for problem in problems)
-
-
-def test_cli_emits_renderer_compatible_both_mode_json(
- tmp_path: Path, monkeypatch
-) -> None:
- selected = tmp_path / "selected.yaml"
- output = tmp_path / "graph.json"
- _write_full_config(selected)
-
- instantiations = 0
- original = config_graph._stages
-
- def count_instantiations(pipeline_config):
- nonlocal instantiations
- instantiations += 1
- return original(pipeline_config)
-
- monkeypatch.setattr(config_graph, "_stages", count_instantiations)
- result = config_graph.main([str(output), str(selected), "--mode", "both", "--lint"])
- payload = json.loads(output.read_text())
-
- assert result == 0
- assert instantiations == 1
- assert set(payload) == {"selected [train]", "selected [inference]"}
- for graph in payload.values():
- assert isinstance(graph["nodes"], list)
- assert isinstance(graph["edges"], list)
- assert all({"i", "t", "in", "out", "p"} <= set(node) for node in graph["nodes"])
- assert all({"a", "b", "k", "s"} <= set(edge) for edge in graph["edges"])
-
-
-def test_cli_resolves_repo_imports_without_pythonpath(tmp_path: Path) -> None:
- output = tmp_path / "action-flow.graph.json"
- experiment = (
- _TOOL.parents[1] / "egomimic/hydra_configs/experiment/sim/pusht/"
- "action_flow_bc_usocket_recon1_s42.yaml"
- )
- environment = os.environ.copy()
- environment.pop("PYTHONPATH", None)
- completed = subprocess.run(
- [
- sys.executable,
- str(_TOOL),
- str(output),
- str(experiment),
- "--mode",
- "both",
- "--lint",
- ],
- cwd=tmp_path,
- check=False,
- capture_output=True,
- text=True,
- env=environment,
- )
-
- assert completed.returncode == 0, completed.stderr
- assert output.is_file()
diff --git a/tests/test_current_family_profiles.py b/tests/test_current_family_profiles.py
deleted file mode 100644
index 8e17c14b32..0000000000
--- a/tests/test_current_family_profiles.py
+++ /dev/null
@@ -1,224 +0,0 @@
-"""Guard owner-selected contracts without launching or reading a training corpus."""
-
-import runpy
-import subprocess
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-
-ROOT = Path(__file__).resolve().parents[1]
-SUITES = [
- ("libero10", 605121),
- ("libero_spatial", 270054),
- ("libero_object", 325065),
- ("libero_goal", 280056),
-]
-
-
-@pytest.mark.parametrize("suite,steps", SUITES)
-def test_libero_current_and_historical_profiles_remain_distinct(suite, steps):
- with initialize_config_dir(
- config_dir=str(ROOT / "egomimic/hydra_configs"), version_base="1.3"
- ):
- af = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- f"+experiment=sim/libero/action_flow_{suite}_oat_dp_matched_s42"
- ],
- )
- dp = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/libero/dp_{suite}_oat_batch16_keep_steps_s42"],
- )
- old = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- f"+experiment=sim/libero_historical/dp_{suite}_oat_dp_matched_s42"
- ],
- )
- assert (
- af.model.action_horizon,
- af.model.num_latent_tokens,
- af.model.latent_dim,
- ) == (32, 16, 16)
- assert (af.benchmark.n_obs_steps, af.benchmark.n_action_steps) == (2, 16)
- assert af.model.pipeline.stages[0]._target_.endswith("OATObservationStage")
- assert af.model.pipeline.stages[-1].action_velocity_weight == 1.0
- assert af.model.pipeline.stages[5].inference_method == "euler"
- assert (af.model.num_inference_steps, af.model.cfg_scale, af.model.hidden_dim) == (
- 50,
- 4.0,
- 240,
- )
- assert af.data.train_datasets.libero_panda.valid_ratio == 0.1
- assert af.data.valid_datasets.libero_panda.valid_ratio == 0.1
- assert (
- af.benchmark.batch_size,
- af.trainer.accumulate_grad_batches,
- af.trainer.gradient_clip_val,
- ) == (32, 1, 3.0)
- assert af.callbacks.ema._target_.endswith("ActionFlowFixedEMACallback")
- assert af.callbacks.ema.decay == 0.9978 and af.callbacks.ema.use_warmup is False
- assert af.model.optimizer._target_.endswith("ReleasedUniteCompositeOptimizer")
- assert (
- dp.benchmark.batch_size,
- dp.trainer.accumulate_grad_batches,
- dp.trainer.max_steps,
- ) == (16, 1, steps)
- assert dp.benchmark.original_optimizer_step_target == steps
- assert dp.callbacks.batch_budget.global_batch_size == 16
- assert dp.norm_stats.precomputed_norm_path is None
- assert dp.callbacks.model_checkpoint.save_top_k == -1
- assert old.benchmark.batch_size * old.trainer.accumulate_grad_batches == 1024
-
-
-def test_dp_multiplier_actual_phase_arguments(monkeypatch):
- for k, v in dict(
- DP_COTRAIN_STANDARD="usocket_chain_manual4919_af_obs_multiplier_261m_v1",
- DP_REPO=str(ROOT),
- DP_WANDB_ID="fixture",
- DP_EXPECTED_HEAD="a" * 40,
- DP_NORM_SHA256="b" * 64,
- DP_SPLIT_SHA256="c" * 64,
- DP_CONTENT_MANIFEST="/fixture/content.json",
- DP_CONTENT_SHA256="d" * 64,
- DP_CONTENT_AGGREGATE_SHA256="e" * 64,
- ).items():
- monkeypatch.setenv(k, v)
- monkeypatch.delenv("DP_SINGLE_SOURCE_STANDARD", raising=False)
- api = runpy.run_path(str(ROOT / "scripts/train/standard_dp_launch.py"))
- for phase in ("normalize", "smoke", "full"):
- argv = api["arguments"](phase, "/fixture/" + phase, "/fixture/norm")
- with initialize_config_dir(
- config_dir=str(ROOT / "egomimic/hydra_configs"), version_base="1.3"
- ):
- cfg = compose(config_name="train_zarr_cartesian", overrides=argv[1:])
- assert cfg.name == "planar_uc_manual4919_dp_261m_af_obs_multiplier"
- assert (
- cfg.planar.batch_size == 32
- and cfg.model.pipeline.stages[5].condition_input_dim == 128
- )
- assert cfg.model.pipeline.stages[2].conditioning_input == "retiming_multiplier"
- assert cfg.model.pipeline.stages[2].get("speed_reference") is None
- assert cfg.callbacks.get("ema") is None
- assert cfg.model.optimizer._target_ == "torch.optim.AdamW"
- assert cfg.model.scheduler.warmup_steps == 500
- assert cfg.norm_stats.sample_frac == 1.0
-
-
-@pytest.mark.parametrize(
- "cotrain,single",
- [
- ("unknown", ""),
- ("", "unknown"),
- (
- "usocket_chain_manual4919_af_obs_multiplier_261m_v1",
- "chain_manual4919_retimed_v1",
- ),
- ],
-)
-def test_unrecognized_or_conflicting_profiles_fail_before_launch(
- monkeypatch, cotrain, single
-):
- monkeypatch.setenv("DP_COTRAIN_STANDARD", cotrain)
- monkeypatch.setenv("DP_SINGLE_SOURCE_STANDARD", single)
- api = runpy.run_path(str(ROOT / "scripts/train/standard_dp_launch.py"))
- with pytest.raises(ValueError):
- api["recipe"]()
- env = dict(__import__("os").environ)
- result = subprocess.run(
- [
- "bash",
- str(ROOT / "scripts/train/planar_v2_dp_cotrain_clean_skynet_v1.sbatch"),
- ],
- env=env,
- capture_output=True,
- )
- assert result.returncode == 64
-
-
-def test_current_pair_helper_composes_all_eight_without_claiming_split_proof(tmp_path):
- api = runpy.run_path(str(ROOT / "tools/libero_oat_pair_config.py"))
- result = api["compose_all"](ROOT, tmp_path)
- assert len(result["rows"]) == 8
- assert result["recipe_contract"] == "current"
- assert (
- result["episode_split_match"]
- == "UNVERIFIED_REQUIRES_MATERIALIZED_EPISODE_LISTS"
- )
- assert result["full_training_ready"] is False
- import json
-
- for suite, steps in SUITES:
- af = json.loads((tmp_path / f"action_flow-{suite}.json").read_text())
- dp = json.loads((tmp_path / f"dp-{suite}.json").read_text())
- assert af["trainer"]["max_steps"] == 120000
- assert dp["trainer"]["max_steps"] == steps
- assert api["validate"](dp, "dp", suite) == 16
- with pytest.raises(AssertionError):
- api["validate"](dp, "dp", suite, recipe_contract="historical-v7")
- af["trainer"]["max_steps"] = 80000
- with pytest.raises(AssertionError):
- api["validate"](af, "action_flow", suite)
- assert (
- api["validate"](af, "action_flow", suite, recipe_contract="historical-v7")
- == 32
- )
-
-
-def test_historical_native_names_preserve_aliases_and_reject_current_recipe():
- from egomimic.benchmarks.libero.native_launch_profiles import (
- HISTORICAL_NATIVE_PROFILES,
- PROFILES,
- historical_profile_for_suite,
- profile_for_experiment,
- profile_for_suite,
- )
-
- assert PROFILES is HISTORICAL_NATIVE_PROFILES
- assert profile_for_suite is historical_profile_for_suite
- for suite, _ in SUITES:
- assert historical_profile_for_suite(suite).experiment.startswith(
- "sim/libero_historical/"
- )
- with pytest.raises(ValueError):
- profile_for_experiment(f"sim/libero/action_flow_{suite}_oat_dp_matched_s42")
-
-
-@pytest.mark.parametrize("receipt", [None, {"recipe_contract": "historical-v7"}])
-def test_current_full_proof_rejects_legacy_gate_before_creating_output(
- monkeypatch, tmp_path, receipt
-):
- import json
- import sys
-
- monkeypatch.syspath_prepend(str(ROOT / "tools"))
- api = runpy.run_path(str(ROOT / "tools/libero_oat_training_proof.py"))
- monkeypatch.setenv("SLURM_STEP_ID", "fixture")
- monkeypatch.setattr(
- api["subprocess"],
- "check_output",
- lambda argv, **kwargs: "a" * 40 if argv[1] == "rev-parse" else "",
- )
- argv = [
- "proof",
- "--root",
- str(tmp_path),
- "--source-commit",
- "a" * 40,
- "--index",
- "0",
- "--phase",
- "full",
- "--attempt",
- "v1",
- ]
- if receipt is not None:
- path = tmp_path / "legacy.json"
- path.write_text(json.dumps(receipt))
- argv += ["--readiness-receipt", str(path)]
- monkeypatch.setattr(sys, "argv", argv)
- with pytest.raises((ValueError, AssertionError)):
- api["main"]()
- assert not (tmp_path / "full-v1").exists()
diff --git a/tests/test_gaussian_paraboloid_dataset.py b/tests/test_gaussian_paraboloid_dataset.py
deleted file mode 100644
index 6c18929a0b..0000000000
--- a/tests/test_gaussian_paraboloid_dataset.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import torch
-
-from egomimic.synthetic import (
- GaussianParaboloidDataset,
- generate_gaussian_paraboloid,
-)
-
-
-def test_gaussian_paraboloid_is_deterministic_and_on_surface():
- first = generate_gaussian_paraboloid(128, seed=7, curvature=0.25)
- second = generate_gaussian_paraboloid(128, seed=7, curvature=0.25)
- torch.testing.assert_close(first.source_latent, second.source_latent)
- torch.testing.assert_close(first.source_gaussian_3d, second.source_gaussian_3d)
- torch.testing.assert_close(first.target_3d, second.target_3d)
- torch.testing.assert_close(first.target_3d[:, :2], first.source_2d)
- torch.testing.assert_close(
- first.target_3d[:, 2], 0.25 * first.source_2d.square().sum(dim=-1)
- )
-
-
-def test_latent_four_preserves_paired_problem_and_adds_nuisance_dimensions():
- wide = generate_gaussian_paraboloid(64, seed=9, source_dim=4)
- reference = generate_gaussian_paraboloid(64, seed=9, source_dim=2)
- assert wide.source_latent.shape == (64, 4)
- torch.testing.assert_close(wide.source_latent[:, :2], reference.source_latent)
- torch.testing.assert_close(wide.target_3d, reference.target_3d)
- torch.testing.assert_close(wide.source_gaussian_3d, reference.source_gaussian_3d)
-
-
-def test_dataset_exposes_exact_linear_bridge():
- dataset = GaussianParaboloidDataset(32, seed=11)
- indices = torch.tensor([0, 3, 7])
- time = torch.tensor([0.0, 0.5, 1.0])
- state, velocity = dataset.cfm_state_velocity(indices, time)
- source = dataset.data.source_3d[indices]
- target = dataset.data.target_3d[indices]
- torch.testing.assert_close(velocity, target - source)
- torch.testing.assert_close(state[0], source[0])
- torch.testing.assert_close(state[1], 0.5 * (source[1] + target[1]))
- torch.testing.assert_close(state[2], target[2])
diff --git a/tests/test_gaussian_torus_dataset.py b/tests/test_gaussian_torus_dataset.py
deleted file mode 100644
index 6591cb47d7..0000000000
--- a/tests/test_gaussian_torus_dataset.py
+++ /dev/null
@@ -1,63 +0,0 @@
-import math
-import subprocess
-import sys
-from pathlib import Path
-
-import torch
-
-from egomimic.synthetic import GaussianTorusDataset, generate_gaussian_torus
-
-
-def test_generator_script_resolves_the_checkout_outside_repo_cwd(tmp_path):
- source = Path(__file__).parents[1]
- output = tmp_path / "generated.npz"
- subprocess.run(
- [
- sys.executable,
- str(source / "scripts/data/generate_gaussian_torus.py"),
- "--output",
- str(output),
- "--count",
- "32",
- "--source-dim",
- "8",
- ],
- cwd=tmp_path,
- check=True,
- )
- assert output.is_file()
-
-
-def test_gaussian_torus_is_deterministic_and_on_surface():
- first = generate_gaussian_torus(128, seed=7)
- second = generate_gaussian_torus(128, seed=7)
- assert torch.equal(first.source_2d, second.source_2d)
- assert torch.equal(first.source_gaussian_3d, second.source_gaussian_3d)
- assert torch.equal(first.target_3d, second.target_3d)
- radial = first.target_3d[:, :2].norm(dim=-1)
- residual = (radial - 2.0).square() + first.target_3d[:, 2].square()
- assert torch.allclose(residual, torch.full_like(residual, 0.65**2), atol=2e-6)
- assert torch.all((first.angles >= 0) & (first.angles <= 2 * math.pi))
-
-
-def test_linear_cfm_bridge_endpoints_and_velocity():
- dataset = GaussianTorusDataset(16, seed=11)
- indices = torch.tensor([1, 4, 9])
- start, velocity = dataset.cfm_state_velocity(indices, torch.zeros(3))
- end, velocity_end = dataset.cfm_state_velocity(indices, torch.ones(3))
- assert torch.equal(start, dataset.data.source_3d[indices])
- assert torch.allclose(end, dataset.data.target_3d[indices])
- assert torch.equal(velocity, velocity_end)
- assert torch.allclose(end - start, velocity)
-
-
-def test_four_dimensional_source_keeps_same_torus_coupling():
- batch = generate_gaussian_torus(64, seed=9, source_dim=4)
- assert batch.source_latent.shape == (64, 4)
- assert torch.equal(batch.source_latent[:, :2], batch.source_2d)
- reference = generate_gaussian_torus(64, seed=9, source_dim=2)
- assert torch.equal(batch.source_2d, reference.source_2d)
- assert torch.equal(batch.target_3d, reference.target_3d)
- assert torch.equal(batch.source_gaussian_3d, reference.source_gaussian_3d)
- assert batch.source_gaussian_latent.shape == (64, 4)
- assert torch.equal(batch.source_gaussian_3d, batch.source_gaussian_latent[:, :3])
diff --git a/tests/test_ice_checkpoint_archive_inventory.py b/tests/test_ice_checkpoint_archive_inventory.py
deleted file mode 100644
index 5c4fff60d4..0000000000
--- a/tests/test_ice_checkpoint_archive_inventory.py
+++ /dev/null
@@ -1,216 +0,0 @@
-#!/usr/bin/env python3
-
-from __future__ import annotations
-
-import hashlib
-import importlib.util
-import json
-import sys
-import tempfile
-import unittest
-from pathlib import Path
-
-ICE_DIR = Path(__file__).parents[1] / "scripts" / "ice"
-INVENTORY_PATH = ICE_DIR / "ice_checkpoint_archive_inventory.py"
-sys.path.insert(0, str(ICE_DIR))
-SPEC = importlib.util.spec_from_file_location(
- "ice_checkpoint_archive_inventory", INVENTORY_PATH
-)
-INVENTORY = importlib.util.module_from_spec(SPEC)
-assert SPEC.loader is not None
-sys.modules[SPEC.name] = INVENTORY
-SPEC.loader.exec_module(INVENTORY)
-
-
-def make_validator(root: Path) -> Path:
- validator = root / "validate.py"
- validator.write_text("#!/bin/sh\nexit 0\n")
- validator.chmod(0o700)
- return validator
-
-
-def write_complete(run: Path, checkpoint: Path, step: int) -> None:
- (run / "COMPLETE.json").write_text(
- json.dumps(
- {
- "schema_version": 1,
- "status": "COMPLETE",
- "checkpoint": {
- "global_step": step,
- "sha256": hashlib.sha256(checkpoint.read_bytes()).hexdigest(),
- },
- }
- )
- )
-
-
-class ArchiveInventoryTest(unittest.TestCase):
- def test_distinguishes_completed_archived_pending_and_unregistered(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw) / "scratch" / "campaign"
- state_dir = root / "service" / "state"
- state_dir.mkdir(parents=True)
- validator = make_validator(Path(raw))
-
- archived = root / "archived"
- archived_ckpt = archived / "checkpoints" / "step-10.ckpt"
- archived_ckpt.parent.mkdir(parents=True)
- archived_ckpt.write_text("archived")
- write_complete(archived, archived_ckpt, 10)
-
- pending = root / "pending"
- pending_ckpt = pending / "checkpoints" / "step-20.ckpt"
- pending_ckpt.parent.mkdir(parents=True)
- pending_ckpt.write_text("pending")
- write_complete(pending, pending_ckpt, 20)
-
- unknown = root / "unknown"
- unknown_ckpt = unknown / "checkpoints" / "step-30.ckpt"
- unknown_ckpt.parent.mkdir(parents=True)
- unknown_ckpt.write_text("unknown")
-
- manifest = Path(raw) / "manifest.json"
- manifest.write_text(
- json.dumps(
- {
- "schema_version": 1,
- "runs": [
- {
- "id": "archived",
- "local_root": str(archived),
- "checkpoint_glob": "checkpoints/*.ckpt",
- "validator": str(validator),
- "remote_dir": "/remote/archive/archived",
- },
- {
- "id": "pending",
- "local_root": str(pending),
- "checkpoint_glob": "checkpoints/*.ckpt",
- "validator": str(validator),
- "remote_dir": "/remote/archive/pending",
- },
- ],
- }
- )
- )
- (state_dir / "mirror-state.json").write_text(
- json.dumps(
- {
- "schema_version": 2,
- "files": {
- str(archived_ckpt.resolve()): {
- "run_id": "archived",
- "remote_verified": True,
- "global_step": 10,
- "sha256": hashlib.sha256(
- archived_ckpt.read_bytes()
- ).hexdigest(),
- }
- },
- }
- )
- )
- output = state_dir / "archive-inventory.json"
- code = INVENTORY.main(
- [
- "--search-root",
- str(root),
- "--manifest",
- str(manifest),
- "--state-dir",
- str(state_dir),
- "--output",
- str(output),
- "--require-clear",
- ]
- )
- self.assertEqual(code, 1)
- report = json.loads(output.read_text())
- statuses = {
- Path(row["run_root"]).name: row["status"] for row in report["runs"]
- }
- self.assertEqual(statuses["archived"], "complete_archived")
- self.assertEqual(statuses["pending"], "complete_needs_transfer")
- self.assertEqual(statuses["unknown"], "unregistered")
-
- def test_slurm_like_metadata_does_not_imply_completion(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw) / "scratch" / "campaign"
- run = root / "run"
- checkpoint = run / "checkpoints" / "step.ckpt"
- checkpoint.parent.mkdir(parents=True)
- checkpoint.write_text("checkpoint")
- (run / "slurm-COMPLETED.txt").write_text("COMPLETED")
- report = INVENTORY.build_inventory(
- root,
- {"schema_version": 1, "runs": []},
- {"schema_version": 2, "files": {}},
- 4,
- )
- self.assertFalse(report["runs"][0]["complete"])
-
- def test_completed_run_with_missing_checkpoint_is_not_invisible(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw) / "scratch" / "campaign"
- run = root / "run"
- run.mkdir(parents=True)
- (run / "COMPLETE.json").write_text(
- json.dumps(
- {
- "schema_version": 1,
- "status": "COMPLETE",
- "checkpoint": {"global_step": 10, "sha256": "a" * 64},
- }
- )
- )
- report = INVENTORY.build_inventory(
- root,
- {"schema_version": 1, "runs": []},
- {"schema_version": 2, "files": {}},
- 4,
- )
- self.assertEqual(report["runs"][0]["status"], "unregistered")
- self.assertTrue(report["runs"][0]["complete"])
- self.assertEqual(report["runs"][0]["checkpoint_count"], 0)
-
- def test_registered_complete_run_with_missing_checkpoint_requires_attention(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw) / "scratch" / "campaign"
- run = root / "run"
- run.mkdir(parents=True)
- sentinel = run / "COMPLETE.json"
- sentinel.write_text(
- json.dumps(
- {
- "schema_version": 1,
- "status": "COMPLETE",
- "checkpoint": {"global_step": 10, "sha256": "a" * 64},
- }
- )
- )
- report = INVENTORY.build_inventory(
- root,
- {
- "schema_version": 1,
- "runs": [
- {
- "id": "run",
- "local_root": str(run),
- "checkpoint_glob": "checkpoints/*.ckpt",
- "validator": "/unused/by_inventory",
- "remote_dir": "/remote/archive/run",
- "completion_sentinel": str(sentinel),
- }
- ],
- },
- {"schema_version": 2, "files": {}},
- 4,
- )
- self.assertEqual(
- report["runs"][0]["status"], "complete_missing_checkpoints"
- )
- self.assertEqual(report["attention_required"], [str(run.resolve())])
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_ice_checkpoint_mirror.py b/tests/test_ice_checkpoint_mirror.py
deleted file mode 100644
index 88f69ee049..0000000000
--- a/tests/test_ice_checkpoint_mirror.py
+++ /dev/null
@@ -1,420 +0,0 @@
-#!/usr/bin/env python3
-
-from __future__ import annotations
-
-import hashlib
-import importlib.util
-import json
-import os
-import subprocess
-import sys
-import tempfile
-import time
-import unittest
-from pathlib import Path
-from unittest import mock
-
-MODULE_PATH = Path(__file__).parents[1] / "scripts" / "ice" / "ice_checkpoint_mirror.py"
-SBATCH_PATH = (
- Path(__file__).parents[1] / "scripts" / "ice" / "ice_checkpoint_mirror.sbatch"
-)
-SPEC = importlib.util.spec_from_file_location("ice_checkpoint_mirror", MODULE_PATH)
-MODULE = importlib.util.module_from_spec(SPEC)
-assert SPEC.loader is not None
-sys.modules[SPEC.name] = MODULE
-SPEC.loader.exec_module(MODULE)
-
-
-def make_validator(root: Path) -> Path:
- validator = root / "validate_checkpoint.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,sys\n"
- "try:\n"
- " data=json.load(open(sys.argv[1]))\n"
- "except Exception:\n"
- " raise SystemExit(2)\n"
- "if not data.get('valid', True): raise SystemExit(3)\n"
- "print(json.dumps({'global_step':data['global_step'],'run_id':data.get('run_id','run')}))\n"
- )
- validator.chmod(0o700)
- return validator
-
-
-def write_checkpoint(path: Path, step: int, *, valid: bool = True) -> None:
- temporary = path.with_suffix(path.suffix + ".tmp")
- temporary.write_text(
- json.dumps({"global_step": step, "valid": valid, "run_id": "run"})
- )
- os.replace(temporary, path)
-
-
-def write_manifest(
- path: Path,
- run: Path,
- validator: Path,
- *,
- retain: int = 2,
- sentinel: Path | None = None,
-) -> None:
- row = {
- "id": "run",
- "local_root": str(run),
- "checkpoint_glob": "*.ckpt",
- "validator": str(validator),
- "remote_dir": "/remote/archive/run",
- "retain_local": retain,
- }
- if sentinel is not None:
- row["completion_sentinel"] = str(sentinel)
- path.write_text(json.dumps({"schema_version": 1, "runs": [row]}))
-
-
-def mirror_args(manifest: Path, state: Path, scratch: Path) -> list[str]:
- return [
- "--manifest",
- str(manifest),
- "--state-dir",
- str(state),
- "--scratch-root",
- str(scratch),
- "--quota-bytes",
- "1000000",
- "--remote-host",
- "Skynet",
- "--prune-after-verify",
- "--once",
- ]
-
-
-class CheckpointMetadataTest(unittest.TestCase):
- def test_cpu_wrapper_forwards_configured_du_timeout(self):
- wrapper = SBATCH_PATH.read_text()
- self.assertIn('"${ICE_MIRROR_PYTHON}" "${ICE_MIRROR_SCRIPT}"', wrapper)
- self.assertIn(
- 'export PATH="$(dirname "${ICE_MIRROR_PYTHON}"):${PATH}"', wrapper
- )
- self.assertIn(
- 'export PYTHONPATH="${ICE_MIRROR_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}"',
- wrapper,
- )
- self.assertIn('if [[ -n "${ICE_MIRROR_DU_TIMEOUT_SECONDS:-}" ]]', wrapper)
- self.assertIn(
- 'ARGS+=(--du-timeout-seconds "${ICE_MIRROR_DU_TIMEOUT_SECONDS}")',
- wrapper,
- )
-
- def test_stable_validation_requires_json_global_step(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- checkpoint = root / "x.ckpt"
- write_checkpoint(checkpoint, 7)
- validator = make_validator(root)
- info = MODULE.stable_valid_info(checkpoint.resolve(), validator)
- self.assertIsNotNone(info)
- self.assertEqual(info.global_step, 7)
- self.assertEqual(
- info.sha256, hashlib.sha256(checkpoint.read_bytes()).hexdigest()
- )
- bad_validator = root / "bad_validator"
- bad_validator.write_text("#!/bin/sh\nprintf 'not-json\\n'\n")
- bad_validator.chmod(0o700)
- rejection = {}
- self.assertIsNone(
- MODULE.stable_valid_info(
- checkpoint.resolve(),
- bad_validator,
- rejection=rejection,
- )
- )
- self.assertEqual(
- rejection["reason"], "invalid_metadata_or_unreadable_checkpoint"
- )
-
- def test_validator_failure_preserves_bounded_stderr_and_returncode(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- checkpoint = root / "x.ckpt"
- write_checkpoint(checkpoint, 7)
- validator = root / "bad_validator"
- validator.write_text(
- "#!/bin/sh\nprintf '%s' '"
- + "x" * 3000
- + "ModuleNotFoundError: No module named torch' >&2\nexit 9\n"
- )
- validator.chmod(0o700)
- rejection = {}
- self.assertIsNone(
- MODULE.stable_valid_info(checkpoint, validator, rejection=rejection)
- )
- self.assertEqual(rejection["reason"], "validator_failed")
- self.assertEqual(rejection["validator_returncode"], 9)
- self.assertEqual(len(rejection["validator_stderr_tail"]), 2048)
- self.assertIn("No module named torch", rejection["validator_stderr_tail"])
- self.assertTrue(checkpoint.exists())
-
- def test_changing_checkpoint_remains_deferred_with_explicit_reason(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- checkpoint = root / "x.ckpt"
- write_checkpoint(checkpoint, 7)
- validator = make_validator(root)
-
- def changing_validator(*args, **kwargs):
- checkpoint.write_text(json.dumps({"global_step": 777777}))
- return subprocess.CompletedProcess(args[0], 0, '{"global_step": 7}', "")
-
- rejection = {}
- with mock.patch.object(
- MODULE.subprocess, "run", side_effect=changing_validator
- ):
- self.assertIsNone(
- MODULE.stable_valid_info(checkpoint, validator, rejection=rejection)
- )
- self.assertEqual(
- rejection, {"reason": "checkpoint_changed_during_validation"}
- )
- self.assertTrue(checkpoint.exists())
-
- def test_specific_absolute_rejects_relative_path(self):
- with self.assertRaisesRegex(SystemExit, "specific absolute"):
- MODULE.specific_absolute(Path("relative/run/path"), "test")
-
-
-class MirrorSafetyTest(unittest.TestCase):
- def test_invalid_newer_files_do_not_count_toward_two_valid_retained(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- scratch = root / "scratch"
- run = scratch / "project" / "run"
- state = scratch / "state" / "mirror"
- run.mkdir(parents=True)
- state.mkdir(parents=True)
- validator = make_validator(root)
- first = run / "step-1.ckpt"
- second = run / "step-2.ckpt"
- corrupt_a = run / "step-100.ckpt"
- corrupt_b = run / "step-101.ckpt"
- write_checkpoint(first, 1)
- write_checkpoint(second, 2)
- write_checkpoint(corrupt_a, 100, valid=False)
- write_checkpoint(corrupt_b, 101, valid=False)
- now = time.time()
- os.utime(corrupt_a, (now + 10, now + 10))
- os.utime(corrupt_b, (now + 20, now + 20))
- manifest = root / "manifest.json"
- write_manifest(manifest, run, validator)
- remote: dict[str, str] = {}
-
- def fake_mirror(path, digest, remote_dir, host, ssh, rsync, rsync_rsh):
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(MODULE, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- MODULE,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(MODULE, "scratch_bytes", return_value=100),
- ):
- code = MODULE.main(mirror_args(manifest, state, scratch))
- self.assertEqual(
- code, 1, "invalid candidates must keep completion unresolved"
- )
- self.assertTrue(first.exists())
- self.assertTrue(second.exists())
- self.assertTrue(corrupt_a.exists())
- self.assertTrue(corrupt_b.exists())
- pressure = json.loads((state / "storage-pressure.json").read_text())
- self.assertEqual(pressure["cycle_errors"], 2)
- rejected = [
- json.loads(line)
- for line in (state / "mirror-events.jsonl").read_text().splitlines()
- if json.loads(line)["event"] == "checkpoint_not_stably_valid"
- ]
- self.assertEqual(len(rejected), 2)
- self.assertTrue(
- all(event["validator_returncode"] == 3 for event in rejected)
- )
- self.assertTrue(
- all(event["reason"] == "validator_failed" for event in rejected)
- )
-
- def test_prune_rechecks_remote_sha_immediately(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- scratch = root / "scratch"
- run = scratch / "project" / "run"
- state = scratch / "state" / "mirror"
- run.mkdir(parents=True)
- state.mkdir(parents=True)
- validator = make_validator(root)
- checkpoints = []
- for step in (1, 2, 3):
- path = run / f"step-{step}.ckpt"
- write_checkpoint(path, step)
- checkpoints.append(path)
- manifest = root / "manifest.json"
- write_manifest(manifest, run, validator)
- files = {}
- remote = {}
- for path in checkpoints:
- digest = hashlib.sha256(path.read_bytes()).hexdigest()
- remote_path = f"/remote/{digest}"
- files[str(path.resolve())] = {
- "sha256": digest,
- "run_id": "run",
- "remote_path": remote_path,
- "remote_verified": True,
- }
- remote[remote_path] = digest
- (state / "mirror-state.json").write_text(
- json.dumps({"schema_version": 1, "files": files})
- )
- oldest_remote = files[str(checkpoints[0].resolve())]["remote_path"]
- calls = {oldest_remote: 0}
-
- def fake_remote_sha(ssh, host, path):
- if path == oldest_remote:
- calls[path] += 1
- return remote[path] if calls[path] == 1 else None
- return remote.get(path)
-
- with (
- mock.patch.object(
- MODULE,
- "mirror_one",
- side_effect=AssertionError("no upload expected"),
- ),
- mock.patch.object(MODULE, "remote_sha", side_effect=fake_remote_sha),
- mock.patch.object(MODULE, "scratch_bytes", return_value=100),
- ):
- code = MODULE.main(mirror_args(manifest, state, scratch))
- self.assertEqual(code, 1)
- self.assertTrue(
- checkpoints[0].exists(),
- "local file was pruned after remote disappeared",
- )
- final_state = json.loads((state / "mirror-state.json").read_text())
- self.assertFalse(
- final_state["files"][str(checkpoints[0].resolve())]["remote_verified"]
- )
-
- def test_du_failure_is_recorded_without_uncaught_exit(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- scratch = root / "scratch"
- run = scratch / "project" / "run"
- state = scratch / "state" / "mirror"
- run.mkdir(parents=True)
- state.mkdir(parents=True)
- validator = make_validator(root)
- manifest = root / "manifest.json"
- write_manifest(manifest, run, validator)
- with mock.patch.object(
- MODULE, "scratch_bytes", side_effect=RuntimeError("du failed")
- ):
- code = MODULE.main(mirror_args(manifest, state, scratch))
- self.assertEqual(code, 1)
- pressure = json.loads((state / "storage-pressure.json").read_text())
- self.assertEqual(pressure["pressure"], "unknown")
- self.assertIn("du failed", pressure["error"])
-
- def test_deterministic_resumable_temp_and_full_digest_name(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- checkpoint = root / "step-9.ckpt"
- checkpoint.write_bytes(b"checkpoint")
- digest = hashlib.sha256(checkpoint.read_bytes()).hexdigest()
- remote_commands = []
- local_commands = []
- sha_results = iter((None, digest, digest))
-
- def fake_remote_command(ssh, host, argv, capture=False):
- remote_commands.append(argv)
- return ""
-
- def fake_run_checked(command, capture=False, timeout=None):
- local_commands.append(command)
- return subprocess.CompletedProcess(command, 0, "", "")
-
- with (
- mock.patch.object(
- MODULE, "remote_sha", side_effect=lambda *args: next(sha_results)
- ),
- mock.patch.object(
- MODULE, "remote_command", side_effect=fake_remote_command
- ),
- mock.patch.object(MODULE, "run_checked", side_effect=fake_run_checked),
- ):
- final = MODULE.mirror_one(
- checkpoint,
- digest,
- "/remote/archive/run",
- "Skynet",
- "ssh",
- "rsync",
- None,
- )
- self.assertIn(digest, final)
- self.assertEqual(local_commands[0][-1], f"Skynet:{final}.partial")
- self.assertNotRegex(local_commands[0][-1], r"\.partial\.\d+$")
- self.assertIn(
- ["mv", "-n", "--", f"{final}.partial", final], remote_commands
- )
-
- def test_completion_sentinel_stops_monitor_after_final_mirror(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- scratch = root / "scratch"
- run = scratch / "project" / "run"
- state = scratch / "state" / "mirror"
- run.mkdir(parents=True)
- state.mkdir(parents=True)
- validator = make_validator(root)
- checkpoint = run / "step-10.ckpt"
- write_checkpoint(checkpoint, 10)
- sentinel = run / "COMPLETE.json"
- sentinel.write_text(
- json.dumps(
- {
- "schema_version": 1,
- "status": "COMPLETE",
- "checkpoint": {
- "global_step": 10,
- "sha256": hashlib.sha256(
- checkpoint.read_bytes()
- ).hexdigest(),
- },
- }
- )
- )
- manifest = root / "manifest.json"
- write_manifest(manifest, run, validator, sentinel=sentinel)
- remote = {}
-
- def fake_mirror(path, digest, remote_dir, host, ssh, rsync, rsync_rsh):
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(MODULE, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- MODULE,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(MODULE, "scratch_bytes", return_value=100),
- ):
- code = MODULE.main(mirror_args(manifest, state, scratch))
- self.assertEqual(code, 0)
- complete = json.loads((state / "mirror-complete.json").read_text())
- self.assertEqual(complete["status"], "COMPLETE")
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_ice_checkpoint_mirror_pool.py b/tests/test_ice_checkpoint_mirror_pool.py
deleted file mode 100644
index beed452522..0000000000
--- a/tests/test_ice_checkpoint_mirror_pool.py
+++ /dev/null
@@ -1,310 +0,0 @@
-#!/usr/bin/env python3
-
-from __future__ import annotations
-
-import hashlib
-import importlib.util
-import json
-import sys
-import tempfile
-import unittest
-from pathlib import Path
-from unittest import mock
-
-ICE_DIR = Path(__file__).parents[1] / "scripts" / "ice"
-POOL_PATH = ICE_DIR / "ice_checkpoint_mirror_pool.py"
-SBATCH_PATH = ICE_DIR / "ice_checkpoint_mirror_pool.sbatch"
-sys.path.insert(0, str(ICE_DIR))
-SPEC = importlib.util.spec_from_file_location("ice_checkpoint_mirror_pool", POOL_PATH)
-POOL = importlib.util.module_from_spec(SPEC)
-assert SPEC.loader is not None
-sys.modules[SPEC.name] = POOL
-SPEC.loader.exec_module(POOL)
-
-
-def make_validator(root: Path) -> Path:
- validator = root / "validate_checkpoint.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,sys\n"
- "data=json.load(open(sys.argv[1]))\n"
- "print(json.dumps({'global_step':data['global_step']}))\n"
- )
- validator.chmod(0o700)
- return validator
-
-
-def make_layout(root: Path) -> tuple[Path, Path, Path, Path, Path]:
- scratch = root / "scratch"
- run = scratch / "project" / "run"
- state = scratch / "state" / "mirror"
- run.mkdir(parents=True)
- state.mkdir(parents=True)
- validator = make_validator(root)
- checkpoint = run / "step-10.ckpt"
- checkpoint.write_text(json.dumps({"global_step": 10}))
- sentinel = run / "COMPLETE.json"
- sentinel.write_text(
- json.dumps(
- {
- "schema_version": 1,
- "status": "COMPLETE",
- "checkpoint": {
- "global_step": 10,
- "sha256": hashlib.sha256(checkpoint.read_bytes()).hexdigest(),
- },
- }
- )
- )
- manifest = root / "manifest.json"
- manifest.write_text(
- json.dumps(
- {
- "schema_version": 1,
- "runs": [
- {
- "id": "run",
- "local_root": str(run),
- "checkpoint_glob": "*.ckpt",
- "validator": str(validator),
- "remote_dir": "/remote/archive/run",
- "retain_local": 2,
- "completion_sentinel": str(sentinel),
- }
- ],
- }
- )
- )
- return scratch, run, state, checkpoint, manifest
-
-
-def worker_args(
- manifest: Path, state: Path, scratch: Path, index: int = 0
-) -> list[str]:
- return [
- "--manifest",
- str(manifest),
- "--state-dir",
- str(state),
- "--scratch-root",
- str(scratch),
- "--quota-bytes",
- "1000000",
- "--remote-host",
- "Skynet",
- "--worker-index",
- str(index),
- "--worker-count",
- "2",
- "--once",
- ]
-
-
-class MirrorPoolTest(unittest.TestCase):
- def test_wrapper_requires_matching_array_and_distinct_nodes(self):
- wrapper = SBATCH_PATH.read_text()
- self.assertIn("#SBATCH --exclusive", wrapper)
- self.assertIn("SLURM_ARRAY_TASK_COUNT", wrapper)
- self.assertIn("SLURM_ARRAY_TASK_ID", wrapper)
- self.assertIn("ICE_MIRROR_POOL_SIZE", wrapper)
- self.assertIn('"${ICE_MIRROR_PYTHON}" "${ICE_MIRROR_POOL_SCRIPT}"', wrapper)
- self.assertIn(
- 'export PATH="$(dirname "${ICE_MIRROR_PYTHON}"):${PATH}"', wrapper
- )
- self.assertIn(
- 'export PYTHONPATH="${ICE_MIRROR_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}"',
- wrapper,
- )
- self.assertIn("ICE_MIRROR_INVENTORY_ROOT", wrapper)
-
- def test_nonblocking_claim_allows_only_one_worker(self):
- with tempfile.TemporaryDirectory() as raw:
- state = Path(raw) / "shared" / "state"
- checkpoint = Path(raw) / "shared" / "run" / "step.ckpt"
- checkpoint.parent.mkdir(parents=True)
- checkpoint.touch()
- with POOL.claim(state, checkpoint) as first:
- self.assertIsNotNone(first)
- with POOL.claim(state, checkpoint) as second:
- self.assertIsNone(second)
- with POOL.claim(state, checkpoint) as third:
- self.assertIsNotNone(third)
-
- def test_pool_mirrors_and_publishes_completion(self):
- with tempfile.TemporaryDirectory() as raw:
- scratch, _run, state, checkpoint, manifest = make_layout(Path(raw))
- remote: dict[str, str] = {}
-
- def fake_mirror(
- path,
- digest,
- remote_dir,
- host,
- ssh,
- rsync,
- rsync_rsh,
- temporary_tag=None,
- ):
- self.assertEqual(temporary_tag, POOL.task_key(path)[:24])
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(POOL.core, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- POOL.core,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(POOL.core, "scratch_bytes", return_value=100),
- ):
- code = POOL.main(worker_args(manifest, state, scratch))
- self.assertEqual(code, 0)
- saved = json.loads((state / "mirror-state.json").read_text())
- entry = saved["files"][str(checkpoint.resolve())]
- self.assertTrue(entry["remote_verified"])
- complete = json.loads((state / "mirror-complete.json").read_text())
- self.assertEqual(complete["worker_count"], 2)
-
- def test_nonmaintenance_worker_never_prunes_or_completes(self):
- with tempfile.TemporaryDirectory() as raw:
- scratch, _run, state, _checkpoint, manifest = make_layout(Path(raw))
- remote: dict[str, str] = {}
-
- def fake_mirror(
- path,
- digest,
- remote_dir,
- host,
- ssh,
- rsync,
- rsync_rsh,
- temporary_tag=None,
- ):
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(POOL.core, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- POOL.core,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(
- POOL,
- "maintain",
- side_effect=AssertionError("worker one ran maintenance"),
- ),
- ):
- code = POOL.main(worker_args(manifest, state, scratch, index=1))
- self.assertEqual(code, 0)
- self.assertFalse((state / "mirror-complete.json").exists())
-
- def test_rejected_candidate_is_reported_and_blocks_completion(self):
- for worker_index in (0, 1):
- with (
- self.subTest(worker_index=worker_index),
- tempfile.TemporaryDirectory() as raw,
- ):
- scratch, run, state, checkpoint, manifest = make_layout(Path(raw))
- invalid = run / "step-100.ckpt"
- invalid.write_text("incomplete checkpoint")
- remote: dict[str, str] = {}
-
- def fake_mirror(path, digest, *args, **kwargs):
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(POOL.core, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- POOL.core,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(POOL.core, "scratch_bytes", return_value=100),
- ):
- code = POOL.main(
- worker_args(manifest, state, scratch, index=worker_index)
- )
- self.assertEqual(code, 1)
- self.assertTrue(checkpoint.exists())
- self.assertTrue(invalid.exists())
- self.assertFalse((state / "mirror-complete.json").exists())
- status = json.loads(
- (state / f"worker-{worker_index}-status.json").read_text()
- )
- self.assertEqual(status["cycle_errors"], 1)
- events = [
- json.loads(line)
- for line in (state / "mirror-events.jsonl").read_text().splitlines()
- ]
- rejected = [
- event
- for event in events
- if event["event"] == "checkpoint_not_stably_valid"
- ]
- self.assertEqual(len(rejected), 1)
- self.assertEqual(rejected[0]["reason"], "validator_failed")
- self.assertEqual(rejected[0]["validator_returncode"], 1)
- self.assertIn("JSONDecodeError", rejected[0]["validator_stderr_tail"])
- if worker_index == 0:
- pressure = json.loads((state / "storage-pressure.json").read_text())
- self.assertEqual(pressure["cycle_errors"], 1)
-
- def test_completion_requires_exact_terminal_checkpoint(self):
- with tempfile.TemporaryDirectory() as raw:
- scratch, run, state, checkpoint, manifest = make_layout(Path(raw))
- sentinel = run / "COMPLETE.json"
- payload = json.loads(sentinel.read_text())
- payload["checkpoint"]["global_step"] = 999
- payload["checkpoint"]["sha256"] = "a" * 64
- sentinel.write_text(json.dumps(payload))
- remote: dict[str, str] = {}
-
- def fake_mirror(
- path,
- digest,
- remote_dir,
- host,
- ssh,
- rsync,
- rsync_rsh,
- temporary_tag=None,
- ):
- destination = f"/remote/{digest}"
- remote[destination] = digest
- return destination
-
- with (
- mock.patch.object(POOL.core, "mirror_one", side_effect=fake_mirror),
- mock.patch.object(
- POOL.core,
- "remote_sha",
- side_effect=lambda ssh, host, path: remote.get(path),
- ),
- mock.patch.object(POOL.core, "scratch_bytes", return_value=100),
- ):
- code = POOL.main(worker_args(manifest, state, scratch))
- self.assertEqual(code, 0)
- self.assertFalse((state / "mirror-complete.json").exists())
- saved = json.loads((state / "mirror-state.json").read_text())
- self.assertTrue(
- saved["files"][str(checkpoint.resolve())]["remote_verified"]
- )
-
- def test_worker_count_and_index_are_fail_closed(self):
- with tempfile.TemporaryDirectory() as raw:
- scratch, _run, state, _checkpoint, manifest = make_layout(Path(raw))
- args = worker_args(manifest, state, scratch)
- args[args.index("--worker-count") + 1] = "1"
- with self.assertRaisesRegex(SystemExit, "worker-count"):
- POOL.main(args)
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_ice_dp_campaign_contract.py b/tests/test_ice_dp_campaign_contract.py
deleted file mode 100644
index c8e40ffc01..0000000000
--- a/tests/test_ice_dp_campaign_contract.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from pathlib import Path
-
-ROOT = Path(__file__).parents[1]
-LAUNCHER = ROOT / "scripts/ice/launch_planar_v2_dp_cotrain_clean.sbatch"
-README = ROOT / "scripts/ice/README.md"
-
-
-def test_dp_ice_launcher_has_one_entrypoint_and_fresh_run_guards():
- text = LAUNCHER.read_text()
- assert "MODE=${MODE:-preflight}" in text
- assert 'case "$MODE" in preflight|smoke|full)' in text
- assert "ckpt_path=null" in text
- assert 'test -z "${ICE_INITIAL_CHECKPOINT:-}"' in text
- assert 'test -z "${ICE_INITIAL_CHECKPOINT_SHA256:-}"' in text
- assert 'test -z "${CKPT_PATH:-}"' in text
- assert "PASSED_SMOKE_DIR" in text
- assert "trainer.max_steps=$max_steps" in text
- assert "FULL_STEPS=240000" in text
- assert "VALIDATE_EVERY=10000" in text
- assert "CHECKPOINT_EVERY=20000" in text
- assert "NVIDIA H100 80GB HBM3" in text and "NVIDIA H200" in text
-
-
-def test_dp_ice_docs_distinguish_smoke_full_and_requeue_from_restore():
- text = README.read_text()
- assert "Do not infer full-run completion from a completed smoke" in text
- assert "scheduler-generated\nHPC checkpoint" in text
- assert "H100 and H200 are\ninterchangeable" in text
diff --git a/tests/test_ice_gpu_probe.py b/tests/test_ice_gpu_probe.py
deleted file mode 100644
index b3fb1b4fec..0000000000
--- a/tests/test_ice_gpu_probe.py
+++ /dev/null
@@ -1,50 +0,0 @@
-import importlib.util
-from pathlib import Path
-
-import pytest
-
-MODULE_PATH = Path(__file__).parents[1] / "scripts" / "ice" / "ice_gpu_probe.py"
-SPEC = importlib.util.spec_from_file_location("ice_gpu_probe", MODULE_PATH)
-assert SPEC is not None and SPEC.loader is not None
-MODULE = importlib.util.module_from_spec(SPEC)
-SPEC.loader.exec_module(MODULE)
-
-
-def test_parse_ecc_health_accepts_zero_counters_and_clear_flags():
- report = """
- Volatile Uncorrectable ECC : 0
- Aggregate Uncorrectable ECC : 0
- Pending Page Blacklist : No
- Reset Required : No
- Remapping Failure Occurred : No
- """
-
- health = MODULE.parse_ecc_health(report)
-
- assert len(health["uncorrectable_ecc_counters"]) == 2
- assert health["unhealthy_flags"] == []
-
-
-@pytest.mark.parametrize(
- "report",
- (
- "Volatile Uncorrectable ECC : 1\nReset Required : No\n",
- "Volatile Uncorrectable ECC : 0\nReset Required : Yes\n",
- "Reset Required : No\n",
- "Volatile Uncorrectable ECC : unknown\nReset Required : No\n",
- ),
-)
-def test_parse_ecc_health_fails_closed(report):
- with pytest.raises(RuntimeError):
- MODULE.parse_ecc_health(report)
-
-
-def test_probe_exercises_backward_and_one_rank_nccl_lifecycle():
- text = MODULE_PATH.read_text()
-
- assert "loss.backward()" in text
- assert "dist.init_process_group(" in text
- assert 'backend="nccl"' in text
- assert "dist.all_reduce(collective)" in text
- assert "dist.destroy_process_group()" in text
- assert '"world_size": world_size' in text
diff --git a/tests/test_ice_planar_bc_launcher.py b/tests/test_ice_planar_bc_launcher.py
deleted file mode 100644
index 44132b4f39..0000000000
--- a/tests/test_ice_planar_bc_launcher.py
+++ /dev/null
@@ -1,110 +0,0 @@
-import re
-import subprocess
-from pathlib import Path
-
-ROOT = Path(__file__).parents[1]
-LAUNCHER = ROOT / "scripts" / "ice" / "launch_planar_bc.sbatch"
-
-
-def test_launcher_is_valid_bash_and_has_single_gpu_requeue_contract():
- subprocess.run(["bash", "-n", str(LAUNCHER)], check=True)
- text = LAUNCHER.read_text()
-
- assert "#SBATCH --account=coc" in text
- assert "#SBATCH --partition=ice-gpu" in text
- assert "#SBATCH --qos=coc-ice" in text
- assert "#SBATCH --nodes=1" in text
- assert "#SBATCH --ntasks-per-node=1" in text
- assert "#SBATCH --gres=gpu:1" in text
- assert "#SBATCH --constraint=H100|H200" in text
- assert "#SBATCH --requeue" in text
- assert "#SBATCH --signal=B:USR1@600" in text
- assert "trainer.strategy=auto" in text
- assert "trainer.devices=1" in text
- assert "trainer.num_nodes=1" in text
- assert "trainer.precision=$ICE_PRECISION" in text
- assert "ICE_EXPECTED_GPU_NAME" not in text
- assert text.count("--allowed-gpu-name") == 2
-
-
-def test_launcher_uses_strict_external_runner_and_completion_sentinel():
- text = LAUNCHER.read_text()
-
- assert "runtime.slurm_requeue_owner=runner" in text
- assert "runtime.slurm_save_signal=SIGUSR2" in text
- assert "--checkpoint-validator" in text
- assert "--checkpoint-signal USR2" in text
- assert "--checkpoint-forwarding slurm-steps" in text
- assert "--requeue-owner runner" in text
- assert "--confirm-child-requeue-disabled" in text
- assert "ICE_OUTPUT_DIR/COMPLETE.json" in text
- assert "allow-unvalidated" not in text
- assert "ICE_RESUME_CHECKPOINT" in text
- assert "ICE_RESUME_CHECKPOINT_SHA256" in text
- assert "ICE_RESUME_GLOBAL_STEP" in text
-
-
-def test_parent_runner_inherits_exact_repo_pythonpath_for_checkpoint_validation():
- text = LAUNCHER.read_text()
-
- parent_export = text.index("export PYTHONPATH=$ICE_REPO")
- preflight_runner = text.index('"$ICE_PYTHON" "$ICE_REQUEUE_RUNNER"')
- final_runner = text.rindex('exec "$ICE_PYTHON" "$ICE_REQUEUE_RUNNER"')
- child_definition = text.index("read -r -d '' CHILD_CODE")
-
- assert parent_export < preflight_runner < child_definition < final_runner
-
-
-def test_launcher_exposes_supported_bc_configs_and_explicit_training_limits():
- text = LAUNCHER.read_text()
-
- assert "sim/pusht/planar_v2_usocket_direct_bc" in text
- assert "sim/pusht/planar_v2_usocket_arc_bc" in text
- assert "sim/pusht/planar_v2_usocket_arc_hybrid_D40_M16_R24deg_bc" in text
- assert "sim/pusht/planar_v2_usocket_arc_paper_uniform_D40_M16_R24deg" in text
- assert "ICE_VALIDATION_DEFERRED" in text
- assert "ICE_LIMIT_VAL_BATCHES must be zero when validation is deferred" in text
- assert '"++run_provenance.validation_deferred=$ICE_VALIDATION_DEFERRED"' in text
- assert (
- text.count("data.train_datasets.pushshapes_sim_u_socket.resolver.folder_path=")
- == 1
- )
- assert (
- text.count("data.valid_datasets.pushshapes_sim_u_socket.resolver.folder_path=")
- == 1
- )
- for value in (
- "trainer.max_steps=$ICE_MAX_STEPS",
- "trainer.val_check_interval=$ICE_VAL_CHECK_INTERVAL",
- "trainer.limit_train_batches=$ICE_LIMIT_TRAIN_BATCHES",
- "trainer.limit_val_batches=$ICE_LIMIT_VAL_BATCHES",
- "+callbacks.model_checkpoint.every_n_train_steps=$ICE_CHECKPOINT_EVERY_N_STEPS",
- ):
- assert value in text
- assert "unique ICE_OUTPUT_DIR already exists" in text
- assert "ICE_LAUNCH_MODE" in text
- assert "--dry-run" in text
- assert "validate_planar_dataset.py" in text
- assert "ICE_EXPECTED_SPLIT_MANIFEST_SHA256" in text
- assert "--expected-manifest-sha256" in text
- assert '"++paths.root_dir=$ICE_REPO"' in text
- assert '"norm_stats.save_cache_dir=$ICE_OUTPUT_DIR"' in text
- assert '"norm_stats.save_cache_dir=$ICE_OUTPUT_DIR/norm_stats"' not in text
- assert "$ICE_OUTPUT_DIR/norm_stats/norm_stats.json" in text
- assert "requeue requires the original run's norm_stats/norm_stats.json" in text
- assert "norm_stats.save_cache_dir=null" in text
-
-
-def test_launcher_contains_no_deployment_identity_or_credentials():
- text = LAUNCHER.read_text()
-
- forbidden = (
- "/coc/",
- "/storage/",
- "paphiwetsa3",
- "WANDB_API_KEY=",
- "wandb.ai/",
- )
- assert not any(value in text for value in forbidden)
- assert re.search(r"\b[0-9]{7}\b", text) is None
- assert "BASH_SOURCE" not in text
diff --git a/tests/test_ice_requeue_runner.py b/tests/test_ice_requeue_runner.py
deleted file mode 100644
index c652e216f3..0000000000
--- a/tests/test_ice_requeue_runner.py
+++ /dev/null
@@ -1,1476 +0,0 @@
-#!/usr/bin/env python3
-
-from __future__ import annotations
-
-import fcntl
-import hashlib
-import importlib.util
-import json
-import os
-import shlex
-import signal
-import subprocess
-import sys
-import tempfile
-import time
-import unittest
-from pathlib import Path
-from unittest import mock
-
-MODULE_PATH = Path(__file__).parents[1] / "scripts" / "ice" / "ice_requeue_runner.py"
-SPEC = importlib.util.spec_from_file_location("ice_requeue_runner", MODULE_PATH)
-MODULE = importlib.util.module_from_spec(SPEC)
-assert SPEC.loader is not None
-sys.modules[SPEC.name] = MODULE
-SPEC.loader.exec_module(MODULE)
-
-
-def make_validator(root: Path) -> Path:
- validator = root / "validate_checkpoint.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,os,sys\n"
- "log=os.environ.get('VALIDATOR_LOG')\n"
- "if log:\n"
- " with open(log,'a') as stream: stream.write(os.path.realpath(sys.argv[1])+'\\n')\n"
- "try:\n"
- " data=json.load(open(sys.argv[1]))\n"
- "except Exception:\n"
- " raise SystemExit(2)\n"
- "if not data.get('valid', True): raise SystemExit(3)\n"
- "metadata={'global_step': data['global_step']}\n"
- "for key in ('run_id','wandb_run_id','config_sha256','source_commit'):\n"
- " if key in data: metadata[key]=data[key]\n"
- "print(json.dumps(metadata, sort_keys=True))\n"
- )
- validator.chmod(0o700)
- return validator
-
-
-def write_checkpoint(
- path: Path, step: int, *, valid: bool = True, run_id: str = "run"
-) -> None:
- temporary = path.with_suffix(path.suffix + ".tmp")
- temporary.write_text(
- json.dumps(
- {
- "global_step": step,
- "valid": valid,
- "run_id": run_id,
- "wandb_run_id": "wandb-run",
- }
- )
- )
- os.replace(temporary, path)
-
-
-def make_signaling_scancel(root: Path) -> Path:
- executable = root / "scancel"
- executable.write_text(
- "#!/bin/sh\n"
- 'printf \'%s\\n\' "$*" > "$ICE_TEST_SCANCEL_LOG"\n'
- 'kill -USR2 "$(cat "$CHILD_PID_FILE")"\n'
- )
- executable.chmod(0o700)
- return executable
-
-
-def base_args(
- state_dir: Path,
- checkpoint_glob: str,
- validator: Path,
- command: list[str],
- *,
- owner: str = "child",
- initial_checkpoint: Path | None = None,
- scancel: Path | str = "/usr/bin/true",
-) -> list[str]:
- result = [
- "--state-dir",
- str(state_dir),
- ]
- if initial_checkpoint is not None:
- result.extend(["--initial-checkpoint", str(initial_checkpoint)])
- result.extend(
- [
- "--checkpoint-glob",
- checkpoint_glob,
- "--checkpoint-validator",
- str(validator),
- "--checkpoint-signal",
- "USR2",
- "--checkpoint-forwarding",
- "slurm-steps",
- "--requeue-owner",
- owner,
- "--scancel",
- str(scancel),
- "--poll-seconds",
- "0.05",
- ]
- )
- if owner == "runner":
- result.append("--confirm-child-requeue-disabled")
- result.extend(["--", *command])
- return result
-
-
-class CheckpointSelectionTest(unittest.TestCase):
- def test_semantic_global_step_beats_newer_mtime(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- validator = make_validator(root)
- step_60 = root / "epoch-6-step-60000.ckpt"
- step_40 = root / "staged-step-40000.ckpt"
- write_checkpoint(step_60, 60_000)
- write_checkpoint(step_40, 40_000)
- now = time.time()
- os.utime(step_60, (now - 20, now - 20))
- os.utime(step_40, (now, now))
- selected = MODULE.select_checkpoint("*.ckpt", root, validator)
- self.assertIsNotNone(selected)
- self.assertEqual(selected.path, step_60.resolve())
- self.assertEqual(selected.global_step, 60_000)
-
- def test_invalid_and_partial_candidates_are_ignored(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- validator = make_validator(root)
- valid = root / "step-10.ckpt"
- invalid = root / "step-20.ckpt"
- partial = root / "step-30.ckpt.partial"
- write_checkpoint(valid, 10)
- write_checkpoint(invalid, 20, valid=False)
- write_checkpoint(partial, 30)
- selected = MODULE.select_checkpoint("*.ckpt*", root, validator)
- self.assertEqual(selected.path, valid.resolve())
-
- def test_multiple_roots_select_highest_semantic_step(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- source = staged / "step-40000.ckpt"
- resumed = live / "step-42000.ckpt"
- write_checkpoint(source, 40_000)
- write_checkpoint(resumed, 42_000)
- selected = MODULE.select_checkpoint(
- [str(staged / "*.ckpt"), str(live / "*.ckpt")],
- root,
- validator,
- )
- self.assertEqual(selected.path, resumed.resolve())
-
- def test_fresh_checkpoint_requires_strictly_higher_global_step(self):
- def info(path: str, step: int, digest: str) -> MODULE.CheckpointInfo:
- return MODULE.CheckpointInfo(
- path=Path(path),
- sha256=digest,
- global_step=step,
- size=1,
- mtime_ns=1,
- inode=1,
- metadata={"global_step": step},
- )
-
- baseline = info("/tmp/run/a.ckpt", 10, "a" * 64)
- self.assertFalse(MODULE.is_fresh(None, baseline))
- self.assertFalse(
- MODULE.is_fresh(info("/tmp/run/b.ckpt", 9, "b" * 64), baseline)
- )
- self.assertFalse(
- MODULE.is_fresh(info("/tmp/run/b.ckpt", 10, "a" * 64), baseline)
- )
- self.assertFalse(
- MODULE.is_fresh(info("/tmp/run/a.ckpt", 10, "b" * 64), baseline)
- )
- self.assertTrue(
- MODULE.is_fresh(info("/tmp/run/a.ckpt", 11, "a" * 64), baseline)
- )
-
- def test_semantic_maximum_same_step_fork_is_rejected(self):
- def info(path: str, digest: str) -> MODULE.CheckpointInfo:
- return MODULE.CheckpointInfo(
- path=Path(path),
- sha256=digest,
- global_step=10,
- size=1,
- mtime_ns=1,
- inode=1,
- metadata={"global_step": 10},
- )
-
- with self.assertRaisesRegex(SystemExit, "semantic-maximum checkpoints"):
- MODULE.newest_checkpoint(
- info("/tmp/run/a.ckpt", "a" * 64),
- info("/tmp/run/b.ckpt", "b" * 64),
- )
-
- def test_checkpoint_forwarding_is_exact_scancel_and_never_killpg(self):
- completed = subprocess.CompletedProcess([], 0, stdout="", stderr="")
- with mock.patch.dict(
- os.environ,
- {"SCANCEL_BATCH": "1", "SCANCEL_FULL": "1", "SLURM_CLUSTERS": "wrong"},
- ):
- with (
- mock.patch.object(
- MODULE.subprocess, "run", return_value=completed
- ) as run,
- mock.patch.object(MODULE.os, "killpg") as killpg,
- ):
- result = MODULE.forward_checkpoint_to_slurm_steps(
- "/mock/bin/scancel",
- "12345",
- signal.SIGUSR2,
- )
- self.assertIs(result, completed)
- run.assert_called_once()
- positional, keyword = run.call_args
- self.assertEqual(positional, (["/mock/bin/scancel", "--signal=USR2", "12345"],))
- self.assertEqual(keyword["text"], True)
- self.assertIs(keyword["stdout"], subprocess.PIPE)
- self.assertIs(keyword["stderr"], subprocess.PIPE)
- self.assertEqual(keyword["check"], False)
- self.assertNotIn("SCANCEL_BATCH", keyword["env"])
- self.assertNotIn("SCANCEL_FULL", keyword["env"])
- self.assertNotIn("SLURM_CLUSTERS", keyword["env"])
- killpg.assert_not_called()
-
-
-class RunnerSafetyTest(unittest.TestCase):
- def test_signal_handlers_are_installed_before_child_spawn(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoint_path = root / "step-1.ckpt"
- write_checkpoint(checkpoint_path, 1)
- checkpoint = MODULE.stable_checkpoint_info(
- checkpoint_path.resolve(), validator
- )
- self.assertIsNotNone(checkpoint)
- installed: list[signal.Signals] = []
- gate_readers: list[int] = []
-
- class FinishedChild:
- pid = 9876
-
- @staticmethod
- def poll() -> int:
- while gate_readers:
- os.close(gate_readers.pop())
- return 0
-
- def install(received: signal.Signals, _handler: object) -> object:
- installed.append(received)
- return signal.SIG_DFL
-
- def spawn(*_args: object, **kwargs: object) -> FinishedChild:
- self.assertEqual(
- installed,
- [signal.SIGUSR1, signal.SIGTERM, signal.SIGINT],
- )
- pass_fds = kwargs["pass_fds"]
- self.assertIsInstance(pass_fds, tuple)
- gate_readers.append(os.dup(pass_fds[0]))
- return FinishedChild()
-
- def validate(*_args: object, **_kwargs: object) -> MODULE.CheckpointInfo:
- self.assertEqual(
- installed,
- [signal.SIGUSR1, signal.SIGTERM, signal.SIGINT],
- )
- return checkpoint
-
- args = base_args(state, str(root / "*.ckpt"), validator, ["unused-child"])
- with (
- mock.patch.dict(
- os.environ,
- {"SLURM_JOB_ID": "990", "SLURM_RESTART_COUNT": "0"},
- clear=False,
- ),
- mock.patch.object(MODULE.signal, "signal", side_effect=install),
- mock.patch.object(
- MODULE, "stable_checkpoint_info", side_effect=validate
- ),
- mock.patch.object(
- MODULE.subprocess, "Popen", side_effect=spawn
- ) as popen,
- ):
- self.assertEqual(MODULE.main(args), 0)
- popen.assert_called_once()
-
- def test_usr1_during_strict_discovery_requeues_without_starting_work(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- checkpoint = root / "step-1.ckpt"
- write_checkpoint(checkpoint, 1)
- early_signal_marker = root / "early-signal-sent"
- validator = root / "early_signal_validator.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,os,signal,sys,time\n"
- "marker=os.environ['EARLY_SIGNAL_MARKER']\n"
- "if not os.path.exists(marker):\n"
- " open(marker,'w').write('sent')\n"
- " os.kill(os.getppid(),signal.SIGUSR1)\n"
- " time.sleep(0.1)\n"
- "with open(sys.argv[1]) as stream: data=json.load(stream)\n"
- "print(json.dumps({'global_step':data['global_step'],'run_id':data['run_id']}))\n"
- )
- validator.chmod(0o700)
- child_started = root / "child-started"
- child = root / "child.py"
- child.write_text(
- "import os\n" "open(os.environ['CHILD_STARTED'],'w').write('unsafe')\n"
- )
- attempt = state / "requeue" / "job-12400-restart-000" / "attempt.json"
- scancel_log = root / "scancel.log"
- fake_scancel = root / "scancel"
- fake_scancel.write_text(
- "#!/bin/sh\n" 'printf \'%s\\n\' "$*" > "$ICE_TEST_SCANCEL_LOG"\n'
- )
- fake_scancel.chmod(0o700)
- scontrol_log = root / "scontrol.log"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text(
- "#!/bin/sh\n" 'printf \'%s\\n\' "$*" > "$SCONTROL_LOG"\n'
- )
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12400",
- "SLURM_RESTART_COUNT": "0",
- "EARLY_SIGNAL_MARKER": str(early_signal_marker),
- "CHILD_STARTED": str(child_started),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- "SCONTROL_LOG": str(scontrol_log),
- }
- )
- args = base_args(
- state,
- str(root / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="runner",
- scancel=fake_scancel,
- )
- args[args.index("--") : args.index("--")] = [
- "--scontrol",
- str(fake_scontrol),
- ]
- result = subprocess.run(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- capture_output=True,
- timeout=10,
- check=False,
- )
- self.assertEqual(result.returncode, 0, result.stderr)
- self.assertTrue(early_signal_marker.is_file())
- self.assertFalse(child_started.exists())
- self.assertFalse(scancel_log.exists())
- self.assertEqual(scontrol_log.read_text().strip(), "requeue 12400")
- final = json.loads(attempt.read_text())
- self.assertEqual(final["status"], "EARLY_BOUNDARY_REQUEUE_ACCEPTED")
- self.assertNotIn("child_pid", final)
- proof_path = Path(final["early_boundary_proof"])
- proof = json.loads(proof_path.read_text())
- self.assertEqual(
- final["early_boundary_proof_sha256"],
- hashlib.sha256(proof_path.read_bytes()).hexdigest(),
- )
- self.assertEqual(
- proof["status"],
- "EARLY_BOUNDARY_NO_WORK_CHECKPOINT_AUTHENTICATED",
- )
- self.assertFalse(proof["child_spawned"])
- self.assertFalse(proof["optimizer_work_performed"])
- self.assertFalse(proof["checkpoint_signal_forwarded"])
- self.assertEqual(proof["checkpoint"]["global_step"], 1)
- events = [
- json.loads(line)
- for line in (attempt.parent / "events.jsonl").read_text().splitlines()
- ]
- self.assertEqual(
- [event["event"] for event in events].count(
- "early_boundary_no_work_checkpoint_authenticated"
- ),
- 1,
- )
-
- def test_incremental_restart_revalidates_high_water_and_only_new_candidates(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoints = [root / f"step-{step}.ckpt" for step in (10, 20, 30)]
- for step, checkpoint in zip((10, 20, 30), checkpoints):
- write_checkpoint(checkpoint, step)
- log = root / "validator.log"
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
-
- with mock.patch.dict(
- os.environ,
- {"SLURM_RESTART_COUNT": "0", "VALIDATOR_LOG": str(log)},
- clear=False,
- ):
- self.assertEqual(MODULE.main(args), 0)
- first = log.read_text().splitlines()
- self.assertEqual(
- set(first), {str(path.resolve()) for path in checkpoints}
- )
- self.assertEqual(len(first), 3)
-
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertEqual(
- log.read_text().splitlines(), [str(checkpoints[-1].resolve())]
- )
-
- newer = root / "step-40.ckpt"
- write_checkpoint(newer, 40)
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertEqual(
- set(log.read_text().splitlines()),
- {str(checkpoints[-1].resolve()), str(newer.resolve())},
- )
-
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertEqual(log.read_text().splitlines(), [str(newer.resolve())])
-
- def test_unbound_observation_index_forces_validator_authoritative_rescan(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoints = [root / f"step-{step}.ckpt" for step in (10, 20, 30)]
- for step, checkpoint in zip((10, 20, 30), checkpoints):
- write_checkpoint(checkpoint, step)
- log = root / "validator.log"
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- with mock.patch.dict(
- os.environ,
- {"SLURM_RESTART_COUNT": "0", "VALIDATOR_LOG": str(log)},
- clear=False,
- ):
- self.assertEqual(MODULE.main(args), 0)
- index = state / MODULE.OBSERVATION_INDEX_NAME
- index.write_text(index.read_text() + "\n")
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertEqual(
- set(log.read_text().splitlines()),
- {str(path.resolve()) for path in checkpoints},
- )
-
- def test_validator_change_invalidates_observation_index(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoints = [root / f"step-{step}.ckpt" for step in (10, 20)]
- for step, checkpoint in zip((10, 20), checkpoints):
- write_checkpoint(checkpoint, step)
- log = root / "validator.log"
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- with mock.patch.dict(
- os.environ,
- {"SLURM_RESTART_COUNT": "0", "VALIDATOR_LOG": str(log)},
- clear=False,
- ):
- self.assertEqual(MODULE.main(args), 0)
- validator.write_text(validator.read_text() + "# authority revision\n")
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertEqual(
- set(log.read_text().splitlines()),
- {str(path.resolve()) for path in checkpoints},
- )
-
- def test_validator_change_during_validation_refuses_state_commit(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- checkpoint = root / "step-10.ckpt"
- write_checkpoint(checkpoint, 10)
- marker = root / "validator-mutated"
- validator = root / "mutating_validator.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,os,sys\n"
- "marker=os.environ['VALIDATOR_MUTATION_MARKER']\n"
- "if not os.path.exists(marker):\n"
- " open(marker,'w').write('mutated')\n"
- " with open(__file__,'a') as stream: stream.write('# changed\\n')\n"
- "with open(sys.argv[1]) as stream: data=json.load(stream)\n"
- "print(json.dumps({'global_step':data['global_step'],'run_id':data['run_id']}))\n"
- )
- validator.chmod(0o700)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- with mock.patch.dict(
- os.environ,
- {
- "SLURM_RESTART_COUNT": "0",
- "VALIDATOR_MUTATION_MARKER": str(marker),
- },
- clear=False,
- ):
- with self.assertRaisesRegex(
- SystemExit, "validator changed after validation"
- ):
- MODULE.main(args)
- self.assertTrue(marker.is_file())
- self.assertFalse((state / MODULE.OBSERVATION_INDEX_NAME).exists())
- self.assertFalse((state / "checkpoint-high-water.json").exists())
-
- def test_failed_candidate_is_not_cached_when_sidecar_appears(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = root / "sidecar_validator.py"
- validator.write_text(
- "#!/usr/bin/env python3\n"
- "import json,os,sys\n"
- "with open(os.environ['VALIDATOR_LOG'],'a') as stream:\n"
- " stream.write(os.path.realpath(sys.argv[1])+'\\n')\n"
- "if not os.path.isfile(sys.argv[1]+'.proof'): raise SystemExit(4)\n"
- "with open(sys.argv[1]) as stream: data=json.load(stream)\n"
- "print(json.dumps({'global_step':data['global_step'],'run_id':data['run_id']}))\n"
- )
- validator.chmod(0o700)
- lower = root / "step-10.ckpt"
- higher = root / "step-20.ckpt"
- write_checkpoint(lower, 10)
- write_checkpoint(higher, 20)
- Path(str(lower) + ".proof").write_text("ready")
- original_higher_stat = higher.stat()
- log = root / "validator.log"
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- with mock.patch.dict(
- os.environ,
- {"SLURM_RESTART_COUNT": "0", "VALIDATOR_LOG": str(log)},
- clear=False,
- ):
- self.assertEqual(MODULE.main(args), 0)
- first = json.loads(
- (
- state / "requeue" / "job-dry-run-restart-000" / "attempt.json"
- ).read_text()
- )
- self.assertEqual(first["checkpoint"]["global_step"], 10)
-
- Path(str(higher) + ".proof").write_text("ready")
- self.assertEqual(
- MODULE.stat_identity(higher.stat()),
- MODULE.stat_identity(original_higher_stat),
- )
- log.write_text("")
- self.assertEqual(MODULE.main(args), 0)
- self.assertIn(str(higher.resolve()), log.read_text().splitlines())
- final = json.loads(
- (
- state / "requeue" / "job-dry-run-restart-000" / "attempt.json"
- ).read_text()
- )
- self.assertEqual(final["checkpoint"]["global_step"], 20)
-
- def test_atomic_replacement_with_preserved_size_and_mtime_is_revalidated(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- historical = root / "step-10.ckpt"
- high_water = root / "step-20.ckpt"
- write_checkpoint(historical, 10)
- write_checkpoint(high_water, 20)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "0"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
-
- original = historical.stat()
- replacement = root / "replacement.ckpt"
- write_checkpoint(replacement, 30)
- self.assertEqual(replacement.stat().st_size, original.st_size)
- os.utime(replacement, ns=(original.st_atime_ns, original.st_mtime_ns))
- os.replace(replacement, historical)
- replaced = historical.stat()
- self.assertEqual(replaced.st_size, original.st_size)
- self.assertEqual(replaced.st_mtime_ns, original.st_mtime_ns)
- self.assertNotEqual(
- (replaced.st_dev, replaced.st_ino, replaced.st_ctime_ns),
- (original.st_dev, original.st_ino, original.st_ctime_ns),
- )
-
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "1"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
- final = json.loads(
- (
- state / "requeue" / "job-dry-run-restart-001" / "attempt.json"
- ).read_text()
- )
- self.assertEqual(final["checkpoint"]["global_step"], 30)
-
- def test_runner_ownership_requires_explicit_child_disable_confirmation(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- args = base_args(
- state,
- str(root / "*.ckpt"),
- validator,
- ["/bin/true"],
- owner="runner",
- )
- args.remove("--confirm-child-requeue-disabled")
- args.insert(args.index("--"), "--dry-run")
- with self.assertRaisesRegex(SystemExit, "disable framework auto-requeue"):
- MODULE.main(args)
-
- def test_required_initial_checkpoint_is_refused(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(-2, "--require-initial-checkpoint")
- args.insert(-2, "--dry-run")
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "0"}, clear=False):
- with self.assertRaisesRegex(
- SystemExit, "requires an exact --initial-checkpoint"
- ):
- MODULE.main(args)
-
- def test_exact_initial_pins_first_attempt_then_live_wins_restart(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-40000.ckpt"
- newer_live = live / "step-60000.ckpt"
- write_checkpoint(initial, 40_000)
- write_checkpoint(newer_live, 60_000)
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- ["/bin/true"],
- initial_checkpoint=initial,
- )
- args.insert(args.index("--"), "--require-initial-checkpoint")
- args.insert(args.index("--"), "--dry-run")
-
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "0"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
- first = json.loads(
- (
- state / "requeue" / "job-dry-run-restart-000" / "attempt.json"
- ).read_text()
- )
- self.assertEqual(first["checkpoint"]["global_step"], 40_000)
- self.assertEqual(first["checkpoint"]["path"], str(initial.resolve()))
-
- initial.unlink()
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "1"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
- restarted = json.loads(
- (
- state / "requeue" / "job-dry-run-restart-001" / "attempt.json"
- ).read_text()
- )
- self.assertEqual(restarted["checkpoint"]["global_step"], 60_000)
- self.assertEqual(restarted["checkpoint"]["path"], str(newer_live.resolve()))
-
- def test_explicit_missing_initial_checkpoint_fails_closed(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- args = base_args(
- state,
- str(root / "live" / "*.ckpt"),
- validator,
- ["/bin/true"],
- initial_checkpoint=root / "staged" / "missing.ckpt",
- )
- args.insert(args.index("--"), "--dry-run")
- with self.assertRaisesRegex(
- SystemExit, "initial checkpoint is unavailable"
- ):
- MODULE.main(args)
-
- def test_restart_without_checkpoint_is_always_refused(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(-2, "--dry-run")
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "1"}, clear=False):
- with self.assertRaisesRegex(SystemExit, "required resume checkpoint"):
- MODULE.main(args)
-
- def test_high_water_refuses_rollback(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- newest = root / "step-60000.ckpt"
- write_checkpoint(newest, 60_000)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(-2, "--dry-run")
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "0"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
- newest.unlink()
- write_checkpoint(root / "step-40000.ckpt", 40_000)
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "1"}, clear=False):
- with self.assertRaisesRegex(SystemExit, "rollback refused"):
- MODULE.main(args)
-
- def test_high_water_refuses_newer_same_step_different_sha(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoint = root / "step-60000.ckpt"
- write_checkpoint(checkpoint, 60_000)
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(-2, "--dry-run")
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "0"}, clear=False):
- self.assertEqual(MODULE.main(args), 0)
- time.sleep(0.01)
- checkpoint.unlink()
- checkpoint.write_text(
- json.dumps(
- {
- "global_step": 60_000,
- "valid": True,
- "run_id": "run",
- "wandb_run_id": "wandb-run",
- "different_payload": True,
- }
- )
- )
- with mock.patch.dict(os.environ, {"SLURM_RESTART_COUNT": "1"}, clear=False):
- with self.assertRaisesRegex(SystemExit, "identity fork refused"):
- MODULE.main(args)
-
- def test_state_directory_lock_refuses_duplicate_runner(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- state.mkdir(parents=True)
- validator = make_validator(root)
- write_checkpoint(root / "step-1.ckpt", 1)
- lock = (state / "runner.lock").open("a+")
- fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
- try:
- args = base_args(state, str(root / "*.ckpt"), validator, ["/bin/true"])
- args.insert(args.index("--"), "--dry-run")
- command = [
- sys.executable,
- str(MODULE_PATH),
- *args,
- ]
- result = subprocess.run(
- command, text=True, capture_output=True, check=False
- )
- finally:
- lock.close()
- self.assertNotEqual(result.returncode, 0)
- self.assertIn("another requeue runner", result.stderr)
-
- def test_exports_checkpoint_path_sha_step_and_metadata(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- checkpoint = root / "step-12.ckpt"
- write_checkpoint(checkpoint, 12)
- output = root / "environment.json"
- child = root / "child.py"
- child.write_text(
- "import json,os\n"
- "keys=['ICE_RESUME_CHECKPOINT','ICE_RESUME_CHECKPOINT_SHA256',"
- "'ICE_RESUME_GLOBAL_STEP','ICE_RESUME_CHECKPOINT_METADATA_JSON',"
- "'ICE_REQUEUE_OWNER']\n"
- "open(os.environ['OUTPUT'],'w').write(json.dumps({k:os.environ[k] for k in keys}))\n"
- )
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "991",
- "SLURM_RESTART_COUNT": "0",
- "OUTPUT": str(output),
- }
- )
- command = [
- sys.executable,
- str(MODULE_PATH),
- *base_args(
- state, str(root / "*.ckpt"), validator, [sys.executable, str(child)]
- ),
- ]
- result = subprocess.run(
- command, env=env, text=True, capture_output=True, check=False
- )
- self.assertEqual(result.returncode, 0, result.stderr)
- exported = json.loads(output.read_text())
- self.assertEqual(
- exported["ICE_RESUME_CHECKPOINT"], str(checkpoint.resolve())
- )
- self.assertEqual(exported["ICE_RESUME_GLOBAL_STEP"], "12")
- self.assertEqual(
- exported["ICE_RESUME_CHECKPOINT_SHA256"],
- hashlib.sha256(checkpoint.read_bytes()).hexdigest(),
- )
- metadata = json.loads(exported["ICE_RESUME_CHECKPOINT_METADATA_JSON"])
- self.assertEqual(metadata["global_step"], 12)
- self.assertEqual(metadata["checkpoint_path"], str(checkpoint.resolve()))
- self.assertEqual(exported["ICE_REQUEUE_OWNER"], "child")
-
- def test_child_signal_has_shell_exit_status_without_requeue(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- write_checkpoint(root / "step-12.ckpt", 12)
- scontrol_log = root / "scontrol.log"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text(
- "#!/bin/sh\n" f"touch {shlex.quote(str(scontrol_log))}\n"
- )
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update({"SLURM_JOB_ID": "993", "SLURM_RESTART_COUNT": "0"})
- args = base_args(
- state,
- str(root / "*.ckpt"),
- validator,
- [
- sys.executable,
- "-c",
- "import os,signal; os.kill(os.getpid(), signal.SIGTERM)",
- ],
- owner="runner",
- )
- args[args.index("--") : args.index("--")] = [
- "--scontrol",
- str(fake_scontrol),
- ]
- result = subprocess.run(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- capture_output=True,
- check=False,
- timeout=10,
- )
- self.assertEqual(result.returncode, 128 + signal.SIGTERM, result.stderr)
- self.assertFalse(scontrol_log.exists())
- attempt = json.loads(
- (state / "requeue" / "job-993-restart-000" / "attempt.json").read_text()
- )
- self.assertEqual(attempt["status"], "CHILD_FAILED_NO_REQUEUE")
- self.assertEqual(attempt["child_exit_code"], -signal.SIGTERM)
- self.assertEqual(attempt["child_exit_signal"], "SIGTERM")
-
- def test_signal_baseline_waits_for_post_signal_checkpoint(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- ready = root / "ready"
- child_pid = root / "child.pid"
- resume_seen = root / "resume-seen"
- child = root / "child.py"
- child.write_text(
- "import json,os,signal,time\n"
- "root=os.environ['ROOT']\n"
- "def write(step):\n"
- " p=os.path.join(root,f'step-{step}.ckpt'); t=p+'.tmp'\n"
- " open(t,'w').write(json.dumps({'global_step':step,'run_id':'run','wandb_run_id':'wandb-run'}))\n"
- " os.replace(t,p)\n"
- "def save(*_):\n"
- " time.sleep(0.6); write(3); raise SystemExit(0)\n"
- "signal.signal(signal.SIGUSR2,save)\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "open(os.environ['RESUME_SEEN'],'w').write(os.environ['ICE_RESUME_CHECKPOINT'])\n"
- "write(2); open(os.environ['READY'],'w').write('ready')\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = make_signaling_scancel(root)
- scontrol_log = root / "scontrol.log"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text(
- '#!/bin/sh\nprintf \'%s\\n\' "$*" > "$SCONTROL_LOG"\n'
- )
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12345",
- "SLURM_RESTART_COUNT": "0",
- "ROOT": str(live),
- "READY": str(ready),
- "CHILD_PID_FILE": str(child_pid),
- "RESUME_SEEN": str(resume_seen),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- "SCONTROL_LOG": str(scontrol_log),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="runner",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- args[args.index("--poll-seconds") : args.index("--poll-seconds") + 2] = [
- "--poll-seconds",
- "0.02",
- ]
- args[args.index("--") : args.index("--")] = [
- "--checkpoint-grace-seconds",
- "4",
- "--scontrol",
- str(fake_scontrol),
- ]
- command = [sys.executable, str(MODULE_PATH), *args]
- process = subprocess.Popen(
- command,
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12345-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if (
- ready.exists()
- and attempt.exists()
- and json.loads(attempt.read_text()).get("status") == "RUNNING"
- ):
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner/child did not become ready")
- process.send_signal(signal.SIGUSR1)
- time.sleep(0.2)
- self.assertFalse(
- scontrol_log.exists(), "runner requeued before post-signal save"
- )
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(process.returncode, 0, f"{stdout}\n{stderr}")
- self.assertEqual(resume_seen.read_text(), str(initial.resolve()))
- self.assertEqual(scancel_log.read_text().strip(), "--signal=USR2 12345")
- self.assertEqual(scontrol_log.read_text().strip(), "requeue 12345")
- final = json.loads(attempt.read_text())
- self.assertEqual(final["checkpoint"]["global_step"], 3)
- self.assertTrue(final["fresh_checkpoint_after_signal"])
- events = [
- json.loads(line)
- for line in (attempt.parent / "events.jsonl").read_text().splitlines()
- ]
- signal_event = next(event for event in events if event["event"] == "signal")
- self.assertEqual(signal_event["baseline"]["global_step"], 2)
-
- def test_runner_refuses_requeue_without_fresh_post_signal_checkpoint(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- ready = root / "ready"
- child_pid = root / "child.pid"
- child = root / "child.py"
- child.write_text(
- "import os,signal,time\n"
- "signal.signal(signal.SIGUSR2, lambda *_: (_ for _ in ()).throw(SystemExit(0)))\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "open(os.environ['READY'],'w').write('ready')\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = make_signaling_scancel(root)
- called = root / "scontrol-called"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text('#!/bin/sh\ntouch "$SCONTROL_CALLED"\n')
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12348",
- "SLURM_RESTART_COUNT": "0",
- "READY": str(ready),
- "CHILD_PID_FILE": str(child_pid),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- "SCONTROL_CALLED": str(called),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="runner",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- args[args.index("--") : args.index("--")] = [
- "--checkpoint-grace-seconds",
- "1",
- "--scontrol",
- str(fake_scontrol),
- ]
- process = subprocess.Popen(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12348-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if (
- ready.exists()
- and attempt.exists()
- and json.loads(attempt.read_text()).get("status") == "RUNNING"
- ):
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner/child did not become ready")
- process.send_signal(signal.SIGUSR1)
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(process.returncode, 74, f"{stdout}\n{stderr}")
- self.assertFalse(called.exists())
- self.assertEqual(scancel_log.read_text().strip(), "--signal=USR2 12348")
- final = json.loads(attempt.read_text())
- self.assertEqual(final["status"], "REQUEUE_REFUSED_NO_FRESH_CHECKPOINT")
- self.assertFalse(final["fresh_checkpoint_after_signal"])
-
- def test_child_requeue_owner_never_calls_scontrol(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- child_pid = root / "child.pid"
- child = root / "child.py"
- child.write_text(
- "import json,os,signal,time\n"
- "def save(*_):\n"
- " p=os.path.join(os.environ['ROOT'],'step-2.ckpt'); t=p+'.tmp'\n"
- " open(t,'w').write(json.dumps({'global_step':2,'run_id':'run','wandb_run_id':'wandb-run'}))\n"
- " os.replace(t,p); raise SystemExit(0)\n"
- "signal.signal(signal.SIGUSR2,save)\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = make_signaling_scancel(root)
- called = root / "scontrol-called"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text('#!/bin/sh\ntouch "$SCONTROL_CALLED"\n')
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12346",
- "SLURM_RESTART_COUNT": "0",
- "ROOT": str(live),
- "CHILD_PID_FILE": str(child_pid),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- "SCONTROL_CALLED": str(called),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="child",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- args[args.index("--") : args.index("--")] = [
- "--checkpoint-grace-seconds",
- "3",
- "--scontrol",
- str(fake_scontrol),
- ]
- process = subprocess.Popen(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12346-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if (
- child_pid.exists()
- and attempt.exists()
- and json.loads(attempt.read_text()).get("status") == "RUNNING"
- ):
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner did not become ready")
- process.send_signal(signal.SIGUSR1)
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(process.returncode, 0, f"{stdout}\n{stderr}")
- self.assertFalse(called.exists())
- self.assertEqual(scancel_log.read_text().strip(), "--signal=USR2 12346")
- final = json.loads(attempt.read_text())
- self.assertEqual(final["status"], "COMPLETE")
-
- def test_duplicate_usr1_boundary_invokes_scancel_only_once(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- ready = root / "ready"
- saved = root / "saved"
- child_pid = root / "child.pid"
- child = root / "child.py"
- child.write_text(
- "import json,os,signal,time\n"
- "def save(*_):\n"
- " p=os.path.join(os.environ['ROOT'],'step-2.ckpt'); t=p+'.tmp'\n"
- " open(t,'w').write(json.dumps({'global_step':2,'run_id':'run','wandb_run_id':'wandb-run'}))\n"
- " os.replace(t,p); open(os.environ['SAVED'],'w').write('saved')\n"
- "signal.signal(signal.SIGUSR2,save)\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "open(os.environ['READY'],'w').write('ready')\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = root / "scancel"
- fake_scancel.write_text(
- "#!/bin/sh\n"
- 'printf \'%s\\n\' "$*" >> "$ICE_TEST_SCANCEL_LOG"\n'
- 'kill -USR2 "$(cat "$CHILD_PID_FILE")"\n'
- )
- fake_scancel.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12350",
- "SLURM_RESTART_COUNT": "0",
- "ROOT": str(live),
- "READY": str(ready),
- "SAVED": str(saved),
- "CHILD_PID_FILE": str(child_pid),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="child",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- process = subprocess.Popen(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12350-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if ready.exists() and attempt.exists():
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner/child did not become ready")
- process.send_signal(signal.SIGUSR1)
- deadline = time.time() + 5
- while time.time() < deadline:
- if (
- saved.exists()
- and json.loads(attempt.read_text()).get("status")
- == "CHILD_REQUEUE_OWNER_CHECKPOINT_READY"
- ):
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("first checkpoint boundary did not complete")
- process.send_signal(signal.SIGUSR1)
- time.sleep(0.2)
- self.assertEqual(
- scancel_log.read_text().splitlines(), ["--signal=USR2 12350"]
- )
- process.send_signal(signal.SIGTERM)
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(
- process.returncode, 128 + signal.SIGTERM, f"{stdout}\n{stderr}"
- )
-
- def test_scancel_failure_is_persisted_and_never_requeues(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- ready = root / "ready"
- child_pid = root / "child.pid"
- child = root / "child.py"
- child.write_text(
- "import os,time\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "open(os.environ['READY'],'w').write('ready')\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = root / "scancel"
- fake_scancel.write_text(
- "#!/bin/sh\n"
- 'printf \'%s\\n\' "$*" > "$ICE_TEST_SCANCEL_LOG"\n'
- 'kill -TERM "$(cat "$CHILD_PID_FILE")"\n'
- "echo signal-refused >&2\n"
- "exit 9\n"
- )
- fake_scancel.chmod(0o700)
- scontrol_called = root / "scontrol-called"
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text('#!/bin/sh\ntouch "$SCONTROL_CALLED"\n')
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12349",
- "SLURM_RESTART_COUNT": "0",
- "READY": str(ready),
- "CHILD_PID_FILE": str(child_pid),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- "SCONTROL_CALLED": str(scontrol_called),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="runner",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- args[args.index("--") : args.index("--")] = [
- "--scontrol",
- str(fake_scontrol),
- ]
- process = subprocess.Popen(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12349-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if ready.exists() and attempt.exists():
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner/child did not become ready")
- process.send_signal(signal.SIGUSR1)
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(process.returncode, 9, f"{stdout}\n{stderr}")
- self.assertEqual(scancel_log.read_text().strip(), "--signal=USR2 12349")
- self.assertFalse(scontrol_called.exists())
- final = json.loads(attempt.read_text())
- self.assertEqual(final["status"], "CHECKPOINT_SIGNAL_FORWARD_FAILED")
- self.assertEqual(final["checkpoint_signal_returncode"], 9)
- self.assertEqual(
- final["checkpoint_signal_argv"],
- [str(fake_scancel.resolve()), "--signal=USR2", "12349"],
- )
- self.assertIn("signal-refused", final["checkpoint_signal_stderr"])
-
- def test_scontrol_failure_is_persisted(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- staged = root / "staged"
- live = root / "live"
- staged.mkdir()
- live.mkdir()
- validator = make_validator(root)
- initial = staged / "step-1.ckpt"
- write_checkpoint(initial, 1)
- child_pid = root / "child.pid"
- child = root / "child.py"
- child.write_text(
- "import json,os,signal,time\n"
- "def save(*_):\n"
- " p=os.path.join(os.environ['ROOT'],'step-2.ckpt'); t=p+'.tmp'\n"
- " open(t,'w').write(json.dumps({'global_step':2,'run_id':'run','wandb_run_id':'wandb-run'}))\n"
- " os.replace(t,p); raise SystemExit(0)\n"
- "signal.signal(signal.SIGUSR2,save)\n"
- "open(os.environ['CHILD_PID_FILE'],'w').write(str(os.getpid()))\n"
- "while True: time.sleep(0.05)\n"
- )
- scancel_log = root / "scancel.log"
- fake_scancel = make_signaling_scancel(root)
- fake_scontrol = root / "scontrol"
- fake_scontrol.write_text("#!/bin/sh\necho refused >&2\nexit 7\n")
- fake_scontrol.chmod(0o700)
- env = os.environ.copy()
- env.update(
- {
- "SLURM_JOB_ID": "12347",
- "SLURM_RESTART_COUNT": "0",
- "ROOT": str(live),
- "CHILD_PID_FILE": str(child_pid),
- "ICE_TEST_SCANCEL_LOG": str(scancel_log),
- }
- )
- args = base_args(
- state,
- str(live / "*.ckpt"),
- validator,
- [sys.executable, str(child)],
- owner="runner",
- initial_checkpoint=initial,
- scancel=fake_scancel,
- )
- args[args.index("--") : args.index("--")] = [
- "--checkpoint-grace-seconds",
- "3",
- "--scontrol",
- str(fake_scontrol),
- ]
- process = subprocess.Popen(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- attempt = state / "requeue" / "job-12347-restart-000" / "attempt.json"
- deadline = time.time() + 5
- while time.time() < deadline:
- if (
- child_pid.exists()
- and attempt.exists()
- and json.loads(attempt.read_text()).get("status") == "RUNNING"
- ):
- break
- time.sleep(0.02)
- else:
- process.kill()
- self.fail("runner did not become ready")
- process.send_signal(signal.SIGUSR1)
- stdout, stderr = process.communicate(timeout=10)
- self.assertEqual(process.returncode, 7, f"{stdout}\n{stderr}")
- final = json.loads(attempt.read_text())
- self.assertEqual(final["status"], "REQUEUE_FAILED")
- self.assertEqual(final["scontrol_returncode"], 7)
- self.assertIn("refused", final["scontrol_stderr"])
-
- def test_completion_sentinel_contains_terminal_checkpoint(self):
- with tempfile.TemporaryDirectory() as raw:
- root = Path(raw)
- state = root / "state" / "run"
- validator = make_validator(root)
- write_checkpoint(root / "step-5.ckpt", 5)
- sentinel = root / "run" / "COMPLETE.json"
- child = root / "child.py"
- child.write_text(
- "import json,os\n"
- "p=os.path.join(os.environ['ROOT'],'step-6.ckpt'); t=p+'.tmp'\n"
- "open(t,'w').write(json.dumps({'global_step':6,'run_id':'run','wandb_run_id':'wandb-run'}))\n"
- "os.replace(t,p)\n"
- )
- env = os.environ.copy()
- env.update(
- {"SLURM_JOB_ID": "992", "SLURM_RESTART_COUNT": "0", "ROOT": str(root)}
- )
- args = base_args(
- state, str(root / "*.ckpt"), validator, [sys.executable, str(child)]
- )
- args[args.index("--") : args.index("--")] = [
- "--completion-sentinel",
- str(sentinel),
- ]
- result = subprocess.run(
- [sys.executable, str(MODULE_PATH), *args],
- env=env,
- text=True,
- capture_output=True,
- check=False,
- )
- self.assertEqual(result.returncode, 0, result.stderr)
- payload = json.loads(sentinel.read_text())
- self.assertEqual(payload["status"], "COMPLETE")
- self.assertEqual(payload["checkpoint"]["global_step"], 6)
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_ice_wandb_resume.py b/tests/test_ice_wandb_resume.py
deleted file mode 100644
index cf04f0d74c..0000000000
--- a/tests/test_ice_wandb_resume.py
+++ /dev/null
@@ -1,118 +0,0 @@
-from pathlib import Path
-
-import pytest
-from lightning import LightningModule
-from omegaconf import OmegaConf
-
-from egomimic.pl_utils.pl_model import ModelWrapper
-from egomimic.utils.logging_utils import configure_runner_wandb
-
-
-def _cfg(tmp_path: Path, *, run_id="stable-run", checkpoint=None):
- return OmegaConf.create(
- {
- "paths": {"output_dir": str(tmp_path)},
- "ckpt_path": checkpoint,
- "logger": {
- "wandb": {
- "_target_": "lightning.pytorch.loggers.wandb.WandbLogger",
- "id": run_id,
- "entity": "rl2-group",
- "project": "abc-yam",
- }
- },
- }
- )
-
-
-def test_explicit_checkpoint_preserves_fresh_wandb_identity(tmp_path):
- checkpoint = tmp_path / "source.ckpt"
- checkpoint.touch()
- cfg = _cfg(tmp_path, checkpoint=str(checkpoint))
- cfg.logger.wandb.resume = "never"
- (tmp_path / "wandb-run-id.txt").write_text("old-run\n")
-
- configure_runner_wandb(cfg, {"SLURM_RESTART_COUNT": "0"})
-
- assert cfg.logger.wandb.id == "stable-run"
- assert cfg.logger.wandb.resume == "never"
-
-
-def test_lightning_restart_recovers_sidecar_and_last_checkpoint(tmp_path):
- checkpoint = tmp_path / "checkpoints" / "last.ckpt"
- checkpoint.parent.mkdir()
- checkpoint.touch()
- (tmp_path / "wandb-run-id.txt").write_text("original-run\n")
- cfg = _cfg(tmp_path, run_id="${now:%Y%m%d}")
-
- configure_runner_wandb(cfg, {"SLURM_RESTART_COUNT": "2"})
-
- assert cfg.ckpt_path == str(checkpoint)
- assert cfg.logger.wandb.id == "original-run"
- assert cfg.logger.wandb.resume == "must"
-
-
-def test_resume_rejects_conflicting_wandb_identities(tmp_path):
- checkpoint = tmp_path / "source.ckpt"
- checkpoint.touch()
- (tmp_path / "wandb-run-id.txt").write_text("sidecar-run\n")
- cfg = _cfg(tmp_path, run_id="configured-run", checkpoint=str(checkpoint))
-
- with pytest.raises(ValueError, match="conflicting W&B run identities"):
- configure_runner_wandb(
- cfg, {"SLURM_RESTART_COUNT": "1", "ICE_RESUME_CHECKPOINT": str(checkpoint)}
- )
-
-
-def test_offline_evaluation_keeps_its_separate_wandb_run(tmp_path):
- cfg = _cfg(tmp_path, run_id="validation-run", checkpoint="source.ckpt")
- cfg.mode = "eval"
- configure_runner_wandb(cfg, {"SLURM_RESTART_COUNT": "0"})
- assert cfg.logger.wandb.id == "validation-run"
- assert "resume" not in cfg.logger.wandb
-
-
-def test_lightning_resume_recovers_actual_wandb_artifact_filename(tmp_path):
- artifact = (
- tmp_path / "wandb/run-20260922_120000-original-run/run-original-run.wandb"
- )
- artifact.parent.mkdir(parents=True)
- artifact.write_bytes(b"wandb")
- cfg = _cfg(tmp_path, run_id="${now:%Y%m%d}", checkpoint="source.ckpt")
- configure_runner_wandb(
- cfg, {"SLURM_RESTART_COUNT": "1", "ICE_RESUME_CHECKPOINT": "source.ckpt"}
- )
- assert cfg.logger.wandb.id == "original-run"
- assert cfg.logger.wandb.resume == "must"
-
-
-def test_model_wrapper_disables_dataloader_suffix_by_default(monkeypatch):
- observed = []
-
- def fake_log(self, *args, **kwargs):
- observed.append(kwargs)
-
- monkeypatch.setattr(LightningModule, "log", fake_log)
- wrapper = ModelWrapper.__new__(ModelWrapper)
- LightningModule.__init__(wrapper)
-
- wrapper.log("metric", 1.0)
- wrapper.log("explicit", 1.0, add_dataloader_idx=True)
-
- assert observed[0]["add_dataloader_idx"] is False
- assert observed[1]["add_dataloader_idx"] is True
-
-
-def test_model_wrapper_log_dict_disables_dataloader_suffix(monkeypatch):
- observed = {}
-
- def fake_log_dict(self, *args, **kwargs):
- observed.update(kwargs)
-
- monkeypatch.setattr(LightningModule, "log_dict", fake_log_dict)
- wrapper = ModelWrapper.__new__(ModelWrapper)
- LightningModule.__init__(wrapper)
-
- wrapper.log_dict({"metric": 1.0})
-
- assert observed["add_dataloader_idx"] is False
diff --git a/tests/test_imported_recipe_framework_parity.py b/tests/test_imported_recipe_framework_parity.py
deleted file mode 100644
index 475fac391f..0000000000
--- a/tests/test_imported_recipe_framework_parity.py
+++ /dev/null
@@ -1,68 +0,0 @@
-"""Imported recipes retain the exact source7c scientific payload.
-
-Hashes were extracted from immutable7c253e5 YAML after removing only the
-framework target. This remains executable in shallow CI without source fetches.
-"""
-
-import hashlib
-import json
-from pathlib import Path
-
-import pytest
-import yaml
-
-
-@pytest.mark.parametrize(
- "name,expected",
- [
- (
- "us_action_flow_bc_latent_fm_sg_200m_adamw",
- "a2007390f6497a04198308190bb259983eee4d75a9beed810926152e5c2f2fb9",
- ),
- (
- "us_action_flow_bc_graph_section",
- "bca57f3f847f55fb4c06e7b53af8ddd22e6d7e9252d288177a18c148e8d91100",
- ),
- (
- "us_action_flow_bc_latent_fm_sg_200m_muon",
- "0d7e4b03b9c5e1321ee009b78a0a89ad5ba9b954402f75833430a3bc5fb0239e",
- ),
- (
- "us_action_flow_bc_latent_fm_sg",
- "4cce3e22ad1a379ec2dfca757631f60d320e903e0a86d2c5cde02def3ab249e3",
- ),
- (
- "us_action_flow_bc_nt16_d8_h512",
- "d909b72505196be3c46d7dfb78b37894e55a9f4f6b812c61af07bb86325a9b00",
- ),
- (
- "cotrain_uc_action_flow_latent_fm_sg_unite_h384d12h12",
- "caba2d2a9432ef36876405a4affa76dc3d32c93ad8728de66cf58d7b9c4fb615",
- ),
- ],
-)
-def test_framework_port_keeps_every_scientific_field(name, expected):
- root = Path(__file__).resolve().parents[1]
- cfg = yaml.safe_load(
- (root / "egomimic/hydra_configs/model/bf" / (name + ".yaml")).read_text()
- )
- assert cfg.pop("_target_") == "egomimic.pl_utils.pl_model.ModelWrapper"
- assert cfg.pop("training_behavior") == {
- "_target_": "egomimic.pl_utils.training_behavior_action_flow.ActionFlowTrainingBehavior"
- }
- assert cfg.pop("diagnostic_provider") == {
- "_target_": "egomimic.eval.pipeline_diagnostics.ActionFlowDiagnosticProvider"
- }
- if name in {
- "us_action_flow_bc_latent_fm_sg_200m_adamw",
- "us_action_flow_bc_latent_fm_sg_200m_muon",
- }:
- # Main adds an explicit non-deployable declaration, not a new adapter.
- assert cfg.pop("inference") == {
- "status": "unsupported",
- "reason": "This simulation recipe has no deployment observation/action adapter. Use its configured simulator evaluator; declare an inference adapter before deployment.",
- }
- actual = hashlib.sha256(
- json.dumps(cfg, sort_keys=True, separators=(",", ":")).encode()
- ).hexdigest()
- assert actual == expected
diff --git a/tests/test_installed_package.py b/tests/test_installed_package.py
deleted file mode 100644
index 9ba6a94ecf..0000000000
--- a/tests/test_installed_package.py
+++ /dev/null
@@ -1,102 +0,0 @@
-"""Build independently and check wheel resources without checkout fallbacks."""
-
-import json
-import os
-import shutil
-import subprocess
-import sys
-import zipfile
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-@pytest.fixture(scope="module")
-def installed_wheel(tmp_path_factory):
- work = tmp_path_factory.mktemp("package")
- source = work / "source"
- source.mkdir()
- shutil.copytree(
- ROOT / "egomimic",
- source / "egomimic",
- ignore=shutil.ignore_patterns("__pycache__", "*.pyc"),
- )
- for name in ("pyproject.toml", "README.md", "LICENSE"):
- if (ROOT / name).exists():
- shutil.copy2(ROOT / name, source / name)
- proc = subprocess.run(
- ["uv", "build", "--wheel", "--out-dir", str(work / "dist")],
- cwd=source,
- text=True,
- capture_output=True,
- timeout=240,
- )
- assert proc.returncode == 0, proc.stderr[-4000:]
- wheel = next((work / "dist").glob("*.whl"))
- with zipfile.ZipFile(wheel) as archive:
- names = archive.namelist()
- subprocess.run(
- ["uv", "venv", str(work / "env"), "--python", sys.executable],
- check=True,
- capture_output=True,
- )
- python = work / "env/bin/python"
- subprocess.run(
- ["uv", "pip", "install", "--python", str(python), "--no-deps", str(wheel)],
- check=True,
- capture_output=True,
- )
- return work, python, names
-
-
-def test_all_declared_resources_are_in_wheel(installed_wheel):
- _, _, names = installed_wheel
- required = set()
- for folder in (
- "hydra_configs",
- "resources",
- "robot/teleop_dashboard_static",
- "robot/rollout_dashboard_static",
- ):
- required.update(
- str(path.relative_to(ROOT))
- for path in (ROOT / "egomimic" / folder).rglob("*")
- if path.suffix
- in {".yaml", ".json", ".xml", ".urdf", ".html", ".css", ".js"}
- )
- required.update(
- str(path.relative_to(ROOT))
- for path in (ROOT / "egomimic/robot/eva/stanford_repo/models/meshes/X5").glob(
- "*.STL"
- )
- )
- assert required <= set(names), sorted(required - set(names))
- assert not any(name.startswith(("tests/", "external/")) for name in names)
-
-
-def test_fresh_install_loads_resources_outside_checkout(installed_wheel):
- work, python, names = installed_wheel
- resources = [name for name in names if name.startswith("egomimic/resources/")]
- code = """
-import egomimic, json, pathlib, sys, xml.etree.ElementTree as ET
-root = pathlib.Path(egomimic.__file__).parent
-assert 'site-packages' in str(root), root
-for name in json.loads(sys.argv[1]):
- path = root.parent / name
- ET.parse(path)
-tree = ET.parse(root / 'resources/model_x5.xml')
-meshdir = root / 'resources' / tree.find('compiler').attrib['meshdir']
-for mesh in tree.findall('asset/mesh'):
- assert (meshdir / mesh.attrib['file']).is_file(), mesh.attrib
-"""
- env = {k: v for k, v in os.environ.items() if k != "PYTHONPATH"}
- result = subprocess.run(
- [str(python), "-c", code, json.dumps(resources)],
- cwd=work,
- env=env,
- capture_output=True,
- text=True,
- )
- assert result.returncode == 0, result.stderr
diff --git a/tests/test_legacy_data_metric_parity.py b/tests/test_legacy_data_metric_parity.py
deleted file mode 100644
index dcf4d86833..0000000000
--- a/tests/test_legacy_data_metric_parity.py
+++ /dev/null
@@ -1,86 +0,0 @@
-"""Numerical references executed in the immutable EgoVerse source checkout."""
-
-import hashlib
-import json
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.utils import instantiate
-
-from scripts.integration.reference_probe import probe
-from tests.fixtures.synthetic_episodes import write_moving_episode
-
-ROOT = Path(__file__).parents[1]
-REFERENCES = Path(__file__).parent / "fixtures/legacy_pipeline"
-CASES = [
- (v, "cartesian", f)
- for v in ("eva", "aria", "mecka", "scale")
- for f in ("camframe", "eef_frame")
-] + [("aria", "keypoints", f) for f in ("camframe", "eef_frame")]
-
-
-@pytest.fixture(scope="module")
-def episodes(tmp_path_factory):
- root = tmp_path_factory.mktemp("moving-legacy-fixtures")
- for vendor in ("eva", "aria", "mecka", "scale"):
- for seed in range(3):
- write_moving_episode(root / vendor, vendor, seed=seed)
- return root
-
-
-@pytest.mark.parametrize("vendor,representation,frame", CASES)
-def test_preprocessing_normalization_and_headline_metrics_match_legacy(
- episodes, vendor, representation, frame
-):
- name = f"{vendor}-{representation}-{frame}"
- archive = REFERENCES / (name + ".npz")
- reference = json.loads(archive.with_suffix(".json").read_text())
- assert reference["source_commit"] == "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2"
- assert (
- hashlib.sha256(archive.read_bytes()).hexdigest() == reference["archive_sha256"]
- )
- actual, metadata = probe(episodes, ROOT, vendor, representation, frame, "graph")
- assert metadata["model_keys"] == reference["model_keys"]
- assert metadata["normalization_frames"] == reference["normalization_frames"]
- with np.load(archive, allow_pickle=False) as expected:
- assert set(actual) == set(expected.files)
- for key, value in actual.items():
- assert value.shape == expected[key].shape, key
- if "observations.images." in key:
- # JPEG decoder rounding can differ between platform wheels;
- # compare routing/shape and pixels to one quantization unit.
- np.testing.assert_allclose(
- value, expected[key], rtol=0, atol=1 / 255, err_msg=key
- )
- else:
- np.testing.assert_allclose(
- value, expected[key], rtol=2e-5, atol=2e-6, err_msg=key
- )
-
-
-@pytest.mark.parametrize("kind", ["pooled", "pertoken"])
-def test_retained_language_schedule_matches_all_source_epochs(kind):
- expected = json.loads(
- (Path(__file__).parent / "fixtures/legacy_language_schedule.json").read_text()
- )
- with initialize_config_dir(
- version_base=None, config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"model=hpt_bc_pickplace_qwen_{kind}"],
- )
- assert cfg.model.scheduler_interval == "epoch"
- optimizer = instantiate(cfg.model.optimizer)(
- params=[torch.nn.Parameter(torch.ones(1))]
- )
- scheduler = instantiate(cfg.model.scheduler)(optimizer=optimizer)
- values = [optimizer.param_groups[0]["lr"]]
- for _ in range(expected["total_epochs"]):
- optimizer.step()
- scheduler.step()
- values.append(optimizer.param_groups[0]["lr"])
- assert values == pytest.approx(expected["values"], rel=1e-12, abs=1e-15)
diff --git a/tests/test_legacy_hpt_trunk_parity.py b/tests/test_legacy_hpt_trunk_parity.py
deleted file mode 100644
index 6c77aac535..0000000000
--- a/tests/test_legacy_hpt_trunk_parity.py
+++ /dev/null
@@ -1,93 +0,0 @@
-"""The generic initializer retains the source HPT trunk namespace and numerics."""
-
-import hashlib
-import json
-from pathlib import Path
-
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.utils import instantiate
-
-from egomimic.pipeline.algo import PipelineAlgo
-from egomimic.pipeline.stages_hpt import HPTTrunkStage
-from scripts.audit_hydra_configs import CONFIGS
-
-FIXTURE = Path(__file__).parent / "fixtures/legacy_hpt_trunk"
-
-
-@pytest.mark.parametrize("transport", ["local", "huggingface"])
-def test_legacy_trunk_loading_outputs_and_gradients(transport, monkeypatch):
- metadata = json.loads((FIXTURE / "metadata.json").read_text())
- assert metadata["source_commit"] == "ec5c903c067bf1b29bbff781bc414c6df2bcf1f2"
- for filename, key in (
- ("trunk.pth", "weights_sha256"),
- ("reference.pt", "reference_sha256"),
- ):
- assert (
- hashlib.sha256((FIXTURE / filename).read_bytes()).hexdigest()
- == metadata[key]
- )
- reference = torch.load(FIXTURE / "reference.pt", weights_only=True)
- with initialize_config_dir(version_base=None, config_dir=str(CONFIGS)):
- cfg = compose(
- config_name="train_zarr_cartesian", overrides=["model=hpt_bc_flow_eva"]
- )
- core = cfg.model.pipeline.stages[1].trunk
- core.embed_dim = core.attn_target.embed_dim = 16
- core.num_blocks = core.attn_target.num_heads = 2
- stage = HPTTrunkStage(instantiate(core), embed_dim=16)
- source = str(FIXTURE / "trunk.pth")
- downloads = []
- if transport == "huggingface":
- # Substitute only transport: tensor loading/validation and the actual
- # source HPT namespace remain the same as the local-file case.
- source = {
- "repo_id": "fixture/hpt",
- "revision": "a" * 40,
- "filename": "trunk.pth",
- }
-
- def fetch(**kwargs):
- downloads.append(kwargs)
- return str(FIXTURE / "trunk.pth")
-
- monkeypatch.setattr("huggingface_hub.hf_hub_download", fetch)
- graph = PipelineAlgo(
- [stage],
- device="cpu",
- stage_ids={"trunk": 0},
- initialization=[
- {
- "stage_id": "trunk",
- "module_path": "trunk",
- "source": source,
- "source_prefix": metadata["source_namespace"],
- "sha256": metadata["weights_sha256"],
- }
- ],
- )
- assert graph.initialization_receipts[0]["sha256"] == metadata["weights_sha256"]
- assert downloads == ([source] if transport == "huggingface" else [])
- trunk = graph.pipeline.stage_by_id("trunk").trunk.eval()
- inputs = reference["input"].clone().requires_grad_()
- output, blocks = trunk(inputs, attn_mask=reference["mask"])
- loss = (output * reference["target"]).mean() + 0.1 * output.square().mean()
- loss.backward()
- torch.testing.assert_close(output, reference["output"], rtol=2e-5, atol=2e-6)
- torch.testing.assert_close(loss, reference["loss"], rtol=2e-5, atol=2e-6)
- torch.testing.assert_close(
- inputs.grad, reference["input_grad"], rtol=2e-5, atol=2e-6
- )
- for actual, expected in zip(blocks, reference["blocks"], strict=True):
- torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6)
- assert {"trunk." + key for key, _ in trunk.named_parameters()} == set(
- reference["parameter_grad"]
- )
- for name, parameter in trunk.named_parameters():
- torch.testing.assert_close(
- parameter.grad,
- reference["parameter_grad"]["trunk." + name],
- rtol=2e-5,
- atol=2e-6,
- )
diff --git a/tests/test_libero_action_flow_artifacts.py b/tests/test_libero_action_flow_artifacts.py
deleted file mode 100644
index 24d6287ce1..0000000000
--- a/tests/test_libero_action_flow_artifacts.py
+++ /dev/null
@@ -1,153 +0,0 @@
-"""Tiny pure-metadata validation; no model or data construction."""
-
-import pytest
-
-from egomimic.benchmarks.libero.action_flow_artifacts import (
- REQUIRED,
- SEED_SHA,
- canonical_sha,
- validate_payload,
-)
-
-
-def fixture():
- identity = {
- "suite": "libero10",
- "source": "libero_panda",
- "action_dim": 7,
- "action_horizon": 16,
- "seed": 42,
- "sample_count": 32,
- "energy_seed_bank_sha256": SEED_SHA,
- "inference_method": "euler",
- "inference_steps": 50,
- "normalization_scope": "training_episodes_only",
- "effective_batch_size": 32,
- "homogeneous": "not_applicable_single_source",
- "source_commit": "a" * 40,
- "resolved_config_sha256": "b" * 64,
- "split_sha256": "c" * 64,
- "normalizer_state_sha256": "d" * 64,
- "dataset_logical_sha256": "e" * 64,
- }
- tensor_metadata = {
- name: {
- "shape": [32, 16, 7],
- "axes": ["batch", "horizon", "action"],
- "finite": True,
- "sha256": "f" * 64,
- }
- for name in (
- "normalized_prediction",
- "native_prediction",
- "normalized_target",
- "native_target",
- )
- }
- for name in ("normalized_samples", "native_samples"):
- tensor_metadata[name] = {
- "shape": [32, 32, 16, 7],
- "axes": ["sample", "batch", "horizon", "action"],
- "finite": True,
- "sha256": "f" * 64,
- }
- for name in ("diagnostic/latent/clean", "diagnostic/decoded/reconstruction"):
- tensor_metadata[name] = {"finite": True, "sha256": "f" * 64}
- return {
- "schema": "libero-native-action-flow-metrics/v2",
- "identity": identity,
- "identity_sha256": canonical_sha(identity),
- "source": "libero_panda",
- "global_step": 2,
- "batches": [
- {
- "batch_index": 0,
- "batch_size": 32,
- "normalized_target_sha256": "f" * 64,
- "metrics": {key: 1.0 for key in REQUIRED},
- # Pure metadata fixture: validates reference contracts only,
- # never proves payload bytes, actual episodes or a run PASS.
- "tensor_payload": {
- "schema": "libero-native-tensors/v1",
- "filename": "fixture.pt",
- "sha256": "f" * 64,
- "tensors": tensor_metadata,
- "alignment": {"fixture": "explicit-synthetic-metadata"},
- "shared_analysis": {
- "sha256": "f" * 64,
- "identity_sha256": "f" * 64,
- "path": "fixture-analysis.pt",
- "global_step": 2,
- },
- },
- }
- ],
- }
-
-
-def test_native_complete():
- assert validate_payload(fixture())
-
-
-def test_legacy_metric_only_schema_does_not_certify_native_tensors():
- payload = fixture()
- payload["schema"] = "libero-native-action-flow-metrics/v1"
- with pytest.raises(ValueError, match="wrong native artifact schema"):
- validate_payload(payload)
-
-
-def test_reference_rejects_wrong_sample_axes():
- payload = fixture()
- payload["batches"][0]["tensor_payload"]["tensors"]["native_samples"]["axes"] = [
- "batch",
- "sample",
- "horizon",
- "action",
- ]
- with pytest.raises(ValueError, match="K32 tensor contract"):
- validate_payload(payload)
-
-
-@pytest.mark.parametrize("key", REQUIRED)
-def test_reject_nonfinite(key):
- payload = fixture()
- payload["batches"][0]["metrics"][key] = float("nan")
- with pytest.raises(ValueError):
- validate_payload(payload)
-
-
-@pytest.mark.parametrize(
- "key",
- (
- "source_commit",
- "resolved_config_sha256",
- "split_sha256",
- "normalizer_state_sha256",
- "dataset_logical_sha256",
- "energy_seed_bank_sha256",
- "normalization_scope",
- "inference_method",
- "inference_steps",
- "effective_batch_size",
- ),
-)
-def test_reject_identity(key):
- payload = fixture()
- payload["identity"][key] = "wrong"
- payload["identity_sha256"] = canonical_sha(payload["identity"])
- with pytest.raises(ValueError):
- validate_payload(payload)
-
-
-def test_missing_first_batch():
- payload = fixture()
- payload["batches"][0]["batch_index"] = 1
- with pytest.raises(ValueError):
- validate_payload(payload)
-
-
-def test_no_implicit_checkpoint_completion():
- payload = fixture()
- payload["global_step"] = 0
- with pytest.raises(ValueError):
- validate_payload(payload)
diff --git a/tests/test_libero_active_recipes.py b/tests/test_libero_active_recipes.py
deleted file mode 100644
index 3a7afa1de5..0000000000
--- a/tests/test_libero_active_recipes.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""Keep active LIBERO choices unambiguous without loading data or checkpoints."""
-
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-
-ROOT = Path(__file__).resolve().parents[1]
-CONFIG = ROOT / "egomimic/hydra_configs"
-TARGETS = {
- "libero10": 605121,
- "libero_spatial": 270054,
- "libero_object": 325065,
- "libero_goal": 280056,
-}
-
-
-def test_active_directory_contains_selected_eight_and_explicit_joint_recipe():
- expected = {
- f"{family}_{suite}_{suffix}_s42.yaml"
- for suite in TARGETS
- for family, suffix in (
- ("action_flow", "oat_dp_matched"),
- ("dp", "oat_batch16_keep_steps"),
- )
- }
- # The joint experiment is separate from the original four AF/DP pairs.
- expected.add("action_flow_libero_spatial_sharpa_gripper_keypoints_s42.yaml")
- expected.add("action_flow_libero_spatial_sharpa_mot_s42.yaml")
- assert {
- p.name for p in (CONFIG / "experiment/sim/libero").glob("*.yaml")
- } == expected
-
-
-@pytest.mark.parametrize("suite,steps", TARGETS.items())
-def test_active_pair_preserves_selected_scientific_contract(suite, steps):
- with initialize_config_dir(config_dir=str(CONFIG), version_base="1.3"):
- af = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- f"+experiment=sim/libero/action_flow_{suite}_oat_dp_matched_s42"
- ],
- )
- dp = compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/libero/dp_{suite}_oat_batch16_keep_steps_s42"],
- )
- assert (af.model.num_latent_tokens, af.model.latent_dim) == (16, 16)
- assert af.model.pipeline.stages[-1].action_velocity_weight == 1
- assert (af.trainer.max_steps, af.trainer.val_check_interval) == (120000, 20000)
- assert af.trainer.check_val_every_n_epoch is None
- assert af.callbacks.ema.decay == 0.9978
- assert af.trainer.gradient_clip_val == 3
- assert af.benchmark.batch_size == 32
- assert (
- dp.benchmark.batch_size,
- dp.trainer.accumulate_grad_batches,
- dp.trainer.max_steps,
- ) == (16, 1, steps)
- assert dp.callbacks.batch_budget.global_batch_size == 16
- if suite == "libero_spatial":
- assert dp.callbacks.model_checkpoint.every_n_train_steps == 15000
- assert dp.callbacks.model_checkpoint.every_n_epochs is None
- assert af.model.pipeline.stages[0] == dp.model.pipeline.stages[0]
- for cfg in (af, dp):
- assert cfg.data.train_datasets.libero_panda.valid_ratio == 0.1
- assert cfg.data.valid_datasets.libero_panda.valid_ratio == 0.1
-
-
-@pytest.mark.parametrize(
- "recipe", sorted((CONFIG / "experiment/sim/libero_historical").glob("*.yaml"))
-)
-def test_historical_profiles_remain_composable(recipe):
- with initialize_config_dir(config_dir=str(CONFIG), version_base="1.3"):
- compose(
- config_name="train_zarr_cartesian",
- overrides=[f"+experiment=sim/libero_historical/{recipe.stem}"],
- )
diff --git a/tests/test_libero_arc_sweep.py b/tests/test_libero_arc_sweep.py
deleted file mode 100644
index 9328e9cce4..0000000000
--- a/tests/test_libero_arc_sweep.py
+++ /dev/null
@@ -1,329 +0,0 @@
-import importlib.util
-import json
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-import yaml
-
-from egomimic.benchmarks.libero.arc_sweep import fixed_replay_spec, profile_settings
-from egomimic.benchmarks.libero.replay import candidates_from_spec, validate_spec
-
-
-@pytest.mark.parametrize(
- "profile,mode",
- [
- ("stk_1", "stk"),
- ("stk_2", "stk"),
- ("dur_1", "dur"),
- ("dur_2", "dur"),
- ("shared", "stk"),
- ("shared", "dur"),
- ],
-)
-def test_frozen_sweep_replay_preserves_choice_and_uses_fresh_confirmation(
- profile, mode
-):
- settings = profile_settings(profile, mode)
- spec = fixed_replay_spec(profile, mode)
- candidates = candidates_from_spec(spec)
- assert len(candidates) == 1
- candidate = next(iter(candidates.values()))
- assert candidate == {
- "mode": mode,
- "num_waypoints": settings["arc_waypoints"],
- "max_rotation_degrees": settings["arc_max_rotation_degrees"],
- "max_translation": settings["arc_max_translation"],
- }
- assert spec["calibration_demos"] == [0, 1]
- assert spec["selection_demos"] == list(range(2, 10))
- assert spec["confirmation_demos"] == list(range(10, 25))
- assert not set(spec["confirmation_demos"]) & set(range(30, 50))
- assert spec["minimum_execution_coverage"] == 0.99
- assert spec["max_raw_repeat_state_error"] == 0
- assert spec["max_dense_action_mse"] == 1e-12
-
-
-def test_profile_rejects_wrong_mode_and_workflow_omissions():
- with pytest.raises(ValueError, match="incompatible"):
- profile_settings("stk_1", "dur")
- path = Path(__file__).parents[1] / "scripts/benchmarks/launch_libero_osmo.py"
- spec = importlib.util.spec_from_file_location("sweep_osmo", path)
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- args = dict(
- commit="a" * 40,
- run_id="sweep-stk1-10",
- suite="libero_10",
- mode="full",
- arc_modes=["stk"],
- arc_profile="stk_1",
- )
- with pytest.raises(ValueError, match="requires calibration"):
- module.workflow(**args)
- result = module.workflow(
- **args, calibration_parent="parent", oat_reference_run="baseline"
- )["workflow"]
- env = result["tasks"][0]["environment"]
- assert env["RUN_KIND"] == "arc_sweep" and env["ARC_PROFILE"] == "stk_1"
- assert env["OAT_REFERENCE_RUN"] == "baseline"
- assert result["resources"]["default"]["gpu"] == 1
- assert result["resources"]["default"]["memory"] == "128Gi"
- assert result["resources"]["default"]["cpu"] == 20
-
-
-def test_fresh_backbone_run_reuses_audited_replay_without_resuming_weights(
- tmp_path, monkeypatch
-):
- from egomimic.benchmarks.libero import arc_sweep, cluster
- from scripts.benchmarks.launch_libero_osmo import workflow
-
- result = workflow(
- "a" * 40,
- "arc-dp-stk1-spatial",
- "libero_spatial",
- mode="full",
- arc_modes=["stk"],
- arc_profile="stk_1",
- arc_backbone="oat_dp",
- arc_replay_runs={"stk": "audited-replay"},
- oat_reference_run="baseline",
- gpus=4,
- )["workflow"]
- env = result["tasks"][0]["environment"]
- assert env["ARC_BACKBONE"] == "oat_dp"
- assert env["RESUME_FROM_RUN"] == ""
- assert result["resources"]["default"]["gpu"] == 4
- for key, value in env.items():
- monkeypatch.setenv(key, value)
- calls = []
- monkeypatch.setattr(cluster, "execute", lambda args, log: calls.append(args))
- monkeypatch.setattr(
- sys,
- "argv",
- [
- "sweep",
- "--root",
- str(tmp_path),
- "--suite",
- "libero_spatial",
- "--run-id",
- "arc-dp-stk1-spatial",
- "--profile",
- "stk_1",
- "--arc-mode",
- "stk",
- ],
- )
- arc_sweep.main()
- assert len(calls) == 1 and calls[0][2].endswith(".cluster")
- assert json.loads(__import__("os").environ["ARC_REPLAY_RUNS_JSON"]) == {
- "stk": "audited-replay"
- }
-
-
-@pytest.mark.parametrize("replay_failure", [False, True])
-def test_policy_launch_follows_replay_and_stops_on_preflight_failure(
- tmp_path, monkeypatch, replay_failure
-):
- from egomimic.benchmarks.libero import arc_sweep, cluster
-
- calls = []
-
- def execute(argv, log):
- calls.append(argv)
- if argv[2].endswith(".replay"):
- # Exercise exactly the disk format read in the real subprocess.
- loaded = yaml.safe_load(Path(argv[argv.index("--spec") + 1]).read_text())
- validate_spec(loaded)
- assert loaded["max_dense_action_mse"] == 1e-12
- if replay_failure:
- raise RuntimeError("preflight failed")
-
- monkeypatch.setattr(cluster, "execute", execute)
- monkeypatch.setenv("REPLAY_CALIBRATION_PARENT", "parent-stk")
- monkeypatch.setattr(
- sys,
- "argv",
- [
- "sweep",
- "--root",
- str(tmp_path),
- "--suite",
- "libero_10",
- "--run-id",
- "sweep-stk1-10",
- "--profile",
- "stk_1",
- "--arc-mode",
- "stk",
- ],
- )
- if replay_failure:
- with pytest.raises(RuntimeError, match="preflight failed"):
- arc_sweep.main()
- assert len(calls) == 1
- else:
- arc_sweep.main()
- assert len(calls) == 2 and "--arc-only" in calls[1]
- assert calls[0][2].endswith(".replay") and calls[1][2].endswith(".cluster")
- assert (
- yaml.safe_load((tmp_path / "fixed-replay-spec.yaml").read_text())[
- "frozen_profile"
- ]
- == "stk_1"
- )
-
-
-@pytest.mark.parametrize(
- "run_mode,bad_replay,gpu_type,backbone",
- [
- ("smoke", False, "L40S", "unet"),
- ("full", False, "L40S", "unet"),
- ("full", True, "L40S", "unet"),
- ("full", False, "H100", "unet"),
- ("smoke", False, "L40S", "oat_dp"),
- ("full", False, "H100", "oat_dp"),
- ],
-)
-def test_arc_only_never_trains_oat_and_checks_frozen_profile(
- tmp_path, monkeypatch, run_mode, bad_replay, gpu_type, backbone
-):
- import torch
-
- from egomimic.benchmarks.libero import cluster
- from egomimic.benchmarks.libero.rollout import rollout_plan, run_rollouts
- from tests.test_libero_benchmark import FakeEnvironment, FakePolicy, make_replay
-
- monkeypatch.setenv("SOURCE_COMMIT", "a" * 40)
- monkeypatch.setenv("BENCHMARK_GPU_TYPE", gpu_type)
- monkeypatch.setenv("ARC_BACKBONE", backbone)
- for name in (
- "CAMPAIGN_ID",
- "CAMPAIGN_RUNS_JSON",
- "RESUME_FROM_RUN",
- "EVALUATE_FROM_RUN",
- "ARC_REPLAY_RUN",
- ):
- monkeypatch.delenv(name, raising=False)
- monkeypatch.setenv("ARC_MODES_JSON", '["stk"]')
- monkeypatch.setenv("ARC_REPLAY_RUNS_JSON", '{"stk":"verified-replay"}')
- monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
- monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
- monkeypatch.setattr(torch.cuda, "get_device_name", lambda _: gpu_type)
- monkeypatch.setattr(cluster.subprocess, "check_output", lambda *a, **k: "a" * 40)
-
- def no_op(*a, **k):
- pass
-
- monkeypatch.setattr(torch.cuda, "synchronize", no_op)
-
- class Uploader:
- def __init__(self, *a):
- self.client = None
- self.prefix = "test/"
- self.thread = SimpleNamespace(start=no_op, join=no_op)
- self.stop = SimpleNamespace(set=no_op)
-
- upload = no_op
-
- monkeypatch.setattr(cluster, "ArtifactUploader", Uploader)
- replay_path = tmp_path / "data.zarr"
- make_replay(replay_path, suite="libero_10")
- monkeypatch.setattr(cluster, "stage_dataset", lambda *a: replay_path)
- monkeypatch.setattr(cluster, "configure_simulator", no_op)
- settings = profile_settings("stk_1", "stk")
- monkeypatch.setattr(
- cluster,
- "load_arc_calibration",
- lambda *a, **k: settings | ({"arc_waypoints": 16} if bad_replay else {}),
- )
- calls = []
-
- def execute(argv, log):
- calls.append(argv)
- if "egomimic.trainHydra" in argv:
- expected_recipe = (
- "libero_arc_oat_dp_policy"
- if backbone == "oat_dp"
- else "libero_arc_stk_policy"
- )
- assert f"+experiment=sim/oat/{expected_recipe}" in argv
- assert "benchmark.arc_waypoints=36" in argv
- checkpoint = tmp_path / "evidence/training/arc_stk/checkpoints/last.ckpt"
- checkpoint.parent.mkdir(parents=True)
- checkpoint.write_bytes(b"checkpoint")
- if "rollout" in argv:
- output = Path(argv[argv.index("--output") + 1])
- run_rollouts(
- FakePolicy(),
- rollout_plan("libero_10")
- if run_mode == "full"
- else rollout_plan("libero_10", trials_per_task=1, repetitions=1),
- output,
- max_episode_steps=550 if run_mode == "full" else 4,
- env_factory=FakeEnvironment,
- metadata={
- "suite": "libero_10",
- "method": "arc",
- "representation": {
- "mode": "stk",
- "waypoints": 36,
- "max_translation": 1.6,
- "max_rotation_degrees": 192,
- },
- "horizon": 32,
- "n_obs_steps": 2,
- "n_action_steps": 16,
- "use_ema": True,
- },
- )
-
- from egomimic.benchmarks.libero import evaluate
-
- parallel_calls = []
-
- def parallel_rollouts(checkpoint, evidence, method, suite):
- parallel_calls.append((method, suite))
- execute(["rollout", "--output", str(evidence / method / suite)], None)
-
- monkeypatch.setattr(evaluate, "parallel_rollouts", parallel_rollouts)
- monkeypatch.setattr(cluster, "execute", execute)
- monkeypatch.setattr(
- sys,
- "argv",
- [
- "cluster",
- "--root",
- str(tmp_path),
- "--suite",
- "libero_10",
- "--run-id",
- "sweep",
- "--mode",
- run_mode,
- "--arc-only",
- "--arc-profile",
- "stk_1",
- ],
- )
- if bad_replay:
- with pytest.raises(ValueError, match="differs from frozen"):
- cluster.main()
- assert not any("egomimic.trainHydra" in c for c in calls)
- else:
- cluster.main()
- assert parallel_calls == (
- [("arc_stk", "libero_10")]
- if backbone == "oat_dp" and run_mode == "full"
- else []
- )
- assert sum("egomimic.trainHydra" in c for c in calls) == 1
- assert not any("reconstruct" in c for c in calls)
- result = json.loads((tmp_path / "evidence/arc-results.json").read_text())
- assert set(result["variants"]) == {"stk"}
- assert result["paired_oat_comparison_complete"] is False
- assert json.loads((tmp_path / "evidence/status.json").read_text())["state"] == (
- "ARC_POLICIES_COMPLETE" if run_mode == "full" else "ARC_SMOKE_PASSED"
- )
diff --git a/tests/test_libero_basis_scaling_audit.py b/tests/test_libero_basis_scaling_audit.py
deleted file mode 100644
index 5b14747014..0000000000
--- a/tests/test_libero_basis_scaling_audit.py
+++ /dev/null
@@ -1,113 +0,0 @@
-import numpy as np
-import pytest
-
-from egomimic.rldb.zarr.libero_arc_global import LiberoArcGlobalCodec
-from scripts.benchmarks.compare_libero_arc_bases import fit_bounded_scaling, noise_audit
-
-
-def test_scaling_is_train_only_bounded_and_reversible():
- rng = np.random.default_rng(8)
- values = rng.normal(size=(8, 32, 12))
- values[:, 0, 0] = 4
- mask = np.arange(8) < 4
- center, scale, std = fit_bounded_scaling(values, mask, 32)
- changed = values.copy()
- changed[~mask] += 1000
- for left, right in zip(
- (center, scale, std), fit_bounded_scaling(changed, mask, 32)
- ):
- np.testing.assert_array_equal(left, right)
- assert np.isfinite(scale).all() and scale.min() >= 1e-4
- assert std[0, 0] == 0
- np.testing.assert_allclose((values - center) / scale * scale + center, values)
- with pytest.raises(ValueError):
- fit_bounded_scaling(values, np.zeros(8, dtype=bool), 32)
-
-
-@pytest.mark.parametrize("basis", ["fourier", "chebyshev"])
-def test_noise_audit_is_deterministic_and_does_not_change_codec(basis):
- codec = LiberoArcGlobalCodec(basis=basis)
- raw = np.random.default_rng(3).uniform(-0.1, 0.1, (4, 32, 7)).astype(np.float32)
- tokens = np.stack([codec.encode(row) for row in raw])
- frozen = tokens.copy()
- identities = [{"split": "train" if i < 2 else "valid"} for i in range(4)]
- first = noise_audit(codec, tokens, raw, identities, sigmas=(0.0, 0.001))
- assert first == noise_audit(codec, tokens, raw, identities, sigmas=(0.0, 0.001))
- np.testing.assert_array_equal(tokens, frozen)
- for variant in first["variants"].values():
- assert variant["perturbations"]["0.0"]["decoded_delta_mse_mean"] < 1e-12
- assert variant["perturbations"]["0.001"]["decoded_delta_mse_mean"] > 0
-
-
-def test_optional_affine_stage_roundtrip_strict_state_and_identity():
- import copy
-
- import torch
-
- from egomimic.pipeline.stages_libero_arc import LiberoArcStage
-
- codec = LiberoArcGlobalCodec(basis="fourier", geometry=64)
- raw = np.random.default_rng(1).uniform(-0.1, 0.1, (6, 32, 7)).astype(np.float32)
- tokens = np.stack([codec.encode(row) for row in raw])
- audit = noise_audit(
- codec,
- tokens,
- raw,
- [{"split": "train" if i < 4 else "valid"} for i in range(6)],
- sigmas=(0.0,),
- )
- affine = audit["affine"]
- stage = LiberoArcStage(codec=codec, token_affine=affine, reconstruction=True)
- stage.action_scale.copy_(torch.arange(1, 8))
- stage.action_offset.copy_(torch.arange(7) / 5)
- actions = torch.from_numpy(raw) * stage.action_scale + stage.action_offset
- result = stage.execute({"actions": actions}, mode="inference")
- expected = np.stack([codec.decode(row) for row in tokens])
- torch.testing.assert_close(
- result["pred_action"],
- torch.from_numpy(expected) * stage.action_scale + stage.action_offset,
- atol=2e-5,
- rtol=1e-5,
- )
- restored = LiberoArcStage(codec=codec, token_affine=affine, operation="decode")
- restored.load_state_dict(stage.state_dict(), strict=True)
- torch.testing.assert_close(
- restored.execute({"pred_arc": result["target"]}, mode="inference")[
- "pred_action"
- ],
- result["pred_action"],
- )
- assert torch.isfinite(
- restored.execute({"pred_arc": result["target"].bfloat16()}, mode="inference")[
- "pred_action"
- ]
- ).all()
- assert stage.representation_context() == restored.representation_context()
- wrong = copy.deepcopy(affine)
- wrong["geometry"] = 128
- with pytest.raises(ValueError, match="identity differ"):
- LiberoArcStage(codec=codec, token_affine=wrong)
- wrong = copy.deepcopy(affine)
- wrong["scale"][0][0] = 0
- with pytest.raises(ValueError, match="affine arrays"):
- LiberoArcStage(codec=codec, token_affine=wrong)
-
-
-def test_default_uniform_has_no_affine_state_or_numeric_change():
- import torch
-
- from egomimic.pipeline.stages_libero_arc import LiberoArcStage
-
- stage = LiberoArcStage(codec=LiberoArcGlobalCodec(basis="uniform"))
- raw = torch.from_numpy(
- np.random.default_rng(2).uniform(-0.1, 0.1, (2, 32, 7)).astype(np.float32)
- )
- expected = torch.tensor(
- np.stack([stage.codec.encode(a) for a in raw.numpy()])
- ) / torch.tensor(stage.codec.token_scale())
- assert not any("token_" in key for key in stage.state_dict())
- assert "token_affine" not in stage.representation_context()
- assert "geometry_fit_grid" not in stage.representation_context()
- assert torch.equal(
- stage.execute({"actions": raw}, mode="train")["target"], expected
- )
diff --git a/tests/test_libero_cluster.py b/tests/test_libero_cluster.py
deleted file mode 100644
index d982dbf3aa..0000000000
--- a/tests/test_libero_cluster.py
+++ /dev/null
@@ -1,750 +0,0 @@
-import importlib.util
-import json
-import zipfile
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-
-from egomimic.benchmarks.libero.cluster import (
- ArtifactUploader,
- CalibrationPending,
- campaign_sources,
- digest,
- extract_replay,
- load_arc_calibration,
- publish_campaign,
- restore_checkpoints,
- stage_dataset,
- training_arguments,
- validate_gpu_allocation,
-)
-
-
-@pytest.mark.parametrize("mode", ["smoke", "full"])
-@pytest.mark.parametrize("method", ["tokenizer", "oat", "arc", "arc_dur", "arc_stk"])
-def test_cluster_arguments_compose_native_graph(mode, method, tmp_path):
- args = training_arguments(
- method, "libero_10", tmp_path / "replay.zarr", tmp_path, mode, 5001
- )
- config_dir = Path(__file__).parents[1] / "egomimic/hydra_configs"
- with initialize_config_dir(config_dir=str(config_dir), version_base="1.3"):
- cfg = compose(config_name="train_zarr_cartesian", overrides=args[3:])
- assert cfg.model.pipeline._target_ == "egomimic.pipeline.algo.PipelineAlgo"
- assert cfg.trainer.devices == 1 and cfg.trainer.precision == "bf16-mixed"
- assert cfg.benchmark.batch_size == (4 if mode == "smoke" else 256)
- assert cfg.trainer.max_epochs == (1 if mode == "smoke" else 5001)
- if mode == "full":
- assert cfg.logger.csv._target_ == "lightning.pytorch.loggers.CSVLogger"
- assert cfg.trainer.accumulate_grad_batches == 4
- assert cfg.callbacks.batch_budget.global_batch_size == 1024
- else:
- assert cfg.trainer.accumulate_grad_batches == 1
- assert cfg.callbacks.batch_budget is None
- assert cfg.callbacks.ema.final_checkpoint_path.endswith("checkpoints/last.ckpt")
- if method == "oat":
- assert cfg.benchmark.tokenizer_checkpoint.endswith(
- "tokenizer/checkpoints/last.ckpt"
- )
- if method in {"arc_dur", "arc_stk"}:
- mode = method.removeprefix("arc_")
- assert cfg.benchmark.arc_mode == mode
- assert cfg.benchmark.arc_action_dim == 12
- stages = cfg.model.pipeline.stages
- assert stages[1].arc_mode == stages[-1].arc_mode == mode
- assert (
- stages[1].velocity_norm_bound
- == stages[-1].velocity_norm_bound
- == pytest.approx(3**0.5)
- )
- assert stages[2].action_dim == stages[3].policy.model.input_dim == 12
-
-
-def test_workflow_pins_source_and_requests_one_gpu():
- path = Path(__file__).parents[1] / "scripts/benchmarks/launch_libero_osmo.py"
- spec = importlib.util.spec_from_file_location("libero_osmo", path)
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- workflow = module.workflow("a" * 40, "libero-smoke-test", "libero_10")["workflow"]
- assert workflow["resources"]["default"]["gpu"] == 1
- assert workflow["resources"]["default"]["platform"] == "ovx-l40s"
- assert workflow["tasks"][0]["environment"]["BENCHMARK_GPU_TYPE"] == "L40S"
- assert workflow["tasks"][0]["environment"]["SOURCE_COMMIT"] == "a" * 40
- assert json.loads(workflow["tasks"][0]["environment"]["ARC_MODES_JSON"]) == [
- "dur",
- "stk",
- ]
- with pytest.raises(ValueError, match="immutable"):
- module.workflow("main", "libero-smoke-test", "libero_10")
- h100 = module.workflow(
- "a" * 40,
- "h100-resume",
- "libero_10",
- mode="full",
- resume_from_run="prior-run",
- gpus=8,
- gpu_type="H100",
- )["workflow"]
- assert h100["resources"]["default"]["platform"] == "dgx-h100"
- assert h100["resources"]["default"]["gpu"] == 8
- assert h100["resources"]["default"]["cpu"] == 80
- h100_env = h100["tasks"][0]["environment"]
- assert h100_env["BENCHMARK_GPU_TYPE"] == "H100"
- assert h100_env["TRAINING_GPUS"] == "8"
- assert h100_env["EPOCHS"] == "5001"
- assert h100_env["RESUME_FROM_RUN"] == "prior-run"
- with pytest.raises(ValueError, match="GPU type"):
- module.workflow("a" * 40, "invalid-gpu", "libero_10", gpu_type="A100")
- full = module.workflow(
- "a" * 40, "study-libero-10", "libero_10", mode="full", campaign_id="study"
- )["workflow"]
- assert full["timeout"]["exec_timeout"] == "60d"
- assert full["tasks"][0]["environment"]["CAMPAIGN_ID"] == "study"
- replay = module.workflow(
- "a" * 40,
- "replay-refine",
- "libero_90",
- replay=True,
- replay_spec="libero_arc_replay_refine",
- )["workflow"]
- assert replay["resources"]["default"]["cpu"] >= 32
- assert replay["resources"]["default"]["gpu"] == 1
- assert replay["resources"]["default"]["memory"] == "192Gi"
- replacements = campaign_sources("study", "a" * 40)
- replacements["libero_90"] = {
- "run_id": "recovered-libero-90",
- "source_commit": "b" * 40,
- }
- replacement = module.workflow(
- "b" * 40,
- "recovered-libero-90",
- "libero_90",
- mode="full",
- campaign_id="study",
- campaign_runs=replacements,
- )["workflow"]
- assert (
- json.loads(replacement["tasks"][0]["environment"]["CAMPAIGN_RUNS_JSON"])
- == replacements
- )
- with pytest.raises(ValueError, match="Campaign"):
- module.workflow(
- "a" * 40,
- "recovered-libero-90",
- "libero_90",
- mode="full",
- campaign_id="study",
- campaign_runs=replacements,
- )
- with pytest.raises(ValueError, match="Campaign"):
- module.workflow(
- "a" * 40, "unrelated-run", "libero_10", mode="full", campaign_id="study"
- )
- variants = module.workflow(
- "a" * 40,
- "both-modes",
- "libero_10",
- mode="full",
- arc_replay_runs={"dur": "dur-replay", "stk": "stk-replay"},
- )["workflow"]
- assert json.loads(variants["tasks"][0]["environment"]["ARC_REPLAY_RUNS_JSON"]) == {
- "dur": "dur-replay",
- "stk": "stk-replay",
- }
- with pytest.raises(ValueError, match="mode-specific"):
- module.workflow(
- "a" * 40,
- "ambiguous-mode",
- "libero_10",
- arc_modes=["dur", "stk"],
- arc_replay_run="legacy-run",
- )
-
-
-@pytest.mark.parametrize(
- "devices,gpu_type,error",
- [
- (["NVIDIA L40S"] * 8, "L40S", None),
- (["NVIDIA H100 80GB HBM3"] * 8, "H100", None),
- (["NVIDIA L40S"] * 8, "H100", RuntimeError),
- (["NVIDIA H100 80GB HBM3"] * 7, "H100", RuntimeError),
- (["NVIDIA H100 80GB HBM3"] * 7 + ["NVIDIA L40S"], "H100", RuntimeError),
- ([], "H100", RuntimeError),
- (["NVIDIA A100"] * 8, "A100", ValueError),
- ],
-)
-def test_gpu_allocation_rejects_missing_mixed_or_wrong_devices(
- monkeypatch, devices, gpu_type, error
-):
- import torch
-
- monkeypatch.setattr(torch.cuda, "is_available", lambda: bool(devices))
- monkeypatch.setattr(torch.cuda, "device_count", lambda: len(devices))
- monkeypatch.setattr(torch.cuda, "get_device_name", devices.__getitem__)
- if error:
- with pytest.raises(error):
- validate_gpu_allocation(8, gpu_type)
- else:
- validate_gpu_allocation(8, gpu_type)
-
-
-@pytest.mark.parametrize(
- "condition", ["complete", "missing", "running", "source", "budget"]
-)
-def test_campaign_only_publishes_complete_matching_results(condition):
- import io
-
- from botocore.exceptions import ClientError
-
- from egomimic.benchmarks.libero.catalog import TASKS
-
- class Storage:
- report = None
-
- def get_object(self, Bucket, Key):
- suite = next(
- suite for suite in TASKS if f"study-{suite.replace('_', '-')}/" in Key
- )
- if suite == "libero_90" and condition == "missing":
- raise ClientError({"Error": {"Code": "NoSuchKey"}}, "GetObject")
- payload = {
- "status.json": {
- "state": "TRAINING" if condition == "running" else "SUITE_COMPLETE"
- },
- "runtime.json": {
- "source_commit": "b" * 40 if condition == "source" else "a" * 40,
- "mode": "full",
- "epochs": 1 if condition == "budget" else 5001,
- "global_batch_size": 1024,
- "suite": suite,
- },
- "comparison.json": {"suite": suite, "complete_protocol": True},
- }[Key.rsplit("/", 1)[1]]
- return {"Body": io.BytesIO(json.dumps(payload).encode())}
-
- def put_object(self, **kwargs):
- self.report = json.loads(kwargs["Body"])
-
- storage = Storage()
- if condition in {"source", "budget"}:
- with pytest.raises(ValueError, match="Campaign"):
- publish_campaign(storage, "study", commit="a" * 40, epochs=5001)
- else:
- assert publish_campaign(storage, "study", commit="a" * 40, epochs=5001) == (
- condition == "complete"
- )
- if condition == "complete":
- assert storage.report["unique_tasks"] == 130
- assert set(storage.report["suites"]) == set(TASKS)
- assert storage.report["full_released_training_budget"] is True
- else:
- assert storage.report is None
-
-
-def test_released_replay_archive_detection_and_traversal_rejection(tmp_path):
- archive = tmp_path / "replay.zip"
- with zipfile.ZipFile(archive, "w") as handle:
- handle.writestr("libero10_N500.zarr/.zgroup", "{}")
- handle.writestr("libero10_N500.zarr/meta/episode_ends/.zarray", "{}")
- assert (
- extract_replay(archive, tmp_path / "out") == tmp_path / "out/libero10_N500.zarr"
- )
- with zipfile.ZipFile(archive, "w") as handle:
- handle.writestr("../outside", "bad")
- with pytest.raises(ValueError, match="Unsafe"):
- extract_replay(archive, tmp_path / "bad")
- assert not (tmp_path / "outside").exists()
-
-
-@pytest.mark.parametrize("bad_source", [False, True])
-def test_campaign_can_pin_a_replacement_without_restarting_other_suites(bad_source):
- import io
-
- from egomimic.benchmarks.libero.catalog import TASKS
-
- pinned = campaign_sources("study", "a" * 40)
- pinned["libero_90"] = {"run_id": "recovered-libero-90", "source_commit": "b" * 40}
-
- class Storage:
- report = None
-
- def get_object(self, Bucket, Key):
- suite = next(s for s, r in pinned.items() if f"/{r['run_id']}/" in Key)
- row = {
- "status.json": {"state": "SUITE_COMPLETE"},
- "runtime.json": {
- "source_commit": "c" * 40
- if bad_source
- else pinned[suite]["source_commit"],
- "suite": suite,
- "mode": "full",
- "epochs": 5001,
- "global_batch_size": 1024,
- },
- "comparison.json": {"suite": suite, "complete_protocol": True},
- }[Key.rsplit("/", 1)[1]]
- return {"Body": io.BytesIO(json.dumps(row).encode())}
-
- def put_object(self, **kwargs):
- self.report = json.loads(kwargs["Body"])
-
- storage = Storage()
- if bad_source:
- with pytest.raises(ValueError, match="source"):
- publish_campaign(
- storage, "study", commit="b" * 40, epochs=5001, run_sources=pinned
- )
- assert storage.report is None
- else:
- assert publish_campaign(
- storage, "study", commit="b" * 40, epochs=5001, run_sources=pinned
- )
- assert storage.report["source_commit"] is None
- assert storage.report["source_commits"]["libero_90"] == "b" * 40
- assert set(storage.report["suites"]) == set(TASKS)
- with pytest.raises(ValueError, match="all suites"):
- campaign_sources("study", "a" * 40, {"libero_90": pinned["libero_90"]})
-
-
-@pytest.mark.parametrize("missing_mode", [False, True])
-def test_dual_mode_campaign_requires_every_variant(missing_mode):
- import io
-
- from egomimic.benchmarks.libero.catalog import TASKS
-
- class Storage:
- report = None
-
- def get_object(self, Bucket, Key):
- suite = next(s for s in TASKS if f"study-{s.replace('_', '-')}/" in Key)
- variants = {
- mode: {"suite": suite, "complete_protocol": True}
- for mode in ("dur", "stk")
- }
- if missing_mode and suite == "libero_90":
- del variants["stk"]
- value = {
- "runtime.json": {
- "suite": suite,
- "mode": "full",
- "epochs": 5001,
- "global_batch_size": 1024,
- "source_commit": "a" * 40,
- "arc_modes": ["dur", "stk"],
- },
- "status.json": {"state": "SUITE_COMPLETE"},
- "comparison.json": {
- "suite": suite,
- "complete_protocol": True,
- "variants": variants,
- },
- }[Key.rsplit("/", 1)[1]]
- return {"Body": io.BytesIO(json.dumps(value).encode())}
-
- def put_object(self, **kwargs):
- self.report = json.loads(kwargs["Body"])
-
- storage = Storage()
- if missing_mode:
- with pytest.raises(ValueError, match="every ARC mode"):
- publish_campaign(
- storage, "study", commit="a" * 40, epochs=5001, arc_modes=["dur", "stk"]
- )
- assert storage.report is None
- else:
- assert publish_campaign(
- storage, "study", commit="a" * 40, epochs=5001, arc_modes=["dur", "stk"]
- )
- assert storage.report["arc_modes"] == ["dur", "stk"]
-
-
-def test_campaign_plan_shares_oat_and_schedules_both_arc_variants():
- from egomimic.benchmarks.libero.campaign import campaign_plan
- from egomimic.benchmarks.libero.catalog import TASKS
-
- plan = campaign_plan("/data", "/results")
- jobs = {job["id"]: job for job in plan["jobs"]}
- assert len(jobs) == len(plan["jobs"]) == 42
- for suite in TASKS:
- assert jobs[f"{suite}/oat"]["requires"] == [f"{suite}/tokenizer"]
- for mode in ("dur", "stk"):
- assert (
- f"+experiment=sim/oat/libero_arc_{mode}_policy"
- in jobs[f"{suite}/arc_{mode}"]["argv"]
- )
- assert (
- f"{suite}/arc_{mode}_rollout"
- in jobs[f"complete_comparison_{mode}"]["requires"]
- )
-
-
-def test_full_training_converts_only_verified_cached_demonstrations(
- tmp_path, monkeypatch
-):
- import hashlib
- import io
-
- from egomimic.benchmarks.libero import cluster, convert
-
- raw = tmp_path / "cache/libero_90/task_demo.hdf5"
- raw.parent.mkdir(parents=True)
- raw.write_bytes(b"official demonstration")
- metadata = {
- "sha": cluster.DATA_REVISION,
- "siblings": [
- {
- "rfilename": "libero_90/task_demo.hdf5",
- "lfs": {"sha256": hashlib.sha256(raw.read_bytes()).hexdigest()},
- }
- ],
- }
- monkeypatch.setenv("LIBERO_RAW_CACHE", str(raw.parent.parent))
- monkeypatch.setattr(cluster, "TASKS", {"libero_90": ["task"]})
- monkeypatch.setattr(
- cluster.urllib.request,
- "urlopen",
- lambda *a, **k: io.BytesIO(json.dumps(metadata).encode()),
- )
- monkeypatch.setattr(
- cluster, "download", lambda *a, **k: pytest.fail("Cache must not download")
- )
- calls = []
- monkeypatch.setattr(convert, "convert_suite", lambda *args: calls.append(args))
- evidence = tmp_path / "evidence"
- evidence.mkdir()
- result = stage_dataset(tmp_path / "root", "libero_90", evidence)
- assert calls == [(raw.parent, result, "libero_90")]
- raw.write_bytes(b"corrupt")
- with pytest.raises(ValueError, match="checksum"):
- stage_dataset(tmp_path / "root", "libero_90", evidence)
- assert raw.read_bytes() == b"corrupt" and len(calls) == 1
-
-
-def test_checkpoint_snapshots_are_content_addressed_and_receipted(
- tmp_path, monkeypatch
-):
- import boto3
-
- class Storage:
- def __init__(self):
- self.objects = {}
-
- def list_objects_v2(self, **kwargs):
- return {"KeyCount": 0}
-
- def upload_file(self, source, bucket, key, ExtraArgs):
- self.objects[key] = Path(source).read_bytes()
- assert ExtraArgs["Metadata"]["sha256"] == digest(source)
-
- storage = Storage()
- monkeypatch.setattr(boto3, "client", lambda *args, **kwargs: storage)
- for name in ("R2_ENDPOINT_URL", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY"):
- monkeypatch.setenv(name, "test")
- evidence = tmp_path / "evidence"
- evidence.mkdir()
- path = evidence / "last.ckpt"
- path.write_bytes(b"first checkpoint")
- uploader = ArtifactUploader(evidence, "test")
- uploader.upload(final=True)
- first_hash = digest(path)
- path.write_bytes(b"final checkpoint")
- uploader.upload(final=True)
- uploader.upload(final=True)
- receipts = json.loads((evidence / "checkpoint-receipts.json").read_text())
- assert receipts["last.ckpt"]["sha256"] == digest(path)
- assert any(first_hash in key for key in storage.objects)
- assert any(digest(path) in key for key in storage.objects)
-
-
-@pytest.mark.parametrize("invalid", [None, "budget", "hash", "incomplete", "outside"])
-def test_recovery_verifies_checkpoints_and_complete_training(tmp_path, invalid):
- import io
- import shutil
-
- import torch
-
- source = tmp_path / "source.ckpt"
- torch.save(
- {
- "global_step": 2,
- "loops": {
- "fit_loop": {
- "epoch_progress": {
- "current": {"completed": 0 if invalid == "incomplete" else 1}
- }
- }
- },
- },
- source,
- )
- prefix = "experiments/arc-oat-20260919/source-run/"
- runtime = {
- "suite": "libero_10",
- "mode": "full" if invalid == "budget" else "smoke",
- "epochs": 1,
- }
- receipt = {
- "sha256": "bad" if invalid == "hash" else digest(source),
- "bytes": source.stat().st_size,
- "uri": "s3://rldb/"
- + ("another-run/" if invalid == "outside" else prefix + "checkpoints/")
- + "last.ckpt",
- }
- receipts = {
- f"training/{method}/checkpoints/last.ckpt": receipt
- for method in ("tokenizer", "oat", "arc")
- }
-
- class Storage:
- def get_object(self, Bucket, Key):
- value = runtime if Key.endswith("runtime.json") else receipts
- return {"Body": io.BytesIO(json.dumps(value).encode())}
-
- def download_file(self, bucket, key, destination):
- shutil.copyfile(source, destination)
-
- evidence = tmp_path / "evidence"
- if invalid is not None:
- with pytest.raises(ValueError):
- restore_checkpoints(
- Storage(),
- "source-run",
- evidence,
- suite="libero_10",
- mode="smoke",
- epochs=5001,
- )
- else:
- restore_checkpoints(
- Storage(),
- "source-run",
- evidence,
- suite="libero_10",
- mode="smoke",
- epochs=5001,
- )
- assert (
- len(
- json.loads((evidence / "recovered-training.json").read_text())[
- "checkpoints"
- ]
- )
- == 3
- )
-
-
-@pytest.mark.parametrize("invalid", [None, "optimizer", "normalizer", "budget"])
-@pytest.mark.parametrize(
- "methods", [("tokenizer", "oat", "arc"), ("tokenizer", "oat", "arc_dur", "arc_stk")]
-)
-def test_partial_resume_preserves_only_existing_verified_checkpoints(
- tmp_path, invalid, methods
-):
- import io
- import shutil
-
- import torch
-
- payload = {
- "global_step": 140,
- "loops": {"fit_loop": {"epoch_progress": {"current": {"completed": 2}}}},
- "training_budget": {
- "global_batch_size": 256 if invalid == "budget" else 1024,
- "epochs": 5001,
- },
- "optimizer_states": [{"state": {}}],
- "normalizer_state": {},
- }
- if invalid == "optimizer":
- del payload["optimizer_states"]
- if invalid == "normalizer":
- del payload["normalizer_state"]
- source = tmp_path / "partial.ckpt"
- torch.save(payload, source)
- runtime = {
- "suite": "libero_10",
- "mode": "full",
- "epochs": 5001,
- "global_batch_size": 1024,
- }
- receipts = {
- "training/tokenizer/checkpoints/last.ckpt": {
- "sha256": digest(source),
- "bytes": source.stat().st_size,
- "uri": "s3://rldb/experiments/arc-oat-20260919/source-run/checkpoints/partial.ckpt",
- }
- }
-
- class Storage:
- def get_object(self, Bucket, Key):
- value = runtime if Key.endswith("runtime.json") else receipts
- return {"Body": io.BytesIO(json.dumps(value).encode())}
-
- def download_file(self, bucket, key, destination):
- shutil.copyfile(source, destination)
-
- if invalid:
- with pytest.raises(ValueError, match="Partial checkpoint"):
- restore_checkpoints(
- Storage(),
- "source-run",
- tmp_path / "evidence",
- suite="libero_10",
- mode="full",
- epochs=5001,
- allow_partial=True,
- methods=methods,
- )
- else:
- restored = restore_checkpoints(
- Storage(),
- "source-run",
- tmp_path / "evidence",
- suite="libero_10",
- mode="full",
- epochs=5001,
- allow_partial=True,
- methods=methods,
- )
- assert set(restored) == {"tokenizer"}
- assert restored["tokenizer"]["global_step"] == 140
- assert restored["tokenizer"]["complete"] is False
- assert not (tmp_path / "evidence/training/oat").exists()
-
-
-@pytest.mark.parametrize(
- "invalid",
- [
- None,
- "pending",
- "suite",
- "protocol",
- "spec",
- "codec",
- "retention",
- "incomplete",
- "evaluated",
- "unpermitted_gap",
- "mode",
- ],
-)
-@pytest.mark.parametrize("arc_mode", ["joint_dur", "stk", "dur"])
-def test_arc_training_requires_matching_confirmed_replay(
- tmp_path, monkeypatch, invalid, arc_mode
-):
- import hashlib
- import io
- import subprocess
-
- import yaml
-
- from egomimic.benchmarks.libero.replay import candidate_id
-
- root = Path(__file__).parents[1]
- spec = yaml.safe_load(
- (root / "egomimic/hydra_configs/benchmark/libero_arc_replay.yaml").read_text()
- )
- if invalid in {"evaluated", "unpermitted_gap"}:
- spec.update(
- selection_objective="success_then_tokens",
- maximum_success_rate_drop=1.0,
- allow_reference_gap=invalid == "evaluated",
- )
- codec = {"num_waypoints": 16, "max_translation": 0.2, "max_rotation_degrees": 48}
- if arc_mode != "joint_dur":
- codec["mode"] = spec["arc_mode"] = arc_mode
- selected = candidate_id(codec)
- result = {
- "confirmed": invalid not in {"pending", "evaluated", "unpermitted_gap"},
- "confirmation_complete": True,
- "arc_mode": "unmatched" if invalid == "mode" else arc_mode,
- "suite": "libero_10",
- "codec": codec,
- "candidate_id": selected,
- "selected_before_confirmation": True,
- "spec_sha256": hashlib.sha256(
- (json.dumps(spec, indent=2) + "\n").encode()
- ).hexdigest(),
- "confirmation": {
- "raw": {"success_rate": 1, "episodes": 50},
- "raw_repeat": {"max_state_l2_vs_raw_mean": 0, "episodes": 50},
- "dense": {"retention": 1, "max_action_mse": 0, "episodes": 50},
- selected: {
- "retention": 0.9 if invalid == "retention" else 1,
- "success_rate": 0.9
- if invalid in {"retention", "evaluated", "unpermitted_gap"}
- else 1,
- "raw_successes": 50,
- "episodes": 50,
- "successes": 50,
- "action_mse": 0.1,
- },
- },
- }
- runtime = {
- "suite": "libero_goal" if invalid == "suite" else "libero_10",
- "source_commit": "a" * 40,
- }
- if invalid == "protocol":
- spec["execute_steps"] = 8
- if invalid == "spec":
- result["spec_sha256"] = "wrong"
- if invalid == "incomplete":
- result["confirmation"][selected]["episodes"] = 49
- content = (root / "egomimic/rldb/zarr/libero_arc.py").read_bytes()
- monkeypatch.setattr(subprocess, "run", lambda *args, **kwargs: None)
- monkeypatch.setattr(
- subprocess,
- "check_output",
- lambda argv, **kwargs: (
- b"different codec"
- if invalid == "codec"
- else (root / argv[-1].split(":", 1)[1]).read_bytes()
- ),
- )
-
- class Storage:
- def get_object(self, Bucket, Key):
- value = {
- "status.json": {
- "state": "REPLAY_EVALUATED"
- if invalid in {"evaluated", "unpermitted_gap"}
- else "CONFIRMED"
- },
- "result.json": result,
- "runtime.json": runtime,
- "spec.json": spec,
- }[Key.rsplit("/", 1)[-1]]
- return {"Body": io.BytesIO(json.dumps(value).encode())}
-
- if invalid and invalid != "evaluated":
- with pytest.raises((ValueError, CalibrationPending)):
- load_arc_calibration(
- Storage(), "replay-run", tmp_path, suite="libero_10", arc_mode=arc_mode
- )
- assert not (tmp_path / "arc-calibration.json").exists()
- else:
- overrides = load_arc_calibration(
- Storage(), "replay-run", tmp_path, suite="libero_10", arc_mode=arc_mode
- )
- expected = {
- "arc_waypoints": 16,
- "arc_max_translation": 0.2,
- "arc_max_rotation_degrees": 48,
- }
- if arc_mode != "joint_dur":
- expected.update(arc_mode=arc_mode, arc_action_dim=12)
- if arc_mode == "stk":
- expected["arc_velocity_norm_bound"] = 3**0.5
- assert overrides == expected
- receipt = (
- "arc-calibration.json"
- if arc_mode == "joint_dur"
- else f"arc-{arc_mode}-calibration.json"
- )
- assert (
- json.loads((tmp_path / receipt).read_text())["verified_codec_sha256"]
- == hashlib.sha256(content).hexdigest()
- )
- with pytest.raises(CalibrationPending, match="requires"):
- load_arc_calibration(Storage(), None, tmp_path, suite="libero_10")
diff --git a/tests/test_libero_multi_embodiment_native.py b/tests/test_libero_multi_embodiment_native.py
deleted file mode 100644
index 17a1932d32..0000000000
--- a/tests/test_libero_multi_embodiment_native.py
+++ /dev/null
@@ -1,1286 +0,0 @@
-import json
-
-import h5py
-import numpy as np
-import pytest
-import torch
-import zarr
-
-from egomimic.benchmarks.libero.catalog import TASK_IDS, get_tasks
-from egomimic.benchmarks.libero.convert import (
- convert_collection_manifest,
- native_episode_values,
-)
-from egomimic.pipeline.stages_sampler import KeyedFeatureProjection
-from egomimic.rldb.zarr.libero_dataset import (
- LiberoDataset,
- LiberoMultiEmbodimentNormalizer,
- LiberoReplayResolver,
- keymap,
-)
-
-
-def replay(path, action_dim, hand_dim, demos_per_task=2):
- tasks = get_tasks("libero_spatial")
- ids = np.repeat([TASK_IDS[t][2] for t in tasks for _ in range(demos_per_task)], 2)
- n = len(ids)
- rng = np.random.default_rng(4)
- group = zarr.open_group(str(path), mode="w", zarr_format=2)
- data = group.create_group("data")
- for key, values in {
- "action": rng.uniform(-1, 1, (n, action_dim)).astype("float32"),
- "robot0_gripper_qpos": rng.uniform(0, 1, (n, hand_dim)).astype("float32"),
- "robot0_eef_pos": rng.uniform(0, 1, (n, 3)).astype("float32"),
- "robot0_eef_quat": np.tile(np.array([0, 0, 0, 1], dtype="float32"), (n, 1)),
- "task_uid": ids[:, None],
- "agentview_rgb": rng.integers(0, 255, (n, 2, 2, 3), dtype="uint8"),
- "robot0_eye_in_hand_rgb": rng.integers(0, 255, (n, 2, 2, 3), dtype="uint8"),
- }.items():
- data.create_array(key, data=values)
- group.create_group("meta").create_array("episode_ends", data=np.arange(2, n + 1, 2))
- return group
-
-
-def resolver(path, sharpa=False):
- return LiberoReplayResolver(
- path,
- keymap(2, 32),
- "libero_spatial",
- **(
- dict(action_dim=28, gripper_state_dim=22, embodiment="libero_panda_sharpa")
- if sharpa
- else {}
- ),
- )
-
-
-def test_distinct_native_widths_and_reject_sharpa_as_original(tmp_path):
- path = tmp_path / "sharpa.zarr"
- replay(path, 28, 22)
- with pytest.raises(ValueError, match="action must"):
- resolver(path).resolve()
- sample = next(iter(resolver(path, True).resolve().values()))[0]
- assert sample["embodiment"] == 22
- assert sample["actions"].shape == (32, 28)
- assert sample["robot0_gripper_qpos"].shape == (2, 22)
- assert torch.equal(sample["actions"][1], sample["actions"][-1])
- with pytest.raises(ValueError, match="Sharpa contract"):
- LiberoReplayResolver(
- path,
- keymap(),
- "libero_spatial",
- action_dim=7,
- gripper_state_dim=22,
- embodiment="libero_panda_sharpa",
- )
-
-
-def test_configured_replay_split_proves_rows_and_rejects_leakage(tmp_path):
- from egomimic.rldb.zarr.libero_saved_state import make_replay_split_proof
-
- path = tmp_path / "replay.zarr"
- replay(path, 7, 2)
- r = resolver(path)
- train = LiberoDataset._from_resolver(r, mode="train", valid_ratio=0.1)
- valid = LiberoDataset._from_resolver(r, mode="valid", valid_ratio=0.1)
- kwargs = dict(expected_episodes=20, valid_ratio=0.1, seed=42)
- proof = make_replay_split_proof(train, valid, **kwargs)
- assert len(proof["train"]) == 18 and len(proof["valid"]) == 2
- rows = sorted(proof["train"] + proof["valid"], key=lambda row: row["start"])
- assert [(row["start"], row["end"]) for row in rows] == [
- (i, i + 2) for i in range(0, 40, 2)
- ]
- assert set(row["physical_id"] for row in proof["train"]).isdisjoint(
- row["physical_id"] for row in proof["valid"]
- )
- with pytest.raises(ValueError, match="coverage"):
- make_replay_split_proof(train, valid, **{**kwargs, "expected_episodes": 21})
- key = next(iter(valid.datasets))
- valid.datasets[key] = next(iter(train.datasets.values()))
- with pytest.raises(ValueError, match="physical rows"):
- make_replay_split_proof(train, valid, **kwargs)
-
-
-def test_joint_operational_split_materializes_450_50_and_rejects_wrong_binding(
- tmp_path, monkeypatch
-):
- import copy
- from pathlib import Path
-
- import hydra
- from hydra.core.hydra_config import HydraConfig
- from omegaconf import OmegaConf
-
- from egomimic.benchmarks.libero.native_launch_profiles import (
- SPATIAL_HAND_JOINT_PROFILE,
- )
- from scripts.train.libero_native_launch_contract import (
- joint_split_contract,
- normalize_joint_context,
- sha,
- )
-
- root = Path(__file__).resolve().parents[1]
- manifests = {}
- for name, width, hand in (("libero_panda", 7, 2), ("libero_panda_sharpa", 69, 22)):
- path = tmp_path / name
- replay(path, width, hand, demos_per_task=50)
- manifest = dict(complete=True, action_dim=width)
- if width == 69:
- manifest.update(
- action_representation="wrist_hand_keypoints",
- keypoint_frame="wrist_relative_Quest_XYZ_m",
- observation_alignment="causal_pre_action_raw_lossless",
- )
- (path / "benchmark_manifest.json").write_text(json.dumps(manifest))
- manifests[name] = sha(path / "benchmark_manifest.json")
- monkeypatch.setenv("LIBERO_ORIGINAL_SPATIAL_ROOT", str(tmp_path / "libero_panda"))
- monkeypatch.setenv(
- "LIBERO_SHARPA_SPATIAL_ROOT", str(tmp_path / "libero_panda_sharpa")
- )
- monkeypatch.setenv("LIBERO_DIAGNOSTIC_SEED_BANK", str(tmp_path / "seeds.json"))
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(root / "egomimic/hydra_configs")
- ):
- cfg = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=[
- "hydra/launcher=basic",
- "+experiment=" + SPATIAL_HAND_JOINT_PROFILE.experiment,
- "trainer.precision=bf16-mixed",
- "trainer.devices=1",
- ],
- )
- cfg.hydra.runtime.output_dir = str(tmp_path / "output")
- HydraConfig.instance().set_config(cfg)
- resolved = OmegaConf.to_container(
- OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]), resolve=True
- )
- result = joint_split_contract(resolved, manifests, expected_episodes=500)
- assert result["status"] == "PASS" and result["launch_ready"] is False
- for proof in result["sources"].values():
- assert len(proof["train"]) == 450 and len(proof["valid"]) == 50
- assert proof["frames"] == 1000 and proof["complete_episode_and_frame_coverage"]
- with pytest.raises(ValueError, match="manifest SHA"):
- joint_split_contract(
- resolved, {**manifests, "libero_panda": "0" * 64}, expected_episodes=500
- )
- wrong = copy.deepcopy(resolved)
- for mode in ("train", "valid"):
- wrong["data"][mode + "_datasets"]["libero_panda_sharpa"]["resolver"][
- "folder_path"
- ] = str(tmp_path / "libero_panda")
- with pytest.raises(ValueError, match="distinct"):
- joint_split_contract(wrong, manifests, expected_episodes=500)
- # Real operational handoff: physical proof from a YAML config -> configured
- # DataModule fitting -> complete context, with no model or trainer.
- original_config = tmp_path / "split-config.yaml"
- OmegaConf.save(OmegaConf.create(resolved), original_config)
- result["resolved_config"] = dict(
- path=str(original_config), sha256=sha(original_config)
- )
- split_path = tmp_path / "split.json"
- split_path.write_text(json.dumps(result))
- fit = copy.deepcopy(resolved)
- fit["norm_stats_only"] = True
- fit["norm_stats"].update(
- precomputed_norm_path=None,
- save_cache_dir=str(tmp_path / "stats"),
- export_context_path=str(tmp_path / "context.json"),
- )
- fit["run_provenance"].update(
- split_manifest_path=str(split_path), split_manifest_sha256=sha(split_path)
- )
- fit["data"]["run_provenance"] = copy.deepcopy(fit["run_provenance"])
- norm = normalize_joint_context(fit)
- assert norm["status"] == "PASS" and norm["launch_ready"] is False
- context = json.loads((tmp_path / "context.json").read_text())["data_context"]
- assert set(context["normalizer_state"]["norm_stats"]) == {"21", "22"}
- assert (
- len(context["normalizer_state"]["norm_stats"]["22"]["actions"]["scale"]) == 69
- )
- # The checkpoint verifier must restore the same complete state without data IO.
- from egomimic.pipeline.checkpoint_binding import checkpoint_binding
- from egomimic.rldb.zarr.data_module import load_data_context
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_checkpoint_context,
- )
-
- checkpoint_config = copy.deepcopy(fit)
- checkpoint_config["norm_stats"]["precomputed_norm_path"] = str(
- tmp_path / "context.json"
- )
- checkpoint_config["run_provenance"].update(
- source_commit="1" * 40, normalization_sha256=sha(tmp_path / "context.json")
- )
- restored_context = load_data_context(tmp_path / "context.json")
- checkpoint = dict(
- hyper_parameters=dict(config_tree=checkpoint_config),
- data_context=restored_context.snapshot(),
- inference_binding=checkpoint_binding(checkpoint_config, restored_context),
- )
- assert (
- validate_joint_checkpoint_context(checkpoint, checkpoint_config)[
- "checkpoint_binding"
- ]
- == checkpoint["inference_binding"]
- )
- # Native checkpoints keep integer embodiment keys/tuples; JSON exports do not.
- native_context = copy.deepcopy(checkpoint["data_context"])
- native_context["normalizer_state"] = restored_context.normalizer.to_state()
- assert native_context != checkpoint["data_context"]
- checkpoint["data_context"] = native_context
- assert (
- validate_joint_checkpoint_context(checkpoint, checkpoint_config)[
- "checkpoint_binding"
- ]
- == checkpoint["inference_binding"]
- )
- wrong_checkpoint = copy.deepcopy(checkpoint)
- wrong_checkpoint["data_context"].pop("preprocessing")
- with pytest.raises(ValueError, match="full data context differs"):
- validate_joint_checkpoint_context(wrong_checkpoint, checkpoint_config)
- # Consume genuine fixture context/split plus explicitly synthetic constructor/runtime metadata.
- from scripts.train.libero_native_launch_contract import (
- contract_sha,
- joint_preflight_evidence,
- reference,
- )
-
- pre = copy.deepcopy(checkpoint_config)
- pre["norm_stats_only"] = False
- pre["trainer"]["accelerator"] = "cpu"
- pre["run_provenance"]["run_kind"] = "preflight"
- config_path = tmp_path / "preflight.yaml"
- OmegaConf.save(OmegaConf.create(pre), config_path)
- constructor_path = tmp_path / "synthetic-constructor.json"
- constructor = dict(
- status="PASS",
- full_model_constructor=True,
- parameter_count=78182614,
- source_commit="1" * 40,
- resolved_config_sha256=sha(config_path),
- model_contract_sha256=contract_sha(pre),
- )
- constructor_path.write_text(json.dumps(constructor))
- runtime_path = tmp_path / "synthetic-runtime.json"
- runtime_path.write_text(
- json.dumps(
- {
- "python": {},
- "installed_distributions": {},
- "core_distribution_versions": {},
- "dependency_inputs": {},
- "torch_runtime": {},
- "source": {"head": "1" * 40, "dirty": False},
- }
- )
- )
- proof = joint_preflight_evidence(
- pre,
- config_path,
- constructor_ref=reference(constructor_path),
- runtime_ref=reference(runtime_path),
- source_head="1" * 40,
- )
- assert proof["launch_ready"] is False
- assert proof["complete_context"]["sha256"] == sha(tmp_path / "context.json")
- constructor["source_commit"] = "0" * 40
- constructor_path.write_text(json.dumps(constructor))
- with pytest.raises(ValueError, match="never relabel cached"):
- joint_preflight_evidence(
- pre,
- config_path,
- constructor_ref=reference(constructor_path),
- runtime_ref=reference(runtime_path),
- source_head="1" * 40,
- )
- with pytest.raises(ValueError, match="fresh absolute"):
- normalize_joint_context(fit)
- fit["norm_stats"]["export_context_path"] = str(tmp_path / "other-context.json")
- result["sources"]["libero_panda"]["episodes"] = 499
- split_path.write_text(json.dumps(result))
- fit["run_provenance"]["split_manifest_sha256"] = sha(split_path)
- with pytest.raises(ValueError, match="split/content binding"):
- normalize_joint_context(fit)
- fit["data"]["train_datasets"]["libero_panda"]["valid_ratio"] = 0.01
- with pytest.raises(ValueError, match="episode split differs"):
- normalize_joint_context(fit)
-
-
-@pytest.mark.parametrize(
- "action", ["joint-normalize", "joint-preflight", "joint-compose", "joint-prepare"]
-)
-@pytest.mark.parametrize("step", ["batch", "extern"])
-def test_joint_data_cli_rejects_non_numeric_scheduler_step(
- tmp_path, monkeypatch, step, action
-):
- import sys
-
- from scripts.train.libero_native_launch_contract import main
-
- monkeypatch.setenv("SLURM_JOB_ID", "1234")
- monkeypatch.setenv("SLURM_STEP_ID", step)
- monkeypatch.setattr(
- sys,
- "argv",
- [
- "contract",
- "--action",
- action,
- "--repo",
- str(tmp_path),
- "--output",
- str(tmp_path / "out"),
- ],
- )
- with pytest.raises(ValueError, match="scheduled numeric srun"):
- main()
- assert not (tmp_path / "out").exists()
-
-
-def test_train_only_limits_separate_bodies_and_roundtrip(tmp_path):
- datasets = {}
- for name, width, hand in [("libero_panda", 7, 2), ("libero_panda_sharpa", 28, 22)]:
- path = tmp_path / (name + ".zarr")
- group = replay(path, width, hand)
- r = resolver(path, width == 28)
- train = LiberoDataset._from_resolver(r, mode="train", valid_ratio=0.1)
- valid = LiberoDataset._from_resolver(r, mode="valid", valid_ratio=0.1)
- assert not set(train.datasets) & set(valid.datasets)
- for episode in valid.datasets.values():
- group["data/action"][episode.start : episode.end] = 1000
- datasets[name] = train
- norm = LiberoMultiEmbodimentNormalizer()
- norm.populate_from_datasets(datasets)
- for name, dataset in datasets.items():
- norm.infer_shapes_from_batch(dataset[0])
- norm.infer_norm_from_dataset(dataset, name)
- ident = dataset.resolver.embodiment_id
- assert norm.norm_stats[ident]["actions"]["max"].max() < 1
- dataset.set_norm_stats_from(norm)
- sample = dataset[0]
- raw = next(iter(dataset.datasets.values()))[0]
- restored = norm.unnormalize({"actions": sample["actions"]}, ident)
- torch.testing.assert_close(restored["actions"], raw["actions"])
- assert norm.norm_stats[21]["actions"]["scale"].shape == (7,)
- assert norm.norm_stats[22]["actions"]["scale"].shape == (28,)
- restored = LiberoMultiEmbodimentNormalizer.from_state(norm.to_state())
- assert restored.tokenizer_context() == norm.tokenizer_context()
- valid = LiberoDataset._from_resolver(
- datasets["libero_panda"].resolver, mode="valid"
- )
- with pytest.raises(ValueError, match="training episodes"):
- LiberoMultiEmbodimentNormalizer().infer_norm_from_dataset(valid, "invalid")
-
-
-def test_complete_joint_context_export_restores_without_refitting(
- tmp_path, monkeypatch
-):
- import copy
-
- from omegaconf import OmegaConf
-
- from egomimic.rldb.zarr.data_module import ZarrDataModule
-
- train, valid = {}, {}
- for name, width, hand in (("libero_panda", 7, 2), ("libero_panda_sharpa", 69, 22)):
- path = tmp_path / (name + ".zarr")
- replay(path, width, hand)
- if width == 69:
- (path / "benchmark_manifest.json").write_text(
- json.dumps(
- dict(
- complete=True,
- action_dim=69,
- action_representation="wrist_hand_keypoints",
- keypoint_frame="wrist_relative_Quest_XYZ_m",
- observation_alignment="causal_pre_action_raw_lossless",
- )
- )
- )
- train[name] = dict(
- _target_="egomimic.rldb.zarr.libero_dataset.LiberoDataset._from_resolver",
- resolver=dict(
- _target_="egomimic.rldb.zarr.libero_dataset.LiberoReplayResolver",
- folder_path=str(path),
- suite="libero_spatial",
- embodiment=name,
- action_dim=width,
- gripper_state_dim=hand,
- action_representation="native"
- if width == 7
- else "wrist_hand_keypoints",
- key_map=dict(
- _target_="egomimic.rldb.zarr.libero_dataset.keymap",
- n_obs_steps=2,
- horizon=32,
- ),
- ),
- mode="train",
- valid_ratio=0.1,
- split_seed=42,
- )
- valid[name] = {**copy.deepcopy(train[name]), "mode": "valid"}
-
- def module():
- return ZarrDataModule(OmegaConf.create(train), OmegaConf.create(valid), {}, {})
-
- normalizer = OmegaConf.create(
- dict(
- _target_="egomimic.rldb.zarr.libero_dataset.LiberoMultiEmbodimentNormalizer"
- )
- )
- exported = tmp_path / "complete-context.json"
- context = module().prepare_context(
- mode="normalization",
- normalizer=normalizer,
- normalization=dict(
- norm_mode="libero_train_limits",
- save_cache_dir=str(tmp_path / "stats"),
- export_context_path=str(exported),
- ),
- )
- payload = json.loads(exported.read_text())["data_context"]
- assert payload["normalizer_config"] == OmegaConf.to_container(normalizer)
- assert set(payload["preprocessing"]) == {"21", "22"}
- assert (
- payload["normalizer_state"]["native_embodiment_contexts"]["22"]["action_dim"]
- == 69
- )
-
- def forbidden(*args, **kwargs):
- raise AssertionError("A complete-context cache must not refit normalization")
-
- monkeypatch.setattr(
- LiberoMultiEmbodimentNormalizer, "infer_norm_from_dataset", forbidden
- )
- restored = module().prepare_context(
- mode="train",
- normalizer=normalizer,
- normalization=dict(
- norm_mode="libero_train_limits", precomputed_norm_path=str(exported)
- ),
- )
- assert (
- restored.normalizer.tokenizer_context()
- == context.normalizer.tokenizer_context()
- )
- assert restored.fingerprint() == context.fingerprint()
- # A changed export target is a hard collision, not an overwrite or refit.
- exported.write_text('{"different": true}\n')
- with pytest.raises(ValueError, match="Refusing to replace"):
- module().prepare_context(
- mode="normalization",
- normalizer=normalizer,
- restored_state=payload,
- normalization=dict(
- norm_mode="libero_train_limits",
- save_cache_dir=str(tmp_path / "other-stats"),
- export_context_path=str(exported),
- ),
- )
- assert exported.read_text() == '{"different": true}\n'
- assert not list(tmp_path.glob(".context-*"))
-
-
-def test_original_state_is_identity_and_sharpa_projector_gets_gradient():
- stage = KeyedFeatureProjection(
- selector_key="embodiment",
- input_key=["eef", "quat", "hand", "task"],
- output_key="proprio_condition",
- output_dim=10,
- selector_aliases={21: "panda", 22: "sharpa"},
- projections={
- "panda": {"input_dim": 10, "identity": True},
- "sharpa": {"input_dim": 30, "hidden_dim": 32},
- },
- )
- for ident, hand_width in [(21, 2), (22, 22)]:
- batch = {
- "embodiment": torch.tensor([ident, ident]),
- "eef": torch.randn(2, 2, 3),
- "quat": torch.randn(2, 2, 4),
- "hand": torch.randn(2, 2, hand_width),
- "task": torch.randn(2, 2, 1),
- }
- expected = torch.cat([batch[k] for k in ("eef", "quat", "hand", "task")], -1)
- result = stage(batch)["proprio_condition"]
- assert result.shape == (2, 2, 10)
- if ident == 21:
- assert torch.equal(result, expected)
- else:
- result.square().mean().backward()
- assert all(
- p.grad is not None and torch.isfinite(p.grad).all()
- for p in stage.projections["sharpa"].parameters()
- )
- with pytest.raises(ValueError, match="must match"):
- KeyedFeatureProjection(
- "body",
- "x",
- "y",
- {"panda": {"input_dim": 9, "identity": True}},
- output_dim=10,
- )
-
-
-def test_native_collection_keeps_upright_rows_quaternion_and_timestamps(tmp_path):
- root = tmp_path / "raw"
- root.mkdir()
- records = []
- image = np.zeros((2, 2, 2, 3), dtype="uint8")
- image[:, 0, :, 0] = 240
- image[:, 1, :, 2] = 170
- for i, task in enumerate(get_tasks("libero_spatial")):
- filename = f"task-{i}.hdf5"
- with h5py.File(root / filename, "w") as handle:
- handle.attrs["metadata_json"] = json.dumps(
- {"suite": "libero_spatial", "task": task, "embodiment": "sharpa"}
- )
- handle.attrs["automatic_success"] = True
- data = handle.create_group("data")
- data.attrs["rgb_row_transform"] = (
- "vertical_flip_from_raw_opengl_to_original_zarr_upright"
- )
- demo = data.create_group("demo_0")
- demo.attrs["observation_alignment"] = "after_action; raw timeline retained"
- demo.create_dataset("actions", data=np.zeros((2, 28), dtype="float32"))
- demo.create_dataset("action_host_ns", data=[10, 30])
- demo.create_dataset("observation_host_ns", data=[20, 40])
- obs = demo.create_group("obs")
- obs.create_dataset("agentview_rgb", data=image)
- obs.create_dataset("eye_in_hand_rgb", data=image)
- obs.create_dataset("ee_pos", data=np.zeros((2, 3)))
- obs.create_dataset("ee_quat", data=[[0, 0, 0, 1]] * 2)
- obs.create_dataset("ee_ori", data=np.zeros((2, 3)))
- obs.create_dataset("gripper_states", data=np.zeros((2, 22)))
- legacy_view = native_episode_values(demo, task, action_dim=28, hand_dim=22)
- np.testing.assert_array_equal(legacy_view["agentview_rgb"], image[:, ::-1])
- records.append(
- {
- "path": filename,
- "task": task,
- "suite": "libero_spatial",
- "action_dimension": 28,
- "automatic_success": True,
- }
- )
- manifest = {
- "suite": "libero_spatial",
- "embodiment": "custom_Panda_SharpaWaveRight_LIBERO",
- "episodes": records,
- "episode_count": 10,
- "per_task": {r["task"]: 1 for r in records},
- }
- m = root / "manifest.json"
- m.write_text(json.dumps(manifest))
- output = tmp_path / "converted.zarr"
- receipt = convert_collection_manifest(root, output, m)
- arrays = zarr.open_group(str(output), mode="r")["data"]
- np.testing.assert_array_equal(arrays["agentview_rgb"][:2], image)
- np.testing.assert_array_equal(arrays["robot0_eef_quat"][:2], [[0, 0, 0, 1]] * 2)
- np.testing.assert_array_equal(arrays["action_host_ns"][:2], [10, 30])
- np.testing.assert_array_equal(arrays["observation_host_ns"][:2], [20, 40])
- assert receipt["complete"] and not receipt["split_assigned"]
- assert len(receipt["episodes"]) == 10
- manifest["episodes"][0]["path"] = "../outside.hdf5"
- m.write_text(json.dumps(manifest))
- with pytest.raises(ValueError, match="outside root"):
- convert_collection_manifest(root, tmp_path / "rejected.zarr", m)
- assert not (tmp_path / "rejected.zarr").exists()
-
-
-def test_keypoint_alignment_rejects_asynchronous_latest_guess(tmp_path):
- from egomimic.benchmarks.libero.convert import aligned_hand_keypoint_actions
-
- metadata = {
- "hand_calibration": {
- "operator_wrist_position_quest_m": [0, 0, 0],
- "position_world_m": [0, 0, 0],
- },
- "world_from_quest": np.eye(3).tolist(),
- "control_settings": {"translation_gain": 1.0, "vertical_gain": 1.0},
- }
- points = np.arange(63).reshape(21, 3) * 0.001
- packets = [
- {
- "source": "quest_webxr",
- "hand": "right",
- "joint_layout": "mediapipe21.v2",
- "tracking_valid": True,
- "joint_valid": [True] * 21,
- "sequence": i,
- "wrist": {"position_m": [x, 0, 0]},
- "joints_m": (points + x).tolist(),
- }
- for i, x in enumerate((0.01, 0.02))
- ]
- with h5py.File(tmp_path / "record.hdf5", "w") as h:
- h.create_dataset("packet_json", data=[json.dumps(p) for p in packets])
- h.create_dataset("packet_receive_ns", data=[10, 20])
- h.create_dataset("action_host_ns", data=[30, 40])
- h.create_dataset(
- "intent_target_position_world_m", data=[[0.01, 0, 0], [0.02, 0, 0]]
- )
- h.create_dataset("actions", data=np.zeros((2, 28)))
- h.create_dataset("tracking_valid", data=[True, True])
- actions, start, sequence = aligned_hand_keypoint_actions(h, metadata)
- assert actions.shape == (2, 69) and start == 0
- assert sequence.tolist() == [
- 0,
- 1,
- ] # Latest-at-action would wrongly choose 1 twice.
- np.testing.assert_allclose(
- actions[0, 6:], (points - points[0]).reshape(63), atol=1e-8
- )
- h["intent_target_position_world_m"][1] = [0.03, 0, 0]
- with pytest.raises(ValueError, match="correspondence at step 1"):
- aligned_hand_keypoint_actions(h, metadata)
- h["intent_target_position_world_m"][1] = [0.02, 0, 0]
- # Recorded IK failure pauses the control after applying a valid-hand
- # command; this is not a missing operator hand target.
- metadata["action_contract"] = {
- "arm_control": {"name": "mink_direct_joint_servo_v4", "mode": "mink"}
- }
- h.create_dataset("control_reason", data=[b"active", b"arm_ik_infeasible"])
- h["tracking_valid"][1] = False
- repaired, _, _ = aligned_hand_keypoint_actions(h, metadata)
- np.testing.assert_array_equal(repaired, actions)
- h["control_reason"][1] = b"tracking_lost"
- with pytest.raises(ValueError, match="hand tracking target"):
- aligned_hand_keypoint_actions(h, metadata)
- h["control_reason"][1] = b"arm_ik_infeasible"
- metadata["action_contract"]["arm_control"]["name"] = "unbound_controller"
- with pytest.raises(ValueError, match="hand tracking target"):
- aligned_hand_keypoint_actions(h, metadata)
- metadata["action_contract"]["arm_control"]["name"] = (
- "mink_direct_joint_servo_v4"
- )
- h.attrs["recenter_events_json"] = "[{}]"
- with pytest.raises(ValueError, match="recenter calibration"):
- aligned_hand_keypoint_actions(h, metadata)
-
-
-def test_keypoint_resolver_requires_declared_frame_manifest(tmp_path):
- path = tmp_path / "sharpa69.zarr"
- replay(path, 69, 22)
- r = LiberoReplayResolver(
- path,
- keymap(2, 32),
- "libero_spatial",
- action_dim=69,
- gripper_state_dim=22,
- embodiment="libero_panda_sharpa",
- action_representation="wrist_hand_keypoints",
- )
- manifest = path / "benchmark_manifest.json"
- manifest.write_text(
- json.dumps(
- {
- "complete": True,
- "action_dim": 69,
- "action_representation": "wrist_hand_keypoints",
- "keypoint_frame": "wrist_relative_Quest_XYZ_m",
- "observation_alignment": "causal_pre_action_raw_lossless",
- }
- )
- )
- assert next(iter(r.resolve().values()))[0]["actions"].shape == (32, 69)
- manifest.write_text(json.dumps({"complete": True, "action_dim": 69}))
- with pytest.raises(ValueError, match="frame manifest"):
- r.resolve()
-
-
-def test_sharpa_raw_pre_action_views_cannot_leak_post_action_frames(tmp_path):
- from egomimic.benchmarks.libero.convert import causal_collection_episode_values
-
- with h5py.File(tmp_path / "timeline.hdf5", "w") as h:
- h.create_dataset("actions", data=np.zeros((3, 28)))
- h.create_dataset("observation_host_ns", data=[10, 30, 50, 70])
- h.create_dataset("action_host_ns", data=[20, 40, 60])
- data = h.create_group("data")
- data.attrs["rgb_row_transform"] = (
- "vertical_flip_from_raw_opengl_to_original_zarr_upright"
- )
- rgb = h.create_group("rgb")
- pixels = np.zeros((4, 2, 2, 3), dtype=np.uint8)
- for t in range(4):
- pixels[t, 0, :, 0] = 10 + t
- pixels[t, 1, :, 2] = 100 + t
- rgb.create_dataset("agentview", data=pixels)
- rgb.create_dataset("robot0_eye_in_hand", data=pixels)
- obs = h.create_group("numeric_observations")
- obs.create_dataset("robot0_eef_pos", data=np.arange(12).reshape(4, 3))
- obs.create_dataset("robot0_eef_quat", data=[[0, 0, 0, 1]] * 4)
- obs.create_dataset(
- "robot0_gripper_qpos", data=np.repeat(np.arange(4)[:, None], 22, axis=1)
- )
- meta = {"observation_order": "before_action_and_terminal"}
- result = causal_collection_episode_values(
- h,
- meta,
- get_tasks("libero_spatial")[0],
- np.ones((2, 69), dtype=np.float32),
- 1,
- np.array([2, 3]),
- )
- np.testing.assert_array_equal(result["agentview_rgb"], pixels[1:3, ::-1])
- np.testing.assert_array_equal(
- result["robot0_eef_pos"], np.arange(12).reshape(4, 3)[1:3]
- )
- np.testing.assert_array_equal(result["observation_host_ns"], [30, 50])
- np.testing.assert_array_equal(result["action_host_ns"], [40, 60])
- assert np.all(result["observation_host_ns"] < result["action_host_ns"])
- assert result["robot0_gripper_qpos"][:, 0].tolist() == [1, 2]
- h.create_dataset(
- "control_reason", data=[b"operator_stop", b"stale_3d_view", b"active"]
- )
- resumed_target = np.ones((1, 69), dtype=np.float32)
- with pytest.raises(ValueError, match="resume proof"):
- causal_collection_episode_values(
- h,
- meta,
- get_tasks("libero_spatial")[0],
- resumed_target,
- 2,
- np.array([3]),
- )
- proof = {
- "status": "PASS_UNIQUE_SOURCE_BACKED_RESUME_CALIBRATION",
- "source_start_step": 2,
- "source_path": h.filename,
- "solutions": [{"recorded_anchor": True}],
- "no_later_gap_discarded": True,
- }
- resumed = causal_collection_episode_values(
- h,
- meta,
- get_tasks("libero_spatial")[0],
- resumed_target,
- 2,
- np.array([3]),
- resume_proof=proof,
- )
- np.testing.assert_array_equal(resumed["observation_host_ns"], [50])
- h["control_reason"][0] = b"active"
- with pytest.raises(ValueError, match="discard active"):
- causal_collection_episode_values(
- h,
- meta,
- get_tasks("libero_spatial")[0],
- resumed_target,
- 2,
- np.array([3]),
- resume_proof=proof,
- )
- h["observation_host_ns"][1] = 45
- with pytest.raises(ValueError, match="not causal"):
- causal_collection_episode_values(
- h,
- meta,
- get_tasks("libero_spatial")[0],
- np.ones((2, 69), dtype=np.float32),
- 1,
- np.array([2, 3]),
- )
-
-
-def test_routed_dit_half_preserves_three_owner_gradients_and_batch_contract():
- from copy import deepcopy
- from types import SimpleNamespace
-
- from torch.utils.checkpoint import checkpoint
-
- from egomimic.utils.libero_dit_half import HalfCheckpointScope, LiberoDiTHalf
-
- torch.manual_seed(4)
- backbones = {
- name: SimpleNamespace(
- blocks=torch.nn.ModuleList([torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)]),
- gradient_checkpointing=True,
- )
- for name in ("encoder/libero_panda", "encoder/libero_panda_sharpa", "velocity")
- }
- baseline = deepcopy(backbones)
- native = SimpleNamespace(checkpoint=checkpoint)
- scope = HalfCheckpointScope(native, backbones, expected_depth=2)
- inputs = torch.randn(2, 4)
-
- def evaluate(owners, dispatcher):
- x = inputs.clone().requires_grad_()
- values = []
- for owner in owners.values():
- value = x
- for block in owner.blocks:
- value = dispatcher(block, value, use_reentrant=False)
- values.append(value)
- output = sum(values)
- output.square().sum().backward()
- return output, x.grad
-
- expected, expected_gradient = evaluate(baseline, checkpoint)
- scope.install()
- try:
- actual, actual_gradient = evaluate(backbones, native.checkpoint)
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
- torch.testing.assert_close(actual_gradient, expected_gradient, rtol=0, atol=0)
- for name in backbones:
- for actual_block, reference_block in zip(
- backbones[name].blocks, baseline[name].blocks
- ):
- torch.testing.assert_close(
- actual_block.weight.grad,
- reference_block.weight.grad,
- rtol=0,
- atol=0,
- )
- assert all(count > 0 for count in scope.counts.values())
- callback = LiberoDiTHalf({"libero_panda": 16, "libero_panda_sharpa": 16})
- callback.scope = scope
- batch = {
- name: {"actions": torch.zeros(16, 32, width)}
- for name, width in [("libero_panda", 7), ("libero_panda_sharpa", 69)]
- }
- callback.on_train_batch_start(None, None, batch, 0)
- batch["libero_panda_sharpa"]["actions"] = torch.zeros(32, 32, 69)
- with pytest.raises(ValueError, match="per-source batch"):
- callback.on_train_batch_start(None, None, batch, 0)
- finally:
- scope.restore()
- assert native.checkpoint is checkpoint
-
-
-def test_collection_image_chunk_layout_keeps_pixels_actions_and_temporal_reads_exact(
- tmp_path,
-):
- from egomimic.benchmarks.libero.convert import append_replay
-
- rng = np.random.default_rng(7)
- values = {
- "action": rng.normal(size=(19, 69)).astype(np.float32),
- "agentview_rgb": rng.integers(0, 256, (19, 3, 4, 3), dtype=np.uint8),
- "robot0_eye_in_hand_rgb": rng.integers(0, 256, (19, 3, 4, 3), dtype=np.uint8),
- }
- outputs = []
- for chunk in (1, 8):
- group = zarr.open_group(str(tmp_path / str(chunk)), mode="w", zarr_format=2)
- offset = 0
- for start, end in ((0, 3), (3, 11), (11, 19)):
- offset = append_replay(
- group,
- {k: v[start:end] for k, v in values.items()},
- offset,
- image_chunk_frames=chunk,
- )
- assert offset == 19 and group["agentview_rgb"].chunks[0] == chunk
- for key, expected in values.items():
- np.testing.assert_array_equal(group[key][:], expected)
- np.testing.assert_array_equal(group[key][7:9], expected[7:9])
- outputs.append(group)
- for key in values:
- np.testing.assert_array_equal(outputs[0][key][:], outputs[1][key][:])
- with pytest.raises(ValueError, match="positive frame count"):
- append_replay(outputs[0], values, 19, image_chunk_frames=0)
-
-
-def test_hand_target_cache_reuses_verified_targets_and_rejects_source_or_cache_drift(
- tmp_path,
-):
- import hashlib
-
- from egomimic.benchmarks.libero.convert import (
- cached_hand_keypoint_actions,
- verified_hand_target_cache,
- )
-
- source = tmp_path / "episode.hdf5"
- with h5py.File(source, "w") as h:
- h.create_dataset("actions", data=np.zeros((3, 28)))
- records = [{"path": source.name, "task": "task"}]
- manifest = tmp_path / "manifest.json"
- manifest.write_text(json.dumps({"episodes": records}))
- cache = tmp_path / "targets.npz"
- targets = np.arange(2 * 69, dtype=np.float32).reshape(2, 69)
- np.savez_compressed(
- cache,
- episode_000000_actions=targets,
- episode_000000_sequence=np.array([10, 11], dtype=np.int64),
- )
- stat = source.stat()
- row = {
- "index": 0,
- **records[0],
- "source_bytes": stat.st_size,
- "source_mtime_ns": stat.st_mtime_ns,
- "source_start_step": 1,
- "source_control_count": 3,
- "target_shape": [2, 69],
- "cache_key": "episode_000000",
- }
- audit = {
- "status": "PASS_ALL_500_EXACT_CAUSAL_HAND_TARGETS",
- "source_episode_count": 1,
- "passed_episode_count": 1,
- "failed_episode_count": 0,
- "source_manifest_sha256": hashlib.sha256(manifest.read_bytes()).hexdigest(),
- "target_cache": str(cache),
- "target_cache_sha256": hashlib.sha256(cache.read_bytes()).hexdigest(),
- "passed": [row],
- }
- audit_path = tmp_path / "audit.json"
- audit_path.write_text(json.dumps(audit))
- _, rows = verified_hand_target_cache(audit_path, manifest, tmp_path, records)
- with np.load(cache, allow_pickle=False) as arrays, h5py.File(source, "r") as h:
- actual, start, sequence = cached_hand_keypoint_actions(arrays, rows[0], h)
- np.testing.assert_array_equal(actual, targets)
- assert start == 1 and sequence.tolist() == [10, 11]
- with source.open("ab") as stream:
- stream.write(b"changed")
- with pytest.raises(ValueError, match="source identity changed"):
- verified_hand_target_cache(audit_path, manifest, tmp_path, records)
- with cache.open("ab") as stream:
- stream.write(b"changed")
- with pytest.raises(ValueError, match="cache SHA differs"):
- verified_hand_target_cache(audit_path, manifest, tmp_path, records)
-
-
-def test_joint_energy_verifier_consumes_real_generic_saved_artifact(
- tmp_path, monkeypatch
-):
- import copy
- from pathlib import Path
- from types import SimpleNamespace
-
- import hydra
- from hydra.core.hydra_config import HydraConfig
- from omegaconf import OmegaConf
-
- from egomimic.eval.energy_score import energy_score
- from egomimic.eval.planar_action_eval import PlanarActionEval
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_energy_artifact,
- )
-
- root = Path(__file__).resolve().parents[1]
- monkeypatch.setenv("LIBERO_ORIGINAL_SPATIAL_ROOT", str(tmp_path / "original"))
- monkeypatch.setenv("LIBERO_SHARPA_SPATIAL_ROOT", str(tmp_path / "sharpa"))
- monkeypatch.setenv(
- "LIBERO_DIAGNOSTIC_SEED_BANK",
- str(
- root / "egomimic/hydra_configs/evaluator/energy_score_seed_bank_k32_v1.json"
- ),
- )
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(root / "egomimic/hydra_configs")
- ):
- cfg = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=[
- "hydra/launcher=basic",
- "+experiment=sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- ],
- )
- cfg.hydra.runtime.output_dir = str(tmp_path)
- HydraConfig.instance().set_config(cfg)
- cfg.evaluator.seed_bank_path = str(
- root / "egomimic/hydra_configs/evaluator/energy_score_seed_bank_k32_v1.json"
- )
- cfg.evaluator.artifact_root = str(tmp_path / "energy")
- cfg.run_provenance.source_commit = "1" * 40
- cfg.run_provenance.normalization_sha256 = "2" * 64
- cfg.run_provenance.split_manifest_sha256 = "3" * 64
- resolved = OmegaConf.to_container(
- OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]), resolve=True
- )
- options = copy.deepcopy(resolved["evaluator"])
- options.pop("_target_")
- options["action_flow_diagnostics"] = None
- evaluator = PlanarActionEval(**options)
- evaluator.trainer = SimpleNamespace(
- global_step=8, current_epoch=0, global_rank=0, precision="bf16-mixed"
- )
- samples, scores, batch = {}, {}, {}
- generator = torch.Generator().manual_seed(42)
- for name, identity, width in (
- ("libero_panda", 21, 7),
- ("libero_panda_sharpa", 22, 69),
- ):
- samples[name] = torch.randn(32, 16, 32, width, generator=generator)
- target = torch.randn(16, 32, width, generator=generator)
- batch[name] = dict(embodiment=torch.full((16,), identity), actions=target)
- scores[name] = energy_score(
- samples[name],
- target,
- resolved["evaluator"]["semantic_blocks_by_embodiment"][name],
- )
- evaluator._save_artifact(0, samples, scores, batch)
- path = next((tmp_path / "energy").rglob("*.pt"))
- from tests.test_verify_action_flow_training_smoke import MODULE
-
- selected, payload = MODULE._step_artifact(
- tmp_path / "energy", label="joint EnergyScore", expected_step=8
- )
- assert selected == path
- with pytest.raises(MODULE.SmokeVerificationError, match="one step-2"):
- MODULE._step_two_artifact(tmp_path / "energy", label="joint EnergyScore")
- result = validate_joint_energy_artifact(payload, resolved)
- assert set(result["sources"]) == {"libero_panda", "libero_panda_sharpa"}
- payload["domains"]["libero_panda_sharpa"]["score_by_condition"][0] += 1
- with pytest.raises(ValueError, match="public math differs"):
- validate_joint_energy_artifact(payload, resolved)
-
- # Exercise the real immutable diagnostic writer with synthetic captured tensors.
- from egomimic.eval.action_flow_diagnostics import ActionFlowDiagnostics
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_diagnostic_artifact,
- )
-
- diagnostic_seeds = tmp_path / "diagnostic-seeds.json"
- diagnostic_seeds.write_text(
- json.dumps({"seeds": [730001 + 7919 * index for index in range(32)]}, indent=2)
- + "\n"
- )
- resolved["evaluator"]["action_flow_diagnostics"]["noise_seed_bank_path"] = str(
- diagnostic_seeds
- )
- inspector = ActionFlowDiagnostics(resolved["evaluator"]["action_flow_diagnostics"])
- diagnostic = dict(
- schema_version=1,
- metric="ActionFlowValidationDiagnostics",
- global_step=8,
- epoch=0,
- rank=0,
- batch_idx=0,
- precision="bf16-mixed",
- identity=inspector.identity,
- identity_sha256=inspector.identity_sha256,
- provenance=inspector.provenance,
- validation_view=inspector.validation_view,
- noise_seed_bank_sha256=inspector.noise_seed_bank_sha256,
- noise_seed=inspector.noise_seeds[0],
- raw_noise_levels=inspector.noise_levels,
- sampler=dict(name="reverse_euler"),
- sources={},
- )
- for name, width in (("libero_panda", 7), ("libero_panda_sharpa", 69)):
- raw = dict(
- source="validation/" + name,
- schema=inspector.schema,
- noise_seed=inspector.noise_seeds[0],
- num_inference_steps=50,
- returned_samples=8,
- )
- for key, shape in (
- ("target", (8, 32, width)),
- ("condition", (8, 276)),
- ("latent/clean", (8, 16, 16)),
- ("latent/noise", (8, 16, 16)),
- ("latent/generated", (8, 16, 16)),
- ("decoded/reconstruction", (8, 32, width)),
- ("decoded/generated", (8, 32, width)),
- ("latent/trajectory", (51, 8, 16, 16)),
- ("decoded/trajectory", (51, 8, 32, width)),
- ):
- raw[key] = torch.zeros(shape)
- diagnostic["sources"][name] = dict(
- source_id=raw["source"],
- diagnostic=raw,
- computed=dict(clean_reconstruction_mse_by_condition=torch.zeros(8)),
- )
- diagnostic_path = tmp_path / "diagnostics.pt"
- inspector._write_immutable_artifact(diagnostic_path, diagnostic)
- assert validate_joint_diagnostic_artifact(diagnostic_path, diagnostic, resolved)[
- "sha256"
- ]
- sidecar = Path(str(diagnostic_path) + ".sha256")
- sidecar.write_text(sidecar.read_text().replace(inspector.identity_sha256, "0" * 64))
- with pytest.raises(ValueError, match="sidecar differs"):
- validate_joint_diagnostic_artifact(diagnostic_path, diagnostic, resolved)
-
- # Metadata-only producer fixture: real saved tensor files, synthetic runtime/preflight.
- (tmp_path / "checkpoints").mkdir()
- (tmp_path / "checkpoints/epoch-0000-step-000000008.ckpt").write_bytes(b"fixture")
- from scripts.train.libero_native_launch_contract import (
- contract_sha,
- joint_smoke_evidence,
- reference,
- )
-
- resolved["run_provenance"]["run_kind"] = "smoke"
- resolved["trainer"]["max_steps"] = 8
- artifact_root = tmp_path / "diag-envelope"
- resolved["evaluator"]["action_flow_diagnostics"]["artifact_root"] = str(
- artifact_root
- )
- nested = artifact_root / "epoch-0-step-8/rank-0-batch-0.pt"
- nested.parent.mkdir(parents=True)
- inspector._write_immutable_artifact(nested, diagnostic)
- (tmp_path / ".hydra").mkdir()
- OmegaConf.save(OmegaConf.create(resolved), tmp_path / ".hydra/config.yaml")
- runtime = {
- key: {}
- for key in (
- "python",
- "installed_distributions",
- "core_distribution_versions",
- "dependency_inputs",
- "source",
- "torch_runtime",
- )
- }
- runtime_path = tmp_path / "runtime-lock.json"
- runtime_path.write_text(json.dumps(runtime))
- synthetic_preflight = dict(
- schema="libero-joint-launch-preflight/v1",
- status="PASS",
- profile="sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- model_contract_sha256=contract_sha(resolved),
- runtime_lock=reference(runtime_path),
- **{
- key: resolved["run_provenance"][key]
- for key in (
- "source_commit",
- "split_manifest_sha256",
- "normalization_sha256",
- )
- },
- )
- preflight_path = tmp_path / "synthetic-preflight.json"
- preflight_path.write_text(json.dumps(synthetic_preflight))
- envelope = joint_smoke_evidence(
- resolved,
- tmp_path,
- preflight_ref=reference(preflight_path),
- runtime_ref=reference(runtime_path),
- shared=MODULE,
- )
- assert envelope["energy_artifact"] == reference(path)
- assert envelope["diagnostic_artifact"] == reference(nested)
- assert (
- "status" not in envelope
- ) # An envelope never claims smoke or launch acceptance.
- resolved["run_provenance"]["run_kind"] = "full"
- with pytest.raises(ValueError, match="eight-update smoke"):
- joint_smoke_evidence(
- resolved,
- tmp_path,
- preflight_ref=reference(preflight_path),
- runtime_ref=reference(runtime_path),
- shared=MODULE,
- )
- resolved["run_provenance"]["run_kind"] = "smoke"
- runtime["source"] = {"wrong": True}
- different_runtime = tmp_path / "wrong-runtime.json"
- different_runtime.write_text(json.dumps(runtime))
- with pytest.raises(AssertionError, match="native runtime differs"):
- joint_smoke_evidence(
- resolved,
- tmp_path,
- preflight_ref=reference(preflight_path),
- runtime_ref=reference(different_runtime),
- shared=MODULE,
- )
-
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_wandb_records,
- )
- from tests.test_verify_action_flow_training_smoke import _history_row
-
- row = _history_row()
- for key, value in list(row.items()):
- if MODULE.SOURCE_LABEL in key:
- for label in ("libero_panda", "libero_panda_sharpa"):
- row[key.replace(MODULE.SOURCE_LABEL, label)] = value
- row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 2.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 28.0
- for prefix in (
- "GradientCosine",
- "GradientCosineDefined",
- "GradientIntersectionParameterCount",
- ):
- row[f"Train/ActionFlow/{prefix}/FM__Reconstruction"] = 0.0
- row["Optimizer/LR/AdamW"] = row["Optimizer/LR/Muon"] = 2.5e-8
- for owner in ("encoder/libero_panda", "encoder/libero_panda_sharpa", "velocity"):
- for mode in ("direct", "checkpoint"):
- row["Train/Execution/DiTHalf/" + owner + "/" + mode] = 6.0
- resolved["logger"] = dict(wandb=dict(mode="online", id="synthetic-history-fixture"))
- histories = dict(
- offline={7: row},
- online={7: row},
- online_config=resolved,
- online_state="finished",
- cfg=resolved,
- shared=MODULE,
- )
- assert validate_joint_wandb_records(**histories)["online"]["valid_step"] == 7
- import copy
-
- logged = copy.deepcopy(resolved)
- logged["model"] = json.loads(json.dumps(logged["model"]))
- logged["model"]["action_horizon"] = "${libero_blueprint.action_horizon}"
- histories["online_config"] = logged
- assert validate_joint_wandb_records(**histories)["online"]["valid_step"] == 7
- logged["model"]["action_horizon"] = 999
- with pytest.raises(ValueError, match="online W&B model differs"):
- validate_joint_wandb_records(**histories)
- histories["online_config"] = resolved
- histories["online_state"] = "running"
- with pytest.raises(ValueError, match="finished online"):
- validate_joint_wandb_records(**histories)
- histories["online_state"] = "finished"
- row["Train/Execution/DiTHalf/encoder/libero_panda_sharpa/checkpoint"] = 0.0
- with pytest.raises(ValueError, match="path inactive"):
- validate_joint_wandb_records(**histories)
-
-
-def test_joint_constructor_explicit_waiver_keeps_cached_evidence_immutable(tmp_path):
- import hashlib
-
- from scripts.train.libero_native_launch_contract import validate_joint_constructor
-
- original = tmp_path / "original.json"
- original.write_text(
- json.dumps(
- dict(
- status="PASS_ACTUAL_CPU_MODEL_CONSTRUCTION_NOT_TRAINING_SMOKE",
- total_parameters=78182614,
- private_codecs_per_embodiment=True,
- single_shared_velocity_stage=True,
- all_parameters_finite=True,
- )
- )
- )
- before = original.read_bytes()
- waiver = tmp_path / "waiver.json"
- data = dict(
- schema="libero-joint-constructor-provenance-waiver/v1",
- authorization="explicit_human_override",
- waived_gate="cached_constructor_source_config_binding",
- source_commit="a" * 40,
- human_message="bypass the provenance gate",
- original_constructor=dict(
- path=str(original), sha256=hashlib.sha256(before).hexdigest()
- ),
- )
-
- def check():
- waiver.write_text(json.dumps(data))
- ref = dict(
- path=str(waiver), sha256=hashlib.sha256(waiver.read_bytes()).hexdigest()
- )
- return validate_joint_constructor(
- ref, source_head="a" * 40, resolved_sha="b" * 64, fingerprint="c" * 64
- )
-
- assert check() == "HUMAN_WAIVED_SOURCE_CONFIG_BINDING"
- assert original.read_bytes() == before
- data["source_commit"] = "d" * 40
- with pytest.raises(ValueError, match="authorization/source"):
- check()
- data["source_commit"] = "a" * 40
- data["authorization"] = ""
- with pytest.raises(ValueError, match="authorization/source"):
- check()
diff --git a/tests/test_libero_native_binding_parser.py b/tests/test_libero_native_binding_parser.py
deleted file mode 100644
index 96973f9df6..0000000000
--- a/tests/test_libero_native_binding_parser.py
+++ /dev/null
@@ -1,75 +0,0 @@
-"""Real Hydra grammar regression; no model construction or smoke claims."""
-
-import importlib.util
-import unittest
-from pathlib import Path
-
-spec = importlib.util.spec_from_file_location(
- "native_schema",
- Path(__file__).parents[1] / "scripts/train/libero_native_schema_contract.py",
-)
-module = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(module)
-KEYS = (
- "path",
- "file_sha256",
- "normalizer_module_sha256",
- "source_commit",
- "dataset_receipt_path",
- "dataset_receipt_sha256",
- "split_receipt_path",
- "split_receipt_sha256",
- "data_root",
- "normalizer_target",
- "physical_proof_path",
- "physical_proof_sha256",
-)
-
-
-class BindingParserTests(unittest.TestCase):
- def test_real_twelve_field_hydra_mapping(self):
- expected = {
- k: (
- "/workspace/" + k
- if k.endswith("path") or k == "data_root"
- else "abc123"
- )
- for k in KEYS
- }
- arg = (
- "++norm_stats.native_saved_state_binding={"
- + ",".join(k + ":" + v for k, v in expected.items())
- + "}"
- )
- self.assertEqual(module.parse_native_saved_state_binding([arg]), expected)
-
- def test_malformed_and_incomplete(self):
- for value in ("{path:/workspace/x}", "{path:/workspace/x", "null", "[x,y]"):
- with self.subTest(value=value), self.assertRaises(Exception):
- module.parse_native_saved_state_binding(
- ["++norm_stats.native_saved_state_binding=" + value]
- )
-
- def test_sweep_and_delete(self):
- arg = (
- "++norm_stats.native_saved_state_binding={"
- + ",".join(k + ":abc" for k in KEYS)
- + "}"
- )
- for anomaly in ("trainer.max_steps=2,3", "~trainer.max_steps"):
- with self.subTest(anomaly=anomaly), self.assertRaises(ValueError):
- module.parse_native_saved_state_binding([arg, anomaly])
-
- def test_missing_and_duplicate(self):
- arg = (
- "++norm_stats.native_saved_state_binding={"
- + ",".join(k + ":abc" for k in KEYS)
- + "}"
- )
- for argv in ([], [arg, arg], [arg, arg.replace("++", "", 1)]):
- with self.subTest(argv=argv), self.assertRaises(ValueError):
- module.parse_native_saved_state_binding(argv)
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_libero_native_half_factory_identity.py b/tests/test_libero_native_half_factory_identity.py
deleted file mode 100644
index f63c645f6e..0000000000
--- a/tests/test_libero_native_half_factory_identity.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""Historical half-ResNet factory identity; no torch import for static guard."""
-
-import ast
-import hashlib
-from pathlib import Path
-
-
-def test_native_half_factory_matches_historical_architecture_and_initialization():
- path = Path(__file__).parents[1] / "egomimic/models/stems/visual_core.py"
- source = path.read_text()
- tree = ast.parse(source)
- branches = [
- n
- for n in ast.walk(tree)
- if isinstance(n, ast.If)
- and ast.unparse(n.test) == "resnet_model == 'resnet18_half'"
- ]
- assert len(branches) == 1, "native historical half-ResNet factory unavailable"
- assert (
- # Pinned from training source7c253e5, including the non-half else arm.
- # Ignore whitespace only; retain every call, initializer and argument.
- hashlib.sha256(
- ast.dump(branches[0], include_attributes=False).encode()
- ).hexdigest()
- == "9521e4594c50d2e98b25fc9c07ef2d0c9aac42c83b57e0d0f55c8128056d8a75"
- )
diff --git a/tests/test_libero_native_launch_contract.py b/tests/test_libero_native_launch_contract.py
deleted file mode 100644
index d595828a2d..0000000000
--- a/tests/test_libero_native_launch_contract.py
+++ /dev/null
@@ -1,102 +0,0 @@
-"""Focused fingerprint invariants only; not native constructor or smoke proof."""
-
-import copy
-import unittest
-
-from scripts.train.libero_native_launch_contract import contract_sha
-
-
-class Contract(unittest.TestCase):
- def setUp(self):
- self.cfg = dict(
- model={
- "pipeline": {
- "stages": [
- {
- "flow_aggregation": "sum_samples",
- "inference_method": "euler",
- "num_inference_steps": 50,
- }
- ]
- },
- "gradient_telemetry_cadence": 100,
- },
- normalizer={"_target_": "native"},
- norm_stats={"native_saved_state_binding": {"source_commit": "fixed"}},
- data={
- "train_datasets": {"libero_panda": {"split_seed": 42}},
- "valid_datasets": {"libero_panda": {"split_seed": 42}},
- "train_dataloader_params": {"libero_panda": {"batch_size": 32}},
- "valid_dataloader_params": {"libero_panda": {"batch_size": 32}},
- },
- trainer={
- "precision": "bf16-mixed",
- "accumulate_grad_batches": 1,
- "max_steps": 80000,
- },
- seed=42,
- callbacks={
- "ema": {"validate_with_ema": True},
- "dit_half": {"_target_": "native"},
- },
- evaluator={
- "energy_sample_count": 32,
- "energy_seed_bank_sha256": "fixed",
- "artifact_root": "/out/a",
- "artifact_identity": {
- "_target_": "factory",
- "config_path": "/out/a/.hydra/config.yaml",
- },
- "native_diagnostic_config": {
- "_target_": "diagfactory",
- "config_path": "/out/a/.hydra/config.yaml",
- "artifact_root": "/out/a/diagnostic",
- },
- },
- )
-
- def test_phase_only_equivalent(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["trainer"]["max_steps"] = 2
- cfg["model"]["gradient_telemetry_cadence"] = 2
- self.assertEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_output_only_equivalent(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["evaluator"]["artifact_root"] = "/out/b"
- cfg["evaluator"]["artifact_identity"]["config_path"] = (
- "/out/b/.hydra/config.yaml"
- )
- cfg["evaluator"]["native_diagnostic_config"]["artifact_root"] = (
- "/out/b/diagnostic"
- )
- self.assertEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_mean_rejected(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["model"]["pipeline"]["stages"][0]["flow_aggregation"] = "mean"
- self.assertNotEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_seedbank_rejected(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["evaluator"]["energy_seed_bank_sha256"] = "changed"
- self.assertNotEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_k_changed_rejected(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["evaluator"]["energy_sample_count"] = 16
- self.assertNotEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_factory_target_rejected(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["evaluator"]["native_diagnostic_config"]["_target_"] = "otherfactory"
- self.assertNotEqual(contract_sha(cfg), contract_sha(self.cfg))
-
- def test_sampling_rejected(self):
- cfg = copy.deepcopy(self.cfg)
- cfg["model"]["pipeline"]["stages"][0]["num_inference_steps"] = 8
- self.assertNotEqual(contract_sha(cfg), contract_sha(self.cfg))
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_libero_native_scheduler_contract.py b/tests/test_libero_native_scheduler_contract.py
deleted file mode 100644
index bd35a9c9a6..0000000000
--- a/tests/test_libero_native_scheduler_contract.py
+++ /dev/null
@@ -1,101 +0,0 @@
-import importlib.util
-from pathlib import Path
-
-import pytest
-
-spec = importlib.util.spec_from_file_location(
- "guard", Path(__file__).parents[1] / "scripts/train/validate_slurm_job_contract.py"
-)
-m = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(m)
-PROFILE = "sim/libero_historical/action_flow_libero10_h240_euler50_dithalf_80k_s42"
-
-
-def inputs():
- fields = m.parse_scontrol_record(
- (
- Path(__file__).parent / "fixtures/libero_native_scontrol_408787.txt"
- ).read_text()
- )
- args = dict(
- expected_job_id="408787",
- expected_account="loaner",
- expected_partition="rl2",
- expected_qos="loaner-normal",
- expected_cpus=8,
- expected_memory="128G",
- expected_time_limit="02:00:00",
- expected_constraint="(null)",
- native_profile=PROFILE,
- gpu_probe=dict(
- status="PASSED",
- gpu_name="NVIDIA H100 80GB HBM3",
- world_size=1,
- rank=0,
- local_rank=0,
- bf16_supported=True,
- bf16_forward_backward=dict(finite=True),
- ),
- )
- return fields, args
-
-
-def test_native_actual_h100_no_constraint():
- f, a = inputs()
- assert not m.evaluate_contract(f, **a)[2]
-
-
-@pytest.mark.parametrize(
- "change",
- [
- "wrongprofile",
- "missingprobe",
- "wronggpu",
- "nonfinite",
- "wrongallocation",
- "constraint",
- "planar",
- ],
-)
-def test_reject_contract_drift(change):
- f, a = inputs()
- if change == "wrongprofile":
- a["native_profile"] = "planar"
- if change == "missingprobe":
- a["gpu_probe"] = None
- if change == "wronggpu":
- a["gpu_probe"]["gpu_name"] = "NVIDIA H200"
- if change == "nonfinite":
- a["gpu_probe"]["bf16_forward_backward"]["finite"] = False
- if change == "wrongallocation":
- f["GRES"] = "gpu:a40:1(IDX:3)"
- if change == "constraint":
- f["Features"] = "H100"
- if change == "planar":
- a["native_profile"] = None
- try:
- failures = m.evaluate_contract(f, **a)[2]
- except m.ContractError:
- return
- assert failures
-
-
-@pytest.mark.parametrize(
- "field,value",
- [
- ("JOB_GRES", "gpu:h100:2"),
- ("TresPerNode", "gres/gpu:h200:1"),
- ("Nodes", "gpu18"),
- ("GRES", "gpu:h100:1"),
- ("AllocTRES", "cpu=4,mem=128G,node=1"),
- ("ReqTRES", "cpu=8,mem=128G,node=1,gres/gpu=1"),
- ],
-)
-def test_actual_allocation_drift(field, value):
- f, a = inputs()
- f[field] = value
- try:
- failures = m.evaluate_contract(f, **a)[2]
- except m.ContractError:
- return
- assert failures
diff --git a/tests/test_libero_native_smoke_contract.py b/tests/test_libero_native_smoke_contract.py
deleted file mode 100644
index 5b35b5da89..0000000000
--- a/tests/test_libero_native_smoke_contract.py
+++ /dev/null
@@ -1,54 +0,0 @@
-"""Focused pure metadata tests; never a full-model smoke proof."""
-
-import unittest
-
-from scripts.train.verify_libero_native_action_flow_smoke import (
- PARAMETER_COUNT,
- PROFILE,
- validate_dit_counts,
- validate_preflight,
-)
-
-
-class Contract(unittest.TestCase):
- def test_missing_preflight(self):
- with self.assertRaises(KeyError):
- validate_preflight({}, {})
-
- def test_tiny_count_not_accepted(self):
- self.assertEqual(PARAMETER_COUNT, 39750391)
- self.assertEqual(
- PROFILE,
- "sim/libero_historical/action_flow_libero10_h240_euler50_dithalf_80k_s42",
- )
-
- def test_preparation_rejected(self):
- with self.assertRaisesRegex(ValueError, "did not pass"):
- validate_preflight(
- {"schema": "libero-native-launch-preflight/v1", "status": "PREPARED"},
- {},
- )
-
- def test_wrong_profile_rejected(self):
- with self.assertRaisesRegex(ValueError, "profile"):
- validate_preflight(
- {
- "schema": "libero-native-launch-preflight/v1",
- "status": "PASS",
- "profile": "invented",
- },
- {},
- )
-
- def test_dit_inactive_rejected(self):
- class Shared:
- @staticmethod
- def _complete_row(rows, required, **kw):
- return 1, {name: 0 for name in required}
-
- with self.assertRaisesRegex(ValueError, "inactive"):
- validate_dit_counts({}, Shared)
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_libero_native_suite_profiles.py b/tests/test_libero_native_suite_profiles.py
deleted file mode 100644
index d65bcbbf8b..0000000000
--- a/tests/test_libero_native_suite_profiles.py
+++ /dev/null
@@ -1,857 +0,0 @@
-"""Suite/seed isolation and portable allocated-resource contracts, without models."""
-
-import copy
-import importlib.util
-import sys
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).resolve().parents[1]
-sys.path.insert(0, str(ROOT))
-sys.path.insert(0, str(ROOT / "tools"))
-from libero_maintained_dispatch_v1 import ( # noqa: E402 - task tool path initialized above
- validate_data_receipt,
-)
-
-
-def test_joint_schedule_reuse_hashes_exact_git_diff_and_rejects_model_changes(tmp_path):
- import hashlib
- import json
- import subprocess
-
- from scripts.train.libero_native_launch_contract import (
- joint_smoke_schedule_reuse,
- reference,
- )
-
- repo = tmp_path / "source"
- repo.mkdir()
-
- def git(*args):
- return subprocess.check_output(
- [
- "git",
- "-C",
- str(repo),
- "-c",
- "user.name=fixture",
- "-c",
- "user.email=fixture@example.invalid",
- *args,
- ]
- )
-
- git("init", "-q")
- source = repo / "egomimic/eval/eval.py"
- source.parent.mkdir(parents=True)
- source.write_text("old schedule\n")
- git("add", ".")
- git("commit", "-qm", "old")
- old = git("rev-parse", "HEAD").decode().strip()
- source.write_text("repaired schedule\n")
- git("commit", "-qam", "repair")
- head = git("rev-parse", "HEAD").decode().strip()
- preflight_path = tmp_path / "preflight.json"
- preflight_path.write_text(json.dumps({"source_commit": old}))
- smoke_path = tmp_path / "smoke.json"
- smoke_path.write_text(
- json.dumps(
- {
- "identities": {"source_commit": old},
- "preflight": reference(preflight_path),
- }
- )
- )
- regression = tmp_path / "regression.log"
- regression.write_text("actual CPU evidence fixture\n")
- proof = dict(
- schema="libero-joint-smoke-schedule-reuse/v1",
- scope="explicit_cross_epoch_validation_schedule",
- from_source_commit=old,
- to_source_commit=head,
- target_repo=str(repo),
- smoke=reference(smoke_path),
- cpu_regression_status="PASS",
- changed_paths=["egomimic/eval/eval.py"],
- audited_diff_sha256=hashlib.sha256(
- git("diff", "--no-ext-diff", "--binary", old, head)
- ).hexdigest(),
- cpu_regression=reference(regression),
- )
- proof_path = tmp_path / "proof.json"
-
- def check():
- proof_path.write_text(json.dumps(proof))
- return joint_smoke_schedule_reuse(
- reference(proof_path),
- smoke_ref=reference(smoke_path),
- target_source=proof["to_source_commit"],
- )
-
- assert check()[0] == old
- proof["audited_diff_sha256"] = "0" * 64
- with pytest.raises(ValueError, match="audited diff differs"):
- check()
- model = repo / "egomimic/model.py"
- model.write_text("changed model\n")
- git("add", ".")
- git("commit", "-qm", "forbidden")
- proof["to_source_commit"] = git("rev-parse", "HEAD").decode().strip()
- proof["changed_paths"] = (
- git("diff", "--name-only", old, "HEAD").decode().splitlines()
- )
- proof["audited_diff_sha256"] = hashlib.sha256(
- git("diff", "--no-ext-diff", "--binary", old, "HEAD")
- ).hexdigest()
- with pytest.raises(ValueError, match="unaudited source paths"):
- check()
-
-
-from egomimic.benchmarks.libero.native_launch_profiles import ( # noqa: E402 - task tool path initialized above
- PROFILES,
- profile_for_argv,
- profile_for_config,
-)
-
-
-def cfg(suite):
- p = PROFILES[suite]
- return dict(
- name=p.name,
- benchmark=dict(suite=suite),
- data={
- mode + "_datasets": {"libero_panda": {"resolver": {"suite": suite}}}
- for mode in ("train", "valid")
- },
- )
-
-
-@pytest.mark.parametrize("suite", list(PROFILES))
-def test_typed_suite_and_receipt_isolation(suite):
- p = PROFILES[suite]
- c = cfg(suite)
- assert profile_for_config(c) == p
- assert profile_for_argv(["+experiment=" + p.experiment]) == p
- receipt = dict(
- suite=suite,
- replay_path="/replay",
- logical_dataset_sha256="a" * 64,
- episodes=500,
- valid_ratio=0.01,
- split_seed=42,
- train_episodes=495,
- valid_episodes=5,
- action_dim=7,
- task_uids=list(p.task_uids),
- )
- # Use exact maintained receipt schema field names.
- receipt.update(
- dataset_logical_sha256="a" * 64,
- frames=5000,
- action_shape=[5000, 7],
- train_episode_indices=list(range(495)),
- valid_episode_indices=list(range(495, 500)),
- )
- assert (
- validate_data_receipt(
- receipt,
- replay_path="/replay",
- expected_logical_sha="a" * 64,
- profile=p.experiment,
- )["id_overlap"]
- == 0
- )
- for other in PROFILES:
- if other == suite:
- continue
- wrong = copy.deepcopy(c)
- wrong["benchmark"]["suite"] = other
- with pytest.raises(ValueError):
- profile_for_config(wrong)
- with pytest.raises((ValueError, KeyError)):
- validate_data_receipt(
- receipt,
- replay_path="/replay",
- expected_logical_sha="a" * 64,
- profile=PROFILES[other].experiment,
- )
-
-
-@pytest.mark.parametrize("suite", ["libero_spatial", "libero_goal", "libero_object"])
-def test_recipe_only_changes_suite_identity(suite):
- text = (
- ROOT
- / "egomimic/hydra_configs/experiment/sim/libero_historical"
- / (PROFILES[suite].name + ".yaml")
- ).read_text()
- assert "action_flow_libero10_h240_euler50_dithalf_80k_s42" in text
- assert "suite: " + suite in text
- assert PROFILES[suite].replay_environment in text
- assert all(
- key not in text
- for key in ("model:", "trainer:", "normalizer:", "optimizer:", "\nseed:")
- )
-
-
-def test_catalog_task_uids_match_registry():
- from egomimic.benchmarks.libero.catalog import TASK_IDS, get_tasks
-
- for suite, p in PROFILES.items():
- assert tuple(TASK_IDS[t][2] for t in get_tasks(suite)) == p.task_uids
-
-
-def scheduler():
- spec = importlib.util.spec_from_file_location(
- "scheduler_suite", ROOT / "scripts/train/validate_slurm_job_contract.py"
- )
- m = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(m)
- return m
-
-
-@pytest.mark.parametrize(
- "constraint,gpu",
- [
- ("H100", "NVIDIA H100 80GB HBM3"),
- ("H200", "NVIDIA H200"),
- ("H100|H200", "NVIDIA H200"),
- ],
-)
-@pytest.mark.parametrize(
- "profile",
- [
- PROFILES["libero_spatial"].experiment,
- "sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- ],
-)
-def test_ice_native_preserves_real_resource_bindings(constraint, gpu, profile):
- m = scheduler()
- fields = dict(
- JobId="123",
- Account="account",
- Partition="partition",
- QOS="qos",
- NumNodes="1",
- NumTasks="1",
- NumCPUs="8",
- MinCPUsNode="8",
- MinMemoryNode="128G",
- Features=constraint,
- TimeLimit="1-00:00:00",
- ReqTRES="cpu=8,mem=128G,node=1,gres/gpu=1",
- )
- fields["CPUs/Task"] = "8"
- args = dict(
- expected_job_id="123",
- expected_account="account",
- expected_partition="partition",
- expected_qos="qos",
- expected_cpus=8,
- expected_memory="128G",
- expected_time_limit="1-00:00:00",
- expected_constraint=constraint,
- native_profile=profile,
- gpu_probe=dict(
- status="PASSED",
- gpu_name=gpu,
- world_size=1,
- rank=0,
- local_rank=0,
- bf16_supported=True,
- bf16_forward_backward=dict(finite=True),
- ),
- )
- assert m.evaluate_contract(fields, **args)[2] == []
- for field in ("Account", "Partition", "QOS", "Features", "NumCPUs"):
- bad = dict(fields)
- bad[field] = "wrong" if field != "NumCPUs" else "4"
- assert m.evaluate_contract(bad, **args)[2]
- args["gpu_probe"]["bf16_forward_backward"]["finite"] = False
- with pytest.raises(m.ContractError):
- m.evaluate_contract(fields, **args)
-
-
-@pytest.mark.parametrize("suite", list(PROFILES))
-@pytest.mark.parametrize("seed", [42, 43])
-def test_metric_training_seed_is_typed_and_not_split_seed(suite, seed):
- from egomimic.benchmarks.libero.action_flow_artifacts import (
- SEED_SHA,
- validate_identity,
- )
-
- identity = dict(
- suite=suite,
- source="libero_panda",
- action_dim=7,
- action_horizon=16,
- seed=seed,
- sample_count=32,
- energy_seed_bank_sha256=SEED_SHA,
- inference_method="euler",
- inference_steps=50,
- normalization_scope="training_episodes_only",
- effective_batch_size=32,
- homogeneous="not_applicable_single_source",
- source_commit="a" * 40,
- resolved_config_sha256="a" * 64,
- split_sha256="a" * 64,
- normalizer_state_sha256="a" * 64,
- dataset_logical_sha256="a" * 64,
- )
- if seed == 43 and suite not in ("libero_goal", "libero_object"):
- with pytest.raises(ValueError):
- validate_identity(identity)
- else:
- assert validate_identity(identity)
-
-
-def test_late_preflight_cross_suite_rejected_before_evidence_reads():
- from scripts.train.verify_libero_native_action_flow_smoke import validate_preflight
-
- pf = dict(
- schema="libero-native-launch-preflight/v1",
- status="PASS",
- profile=PROFILES["libero10"].experiment,
- )
- for suite in ("libero_spatial", "libero_goal", "libero_object"):
- with pytest.raises(ValueError, match="native profile mismatch"):
- validate_preflight(pf, {}, PROFILES[suite].experiment)
-
-
-@pytest.mark.parametrize(
- "suite", ["libero_spatial", "libero_goal", "libero_object", "libero10"]
-)
-def test_actual_hydra_recipe_composes_same_science_and_fixed_split(monkeypatch, suite):
- from hydra import compose, initialize_config_dir
- from omegaconf import OmegaConf
-
- p = PROFILES[suite]
- monkeypatch.setenv(p.replay_environment, "/replay")
- for seed in [42, 43] if suite in ("libero_goal", "libero_object") else [42]:
- with initialize_config_dir(
- version_base="1.3", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- c = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "hydra/launcher=basic",
- "+experiment=" + p.experiment,
- "seed=" + str(seed),
- ],
- )
- assert c.name == p.name and c.benchmark.suite == suite and c.seed == seed
- assert c.data.train_datasets.libero_panda.split_seed == 42
- assert c.data.valid_datasets.libero_panda.split_seed == 42
- assert (
- c.model.hidden_dim == 240
- and c.model.action_horizon == 16
- and c.model.action_dim == 7
- )
- assert c.data.train_dataloader_params.libero_panda.batch_size == 32
- assert c.model.pipeline.stages[6].inference_method == "euler"
- monkeypatch.setenv("LIBERO10_REPLAY_ROOT", "/reference")
- with initialize_config_dir(
- version_base="1.3", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- reference = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "hydra/launcher=basic",
- "+experiment=" + PROFILES["libero10"].experiment,
- ],
- )
- for field in ("model", "normalizer", "trainer", "callbacks", "evaluator"):
- assert OmegaConf.to_container(
- c[field], resolve=False
- ) == OmegaConf.to_container(reference[field], resolve=False)
-
-
-def test_explicit_joint_profile_and_contract_reject_single_source_or_semantic_drift(
- tmp_path, monkeypatch
-):
- import hydra
- from hydra.core.hydra_config import HydraConfig
- from omegaconf import OmegaConf
-
- from egomimic.benchmarks.libero.native_launch_profiles import (
- joint_profile_for_config,
- )
-
- spec = importlib.util.spec_from_file_location(
- "joint_contract_test", ROOT / "scripts/train/libero_native_launch_contract.py"
- )
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- monkeypatch.setenv("LIBERO_ORIGINAL_SPATIAL_ROOT", str(tmp_path / "original"))
- monkeypatch.setenv("LIBERO_SHARPA_SPATIAL_ROOT", str(tmp_path / "sharpa"))
- monkeypatch.setenv("LIBERO_DIAGNOSTIC_SEED_BANK", str(tmp_path / "SEEDS.json"))
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- config = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=[
- "hydra/launcher=basic",
- "+experiment=sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- "paths.output_dir=" + str(tmp_path),
- "paths.work_dir=" + str(ROOT),
- ],
- )
- config.hydra.runtime.output_dir = str(tmp_path)
- HydraConfig.instance().set_config(config)
- c = OmegaConf.to_container(
- OmegaConf.masked_copy(config, [k for k in config if k != "hydra"]),
- resolve=True,
- )
- assert joint_profile_for_config(c).sources[1] == ("libero_panda_sharpa", 22, 69, 22)
- missing_latent = copy.deepcopy(c)
- missing_latent["evaluator"]["action_flow_diagnostics"]["provenance"].pop(
- "latent_shape"
- )
- with pytest.raises(ValueError, match="diagnostic latent shape"):
- joint_profile_for_config(missing_latent)
- with pytest.raises(ValueError, match="unsupported native"):
- profile_for_config(c)
- contract = module.model_contract(c)
- assert contract["schema"] == "libero-joint-scientific-contract/v1"
- assert contract["train_batch"] == {"libero_panda": 16, "libero_panda_sharpa": 16}
- assert module.contract_sha(c) == module.contract_sha(copy.deepcopy(c))
- relocated = copy.deepcopy(c)
- relocated["callbacks"]["ema"]["final_checkpoint_path"] = (
- "/other/output/checkpoints/last.ckpt"
- )
- relocated["evaluator"]["action_flow_diagnostics"]["artifact_root"] = (
- "/other/output/diagnostics"
- )
- assert module.contract_sha(c) == module.contract_sha(relocated)
- changed_ema = copy.deepcopy(c)
- changed_ema["callbacks"]["ema"]["decay"] = 0.99
- assert module.contract_sha(c) != module.contract_sha(changed_ema)
- changes = [
- (
- "model.pipeline.stages.4.encoders.libero_panda_sharpa.backbone.hidden_dim",
- 384,
- ),
- ("model.pipeline.stages.6.field.backbone.depth", 24),
- ("model.pipeline.stages.4.encoders.libero_panda_sharpa.action_dim", 7),
- ("model.pipeline.stages.7.decoders.libero_panda_sharpa.condition_dim", 276),
- ("model.pipeline.stages.6.experts", {}),
- ("model.pipeline.stages.3.num_tokens", 8),
- ("data.train_dataloader_params.libero_panda_sharpa.batch_size", 32),
- ("data.train_datasets.libero_panda_sharpa.valid_ratio", 0.01),
- (
- "data.train_datasets.libero_panda_sharpa.resolver.action_representation",
- "native",
- ),
- ("model.pipeline.stages.6.num_inference_steps", 8),
- ("normalizer._target_", "egomimic.rldb.zarr.libero_dataset.LiberoNormalizer"),
- ]
- for path, replacement in changes:
- wrong = copy.deepcopy(c)
- keys = path.split(".")
- node = wrong
- for key in keys[:-1]:
- node = node[int(key)] if isinstance(node, list) else node[key]
- node[keys[-1]] = replacement
- with pytest.raises(ValueError, match="joint LIBERO"):
- module.contract_sha(wrong)
-
-
-@pytest.mark.parametrize("phase", ["normalization", "preflight", "smoke", "full"])
-def test_joint_phase_factory_composes_actual_paths_and_private_shared_contract(
- tmp_path, monkeypatch, phase
-):
- import hydra
- from hydra.core.hydra_config import HydraConfig
- from omegaconf import OmegaConf
-
- from egomimic.benchmarks.libero.native_launch_profiles import (
- joint_phase_argv,
- joint_profile_for_config,
- )
-
- # These env values must be replaced by the exact typed arguments, including
- # spaces and '=' which must survive the shell/Hydra boundary literally.
- for key in (
- "LIBERO_ORIGINAL_SPATIAL_ROOT",
- "LIBERO_SHARPA_SPATIAL_ROOT",
- "LIBERO_DIAGNOSTIC_SEED_BANK",
- ):
- monkeypatch.setenv(key, "/wrong/unselected")
- original = str(tmp_path / "original data=1")
- sharpa = str(tmp_path / "hand data=2")
- args = dict(
- source_root=str(ROOT),
- original_root=original,
- sharpa_root=sharpa,
- output_root=str(tmp_path / phase),
- context_path=str(tmp_path / "context.json"),
- split_path=str(tmp_path / "split.json"),
- split_sha256="a" * 64,
- normalization_sha256=None if phase == "normalization" else "b" * 64,
- source_commit="c" * 40,
- diagnostic_seed_bank=str(tmp_path / "seeds.json"),
- wandb_entity="test",
- wandb_project="test",
- wandb_run_id="joint-" + phase,
- )
- argv = joint_phase_argv(phase=phase, **args)
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- cfg = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=argv[1:],
- )
- cfg.hydra.runtime.output_dir = str(tmp_path / phase)
- HydraConfig.instance().set_config(cfg)
- resolved = OmegaConf.to_container(
- OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]), resolve=True
- )
- joint_profile_for_config(resolved)
- assert resolved["evaluator"]["action_flow_diagnostics"]["provenance"][
- "latent_shape"
- ] == [16, 16]
- assert resolved["trainer"]["log_every_n_steps"] == (1 if phase == "smoke" else 10)
- if phase != "normalization":
- from scripts.train import libero_native_launch_contract as adapter
-
- def compose_exact(command, *, cwd, text):
- assert command[1:3] == ["-m", "egomimic.trainHydra"]
- assert command[3:-3] == argv
- assert command[-3:] == ["--cfg", "job", "--resolve"]
- assert cwd == ROOT and text is True
- return OmegaConf.to_yaml(OmegaConf.create(resolved))
-
- monkeypatch.setattr(adapter.subprocess, "check_output", compose_exact)
- destination = tmp_path / (phase + "-composition")
- destination.mkdir()
- composed = adapter.compose_joint_phase(
- args, phase=phase, output=destination, repo=ROOT
- )
- assert composed["launch_ready"] is False
- assert (destination / "exact-phase.argv0").read_bytes() == b"\0".join(
- arg.encode() for arg in argv
- ) + b"\0"
- with pytest.raises(FileExistsError):
- adapter.compose_joint_phase(
- args, phase=phase, output=destination, repo=ROOT
- )
- if phase in {"smoke", "full"}:
- import json
-
- from scripts.train.libero_native_launch_contract import (
- contract_sha,
- joint_prepare_evidence,
- reference,
- )
-
- runtime = {
- key: {}
- for key in (
- "python",
- "installed_distributions",
- "core_distribution_versions",
- "dependency_inputs",
- "source",
- "torch_runtime",
- )
- }
- runtime["dependency_inputs"] = []
- runtime_path = tmp_path / "synthetic-runtime.json"
- runtime_path.write_text(json.dumps(runtime))
- identity = {
- key: resolved["run_provenance"][key]
- for key in (
- "source_commit",
- "normalization_sha256",
- "split_manifest_sha256",
- )
- }
- preflight = dict(
- schema="libero-joint-launch-preflight/v1",
- status="PASS",
- profile="sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- model_contract_sha256=contract_sha(resolved),
- runtime_lock=reference(runtime_path),
- **identity,
- )
- preflight_path = tmp_path / "synthetic-preflight.json"
- preflight_path.write_text(json.dumps(preflight))
- kwargs = dict(
- phase=phase,
- preflight_ref=reference(preflight_path),
- runtime_ref=reference(runtime_path),
- )
- if phase == "full":
- with pytest.raises(ValueError, match="genuine eight-update"):
- joint_prepare_evidence(resolved, **kwargs)
- smoke = dict(
- schema="libero-joint-action-flow-smoke/v1",
- status="PASS",
- profile=preflight["profile"],
- model_contract_sha256=contract_sha(resolved),
- identities=identity,
- preflight=reference(preflight_path),
- strict_online_and_alias_safe_ema_reload=True,
- checkpoint=dict(
- parameter_count=78182614,
- global_step=8,
- joint_ema=dict(alias_safe_strict_reload=True),
- ),
- )
- smoke_path = tmp_path / "synthetic-smoke.json"
- smoke_path.write_text(json.dumps(smoke))
- kwargs["smoke_ref"] = reference(smoke_path)
- assert joint_prepare_evidence(resolved, **kwargs)["launch_ready"] is False
- if phase == "full":
- from scripts.train import libero_native_launch_contract as adapter
-
- original_smoke = copy.deepcopy(smoke)
- previous = dict(preflight, source_commit="d" * 40)
- previous["model_contract_sha256"] = adapter.joint_contract_at_source(
- resolved, "d" * 40
- )
- previous_runtime = dict(runtime, source={"head": "d" * 40, "dirty": False})
- old_runtime_path = tmp_path / "previous-runtime.json"
- old_runtime_path.write_text(json.dumps(previous_runtime))
- previous["runtime_lock"] = reference(old_runtime_path)
- runtime["source"] = {"head": identity["source_commit"], "dirty": False}
- runtime_path.write_text(json.dumps(runtime))
- preflight["runtime_lock"] = reference(runtime_path)
- preflight_path.write_text(json.dumps(preflight))
- kwargs["runtime_ref"] = reference(runtime_path)
- kwargs["preflight_ref"] = reference(preflight_path)
- previous_path = tmp_path / "previous-preflight.json"
- previous_path.write_text(json.dumps(previous))
- smoke["identities"] = dict(identity, source_commit="d" * 40)
- smoke["model_contract_sha256"] = previous["model_contract_sha256"]
- smoke["preflight"] = reference(previous_path)
- smoke_path.write_text(json.dumps(smoke))
- kwargs["smoke_ref"] = reference(smoke_path)
- with pytest.raises(ValueError, match="smoke proof differs"):
- joint_prepare_evidence(resolved, **kwargs)
- repair_ref = {"path": "separate-repair-proof", "sha256": "e" * 64}
- monkeypatch.setattr(
- adapter,
- "joint_smoke_schedule_reuse",
- lambda *a, **k: ("d" * 40, previous),
- )
- reused = joint_prepare_evidence(
- resolved, **kwargs, smoke_source_repair_ref=repair_ref
- )
- assert reused["smoke_source_repair"] == repair_ref
- assert (
- json.loads(smoke_path.read_text())["identities"]["source_commit"]
- == "d" * 40
- )
- # Source-local assets move with immutable checkout roots. Compare
- # their actual bytes against the configured SHA before rebinding
- # this one path, keeping every other scientific field exact.
- asset = Path(resolved["evaluator"]["seed_bank_path"])
- payload = asset.read_bytes()
- old_repo, new_repo = tmp_path / "old-source", tmp_path / "new-source"
- old_repo.mkdir()
- new_repo.mkdir()
- old_seed, new_seed = old_repo / "seeds.json", new_repo / "seeds.json"
- old_seed.write_bytes(payload)
- new_seed.write_bytes(payload)
- before = copy.deepcopy(resolved)
- before["evaluator"]["seed_bank_path"] = str(old_seed)
- expected = adapter.joint_contract_at_source(before, "d" * 40)
- after = copy.deepcopy(resolved)
- after["evaluator"]["seed_bank_path"] = str(new_seed)
- rebound = dict(source_repo=old_repo, current_repo=new_repo)
- assert (
- adapter.joint_contract_at_source(after, "d" * 40, **rebound) == expected
- )
- assert after["evaluator"]["seed_bank_path"] == str(new_seed)
- for changed in (old_seed, new_seed):
- changed.write_bytes(b"changed seed bank")
- with pytest.raises(ValueError, match="seed asset SHA differs"):
- adapter.joint_contract_at_source(after, "d" * 40, **rebound)
- changed.write_bytes(payload)
- after["model"]["cfg_scale"] += 1
- assert (
- adapter.joint_contract_at_source(after, "d" * 40, **rebound) != expected
- )
- previous["model_contract_sha256"] = "wrong"
- with pytest.raises(ValueError, match="scientific contract differs"):
- joint_prepare_evidence(
- resolved, **kwargs, smoke_source_repair_ref=repair_ref
- )
- smoke = original_smoke
- smoke["checkpoint"]["global_step"] = 2
- smoke_path.write_text(json.dumps(smoke))
- kwargs["smoke_ref"] = reference(smoke_path)
- with pytest.raises(ValueError, match="smoke proof differs"):
- joint_prepare_evidence(resolved, **kwargs)
- if phase == "preflight":
- from scripts.train.libero_native_launch_contract import contract_sha
-
- for other_phase in ("smoke", "full"):
- other_args = {
- **args,
- "output_root": str(tmp_path / other_phase),
- "wandb_run_id": "joint-" + other_phase,
- }
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- other = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=joint_phase_argv(phase=other_phase, **other_args)[1:],
- )
- other.hydra.runtime.output_dir = other_args["output_root"]
- HydraConfig.instance().set_config(other)
- other_resolved = OmegaConf.to_container(
- OmegaConf.masked_copy(other, [k for k in other if k != "hydra"]),
- resolve=True,
- )
- assert contract_sha(resolved) == contract_sha(other_resolved)
- if phase == "smoke":
- from scripts.train.libero_native_launch_contract import joint_prepare_evidence
-
- assert resolved["model"]["gradient_telemetry_cadence"] == 1
- missing_telemetry = copy.deepcopy(resolved)
- missing_telemetry["model"]["gradient_telemetry_cadence"] = 100
- with pytest.raises(ValueError, match="per-step gradient telemetry"):
- joint_prepare_evidence(
- missing_telemetry, phase="smoke", preflight_ref={}, runtime_ref={}
- )
-
- from scripts.train.verify_libero_native_action_flow_smoke import (
- validate_joint_smoke_config,
- )
-
- assert (
- validate_joint_smoke_config(resolved, args["output_root"], "c" * 40)[
- "optimizer_updates"
- ]
- == 8
- )
- for key, value in (
- ("trainer.max_steps", 2),
- ("trainer.accumulate_grad_batches", 2),
- ("run_provenance.run_kind", "full"),
- ("logger.wandb.mode", "offline"),
- ("callbacks.model_checkpoint.save_last", True),
- ("run_provenance.source_commit", None),
- ):
- altered = OmegaConf.create(copy.deepcopy(resolved))
- OmegaConf.update(altered, key, value)
- with pytest.raises(
- ValueError, match="joint (smoke config mismatch|LIBERO:)"
- ):
- validate_joint_smoke_config(altered, args["output_root"], "c" * 40)
- for mode in ("train", "valid"):
- for source, root in (
- ("libero_panda", original),
- ("libero_panda_sharpa", sharpa),
- ):
- assert (
- resolved["data"][mode + "_datasets"][source]["resolver"]["folder_path"]
- == root
- )
- assert resolved["callbacks"]["model_checkpoint"]["save_last"] is False
- assert resolved["trainer"]["default_root_dir"] == args["output_root"]
- assert (
- resolved["callbacks"]["model_checkpoint"]["dirpath"]
- == args["output_root"] + "/checkpoints"
- )
- assert (
- resolved["callbacks"]["ema"]["final_checkpoint_path"]
- == args["output_root"] + "/checkpoints/last.ckpt"
- )
- assert resolved["runtime"]["slurm_requeue_owner"] == "runner"
- assert resolved["callbacks"]["early_recovery"]["first_step"] == 1000
- assert resolved["trainer"]["max_steps"] == (8 if phase == "smoke" else 120000)
- if phase == "normalization":
- assert "logger" not in resolved
- # The experiment itself extends the evaluator body; group deletion
- # would leave an invalid target-less partial mapping. Normalization-only
- # trainHydra returns before model/trainer construction or evaluation.
- assert resolved["evaluator"]["_target_"].endswith(".PlanarActionEval")
- assert resolved["norm_stats"]["precomputed_norm_path"] is None
- assert resolved["norm_stats_only"] is True
- assert resolved["norm_stats"]["export_context_path"] == args["context_path"]
- assert resolved["run_provenance"]["normalization_sha256"] is None
- else:
- assert resolved["norm_stats"]["precomputed_norm_path"] == args["context_path"]
- assert resolved["logger"]["wandb"]["id"] == args["wandb_run_id"]
- assert resolved["logger"]["wandb"]["save_dir"] == args["output_root"]
- assert resolved["trainer"]["val_check_interval"] == (
- 8 if phase == "smoke" else 15000
- )
- with pytest.raises(ValueError, match="normalization_sha256"):
- joint_phase_argv(phase=phase, **{**args, "normalization_sha256": None})
- for key, bad_value in (
- ("original_root", "relative/root"),
- ("sharpa_root", "/path\nwith-newline"),
- ("source_commit", "unknown"),
- ("wandb_run_id", "unsafe id"),
- ):
- with pytest.raises(ValueError):
- joint_phase_argv(phase=phase, **{**args, key: bad_value})
-
-
-def test_joint_phase_factory_data_only_snapshot_never_authorizes_training(tmp_path):
- import hydra
- from hydra.core.hydra_config import HydraConfig
- from omegaconf import OmegaConf
-
- from egomimic.benchmarks.libero.native_launch_profiles import (
- joint_phase_argv,
- joint_profile_for_config,
- )
-
- args = dict(
- source_root=str(ROOT),
- original_root=str(tmp_path / "original"),
- sharpa_root=str(tmp_path / "sharpa"),
- output_root=str(tmp_path / "out"),
- context_path=str(tmp_path / "context.json"),
- split_path=str(tmp_path / "split.json"),
- split_sha256="a" * 64,
- normalization_sha256=None,
- source_commit=None,
- diagnostic_seed_bank=str(tmp_path / "bank.json"),
- wandb_entity="test",
- wandb_project="test",
- wandb_run_id="data-only",
- data_preparation_source_sha256="d" * 64,
- )
- argv = joint_phase_argv(phase="normalization", **args)
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- cfg = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=argv[1:],
- )
- cfg.hydra.runtime.output_dir = str(tmp_path / "out")
- HydraConfig.instance().set_config(cfg)
- resolved = OmegaConf.to_container(
- OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]), resolve=True
- )
- joint_profile_for_config(resolved)
- assert (
- resolved["norm_stats_only"] is True
- and resolved["run_provenance"]["source_commit"] is None
- )
- assert (
- resolved["run_provenance"]["data_preparation_source_snapshot_sha256"]
- == "d" * 64
- )
- for phase in ("preflight", "smoke", "full"):
- with pytest.raises(ValueError, match="data-only source snapshots"):
- joint_phase_argv(phase=phase, **args)
- with pytest.raises(ValueError, match="exact SHA"):
- joint_phase_argv(
- phase="normalization", **{**args, "data_preparation_source_sha256": None}
- )
diff --git a/tests/test_libero_native_velocity_ablation.py b/tests/test_libero_native_velocity_ablation.py
deleted file mode 100644
index 79879d24b0..0000000000
--- a/tests/test_libero_native_velocity_ablation.py
+++ /dev/null
@@ -1,276 +0,0 @@
-"""Explicit objective identity and single-variable resolved recipe regression."""
-
-import copy
-import sys
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-from hydra.core.hydra_config import HydraConfig
-from omegaconf import OmegaConf
-
-ROOT = Path(__file__).resolve().parents[1]
-sys.path[:0] = [str(ROOT), str(ROOT / "tools")]
-from typed_libero_profile_v3 import ( # noqa: E402 - task tool path initialized above
- COMMON,
- PHASES,
- validate_flat_config,
-)
-
-from egomimic.benchmarks.libero.native_launch_profiles import ( # noqa: E402 - task tool path initialized above
- AV0_PROFILES,
- PROFILES,
- profile_for_argv,
- profile_for_config,
-)
-
-
-def recipe(profile, monkeypatch):
- monkeypatch.setenv(profile.replay_environment, "/unchanged/replay")
- with initialize_config_dir(
- version_base="1.3", config_dir=str(ROOT / "egomimic/hydra_configs")
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=["hydra/launcher=basic", "+experiment=" + profile.experiment],
- return_hydra_config=True,
- )
- cfg.hydra.runtime.output_dir = "/unchanged/output"
- HydraConfig.instance().set_config(cfg)
- OmegaConf.set_struct(cfg, False)
- del cfg["hydra"]
- return OmegaConf.to_container(cfg, resolve=True)
-
-
-@pytest.mark.parametrize("suite", list(PROFILES))
-def test_only_resolved_model_velocity_weight_changes(suite, monkeypatch):
- base = recipe(PROFILES[suite], monkeypatch)
- ablated = recipe(AV0_PROFILES[suite], monkeypatch)
- assert profile_for_config(base) == PROFILES[suite]
- assert profile_for_config(ablated) == AV0_PROFILES[suite]
- assert (
- profile_for_argv(["+experiment=" + AV0_PROFILES[suite].experiment])
- == AV0_PROFILES[suite]
- )
- expected = copy.deepcopy(base["model"])
- expected["pipeline"]["stages"][8]["action_velocity_weight"] = 0.0
- assert ablated["model"] == expected
- expected_data = copy.deepcopy(base["data"])
- expected_data["run_provenance"]["objective"]["action_velocity_weight"] = 0.0
- expected_data["run_provenance"]["ablation"] = "action_velocity_off"
- assert ablated["data"] == expected_data
- for key in ("trainer", "normalizer", "callbacks", "evaluator", "seed"):
- assert ablated[key] == base[key], key
- assert ablated["run_provenance"]["objective"] == {
- **base["run_provenance"]["objective"],
- "action_velocity_weight": 0.0,
- }
- bad = copy.deepcopy(ablated)
- bad["run_provenance"]["objective"]["action_velocity_weight"] = 1.0
- with pytest.raises(ValueError, match="objective provenance"):
- profile_for_config(bad)
- bad = copy.deepcopy(base)
- bad["run_provenance"]["ablation"] = "action_velocity_off"
- with pytest.raises(ValueError, match="ablation provenance"):
- profile_for_config(bad)
-
-
-@pytest.mark.parametrize("weight", [0.0, 1.0])
-def test_typed_flat_guard_rejects_opposite_objective(weight):
- flat = {
- **COMMON,
- **PHASES["smoke"],
- "model.pipeline.stages.8.action_velocity_weight": weight,
- }
- for idx, owner in [(4, "encoder"), (6, "field")]:
- for key, value in [("hidden_dim", 240), ("depth", 12), ("num_heads", 8)]:
- flat[f"model.pipeline.stages.{idx}.{owner}.backbone.{key}"] = value
- flat.update(
- {
- "data.train_source_names": ["libero_panda"],
- "data.valid_source_names": ["libero_panda"],
- "typed_profile.velocity_augmentation": False,
- "typed_profile.homogeneous": "not_applicable_single_source",
- "typed_profile.checkpoint_policy": "dit-half",
- }
- )
- assert validate_flat_config(flat, "smoke", action_velocity_weight=weight)
- with pytest.raises(ValueError, match="resolved profile mismatch"):
- validate_flat_config(flat, "smoke", action_velocity_weight=1.0 - weight)
-
-
-def verifier_functions():
- # Execute exact pure verifier functions without importing GPU/model dependencies.
- import ast
- import math
- from collections.abc import Mapping, Sequence
- from typing import Any
-
- source = ast.parse(
- (ROOT / "scripts/train/verify_action_flow_training_smoke.py").read_text()
- )
- names = {
- "SmokeVerificationError",
- "_require",
- "_metric",
- "_complete_row",
- "_validate_history",
- "_validate_checkpoint_loss_schedule",
- }
- constants = {
- "LEGACY_METHOD",
- "GRAPH_METHOD",
- "LIKELIHOOD_METHOD",
- "STOPGRAD_METHOD",
- "STOPGRAD_UNITE_METHOD",
- "SOURCE_LABEL",
- }
- nodes = [
- ast.ImportFrom(
- module="__future__", names=[ast.alias(name="annotations")], level=0
- )
- ]
- namespace = dict(
- math=math, Any=Any, Mapping=Mapping, Sequence=Sequence, OmegaConf=OmegaConf
- )
- for node in ast.parse(
- (ROOT / "tools/validate_action_flow_config.py").read_text()
- ).body:
- if isinstance(node, ast.Assign):
- for target in node.targets:
- if isinstance(target, ast.Name) and target.id in constants:
- namespace[target.id] = ast.literal_eval(node.value)
- for node in source.body:
- if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names:
- nodes.append(node)
- if isinstance(node, ast.Assign):
- for target in node.targets:
- if isinstance(target, ast.Name) and target.id in constants:
- namespace[target.id] = ast.literal_eval(node.value)
- exec(
- compile(
- ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])),
- str(ROOT / "scripts/train/verify_action_flow_training_smoke.py"),
- "exec",
- ),
- namespace,
- )
- fixture = ast.parse(
- (ROOT / "tests/test_verify_action_flow_training_smoke.py").read_text()
- )
- import types
-
- namespace["MODULE"] = types.SimpleNamespace(**namespace)
- row_node = next(
- node
- for node in fixture.body
- if isinstance(node, ast.FunctionDef) and node.name == "_history_row"
- )
- exec(
- compile(
- ast.fix_missing_locations(ast.Module(body=[row_node], type_ignores=[])),
- "existing_history_fixture",
- "exec",
- ),
- namespace,
- )
- return namespace
-
-
-def test_disabled_velocity_must_be_zero_weight_in_schedule_and_total():
- v = verifier_functions()
- row = v["_history_row"]()
- row["Train/ActionFlow/Schedule/EffectiveActionVelocityWeight"] = 0.0
- for suffix in ("", "/" + v["SOURCE_LABEL"]):
- row[f"Train/ActionFlow/TotalLoss{suffix}_step"] = 2.5
- row["Valid/ActionFlow/TotalLoss"] = 3.0
- assert (
- v["_validate_history"]({2: row}, action_velocity_weight=0.0)["train_step"] == 2
- )
- bad = dict(row)
- bad["Train/ActionFlow/TotalLoss_step"] = 3.75
- with pytest.raises(v["SmokeVerificationError"], match="weighted total"):
- v["_validate_history"]({2: bad}, action_velocity_weight=0.0)
- bad = dict(row)
- bad["Train/ActionFlow/Schedule/EffectiveActionVelocityWeight"] = 1.0
- with pytest.raises(v["SmokeVerificationError"], match="joint smoke step"):
- v["_validate_history"]({2: bad}, action_velocity_weight=0.0)
- stages = [{} for _ in range(9)]
- stages[8] = {"action_velocity_weight": 0.0}
- cfg = OmegaConf.create(
- {
- "name": AV0_PROFILES["libero10"].name,
- "benchmark": {"suite": "libero10"},
- "data": {
- mode + "_datasets": {
- "libero_panda": {"resolver": {"suite": "libero10"}}
- }
- for mode in ("train", "valid")
- },
- "run_provenance": {
- "ablation": "action_velocity_off",
- "objective": {"action_velocity_weight": 0.0},
- },
- "model": {
- "reconstruction_only_warmup_steps": 0,
- "pipeline": {"stages": stages},
- },
- }
- )
- schedule = {
- "joint_objective_begins_at_global_step": 0,
- "reconstruction_only_optimizer_steps": 0,
- "joint_flow_weight": 1.0,
- "joint_reconstruction_weight": 1.0,
- "joint_action_velocity_weight": 0.0,
- "schema_version": 1,
- }
- assert (
- v["_validate_checkpoint_loss_schedule"](
- schedule, cfg, reconstruction_weight=1.0, flow_weight=1.0
- )
- == 0
- )
- with pytest.raises(v["SmokeVerificationError"], match="loss schedule"):
- v["_validate_checkpoint_loss_schedule"](
- {**schedule, "joint_action_velocity_weight": 1.0},
- cfg,
- reconstruction_weight=1.0,
- flow_weight=1.0,
- )
-
- # A nonnative config cannot weaken the historical checkpoint guard by setting stage7=0.
- nonnative = OmegaConf.create(
- {
- "name": "unrelated",
- "model": {
- "reconstruction_only_warmup_steps": 0,
- "pipeline": {"stages": [{}] * 7 + [{"action_velocity_weight": 0.0}]},
- },
- }
- )
- with pytest.raises(
- v["_validate_checkpoint_loss_schedule"].__globals__["SmokeVerificationError"],
- match="loss schedule",
- ):
- v["_validate_checkpoint_loss_schedule"](
- schedule, nonnative, reconstruction_weight=1.0, flow_weight=1.0
- )
- assert (
- v["_validate_checkpoint_loss_schedule"](
- {**schedule, "joint_action_velocity_weight": 1.0},
- nonnative,
- reconstruction_weight=1.0,
- flow_weight=1.0,
- )
- == 0
- )
- baseline = OmegaConf.create(OmegaConf.to_container(cfg, resolve=True))
- baseline.name = PROFILES["libero10"].name
- baseline.run_provenance.objective.action_velocity_weight = 1.0
- baseline.run_provenance.ablation = None
- baseline.model.pipeline.stages[8].action_velocity_weight = 1.0
- with pytest.raises(v["SmokeVerificationError"], match="loss schedule"):
- v["_validate_checkpoint_loss_schedule"](
- schedule, baseline, reconstruction_weight=1.0, flow_weight=1.0
- )
diff --git a/tests/test_libero_native_whole_launcher.py b/tests/test_libero_native_whole_launcher.py
deleted file mode 100644
index 7968d6454c..0000000000
--- a/tests/test_libero_native_whole_launcher.py
+++ /dev/null
@@ -1,417 +0,0 @@
-"""Execute the complete real Bash launcher through safe native tool handoff.
-
-No model, data, Python tool, GPU, scheduler, credential or training execution.
-Every preceding guard/hash/path/clean-Git check and actual array remains real.
-"""
-
-import hashlib
-import json
-import os
-import shutil
-import subprocess
-import tempfile
-import unittest
-from pathlib import Path
-
-R = Path(__file__).parents[1]
-PROFILE = "sim/libero_historical/action_flow_libero10_h240_euler50_dithalf_80k_s42"
-PLANAR = "sim/pusht/action_flow_bc_usocket_recon1_s42"
-
-
-def sha(p):
- return hashlib.sha256(Path(p).read_bytes()).hexdigest()
-
-
-class WholeLauncher(unittest.TestCase):
- def setUp(self):
- self.temp = tempfile.TemporaryDirectory()
- self.root = Path(self.temp.name)
- self.repo = self.root / "source"
- self.repo.mkdir()
- files = [
- "scripts/train/launch_action_flow_usocket.sbatch",
- "scripts/train/libero_native_launch_phases.sh",
- "scripts/train/libero_native_schema_contract.py",
- "scripts/train/verify_action_flow_training_smoke.py",
- "scripts/train/validate_slurm_job_contract.py",
- "scripts/train/check_checkpoint_storage.py",
- "scripts/ice/ice_gpu_probe.py",
- "scripts/ice/ice_requeue_runner.py",
- "scripts/ice/relay_batch_signals.sh",
- "scripts/ice/validate_lightning_checkpoint.py",
- "scripts/ice/capture_runtime_lock.py",
- "tools/libero_maintained_dispatch_v1.py",
- "tools/validate_action_flow_config.py",
- "scripts/ice/validate_planar_dataset.py",
- "egomimic/hydra_configs/experiment/" + PROFILE + ".yaml",
- "egomimic/hydra_configs/experiment/" + PLANAR + ".yaml",
- "egomimic/hydra_configs/experiment/sim/libero_historical/action_flow_libero_spatial_h240_euler50_dithalf_80k_s42.yaml",
- "egomimic/hydra_configs/experiment/sim/libero_historical/action_flow_libero_goal_h240_euler50_dithalf_80k_s42.yaml",
- "egomimic/hydra_configs/experiment/sim/libero_historical/action_flow_libero_object_h240_euler50_dithalf_80k_s42.yaml",
- ]
- from egomimic.benchmarks.libero.native_launch_profiles import AV0_PROFILES
-
- files += [
- "egomimic/hydra_configs/experiment/" + p.experiment + ".yaml"
- for p in AV0_PROFILES.values()
- ]
- for rel in files:
- p = self.repo / rel
- p.parent.mkdir(parents=True, exist_ok=True)
- shutil.copy2(R / rel, p)
- self.launcher = self.repo / files[0]
- for name in ["data", "output-parent"]:
- (self.root / name).mkdir()
- self.split = self.root / "split.json"
- self.split.write_text("{}\n")
- self.receipt = self.root / "receipt.json"
- self.receipt.write_text("{}\n")
- self.norm = self.root / "norm.json"
- self.norm.write_text("{}\n")
- self.preflight = self.root / "preflight.json"
- self.preflight.write_text("{}\n")
- self.override = self.root / "binding.override"
- self.override.write_text(
- '++norm_stats.native_saved_state_binding={test:"fixture-only-no-execution"}\n'
- )
- self.mock = self.repo / "safe-tools/mock-python"
- self.mock.parent.mkdir()
- self.mock.write_text("""#!/bin/bash
-set -Eeuo pipefail
-case "$1" in
- */capture_runtime_lock.py)
- while test "$#" -gt 0; do
- if test "$1" = --output; then shift; printf '{"scope":"mocked-runtime-no-python-execution"}\\n' > "$1"; break; fi
- shift
- done ;;
- */libero_native_launch_contract.py)
- printf 'NATIVE_DISPATCH_PREPARE_STOP\\n'
- printf '%s\\n' "$*" > "$SAFE_HANDOFF_LOG"
- replay_variable=${SAFE_REPLAY_VARIABLE:-LIBERO_SPATIAL_REPLAY_ROOT}
- printf '%s\\n' "${!replay_variable:-}" > "$SAFE_HANDOFF_LOG.replay"
- exit 91 ;;
- *) printf 'UNEXPECTED_EXTERNAL_TOOL\\n' >&2; exit 92 ;;
-esac
-""")
- self.mock.chmod(0o755)
- cp = self.repo / "scripts/ice/validate_lightning_checkpoint.py"
- cp.chmod(0o755)
- subprocess.run(["git", "init", "-q", str(self.repo)], check=True)
- subprocess.run(["git", "-C", str(self.repo), "add", "."], check=True)
- subprocess.run(
- [
- "git",
- "-C",
- str(self.repo),
- "-c",
- "user.name=Fixture",
- "-c",
- "user.email=fixture@example.invalid",
- "commit",
- "-qm",
- "safe launcher fixture",
- ],
- check=True,
- )
- head = subprocess.check_output(
- ["git", "-C", str(self.repo), "rev-parse", "HEAD"], text=True
- ).strip()
- env = {
- k: v
- for k, v in os.environ.items()
- if not k.startswith(("AF_", "ICE_", "SLURM_", "WANDB_", "LIBERO_"))
- }
- env.update(
- AF_REPO=str(self.repo),
- AF_EXPECTED_HEAD=head,
- AF_PYTHON=str(self.mock),
- AF_DATASET_DIR=str(self.root / "data"),
- AF_OUTPUT_DIR=str(self.root / "output-parent/run"),
- AF_EXPERIMENT=PROFILE,
- AF_WANDB_ENTITY="rl2-group",
- AF_WANDB_PROJECT="pushshapes-action-flow",
- AF_WANDB_RUN_ID="lib10-noaug-smoke-ce2c-20261007",
- AF_WANDB_GROUP="cotrain-velocity-first-wave-20261007",
- AF_WANDB_TAGS="[libero10,noaug,action-flow,h240,dithalf,bf16,euler50,smoke]",
- AF_EXPECTED_SPLIT_MANIFEST_SHA256=sha(self.split),
- AF_EXPECTED_LAUNCHER_SHA256=sha(self.launcher),
- AF_CONTENT_MANIFEST=str(self.receipt),
- AF_EXPECTED_CONTENT_MANIFEST_SHA256=sha(self.receipt),
- AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256="3ed2b86cf97eaa27045af81e6d6c4834c8659d15eda69a2394d835e602014e1d",
- AF_NATIVE_DATASET_SHA256="3ed2b86cf97eaa27045af81e6d6c4834c8659d15eda69a2394d835e602014e1d",
- AF_NATIVE_HISTORICAL_SPLIT_SHA256="c588c831ef0deab61a59a21840ea9e11dafe92a170d9023a87e0acee831df236",
- AF_CLUSTER_LABEL="lambda",
- AF_EXPECTED_ACCOUNT="loaner",
- AF_EXPECTED_PARTITION="rl2",
- AF_EXPECTED_QOS="loaner-normal",
- AF_SPLIT_MANIFEST=str(self.split),
- AF_NORM_STATS_PATH=str(self.norm),
- AF_EXPECTED_NORM_SHA256=sha(self.norm),
- AF_NATIVE_SCHEMA_DRIVER=str(
- self.repo / "scripts/train/libero_native_schema_contract.py"
- ),
- AF_NATIVE_SCHEMA_DRIVER_SHA256=sha(
- self.repo / "scripts/train/libero_native_schema_contract.py"
- ),
- AF_NATIVE_BINDING_OVERRIDE=str(self.override),
- AF_NATIVE_BINDING_OVERRIDE_SHA256=sha(self.override),
- AF_NATIVE_DATASET_VALIDATOR=str(
- self.repo / "tools/libero_maintained_dispatch_v1.py"
- ),
- AF_NATIVE_CONFIG_VALIDATOR=str(
- self.repo / "tools/libero_maintained_dispatch_v1.py"
- ),
- AF_NATIVE_SMOKE_VERIFIER=str(
- self.repo / "scripts/train/verify_action_flow_training_smoke.py"
- ),
- AF_NATIVE_PHASE="smoke",
- AF_PRECISION="bf16-mixed",
- AF_LAUNCH_MODE="run",
- AF_RUN_KIND="smoke",
- AF_MAX_RESTARTS="0",
- AF_EXPECTED_GPU_CONSTRAINT="(null)",
- AF_EXPECTED_MEMORY="128G",
- AF_EXPECTED_TIME_LIMIT="02:00:00",
- AF_PREFLIGHT_RESULT=str(self.preflight),
- AF_EXPECTED_PREFLIGHT_SHA256=sha(self.preflight),
- AF_NATIVE_SCHEMA_ALLOW_PENDING_PROOFS="false",
- SLURM_STEP_ID="fixture-step",
- SLURM_JOB_ID="123456",
- SLURM_JOB_ACCOUNT="loaner",
- SLURM_JOB_PARTITION="rl2",
- SLURM_JOB_QOS="loaner-normal",
- SLURM_JOB_NUM_NODES="1",
- SLURM_NTASKS="1",
- SLURM_CPUS_PER_TASK="8",
- SLURM_RESTART_COUNT="0",
- SAFE_HANDOFF_LOG=str(self.root / "handoff.log"),
- )
- self.env = env
-
- def tearDown(self):
- self.temp.cleanup()
-
- def invoke(self, **changes):
- env = dict(self.env)
- env.update(changes)
- return subprocess.run(
- ["bash", str(self.launcher)],
- env=env,
- capture_output=True,
- text=True,
- timeout=15,
- )
-
- def assertRejected(self, result, text):
- self.assertEqual(result.returncode, 64, result.stdout + result.stderr)
- self.assertIn(text, result.stderr)
- self.assertFalse((self.root / "handoff.log").exists())
-
- def test_exact_v7_native_smoke_envelope_reaches_real_native_dispatch(self):
- result = self.invoke()
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
- self.assertIn("NATIVE_DISPATCH_PREPARE_STOP", result.stdout)
- self.assertIn("--action prepare", (self.root / "handoff.log").read_text())
- argv = (
- self.root / "output-parent/run/provenance/restart-0/exact-phase.argv0"
- ).read_bytes()
- self.assertIn(b"trainer.max_steps=2\0", argv)
- self.assertIn(b"trainer.precision=bf16-mixed\0", argv)
- self.assertIn(b"++run_provenance.preflight_result_sha256=", argv)
-
- def suite_values(self, suite, seed):
- record = json.loads(
- (R / "tests/fixtures/libero_native_suite_corpus_authority.json").read_text()
- )["suites"][suite]
- return dict(
- AF_EXPERIMENT="sim/libero_historical/action_flow_"
- + suite
- + "_h240_euler50_dithalf_80k_s42",
- AF_SEED=str(seed),
- AF_NATIVE_DATASET_SHA256=record["dataset_logical_sha256"],
- AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256=record[
- "dataset_logical_sha256"
- ],
- AF_NATIVE_HISTORICAL_SPLIT_SHA256=record["historical_split_sha256"],
- AF_EXPECTED_GPU_CONSTRAINT="H100|H200"
- if suite == "libero_spatial"
- else "(null)",
- AF_CLUSTER_LABEL="ice" if suite == "libero_spatial" else "lambda",
- )
-
- def test_actual_goal_object_corpus_and_both_training_seeds_reach_dispatch(self):
- for suite in ("libero_goal", "libero_object"):
- for seed in (42, 43):
- with self.subTest(suite=suite, seed=seed):
- values = self.suite_values(suite, seed)
- values["AF_OUTPUT_DIR"] = str(
- self.root / "output-parent" / (suite + "-" + str(seed))
- )
- result = self.invoke(**values)
- self.assertEqual(
- result.returncode, 91, result.stdout + result.stderr
- )
- argv = (
- Path(values["AF_OUTPUT_DIR"])
- / "provenance/restart-0/exact-phase.argv0"
- ).read_bytes()
- self.assertIn(("seed=" + str(seed) + "\0").encode(), argv)
- self.assertIn(
- (
- "++run_provenance.dataset_content_aggregate_sha256="
- + values["AF_NATIVE_DATASET_SHA256"]
- + "\0"
- ).encode(),
- argv,
- )
- self.assertIn(
- (
- "++run_provenance.historical_split_manifest_sha256="
- + values["AF_NATIVE_HISTORICAL_SPLIT_SHA256"]
- + "\0"
- ).encode(),
- argv,
- )
-
- def test_cross_suite_corpus_and_historical_split_rejected(self):
- for suite in ("libero_spatial", "libero_goal", "libero_object"):
- for field, message in [
- ("AF_NATIVE_DATASET_SHA256", "native dataset identity mismatch"),
- (
- "AF_NATIVE_HISTORICAL_SPLIT_SHA256",
- "historical split identity mismatch",
- ),
- (
- "AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256",
- "native dataset aggregate identity mismatch",
- ),
- ]:
- with self.subTest(suite=suite, field=field):
- values = self.suite_values(suite, 42)
- values[field] = self.env[field]
- self.assertRejected(self.invoke(**values), message)
-
- def test_spatial_ice_actual_whole_path_constrained_dispatch(self):
- result = self.invoke(**self.suite_values("libero_spatial", 42))
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
- argv = (
- self.root / "output-parent/run/provenance/restart-0/exact-phase.argv0"
- ).read_bytes()
- self.assertIn(
- b"+experiment=sim/libero_historical/action_flow_libero_spatial_h240_euler50_dithalf_80k_s42\0",
- argv,
- )
- self.assertEqual(
- (self.root / "handoff.log.replay").read_text().strip(),
- str(self.root / "data"),
- )
-
- def test_spatial_native_actual_whole_path_null_reaches_hardware_guard(self):
- result = self.invoke(
- **{
- **self.suite_values("libero_spatial", 42),
- "AF_EXPECTED_GPU_CONSTRAINT": "(null)",
- }
- )
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
-
- def test_native_full_positive_restart_cap_and_null_constraint_reach_dispatch(self):
- result = self.invoke(
- AF_RUN_KIND="full",
- AF_MAX_RESTARTS="4",
- AF_SMOKE_RESULT=str(self.preflight),
- AF_EXPECTED_SMOKE_SHA256=sha(self.preflight),
- )
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
- self.assertIn("--phase full", (self.root / "handoff.log").read_text())
- argv = (
- self.root / "output-parent/run/provenance/restart-0/exact-phase.argv0"
- ).read_bytes()
- self.assertIn(b"trainer.max_steps=80000\0", argv)
- self.assertIn(b"trainer.val_check_interval=15000\0", argv)
- self.assertIn(b"++callbacks.model_checkpoint.every_n_train_steps=5000\0", argv)
-
- def test_native_unsupported_gpu_constraint_rejected(self):
- self.assertRejected(
- self.invoke(AF_EXPECTED_GPU_CONSTRAINT="RTX4090"),
- "AF_EXPECTED_GPU_CONSTRAINT must be H100",
- )
-
- def test_native_invalid_restart_counter_rejected(self):
- self.assertRejected(
- self.invoke(SLURM_RESTART_COUNT="bad"), "invalid SLURM_RESTART_COUNT"
- )
-
- def test_native_smoke_requeue_cap_must_be_zero(self):
- self.assertRejected(
- self.invoke(AF_MAX_RESTARTS="1"), "native smoke requires AF_MAX_RESTARTS=0"
- )
-
- def test_negative_restart_cap_rejected(self):
- self.assertRejected(
- self.invoke(AF_MAX_RESTARTS="-1"), "native smoke requires AF_MAX_RESTARTS=0"
- )
-
- def test_native_restart_attempt_rejected(self):
- self.assertRejected(
- self.invoke(SLURM_RESTART_COUNT="1"), "restart limit exceeded"
- )
-
- def test_native_h100_constraint_reaches_hardware_guard(self):
- result = self.invoke(AF_EXPECTED_GPU_CONSTRAINT="H100")
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
-
- def test_native_wrong_precision_rejected(self):
- self.assertRejected(
- self.invoke(AF_PRECISION="bf16"), "native BF16-mixed required"
- )
-
- def test_native_wrong_cpu_envelope_rejected_before_tool(self):
- self.assertRejected(
- self.invoke(SLURM_CPUS_PER_TASK="16"), "exact eight CPUs required"
- )
-
- def planar(self, **changes):
- params = dict(
- AF_EXPERIMENT=PLANAR,
- AF_PRECISION="bf16",
- AF_MAX_RESTARTS="4",
- AF_EXPECTED_CONTENT_MANIFEST_SHA256="a1c81fb0ce8967aba795383a293180f9ba08a0ecfdd6f4a878afb20b39733761",
- AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256="80f835ad37c3d5c5b7b2d5c3e1656c307ee567a1f63f51081165bf404b8ceb52",
- )
- params.update(changes)
- return self.invoke(**params)
-
- def test_av0_suites_reach_same_canonical_phase_with_distinct_identity(self):
- from egomimic.benchmarks.libero.native_launch_profiles import AV0_PROFILES
-
- for suite, profile in AV0_PROFILES.items():
- with self.subTest(suite=suite):
- values = self.suite_values(suite, 42)
- values["AF_EXPERIMENT"] = profile.experiment
- values["SAFE_REPLAY_VARIABLE"] = profile.replay_environment
- output = self.root / "output-parent" / ("run-" + suite)
- values["AF_OUTPUT_DIR"] = str(output)
- result = self.invoke(**values)
- self.assertEqual(result.returncode, 91, result.stdout + result.stderr)
- argv = (output / "provenance/restart-0/exact-phase.argv0").read_bytes()
- self.assertIn(
- ("+experiment=" + profile.experiment + "\0").encode(), argv
- )
- self.assertEqual(
- (self.root / "handoff.log.replay").read_text().strip(),
- str(self.root / "data"),
- )
-
- def test_general_zero_restart_gate_preserved(self):
- self.assertRejected(self.planar(AF_MAX_RESTARTS="0"), "invalid AF_MAX_RESTARTS")
-
- def test_general_null_gpu_constraint_gate_preserved(self):
- self.assertRejected(self.planar(), "AF_EXPECTED_GPU_CONSTRAINT must be H100")
-
- def test_native_full_zero_restart_gate_unchanged(self):
- self.assertRejected(self.invoke(AF_RUN_KIND="full"), "invalid AF_MAX_RESTARTS")
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_libero_performance.py b/tests/test_libero_performance.py
deleted file mode 100644
index 8a2aa92672..0000000000
--- a/tests/test_libero_performance.py
+++ /dev/null
@@ -1,313 +0,0 @@
-"""Performance paths must preserve exact targets, data and training budgets."""
-
-import importlib.util
-import json
-import os
-import pickle
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-
-from egomimic.benchmarks.libero.cluster import (
- arc_checkpoint_settings,
- digest,
- restore_checkpoints,
- training_arguments,
- training_layout,
-)
-from egomimic.pipeline.stages_libero_arc import LiberoArcStage
-
-
-@pytest.mark.parametrize("mode", ["joint_dur", "stk", "dur"])
-def test_target_cache_is_exact_bounded_and_invalidated_by_codec_changes(mode):
- options = dict(
- arc_mode=mode,
- num_waypoints=24,
- max_translation=0.8,
- max_rotation_degrees=192,
- velocity_norm_bound=3**0.5,
- )
- cached = LiberoArcStage(**options, encode_cache_size=2)
- reference = LiberoArcStage(**options)
- actions = torch.from_numpy(
- np.random.default_rng(5).uniform(-1, 1, (3, 32, 7)).astype(np.float32)
- )
- actions[0, 4:12, :6] = 0 # Exercise duration dwells and velocity limitations.
- calls = []
- original = cached.codec.encode
-
- def count(row):
- calls.append(row.copy())
- return original(row)
-
- cached.codec.encode = count
- for indices, count_expected in [([0, 1, 0], 2), ([1, 0], 2), ([2, 0], 3), ([1], 4)]:
- batch = actions[indices]
- expected = reference.execute({"actions": batch}, mode="train")["target"]
- actual = cached.execute({"actions": batch}, mode="train")["target"]
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
- assert len(calls) == count_expected
- assert len(cached._encode_cache) <= 2
- actual.zero_() # Returned values must never alias cached physical tokens.
- cached.codec.max_translation = reference.codec.max_translation = 0.01
- actual = cached.execute({"actions": actions[1:2]}, mode="train")["target"]
- expected = reference.execute({"actions": actions[1:2]}, mode="train")["target"]
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
- assert len(calls) == 5
- assert cached.state_dict().keys() == reference.state_dict().keys()
- with pytest.raises(ValueError, match="batch"):
- cached._encode(actions[1:2].numpy().reshape(1, 7, 32))
-
-
-def test_target_cache_preserves_normalizer_changes_and_rejects_nonfinite():
- cached = LiberoArcStage(arc_mode="stk", encode_cache_size=8)
- reference = LiberoArcStage(arc_mode="stk")
- actions = torch.ones(1, 32, 7) * 0.25
- cached.execute({"actions": actions}, mode="train")
- for stage in (cached, reference):
- stage.action_scale.fill_(2)
- stage.action_offset.fill_(0.1)
- torch.testing.assert_close(
- cached.execute({"actions": actions}, mode="train")["target"],
- reference.execute({"actions": actions}, mode="train")["target"],
- rtol=0,
- atol=0,
- )
- assert len(cached._encode_cache) == 2
- with pytest.raises(ValueError, match="finite"):
- cached.execute(
- {"actions": torch.full_like(actions, float("nan"))}, mode="train"
- )
-
-
-def test_replay_array_handles_survive_repeated_reads_and_reset_on_spawn(
- tmp_path, monkeypatch
-):
- import zarr
-
- from egomimic.rldb.zarr.libero_dataset import LiberoReplayResolver, keymap
- from tests.test_libero_benchmark import make_replay
-
- path = tmp_path / "replay.zarr"
- make_replay(path)
- leaf = next(
- iter(LiberoReplayResolver(path, keymap(), "libero10").resolve().values())
- )
- expected = leaf[0]
- original = zarr.Group.__getitem__
- calls = []
-
- def tracked(group, key):
- calls.append(key)
- return original(group, key)
-
- monkeypatch.setattr(zarr.Group, "__getitem__", tracked)
- actual = leaf[0]
- assert calls == []
- restored = pickle.loads(pickle.dumps(leaf))
- assert restored._arrays is None
- restored_sample = restored[0]
- assert calls
- for key, value in expected.items():
- if torch.is_tensor(value):
- torch.testing.assert_close(actual[key], value, rtol=0, atol=0)
- torch.testing.assert_close(restored_sample[key], value, rtol=0, atol=0)
-
-
-@pytest.mark.parametrize("gpus", [1, 2, 4, 8])
-@pytest.mark.parametrize("mode", ["full", "smoke"])
-def test_distributed_arguments_preserve_global_batch(gpus, mode, tmp_path):
- args = training_arguments(
- "arc_stk", "libero_10", tmp_path / "replay", tmp_path, mode, 5001, gpus=gpus
- )
- with initialize_config_dir(
- version_base=None,
- config_dir=str(Path(__file__).parents[1] / "egomimic/hydra_configs"),
- ):
- cfg = compose(config_name="train_zarr_cartesian", overrides=args[3:])
- assert cfg.trainer.devices == gpus
- effective = cfg.benchmark.batch_size * gpus * cfg.trainer.accumulate_grad_batches
- assert effective == (1024 if mode == "full" else 4 * gpus)
- if mode == "full":
- assert cfg.callbacks.batch_budget.global_batch_size == effective
- assert cfg.trainer.max_epochs == 5001
- if gpus > 1:
- assert cfg.trainer.strategy == "ddp_find_unused_parameters_true"
- with pytest.raises(ValueError, match="GPUs"):
- training_layout(3, mode)
-
-
-def test_sweep_resume_reuses_confirmed_replay_and_resources(tmp_path, monkeypatch):
- from egomimic.benchmarks.libero import arc_sweep, cluster
-
- path = Path(__file__).parents[1] / "scripts/benchmarks/launch_libero_osmo.py"
- spec = importlib.util.spec_from_file_location("ddp_launcher", path)
- launcher = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(launcher)
- args = dict(
- commit="a" * 40,
- run_id="ddp-stk1",
- suite="libero_10",
- mode="full",
- arc_modes=["stk"],
- arc_profile="stk_1",
- resume_from_run="old-stk1",
- oat_reference_run="baseline",
- gpus=4,
- )
- with pytest.raises(ValueError, match="completed mode-specific replay"):
- launcher.workflow(**args)
- workflow = launcher.workflow(**args, arc_replay_runs={"stk": "confirmed-stk1"})[
- "workflow"
- ]
- assert workflow["resources"]["default"]["gpu"] == 4
- assert workflow["resources"]["default"]["cpu"] == 48
- assert workflow["resources"]["default"]["memory"] == "256Gi"
- env = workflow["tasks"][0]["environment"]
- for name, value in env.items():
- monkeypatch.setenv(name, value)
- assert env["TRAINING_GPUS"] == "4"
- calls = []
- monkeypatch.setattr(cluster, "execute", lambda argv, log: calls.append(argv))
- monkeypatch.setattr(
- sys,
- "argv",
- [
- "sweep",
- "--root",
- str(tmp_path),
- "--suite",
- "libero_10",
- "--run-id",
- "ddp-stk1",
- "--profile",
- "stk_1",
- "--arc-mode",
- "stk",
- ],
- )
- arc_sweep.main()
- assert len(calls) == 1 and calls[0][2].endswith(".cluster")
- assert "--arc-only" in calls[0]
- assert os.environ["ARC_REPLAY_RUNS_JSON"] == '{"stk": "confirmed-stk1"}'
-
-
-def test_resume_changes_world_size_without_resetting_optimizer_or_ema(tmp_path):
- script = Path(__file__).with_name("libero_ddp_resume_worker.py")
- env = dict(
- os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1"
- )
- # Execute against this audited checkout, not a globally installed package or
- # the worker script directory. This also survives a task-local test venv.
- env["PYTHONPATH"] = os.pathsep.join(
- [str(Path(__file__).resolve().parents[1]), env.get("PYTHONPATH", "")]
- )
- for suffix in ([], ["--resume"]):
- result = subprocess.run(
- [sys.executable, str(script), "--root", str(tmp_path), *suffix],
- env=env,
- capture_output=True,
- text=True,
- timeout=180,
- )
- assert result.returncode == 0, result.stdout[-8000:] + result.stderr[-8000:]
- initial = json.loads((tmp_path / "initial.json").read_text())
- resumed = json.loads((tmp_path / "resumed.json").read_text())
- assert initial["global_step"] == initial["ema_num_updates"] == 11
- assert resumed["global_step"] == resumed["ema_num_updates"] == 22
- assert resumed["optimizer_states"]
- budget = resumed["training_budget"]
- assert budget["world_size"] == budget["gradient_accumulation"] == 2
- assert budget["global_batch_size"] == 8
- assert budget["optimizer_steps_per_epoch"] == 11
-
-
-@pytest.mark.parametrize("method", ["arc", "arc_stk", "arc_dur"])
-@pytest.mark.parametrize("backbone", ["unet", "oat_dp"])
-def test_arc_resume_reads_actual_saved_graph_config(tmp_path, method, backbone):
- import io
- import shutil
-
- from egomimic.trainHydra import _build_model_config_tree
-
- args = training_arguments(
- method,
- "libero_10",
- tmp_path,
- tmp_path,
- "full",
- 5001,
- arc_backbone=backbone,
- )
- with initialize_config_dir(
- version_base=None,
- config_dir=str(Path(__file__).parents[1] / "egomimic/hydra_configs"),
- ):
- cfg = compose(config_name="train_zarr_cartesian", overrides=args[3:])
- config = _build_model_config_tree(cfg)
- assert "benchmark" not in config # Actual trainHydra checkpoint format.
- source = tmp_path / "source.ckpt"
- torch.save(
- {
- "global_step": 140,
- "loops": {"fit_loop": {"epoch_progress": {"current": {"completed": 2}}}},
- "training_budget": {"global_batch_size": 1024, "epochs": 5001},
- "optimizer_states": [{"state": {}}],
- "normalizer_state": {},
- "hyper_parameters": {"config_tree": config},
- },
- source,
- )
- runtime = {
- "suite": "libero_10",
- "mode": "full",
- "epochs": 5001,
- "global_batch_size": 1024,
- }
- receipts = {
- f"training/{method}/checkpoints/last.ckpt": {
- "sha256": digest(source),
- "bytes": source.stat().st_size,
- "uri": "s3://rldb/experiments/arc-oat-20260919/source/checkpoints/last.ckpt",
- }
- }
-
- class Storage:
- def get_object(self, Bucket, Key):
- value = runtime if Key.endswith("runtime.json") else receipts
- return {"Body": io.BytesIO(json.dumps(value).encode())}
-
- def download_file(self, bucket, key, destination):
- shutil.copyfile(source, destination)
-
- restored = restore_checkpoints(
- Storage(),
- "source",
- tmp_path / "evidence",
- suite="libero_10",
- mode="full",
- epochs=5001,
- allow_partial=True,
- methods=[method],
- )[method]
- assert restored["global_step"] == 140
- assert restored["benchmark"].get("arc_backbone", "unet") == backbone
- for key, value in restored["benchmark"].items():
- assert value == cfg.benchmark[key]
- with pytest.raises(ValueError, match="protocol"):
- arc_checkpoint_settings(config, suite="libero_spatial")
- from omegaconf import open_dict
-
- with open_dict(config.model.benchmark_protocol):
- config.model.benchmark_protocol.arc_backbone = "unknown"
- with pytest.raises(ValueError, match="unknown backbone"):
- arc_checkpoint_settings(config, suite="libero_10")
- config.model.benchmark_protocol.arc_backbone = backbone
- config.model.pipeline.stages[-1].num_waypoints += 1
- with pytest.raises(ValueError, match="parameters differ"):
- arc_checkpoint_settings(config, suite="libero_10")
diff --git a/tests/test_libero_saved_hydra_context.py b/tests/test_libero_saved_hydra_context.py
deleted file mode 100644
index 3ee9521d08..0000000000
--- a/tests/test_libero_saved_hydra_context.py
+++ /dev/null
@@ -1,341 +0,0 @@
-"""Real Hydra/OmegaConf restoration, no model/data/cluster execution."""
-
-import hashlib
-import json
-import os
-from pathlib import Path
-from unittest.mock import patch
-
-import pytest
-from hydra import __version__ as hydra_version
-from hydra.conf import HydraConf
-from hydra.core.hydra_config import HydraConfig
-from omegaconf import OmegaConf
-
-from egomimic.benchmarks.libero.native_launch_profiles import PROFILES
-from egomimic.benchmarks.libero.saved_hydra_context import load_saved_native_config
-
-
-def setup(tmp_path, suite):
- repo = tmp_path / "source"
- repo.mkdir()
- run = tmp_path / "run"
- (run / ".hydra").mkdir(parents=True)
- p = PROFILES[suite]
- receipt = tmp_path / "receipt.json"
- logical_sha256 = "a" * 64
- receipt.write_text(
- json.dumps(
- dict(
- suite=suite,
- replay_path="/replay",
- dataset_logical_sha256=logical_sha256,
- )
- )
- )
- proof = tmp_path / "physical.json"
- proof.write_text("{}")
-
- def sha(path):
- return hashlib.sha256(path.read_bytes()).hexdigest()
-
- binding = dict(
- source_commit="head",
- data_root="/replay",
- dataset_receipt_path=str(receipt),
- dataset_receipt_sha256=sha(receipt),
- physical_proof_path=str(proof),
- physical_proof_sha256=sha(proof),
- )
- config = dict(
- name=p.name,
- benchmark=dict(suite=suite, dataset="${oc.env:" + p.replay_environment + "}"),
- run_provenance=dict(
- source_commit="head",
- dataset_sha256=logical_sha256,
- dataset_content_aggregate_sha256=logical_sha256,
- ),
- norm_stats=dict(native_saved_state_binding=binding),
- data={
- group: {
- "libero_panda": {
- "resolver": dict(
- suite="${benchmark.suite}",
- folder_path="${oc.env:" + p.replay_environment + "}",
- )
- }
- }
- for group in ("train_datasets", "valid_datasets")
- },
- evaluator=dict(
- energy_seed_bank_path="${hydra:runtime.cwd}/egomimic/hydra_configs/evaluator/energy_score_seed_bank_k32_v1.json"
- ),
- )
- OmegaConf.save(OmegaConf.create(config), run / ".hydra/config.yaml")
- hydra = OmegaConf.structured(HydraConf)
- hydra.runtime.cwd = str(repo)
- hydra.runtime.output_dir = str(run)
- hydra.runtime.version = hydra_version
- saved = OmegaConf.to_container(hydra, resolve=False)
- saved["env"] = {}
- OmegaConf.save(OmegaConf.create({"hydra": saved}), run / ".hydra/hydra.yaml")
- return repo, run, p
-
-
-@pytest.mark.parametrize("suite", list(PROFILES))
-def test_saved_context_preserves_all_fields_and_missing_env(
- tmp_path, monkeypatch, suite
-):
- repo, run, p = setup(tmp_path, suite)
- monkeypatch.delenv(p.replay_environment, raising=False)
- with patch("subprocess.check_output", return_value="head\n"):
- cfg = load_saved_native_config(run, repo)
- assert os.environ[p.replay_environment] == "/replay"
- assert cfg.benchmark.dataset == "/replay"
- assert OmegaConf.get_type(HydraConfig.get()) is HydraConf
- assert dict(HydraConfig.get().env) == {}
- saved = OmegaConf.to_container(
- OmegaConf.load(run / ".hydra/hydra.yaml").hydra, resolve=False, enum_to_str=True
- )
- actual = OmegaConf.to_container(HydraConfig.get(), resolve=False, enum_to_str=True)
- assert all(actual[k] == v for k, v in saved.items())
-
-
-@pytest.mark.parametrize(
- "fault",
- ["conflict", "unknown_env", "cwd", "output", "source", "receipt", "resolver"],
-)
-def test_context_rejects_identity_drift(tmp_path, monkeypatch, fault):
- repo, run, p = setup(tmp_path, "libero_spatial")
- monkeypatch.delenv(p.replay_environment, raising=False)
- cfg = OmegaConf.load(run / ".hydra/config.yaml")
- saved = OmegaConf.load(run / ".hydra/hydra.yaml")
- if fault == "conflict":
- monkeypatch.setenv(p.replay_environment, "/wrong")
- if fault == "unknown_env":
- cfg.extra = "${oc.env:UNEXPECTED}"
- if fault == "cwd":
- saved.hydra.runtime.cwd = "/wrong"
- if fault == "output":
- saved.hydra.runtime.output_dir = "/wrong"
- if fault == "source":
- cfg.run_provenance.source_commit = "wrong"
- if fault == "receipt":
- Path(cfg.norm_stats.native_saved_state_binding.dataset_receipt_path).write_text(
- "{}"
- )
- if fault == "resolver":
- cfg.data.valid_datasets.libero_panda.resolver.folder_path = "/wrong"
- OmegaConf.save(cfg, run / ".hydra/config.yaml")
- OmegaConf.save(saved, run / ".hydra/hydra.yaml")
- with (
- patch("subprocess.check_output", return_value="head\n"),
- pytest.raises(ValueError),
- ):
- load_saved_native_config(run, repo)
-
-
-@pytest.mark.parametrize("fault", [None, "context", "source", "resolver"])
-def test_joint_saved_context_uses_complete_context_without_historical_binding(
- tmp_path, monkeypatch, fault
-):
- import hydra
-
- from egomimic.benchmarks.libero.native_launch_profiles import joint_phase_argv
-
- repo = Path(__file__).resolve().parents[1]
- run = tmp_path / "joint"
- (run / ".hydra").mkdir(parents=True)
- context = tmp_path / "context.json"
- split = tmp_path / "split.json"
- context.write_text("{}")
- split.write_text("{}")
-
- def sha(path):
- return hashlib.sha256(path.read_bytes()).hexdigest()
-
- args = dict(
- source_root=str(repo),
- original_root=str(tmp_path / "original"),
- sharpa_root=str(tmp_path / "sharpa"),
- output_root=str(run),
- context_path=str(context),
- split_path=str(split),
- split_sha256=sha(split),
- normalization_sha256=sha(context),
- source_commit="c" * 40,
- diagnostic_seed_bank=str(tmp_path / "seeds.json"),
- wandb_entity="test",
- wandb_project="test",
- wandb_run_id="joint",
- )
- with hydra.initialize_config_dir(
- version_base="1.2", config_dir=str(repo / "egomimic/hydra_configs")
- ):
- cfg = hydra.compose(
- config_name="train_zarr_cartesian",
- return_hydra_config=True,
- overrides=joint_phase_argv(phase="smoke", **args)[1:]
- + ["hydra/launcher=basic"],
- )
- cfg.hydra.runtime.cwd = str(repo)
- cfg.hydra.runtime.output_dir = str(run)
- HydraConfig.instance().set_config(cfg)
- payload = OmegaConf.to_container(
- OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]), resolve=True
- )
- saved = OmegaConf.to_container(cfg.hydra, resolve=False)
- if fault == "context":
- context.write_text("changed")
- if fault == "source":
- payload["run_provenance"]["source_commit"] = "wrong"
- if fault == "resolver":
- payload["data"]["valid_datasets"]["libero_panda_sharpa"]["resolver"][
- "folder_path"
- ] = "/wrong"
- OmegaConf.save(OmegaConf.create(payload), run / ".hydra/config.yaml")
- OmegaConf.save(OmegaConf.create({"hydra": saved}), run / ".hydra/hydra.yaml")
- with patch("subprocess.check_output", return_value="c" * 40 + "\n"):
- if fault:
- with pytest.raises(ValueError):
- load_saved_native_config(run, repo)
- else:
- loaded = load_saved_native_config(run, repo)
- assert OmegaConf.to_container(loaded, resolve=True) == payload
-
-
-@pytest.mark.parametrize("fault", [None, "rolling", "wrong_callback", "wrong_path"])
-def test_joint_checkpoint_distinguishes_configured_terminal_from_rolling_alias(
- tmp_path, fault
-):
- from scripts.train.verify_action_flow_training_smoke import (
- _joint_smoke_checkpoint_path,
- )
-
- cadence = tmp_path / "epoch-0000-step-000000008.ckpt"
- terminal = tmp_path / "last.ckpt"
- cadence.write_bytes(b"cadence")
- terminal.write_bytes(b"terminal")
- cfg = OmegaConf.create(
- {
- "callbacks": {
- "model_checkpoint": {"save_last": False},
- "ema": {
- "_target_": "egomimic.pl_utils.oat_matched_training.ActionFlowFixedEMACallback",
- "final_checkpoint_path": str(terminal),
- },
- }
- }
- )
- if fault == "rolling":
- cfg.callbacks.model_checkpoint.save_last = True
- if fault == "wrong_callback":
- cfg.callbacks.ema._target_ = "unknown.Callback"
- if fault == "wrong_path":
- cfg.callbacks.ema.final_checkpoint_path = "/wrong"
- if fault:
- with pytest.raises(Exception, match="rolling last"):
- _joint_smoke_checkpoint_path(tmp_path, cfg)
- else:
- assert _joint_smoke_checkpoint_path(tmp_path, cfg) == cadence
- assert terminal.read_bytes() == b"terminal"
- with pytest.raises(Exception, match="rolling last"):
- _joint_smoke_checkpoint_path(tmp_path)
-
-
-def test_separate_verifier_finalizes_through_same_loaded_helpers(tmp_path, monkeypatch):
- from egomimic.benchmarks.libero import saved_hydra_context as loader
- from scripts.train import libero_native_launch_contract as producer
- from scripts.train import verify_action_flow_training_smoke as shared
-
- training, verifier, run = [tmp_path / p for p in ("training", "verifier", "run")]
- for p in (training, verifier, run):
- p.mkdir()
- preflight = tmp_path / "preflight.json"
- preflight.write_text("{}")
- runtime = tmp_path / "runtime.json"
- runtime.write_text("{}")
- monkeypatch.setattr(shared, "REPOSITORY_ROOT", verifier)
- monkeypatch.setattr(shared, "VERIFIER_PROVENANCE", None)
- monkeypatch.setattr(
- shared,
- "_git_head",
- lambda repository_root=None: "a" * 40
- if repository_root == verifier
- else "b" * 40,
- )
- monkeypatch.setattr(shared.subprocess, "check_output", lambda *a, **k: "")
- monkeypatch.setattr(shared.subprocess, "run", lambda *a, **k: None)
- monkeypatch.setenv("SLURM_STEP_ID", "0")
- monkeypatch.setenv("AF_PREFLIGHT_RESULT", str(preflight))
- calls = []
-
- def load(root, repo):
- assert root == run and repo == training
- return OmegaConf.create({"original": True})
-
- def evidence(cfg, root, **kwargs):
- assert isinstance(cfg, dict) and cfg["original"] and root == run
- json.dumps(cfg, allow_nan=False)
- assert kwargs["shared"] is shared
- calls.append("finalize")
- return {"actual": True}
-
- monkeypatch.setattr(loader, "load_saved_native_config", load)
- monkeypatch.setattr(producer, "joint_smoke_evidence", evidence)
- monkeypatch.setattr(producer, "write", lambda path, value: calls.append("write"))
- monkeypatch.setattr(
- shared,
- "verify_smoke",
- lambda *a, **k: calls.append("verify") or {"run_dir": str(run)},
- )
- assert (
- shared.main(
- [
- str(run),
- "--experiment",
- "sim/libero/action_flow_libero_spatial_sharpa_gripper_keypoints_s42",
- "--expected-head",
- "b" * 40,
- "--training-repo",
- str(training),
- "--expected-verifier-head",
- "a" * 40,
- "--finalize-runtime-lock",
- str(runtime),
- "--expected-preflight-sha256",
- hashlib.sha256(preflight.read_bytes()).hexdigest(),
- ]
- )
- == 0
- )
- assert calls == ["finalize", "write", "verify"]
-
-
-def test_joint_artifacts_bind_cadence_epoch_and_preserve_postfit(tmp_path, monkeypatch):
- import torch
-
- from scripts.train import verify_action_flow_training_smoke as shared
-
- monkeypatch.delenv("SLURM_JOB_ID", raising=False)
- run = tmp_path / "run"
- checkpoint = run / "checkpoints/epoch-0000-step-000000008.ckpt"
- checkpoint.parent.mkdir(parents=True)
- checkpoint.write_bytes(b"filename-only-selection")
- root = run / "artifacts"
- for epoch in (0, 1):
- path = root / f"epoch-{epoch}-step-8/rank-0-batch-0.pt"
- path.parent.mkdir(parents=True)
- torch.save({"global_step": 8, "epoch": epoch}, path)
- epoch = shared._joint_smoke_artifact_epoch(run, None)
- assert epoch == 0
- selected, payload = shared._step_artifact(
- root, label="joint", expected_step=8, expected_epoch=epoch
- )
- assert selected.parent.name == "epoch-0-step-8" and payload["epoch"] == 0
- assert (root / "epoch-1-step-8/rank-0-batch-0.pt").is_file()
- with pytest.raises(shared.SmokeVerificationError, match="expected one step-8"):
- shared._step_artifact(root, label="joint", expected_step=8)
- torch.save({"global_step": 8, "epoch": 1}, selected)
- with pytest.raises(shared.SmokeVerificationError, match="epoch differs"):
- shared._step_artifact(root, label="joint", expected_step=8, expected_epoch=0)
diff --git a/tests/test_likelihood_training_behavior_parity.py b/tests/test_likelihood_training_behavior_parity.py
deleted file mode 100644
index f15916de2c..0000000000
--- a/tests/test_likelihood_training_behavior_parity.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""Compare likelihood mechanics against the retained source-bound wrapper."""
-
-import ast
-import copy
-import hashlib
-import json
-from pathlib import Path
-
-import pytest
-import torch
-import yaml
-
-from egomimic.pipeline.algo import PipelineAlgo
-from egomimic.pl_utils.pl_model import ModelWrapper
-from egomimic.pl_utils.pl_model_action_flow_likelihood import (
- ActionFlowLikelihoodModelWrapper,
-)
-from egomimic.pl_utils.training_behavior_action_flow_likelihood import (
- ActionFlowLikelihoodTrainingBehavior,
-)
-from tests.test_action_flow_likelihood import batch, stages
-
-
-def test_likelihood_recipe_changes_only_framework_selection():
- root = Path(__file__).resolve().parents[1]
- path = "egomimic/hydra_configs/model/bf/us_action_flow_bc_bridge_likelihood.yaml"
- candidate = yaml.safe_load((root / path).read_text())
- assert candidate.pop("_target_") == "egomimic.pl_utils.pl_model.ModelWrapper"
- assert candidate.pop("training_behavior") == {
- "_target_": "egomimic.pl_utils.training_behavior_action_flow_likelihood.ActionFlowLikelihoodTrainingBehavior"
- }
- # Canonical YAML payload (without the framework target), independently
- # extracted from exact training source7c253e5. No full-history checkout is
- # required in shallow CI, and every scientific field is included.
- digest = hashlib.sha256(
- json.dumps(candidate, sort_keys=True, separators=(",", ":")).encode()
- ).hexdigest()
- assert digest == "e55a71bec35d62e694119ea3c42fed1a8e7855f2d97fc3e7310c21a37831e033"
-
-
-def test_port_preserves_historical_tensor_and_checkpoint_algorithms():
- root = Path(__file__).resolve().parents[1] / "egomimic/pl_utils"
-
- def methods(path, name):
- tree = ast.parse((root / path).read_text())
- cls = next(
- node
- for node in tree.body
- if isinstance(node, ast.ClassDef) and node.name == name
- )
- return {
- node.name: node for node in cls.body if isinstance(node, ast.FunctionDef)
- }
-
- historical = methods(
- "pl_model_action_flow_likelihood.py", "ActionFlowLikelihoodModelWrapper"
- )
- candidate = methods(
- "training_behavior_action_flow_likelihood.py",
- "ActionFlowLikelihoodTrainingBehavior",
- )
-
- class ContextToHistorical(ast.NodeTransformer):
- def visit_Attribute(self, node):
- node = self.generic_visit(node)
- if (
- isinstance(node.value, ast.Attribute)
- and node.value.attr == "context"
- and isinstance(node.value.value, ast.Name)
- and node.value.value.id == "self"
- ):
- node.value = ast.Name(id="self", ctx=ast.Load())
- return node
-
- for name in (
- "_log_telemetry",
- "_capture_gradient_routes",
- "_log_prediction_metrics",
- "validation_step",
- "on_save_checkpoint",
- ):
- port_name = (
- "log_prediction_metrics" if name == "_log_prediction_metrics" else name
- )
- port = ContextToHistorical().visit(copy.deepcopy(candidate[port_name]))
- port.name = name
- assert ast.dump(port, include_attributes=False) == ast.dump(
- historical[name], include_attributes=False
- ), name
-
-
-@pytest.mark.parametrize("cadence", [0, 1])
-def test_training_validation_checkpoint_and_rng_match_historical(monkeypatch, cadence):
- def observe(shared):
- torch.manual_seed(73)
- algo = PipelineAlgo(stages(dropout=0.3), device="cpu")
- wrapper = (
- ModelWrapper(
- pipeline=algo,
- enable_grad_norm=False,
- training_behavior=ActionFlowLikelihoodTrainingBehavior(
- gradient_telemetry_cadence=cadence
- ),
- )
- if shared
- else ActionFlowLikelihoodModelWrapper(
- pipeline=algo,
- enable_grad_norm=False,
- gradient_telemetry_cadence=cadence,
- )
- )
- logged = {}
- monkeypatch.setattr(
- wrapper,
- "log",
- lambda name, value, **kw: logged.update(
- {name: (torch.as_tensor(value).detach().clone(), kw)}
- ),
- )
- optimizer = torch.optim.SGD(wrapper.parameters(), lr=0.001)
- losses, gradients = [], []
- for index in range(2):
- optimizer.zero_grad()
- loss = wrapper.training_step({"example": batch()}, index)
- losses.append(loss.detach().clone())
- loss.backward()
- gradients.append(
- [
- None if p.grad is None else p.grad.detach().clone()
- for p in wrapper.parameters()
- ]
- )
- wrapper.on_after_backward()
- wrapper.on_before_optimizer_step(optimizer)
- optimizer.step()
- wrapper.eval()
- before = torch.get_rng_state().clone()
- with torch.inference_mode():
- wrapper.validation_step({"example": batch(), "inactive": None}, 2, 1)
- # batch() constructs random data, so compare resulting RNG across sources.
- after = torch.get_rng_state().clone()
- checkpoint = {}
- wrapper.on_save_checkpoint(checkpoint)
- return (
- copy.deepcopy(wrapper.nets.state_dict()),
- losses,
- gradients,
- {k: v for k, v in logged.items() if not k.startswith("Timing/")},
- before,
- after,
- {k: v for k, v in checkpoint.items() if k.startswith("action_flow_")},
- )
-
- historical, candidate = observe(False), observe(True)
- assert historical[0].keys() == candidate[0].keys()
- for key in historical[0]:
- assert torch.equal(historical[0][key], candidate[0][key]), key
- for a, b in zip(historical[1], candidate[1]):
- assert torch.equal(a, b)
- for left, right in zip(historical[2], candidate[2]):
- for a, b in zip(left, right):
- assert (a is None) == (b is None)
- if a is not None:
- assert torch.equal(a, b)
- assert historical[3].keys() == candidate[3].keys()
- for key, (value, options) in historical[3].items():
- assert torch.equal(value, candidate[3][key][0]), key
- assert options == candidate[3][key][1], key
- assert torch.equal(historical[4], candidate[4])
- assert torch.equal(historical[5], candidate[5])
- assert historical[6] == candidate[6]
diff --git a/tests/test_local_cpu_storage_guard.py b/tests/test_local_cpu_storage_guard.py
deleted file mode 100644
index dd11a6d4d2..0000000000
--- a/tests/test_local_cpu_storage_guard.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from types import SimpleNamespace
-
-import pytest
-
-from scripts.local_cpu_storage_guard import pytest_runtest_setup
-
-
-@pytest.mark.parametrize("free_gib", [3, 4, 8])
-def test_per_test_headroom_gate(monkeypatch, free_gib):
- monkeypatch.delenv("LOCAL_CPU_TEST_MIN_FREE_KIB", raising=False)
- monkeypatch.setattr(
- "scripts.local_cpu_storage_guard.shutil.disk_usage",
- lambda path: SimpleNamespace(free=free_gib * 1024**3),
- )
- if free_gib < 4:
- with pytest.raises(pytest.UsageError, match="blocked before test"):
- pytest_runtest_setup(None)
- else:
- pytest_runtest_setup(None)
-
-
-def test_per_test_floor_cannot_be_lowered(monkeypatch):
- monkeypatch.setenv("LOCAL_CPU_TEST_MIN_FREE_KIB", "1")
- with pytest.raises(pytest.UsageError, match="cannot be below"):
- pytest_runtest_setup(None)
diff --git a/tests/test_local_cpu_validation_guard.py b/tests/test_local_cpu_validation_guard.py
deleted file mode 100644
index 3b17dc3fa0..0000000000
--- a/tests/test_local_cpu_validation_guard.py
+++ /dev/null
@@ -1,40 +0,0 @@
-"""Local validation limits cannot be lowered into misleading storage failures."""
-
-import os
-import subprocess
-from pathlib import Path
-
-import pytest
-
-
-@pytest.mark.parametrize("value", ["0", "4194303", "invalid", "999999999999"])
-def test_low_invalid_or_unavailable_space_blocks_before_python(value):
- script = Path(__file__).resolve().parents[1] / "scripts/run_local_cpu_validation.sh"
- env = dict(
- os.environ,
- LOCAL_CPU_TEST_MIN_FREE_KIB=value,
- LOCAL_CPU_TEST_PYTHON="/does-not-exist",
- )
- result = subprocess.run(
- ["bash", str(script), "--version"], env=env, capture_output=True, text=True
- )
- assert result.returncode == 64
- assert "Local CPU validation" in result.stderr
- assert "/does-not-exist" not in result.stderr
-
-
-@pytest.mark.parametrize(
- ("target", "timeout"),
- [("tests/test_recursive_config_audit.py", 600), ("tests/test_oat_native.py", 60)],
-)
-def test_long_constructor_audit_has_its_own_diagnostic_threshold(target, timeout):
- script = Path(__file__).resolve().parents[1] / "scripts/run_local_cpu_validation.sh"
- # Print the runner's invocation, without executing or skipping any real test.
- result = subprocess.run(
- ["bash", str(script), target],
- env=dict(os.environ, LOCAL_CPU_TEST_PYTHON="/bin/echo"),
- capture_output=True,
- text=True,
- )
- assert result.returncode == 0, result.stderr
- assert f"faulthandler_timeout={timeout}" in result.stdout
diff --git a/tests/test_measured_speed_removed.py b/tests/test_measured_speed_removed.py
deleted file mode 100644
index 387c844836..0000000000
--- a/tests/test_measured_speed_removed.py
+++ /dev/null
@@ -1,76 +0,0 @@
-"""Retired measured-speed paths must never become selectable again."""
-
-from pathlib import Path
-
-import pytest
-from omegaconf import OmegaConf
-
-from egomimic.pipeline.stages_speed import (
- SharedSpeedCondition,
- build_speed_conditioned_pipeline,
- requested_rollout_condition,
-)
-from egomimic.rldb.zarr.physical_retiming import PhysicalWindowRetiming
-
-ROOT = Path(__file__).parents[1]
-
-
-@pytest.mark.parametrize("mode", ["native_speed", "physical_speed", "unknown"])
-def test_conditioner_and_factory_reject_retired_modes(mode):
- with pytest.raises(ValueError, match="removed"):
- SharedSpeedCondition(conditioning_input=mode)
- with pytest.raises(ValueError, match="removed"):
- build_speed_conditioned_pipeline([], conditioning_input=mode)
-
-
-def test_default_is_direct_multiplier_and_reference_cannot_select_old_mode():
- model = SharedSpeedCondition()
- assert model.speed_key == "retiming_rate"
- assert model.conditioning_input == "retiming_multiplier"
- assert "speed_reference" not in model.state_dict()
- with pytest.raises(ValueError, match="removed"):
- SharedSpeedCondition(1.0)
- with pytest.raises(ValueError, match="retiming_rate"):
- SharedSpeedCondition(speed_key="requested_speed")
-
-
-def test_retimer_rejects_retired_mode_before_loading_data():
- with pytest.raises(ValueError, match="removed"):
- PhysicalWindowRetiming(
- [1.0],
- {"left": "pose_wxyz", "right": "pose_wxyz"},
- ["left", "right"],
- 30,
- timestamp_key="clock",
- conditioning_input="native_speed",
- )
-
-
-@pytest.mark.parametrize("value", [None, "native_speed"])
-def test_rollout_never_falls_back_to_measured_speed(value):
- pipeline = {
- "_target_": "egomimic.pipeline.stages_speed.build_speed_conditioned_pipeline"
- }
- if value is not None:
- pipeline["conditioning_input"] = value
- cfg = OmegaConf.create(
- {"model": {"pipeline": pipeline}, "deployment": {"requested_speed": 1.0}}
- )
- with pytest.raises(ValueError):
- requested_rollout_condition(cfg)
-
-
-def test_configs_and_launchers_have_no_retired_routes():
- retired = (
- "human_speed_v1",
- "action_flow_cotrain_uc_speed_interpolation",
- "yam_human_keypoints_speed_h816_private512_s42",
- "AF_STATIONARY_SPEED_REFERENCE",
- "AF_SPEED_REFERENCE",
- )
- for root in (ROOT / "egomimic/hydra_configs", ROOT / "scripts/train"):
- for file in root.rglob("*"):
- if file.suffix not in {".yaml", ".sh", ".sbatch", ".py"}:
- continue
- body = file.read_text()
- assert not any(name in body for name in retired), file
diff --git a/tests/test_multiplier_training_launch.py b/tests/test_multiplier_training_launch.py
deleted file mode 100644
index df8cdd827b..0000000000
--- a/tests/test_multiplier_training_launch.py
+++ /dev/null
@@ -1,93 +0,0 @@
-"""Typed multiplier route preserves native per-source execution and full batches."""
-
-import runpy
-import subprocess
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-
-from egomimic.pipeline.stages_speed import configure_dit_checkpoint_policy
-
-
-def test_multiplier_training_contract():
- root = Path(__file__).resolve().parents[1]
- with initialize_config_dir(
- config_dir=str(root / "egomimic/hydra_configs"), version_base="1.3"
- ):
- cfg = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "hydra/launcher=basic",
- "+experiment=sim/pusht/action_flow_cotrain_uc_multiplier_interpolation",
- "trainer.precision=bf16",
- "trainer.devices=1",
- "trainer.num_nodes=1",
- "model.pipeline.stages.6.inference_method=euler",
- "model.pipeline.stages.6.num_inference_steps=50",
- ],
- )
- validator = runpy.run_path(
- str(root / "scripts/train/validate_multiplier_config.py")
- )
- validator["validate_speed_contract"](cfg, "scalar")
- assert cfg.callbacks.ema.decay == 0.9978
- assert cfg.callbacks.ema.use_warmup is False
- assert "homogeneous_dithalf" not in cfg.callbacks
- assert cfg.model.pipeline.dit_checkpoint_policy == "dit_half"
- assert not (root / "egomimic/pl_utils/homogeneous_training/callback.py").exists()
- assert cfg.norm_stats.sample_frac == 1.0
- assert cfg.norm_stats.norm_mode == "minmax"
- # The typed validator above requires one full batch32 per source.
- assert len(cfg.data.train_datasets) == 2
-
-
-def _owners():
- def owner():
- return {
- "backbone": {
- "_target_": "egomimic.models.unite_dit.UniteDiTBackbone",
- "gradient_checkpointing": True,
- "depth": 12,
- }
- }
-
- return {"encoders": {"u": owner(), "chain": owner()}}, {"field": owner()}
-
-
-def test_native_dit_half_does_not_change_image_or_decoder_execution():
- encoder, field = _owners()
- configure_dit_checkpoint_policy(encoder, field, "dit_half")
- for owner in [*encoder["encoders"].values(), field["field"]]:
- assert owner["backbone"]["checkpoint_policy"] == "dit_half"
- assert set(encoder) == {"encoders"}
- assert set(field) == {"field"}
-
-
-@pytest.mark.parametrize("invalid", ["disabled", "odd_depth", "wrong_owner", "unknown"])
-def test_native_dit_half_rejects_mismatch_atomically(invalid):
- encoder, field = _owners()
- backbone = field["field"]["backbone"]
- policy = "dit_half"
- if invalid == "disabled":
- backbone["gradient_checkpointing"] = False
- elif invalid == "odd_depth":
- backbone["depth"] = 11
- elif invalid == "wrong_owner":
- backbone["_target_"] = "unreviewed.Backbone"
- else:
- policy = "guess"
- with pytest.raises(ValueError):
- configure_dit_checkpoint_policy(encoder, field, policy)
- assert all(
- "checkpoint_policy" not in owner["backbone"]
- for owner in [*encoder["encoders"].values(), field["field"]]
- )
-
-
-def test_launcher_child_helper_exists_and_parses():
- root = Path(__file__).resolve().parents[1]
- helper = root / "scripts/train/action_flow_run_child.sh"
- assert helper.is_file()
- subprocess.run(["bash", "-n", str(helper)], check=True)
- assert "--cpu-bind=none" in helper.read_text()
diff --git a/tests/test_native_final_validation_lifecycle.py b/tests/test_native_final_validation_lifecycle.py
deleted file mode 100644
index d1a0392d5d..0000000000
--- a/tests/test_native_final_validation_lifecycle.py
+++ /dev/null
@@ -1,135 +0,0 @@
-"""Execute maintained lifecycle methods without model/data/runtime imports."""
-
-import ast
-import sys
-import types
-import unittest
-from pathlib import Path
-from unittest.mock import patch
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-def functions(path, names):
- tree = ast.parse((ROOT / path).read_text())
- nodes = []
- for node in ast.walk(tree):
- if isinstance(node, ast.FunctionDef) and node.name in names:
- node.decorator_list = []
- nodes.append(node)
- ns = {}
- exec(compile(ast.Module(body=nodes, type_ignores=[]), str(path), "exec"), ns)
- return ns
-
-
-F = functions("egomimic/trainHydra.py", {"_needs_final_validation"})[
- "_needs_final_validation"
-]
-N = functions(
- "egomimic/eval/libero_action_flow_eval.py",
- {"on_validation_end", "has_completed_validation"},
-)
-
-
-class Native:
- on_validation_end = N["on_validation_end"]
- has_completed_validation = N["has_completed_validation"]
-
- def __init__(self, group="valid", batches=None, sanity=False):
- self.artifact_root = "unused"
- self.artifact_identity = {}
- self.group = group
- self.artifact_batches = [{"batch_index": 0}] if batches is None else batches
- self.model = types.SimpleNamespace(global_step=2)
- self.trainer = types.SimpleNamespace(sanity_checking=sanity)
- self._completed_validation_step = None
-
-
-class Lifecycle(unittest.TestCase):
- def publication(self, e, error=None):
- module = types.ModuleType("egomimic.benchmarks.libero.action_flow_artifacts")
- module.canonical_sha = lambda identity: "identity"
-
- def writer(root, payload):
- if error:
- raise error
-
- module.write_artifact = writer
- with patch.dict(sys.modules, {module.__name__: module}):
- e.on_validation_end()
-
- def test_periodic_terminal_completion_skips_only_duplicate(self):
- e = Native()
- self.assertTrue(F(e, 2))
- self.publication(e)
- self.assertFalse(F(e, 2))
- self.assertTrue(F(e, 3))
- self.assertTrue(F(e, 80000))
-
- def test_missing_generic_and_nonboolean_capability_keep_final(self):
- for e in (
- None,
- object(),
- types.SimpleNamespace(has_completed_validation=lambda step: 1),
- ):
- self.assertTrue(F(e, 2))
-
- def test_empty_sanity_and_other_group_keep_final(self):
- for e in (Native(batches=[]), Native(sanity=True), Native(group="train_viz")):
- self.publication(e)
- self.assertTrue(F(e, 2))
-
- def test_new_empty_or_failed_pass_invalidates_prior_completion(self):
- tree = ast.parse(
- (ROOT / "egomimic/eval/libero_action_flow_eval.py").read_text()
- )
- method = next(
- n
- for n in ast.walk(tree)
- if isinstance(n, ast.FunctionDef) and n.name == "on_validation_start"
- )
- e = Native()
- e._completed_validation_step = 2
- exec(
- compile(
- ast.Module(body=method.body[:1], type_ignores=[]),
- "native-start",
- "exec",
- ),
- {"self": e},
- )
- e.artifact_batches = []
- self.publication(e)
- self.assertTrue(F(e, 2))
-
- def test_failed_immutable_publication_does_not_complete(self):
- e = Native()
- with self.assertRaises(FileExistsError):
- self.publication(e, FileExistsError("immutable collision"))
- self.assertTrue(F(e, 2))
- self.assertTrue(e.artifact_batches)
-
- def test_completion_reset_precedes_replayed_valid_batch(self):
- tree = ast.parse(
- (ROOT / "egomimic/eval/libero_action_flow_eval.py").read_text()
- )
- method = next(
- n
- for n in ast.walk(tree)
- if isinstance(n, ast.FunctionDef) and n.name == "on_validation_step"
- )
- # Execute the actual initial statements, stopping before model inference.
- statements = method.body[:2]
- e = Native()
- e._completed_validation_step = 2
- exec(
- compile(
- ast.Module(body=statements, type_ignores=[]), "native-reset", "exec"
- ),
- {"self": e, "batch_idx": 0, "dataloader_idx": 0},
- )
- self.assertTrue(F(e, 2))
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/tests/test_oat_step_budget.py b/tests/test_oat_step_budget.py
deleted file mode 100644
index 45b12d0023..0000000000
--- a/tests/test_oat_step_budget.py
+++ /dev/null
@@ -1,214 +0,0 @@
-"""Step-limited pilot provenance and weights-only callback metadata."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from egomimic.pl_utils.oat_training import (
- OATBatchBudgetCallback,
- OATEMACallback,
- OATTrainingBehavior,
-)
-
-
-@pytest.mark.parametrize(
- "epochs,max_steps,expected",
- [(3, -1, 30), (3, 17, 17), (-1, 17, 17), (3, 45, 30)],
-)
-def test_step_budget_honors_first_limit(epochs, max_steps, expected):
- trainer = SimpleNamespace(
- datamodule=SimpleNamespace(
- train_datasets={"x": range(83)},
- train_dataloader_params={"x": {"batch_size": 2, "drop_last": True}},
- ),
- accumulate_grad_batches=4,
- world_size=1,
- limit_train_batches=1.0,
- max_epochs=epochs,
- max_steps=max_steps,
- is_global_zero=False,
- )
- callback = OATBatchBudgetCallback(global_batch_size=8)
- callback.on_fit_start(trainer, None)
- assert callback.budget["total_optimizer_steps"] == expected
- assert callback.budget["max_steps"] == max_steps
- assert callback.budget["dropped_examples_per_epoch"] == 3
- assert trainer.limit_train_batches == 40
-
-
-def test_joint_budget_matches_actual_cycling_loader_and_preserves_both_sources():
- from lightning.pytorch.utilities.combined_loader import CombinedLoader
- from torch.utils.data import DataLoader
-
- datasets = {"original": range(83), "sharpa": range(37)}
- params = {name: {"batch_size": 16, "drop_last": True} for name in datasets}
- trainer = SimpleNamespace(
- datamodule=SimpleNamespace(
- train_datasets=datasets,
- train_dataloader_params=params,
- train_loader_mode="max_size_cycle",
- ),
- accumulate_grad_batches=1,
- world_size=1,
- limit_train_batches=1.0,
- max_epochs=-1,
- max_steps=120000,
- is_global_zero=False,
- )
- callback = OATBatchBudgetCallback(
- global_batch_size=32, expected_source_batches={"original": 16, "sharpa": 16}
- )
- callback.on_fit_start(trainer, None)
- loader = CombinedLoader(
- {name: DataLoader(data, **params[name]) for name, data in datasets.items()},
- "max_size_cycle",
- )
- seen = list(loader)
- assert trainer.limit_train_batches == len(seen) == 5
- assert all(
- {name: len(batch) for name, batch in item[0].items()}
- == {"original": 16, "sharpa": 16}
- for item in seen
- )
- assert callback.budget["total_optimizer_steps"] == 120000
- assert callback.budget["global_batch_size"] == 32
- params["sharpa"]["batch_size"] = 8
- trainer.limit_train_batches = 1.0
- with pytest.raises(ValueError, match="preserve batch size"):
- callback.on_fit_start(trainer, None)
-
-
-@pytest.mark.parametrize("weights_only", [True, False])
-def test_weights_only_retains_ema_and_budget_without_optimizer(
- monkeypatch, weights_only
-):
- from egomimic.models.oat import checkpoint as contract
-
- monkeypatch.setattr(contract, "validate_input_representation", lambda stages: {})
- stage = SimpleNamespace(normalizer_state={"scale": 1}, data_context={"split": 42})
-
- class Model(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.weight = torch.nn.Parameter(torch.tensor([2.0]))
- self.model = SimpleNamespace(pipeline=SimpleNamespace(stages=[stage]))
-
- model = Model()
- budget = OATBatchBudgetCallback()
- budget.budget = {"total_optimizer_steps": 10000}
- ema = OATEMACallback()
- trainer = SimpleNamespace(callbacks=[budget, ema], global_step=0)
- model.trainer = trainer
- ema.on_fit_start(trainer, model)
- ema._num_updates = 3
- ema._shadow["weight"].fill_(1.0)
- behavior = OATTrainingBehavior()
- behavior.bind(model)
- checkpoint = {"state_dict": model.state_dict(), "global_step": 3}
- if not weights_only:
- checkpoint["optimizer_states"] = [{"sentinel": True}]
- for callback in trainer.callbacks:
- callback.on_save_checkpoint(trainer, model, checkpoint)
- # Full-state model hook must not call callbacks a second time.
- monkeypatch.setattr(
- ema, "on_save_checkpoint", lambda *args: pytest.fail("twice")
- )
- behavior.on_save_checkpoint(checkpoint)
- assert checkpoint["ema_num_updates"] == 3
- assert checkpoint["ema_state_dict"]["weight"].item() == 1.0
- assert checkpoint["state_dict"]["weight"].item() == 2.0
- assert checkpoint["training_budget"]["total_optimizer_steps"] == 10000
- assert checkpoint["normalizer_state"] == {"scale": 1}
- assert ("optimizer_states" not in checkpoint) == weights_only
-
-
-@pytest.mark.parametrize("basis", ["uniform", "fourier", "chebyshev"])
-@pytest.mark.parametrize("mode", ["smoke", "full"])
-@pytest.mark.parametrize("gpus", [1, 4])
-def test_pilot_arguments_compose_matched_budget_and_checkpoint_policy(
- tmp_path, basis, mode, gpus
-):
- from pathlib import Path
-
- from hydra import compose, initialize_config_dir
- from hydra.utils import instantiate
-
- from egomimic.benchmarks.libero.cluster import global_basis_training_arguments
-
- args = global_basis_training_arguments(
- basis, tmp_path / "data.zarr", tmp_path, mode, f"test-{basis}-{mode}", gpus=gpus
- )
- with initialize_config_dir(
- config_dir=str(Path(__file__).parents[1] / "egomimic/hydra_configs"),
- version_base="1.3",
- ):
- cfg = compose(config_name="train_zarr_cartesian", overrides=args[3:])
- assert cfg.benchmark.arc_basis == basis
- assert cfg.trainer.max_steps == (16 if mode == "smoke" else 10000)
- assert cfg.trainer.max_epochs == -1
- assert cfg.trainer.devices == gpus
- assert cfg.model.train_log_on_step
- assert cfg.trainer.check_val_every_n_epoch is None
- cp = cfg.callbacks.model_checkpoint
- assert cp.every_n_train_steps == (8 if mode == "smoke" else 2500)
- assert (
- cfg.trainer.val_check_interval
- == cp.every_n_train_steps * cfg.trainer.accumulate_grad_batches
- )
- assert cp.every_n_epochs == 0
- assert cp.save_weights_only and not cp.save_last and cp.save_top_k == -1
- assert not cp.save_on_train_epoch_end
- assert cfg.callbacks.ema.final_checkpoint_path is None
- assert "{epoch:" in cp.filename and "{step:" in cp.filename
- instantiate(cp) # Real callback argument validation, no training.
- expected_batch = 4 * gpus if mode == "smoke" else 1024
- assert cfg.callbacks.batch_budget.global_batch_size == expected_batch
- assert (
- cfg.benchmark.batch_size * gpus * cfg.trainer.accumulate_grad_batches
- == expected_batch
- )
- assert cfg.logger.wandb.resume == "never" and not cfg.logger.wandb.offline
-
-
-def test_early_recovery_save_has_full_state_and_refuses_overwrite(tmp_path):
- import json
- from types import SimpleNamespace
-
- import torch
-
- from egomimic.pl_utils.oat_training import EarlyRecoveryCheckpoint
-
- callback = EarlyRecoveryCheckpoint(str(tmp_path), first_step=2)
- state = {
- "global_step": 2,
- "ema_num_updates": 2,
- "state_dict": {"weight": torch.ones(2)},
- "optimizer_states": [{"state": {0: {"moment": torch.zeros(2)}}}],
- "lr_schedulers": [{"last_epoch": 2}],
- "data_context": {"bound": True},
- }
-
- def save(path, weights_only):
- assert weights_only is False
- torch.save(state, path)
-
- trainer = SimpleNamespace(
- global_step=1, current_epoch=0, world_size=1, save_checkpoint=save
- )
- callback.on_train_batch_end(trainer, None, None, None, 0)
- assert not list(tmp_path.iterdir())
- trainer.global_step = 2
- callback.on_train_batch_end(trainer, None, None, None, 1)
- path = next((tmp_path / "checkpoints").glob("*.ckpt"))
- proof = json.loads(path.with_suffix(".ckpt.recovery.json").read_text())
- assert (
- proof["global_step"] == 2
- and proof["full_state"]
- and proof["bytes"] == path.stat().st_size
- )
- callback.on_train_batch_end(trainer, None, None, None, 1)
- with pytest.raises(FileExistsError, match="overwrite"):
- EarlyRecoveryCheckpoint(str(tmp_path), 2).on_train_batch_end(
- trainer, None, None, None, 1
- )
diff --git a/tests/test_offline_constructor_inputs.py b/tests/test_offline_constructor_inputs.py
deleted file mode 100644
index 705391cf92..0000000000
--- a/tests/test_offline_constructor_inputs.py
+++ /dev/null
@@ -1,119 +0,0 @@
-"""Structural fixtures cannot become learned initialization or run evidence."""
-
-import json
-
-import pytest
-import torch
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-
-from scripts.audit_components import offline_static_buffer_construction
-from scripts.audit_constructor_inputs import evaluator_inputs, model_inputs
-from scripts.audit_hydra_configs import CONFIGS, compose_for_audit
-
-
-def test_native_fixture_is_explicit_and_does_not_read_runtime_snapshot(monkeypatch):
- monkeypatch.setattr(
- "egomimic.benchmarks.libero.native_artifact_identity._snapshot",
- lambda *args: pytest.fail("offline constructor read a runtime snapshot"),
- )
- path = (
- CONFIGS
- / "experiment/sim/libero_historical/action_flow_libero10_h240_euler50_dithalf_80k_s42.yaml"
- )
- with compose_for_audit(path) as cfg:
- before = OmegaConf.to_container(cfg.evaluator, resolve=True)
- inputs, evidence = evaluator_inputs(cfg.evaluator)
- evaluator = instantiate(inputs)
- assert evaluator.shared_diagnostics.max_samples == 8
- assert evaluator.shared_diagnostics.activation_layer_map == tuple(
- (i, i) for i in range(12)
- )
- assert evidence["external_binding"].startswith("UNVERIFIED")
- assert OmegaConf.to_container(cfg.evaluator, resolve=True) == before
- invalid = OmegaConf.create(before)
- invalid.artifact_identity._target_ = "not.a.canonical.factory"
- with pytest.raises(ValueError, match="canonical descriptors"):
- evaluator_inputs(invalid)
-
-
-@pytest.mark.parametrize(
- ("recipe", "dimension", "horizon"),
- [
- ("libero_oatpolicy", 7, 32),
- # Use the declared ARC recipe geometry, not the native-action horizon.
- ("libero_arc_oatpolicy", 12, 32),
- ],
-)
-def test_meta_tokenizer_comes_from_declared_template_and_scope_restores(
- recipe, dimension, horizon
-):
- from egomimic.models.oat import factory
-
- original = factory.load_tokenizer
- with compose_for_audit(CONFIGS / f"experiment/sim/oat/{recipe}.yaml") as cfg:
- before = json.dumps(OmegaConf.to_container(cfg, resolve=True), sort_keys=True)
- with (
- torch.device("meta"),
- offline_static_buffer_construction(),
- model_inputs(cfg, CONFIGS) as evidence,
- ):
- tokenizer = factory.load_tokenizer(None)
- assert tokenizer.decoder.sample_dim == dimension
- assert tokenizer.decoder.sample_horizon == horizon
- assert all(
- p.device.type == "meta" and not p.requires_grad
- for p in tokenizer.parameters()
- )
- assert evidence["external_binding"].startswith("UNVERIFIED")
- assert not hasattr(tokenizer, "_training_data_context")
- if recipe == "libero_arc_oatpolicy":
- from egomimic.pipeline.stages_libero_arc import LiberoArcStage
-
- codec = instantiate(cfg.model.pipeline.stages[0])
- assert isinstance(codec, LiberoArcStage)
- assert (
- tokenizer._training_input_representation
- == codec.representation_context()
- )
- else:
- assert tokenizer._training_input_representation is None
- assert factory.load_tokenizer is original
-
- assert (
- json.dumps(OmegaConf.to_container(cfg, resolve=True), sort_keys=True)
- == before
- )
- with pytest.raises(ValueError, match="trained OAT tokenizer"):
- factory.load_tokenizer(None)
-
-
-def test_initializer_rejects_cpu_use_and_explicit_checkpoint_and_restores():
- from egomimic.models.oat import factory
-
- original = factory.load_tokenizer
- with compose_for_audit(CONFIGS / "experiment/sim/oat/libero_oatpolicy.yaml") as cfg:
- with pytest.raises(ValueError, match="restricted to meta"):
- with model_inputs(cfg, CONFIGS):
- factory.load_tokenizer(None)
- assert factory.load_tokenizer is original
- cfg.benchmark.tokenizer_checkpoint = "/real/run/tokenizer.ckpt"
- with pytest.raises(ValueError, match="explicit learned checkpoint"):
- with model_inputs(cfg, CONFIGS):
- pytest.fail("explicit checkpoint replaced")
- assert factory.load_tokenizer is original
-
-
-def test_arc_structural_metadata_preserves_real_values_rng_and_method():
- from egomimic.pipeline.stages_libero_arc import LiberoArcStage
-
- original = LiberoArcStage.representation_context
- with torch.device("cpu"):
- expected = LiberoArcStage(num_waypoints=32, horizon=32).representation_context()
- rng = torch.get_rng_state().clone()
- with torch.device("meta"), offline_static_buffer_construction():
- codec = LiberoArcStage(num_waypoints=32, horizon=32)
- assert codec.action_scale.device.type == "meta"
- assert codec.representation_context() == expected
- assert LiberoArcStage.representation_context is original
- assert torch.equal(torch.get_rng_state(), rng)
diff --git a/tests/test_offline_static_buffer_audit.py b/tests/test_offline_static_buffer_audit.py
deleted file mode 100644
index a45daddc55..0000000000
--- a/tests/test_offline_static_buffer_audit.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import pytest
-import torch
-
-from egomimic.models.oat.tokenizer.oat.quantizer.fsq import FSQ
-from scripts.audit_components import offline_static_buffer_construction
-
-
-@pytest.mark.parametrize("levels", [(8, 5, 5, 5), (3, 4)])
-def test_static_lookup_buffers_match_real_constructor_and_rng(levels):
- original = FSQ.__init__
- baseline = FSQ(list(levels))
- before = torch.get_rng_state().clone()
- with offline_static_buffer_construction():
- candidate = FSQ(list(levels))
- with torch.device("meta"):
- meta = FSQ(list(levels))
- assert FSQ.__init__ is original
- assert torch.equal(before, torch.get_rng_state())
- assert candidate.codebook_size == meta.codebook_size == baseline.codebook_size
- assert not tuple(candidate.parameters())
- for name, value in baseline.named_buffers():
- actual = dict(candidate.named_buffers())[name]
- metadata = dict(meta.named_buffers())[name]
- assert torch.equal(value, actual)
- assert metadata.device.type == "meta"
- assert metadata.shape == value.shape and metadata.dtype == value.dtype
-
-
-def test_scope_restores_constructor_on_failure():
- original = FSQ.__init__
- linspace = torch.linspace
- with pytest.raises(RuntimeError, match="fixture failure"):
- with offline_static_buffer_construction():
- raise RuntimeError("fixture failure")
- assert FSQ.__init__ is original
- assert torch.linspace is linspace
-
-
-def test_full_tokenizer_layout_and_scalar_schedule_scope():
- from egomimic.models.oat.factory import make_tokenizer
-
- baseline = make_tokenizer()
- before = torch.get_rng_state().clone()
- with offline_static_buffer_construction(), torch.device("meta"):
- meta = make_tokenizer()
- assert torch.linspace(0, 1, 3).device.type == "meta"
- assert torch.equal(before, torch.get_rng_state())
- assert baseline.state_dict().keys() == meta.state_dict().keys()
- for name, value in baseline.state_dict().items():
- candidate = meta.state_dict()[name]
- assert candidate.shape == value.shape and candidate.dtype == value.dtype
diff --git a/tests/test_original_size_action_flow_cotrain_config.py b/tests/test_original_size_action_flow_cotrain_config.py
deleted file mode 100644
index 9b73022548..0000000000
--- a/tests/test_original_size_action_flow_cotrain_config.py
+++ /dev/null
@@ -1,72 +0,0 @@
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-
-CONFIG_DIR = Path(__file__).parents[1] / "egomimic" / "hydra_configs"
-ROW = "parents/action_flow_cotrain_uc_manual4919"
-
-
-def test_original_size_cotrain_contract(monkeypatch):
- monkeypatch.setenv("PUSHSHAPES_USOCKET_ROOT", "/verified/usocket")
- monkeypatch.setenv("PUSHSHAPES_CHAIN_GRIPPER_ROOT", "/verified/chain")
- with initialize_config_dir(version_base=None, config_dir=str(CONFIG_DIR.resolve())):
- cfg = compose(
- config_name="train_zarr_cartesian",
- # Component compatibility fixture, not a selectable training recipe.
- # Keep the historical schedule assertions explicit after leaf retirement.
- overrides=[
- f"+experiment=sim/pusht/{ROW}",
- "++paths.root_dir=.",
- "trainer.max_steps=150000",
- "trainer.val_check_interval=10000",
- "callbacks.model_checkpoint.every_n_train_steps=30000",
- ],
- )
-
- stages = cfg.model.pipeline.stages
- encoder = next(
- s for s in stages if s._target_.endswith("RoutedContentEncoderStage")
- )
- decoder = next(
- s for s in stages if s._target_.endswith("RoutedContentDecoderStage")
- )
- fields = [s for s in stages if s._target_.endswith("ConditionalVelocityStage")]
- field = fields[0]
- objective = stages[-1]
- routes = {"pushshapes_sim_u_socket", "pushshapes_sim_chain_gripper"}
-
- assert cfg.model._target_ == "egomimic.pl_utils.pl_model.ModelWrapper"
- assert cfg.model.training_behavior._target_.endswith("ActionFlowTrainingBehavior")
- assert cfg.model.hidden_dim == 384
- assert cfg.model.flow_samples_per_content == 14
- assert cfg.model.flow_mini_batch == 14
- assert cfg.model.cfg_scale == pytest.approx(4.0)
- assert cfg.model.num_inference_steps == 50
- assert field.inference_method == "dopri5"
- assert field.field.backbone.hidden_dim == 384
- assert field.field.backbone.depth == 12
- assert field.field.backbone.num_heads == 12
- assert len(fields) == 1
- assert set(encoder.encoders) == routes
- assert set(decoder.decoders) == routes
- assert encoder.encoders.pushshapes_sim_u_socket.action_dim == 4
- assert encoder.encoders.pushshapes_sim_chain_gripper.action_dim == 6
- assert decoder.decoders.pushshapes_sim_u_socket.action_dim == 4
- assert decoder.decoders.pushshapes_sim_chain_gripper.action_dim == 6
- assert objective.flow_aggregation == "sum_samples"
- assert objective.action_velocity_weight == pytest.approx(1.0)
- assert cfg.model.optimizer.lr == pytest.approx(1e-4)
- assert cfg.model.scheduler.base_lr_1 == pytest.approx(1e-4)
- assert cfg.model.scheduler.base_lr_2 == pytest.approx(5e-5)
- assert cfg.model.scheduler.final_lr == pytest.approx(5e-5)
- assert cfg.trainer.max_steps == 150_000
- assert cfg.trainer.val_check_interval == 10_000
- assert cfg.callbacks.model_checkpoint.every_n_train_steps == 30_000
- assert cfg.data.train_dataloader_params.pushshapes_sim_u_socket.batch_size == 32
- assert (
- cfg.data.train_dataloader_params.pushshapes_sim_chain_gripper.batch_size == 32
- )
- assert cfg.run_provenance.objective.domain_aggregation == (
- "equal_mean_from_one_batch32_per_domain"
- )
diff --git a/tests/test_osmo_gate.py b/tests/test_osmo_gate.py
deleted file mode 100644
index f87599913f..0000000000
--- a/tests/test_osmo_gate.py
+++ /dev/null
@@ -1,122 +0,0 @@
-"""Exercise the same five-update/save/strict-reload harness before OSMO allocation."""
-
-import hashlib
-import io
-import json
-
-import pytest
-from omegaconf import OmegaConf, open_dict
-
-from scripts.integration import run_gate
-from scripts.integration.prepare_inputs import copy_stream, relative_path
-from tests.fixtures.synthetic_episodes import write_episode
-from tests.test_retained_recipe_steps import small_cpu_model
-
-
-def test_input_staging_checks_hashes_and_never_replaces_existing_data(tmp_path):
- path = tmp_path / "verified.bin"
- value = b"immutable-input"
- sha = hashlib.sha256(value).hexdigest()
- receipt = copy_stream(
- io.BytesIO(value), path, expected_bytes=len(value), expected_sha256=sha
- )
- assert receipt["sha256"] == sha and path.read_bytes() == value
- with pytest.raises(FileExistsError):
- copy_stream(
- io.BytesIO(value), path, expected_bytes=len(value), expected_sha256=sha
- )
- assert path.read_bytes() == value
- with pytest.raises(ValueError, match="SHA-256"):
- copy_stream(
- io.BytesIO(b"bad"),
- tmp_path / "changed.bin",
- expected_bytes=3,
- expected_sha256=sha,
- )
- assert not (tmp_path / "changed.bin").exists()
- with pytest.raises(ValueError, match="relative"):
- relative_path("../../outside")
-
-
-def harness_inputs(monkeypatch, tmp_path):
- inputs, output = tmp_path / "inputs", tmp_path / "result"
- output.mkdir()
- data = inputs / "data/eva"
- episodes = []
- for index in range(3):
- write_episode(data, "eva", T=8, H=32, W=32, seed=index)
- episodes.append(
- {
- "source": "eva",
- "split": "train" if index < 2 else "valid",
- "episode_hash": f"eva_{index:02}",
- }
- )
- (inputs / "data-manifest.json").write_text(json.dumps({"episodes": episodes}))
- (inputs / "input-receipt.json").write_text('{"scope":"synthetic CPU fixture"}')
- monkeypatch.setenv("SMOKE_INPUT_ROOT", str(inputs))
- matrix = OmegaConf.load(run_gate.ROOT / "scripts/integration/matrix.yaml")
- matrix.trainer.accelerator = "cpu"
- matrix.trainer.precision = "32-true"
- # Tiny 32px fixtures collapse ResNet's last feature map to 1x1; BatchNorm
- # still needs two samples. Real GPU cases retain their full image sizes.
- matrix.cases["hpt-eva"].batch_size = 2
- matrix_path = tmp_path / "matrix.yaml"
- OmegaConf.save(matrix, matrix_path)
- return inputs, output, matrix_path
-
-
-@pytest.mark.parametrize("export", [True, False])
-def test_wrong_target_frame_rejected_before_network_construction(
- monkeypatch, tmp_path, export
-):
- import egomimic.trainHydra as entrypoint
-
- inputs, output, matrix_path = harness_inputs(monkeypatch, tmp_path)
- cfg, _, _ = run_gate.configured_case(matrix_path, "hpt-eva", inputs, output)
- cfg.inference_config.enabled = export
- for split in ("train_datasets", "valid_datasets"):
- ds = cfg.data[split].eva_bimanual
- with open_dict(ds):
- ds.bounds_check = False
- ds.resolver.transform_list.coord_frame = "eef_frame"
-
- def forbidden(*args, **kwargs):
- raise AssertionError("Network was constructed before frame preflight")
-
- monkeypatch.setattr(entrypoint, "_instantiate_model_wrapper", forbidden)
- with pytest.raises(ValueError, match="coord_frame.*camframe.*eef_frame"):
- entrypoint.train(cfg)
- assert not (output / "checkpoints/inference-config.yaml").exists()
-
-
-def test_training_harness_runs_five_updates_validation_and_bound_reload(
- monkeypatch, tmp_path
-):
- inputs, output, matrix_path = harness_inputs(monkeypatch, tmp_path)
- original = run_gate.configured_case
-
- def configured(*args, **kwargs):
- cfg, _, spec = original(*args, **kwargs)
- cfg = small_cpu_model(cfg, "hpt")
- with open_dict(cfg):
- cfg.model.pipeline.device = "cpu"
- for split in ("train_datasets", "valid_datasets"):
- cfg.data[split].eva_bimanual.bounds_check = False
- saved = OmegaConf.masked_copy(cfg, [key for key in cfg if key != "hydra"])
- saved = OmegaConf.create(OmegaConf.to_container(saved, resolve=True))
- return cfg, saved, spec
-
- monkeypatch.setattr(run_gate, "configured_case", configured)
- run_gate.training_gate(matrix_path, "hpt-eva", inputs, output)
- receipt = json.loads((output / "training-receipt.json").read_text())
- assert receipt["steps"] == 5 and len(receipt["gradients"]) == 5
- assert receipt["gpu"] is None
- assert (
- receipt["validation"]
- and receipt["initial_parameter_probe"] != receipt["final_parameter_probe"]
- )
- run_gate.verification_gate(output, device="cpu")
- verified = json.loads((output / "inference-receipt.json").read_text())
- assert verified["checkpoint_sha256"] == receipt["checkpoint_sha256"]
- assert [case["executed_shape"][1] for case in verified["cases"]] == [1, 2]
diff --git a/tests/test_pi_source_patch.py b/tests/test_pi_source_patch.py
deleted file mode 100644
index 9cf7d3b088..0000000000
--- a/tests/test_pi_source_patch.py
+++ /dev/null
@@ -1,50 +0,0 @@
-"""Pinned OpenPI patching must never modify another environment's hardlinks."""
-
-import hashlib
-import os
-
-import pytest
-
-from scripts.install_pi05_source import apply_transformers_patch
-
-
-def test_verified_patch_preserves_cache_hardlink_and_original(tmp_path):
- cache = tmp_path / "cached.py"
- cache.write_bytes(b"original code")
- root = tmp_path / "transformers"
- root.mkdir()
- target = root / "model.py"
- os.link(cache, target)
- content = b"pinned patched code"
- manifest = {"model.py": hashlib.sha256(content).hexdigest()}
- receipt = apply_transformers_patch(
- root, manifest, lambda path: content, tmp_path / "backups"
- )
- assert target.read_bytes() == content
- assert cache.read_bytes() == b"original code"
- assert target.stat().st_ino != cache.stat().st_ino
- from pathlib import Path
-
- assert Path(receipt[0]["backup"]).read_bytes() == b"original code"
- assert (
- apply_transformers_patch(
- root, manifest, lambda path: content, tmp_path / "backups"
- )[0]["sha256"]
- == receipt[0]["sha256"]
- )
-
-
-def test_every_hash_is_preflighted_before_patch_mutation(tmp_path):
- root = tmp_path / "transformers"
- root.mkdir()
- (root / "a.py").write_bytes(b"original")
- expected = hashlib.sha256(b"valid").hexdigest()
- with pytest.raises(ValueError, match="hash mismatch"):
- apply_transformers_patch(
- root,
- {"a.py": expected, "b.py": "wrong"},
- lambda path: b"valid",
- tmp_path / "backups",
- )
- assert (root / "a.py").read_bytes() == b"original"
- assert not (root / "b.py").exists()
diff --git a/tests/test_pipeline_only_surface.py b/tests/test_pipeline_only_surface.py
deleted file mode 100644
index f183f89b74..0000000000
--- a/tests/test_pipeline_only_surface.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import tomllib
-from pathlib import Path
-
-ROOT = Path(__file__).parents[1]
-
-
-def test_legacy_first_party_policy_surface_is_removed():
- removed = [
- "egomimic/models/act_nets.py",
- "egomimic/models/hpt_nets.py",
- "egomimic/models/preprocess_pi_obs.py",
- "egomimic/models/denoising_policy.py",
- "egomimic/models/fm_policy.py",
- "egomimic/eval/eval_act.py",
- "egomimic/eval/eval_hpt.py",
- "egomimic/eval/eval_pi.py",
- "egomimic/eval/eval_latent.py",
- "egomimic/eval/eval_video.py",
- "egomimic/eval/latent_dataset.py",
- "egomimic/utils/hydra_resolvers.py",
- "egomimic/utils/scheduler_utils.py",
- "egomimic/utils/tensor_utils.py",
- "external/openpi",
- ]
- leftovers = [relative for relative in removed if (ROOT / relative).exists()]
- assert leftovers == []
- assert list((ROOT / "egomimic/algo").glob("*.py")) == []
-
-
-def test_surviving_model_configs_are_pipeline_only():
- config_dir = ROOT / "egomimic/hydra_configs/model"
- configs = sorted(config_dir.rglob("*.yaml"))
- for config in configs:
- text = config.read_text()
- assert "egomimic.algo." not in text, config
- assert "robomimic_model" not in text, config
-
-
-def test_distribution_discovers_only_the_egomimic_package():
- project = tomllib.loads((ROOT / "pyproject.toml").read_text())
- package_find = project["tool"]["setuptools"]["packages"]["find"]
- assert package_find["include"] == ["egomimic*"]
- assert not (ROOT / "setup.py").exists()
-
-
-def test_training_entrypoint_has_one_strict_evaluation_path():
- entrypoint = (ROOT / "egomimic/trainHydra.py").read_text()
- assert 'hasattr(eval_obj, "run")' not in entrypoint
- assert 'hasattr(eval_obj, "bind_data_context")' not in entrypoint
- assert "strict=False" not in entrypoint
- assert "register_new_resolver" not in entrypoint
diff --git a/tests/test_pusht_parent_cleanup.py b/tests/test_pusht_parent_cleanup.py
deleted file mode 100644
index 01250053c1..0000000000
--- a/tests/test_pusht_parent_cleanup.py
+++ /dev/null
@@ -1,86 +0,0 @@
-"""Keep every PushT recipe's science identical when deduplicating its parents."""
-
-import hashlib
-import json
-from pathlib import Path
-
-import pytest
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-
-from scripts.audit_hydra_configs import CONFIGS, compose_for_audit
-
-BASELINE = json.loads(
- (
- Path(__file__).parent / "fixtures/pusht_config_before_parent_cleanup.json"
- ).read_text()
-)
-
-
-@pytest.mark.parametrize(
- "parent",
- [
- "episode_split",
- "action_flow_uc_h16",
- "action_flow_points6_uc_h16",
- "standard_dp_uc_h16",
- ],
-)
-def test_shared_parents_construct_their_data_boundary(parent):
- with compose_for_audit(CONFIGS / "data/pusht/parents" / f"{parent}.yaml") as cfg:
- data = instantiate(cfg.data, _recursive_=False)
- assert data.preflight_configuration()
-
-
-MOVE_BASELINE = json.loads(
- (Path(__file__).parent / "fixtures/pusht_before_real_sim_cleanup.json").read_text()
-)
-
-
-@pytest.mark.parametrize("recipe,expected", sorted(MOVE_BASELINE["recipes"].items()))
-def test_existing_pusht_recipe_keeps_exact_configuration(recipe, expected):
- with compose_for_audit(CONFIGS / "experiment/sim/pusht" / f"{recipe}.yaml") as cfg:
- actual = OmegaConf.to_container(cfg, resolve=False)
- actual.pop("hydra", None) # A parent's selection path intentionally changed.
- encoded = json.dumps(actual, sort_keys=True, separators=(",", ":")).encode()
- assert hashlib.sha256(encoded).hexdigest() == expected
-
-
-def test_retired_parent_paths_have_no_selectable_config_references():
- for group, retired in (
- ("data", BASELINE["retired_data"]),
- ("experiment", BASELINE["retired_recipes"]),
- ):
- for old in retired:
- assert not (CONFIGS / group / "pusht" / f"{old}.yaml").exists()
- for path in CONFIGS.rglob("*.yaml"):
- assert f"pusht/{old}" not in path.read_text()
- assert not (
- CONFIGS / "data/pusht/planar_usocket_chain_manual4919_standard_retimed.yaml"
- ).exists()
-
-
-@pytest.mark.parametrize(
- "recipe",
- [
- "planar_uc_manual4919_dp_261m_af_obs_multiplier",
- "action_flow_cotrain_uc_multiplier_interpolation",
- ],
-)
-def test_selected_pair_still_has_2999_and_4919_with_two_full_batches(recipe):
- with compose_for_audit(CONFIGS / "experiment/sim/pusht" / f"{recipe}.yaml") as cfg:
- for group in ("train_datasets", "valid_datasets"):
- datasets = cfg.data[group]
- assert (
- datasets.pushshapes_sim_u_socket.resolver.expected_episode_count == 2999
- )
- assert (
- datasets.pushshapes_sim_chain_gripper.resolver.expected_episode_count
- == 4919
- )
- assert all(
- d.valid_ratio == 0.01 and d.split_seed == 42 for d in datasets.values()
- )
- assert all(
- d.batch_size == 32 for d in cfg.data.train_dataloader_params.values()
- )
diff --git a/tests/test_rank_synthetic_trajectory_nn.py b/tests/test_rank_synthetic_trajectory_nn.py
deleted file mode 100644
index f22602274e..0000000000
--- a/tests/test_rank_synthetic_trajectory_nn.py
+++ /dev/null
@@ -1,207 +0,0 @@
-import hashlib
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-
-SCRIPT = (
- Path(__file__).parents[1]
- / "scripts"
- / "experiments"
- / "rank_synthetic_trajectory_nn.py"
-)
-
-
-def _matrix(tmp_path, *, variants=("a", "b"), seeds=(42, 43, 44), particles=16):
- rng = np.random.default_rng(7)
- targets = rng.normal(size=(particles, 3)).astype(np.float32)
- dataset = tmp_path / "eval.npz"
- np.savez_compressed(
- dataset,
- split=np.ones(particles, dtype=np.int64),
- target_3d=targets,
- )
- config_dir = tmp_path / "config"
- config_dir.mkdir()
- trajectory_root = tmp_path / "trajectories"
- trajectory_root.mkdir()
- configs = []
- trajectories = {}
- for variant in variants:
- for seed in seeds:
- run_id = f"run-{variant}-{seed}"
- output = tmp_path / "runs" / run_id
- (output / "checkpoints").mkdir(parents=True)
- (output / "summary.json").write_text("{}\n")
- checkpoint = (
- output / "checkpoints" / "epoch-equivalent-000001-global-step-000010.pt"
- )
- checkpoint.write_bytes(f"checkpoint-{variant}-{seed}".encode())
- config = {
- "variant": variant,
- "seed": seed,
- "evaluation_dataset": str(dataset),
- "output_dir": str(output),
- "max_steps": 10,
- "wandb": {"id": run_id},
- }
- config_path = config_dir / f"{run_id}.json"
- config_path.write_text(json.dumps(config))
- configs.append(str(config_path))
- final = targets if variant == "a" else targets + 5.0
- points = np.stack([np.zeros_like(targets), final])
- trajectory = trajectory_root / f"{run_id}.npz"
- np.savez_compressed(
- trajectory,
- times=np.array([0.0, 1.0], dtype=np.float32),
- points=points,
- target=targets,
- )
- trajectories[run_id] = trajectory
- manifest = {
- "schema_version": 1,
- "variants": list(variants),
- "seeds": list(seeds),
- "datasets": {
- "evaluation": {
- "path": str(dataset),
- "sha256": hashlib.sha256(dataset.read_bytes()).hexdigest(),
- }
- },
- "configs": configs,
- }
- manifest_path = config_dir / "manifest.json"
- manifest_path.write_text(json.dumps(manifest))
- return manifest_path, dataset, trajectory_root, trajectories, targets
-
-
-def _command(manifest, output_json, output_markdown, particles):
- return [
- sys.executable,
- str(SCRIPT),
- "--manifest",
- str(manifest),
- "--expected-particles",
- str(particles),
- "--distance-chunk-size",
- "5",
- "--output-json",
- str(output_json),
- "--output-markdown",
- str(output_markdown),
- ]
-
-
-def test_ranks_complete_matrix_and_records_exact_provenance(tmp_path):
- manifest, dataset, root, trajectories, _ = _matrix(tmp_path)
- output_json = tmp_path / "ranking.json"
- output_markdown = tmp_path / "ranking.md"
- subprocess.run(
- _command(manifest, output_json, output_markdown, 16)
- + ["--trajectory-root", str(root)],
- check=True,
- )
-
- result = json.loads(output_json.read_text())
- assert [row["variant"] for row in result["ranking"]] == ["a", "b"]
- assert result["ranking"][0]["mean"] == 0.0
- assert result["ranking"][0]["sample_std"] == 0.0
- assert [row["seed"] for row in result["ranking"][0]["per_seed"]] == [
- 42,
- 43,
- 44,
- ]
- assert result["comparison"]["particles"] == 16
- assert (
- result["comparison"]["dataset"]["sha256"]
- == hashlib.sha256(dataset.read_bytes()).hexdigest()
- )
- run = next(row for row in result["runs"] if row["run_id"] == "run-a-42")
- assert (
- run["trajectory"]["sha256"]
- == hashlib.sha256(trajectories["run-a-42"].read_bytes()).hexdigest()
- )
- assert run["checkpoint"]["global_step"] == 10
- assert len(run["checkpoint"]["sha256"]) == 64
- markdown = output_markdown.read_text()
- assert "Symmetric NN MSE" in markdown
- assert "42: 0" in markdown
-
-
-def test_accepts_only_a_complete_set_of_exact_npz_paths(tmp_path):
- manifest, _, _, trajectories, _ = _matrix(
- tmp_path, variants=("a",), seeds=(42, 43), particles=8
- )
- output_json = tmp_path / "ranking.json"
- output_markdown = tmp_path / "ranking.md"
- explicit = []
- for run_id, path in trajectories.items():
- explicit.extend(["--trajectory-npz", f"{run_id}={path}"])
- subprocess.run(
- _command(manifest, output_json, output_markdown, 8) + explicit,
- check=True,
- )
- assert len(json.loads(output_json.read_text())["runs"]) == 2
-
- incomplete = subprocess.run(
- _command(manifest, output_json, output_markdown, 8)
- + ["--trajectory-npz", f"run-a-42={trajectories['run-a-42']}"],
- text=True,
- capture_output=True,
- )
- assert incomplete.returncode != 0
- assert "explicit trajectory set does not match complete matrix" in incomplete.stderr
-
-
-def test_rejects_target_or_time_grid_mismatch(tmp_path):
- manifest, _, root, trajectories, targets = _matrix(
- tmp_path, variants=("a",), seeds=(42, 43), particles=8
- )
- mismatched = trajectories["run-a-43"]
- np.savez_compressed(
- mismatched,
- times=np.array([0.0, 0.4, 1.0], dtype=np.float32),
- points=np.stack([targets, targets, targets]),
- target=targets,
- )
- result = subprocess.run(
- _command(manifest, tmp_path / "ranking.json", tmp_path / "ranking.md", 8)
- + ["--trajectory-root", str(root)],
- text=True,
- capture_output=True,
- )
- assert result.returncode != 0
- assert "time grid differs" in result.stderr
-
- np.savez_compressed(
- mismatched,
- times=np.array([0.0, 1.0], dtype=np.float32),
- points=np.stack([targets, targets]),
- target=targets + 1.0,
- )
- result = subprocess.run(
- _command(manifest, tmp_path / "ranking.json", tmp_path / "ranking.md", 8)
- + ["--trajectory-root", str(root)],
- text=True,
- capture_output=True,
- )
- assert result.returncode != 0
- assert "does not exactly match evaluation dataset" in result.stderr
-
-
-def test_rejects_run_without_terminal_completion_artifacts(tmp_path):
- manifest, _, root, _, _ = _matrix(
- tmp_path, variants=("a",), seeds=(42,), particles=8
- )
- config = json.loads(next((tmp_path / "config").glob("run-*.json")).read_text())
- (Path(config["output_dir"]) / "summary.json").unlink()
- result = subprocess.run(
- _command(manifest, tmp_path / "ranking.json", tmp_path / "ranking.md", 8)
- + ["--trajectory-root", str(root)],
- text=True,
- capture_output=True,
- )
- assert result.returncode != 0
- assert "run is incomplete because summary.json is missing" in result.stderr
diff --git a/tests/test_speed_retiming_diagnostic.py b/tests/test_speed_retiming_diagnostic.py
deleted file mode 100644
index ba1f02022d..0000000000
--- a/tests/test_speed_retiming_diagnostic.py
+++ /dev/null
@@ -1,123 +0,0 @@
-import numpy as np
-import pytest
-import torch
-
-from egomimic.pipeline.stages_speed import SharedSpeedCondition
-from egomimic.rldb.zarr.planar_retiming import PlanarCommandRetiming
-
-
-@pytest.mark.parametrize("view,rate", list(enumerate((1, 1.25, 1.5, 1.75, 2))))
-def test_constant_speed_native_retiming(view, rate):
- t = PlanarCommandRetiming((1, 1.25, 1.5, 1.75, 2))
- native = np.stack(
- (np.arange(31), np.zeros(31), np.zeros(31), np.linspace(0, 1, 31)), axis=-1
- )
- obs = np.array([5, 6, 7])
- b = t.transform({"actions": native, "_retiming_view": view, "obs": obs})
- np.testing.assert_allclose(b["actions"][:, 0], np.arange(16) * rate)
- np.testing.assert_allclose(b["actions"][:, 3], np.arange(16) * rate / 30)
- np.testing.assert_allclose(b["retiming_rate"], [rate])
- assert "requested_speed" not in b
- assert b["obs"] is obs
- assert np.array_equal(native[:, 0], np.arange(31))
-
-
-def test_wrap_stationary_and_tail_rejection():
- t = PlanarCommandRetiming((1.5,), horizon=3)
- a = np.zeros((4, 3))
- a[:, 2] = [3.0, 3.1, -3.1, -3.0]
- b = t.transform({"actions": a, "_retiming_view": 0})
- assert abs(abs(b["actions"][1, 2]) - np.pi) < 1e-6
- assert b["retiming_rate"].item() == 1.5
- assert "requested_speed" not in b
- with pytest.raises(ValueError, match="unpadded"):
- t.transform({"actions": a[:3], "_retiming_view": 0})
-
-
-@pytest.mark.parametrize("encoding", ("scalar", "fourier"))
-def test_condition_modes_rng_and_learning(encoding):
- torch.manual_seed(123)
- rng = torch.get_rng_state().clone()
- stage = SharedSpeedCondition(None, encoding)
- assert torch.equal(rng, torch.get_rng_state())
- assert stage.contract("train") == stage.contract("inference")
- cond = torch.zeros(2, 128)
- speed = torch.tensor([[0.25], [2.0]])
- out = stage({"condition": cond, "retiming_rate": speed})["speed_condition"]
- assert torch.equal(out, cond)
- opt = torch.optim.SGD(stage.parameters(), lr=0.01)
- for _ in range(2):
- opt.zero_grad()
- out = stage({"condition": cond, "retiming_rate": speed})["speed_condition"]
- (out - 1).square().mean().backward()
- opt.step()
- assert stage.mlp[0].weight.grad.abs().sum() > 0
- assert not torch.equal(out[0], out[1])
- with pytest.raises(KeyError):
- stage({"condition": cond})
- with pytest.raises(ValueError):
- stage({"condition": cond, "retiming_rate": -torch.ones(2, 1)})
-
-
-def test_unchanged_initialization_across_arms():
- weights = []
- for encoding in ("scalar", "fourier"):
- torch.manual_seed(42)
- before = torch.nn.Linear(10, 10)
- SharedSpeedCondition(None, encoding)
- after = torch.nn.Linear(10, 10)
- weights.append((before.weight.detach(), after.weight.detach()))
- assert all(torch.equal(a, b) for a, b in zip(*weights))
-
-
-def test_condition_initializer_never_reseeds_cuda(monkeypatch):
- def forbidden(*args, **kwargs):
- raise AssertionError("Condition initializer must not reseed CUDA")
-
- monkeypatch.setattr(torch, "manual_seed", forbidden)
- monkeypatch.setattr(torch.cuda, "manual_seed_all", forbidden)
- rng = torch.get_rng_state().clone()
- for encoding in ("scalar", "fourier"):
- SharedSpeedCondition(None, encoding)
- assert torch.equal(rng, torch.get_rng_state())
-
-
-@pytest.mark.parametrize("encoding", ("scalar", "fourier"))
-def test_real_graph_consumes_speed_in_both_modes(encoding):
- from egomimic.pipeline.core import Pipeline
- from egomimic.pipeline.stages_action_flow import (
- ConditionalVelocityStage,
- LatentBridgeStage,
- )
-
- class Field(torch.nn.Module):
- def forward(self, x, t, condition, **kwargs):
- return condition[:, : x.shape[-1]].unsqueeze(1).expand_as(x)
-
- speed = SharedSpeedCondition(None, encoding, condition_dim=4)
- torch.nn.init.constant_(speed.mlp[-1].weight, 0.1)
- bridge = LatentBridgeStage(
- samples_per_content=1,
- condition_key="speed_condition",
- condition_dropout_probability=0,
- )
- velocity = ConditionalVelocityStage(
- Field(), num_inference_steps=2, inference_condition_key="speed_condition"
- )
- graph = Pipeline([speed, bridge, velocity])
- b = {
- "condition": torch.zeros(2, 4),
- "retiming_rate": torch.tensor([[0.25], [2.0]]),
- "action_flow/clean_latent": torch.zeros(2, 2, 4),
- "sampler/noise": torch.zeros(2, 2, 4),
- }
- for mode in ("train", "inference"):
- _, excluded = graph.plan(b.keys(), mode)
- assert all(missing == [""] for _, missing in excluded), excluded
- result = graph.execute(b, mode=mode)
- key = (
- "action_flow/predicted_velocity"
- if mode == "train"
- else "action_flow/generated_latent"
- )
- assert not torch.equal(result[key][0], result[key][1])
diff --git a/tests/test_standard_dp_launch_args.py b/tests/test_standard_dp_launch_args.py
deleted file mode 100644
index ea31518103..0000000000
--- a/tests/test_standard_dp_launch_args.py
+++ /dev/null
@@ -1,16 +0,0 @@
-"""The maintained DP helper cannot silently select a retired single-source recipe."""
-
-import runpy
-from pathlib import Path
-
-import pytest
-
-
-@pytest.mark.parametrize("profile", ["", "chain_manual4919_retimed_v1"])
-def test_retired_standard_dp_default_rejected(monkeypatch, profile):
- monkeypatch.delenv("DP_COTRAIN_STANDARD", raising=False)
- monkeypatch.setenv("DP_SINGLE_SOURCE_STANDARD", profile)
- root = Path(__file__).resolve().parents[1]
- recipe = runpy.run_path(str(root / "scripts/train/standard_dp_launch.py"))["recipe"]
- with pytest.raises(ValueError, match="retired"):
- recipe()
diff --git a/tests/test_sync_runtime_script.py b/tests/test_sync_runtime_script.py
deleted file mode 100644
index 49b15a0ee0..0000000000
--- a/tests/test_sync_runtime_script.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import subprocess
-from pathlib import Path
-
-SCRIPT = (
- Path(__file__).parents[1] / "scripts" / "clusters" / "common" / "sync_runtime.sh"
-)
-
-
-def test_sync_runtime_is_valid_bash_and_rejects_relative_paths():
- subprocess.run(["bash", "-n", str(SCRIPT)], check=True)
- result = subprocess.run(
- [str(SCRIPT), "relative/source", "/tmp/runtime"],
- check=False,
- capture_output=True,
- text=True,
- )
-
- assert result.returncode == 2
- assert "must be absolute paths" in result.stderr
-
-
-def test_sync_runtime_uses_frozen_lock_and_verifies_torchdiffeq():
- text = SCRIPT.read_text()
-
- assert 'uv sync --frozen --project "$repo"' in text
- assert 'version("torchdiffeq")' in text
- assert 'item.startswith("torchdiffeq==")' in text
- assert "/coc/" not in text
- assert "/storage/" not in text
diff --git a/tests/test_synthetic_action_adapter_flow.py b/tests/test_synthetic_action_adapter_flow.py
deleted file mode 100644
index e415ea383a..0000000000
--- a/tests/test_synthetic_action_adapter_flow.py
+++ /dev/null
@@ -1,234 +0,0 @@
-import pytest
-import torch
-
-from egomimic.synthetic.action_adapter_flow import SyntheticActionAdapterFlow
-
-
-def _perturb_residual_outputs(model: SyntheticActionAdapterFlow) -> None:
- with torch.no_grad():
- for adapter in (model.encoder, model.decoder):
- adapter.residual[-1].weight.normal_(std=0.02)
- adapter.residual[-1].bias.normal_(std=0.02)
-
-
-def test_fixed_lift_reconstructs_and_projects_unit_noise_exactly():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="fixed_affine")
- action = torch.randn(11, 3)
- torch.testing.assert_close(model.decoder(model.encoder(action)), action)
- torch.testing.assert_close(
- model.decoder.weight @ model.decoder.weight.T, torch.eye(3)
- )
- torch.testing.assert_close(model.scale_loss(torch.randn(16, 8)), torch.tensor(0.0))
- assert not any(parameter.requires_grad for parameter in model.encoder.parameters())
- assert not any(parameter.requires_grad for parameter in model.decoder.parameters())
-
-
-def test_affine_path_identity_equals_reconstruction():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="joint_affine")
- with torch.no_grad():
- model.encoder.weight.add_(0.05 * torch.randn_like(model.encoder.weight))
- model.decoder.weight.add_(0.05 * torch.randn_like(model.decoder.weight))
- model.encoder.bias.add_(0.02 * torch.randn_like(model.encoder.bias))
- model.decoder.bias.add_(0.02 * torch.randn_like(model.decoder.bias))
- action = torch.randn(13, 3)
- noise = torch.randn(13, 8)
- time = torch.rand(13, 1)
- torch.testing.assert_close(
- model.path_consistency_loss(action, noise, time),
- model.reconstruction_loss(action),
- rtol=2e-5,
- atol=2e-6,
- )
-
-
-def test_nonlinear_adapters_start_at_the_same_fixed_lift():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="nonlinear")
- action = torch.randn(9, 3)
- torch.testing.assert_close(model.decoder(model.encoder(action)), action)
-
-
-def test_nonlinear_decoder_jvp_matches_centered_finite_difference():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="nonlinear")
- _perturb_residual_outputs(model)
- state = torch.randn(7, 8)
- tangent = torch.randn(7, 8)
- epsilon = 1e-3
- finite_difference = (
- model.decoder(state + epsilon * tangent)
- - model.decoder(state - epsilon * tangent)
- ) / (2 * epsilon)
- torch.testing.assert_close(
- model.decoder_jvp(state, tangent),
- finite_difference,
- rtol=3e-3,
- atol=3e-4,
- )
-
-
-def test_flow_gradient_reaches_trainable_encoder_through_state_and_target():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="joint_affine")
- losses = model.losses(
- torch.randn(12, 3), objective="reconstruction", flow_samples=3
- )
- losses["flow_loss"].backward()
- assert all(parameter.grad is not None for parameter in model.encoder.parameters())
- assert (
- sum(
- float(parameter.grad.abs().sum())
- for parameter in model.encoder.parameters()
- )
- > 0
- )
-
-
-def test_loss_uses_the_supplied_fixed_noise_cloud_for_all_flow_samples():
- model = SyntheticActionAdapterFlow(
- latent_dim=8, adapter_family="fixed_affine", field_width=16, field_depth=2
- )
- action = torch.randn(5, 3)
- noise = torch.randn(5, 8)
- time = torch.rand(15, 1)
- losses = model.losses(
- action,
- objective="none",
- flow_samples=3,
- noise=noise,
- time=time,
- )
- clean = model.encoder(action)
- clean_many = clean[:, None].expand(-1, 3, -1).reshape(-1, 8)
- noise_many = noise[:, None].expand(-1, 3, -1).reshape(-1, 8)
- state = (1.0 - time) * clean_many + time * noise_many
- expected = (model.velocity(state, time) - (noise_many - clean_many)).square().mean()
- torch.testing.assert_close(losses["flow_loss"], expected)
-
-
-def test_clean_gradient_modes_keep_forward_loss_identical_and_change_encoder_route():
- model = SyntheticActionAdapterFlow(
- latent_dim=8, adapter_family="nonlinear", field_width=16, field_depth=2
- )
- action, noise, time = torch.randn(12, 3), torch.randn(12, 8), torch.rand(36, 1)
- losses = {}
- gradients = {}
- for mode in ("full", "target_stopgrad", "all_stopgrad"):
- model.zero_grad(set_to_none=True)
- losses[mode] = model.losses(
- action,
- objective="action_velocity",
- flow_samples=3,
- noise=noise,
- time=time,
- clean_gradient_mode=mode,
- )["flow_loss"]
- losses[mode].backward()
- gradients[mode] = sum(
- float(parameter.grad.abs().sum())
- for parameter in model.encoder.parameters()
- if parameter.grad is not None
- )
- torch.testing.assert_close(losses["full"], losses["target_stopgrad"])
- torch.testing.assert_close(losses["full"], losses["all_stopgrad"])
- assert gradients["full"] > 0
- assert gradients["target_stopgrad"] > 0
- assert gradients["all_stopgrad"] == 0
-
-
-def test_unknown_clean_gradient_mode_is_rejected():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="nonlinear")
- with pytest.raises(ValueError, match="unknown clean gradient mode"):
- model.losses(
- torch.randn(8, 3), objective="action_velocity", clean_gradient_mode="bad"
- )
-
-
-def test_path_loss_gradients_reach_both_nonlinear_adapters():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="nonlinear")
- _perturb_residual_outputs(model)
- losses = model.losses(torch.randn(10, 3), objective="path", flow_samples=2)
- losses["path_loss"].backward()
- for adapter in (model.encoder, model.decoder):
- gradients = [
- parameter.grad
- for parameter in adapter.parameters()
- if parameter.requires_grad
- ]
- assert any(
- gradient is not None and bool(gradient.abs().sum())
- for gradient in gradients
- )
-
-
-def test_action_velocity_loss_matches_explicit_decoder_jacobian_product():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="nonlinear")
- _perturb_residual_outputs(model)
- state = torch.randn(7, 8)
- residual = torch.randn(7, 8)
-
- jacobians = torch.func.vmap(torch.func.jacrev(model.decoder))(state)
- explicit_action_residual = torch.einsum("bai,bi->ba", jacobians, residual)
- expected = explicit_action_residual.square().sum(dim=-1).mean() / 3.0
-
- torch.testing.assert_close(model.action_velocity_loss(state, residual), expected)
-
-
-def test_action_velocity_objective_reaches_encoder_decoder_and_field():
- model = SyntheticActionAdapterFlow(
- latent_dim=8,
- adapter_family="nonlinear",
- field_width=16,
- field_depth=2,
- )
- _perturb_residual_outputs(model)
- losses = model.losses(
- torch.randn(10, 3),
- objective="action_velocity",
- flow_samples=2,
- lambda_reconstruction=0.7,
- lambda_scale=0.2,
- lambda_action_velocity=0.4,
- )
- losses["action_velocity_loss"].backward()
-
- for module in (model.encoder, model.decoder, model.field):
- gradients = [
- parameter.grad
- for parameter in module.parameters()
- if parameter.requires_grad
- ]
- assert any(
- gradient is not None and bool(gradient.abs().sum())
- for gradient in gradients
- )
-
-
-def test_action_velocity_objective_combines_reconstruction_scale_and_metric_loss():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="joint_affine")
- losses = model.losses(
- torch.randn(9, 3),
- objective="action_velocity",
- lambda_reconstruction=0.7,
- lambda_scale=0.2,
- lambda_action_velocity=0.4,
- )
- expected = (
- losses["flow_loss"]
- + 0.7 * losses["reconstruction_loss"]
- + 0.2 * losses["scale_loss"]
- + 0.4 * losses["action_velocity_loss"]
- )
- torch.testing.assert_close(losses["loss"], expected)
-
-
-def test_affine_path_objective_is_rejected_as_duplicate():
- model = SyntheticActionAdapterFlow(latent_dim=8, adapter_family="joint_affine")
- with pytest.raises(ValueError, match="duplicates reconstruction"):
- model.losses(torch.randn(8, 3), objective="path")
-
-
-def test_reverse_time_trajectory_decodes_every_state():
- model = SyntheticActionAdapterFlow(
- latent_dim=8, adapter_family="fixed_affine", field_width=16, field_depth=2
- )
- trajectory = model.trajectory(torch.randn(6, 8), steps=4)
- assert trajectory.shape == (5, 6, 3)
- assert torch.isfinite(trajectory).all()
diff --git a/tests/test_synthetic_multi_action_adapter_flow.py b/tests/test_synthetic_multi_action_adapter_flow.py
deleted file mode 100644
index 649c565748..0000000000
--- a/tests/test_synthetic_multi_action_adapter_flow.py
+++ /dev/null
@@ -1,125 +0,0 @@
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import torch
-
-from egomimic.synthetic.multi_action_adapter_flow import (
- SyntheticMultiActionAdapterFlow,
-)
-
-
-def test_private_adapters_and_shared_field_receive_expected_gradients():
- model = SyntheticMultiActionAdapterFlow(
- embodiments=["shallow", "steep"], field_width=16, field_depth=2
- )
- losses = model.losses_for_embodiment("shallow", torch.randn(8, 3), flow_samples=2)
- losses["loss"].backward()
- assert any(
- p.grad is not None and bool(p.grad.abs().sum())
- for p in model.field.parameters()
- )
- assert any(
- p.grad is not None and bool(p.grad.abs().sum())
- for p in model.encoders["shallow"].parameters()
- )
- assert any(
- p.grad is not None and bool(p.grad.abs().sum())
- for p in model.decoders["shallow"].parameters()
- )
- assert all(p.grad is None for p in model.encoders["steep"].parameters())
- assert all(p.grad is None for p in model.decoders["steep"].parameters())
-
-
-def test_private_decoders_share_one_generated_latent_path():
- model = SyntheticMultiActionAdapterFlow(
- embodiments=["shallow", "steep"], field_width=16, field_depth=2
- )
- noise = torch.randn(5, 8)
- shallow = model.trajectory(noise, embodiment="shallow", steps=3)
- steep = model.trajectory(noise, embodiment="steep", steps=3)
- assert shallow.shape == steep.shape == (4, 5, 3)
- torch.testing.assert_close(shallow, steep)
-
-
-def test_paraboloid_surface_metric_is_zero_on_surface():
- from egomimic.eval.synthetic_trajectory_eval import SyntheticTrajectoryEval
-
- xy = torch.randn(20, 2)
- curvature = 0.5
- points = torch.cat([xy, curvature * xy.square().sum(dim=-1, keepdim=True)], dim=-1)
- torch.testing.assert_close(
- SyntheticTrajectoryEval.paraboloid_surface_rmse(points, curvature=curvature),
- torch.tensor(0.0),
- )
-
-
-def test_multi_trainer_runs_optimizer_validation_and_immutable_checkpoint(tmp_path):
- repository = Path(__file__).parents[1]
- source = tmp_path / "source"
- source.mkdir()
- datasets = {}
- evaluations = {}
- rng = np.random.default_rng(5)
- for name, curvature in (("shallow", 0.125), ("steep", 0.5)):
- latent = rng.normal(size=(24, 8)).astype(np.float32)
- xy = rng.normal(size=(24, 2)).astype(np.float32)
- target = np.concatenate(
- [xy, curvature * np.square(xy).sum(axis=1, keepdims=True)], axis=1
- ).astype(np.float32)
- split = np.array([0] * 20 + [1] * 4, dtype=np.uint8)
- path = source / f"{name}.npz"
- np.savez_compressed(
- path,
- source_gaussian_latent=latent,
- target_3d=target,
- split=split,
- )
- datasets[name] = str(path)
- evaluations[name] = str(path)
- output = tmp_path / "run"
- config = {
- "seed": 42,
- "output_dir": str(output),
- "datasets": datasets,
- "evaluation_datasets": evaluations,
- "curvatures": {"shallow": 0.125, "steep": 0.5},
- "source_key": "source_gaussian_latent",
- "batch_size_per_embodiment": 4,
- "flow_samples": 2,
- "learning_rate": 3e-4,
- "max_steps": 1,
- "log_every": 1,
- "checkpoint_every": 1,
- "inference_steps": 2,
- "evaluation_particles": 4,
- "model": {
- "embodiments": ["shallow", "steep"],
- "latent_dim": 8,
- "residual_width": 8,
- "residual_depth": 1,
- "field_width": 8,
- "field_depth": 1,
- },
- }
- config_path = tmp_path / "config.json"
- config_path.write_text(json.dumps(config))
- subprocess.run(
- [
- sys.executable,
- str(repository / "scripts/train/train_synthetic_multi_action_flow.py"),
- "--config",
- str(config_path),
- ],
- check=True,
- )
- assert list(
- (output / "checkpoints").glob("epoch-equivalent-*-global-step-000001.pt")
- )
- summary = json.loads((output / "summary.json").read_text())
- assert set(summary["embodiments"]) == {"shallow", "steep"}
- assert summary["shared_field_parameters"] > 0
- for name in ("shallow", "steep"):
- assert (output / f"validation_trajectory_{name}.npz").is_file()
diff --git a/tests/test_synthetic_shared_latent_flow.py b/tests/test_synthetic_shared_latent_flow.py
deleted file mode 100644
index 7e51b47421..0000000000
--- a/tests/test_synthetic_shared_latent_flow.py
+++ /dev/null
@@ -1,55 +0,0 @@
-import torch
-
-from egomimic.synthetic.shared_latent_flow import (
- SyntheticDirectFlow,
- SyntheticSharedLatentFlow,
-)
-from scripts.train.train_synthetic_manifold import energy_distance
-
-
-def test_direct_flow_has_no_encoder_or_decoder_and_updates_field():
- model = SyntheticDirectFlow(data_dim=3, field_width=16, field_depth=2)
- assert not hasattr(model, "encoder")
- assert not hasattr(model, "decoder")
- source = torch.randn(8, 3)
- target = torch.randn(8, 3)
- losses = model.losses(source, target, flow_samples=2)
- losses["loss"].backward()
- assert set(losses) == {"loss", "flow_loss"}
- assert all(parameter.grad is not None for parameter in model.parameters())
- assert model.integrate(source, steps=2).shape == source.shape
-
-
-def test_energy_distance_is_zero_for_identical_empirical_distributions():
- points = torch.randn(12, 3)
- torch.testing.assert_close(energy_distance(points, points), torch.tensor(0.0))
-
-
-def test_both_objectives_are_finite_and_differentiable():
- for latent_dim in (2, 4):
- for method in ("unite", "vfm"):
- model = SyntheticSharedLatentFlow(latent_dim=latent_dim)
- losses = model.losses(
- torch.randn(8, latent_dim), torch.randn(8, 3), method=method
- )
- assert all(torch.isfinite(value) for value in losses.values())
- losses["loss"].backward()
- assert all(parameter.grad is not None for parameter in model.parameters())
-
-
-def test_vfm_reconstruction_can_freeze_only_field_parameters():
- model = SyntheticSharedLatentFlow(latent_dim=4)
- losses = model.losses(
- torch.randn(8, 4),
- torch.randn(8, 3),
- method="vfm",
- reconstruction_updates_field=False,
- )
- losses["reconstruction_loss"].backward(retain_graph=True)
- assert all(parameter.grad is None for parameter in model.field.parameters())
- assert all(parameter.grad is not None for parameter in model.encoder.parameters())
- assert all(parameter.grad is not None for parameter in model.decoder.parameters())
-
- model.zero_grad(set_to_none=True)
- losses["flow_loss"].backward()
- assert all(parameter.grad is not None for parameter in model.field.parameters())
diff --git a/tests/test_synthetic_trainer_resume.py b/tests/test_synthetic_trainer_resume.py
deleted file mode 100644
index f5eeedb45d..0000000000
--- a/tests/test_synthetic_trainer_resume.py
+++ /dev/null
@@ -1,188 +0,0 @@
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from scripts.train.train_synthetic_manifold import (
- _capture_rng_state,
- _restore_rng_state,
-)
-
-
-def _run(source: Path, config: Path, resume: Path | None = None) -> None:
- command = [
- sys.executable,
- str(source / "scripts/train/train_synthetic_manifold.py"),
- "--config",
- str(config),
- ]
- if resume is not None:
- command.extend(("--resume", str(resume)))
- subprocess.run(command, cwd=source, check=True)
-
-
-def test_training_resume_preserves_checkpoints_and_appends_metrics(tmp_path):
- source = Path(__file__).parents[1]
- dataset = tmp_path / "dataset.npz"
- rng = np.random.default_rng(7)
- split = np.array([0] * 30 + [1] * 5 + [2] * 5, dtype=np.uint8)
- np.savez_compressed(
- dataset,
- source_latent=rng.normal(size=(40, 2)).astype(np.float32),
- target_3d=rng.normal(size=(40, 3)).astype(np.float32),
- split=split,
- )
- output = tmp_path / "run"
- config_path = tmp_path / "config.json"
- config = {
- "method": "unite",
- "seed": 42,
- "dataset": str(dataset),
- "source_key": "source_latent",
- "output_dir": str(output),
- "model": {
- "latent_dim": 2,
- "codec_width": 8,
- "codec_depth": 1,
- "field_width": 8,
- "field_depth": 1,
- },
- "flow_samples": 2,
- "reconstruction_noise_range": [0.5, 1.0],
- "learning_rate": 0.0003,
- "batch_size": 4,
- "max_steps": 2,
- "inference_steps": 2,
- "log_every": 1,
- "checkpoint_every": 1,
- }
- config_path.write_text(json.dumps(config))
- _run(source, config_path)
- checkpoint = next((output / "checkpoints").glob("*global-step-000002.pt"))
- legacy_state = torch.load(checkpoint, map_location="cpu", weights_only=False)
- legacy_state.pop("rng")
- torch.save(legacy_state, checkpoint)
-
- config["max_steps"] = 4
- config_path.write_text(json.dumps(config))
- _run(source, config_path, checkpoint)
-
- steps = sorted(
- int(path.stem.rsplit("-", 1)[-1])
- for path in (output / "checkpoints").glob("*.pt")
- )
- assert steps == [1, 2, 3, 4]
- metric_steps = [
- json.loads(row)["step"]
- for row in (output / "metrics.jsonl").read_text().splitlines()
- ]
- assert metric_steps == [1, 2, 3, 4]
- final = torch.load(
- next((output / "checkpoints").glob("*global-step-000004.pt")),
- map_location="cpu",
- weights_only=False,
- )
- assert final["step"] == 4
- assert final["resume"]["checkpoint"] == str(checkpoint)
- assert final["resume"]["rng_restored"] is False
- assert "rng" in final
-
-
-def test_restore_rng_state_moves_saved_cuda_states_to_cpu(monkeypatch):
- generator = torch.Generator().manual_seed(9)
- state = _capture_rng_state(generator)
-
- class DeviceState:
- def __init__(self):
- self.cpu_called = False
-
- def cpu(self):
- self.cpu_called = True
- return torch.tensor([1, 2, 3], dtype=torch.uint8)
-
- device_state = DeviceState()
- state["cuda"] = [device_state]
- captured = []
- monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
- monkeypatch.setattr(torch.cuda, "set_rng_state_all", captured.append)
-
- _restore_rng_state(state, generator)
-
- assert device_state.cpu_called
- assert captured[0][0].device.type == "cpu"
- assert captured[0][0].dtype == torch.uint8
-
-
-@pytest.mark.parametrize(
- ("adapter_family", "adapter_objective"),
- (
- ("fixed_affine", "none"),
- ("nonlinear", "path"),
- ("nonlinear", "action_velocity"),
- ),
-)
-def test_action_adapter_training_uses_50k_default_and_real_validation(
- tmp_path, adapter_family, adapter_objective
-):
- source = Path(__file__).parents[1]
- dataset = tmp_path / "dataset.npz"
- rng = np.random.default_rng(17)
- split = np.array([0] * 30 + [1] * 5 + [2] * 5, dtype=np.uint8)
- np.savez_compressed(
- dataset,
- source_gaussian_latent=rng.normal(size=(40, 8)).astype(np.float32),
- target_3d=rng.normal(size=(40, 3)).astype(np.float32),
- split=split,
- )
- output = tmp_path / f"action-adapter-{adapter_family}-{adapter_objective}"
- config_path = tmp_path / f"action-adapter-{adapter_family}-{adapter_objective}.json"
- config = {
- "architecture": "action_adapter_flow",
- "adapter_objective": adapter_objective,
- "seed": 42,
- "dataset": str(dataset),
- "evaluation_dataset": str(dataset),
- "evaluation_particles": 5,
- "source_key": "source_gaussian_latent",
- "output_dir": str(output),
- "model": {
- "latent_dim": 8,
- "adapter_family": adapter_family,
- "residual_width": 8,
- "residual_depth": 1,
- "field_width": 8,
- "field_depth": 1,
- },
- "flow_samples": 2,
- "lambda_reconstruction": 0.0,
- "lambda_scale": 0.0,
- "lambda_path": 0.0,
- "lambda_action_velocity": 1.0,
- "learning_rate": 0.0003,
- "batch_size": 4,
- "max_steps": 2,
- "inference_steps": 2,
- "diagnostic_noise_samples": 16,
- "angular_bins": 4,
- "log_every": 1,
- }
- config_path.write_text(json.dumps(config))
- _run(source, config_path)
- checkpoints = list((output / "checkpoints").glob("*.pt"))
- assert len(checkpoints) == 1
- assert checkpoints[0].name.endswith("global-step-000002.pt")
- summary = json.loads((output / "summary.json").read_text())
- for key in (
- "validation_generation_symmetric_nn_mse",
- "validation_path_consistency_mse",
- "validation_action_velocity_mse",
- "validation_scale_loss",
- "validation_torus_surface_rmse",
- "validation_angular_histogram_l1",
- "validation_decoder_jacobian_singular_max",
- ):
- assert np.isfinite(summary[key])
diff --git a/tests/test_synthetic_trainer_signals.py b/tests/test_synthetic_trainer_signals.py
deleted file mode 100644
index 46cad3b9e7..0000000000
--- a/tests/test_synthetic_trainer_signals.py
+++ /dev/null
@@ -1,181 +0,0 @@
-"""Regression coverage for checkpoint requests during startup and saving."""
-
-import json
-import os
-import signal
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from scripts.train.train_synthetic_manifold import (
- _atomic_torch_save,
- _CheckpointRequests,
-)
-
-
-def test_request_received_during_save_stays_pending():
- requests = _CheckpointRequests()
- requests.request(0, None)
- snapshot = requests.received
- requests.request(0, None)
- requests.saved = snapshot
- assert requests.received != requests.saved
-
-
-def test_failed_atomic_save_preserves_checkpoint_and_removes_temporary(
- tmp_path, monkeypatch
-):
- checkpoint = tmp_path / "checkpoint.pt"
- torch.save({"step": 1}, checkpoint)
- original = checkpoint.read_bytes()
-
- def fail_save(state, temporary):
- Path(temporary).write_bytes(b"incomplete")
- raise OSError("simulated interrupted write")
-
- monkeypatch.setattr(torch, "save", fail_save)
- with pytest.raises(OSError, match="interrupted write"):
- _atomic_torch_save({"step": 2}, checkpoint)
- assert checkpoint.read_bytes() == original
- assert list(tmp_path.iterdir()) == [checkpoint]
-
-
-def _assert_equal(left, right):
- if isinstance(left, torch.Tensor):
- torch.testing.assert_close(left, right, rtol=0, atol=0)
- elif isinstance(left, np.ndarray):
- np.testing.assert_array_equal(left, right)
- elif isinstance(left, dict):
- assert left.keys() == right.keys()
- for key in left:
- _assert_equal(left[key], right[key])
- elif isinstance(left, (tuple, list)):
- assert len(left) == len(right)
- for lhs, rhs in zip(left, right, strict=True):
- _assert_equal(lhs, rhs)
- else:
- assert left == right
-
-
-@pytest.mark.skipif(not hasattr(signal, "SIGUSR2"), reason="requires POSIX signals")
-def test_startup_signal_checkpoints_and_resumes_exactly(tmp_path):
- source = Path(__file__).resolve().parents[1]
- trainer = source / "scripts/train/train_synthetic_manifold.py"
- dataset = tmp_path / "dataset.npz"
- rng = np.random.default_rng(17)
- np.savez_compressed(
- dataset,
- source_gaussian_latent=rng.normal(size=(40, 8)).astype(np.float32),
- target_3d=rng.normal(size=(40, 3)).astype(np.float32),
- split=np.array([0] * 30 + [1] * 5 + [2] * 5, dtype=np.uint8),
- )
- config = {
- "architecture": "action_adapter_flow",
- "adapter_objective": "action_velocity",
- "seed": 42,
- "dataset": str(dataset),
- "source_key": "source_gaussian_latent",
- "output_dir": str(tmp_path / "interrupted"),
- "model": {
- "latent_dim": 8,
- "adapter_family": "nonlinear",
- "residual_width": 8,
- "residual_depth": 1,
- "field_width": 8,
- "field_depth": 1,
- },
- "flow_samples": 2,
- "lambda_reconstruction": 100.0,
- "learning_rate": 0.0003,
- "batch_size": 4,
- "max_steps": 4,
- "inference_steps": 2,
- "diagnostic_noise_samples": 16,
- "angular_bins": 4,
- "log_every": 1,
- "checkpoint_every": 50_000,
- }
- config_path = tmp_path / "config.json"
- config_path.write_text(json.dumps(config))
- # Inject a real OS signal precisely at the first heavyweight import. There
- # is no sleep-based startup race and no production-only test hook. Exit only
- # after the requested atomic checkpoint is published, simulating a restart.
- bootstrap = """
-import builtins
-import os
-import runpy
-import signal
-import sys
-
-original_import = builtins.__import__
-sent = False
-def signal_before_numpy(name, *args, **kwargs):
- global sent
- if name == 'numpy' and not sent:
- sent = True
- assert callable(signal.getsignal(signal.SIGUSR2)), 'handler installed too late'
- os.kill(os.getpid(), signal.SIGUSR2)
- return original_import(name, *args, **kwargs)
-builtins.__import__ = signal_before_numpy
-original_replace = os.replace
-def interrupt_after_checkpoint(src, dst):
- original_replace(src, dst)
- if str(dst).endswith('.pt'):
- os._exit(75)
-os.replace = interrupt_after_checkpoint
-sys.argv = sys.argv[1:]
-runpy.run_path(sys.argv[0], run_name='__main__')
-"""
- env = {**os.environ, "CUDA_VISIBLE_DEVICES": "", "OMP_NUM_THREADS": "1"}
-
- def run(command, expected=0):
- result = subprocess.run(
- command, cwd=tmp_path, env=env, capture_output=True, text=True, timeout=60
- )
- assert result.returncode == expected, result.stdout + result.stderr
-
- run(
- [sys.executable, "-c", bootstrap, str(trainer), "--config", str(config_path)],
- expected=75,
- )
- output = Path(config["output_dir"])
- checkpoints = list((output / "checkpoints").glob("*.pt"))
- assert len(checkpoints) == 1
- checkpoint = checkpoints[0]
- assert checkpoint.name.endswith("global-step-000001.pt")
- checkpoint_bytes = checkpoint.read_bytes()
- run(
- [
- sys.executable,
- str(trainer),
- "--config",
- str(config_path),
- "--resume",
- str(checkpoint),
- ]
- )
- assert checkpoint.read_bytes() == checkpoint_bytes
- final_path = next((output / "checkpoints").glob("*global-step-000004.pt"))
- resumed = torch.load(final_path, map_location="cpu", weights_only=False)
- assert resumed["resume"]["rng_restored"] is True
- assert [
- json.loads(row)["step"]
- for row in (output / "metrics.jsonl").read_text().splitlines()
- ] == [1, 2, 3, 4]
- summary = json.loads((output / "summary.json").read_text())
- assert np.isfinite(summary["validation_generation_symmetric_nn_mse"])
- assert np.isfinite(summary["validation_action_velocity_mse"])
-
- config["output_dir"] = str(tmp_path / "uninterrupted")
- config_path.write_text(json.dumps(config))
- run([sys.executable, str(trainer), "--config", str(config_path)])
- control_path = next((tmp_path / "uninterrupted/checkpoints").glob("*.pt"))
- control = torch.load(control_path, map_location="cpu", weights_only=False)
- for key in ("model", "optimizer", "rng", "step"):
- _assert_equal(resumed[key], control[key])
- control_summary = json.loads((tmp_path / "uninterrupted/summary.json").read_text())
- assert summary == control_summary
diff --git a/tests/test_synthetic_trajectory_eval.py b/tests/test_synthetic_trajectory_eval.py
deleted file mode 100644
index a26ccdb389..0000000000
--- a/tests/test_synthetic_trajectory_eval.py
+++ /dev/null
@@ -1,120 +0,0 @@
-import numpy as np
-import pytest
-import torch
-
-from egomimic.eval.synthetic_trajectory_eval import SyntheticTrajectoryEval
-from egomimic.synthetic.action_adapter_flow import SyntheticActionAdapterFlow
-from egomimic.synthetic.multi_action_adapter_flow import (
- SyntheticMultiActionAdapterFlow,
-)
-from egomimic.synthetic.shared_latent_flow import (
- SyntheticDirectFlow,
- SyntheticSharedLatentFlow,
-)
-from scripts.eval.export_synthetic_trajectory_npz import (
- _EmbodimentTrajectoryView,
- resolve_device,
-)
-
-
-def test_shared_eval_exports_same_npz_contract_for_both_model_families(tmp_path):
- target = torch.randn(7, 3)
- models_and_sources = (
- (SyntheticSharedLatentFlow(latent_dim=2), torch.randn(7, 2)),
- (SyntheticDirectFlow(data_dim=3), torch.randn(7, 3)),
- (
- SyntheticActionAdapterFlow(latent_dim=8, adapter_family="fixed_affine"),
- torch.randn(7, 8),
- ),
- )
- for index, (model, source) in enumerate(models_and_sources):
- output = tmp_path / f"trajectory-{index}.npz"
- points = SyntheticTrajectoryEval.export(model, source, target, output, steps=3)
- archive = np.load(output, allow_pickle=False)
- assert set(archive.files) == {"times", "points", "target"}
- assert points.shape == (4, 7, 3)
- assert archive["points"].shape == (4, 7, 3)
- np.testing.assert_array_equal(archive["target"], target.numpy())
-
-
-def test_ground_truth_uses_same_npz_contract(tmp_path):
- source = torch.randn(5, 2)
- target = torch.randn(5, 3)
- output = tmp_path / "ground-truth.npz"
- points = SyntheticTrajectoryEval.export_linear_ground_truth(
- source, target, output, steps=2
- )
- torch.testing.assert_close(points[0, :, :2], source)
- torch.testing.assert_close(points[0, :, 2], torch.zeros(5))
- torch.testing.assert_close(points[-1], target)
-
-
-def test_validation_loader_requires_exact_requested_particle_count(tmp_path):
- dataset = tmp_path / "small.npz"
- np.savez_compressed(
- dataset,
- source_latent=np.zeros((5, 2), dtype=np.float32),
- target_3d=np.zeros((5, 3), dtype=np.float32),
- split=np.array([0, 1, 1, 2, 2], dtype=np.uint8),
- )
- source, target = SyntheticTrajectoryEval.load_validation_data(
- dataset, "source_latent", 2
- )
- assert source.shape == (2, 2)
- assert target.shape == (2, 3)
- with pytest.raises(ValueError, match="requested 3 validation particles"):
- SyntheticTrajectoryEval.load_validation_data(dataset, "source_latent", 3)
-
-
-def test_symmetric_nearest_neighbor_mse_checks_both_directions():
- samples = torch.tensor([[0.0, 0.0], [2.0, 0.0]])
- targets = torch.tensor([[0.0, 0.0], [4.0, 0.0]])
- loss = SyntheticTrajectoryEval.symmetric_nearest_neighbor_mse(samples, targets)
- torch.testing.assert_close(loss, torch.tensor(2.0))
-
- collapsed = torch.zeros(2, 2)
- collapsed_loss = SyntheticTrajectoryEval.symmetric_nearest_neighbor_mse(
- collapsed, targets
- )
- assert collapsed_loss > 0
-
-
-def test_torus_metrics_are_exact_on_uniform_surface_samples():
- theta = torch.linspace(0, 2 * torch.pi, 32)[:-1]
- phi = torch.linspace(0, 2 * torch.pi, 32)[:-1]
- theta, phi = torch.meshgrid(theta, phi, indexing="ij")
- tube = 2.0 + 0.65 * phi.cos()
- points = torch.stack(
- (tube * theta.cos(), tube * theta.sin(), 0.65 * phi.sin()), dim=-1
- ).reshape(-1, 3)
- torch.testing.assert_close(
- SyntheticTrajectoryEval.torus_surface_rmse(
- points, major_radius=2.0, minor_radius=0.65
- ),
- torch.tensor(0.0),
- atol=2e-7,
- rtol=0,
- )
- coverage = SyntheticTrajectoryEval.torus_angular_coverage(points, points, bins=8)
- torch.testing.assert_close(coverage["angular_histogram_l1"], torch.tensor(0.0))
- torch.testing.assert_close(coverage["angular_support_recall"], torch.tensor(1.0))
-
-
-def test_trajectory_export_device_resolution_rejects_missing_cuda(monkeypatch):
- assert resolve_device("cpu") == torch.device("cpu")
- monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
- assert resolve_device("auto") == torch.device("cpu")
- with pytest.raises(RuntimeError, match="CUDA was requested"):
- resolve_device("cuda")
-
-
-def test_multi_action_trajectory_view_selects_private_decoder():
- model = SyntheticMultiActionAdapterFlow(
- embodiments=["shallow", "steep"], field_width=8, field_depth=1
- )
- source = torch.randn(5, 8)
- view = _EmbodimentTrajectoryView(model, "steep")
- torch.testing.assert_close(
- view.trajectory(source, steps=2),
- model.trajectory(source, embodiment="steep", steps=2),
- )
diff --git a/tests/test_token_diagnostics.py b/tests/test_token_diagnostics.py
deleted file mode 100644
index 85817625e4..0000000000
--- a/tests/test_token_diagnostics.py
+++ /dev/null
@@ -1,297 +0,0 @@
-"""Provider-scoped attention capture, source identities and reusable archives."""
-
-from functools import wraps
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-import torch
-
-from egomimic.eval.eval import EvaluationDataRequirements
-from egomimic.eval.latent_archive import (
- read_archive,
- rebuild_archives,
- reduce_features,
- write_archive,
-)
-from egomimic.eval.token_diagnostics import TokenDiagnosticsEval
-from egomimic.models.pi05.diagnostics import PI05AttentionProvider, capture_attention
-
-
-def attention_layers():
- layer = torch.nn.Module()
- layer.self_attn = torch.nn.Module()
- layer.self_attn.k_proj = torch.nn.Linear(4, 4)
- return torch.nn.ModuleList([layer])
-
-
-class AttentionBackend(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.paligemma_with_expert = torch.nn.Module()
- self.paligemma_with_expert.paligemma = torch.nn.Module()
- self.paligemma_with_expert.paligemma.language_model = torch.nn.Module()
- self.paligemma_with_expert.paligemma.language_model.layers = attention_layers()
- self.paligemma_with_expert.gemma_expert = torch.nn.Module()
- self.paligemma_with_expert.gemma_expert.model = torch.nn.Module()
- self.paligemma_with_expert.gemma_expert.model.layers = attention_layers()
-
- def embed_prefix(self, images, image_masks, language_tokens, language_masks):
- embeddings = torch.ones(len(language_tokens), 5, 4)
- padding = torch.cat(
- [torch.ones(len(language_tokens), 2, dtype=torch.bool), language_masks], 1
- )
- return embeddings, padding, torch.zeros_like(padding)
-
- def sample(self):
- prefix, _, _ = self.embed_prefix(
- None,
- None,
- torch.ones(2, 3),
- torch.tensor([[True, False, False], [True, True, False]]),
- )
- language = self.paligemma_with_expert.paligemma.language_model.layers[
- 0
- ].self_attn.k_proj(prefix)
- expert = self.paligemma_with_expert.gemma_expert.model.layers[
- 0
- ].self_attn.k_proj
- first = expert(torch.ones(2, 2, 4))
- expert(torch.ones(2, 2, 4) * 5)
- return language, first
-
-
-def test_capture_slices_first_step_masks_and_cleanup():
- backend = AttentionBackend()
- with capture_attention(backend, mask_padding=True) as output:
- language, first = backend.sample()
- torch.testing.assert_close(output["expert_layer_00"]["tokens"], first)
- torch.testing.assert_close(
- output["paligemma_layer_00_img"]["tokens"], language[:, :2]
- )
- assert output["paligemma_layer_00_lang"]["mask"].sum() == 3
- assert "embed_prefix" not in backend.__dict__
- for module in backend.modules():
- assert not module._forward_hooks
- with pytest.raises(RuntimeError, match="sample failure"):
- with capture_attention(backend):
- raise RuntimeError("sample failure")
- assert "embed_prefix" not in backend.__dict__
- assert all(not module._forward_hooks for module in backend.modules())
-
-
-def test_provider_uses_declared_id_and_normal_graph_forward():
- backend = AttentionBackend()
- seen = []
- stage = SimpleNamespace(backend=True, policy_nets={"policy": backend})
-
- def lookup(name):
- assert name == "explicit-policy-id"
- return stage
-
- def forward(batch):
- seen.append(next(iter(batch)))
- backend.sample()
- return {source: {"pred_action": torch.ones(2, 4, 3)} for source in batch}
-
- graph = SimpleNamespace(
- pipeline=SimpleNamespace(stage_by_id=lookup),
- forward_eval=forward,
- process_batch_for_training=lambda batch: batch,
- )
- provider = PI05AttentionProvider(stage_id="explicit-policy-id")
- result = provider.run(
- graph, {"source_a": {}, "source_b": {}}, already_processed=False
- )
- assert seen == ["source_a", "source_b"]
- assert result["source_b"]["predictions"]["pred_action"].shape == (2, 4, 3)
- assert len(result["source_a"]["activations"]) == 4
-
-
-def test_capture_bypasses_compiled_sampler_and_restores_it():
- backend = AttentionBackend()
-
- @wraps(backend.sample)
- def compiled():
- raise AssertionError("Capture must use the original eager sampler")
-
- backend.sample_actions = compiled
- with capture_attention(backend) as output:
- backend.sample_actions()
- assert output and backend.sample_actions is compiled
- with pytest.raises(RuntimeError, match="inference failed"):
- with capture_attention(backend):
- raise RuntimeError("inference failed")
- assert backend.sample_actions is compiled
-
-
-def test_prediction_requirements_are_preserved_and_conflicts_rejected():
- prediction = SimpleNamespace(
- data_requirements=lambda: EvaluationDataRequirements(
- ordered=True,
- complete_episodes=True,
- source_fps=20,
- max_episodes=2,
- sample_id_key="episode_hash",
- frame_index_key="frame_index",
- ),
- trainer_overrides=lambda: {"inference_mode": False},
- )
- evaluator = TokenDiagnosticsEval(
- prediction_evaluator=prediction, limit_val_batches=1.0
- )
- assert evaluator.data_requirements() == prediction.data_requirements()
- assert evaluator.trainer_overrides() == {
- "inference_mode": False,
- "limit_val_batches": 1.0,
- "num_sanity_val_steps": 0,
- }
- evaluator.sample_id_key = "different_episode_id"
- with pytest.raises(ValueError, match="sample_id_key"):
- evaluator.data_requirements()
-
-
-def test_diagnostics_score_the_captured_prediction_through_shared_metrics(tmp_path):
- from egomimic.eval.bimanual_cartesian_eval import BimanualCartesianEval
-
- prediction_evaluator = BimanualCartesianEval(
- viz_every_n_epochs=0, pose_metrics=False
- )
- evaluator = TokenDiagnosticsEval(
- prediction_evaluator=prediction_evaluator, methods=["pca"], save_plots=False
- )
- logged, calls = [], []
- evaluator.trainer = SimpleNamespace(
- default_root_dir=str(tmp_path),
- current_epoch=0,
- global_step=1,
- global_rank=0,
- is_global_zero=True,
- world_size=1,
- lightning_module=SimpleNamespace(
- log_dict=lambda values, **kw: logged.append(values)
- ),
- )
- evaluator.bind_data_context(
- normalizer=SimpleNamespace(
- unnormalize=lambda values, identity: {k: v * 2 for k, v in values.items()}
- )
- )
- batch = {
- "opaque": {
- "embodiment": torch.tensor([7]),
- "episode_hash": ["episode"],
- "frame_index": torch.tensor([0]),
- "actions_cartesian": torch.zeros(1, 4, 14),
- }
- }
-
- def run(capability, values):
- calls.append(capability)
- return {
- "opaque": {
- "predictions": {"pred_action": torch.ones(1, 4, 14)},
- "activations": {
- "layer": {
- "tokens": torch.eye(4)[None],
- "mask": torch.ones(1, 4, dtype=torch.bool),
- }
- },
- }
- }
-
- evaluator.model = SimpleNamespace(run_diagnostic=run)
- evaluator.on_validation_start()
- metrics = evaluator.on_validation_step(batch, 0)
- evaluator.on_validation_end()
- assert calls == ["token_activations"]
- assert metrics["Valid/MSE"] == 1
- assert metrics["Valid/Native_MSE"] == 4
- assert len(logged) == 1
-
-
-def test_archives_preserve_frames_balance_sources_and_rebuild(tmp_path):
- evaluator = TokenDiagnosticsEval(
- methods=["pca"], max_tokens_per_source_layer=6, save_plots=True
- )
- evaluator.trainer = SimpleNamespace(
- default_root_dir=str(tmp_path), current_epoch=0, global_step=10, global_rank=1
- )
- batch = {
- name: {
- "episode_hash": [f"{name}-episode"] * 2,
- "frame_index": torch.tensor([4, 17]),
- }
- for name in ["A", "B"]
- }
-
- def diagnostics(capability, values):
- assert capability == "token_activations"
- return {
- name: {
- "activations": {
- "layer_00": {
- "tokens": torch.arange(32).reshape(2, 4, 4).float(),
- "mask": torch.ones(2, 4, dtype=torch.bool),
- }
- }
- }
- for name in values
- }
-
- evaluator.model = SimpleNamespace(run_diagnostic=diagnostics)
- evaluator.on_validation_start()
- evaluator.on_validation_step(batch, 0)
- evaluator.on_validation_step(
- batch, 1
- ) # repeated distributed/combined-loader samples
- evaluator.on_validation_end()
- path = next(evaluator.output_dir.rglob("*.csv"))
- keys, rows = read_archive(path)
- assert keys.shape == (12, 4)
- assert {row["frame_idx"] for row in rows} == {"4", "17"}
- assert sum(row["embodiment"] == "A" for row in rows) == 6
- assert sum(row["embodiment"] == "B" for row in rows) == 6
- assert (
- len(
- {
- tuple(row[key] for key in ("video_hash", "frame_idx", "token_idx"))
- for row in rows
- }
- )
- == 12
- )
- rebuilt = tmp_path / "rebuild"
- rebuild_archives(path.parent, rebuilt, methods=["pca"], seed=0)
- restored, metadata = read_archive(rebuilt / path.name)
- np.testing.assert_array_equal(restored, keys)
- assert metadata == rows
- assert list(rebuilt.glob("*.png"))
- with pytest.raises(FileExistsError):
- rebuild_archives(path.parent, rebuilt, methods=["pca"])
- with pytest.raises(ValueError, match="differ"):
- rebuild_archives(path.parent, path.parent, methods=["pca"])
-
-
-def test_all_reduction_methods_are_finite_and_reproducible():
- pytest.importorskip("umap")
- features = np.random.default_rng(42).normal(size=(16, 8)).astype(np.float32)
- methods = ["pca", "umap", "pca_umap", "tsne2d", "tsne3d"]
- result, receipt = reduce_features(features, methods=methods, pca_components=5)
- assert set(result) == set(methods)
- assert len(receipt["pca_explained_variance"]) == 5
- for name, coordinates in result.items():
- assert coordinates.shape == (16, 2 if name == "tsne2d" else 3)
- assert np.isfinite(coordinates).all()
-
-
-def test_old_inline_csv_rebuild_and_bad_row_count(tmp_path):
- source = tmp_path / "old.csv"
- source.write_text(
- "video_hash,embodiment,frame_idx,token_idx,k0,k1\nhash,A,7,0,1.0,2.0\n"
- )
- keys, rows = read_archive(source)
- np.testing.assert_array_equal(keys, [[1, 2]])
- assert rows[0]["frame_idx"] == "7"
- with pytest.raises(ValueError, match="row counts"):
- write_archive(tmp_path / "bad", "layer", keys, [], reductions={}, receipt={})
diff --git a/tests/test_unite_diagnostics.py b/tests/test_unite_diagnostics.py
deleted file mode 100644
index bcdbc08880..0000000000
--- a/tests/test_unite_diagnostics.py
+++ /dev/null
@@ -1,117 +0,0 @@
-import json
-
-import numpy as np
-import pytest
-import torch
-
-from egomimic.eval.diagnostics.gradients import gradient_cosine_similarity
-from egomimic.eval.diagnostics.latent_projection import project_pca
-from egomimic.eval.diagnostics.trajectory import (
- chunk_seam_metrics,
- validate_chunk_seam_artifact,
- write_chunk_seam_artifact,
-)
-from egomimic.eval.diagnostics.unite_artifact import (
- UniteDiagnosticArtifactWriter,
- validate_unite_artifact,
-)
-
-
-def test_gradient_cosine_detects_exact_conflict_without_populating_grad():
- parameter = torch.nn.Parameter(torch.tensor([1.0, -2.0]))
- first = parameter.sum()
- second = -parameter.sum()
-
- cosine = gradient_cosine_similarity(first, second, [parameter])
-
- assert cosine.item() == pytest.approx(-1.0)
- assert parameter.grad is None
-
-
-def test_chunk_seam_metrics_use_circular_rotation_and_derivatives():
- old = np.asarray(
- [
- [0.0, 0.0, 3.10, 0.0],
- [1.0, 0.0, 3.11, 0.5],
- [2.0, 0.0, 3.12, 1.0],
- [3.0, 0.0, 3.13, 1.0],
- ],
- dtype=np.float32,
- )
- new = old.copy()
- new[:, 2] = -3.10
- metrics = chunk_seam_metrics(old, new)
-
- assert metrics["aligned_steps"] == 4
- assert metrics["position_mean_l2"] == pytest.approx(0.0)
- assert metrics["rotation_mean_abs_radians"] < 0.1
- assert metrics["velocity_rms_difference"] == pytest.approx(0.0)
- assert metrics["jerk_rms_difference"] == pytest.approx(0.0)
-
-
-def test_chunk_seam_artifact_is_immutable_and_hash_validated(tmp_path):
- event = {
- "t": 8,
- "embodiment_id": 19,
- "episode_index": 0,
- "previous_tail": np.zeros((8, 4), dtype=np.float32),
- "new_head": np.ones((8, 4), dtype=np.float32),
- "executed_head": np.ones((8, 4), dtype=np.float32),
- }
- artifact = write_chunk_seam_artifact([event], tmp_path / "seams")
-
- manifest = validate_chunk_seam_artifact(artifact)
-
- assert manifest["event_count"] == 1
- with pytest.raises(FileExistsError):
- write_chunk_seam_artifact([event], artifact)
-
-
-def test_shared_latent_pca_has_stable_shape():
- features = np.arange(48, dtype=np.float32).reshape(8, 6)
-
- projected, ratio = project_pca(features, 3)
-
- assert projected.shape == (8, 3)
- assert ratio.shape == (3,)
- assert np.all(np.isfinite(projected))
-
-
-def test_unite_artifact_round_trip(tmp_path):
- request = tmp_path / "request.json"
- request.write_text(json.dumps({"schema_version": 1, "checkpoint_sha256": "abc"}))
- writer = UniteDiagnosticArtifactWriter(
- str(tmp_path / "unite"), request_path=str(request)
- )
- batch, horizon, action_dim, tokens, latent_dim = 2, 4, 3, 4, 8
- steps, samples = 3, 5
- writer.append(
- embodiment_id=17,
- embodiment_name="pushshapes_sim_u_socket",
- action_key="actions",
- target_action=np.zeros((batch, horizon, action_dim), dtype=np.float32),
- reconstructed_action=np.zeros((batch, horizon, action_dim), dtype=np.float32),
- generated_action=np.zeros((batch, horizon, action_dim), dtype=np.float32),
- clean_latent=np.zeros((batch, tokens, latent_dim), dtype=np.float32),
- generated_latent=np.zeros((batch, tokens, latent_dim), dtype=np.float32),
- denoising_latents=np.zeros(
- (batch, steps, tokens, latent_dim), dtype=np.float32
- ),
- denoising_actions=np.zeros(
- (batch, steps, horizon, action_dim), dtype=np.float32
- ),
- diversity_latents=np.zeros(
- (batch, samples, tokens, latent_dim), dtype=np.float32
- ),
- diversity_actions=np.zeros(
- (batch, samples, horizon, action_dim), dtype=np.float32
- ),
- episode_hashes=["episode-a", "episode-b"],
- )
- artifact = writer.finalize()
-
- manifest = validate_unite_artifact(artifact, request)
-
- assert manifest["sample_count"] == 2
- assert manifest["parts"][0]["denoising_steps"] == steps
- assert manifest["parts"][0]["diversity_samples"] == samples
diff --git a/tests/test_unite_fixed_step_requeue_boundary.py b/tests/test_unite_fixed_step_requeue_boundary.py
deleted file mode 100644
index 7c50632e2d..0000000000
--- a/tests/test_unite_fixed_step_requeue_boundary.py
+++ /dev/null
@@ -1,121 +0,0 @@
-import json
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-
-from egomimic.utils.slurm_requeue import SaveOnlySignalCheckpoint
-from scripts.ice.unite.fixed_step_requeue_boundary import FixedStepRequeueBoundary
-
-
-class _Strategy:
- world_size = 2
-
- def __init__(self):
- self.barriers = []
-
- def barrier(self, name):
- self.barriers.append(name)
-
- @staticmethod
- def broadcast(value, src=0):
- return value
-
-
-def _callback(tmp_path: Path, scancel: Path) -> FixedStepRequeueBoundary:
- return FixedStepRequeueBoundary(
- initial_global_step=3,
- trigger_global_step=8,
- slurm_job_id="1234",
- scancel=str(scancel),
- record_path=str(tmp_path / "boundary.json"),
- signal_ack_timeout_seconds=0.1,
- requeue_wait_timeout_seconds=0.01,
- )
-
-
-def test_boundary_requires_exactly_five_steps(tmp_path):
- executable = tmp_path / "scancel"
- executable.write_text("#!/bin/sh\n", encoding="utf-8")
- executable.chmod(0o700)
- with pytest.raises(ValueError, match="exactly five"):
- FixedStepRequeueBoundary(
- initial_global_step=3,
- trigger_global_step=9,
- slurm_job_id="1234",
- scancel=str(executable),
- record_path=str(tmp_path / "boundary.json"),
- )
-
-
-def test_step_eight_requests_boundary_and_records_all_rank_ack(monkeypatch, tmp_path):
- executable = tmp_path / "scancel"
- executable.write_text("#!/bin/sh\n", encoding="utf-8")
- executable.chmod(0o700)
- callback = _callback(tmp_path, executable)
- save_callback = SaveOnlySignalCheckpoint(tmp_path / "checkpoints")
- strategy = _Strategy()
- trainer = SimpleNamespace(
- callbacks=[callback, save_callback],
- global_step=8,
- is_global_zero=True,
- strategy=strategy,
- )
- callback.setup(trainer, SimpleNamespace(), "fit")
- monkeypatch.setenv("ICE_REQUEUE_OWNER", "runner")
- monkeypatch.setenv("ICE_CHILD_REQUEUE_DISABLED", "1")
- monkeypatch.setenv("SLURM_JOB_ID", "1234")
- monkeypatch.delenv("SLURM_RESTART_COUNT", raising=False)
-
- def request(*_args, **_kwargs):
- save_callback._request_generation += 1
- return SimpleNamespace(returncode=0, stdout="", stderr="")
-
- monkeypatch.setattr("subprocess.run", request)
- callback.on_train_batch_end(trainer, SimpleNamespace(), None, None, 0)
-
- proof = json.loads((tmp_path / "boundary.json").read_text(encoding="utf-8"))
- assert proof["observed_global_step"] == 8
- assert proof["optimizer_steps_advanced"] == 5
- assert proof["step_after_signal_executed"] is False
- assert proof["save_acknowledged_ranks"] == [0, 1]
- assert proof["world_size"] == 2
- assert strategy.barriers == [
- "unite-fixed-step-boundary-ready",
- "unite-fixed-step-boundary-signal-acknowledged",
- ]
-
-
-def test_boundary_rejects_wrong_owner_before_signalling(monkeypatch, tmp_path):
- executable = tmp_path / "scancel"
- executable.write_text("#!/bin/sh\n", encoding="utf-8")
- executable.chmod(0o700)
- callback = _callback(tmp_path, executable)
- save_callback = SaveOnlySignalCheckpoint(tmp_path / "checkpoints")
- trainer = SimpleNamespace(
- callbacks=[callback, save_callback],
- global_step=8,
- is_global_zero=True,
- strategy=_Strategy(),
- )
- callback.setup(trainer, SimpleNamespace(), "fit")
- monkeypatch.setenv("ICE_REQUEUE_OWNER", "child")
- monkeypatch.setenv("ICE_CHILD_REQUEUE_DISABLED", "1")
- monkeypatch.setenv("SLURM_JOB_ID", "1234")
-
- with pytest.raises(RuntimeError, match="sole requeue owner"):
- callback.on_train_batch_end(trainer, SimpleNamespace(), None, None, 0)
- assert not (tmp_path / "boundary.json").exists()
-
-
-def test_step_after_boundary_is_blocked(monkeypatch, tmp_path):
- executable = tmp_path / "scancel"
- executable.write_text("#!/bin/sh\n", encoding="utf-8")
- executable.chmod(0o700)
- callback = _callback(tmp_path, executable)
- callback._sent = True
- monotonic_values = iter((0.0, 0.0, 0.02))
- monkeypatch.setattr("time.monotonic", lambda: next(monotonic_values))
- monkeypatch.setattr("time.sleep", lambda _seconds: None)
- with pytest.raises(RuntimeError, match="did not terminate"):
- callback.on_train_batch_start(SimpleNamespace(), SimpleNamespace(), None, 0)
diff --git a/tests/test_unite_ice_contract.py b/tests/test_unite_ice_contract.py
deleted file mode 100644
index 52d7e0737d..0000000000
--- a/tests/test_unite_ice_contract.py
+++ /dev/null
@@ -1,106 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-from hydra import compose, initialize_config_dir
-from omegaconf import OmegaConf
-
-from scripts.ice.unite.unite_contract import validate_unite_config
-
-
-@pytest.fixture
-def resolved_config():
- config_dir = Path("egomimic/hydra_configs").resolve()
- with initialize_config_dir(config_dir=str(config_dir), version_base="1.3"):
- return compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "model=bf/us_unite_register_shared_nt4_s42",
- "+experiment=sim/pusht/unite_usocket_register_sweep_val01_h16",
- "trainer.devices=2",
- "++paths.root_dir=.",
- ],
- )
-
-
-@pytest.mark.parametrize(
- ("shared", "tokens"),
- [(True, 4), (True, 8), (False, 4), (False, 8)],
-)
-def test_clean_rows_bind_energy_view_to_world_size(resolved_config, shared, tokens):
- config = OmegaConf.create(OmegaConf.to_container(resolved_config, resolve=False))
- config.model.share_encoder_denoiser = shared
- config.model.num_latent_tokens = tokens
- result = validate_unite_config(config, expected_world_size=2)
- assert result["status"] == "READY"
- assert result["row"] == {
- "share_encoder_denoiser": shared,
- "num_latent_tokens": tokens,
- "latent_dim": 16,
- }
- assert result["energy_score_conditions"] == 32
- assert result["diagnostics"] == "ABSENT"
-
-
-def test_single_gpu_view_is_supported_when_total_is_32(resolved_config):
- config = OmegaConf.create(OmegaConf.to_container(resolved_config, resolve=False))
- config.trainer.devices = 1
- config.evaluator.energy_score_validation_view.world_size = 1
- config.evaluator.energy_score_validation_view.per_rank_batch_size = 32
- assert validate_unite_config(config, expected_world_size=1)["world_size"] == 1
-
-
-def test_energy_view_rejects_world_size_or_total_mismatch(resolved_config):
- config = OmegaConf.create(OmegaConf.to_container(resolved_config, resolve=False))
- config.evaluator.energy_score_validation_view.per_rank_batch_size = 8
- with pytest.raises(ValueError, match="exactly 32 conditions"):
- validate_unite_config(config, expected_world_size=2)
-
-
-def test_optional_diagnostics_use_the_same_world_size_contract(resolved_config):
- config = OmegaConf.create(OmegaConf.to_container(resolved_config, resolve=False))
- config.evaluator.unite_diagnostics = {
- "validation_view": {"world_size": 2, "per_rank_batch_size": 16}
- }
- assert (
- validate_unite_config(config, expected_world_size=2)["diagnostics"]
- == "BOUND_TO_WORLD_SIZE"
- )
- config.evaluator.unite_diagnostics.validation_view.per_rank_batch_size = 15
- with pytest.raises(ValueError, match="diagnostic validation view"):
- validate_unite_config(config, expected_world_size=2)
-
-
-def test_contract_rejects_scientific_drift(resolved_config):
- config = OmegaConf.create(OmegaConf.to_container(resolved_config, resolve=False))
- policy = next(
- stage
- for stage in config.model.pipeline.stages
- if str(stage.get("_target_", "")).endswith("ReleasedRecipeUniteLatentPolicy")
- )
- policy.flow_steps_per_reconstruction = 13
- with pytest.raises(ValueError, match="unsupported UNITE"):
- validate_unite_config(config, expected_world_size=2)
-
-
-def test_n8_flow42_ablation_is_explicit_and_contract_valid():
- config_dir = Path("egomimic/hydra_configs").resolve()
- with initialize_config_dir(config_dir=str(config_dir), version_base="1.3"):
- config = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "model=bf/us_unite_register_shared_nt8_flow42_s42",
- "+experiment=sim/pusht/unite_usocket_register_sweep_val01_h16",
- "trainer.devices=1",
- "evaluator.energy_score_validation_view.per_rank_batch_size=32",
- "++paths.root_dir=.",
- ],
- )
- result = validate_unite_config(config, expected_world_size=1)
- assert result["row"]["num_latent_tokens"] == 8
- assert result["flow_steps_per_reconstruction"] == 42
- assert (
- config.run_provenance.unite_update_schedule.flow_samples_per_reconstruction
- == 42
- )
diff --git a/tests/test_validate_action_flow_config.py b/tests/test_validate_action_flow_config.py
deleted file mode 100644
index b532df776b..0000000000
--- a/tests/test_validate_action_flow_config.py
+++ /dev/null
@@ -1,421 +0,0 @@
-from __future__ import annotations
-
-import copy
-import hashlib
-import importlib.util
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).resolve().parents[1]
-CONFIG_ROOT = ROOT / "egomimic" / "hydra_configs"
-TOOL = ROOT / "tools" / "validate_action_flow_config.py"
-SPEC = importlib.util.spec_from_file_location("action_flow_preflight", TOOL)
-assert SPEC is not None and SPEC.loader is not None
-preflight = importlib.util.module_from_spec(SPEC)
-SPEC.loader.exec_module(preflight)
-
-
-def test_real_action_flow_config_instantiates_and_passes_full_preflight():
- report, _ = preflight.validate_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42",
- config_root=CONFIG_ROOT,
- )
-
- assert report["status"] == "PASS"
- assert report["config_name"] == "action_flow_bc_usocket_recon1_s42"
- assert len(report["resolved_config_sha256"]) == 64
- assert report["dimensions"] == {
- "action": [16, 4],
- "condition": 67,
- "image_feature": 64,
- "latent": [16, 8],
- "normalized_state": 3,
- }
- expected_counts = {
- "decoder_g": 10_744,
- "encoder_e": 10_748,
- "field_v": 39_506_641,
- "observation_encoder": 11_197_088,
- "pipeline_total": 50_725_221,
- }
- assert set(report["parameters"]) == set(expected_counts)
- for name, expected in expected_counts.items():
- manifest = report["parameters"][name]
- assert manifest["total"] == manifest["trainable"] == expected
- assert len(manifest["manifest_sha256"]) == 64
- assert sum(entry["numel"] for entry in manifest["entries"]) == expected
- assert len({entry["name"] for entry in manifest["entries"]}) == len(
- manifest["entries"]
- )
- assert report["topology"]["shared_field_instance"] is True
- assert report["topology"]["shared_decoder_instance"] is True
- assert report["topology"]["inference_order"] == [
- "FusedObsEncoder",
- "GaussianLatentNoise",
- "ConditionalVelocityStage",
- "ContentDecoderStage",
- ]
- assert report["launch"]["effective_global_batch"] == 32
- assert report["split"]["zero_id_overlap"] is True
- assert report["split"]["zero_resolved_path_overlap"] is True
- content_manifest_path = ROOT / preflight.CANONICAL_CONTENT_MANIFEST_RELATIVE_PATH
- content_manifest = json.loads(content_manifest_path.read_text())
- assert report["split"]["content_manifest_path"] == (
- preflight.CANONICAL_CONTENT_MANIFEST_RELATIVE_PATH.as_posix()
- )
- assert (
- report["split"]["content_manifest_sha256"]
- == hashlib.sha256(content_manifest_path.read_bytes()).hexdigest()
- )
- assert (
- report["split"]["dataset_content_aggregate_sha256"]
- == content_manifest["aggregate_sha256"]
- )
-
-
-def test_option_a_200m_muon_config_has_exact_capacity_optimizer_and_schedule():
- report, _ = preflight.validate_experiment(
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42",
- config_root=CONFIG_ROOT,
- )
-
- assert report["status"] == "PASS"
- assert report["parameters"]["encoder_e"]["total"] == 1_010_420
- assert report["parameters"]["decoder_g"]["total"] == 1_010_416
- assert report["parameters"]["field_v"]["total"] == 186_536_913
- assert report["parameters"]["pipeline_total"]["total"] == 199_754_837
- optimization = report["optimization"]
- assert optimization["optimizer"]["target"].endswith(
- "ReleasedUniteCompositeOptimizer"
- )
- assert optimization["optimizer"]["lr"] == pytest.approx(1.0e-5)
- assert optimization["scheduler"]["eta_min"] == pytest.approx(1.0e-6)
- assert optimization["parameter_groups"]["complete"] is True
- assert optimization["parameter_groups"]["disjoint"] is True
-
-
-def test_option_a_200m_adamw_changes_only_optimizer_family():
- report, _ = preflight.validate_experiment(
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42",
- config_root=CONFIG_ROOT,
- )
-
- assert report["status"] == "PASS"
- assert report["parameters"]["encoder_e"]["total"] == 1_010_420
- assert report["parameters"]["decoder_g"]["total"] == 1_010_416
- assert report["parameters"]["field_v"]["total"] == 186_536_913
- assert report["parameters"]["pipeline_total"]["total"] == 199_754_837
- optimization = report["optimization"]
- assert optimization["optimizer"]["target"] == "torch.optim.AdamW"
- assert optimization["optimizer"]["lr"] == pytest.approx(1.0e-5)
- assert optimization["scheduler"]["eta_min"] == pytest.approx(1.0e-6)
- assert optimization["parameter_groups"]["adamw_parameters"] == 199_754_837
- assert optimization["parameter_groups"]["muon_parameters"] == 0
-
-
-def test_resolved_hash_is_stable_and_uses_runtime_sentinels():
- first = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- second = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- first_payload, first_hash = preflight.resolved_config_payload(first)
- second_payload, second_hash = preflight.resolved_config_payload(second)
-
- assert first_hash == second_hash
- assert first_payload == second_payload
- assert first_payload["paths"]["output_dir"] == ""
- assert first_payload["logger"]["wandb"]["id"] == ""
-
-
-def test_pair_diff_is_only_the_declared_reconstruction_setting():
- report = preflight.validate_pair(
- "sim/pusht/action_flow_bc_usocket_recon1_s42",
- "sim/pusht/action_flow_bc_usocket_recon10_s42",
- config_root=CONFIG_ROOT,
- )
-
- assert report["status"] == "PASS"
- assert report["comparison"] == {
- "differing_paths": sorted(preflight.ALLOWED_PAIR_DIFFERENCES),
- "only_declared_reconstruction_differences": True,
- "status": "PASS",
- }
- assert {
- item["objective"]["reconstruction_weight"] for item in report["experiments"]
- } == {1.0, 10.0}
-
-
-def test_recon100_diff_is_only_the_declared_reconstruction_setting():
- report = preflight.validate_pair(
- "sim/pusht/action_flow_bc_usocket_recon1_s42",
- "sim/pusht/action_flow_bc_usocket_recon100_s42",
- config_root=CONFIG_ROOT,
- )
-
- assert report["status"] == "PASS"
- assert report["comparison"]["only_declared_reconstruction_differences"] is True
- assert {
- item["objective"]["reconstruction_weight"] for item in report["experiments"]
- } == {1.0, 100.0}
-
-
-@pytest.mark.parametrize(
- ("path", "value", "message"),
- [
- ("model.pipeline.stages.5.field.time_scale", 1.0, "time scale"),
- ("model.pipeline.stages.4.samples_per_content", 13, "bridge samples"),
- ("callbacks.model_checkpoint.every_n_train_steps", 20_000, "checkpoint"),
- ("launch_params.gpus_per_node", 2, "GPUs per node"),
- (
- "evaluator.action_flow_diagnostics.enabled",
- False,
- "diagnostics enabled",
- ),
- (
- "evaluator.action_flow_diagnostics.native_error.native_theta_index",
- 1,
- "unsupported USocket native error contract",
- ),
- (
- "run_provenance.content_manifest_path",
- "egomimic/hydra_configs/data/pusht/manifests/not-canonical.json",
- "canonical dataset-content manifest path",
- ),
- (
- "run_provenance.content_manifest_sha256",
- "0" * 64,
- "canonical dataset-content manifest hash",
- ),
- (
- "run_provenance.dataset_content_aggregate_sha256",
- "0" * 64,
- "canonical dataset aggregate content hash",
- ),
- (
- "evaluator.energy_score_provenance.source_commit",
- "0" * 40,
- "EnergyScore source commit provenance",
- ),
- (
- "evaluator.energy_score_provenance.normalization_sha256",
- "0" * 64,
- "EnergyScore normalization provenance",
- ),
- (
- "evaluator.energy_score_provenance.split_manifest_sha256",
- "0" * 64,
- "EnergyScore split provenance",
- ),
- (
- "evaluator.energy_score_validation_view.split_manifest_sha256",
- "0" * 64,
- "EnergyScore validation-view split provenance",
- ),
- (
- "evaluator.energy_score_provenance.dataset_content.manifest_path",
- "not-the-run-manifest.json",
- "evaluator dataset-content manifest path provenance",
- ),
- (
- "evaluator.energy_score_provenance.dataset_content.manifest_sha256",
- "0" * 64,
- "evaluator dataset-content manifest hash provenance",
- ),
- (
- "evaluator.energy_score_provenance.dataset_content.aggregate_sha256",
- "0" * 64,
- "evaluator dataset aggregate content hash provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.provenance.source_commit",
- "0" * 40,
- "diagnostic source commit provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.provenance.normalization_sha256",
- "0" * 64,
- "diagnostic normalization provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.provenance.split_manifest_sha256",
- "0" * 64,
- "diagnostic split provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.validation_view.split_manifest_sha256",
- "0" * 64,
- "diagnostic validation-view split provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.provenance.dataset_content.manifest_sha256",
- "0" * 64,
- "diagnostic dataset-content manifest provenance",
- ),
- (
- "evaluator.action_flow_diagnostics.provenance.dataset_content.aggregate_sha256",
- "0" * 64,
- "diagnostic dataset aggregate content provenance",
- ),
- ],
-)
-def test_preflight_fails_closed_before_training(path, value, message):
- config = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- from omegaconf import OmegaConf
-
- OmegaConf.update(config, path, value)
- with pytest.raises(preflight.PreflightError, match=message):
- preflight.validate_config(
- config,
- experiment="sim/pusht/action_flow_bc_usocket_recon1_s42",
- config_root=CONFIG_ROOT,
- )
-
-
-@pytest.mark.parametrize(
- "path",
- [
- "evaluator.energy_score_distance",
- "run_provenance.energy_score_contract.distance",
- "evaluator.energy_score_provenance.distance_contract",
- ],
-)
-def test_every_energy_distance_surface_requires_the_typed_usocket_contract(path):
- config = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- from omegaconf import OmegaConf
-
- invalid = copy.deepcopy(preflight.USOCKET_ENERGY_DISTANCE_CONFIG)
- invalid["rotation_scale_radians"] = 1.0
- OmegaConf.update(config, path, invalid, merge=False)
-
- with pytest.raises(
- preflight.PreflightError,
- match="unsupported USocket EnergyScore distance contract",
- ):
- preflight.validate_config(
- config,
- experiment="sim/pusht/action_flow_bc_usocket_recon1_s42",
- config_root=CONFIG_ROOT,
- )
-
-
-def test_native_error_surface_rejects_undeclared_contract_keys():
- config = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- from omegaconf import OmegaConf
-
- invalid = copy.deepcopy(preflight.USOCKET_NATIVE_ERROR_CONFIG)
- invalid["undeclared"] = True
- OmegaConf.update(
- config,
- "evaluator.action_flow_diagnostics.native_error",
- invalid,
- merge=False,
- )
-
- with pytest.raises(preflight.PreflightError, match="native error keys differ"):
- preflight.validate_config(
- config,
- experiment="sim/pusht/action_flow_bc_usocket_recon1_s42",
- config_root=CONFIG_ROOT,
- )
-
-
-@pytest.mark.parametrize(
- "path",
- [
- "evaluator.energy_score_provenance.dataset_content",
- "evaluator.action_flow_diagnostics.provenance.dataset_content",
- ],
-)
-def test_dataset_content_provenance_rejects_undeclared_keys(path):
- config = preflight.compose_experiment(
- "sim/pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT
- )
- from omegaconf import OmegaConf
-
- payload = OmegaConf.to_container(OmegaConf.select(config, path), resolve=True)
- assert isinstance(payload, dict)
- payload["undeclared"] = True
- OmegaConf.update(config, path, payload, merge=False)
-
- with pytest.raises(preflight.PreflightError, match="provenance keys"):
- preflight.validate_config(
- config,
- experiment="sim/pusht/action_flow_bc_usocket_recon1_s42",
- config_root=CONFIG_ROOT,
- )
-
-
-def test_cli_emits_deterministic_machine_readable_failure(tmp_path, monkeypatch):
- output = tmp_path / "preflight.json"
-
- def fail(*_args, **_kwargs):
- raise preflight.PreflightError("deliberate contract failure")
-
- monkeypatch.setattr(preflight, "validate_experiment", fail)
- status = preflight.main(
- [
- "--experiment",
- "sim/pusht/action_flow_bc_usocket_recon1_s42",
- "--output",
- str(output),
- ]
- )
-
- assert status == 1
- assert json.loads(output.read_text()) == {
- "error": {
- "message": "deliberate contract failure",
- "type": "PreflightError",
- },
- "schema_version": 1,
- "status": "FAIL",
- }
- assert output.read_text().endswith("\n")
-
-
-def test_direct_script_cli_resolves_repository_package(tmp_path):
- output = tmp_path / "preflight.json"
- completed = subprocess.run(
- [
- sys.executable,
- str(TOOL),
- "--experiment",
- "sim/pusht/action_flow_bc_usocket_recon1_s42",
- "--output",
- str(output),
- ],
- cwd=ROOT,
- check=False,
- capture_output=True,
- text=True,
- )
-
- assert completed.returncode == 0, completed.stderr or completed.stdout
- payload = json.loads(output.read_text())
- assert payload["status"] == "PASS"
- assert payload["parameters"]["field_v"]["total"] == 39_506_641
-
-
-@pytest.mark.parametrize(
- "selector",
- [
- "pusht/action_flow_chain_manual4919_avln_80k_s42",
- "sim/pusht/parents/action_flow_usocket_h384",
- "pusht_historical/action_flow_cotrain_uc_speed_interpolation",
- ],
-)
-def test_fresh_composition_rejects_retired_or_non_recipe_selectors(selector):
- with pytest.raises(preflight.PreflightError, match="retired|historical"):
- preflight.compose_experiment(selector, config_root=CONFIG_ROOT)
diff --git a/tests/test_validate_planar_dataset.py b/tests/test_validate_planar_dataset.py
deleted file mode 100644
index b93f493bde..0000000000
--- a/tests/test_validate_planar_dataset.py
+++ /dev/null
@@ -1,241 +0,0 @@
-import copy
-import hashlib
-import importlib.util
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-import zarr
-
-ROOT = Path(__file__).parents[1]
-MODULE_PATH = ROOT / "scripts" / "ice" / "validate_planar_dataset.py"
-MANIFEST_PATH = (
- ROOT
- / "egomimic"
- / "hydra_configs"
- / "data"
- / "pusht"
- / "planar_v2_usocket_dp_3k_split_seed42_v1.json"
-)
-SPEC = importlib.util.spec_from_file_location("validate_planar_dataset", MODULE_PATH)
-assert SPEC is not None and SPEC.loader is not None
-MODULE = importlib.util.module_from_spec(SPEC)
-SPEC.loader.exec_module(MODULE)
-
-
-def _manifest(names, source_root, domain=MODULE.DOMAIN):
- train, valid = MODULE.split_names(names)
- source_paths = {name: source_root / name for name in names}
- train_paths = {source_paths[name] for name in train}
- valid_paths = {source_paths[name] for name in valid}
- return {
- "schema_version": 1,
- "status": "PASS",
- "split_algorithm": MODULE.SPLIT_ALGORITHM,
- "split_seed": 42,
- "valid_ratio": 0.01,
- "generator_sha256": "a" * 64,
- "cross_domain_train_valid_resolved_path_overlap_count": 0,
- "domains": {
- domain: {
- "folder_path": str(source_root),
- "total_count": len(names),
- "inventory_names_sha256": MODULE.names_sha256(names),
- "train_count": len(train),
- "train_names_sha256": MODULE.names_sha256(train),
- "train_resolved_paths_sha256": MODULE.paths_sha256(train_paths),
- "train_ids": sorted(train),
- "valid_count": len(valid),
- "valid_names_sha256": MODULE.names_sha256(valid),
- "valid_resolved_paths_sha256": MODULE.paths_sha256(valid_paths),
- "valid_ids": sorted(valid),
- "id_overlap_count": 0,
- "resolved_path_overlap_count": 0,
- "union_count": len(names),
- "union_matches_inventory": True,
- }
- },
- }
-
-
-def test_committed_manifest_has_exact_hash_and_contract():
- expected = "3683e3461596eef8df2432fa865779b3c77b2a2057dabd0fea125595729cf313"
- assert hashlib.sha256(MANIFEST_PATH.read_bytes()).hexdigest() == expected
-
- result = MODULE.validate_manifest_structure(json.loads(MANIFEST_PATH.read_text()))
-
- assert result["total_count"] == 2999
- assert result["train_count"] == 2970
- assert result["valid_count"] == 29
- assert result["id_overlap_count"] == 0
- assert result["union_matches_inventory"] is True
-
-
-def test_committed_chain_manifest_has_exact_hash_and_contract():
- path = MANIFEST_PATH.with_name("planar_v2_chain_gripper_dp_3k_split_seed42_v1.json")
- expected = "3ced944ea3af8e875ea88fc5c2df3a5d2865a9f95223d109fb4bd28c8be7cf69"
- assert hashlib.sha256(path.read_bytes()).hexdigest() == expected
- result = MODULE.validate_manifest_structure(json.loads(path.read_text()))
- assert result["domain"] == "pushshapes_sim_chain_gripper"
- assert (result["total_count"], result["train_count"], result["valid_count"]) == (
- 3000,
- 2970,
- 30,
- )
-
-
-def test_unsupported_domain_still_fails_closed():
- manifest = _manifest({"episode_000"}, Path("/original/dataset"), "other_domain")
- with pytest.raises(RuntimeError, match="supported PushT domain"):
- MODULE.validate_manifest_structure(manifest)
-
-
-def test_portable_inventory_reproduces_split_and_has_no_path_overlap(tmp_path):
- names = {f"episode_{index:03d}" for index in range(100)}
- source_root = Path("/original/dataset")
- manifest = _manifest(names, source_root)
- dataset_root = tmp_path / "portable-dataset"
- dataset_root.mkdir()
- entries = []
- for name in sorted(names):
- path = dataset_root / f"{name}.zarr"
- path.mkdir()
- entries.append((path, name))
-
- result = MODULE.validate_physical_inventory(manifest, entries, dataset_root)
-
- assert result["domain"] == MODULE.DOMAIN
- assert result["total_count"] == 100
- assert result["train_count"] == 99
- assert result["valid_count"] == 1
- assert result["resolved_path_overlap_count"] == 0
- assert result["union_count"] == 100
- assert result["union_matches_inventory"] is True
-
-
-def test_inventory_mismatch_fails_closed(tmp_path):
- names = {f"episode_{index:03d}" for index in range(10)}
- manifest = _manifest(names, Path("/original/dataset"))
- dataset_root = tmp_path / "dataset"
- dataset_root.mkdir()
- entries = []
- for name in sorted(names - {"episode_009"}):
- path = dataset_root / name
- path.mkdir()
- entries.append((path, name))
-
- with pytest.raises(RuntimeError, match="physical episode inventory mismatch"):
- MODULE.validate_physical_inventory(manifest, entries, dataset_root)
-
-
-def test_manifest_split_tampering_fails_closed():
- names = {f"episode_{index:03d}" for index in range(100)}
- manifest = _manifest(names, Path("/original/dataset"))
- tampered = copy.deepcopy(manifest)
- domain = tampered["domains"][MODULE.DOMAIN]
- domain["train_ids"][0], domain["valid_ids"][0] = (
- domain["valid_ids"][0],
- domain["train_ids"][0],
- )
- domain["train_ids"].sort()
- domain["valid_ids"].sort()
- domain["train_names_sha256"] = MODULE.names_sha256(domain["train_ids"])
- domain["valid_names_sha256"] = MODULE.names_sha256(domain["valid_ids"])
- domain["inventory_names_sha256"] = MODULE.names_sha256(
- domain["train_ids"] + domain["valid_ids"]
- )
-
- with pytest.raises(RuntimeError, match="do not reproduce"):
- MODULE.validate_manifest_structure(tampered)
-
-
-def test_two_ids_resolving_to_one_physical_directory_fail(tmp_path):
- names = {"episode_a", "episode_b"}
- manifest = _manifest(names, Path("/original/dataset"))
- dataset_root = tmp_path / "dataset"
- dataset_root.mkdir()
- first = dataset_root / "episode_a"
- first.mkdir()
- second = dataset_root / "episode_b"
- second.symlink_to(first, target_is_directory=True)
-
- with pytest.raises(RuntimeError, match="same physical path"):
- MODULE.validate_physical_inventory(
- manifest,
- [(first, "episode_a"), (second, "episode_b")],
- dataset_root,
- )
-
-
-def test_explicit_deleted_directory_is_excluded_but_accounted(tmp_path):
- names = {"episode_a", "episode_b"}
- manifest = _manifest(names, Path("/original/dataset"))
- dataset_root = tmp_path / "dataset"
- dataset_root.mkdir()
- entries = []
- for name in sorted(names):
- path = dataset_root / name
- path.mkdir()
- entries.append((path, name))
- deleted = dataset_root / "deleted_episode"
- deleted.mkdir()
-
- result = MODULE.validate_physical_inventory(
- manifest,
- entries,
- dataset_root,
- excluded_deleted_paths=[deleted],
- )
-
- assert result["excluded_deleted_directory_count"] == 1
-
-
-def test_cli_reuses_training_resolver_and_allows_only_explicitly_deleted_store(
- tmp_path,
-):
- names = {f"episode_{index:03d}" for index in range(10)}
- dataset_root = tmp_path / "dataset"
- dataset_root.mkdir()
- for name in names:
- zarr.open_group(str(dataset_root / f"{name}.zarr"), mode="w")
- deleted = zarr.open_group(str(dataset_root / "deleted.zarr"), mode="w")
- deleted.attrs["is_deleted"] = True
-
- manifest = _manifest(names, Path("/original/dataset"))
- generator = ROOT / "scripts" / "data" / "materialize_local_episode_split.py"
- manifest["generator_sha256"] = hashlib.sha256(generator.read_bytes()).hexdigest()
- manifest_path = tmp_path / "manifest.json"
- manifest_path.write_text(
- json.dumps(manifest, sort_keys=True, separators=(",", ":")) + "\n"
- )
- expected_sha = hashlib.sha256(manifest_path.read_bytes()).hexdigest()
- output = tmp_path / "validation.json"
-
- completed = subprocess.run(
- [
- sys.executable,
- str(MODULE_PATH),
- "--repo",
- str(ROOT),
- "--dataset-root",
- str(dataset_root),
- "--manifest",
- str(manifest_path),
- "--expected-manifest-sha256",
- expected_sha,
- "--output",
- str(output),
- ],
- check=True,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
-
- result = json.loads(output.read_text())
- assert result["status"] == "DATASET_VALIDATED"
- assert result["physical_inventory"]["total_count"] == 10
- assert result["physical_inventory"]["excluded_deleted_directory_count"] == 1
- assert json.loads(completed.stdout.splitlines()[-1]) == result
diff --git a/tests/test_validate_slurm_job_contract.py b/tests/test_validate_slurm_job_contract.py
deleted file mode 100644
index 138581fd3c..0000000000
--- a/tests/test_validate_slurm_job_contract.py
+++ /dev/null
@@ -1,240 +0,0 @@
-import importlib.util
-import json
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-ROOT = Path(__file__).parents[1]
-SCRIPT = ROOT / "scripts" / "train" / "validate_slurm_job_contract.py"
-
-
-def _load_module():
- spec = importlib.util.spec_from_file_location("validate_slurm_job_contract", SCRIPT)
- assert spec is not None and spec.loader is not None
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- return module
-
-
-def _record(**replacements):
- fields = {
- "JobId": "12345",
- "JobName": "action-flow",
- "Account": "gts-dxu345-rl2",
- "QOS": "inferno",
- "JobState": "RUNNING",
- "TimeLimit": "3-00:00:00",
- "Partition": "gpu-h200",
- "NumNodes": "1",
- "NumCPUs": "8",
- "NumTasks": "1",
- "CPUs/Task": "8",
- "ReqTRES": "cpu=8,mem=250G,node=1,billing=8,gres/gpu=1",
- "MinCPUsNode": "8",
- "MinMemoryNode": "250G",
- "Features": "H100|H200",
- "Command": "/task/launch.sh",
- }
- fields.update(replacements)
- return " ".join(f"{key}={value}" for key, value in fields.items()) + "\n"
-
-
-def _expectations():
- return {
- "expected_job_id": "12345",
- "expected_account": "gts-dxu345-rl2",
- "expected_partition": "gpu-h200",
- "expected_qos": "inferno",
- "expected_cpus": 8,
- "expected_memory": "250G",
- "expected_time_limit": "3-00:00:00",
- "expected_constraint": "H100|H200",
- }
-
-
-def test_contract_cli_writes_immutable_pass_evidence(tmp_path):
- record = tmp_path / "scontrol.txt"
- record.write_text(_record())
- output = tmp_path / "contract.json"
- completed = subprocess.run(
- [
- sys.executable,
- str(SCRIPT),
- "--record",
- str(record),
- "--expected-job-id",
- "12345",
- "--expected-account",
- "gts-dxu345-rl2",
- "--expected-partition",
- "gpu-h200",
- "--expected-qos",
- "inferno",
- "--expected-cpus",
- "8",
- "--expected-memory",
- "250G",
- "--expected-time-limit",
- "3-00:00:00",
- "--expected-constraint",
- "H100|H200",
- "--output",
- str(output),
- ],
- check=True,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- payload = json.loads(output.read_text())
- assert payload["status"] == "SLURM_JOB_CONTRACT_VALIDATED"
- assert payload["failures"] == []
- assert payload["observed"]["memory_bytes"] == 250 * 1024**3
- assert payload["observed"]["requested_tres_gpus"] == 1
-
- repeated = subprocess.run(
- completed.args,
- check=False,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- assert repeated.returncode != 0
- assert json.loads(output.read_text())["status"] == "SLURM_JOB_CONTRACT_VALIDATED"
-
-
-@pytest.mark.parametrize("constraint", ["H100", "H200", "H100|H200"])
-def test_exact_supported_gpu_constraints_pass(constraint):
- module = _load_module()
- fields = module.parse_scontrol_record(_record(Features=constraint))
- expectations = _expectations()
- expectations["expected_constraint"] = constraint
- expected, _, failures = module.evaluate_contract(fields, **expectations)
- assert failures == []
- assert expected["constraint"] == constraint
-
-
-@pytest.mark.parametrize(
- ("replacement", "failure_field"),
- [
- ({"JobId": "999"}, "job_id"),
- ({"Account": "coc"}, "account"),
- ({"Partition": "ice-gpu"}, "partition"),
- ({"QOS": "coc-ice"}, "qos"),
- ({"NumNodes": "2"}, "nodes"),
- ({"NumTasks": "2"}, "tasks"),
- ({"NumCPUs": "16"}, "cpus"),
- ({"CPUs/Task": "4"}, "cpus_per_task"),
- ({"MinCPUsNode": "4"}, "minimum_cpus_per_node"),
- ({"MinMemoryNode": "200G"}, "memory_bytes"),
- ({"Features": "H100"}, "constraint"),
- ({"TimeLimit": "2-00:00:00"}, "time_limit_seconds"),
- (
- {"ReqTRES": "cpu=8,mem=250G,node=2,billing=8,gres/gpu=1"},
- "requested_tres_nodes",
- ),
- (
- {"ReqTRES": "cpu=4,mem=250G,node=1,billing=8,gres/gpu=1"},
- "requested_tres_cpus",
- ),
- (
- {"ReqTRES": "cpu=8,mem=200G,node=1,billing=8,gres/gpu=1"},
- "requested_tres_memory_bytes",
- ),
- (
- {"ReqTRES": "cpu=8,mem=250G,node=1,billing=8,gres/gpu=2"},
- "requested_tres_gpus",
- ),
- (
- {
- "ReqTRES": (
- "cpu=8,mem=250G,node=1,billing=8," "gres/gpu=1,gres/gpu:h100=1"
- )
- },
- "typed_gpu_tres",
- ),
- ],
-)
-def test_each_scheduler_contract_mismatch_fails(replacement, failure_field):
- module = _load_module()
- fields = module.parse_scontrol_record(_record(**replacement))
- _, _, failures = module.evaluate_contract(fields, **_expectations())
- assert failure_field in {failure["field"] for failure in failures}
-
-
-def test_contract_cli_records_failure_and_exits_nonzero(tmp_path):
- record = tmp_path / "scontrol.txt"
- record.write_text(_record(Features="H200"))
- output = tmp_path / "contract.json"
- completed = subprocess.run(
- [
- sys.executable,
- str(SCRIPT),
- "--record",
- str(record),
- "--expected-job-id",
- "12345",
- "--expected-account",
- "gts-dxu345-rl2",
- "--expected-partition",
- "gpu-h200",
- "--expected-qos",
- "inferno",
- "--expected-cpus",
- "8",
- "--expected-memory",
- "250G",
- "--expected-time-limit",
- "3-00:00:00",
- "--expected-constraint",
- "H100|H200",
- "--output",
- str(output),
- ],
- check=False,
- text=True,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- )
- assert completed.returncode == 1
- payload = json.loads(output.read_text())
- assert payload["status"] == "SLURM_JOB_CONTRACT_FAILED"
- assert payload["failures"] == [
- {"field": "constraint", "expected": "H100|H200", "observed": "H200"}
- ]
-
-
-def test_parser_rejects_multiline_and_duplicate_fields():
- module = _load_module()
- with pytest.raises(module.ContractError, match="one-line"):
- module.parse_scontrol_record("JobId=1\nAccount=x\n")
- with pytest.raises(module.ContractError, match="repeats field JobId"):
- module.parse_scontrol_record("JobId=1 JobId=2")
-
-
-def test_parser_accepts_empty_unrelated_slurm_fields():
- module = _load_module()
- fields = module.parse_scontrol_record(_record(StdErr=""))
- assert fields["StdErr"] == ""
- _, _, failures = module.evaluate_contract(fields, **_expectations())
- assert failures == []
-
-
-def test_contract_still_rejects_empty_required_fields():
- module = _load_module()
- fields = module.parse_scontrol_record(_record(Account=""))
- with pytest.raises(
- module.ContractError, match="required scontrol field Account is empty"
- ):
- module.evaluate_contract(fields, **_expectations())
-
-
-@pytest.mark.parametrize(
- ("raw", "seconds"),
- [("3-00:00:00", 3 * 86400), ("16:00:00", 16 * 3600), ("30", 1800)],
-)
-def test_slurm_duration_forms(raw, seconds):
- module = _load_module()
- assert module.parse_slurm_duration(raw) == seconds
diff --git a/tests/test_verify_action_flow_training_smoke.py b/tests/test_verify_action_flow_training_smoke.py
deleted file mode 100644
index 3505141301..0000000000
--- a/tests/test_verify_action_flow_training_smoke.py
+++ /dev/null
@@ -1,1028 +0,0 @@
-from __future__ import annotations
-
-import hashlib
-import importlib.util
-import json
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.core.hydra_config import HydraConfig
-from omegaconf import OmegaConf, open_dict
-
-from tools.validate_action_flow_config import PreflightError, compose_experiment
-
-SCRIPT = (
- Path(__file__).parents[1]
- / "scripts"
- / "train"
- / "verify_action_flow_training_smoke.py"
-)
-SPEC = importlib.util.spec_from_file_location(
- "verify_action_flow_training_smoke", SCRIPT
-)
-MODULE = importlib.util.module_from_spec(SPEC)
-assert SPEC.loader is not None
-sys.modules[SPEC.name] = MODULE
-SPEC.loader.exec_module(MODULE)
-
-HEAD = "a" * 40
-
-
-def _resolved_smoke_config(
- tmp_path: Path,
- *,
- reconstruction_weight: float = 1.0,
- experiment: str | None = None,
-):
- if experiment is None:
- suffix = {1.0: "1", 10.0: "10", 100.0: "100"}[reconstruction_weight]
- experiment = f"sim/pusht/action_flow_bc_usocket_recon{suffix}_s42"
- cfg = compose_experiment(experiment)
- run_dir = tmp_path / "run"
- run_dir.mkdir()
- normalization = tmp_path / "normalization.pt"
- normalization.write_bytes(b"train-only-normalization")
- normalization_hash = hashlib.sha256(normalization.read_bytes()).hexdigest()
-
- with open_dict(cfg):
- cfg.trainer.max_steps = 2
- cfg.trainer.val_check_interval = 1
- cfg.trainer.limit_val_batches = 1
- cfg.trainer.log_every_n_steps = 1
- cfg.trainer.precision = "bf16"
- cfg.callbacks.model_checkpoint.every_n_train_steps = 1
- cfg.model.gradient_telemetry_cadence = 2
- cfg.norm_stats.precomputed_norm_path = str(normalization)
- cfg.run_provenance.source_commit = HEAD
- cfg.run_provenance.normalization_sha256 = normalization_hash
- cfg.evaluator.artifact_root = str(run_dir / "energy")
- cfg.evaluator.action_flow_diagnostics.artifact_root = str(
- run_dir / "diagnostics"
- )
- cfg.hydra.runtime.cwd = str(MODULE.REPOSITORY_ROOT)
- cfg.hydra.runtime.output_dir = str(run_dir)
- HydraConfig.instance().set_config(cfg)
- selected = OmegaConf.masked_copy(cfg, [key for key in cfg.keys() if key != "hydra"])
- resolved = OmegaConf.create(OmegaConf.to_container(selected, resolve=True))
- config_path = run_dir / ".hydra/config.yaml"
- config_path.parent.mkdir()
- OmegaConf.save(resolved, config_path)
- return experiment, run_dir, config_path, normalization_hash
-
-
-def test_config_gate_accepts_only_exact_two_step_contract(tmp_path, monkeypatch):
- experiment, run_dir, config_path, normalization_hash = _resolved_smoke_config(
- tmp_path
- )
- monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD)
-
- _, identities = MODULE._validate_config(
- config_path=config_path,
- experiment=experiment,
- run_dir=run_dir,
- expected_head=HEAD,
- expected_config_sha256=hashlib.sha256(config_path.read_bytes()).hexdigest(),
- expected_split_sha256=None,
- expected_normalization_sha256=normalization_hash,
- )
-
- assert identities["repo_head"] == HEAD
- assert identities["normalization_sha256"] == normalization_hash
- assert (
- identities["config_sha256"]
- == hashlib.sha256(config_path.read_bytes()).hexdigest()
- )
-
-
-def test_config_gate_accepts_option_a_200m_muon_contract(tmp_path, monkeypatch):
- experiment, run_dir, config_path, normalization_hash = _resolved_smoke_config(
- tmp_path,
- experiment=(
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"
- ),
- )
- monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD)
-
- config, _ = MODULE._validate_config(
- config_path=config_path,
- experiment=experiment,
- run_dir=run_dir,
- expected_head=HEAD,
- expected_config_sha256=hashlib.sha256(config_path.read_bytes()).hexdigest(),
- expected_split_sha256=None,
- expected_normalization_sha256=normalization_hash,
- )
-
- assert config.model.pipeline.stages[5].field.hidden_dim == 1_024
- assert config.model.pipeline.stages[5].field.depth == 14
- assert config.model.optimizer.lr == pytest.approx(1.0e-5)
-
-
-def test_config_gate_accepts_option_a_200m_adamw_contract(tmp_path, monkeypatch):
- experiment, run_dir, config_path, normalization_hash = _resolved_smoke_config(
- tmp_path,
- experiment=(
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42"
- ),
- )
- monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD)
-
- config, _ = MODULE._validate_config(
- config_path=config_path,
- experiment=experiment,
- run_dir=run_dir,
- expected_head=HEAD,
- expected_config_sha256=hashlib.sha256(config_path.read_bytes()).hexdigest(),
- expected_split_sha256=None,
- expected_normalization_sha256=normalization_hash,
- )
-
- assert config.model.pipeline.stages[5].field.hidden_dim == 1_024
- assert config.model.pipeline.stages[5].field.depth == 14
- assert config.model.optimizer._target_ == "torch.optim.AdamW"
- assert config.model.optimizer.lr == pytest.approx(1.0e-5)
- assert config.model.optimizer_named_parameters is False
-
-
-@pytest.mark.parametrize(
- "experiment",
- [
- "pusht/action_flow_usocket_latent_fm_sg_unite_h384_s42",
- "pusht/action_flow_usocket_latent_fm_sg_unite_h384_sum14_cfg4_val8_s42",
- ],
-)
-def test_config_gate_rejects_retired_unite_h384_recipes(experiment):
- with pytest.raises(PreflightError, match="absent or retired"):
- compose_experiment(experiment)
-
-
-def test_unite_h384_parent_preserves_component_contract():
- # Direct component composition is not authorization for a fresh smoke/run.
- config_dir = MODULE.REPOSITORY_ROOT / "egomimic" / "hydra_configs"
- with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
- config = compose(
- config_name="train_zarr_cartesian",
- overrides=[
- "+experiment=sim/pusht/parents/action_flow_usocket_h384",
- "++paths.root_dir=.",
- "model.cfg_scale=4.0",
- "model.flow_loss_aggregation=sum_samples",
- "data.valid_dataloader_params.pushshapes_sim_u_socket.batch_size=32",
- ],
- )
- assert config.model.latent_dim == 16
- assert config.model.condition_dim == 128
- assert config.model.pipeline.stages[6].field.backbone.hidden_dim == 384
- assert config.model.pipeline.stages[6].field.backbone.depth == 12
- assert config.model.optimizer.lr == pytest.approx(1.0e-4)
- assert config.norm_stats.norm_mode == "minmax"
- assert config.model.pipeline.stages[6].cfg_scale == pytest.approx(4.0)
- assert config.model.pipeline.stages[8].flow_aggregation == "sum_samples"
- assert config.data.valid_dataloader_params.pushshapes_sim_u_socket.batch_size == 32
-
-
-@pytest.mark.parametrize(
- ("path", "value", "message"),
- [
- ("model.flow_samples_per_content", 1, "flow_samples_per_content"),
- ("model.gradient_telemetry_cadence", 1, "gradient_telemetry_cadence"),
- ("trainer.devices", 2, "trainer.devices"),
- ("callbacks.model_checkpoint.every_n_train_steps", 2, "every_n_train_steps"),
- ],
-)
-def test_config_gate_rejects_smoke_shortcuts(
- tmp_path, monkeypatch, path, value, message
-):
- experiment, run_dir, config_path, _ = _resolved_smoke_config(tmp_path)
- cfg = OmegaConf.load(config_path)
- OmegaConf.update(cfg, path, value)
- OmegaConf.save(cfg, config_path)
- monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD)
-
- with pytest.raises(MODULE.SmokeVerificationError, match=message):
- MODULE._validate_config(
- config_path=config_path,
- experiment=experiment,
- run_dir=run_dir,
- expected_head=HEAD,
- expected_config_sha256=None,
- expected_split_sha256=None,
- expected_normalization_sha256=None,
- )
-
-
-def _history_row():
- row = {}
- for name in (
- "TotalLoss",
- "FlowMatchingLoss",
- "ReconstructionLoss",
- "ReconstructionL1",
- "ActionVelocityLoss",
- ):
- train_value = 3.75 if name == "TotalLoss" else 1.25
- valid_value = 4.5 if name == "TotalLoss" else 1.5
- row[f"Train/ActionFlow/{name}_step"] = train_value
- row[f"Train/ActionFlow/{name}/{MODULE.SOURCE_LABEL}_step"] = train_value
- row[f"Valid/ActionFlow/{name}"] = valid_value
- for label in ("FM", "Reconstruction", "ActionVelocity"):
- row[f"Train/ActionFlow/GradientNorm/{label}"] = 0.5
- row[f"Train/ActionFlow/GradientParameterCount/{label}"] = 100.0
- for pair in (
- "FM__Reconstruction",
- "FM__ActionVelocity",
- "Reconstruction__ActionVelocity",
- ):
- row[f"Train/ActionFlow/GradientCosine/{pair}"] = 0.25
- row[f"Train/ActionFlow/GradientCosineDefined/{pair}"] = 1.0
- row[f"Train/ActionFlow/GradientIntersectionParameterCount/{pair}"] = 50.0
- row["Train/MSE"] = 1.25
- row[f"Train/MSE/{MODULE.SOURCE_LABEL}"] = 1.25
- row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 1.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 14.0
- row["Train/ActionFlow/Compute/DecoderJVPCallsPerStep"] = 1.0
- row["Train/ActionFlow/Compute/PeakAllocatedBytes"] = 1024.0
- row["Train/ActionFlow/Schedule/ReconstructionOnly"] = 0.0
- row["Train/ActionFlow/Schedule/EffectiveFlowWeight"] = 1.0
- row["Train/ActionFlow/Schedule/EffectiveActionVelocityWeight"] = 1.0
- for base in (
- "Valid/MSE",
- "Valid/Native_MSE",
- "Valid/EnergyScore@32",
- "Valid/EnergyScoreAccuracy@32",
- "Valid/EnergyScoreDiversity@32",
- ):
- row[base] = 0.75
- row[f"{base}/{MODULE.SOURCE_LABEL}"] = 0.75
- for name in (
- "Valid/ActionFlow/CleanReconstructionMSE",
- "Valid/ActionFlow/CleanReconstructionNativeMSE",
- "Valid/ActionFlow/Latent/clean/RMS",
- "Valid/ActionFlow/Latent/generated/EffectiveRank",
- "Valid/ActionFlow/DecoderJacobian/clean/spectral_norm_mean",
- "Valid/ActionFlow/DecodedFullNoise/token_radius_mean",
- "Valid/ActionFlow/DenoisingTrajectory/LatentMSE/t1000",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedMSE/t0000",
- "Valid/ActionFlow/DenoisingTrajectory/DecodedNativeMSE/t0000",
- "Valid/ActionFlow/Alignment/FinalLatentCosine/t1000",
- "Valid/ActionFlow/Alignment/CKA/encoder_00__field_00/t0500",
- "Valid/ActionFlow/Alignment/CKNNA/encoder_00__field_00/t0500",
- ):
- row[name] = 0.5
- if "NativeMSE" in name or "decoded_native_mse" in name:
- row[f"{name}/{MODULE.SOURCE_LABEL}"] = 0.5
- return row
-
-
-def test_history_gate_requires_components_gradients_and_scheduled_validation():
- result = MODULE._validate_history({0: {}, 2: _history_row()})
-
- assert result["train_step"] == 2
- assert result["valid_step"] == 2
- assert result["train"]["Train/ActionFlow/FlowMatchingLoss"] == 1.25
-
-
-def test_history_gate_accepts_float32_flow_weight_telemetry():
- row = _history_row()
- row["Train/ActionFlow/Schedule/EffectiveFlowWeight"] = float(torch.tensor(0.01))
- for suffix in ("", f"/{MODULE.SOURCE_LABEL}"):
- row[f"Train/ActionFlow/TotalLoss{suffix}_step"] = 2.5125
- row["Valid/ActionFlow/TotalLoss"] = 3.015
-
- result = MODULE._validate_history(
- {2: row}, reconstruction_weight=1.0, flow_weight=0.01
- )
-
- assert result["train_step"] == 2
-
-
-def test_unite_history_gate_requires_both_optimizer_family_rates():
- row = _history_row()
- row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 2.0
- row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 28.0
- row["Train/ActionFlow/GradientCosine/FM__Reconstruction"] = 0.0
- row["Train/ActionFlow/GradientCosineDefined/FM__Reconstruction"] = 0.0
- row["Train/ActionFlow/GradientIntersectionParameterCount/FM__Reconstruction"] = 0.0
- row["Optimizer/LR/AdamW"] = 2.5e-8
- row["Optimizer/LR/Muon"] = 2.5e-8
-
- result = MODULE._validate_history({2: row}, method=MODULE.STOPGRAD_UNITE_METHOD)
-
- assert result["train"]["Optimizer/LR/AdamW"] == pytest.approx(2.5e-8)
- assert result["train"]["Optimizer/LR/Muon"] == pytest.approx(2.5e-8)
-
- del row["Optimizer/LR/Muon"]
- with pytest.raises(MODULE.SmokeVerificationError, match="gradient telemetry"):
- MODULE._validate_history({2: row}, method=MODULE.STOPGRAD_UNITE_METHOD)
-
-
-def test_history_gate_rejects_missing_gradient_telemetry():
- row = _history_row()
- del row["Train/ActionFlow/GradientCosine/FM__ActionVelocity"]
-
- with pytest.raises(MODULE.SmokeVerificationError, match="gradient telemetry"):
- MODULE._validate_history({2: row})
-
-
-def _write_artifacts(tmp_path: Path):
- split_hash = "b" * 64
- norm_hash = "d" * 64
- content_hash = "e" * 64
- aggregate_hash = "f" * 64
- energy_root = tmp_path / "energy"
- diagnostic_root = tmp_path / "diagnostics"
- config_path = tmp_path / ".hydra/config.yaml"
- config_path.parent.mkdir()
- config_path.write_text("name: test\n")
- content_path = tmp_path / "content.json"
- content_path.write_text("{}\n")
- leaf = Path("epoch-0-step-2/rank-0-batch-0.pt")
- energy_path = energy_root / leaf
- diagnostic_path = diagnostic_root / leaf
- energy_path.parent.mkdir(parents=True)
- diagnostic_path.parent.mkdir(parents=True)
- seed_hash = "c" * 64
- targets = torch.zeros(16, 16, 4)
- condition_ids = [
- {
- "batch_position": index,
- "episode_hash": f"episode-{index:04d}",
- "frame_index": index,
- "normalized_target_sha256": MODULE._target_tensor_sha256(targets[index]),
- }
- for index in range(16)
- ]
- energy_view = {
- "definition": "first_deterministic_validation_batch_per_ddp_rank",
- "split_manifest_sha256": split_hash,
- "per_rank_batch_size": 16,
- "world_size": 1,
- }
- distance = MODULE.usocket_energy_distance_metadata(
- MODULE.USOCKET_ENERGY_DISTANCE_CONFIG
- )
- wandb = {"entity": "rl2-group", "project": "test", "run_id": "smoke"}
- energy_identity = {
- "schema": "energy-score-validation-artifact/v2",
- "metric": {
- "name": "EnergyScore@32",
- "sample_count": 32,
- "deterministic_seed": 420042,
- },
- "source_commit": HEAD,
- "normalization_sha256": norm_hash,
- "split_manifest_sha256": split_hash,
- "resolved_config_path": str(config_path.resolve()),
- "resolved_config_sha256": hashlib.sha256(config_path.read_bytes()).hexdigest(),
- "wandb": wandb,
- "dataset_content": {
- "manifest_path": str(content_path.resolve()),
- "manifest_sha256": content_hash,
- "aggregate_sha256": aggregate_hash,
- },
- "distance": distance,
- "seed_bank_sha256": seed_hash,
- "validation_view": energy_view,
- "validation_conditions": {MODULE.SOURCE_LABEL: condition_ids},
- "checkpoint_binding": {
- "global_step": 2,
- "checkpoint_sha256": None,
- "sha256_status": MODULE.ENERGY_CHECKPOINT_STATUS,
- },
- }
- torch.save(
- {
- "schema_version": 2,
- "metric": "EnergyScore@32",
- "sample_count": 32,
- "seed_bank": list(range(32)),
- "seed_bank_sha256": seed_hash,
- "deterministic_seed": 420042,
- "distance": distance,
- "global_step": 2,
- "validation_view": energy_view,
- "identity": energy_identity,
- "identity_sha256": MODULE._canonical_json_sha256(energy_identity),
- "domains": {
- MODULE.SOURCE_LABEL: {
- "predictions": torch.zeros(32, 16, 16, 4),
- "targets": targets,
- "native_predictions": torch.zeros(32, 16, 16, 3),
- "native_targets": torch.zeros(16, 16, 3),
- "condition_ids": condition_ids,
- }
- },
- },
- energy_path,
- )
- diagnostic_provenance = {
- "source_commit": HEAD,
- "normalization_sha256": norm_hash,
- "split_manifest_sha256": split_hash,
- "sampler_steps": 16,
- "dataset_content": {
- "manifest_sha256": content_hash,
- "aggregate_sha256": aggregate_hash,
- },
- }
- identity = {
- "schema": "action-flow-validation-diagnostics/v1",
- "validation_view": energy_view,
- "provenance": diagnostic_provenance,
- "native_error": MODULE.USOCKET_NATIVE_ERROR_CONFIG,
- }
- identity_hash = hashlib.sha256(
- json.dumps(
- identity,
- sort_keys=True,
- separators=(",", ":"),
- ensure_ascii=True,
- allow_nan=False,
- ).encode()
- ).hexdigest()
- torch.save(
- {
- "schema_version": 1,
- "metric": "ActionFlowValidationDiagnostics",
- "global_step": 2,
- "identity": identity,
- "identity_sha256": identity_hash,
- "sampler": {"name": "reverse_euler"},
- "validation_view": energy_view,
- "provenance": diagnostic_provenance,
- "statistics": {"native_action_error": MODULE.USOCKET_NATIVE_ERROR_CONFIG},
- "sources": {
- MODULE.SOURCE_LABEL: {
- "computed": {
- "clean_reconstruction_native_mse_by_condition": torch.ones(16),
- "trajectory_decoded_native_mse_by_condition": torch.ones(
- 17, 16
- ),
- "fixed_level_metrics": {
- label: {
- metric: torch.ones(16)
- for metric in (
- "state_decoded_native_mse",
- "predicted_clean_decoded_native_mse",
- "final_decoded_native_mse",
- )
- }
- for label in MODULE.NATIVE_LEVELS
- },
- }
- }
- },
- },
- diagnostic_path,
- )
- diagnostic_hash = hashlib.sha256(diagnostic_path.read_bytes()).hexdigest()
- Path(f"{diagnostic_path}.sha256").write_text(
- json.dumps(
- {
- "artifact": diagnostic_path.name,
- "identity_sha256": identity_hash,
- "sha256": diagnostic_hash,
- }
- )
- )
- config = OmegaConf.create(
- {
- "evaluator": {
- "artifact_root": str(energy_root),
- "seed_bank_sha256": seed_hash,
- "energy_score_validation_view": energy_view,
- "action_flow_diagnostics": {
- "artifact_root": str(diagnostic_root),
- "validation_view": energy_view,
- "provenance": {"sampler_steps": 16},
- },
- },
- "logger": {
- "wandb": {
- "entity": wandb["entity"],
- "project": wandb["project"],
- "id": wandb["run_id"],
- }
- },
- }
- )
- identities = {
- "config_sha256": hashlib.sha256(config_path.read_bytes()).hexdigest(),
- "content_manifest_path": str(content_path.resolve()),
- "content_manifest_sha256": content_hash,
- "dataset_content_aggregate_sha256": aggregate_hash,
- "normalization_sha256": norm_hash,
- "repo_head": HEAD,
- "split_manifest_sha256": split_hash,
- }
- checkpoint = {
- "checkpoint_sha256": "1" * 64,
- "file_size_bytes": 100,
- "global_step": 2,
- }
- return config, identities, checkpoint, diagnostic_path
-
-
-def test_artifact_gate_verifies_energy_and_diagnostic_immutability(tmp_path):
- config, identities, checkpoint, diagnostic_path = _write_artifacts(tmp_path)
-
- result = MODULE._validate_artifacts(
- config=config,
- run_dir=tmp_path,
- identities=identities,
- checkpoint=checkpoint,
- )
-
- assert result["energy_score"]["sha256"]
- assert result["action_flow_diagnostics"]["path"] == str(diagnostic_path)
-
-
-def test_artifact_gate_selects_exact_slurm_attempt_and_checks_execution(
- tmp_path, monkeypatch
-):
- root = tmp_path / "artifacts"
- for restart in (0, 1):
- artifact = (
- root / f"job-5714540-restart-{restart}" / "epoch-0-step-2/rank-0-batch-0.pt"
- )
- artifact.parent.mkdir(parents=True)
- torch.save(
- {
- "global_step": 2,
- "execution": {
- "slurm_job_id": "5714540",
- "slurm_restart_count": restart,
- },
- },
- artifact,
- )
- monkeypatch.setenv("SLURM_JOB_ID", "5714540")
- monkeypatch.setenv("SLURM_RESTART_COUNT", "1")
- selected, payload = MODULE._step_two_artifact(root, label="test")
- assert "job-5714540-restart-1" in str(selected)
- payload["execution"]["slurm_restart_count"] = 0
- torch.save(payload, selected)
- with pytest.raises(
- MODULE.SmokeVerificationError, match="execution identity mismatch"
- ):
- MODULE._step_two_artifact(root, label="test")
- monkeypatch.delenv("SLURM_JOB_ID")
- with pytest.raises(MODULE.SmokeVerificationError, match="expected one step-2"):
- MODULE._step_two_artifact(root, label="test")
-
-
-def test_artifact_gate_rejects_tampered_diagnostic_sidecar(tmp_path):
- config, identities, checkpoint, diagnostic_path = _write_artifacts(tmp_path)
- sidecar = Path(f"{diagnostic_path}.sha256")
- payload = json.loads(sidecar.read_text())
- payload["sha256"] = "0" * 64
- sidecar.write_text(json.dumps(payload))
-
- with pytest.raises(MODULE.SmokeVerificationError, match="sidecar mismatch"):
- MODULE._validate_artifacts(
- config=config,
- run_dir=tmp_path,
- identities=identities,
- checkpoint=checkpoint,
- )
-
-
-def test_artifact_gate_accepts_unique_origin_from_separate_validation_job(
- tmp_path, monkeypatch
-):
- root = tmp_path / "artifacts"
- artifact = root / "job-5714540-restart-1/epoch-0-step-2/rank-0-batch-0.pt"
- artifact.parent.mkdir(parents=True)
- torch.save(
- {
- "global_step": 2,
- "execution": {"slurm_job_id": "5714540", "slurm_restart_count": 1},
- },
- artifact,
- )
- monkeypatch.setenv("SLURM_JOB_ID", "5719999")
- monkeypatch.setenv("SLURM_RESTART_COUNT", "0")
- selected, payload = MODULE._step_two_artifact(root, label="test")
- assert selected == artifact
- assert payload["execution"]["slurm_job_id"] == "5714540"
-
-
-def test_preflight_gate_binds_source_experiment_split_and_normalization(tmp_path):
- experiment = "sim/pusht/action_flow_bc_usocket_recon1_s42"
- split_hash = "b" * 64
- normalization_hash = "c" * 64
- content_hash = "d" * 64
- aggregate_hash = "e" * 64
- destination = tmp_path / "provenance/restart-0/PREFLIGHT_RESULT.json"
- destination.parent.mkdir(parents=True)
- destination.write_text(
- json.dumps(
- {
- "status": "PASS",
- "source": {"head": HEAD},
- "experiment": experiment,
- "split_manifest_sha256": split_hash,
- "normalization_sha256": normalization_hash,
- "content_manifest_sha256": content_hash,
- "dataset_content_aggregate_sha256": aggregate_hash,
- },
- sort_keys=True,
- )
- )
- digest = hashlib.sha256(destination.read_bytes()).hexdigest()
-
- result = MODULE._validate_preflight(
- run_dir=tmp_path,
- expected_sha256=digest,
- expected_head=HEAD,
- experiment=experiment,
- split_sha256=split_hash,
- normalization_sha256=normalization_hash,
- content_manifest_sha256=content_hash,
- dataset_content_aggregate_sha256=aggregate_hash,
- )
-
- assert result == {"path": str(destination), "sha256": digest, "copy_count": 1}
-
-
-def test_cli_accepts_launcher_compatibility_names():
- args = MODULE._parser().parse_args(
- [
- "/tmp/run",
- "--expected-head",
- HEAD,
- "--expected-experiment",
- "sim/pusht/action_flow_bc_usocket_recon10_s42",
- "--expected-reconstruction-weight",
- "10",
- "--expected-preflight-sha256",
- "d" * 64,
- ]
- )
-
- assert args.experiment.endswith("recon10_s42")
- assert args.expected_reconstruction_weight == 10.0
-
-
-def test_cli_accepts_codec98k_experiment():
- args = MODULE._parser().parse_args(
- [
- "/tmp/run",
- "--expected-head",
- HEAD,
- "--expected-experiment",
- "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_codec98k_s42",
- "--expected-preflight-sha256",
- "d" * 64,
- ]
- )
-
- assert args.experiment.endswith("recon1_codec98k_s42")
-
-
-def test_cli_accepts_option_a_200m_muon_experiment():
- args = MODULE._parser().parse_args(
- [
- "/tmp/run",
- "--expected-head",
- HEAD,
- "--expected-experiment",
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42",
- "--expected-preflight-sha256",
- "d" * 64,
- ]
- )
-
- assert args.experiment.endswith("recon1_200m_muon_lr1e5_s42")
-
-
-def test_cli_accepts_option_a_200m_adamw_experiment():
- args = MODULE._parser().parse_args(
- [
- "/tmp/run",
- "--expected-head",
- HEAD,
- "--expected-experiment",
- "sim/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_adamw_lr1e5_s42",
- "--expected-preflight-sha256",
- "d" * 64,
- ]
- )
-
- assert args.experiment.endswith("recon1_200m_adamw_lr1e5_s42")
-
-
-def test_gpu_probe_gate_requires_real_single_h100_or_h200_bf16(tmp_path):
- path = tmp_path / "provenance/restart-0/gpu_probe.json"
- path.parent.mkdir(parents=True)
- path.write_text(
- json.dumps(
- {
- "status": "PASSED",
- "world_size": 1,
- "gpu_name": "NVIDIA H200",
- "bf16_supported": True,
- "bf16_forward_backward": {
- "dtype": "torch.bfloat16",
- "gradient_dtype": "torch.bfloat16",
- "finite": True,
- },
- "nccl": {"world_size": 1, "all_reduce": 1.0, "destroyed": True},
- }
- )
- )
-
- records = MODULE._validate_gpu_probes(tmp_path)
-
- assert records[0]["gpu_name"] == "NVIDIA H200"
-
-
-def test_gpu_probe_gate_rejects_non_target_gpu(tmp_path):
- path = tmp_path / "provenance/restart-0/gpu_probe.json"
- path.parent.mkdir(parents=True)
- path.write_text(
- json.dumps(
- {
- "status": "PASSED",
- "world_size": 1,
- "gpu_name": "NVIDIA A100-SXM4-80GB",
- "bf16_supported": True,
- "bf16_forward_backward": {
- "dtype": "torch.bfloat16",
- "gradient_dtype": "torch.bfloat16",
- "finite": True,
- },
- "nccl": {"world_size": 1, "all_reduce": 1.0, "destroyed": True},
- }
- )
- )
-
- with pytest.raises(MODULE.SmokeVerificationError, match="H100 or H200"):
- MODULE._validate_gpu_probes(tmp_path)
-
-
-def _scaled_muon_optimizer_state():
- return {
- "adamw": {
- "state": {0: {"exp_avg": torch.zeros(1)}},
- "param_groups": [{"lr": 1.0e-5, "params": [0]}],
- },
- "muon": {
- "state": {1: {"momentum_buffer": torch.zeros(1)}},
- "param_groups": [{"lr": 1.0e-5, "params": [1]}],
- },
- "group_manifest": {
- "adamw_parameter_names": ("encoder.bias",),
- "muon_parameter_names": ("field.weight",),
- "muon_source_commit": "a" * 40,
- },
- }
-
-
-def test_optimizer_state_gate_accepts_scaled_muon_composite_state():
- config = OmegaConf.create(
- {"name": "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"}
- )
-
- MODULE._validate_optimizer_state(_scaled_muon_optimizer_state(), config)
-
-
-@pytest.mark.parametrize("group", ["adamw", "muon"])
-def test_optimizer_state_gate_rejects_empty_scaled_muon_nested_state(group):
- config = OmegaConf.create(
- {"name": "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"}
- )
- optimizer_state = _scaled_muon_optimizer_state()
- optimizer_state[group]["state"] = {}
-
- with pytest.raises(
- MODULE.SmokeVerificationError, match=f"{group} optimizer state is empty"
- ):
- MODULE._validate_optimizer_state(optimizer_state, config)
-
-
-def _checkpoint_payload():
- return {
- "global_step": 2,
- "state_dict": {"model.weight": torch.ones(1)},
- "optimizer_states": [
- {
- "state": {0: {"exp_avg": torch.zeros(1)}},
- "param_groups": [{"lr": 3.0e-5, "params": [0]}],
- }
- ],
- "lr_schedulers": [{"last_epoch": 2, "_last_lr": [3.0e-6]}],
- "loops": {"fit_loop": {"completed": 2}},
- }
-
-
-def test_checkpoint_gate_strictly_reloads_exact_action_flow_wrapper(
- tmp_path, monkeypatch
-):
- checkpoint_dir = tmp_path / "checkpoints"
- checkpoint_dir.mkdir()
- payload = _checkpoint_payload()
- immutable = checkpoint_dir / "epoch-epoch=0-step-step=2.ckpt"
- torch.save(payload, immutable)
- (checkpoint_dir / "last.ckpt").symlink_to(immutable.name)
-
- class _Parameter:
- @staticmethod
- def numel():
- return MODULE.EXPECTED_PARAMETER_COUNT
-
- class _Wrapper:
- encoder_e = object()
- field_v = object()
- decoder_g = object()
-
- class _Nets:
- @staticmethod
- def named_parameters(**_kwargs):
- return []
-
- nets = _Nets()
-
- @classmethod
- def load_from_checkpoint(cls, *_args, **_kwargs):
- return cls()
-
- @staticmethod
- def parameters():
- return [_Parameter()]
-
- monkeypatch.setattr(MODULE, "ActionFlowModelWrapper", _Wrapper)
- monkeypatch.setattr(
- MODULE,
- "_validate_gradient_route_manifest",
- lambda *_args: {"manifest_sha256": "f" * 64},
- )
-
- result = MODULE._validate_checkpoint(
- tmp_path, reconstruction_weight=1.0, flow_weight=1.0
- )
-
- assert result["global_step"] == 2
- assert result["parameter_count"] == MODULE.EXPECTED_PARAMETER_COUNT
- assert result["strict_checkpoint_reload"] == "passed"
- assert result["file_size_bytes"] > 0
-
-
-def _gradient_manifest(named_parameters):
- route_names = {
- "FM": [name for name, _ in named_parameters[:3]],
- "Reconstruction": [named_parameters[1][0], named_parameters[3][0]],
- "ActionVelocity": [name for name, _ in named_parameters],
- }
- routes = {}
- route_hashes = {}
- for label, names in route_names.items():
- parameters = dict(named_parameters)
- route = [
- {
- "dtype": str(parameters[name].dtype),
- "name": name,
- "numel": parameters[name].numel(),
- "shape": list(parameters[name].shape),
- }
- for name in names
- ]
- routes[label] = route
- route_hashes[label] = MODULE._canonical_json_sha256(route)
- intersections = {
- "FM__Reconstruction": [named_parameters[1][0]],
- "FM__ActionVelocity": [name for name, _ in named_parameters[:3]],
- "Reconstruction__ActionVelocity": [
- named_parameters[1][0],
- named_parameters[3][0],
- ],
- }
- core = {
- "routes": routes,
- "route_sha256": route_hashes,
- "intersections": intersections,
- "schema_version": 1,
- }
- return {**core, "manifest_sha256": MODULE._canonical_json_sha256(core)}
-
-
-def test_gradient_route_manifest_binds_exact_named_shared_paths():
- named = (
- ("nets.pipeline.stages.0.observation", torch.nn.Parameter(torch.ones(2))),
- ("nets.pipeline.stages.3.encoder", torch.nn.Parameter(torch.ones(3))),
- ("nets.pipeline.stages.5.field", torch.nn.Parameter(torch.ones(4))),
- ("nets.pipeline.stages.6.decoder", torch.nn.Parameter(torch.ones(5))),
- )
- manifest = _gradient_manifest(named)
-
- result = MODULE._validate_gradient_route_manifest(manifest, named)
-
- assert result["manifest_sha256"] == manifest["manifest_sha256"]
- assert result["route_parameter_counts"] == {
- "FM": 9,
- "Reconstruction": 8,
- "ActionVelocity": 14,
- }
- assert result["shared_path_parameter_counts"]["FM__Reconstruction"] == 3
-
-
-def test_gradient_route_manifest_rejects_tampered_named_route():
- named = (
- ("nets.pipeline.stages.0.observation", torch.nn.Parameter(torch.ones(2))),
- ("nets.pipeline.stages.3.encoder", torch.nn.Parameter(torch.ones(3))),
- ("nets.pipeline.stages.5.field", torch.nn.Parameter(torch.ones(4))),
- ("nets.pipeline.stages.6.decoder", torch.nn.Parameter(torch.ones(5))),
- )
- manifest = _gradient_manifest(named)
- manifest["routes"]["FM"][0]["numel"] = 99
-
- with pytest.raises(MODULE.SmokeVerificationError, match="size mismatch"):
- MODULE._validate_gradient_route_manifest(manifest, named)
-
-
-def test_checkpoint_gate_rejects_missing_scheduler(tmp_path):
- checkpoint_dir = tmp_path / "checkpoints"
- checkpoint_dir.mkdir()
- payload = _checkpoint_payload()
- payload["lr_schedulers"] = []
- immutable = checkpoint_dir / "epoch-epoch=0-step-step=2.ckpt"
- torch.save(payload, immutable)
- (checkpoint_dir / "last.ckpt").symlink_to(immutable.name)
-
- with pytest.raises(MODULE.SmokeVerificationError, match="scheduler state"):
- MODULE._validate_checkpoint(
- tmp_path, reconstruction_weight=1.0, flow_weight=1.0
- )
-
-
-def _raw_multiplier_verifier():
- import importlib.util
-
- path = (
- Path(__file__).resolve().parents[1]
- / "scripts/train/verify_scaled_action_flow_training_smoke.py"
- )
- spec = importlib.util.spec_from_file_location("raw_multiplier_verifier", path)
- module = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(module)
- return module
-
-
-@pytest.mark.parametrize(
- "damage", [None, "missing", "decay", "updates", "shape", "nan"]
-)
-def test_raw_multiplier_native_ema_verification(damage):
- from types import SimpleNamespace
-
- from omegaconf import OmegaConf
-
- from egomimic.utils.ema_callback import EMACallback
-
- verifier = _raw_multiplier_verifier()
- row = verifier.EXPERIMENTS[
- "sim/pusht/action_flow_cotrain_uc_multiplier_interpolation"
- ]
- assert row["parameter_count"] == 471549402 and row["hidden_dim"] == 696
- model = torch.nn.Linear(2, 2)
- callback = EMACallback(decay=0.9978)
- trainer = SimpleNamespace(global_step=0)
- callback.on_fit_start(trainer, model)
- callback._update(model)
- callback._update(model)
- checkpoint = {"global_step": 2, "state_dict": model.state_dict()}
- callback.on_save_checkpoint(trainer, model, checkpoint)
- cfg = OmegaConf.create(
- {
- "callbacks": {
- "ema": {
- "_target_": "egomimic.utils.ema_callback.EMACallback",
- "decay": 0.9978,
- "validate_with_ema": True,
- "use_warmup": False,
- }
- }
- }
- )
- if damage == "missing":
- checkpoint.pop("ema_state_dict")
- elif damage == "decay":
- checkpoint["ema_config"]["decay"] = 0.9
- elif damage == "updates":
- checkpoint["ema_num_updates"] = 1
- elif damage == "shape":
- checkpoint["ema_state_dict"]["weight"] = torch.zeros(3, 3)
- elif damage == "nan":
- checkpoint["ema_state_dict"]["weight"].fill_(float("nan"))
- if damage:
- with pytest.raises((RuntimeError, verifier.VerificationError)):
- verifier.validate_raw_multiplier_ema(cfg, model, checkpoint)
- else:
- assert (
- verifier.validate_raw_multiplier_ema(cfg, model, checkpoint)[
- "strict_ema_reload"
- ]
- == "passed"
- )
diff --git a/tests/test_wandb_health.py b/tests/test_wandb_health.py
deleted file mode 100644
index 4c1dc3b70c..0000000000
--- a/tests/test_wandb_health.py
+++ /dev/null
@@ -1,120 +0,0 @@
-import importlib.util
-from pathlib import Path
-
-import pytest
-
-PATH = Path(__file__).parents[1] / "scripts/train/check_wandb_health.py"
-SPEC = importlib.util.spec_from_file_location("check_wandb_health", PATH)
-health = importlib.util.module_from_spec(SPEC)
-SPEC.loader.exec_module(health)
-
-
-class FakeRun:
- path = ["team", "project", "run"]
- url = "https://wandb.ai/team/project/runs/run"
- state = "running"
- lastHistoryStep = 20545
- config = {"run_provenance": {"source_commit": "exact-source"}}
-
- def __init__(self, rows, summary=None):
- self.rows = rows
- self.summary = summary or {}
- self.scan_arguments = None
-
- def scan_history(self, **kwargs):
- self.scan_arguments = kwargs
- assert kwargs["keys"] is None # Sparse keys must never be intersected.
- return iter(self.rows)
-
-
-class FakeApi:
- def __init__(self, run):
- self.run_object = run
-
- def run(self, path):
- assert path == "team/project/run"
- return self.run_object
-
-
-def inspect(run, **kwargs):
- return health.collect_health(FakeApi(run), "team/project/run", **kwargs)
-
-
-def test_sparse_latest_metrics_and_optimizer_axis_use_bounded_history():
- run = FakeRun(
- [
- {"_step": 20544, "trainer/global_step": 200000, "Train/loss": 0.2},
- {"_step": 20545, "trainer/global_step": 200010, "Schedule/flow": 0.01},
- ]
- )
- result = inspect(
- run,
- metrics=["Train/loss", "Schedule/flow"],
- window=20,
- expected_config={"run_provenance.source_commit": "exact-source"},
- )
- assert result["status"] == "PASS"
- assert result["optimizer_step"] == 200010
- assert result["wandb_history_step"] == 20545
- assert result["metrics"]["Train/loss"]["optimizer_step"] == 200000
- assert run.scan_arguments == {
- "keys": None,
- "min_step": 20526,
- "max_step": 20546,
- "page_size": 20,
- }
-
-
-def test_latest_nan_is_not_hidden_by_an_earlier_finite_value():
- run = FakeRun(
- [
- {"_step": 20544, "trainer/global_step": 20, "loss": 0.1},
- {"_step": 20545, "trainer/global_step": 30, "loss": float("nan")},
- ]
- )
- result = inspect(run, metrics=["loss"])
- assert result["status"] == "ERROR"
- assert result["metrics"]["loss"]["finite"] is False
-
-
-def test_summary_only_metric_is_reported_without_inventing_a_step():
- run = FakeRun([{"_step": 20545, "trainer/global_step": 20}], {"Valid/loss": 0.3})
- result = inspect(run, metrics=["Valid/loss"])
- assert result["status"] == "WARN"
- assert result["metrics"]["Valid/loss"]["source"] == "summary_only"
- assert result["metrics"]["Valid/loss"]["optimizer_step"] is None
-
-
-def test_missing_and_wrong_config_identity_fail_closed():
- run = FakeRun([{"trainer/global_step": 20}])
- result = inspect(
- run, expected_config={"run_provenance.source_commit": "wrong", "missing": 1}
- )
- assert result["status"] == "ERROR"
- assert len(result["errors"]) == 2
-
-
-def test_history_counter_is_never_substituted_for_missing_optimizer_step():
- result = inspect(FakeRun([{"_step": 20545, "loss": 0.1}]), metrics=["loss"])
- assert result["status"] == "INCOMPLETE"
- assert result["optimizer_step"] is None
-
-
-def test_recent_resume_regression_and_stale_metrics_are_visible():
- run = FakeRun(
- [
- {"_step": 20544, "trainer/global_step": 200000, "_timestamp": 10},
- {"_step": 20545, "trainer/global_step": 160000, "_timestamp": 20},
- ]
- )
- result = inspect(run, max_age_seconds=30, now=100)
- assert result["status"] == "WARN"
- assert result["optimizer_step_regressions"] == [
- {"from": 200000, "to": 160000, "wandb_step": 20545}
- ]
- assert result["latest_history_age_seconds"] == 80
-
-
-def test_no_unbounded_history_request():
- with pytest.raises(ValueError, match="window"):
- inspect(FakeRun([]), window=1000000)
diff --git a/tests/test_yam_gpu_uuid_contract.py b/tests/test_yam_gpu_uuid_contract.py
deleted file mode 100644
index be82ec9e1d..0000000000
--- a/tests/test_yam_gpu_uuid_contract.py
+++ /dev/null
@@ -1,67 +0,0 @@
-"""CPU-only contract checks; no CUDA or dataset initialization."""
-
-import ast
-import re
-from pathlib import Path
-
-import pytest
-
-
-def helper():
- path = Path(__file__).resolve().parents[1] / "scripts/ice/ice_gpu_probe.py"
- tree = ast.parse(path.read_text())
- function = next(
- n
- for n in tree.body
- if isinstance(n, ast.FunctionDef) and n.name == "nvml_gpu_uuid"
- )
- namespace = {"re": re}
- exec(
- compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"),
- namespace,
- )
- return namespace["nvml_gpu_uuid"]
-
-
-@pytest.mark.parametrize("prefix", ["", "GPU-"])
-def test_actual_cuda_uuid_is_normalized(prefix):
- bare = "12345678-1234-1234-1234-123456789abc"
- assert helper()(prefix + bare) == "GPU-" + bare
-
-
-@pytest.mark.parametrize(
- "value",
- [
- "0",
- "MIG-1234",
- "GPU-",
- None,
- b"1234",
- "GPU-12345678-1234-1234-1234-123456789abc; x",
- ],
-)
-def test_physical_indices_or_unknown_uuid_rejected(value):
- with pytest.raises(RuntimeError):
- helper()(value)
-
-
-@pytest.mark.parametrize("text", ["f8ae3ac1-0e5a-3e23-c619-d06dd44000a2", "invalid"])
-def test_torch_cuuuid_wrapper_keeps_strict_uuid_validation(text):
- wrapper = type(
- "_CUuuid", (), {"__module__": "torch._C", "__str__": lambda self: text}
- )
- if text == "invalid":
- with pytest.raises(RuntimeError):
- helper()(wrapper())
- else:
- assert helper()(wrapper()) == "GPU-" + text
-
-
-def test_unrelated_stringifiable_object_is_rejected():
- wrapper = type(
- "OtherUUID",
- (),
- {"__str__": lambda self: "f8ae3ac1-0e5a-3e23-c619-d06dd44000a2"},
- )
- with pytest.raises(RuntimeError):
- helper()(wrapper())