test_components.py

1import unittest,tempfile,copy,sys,random
2from pathlib import Path
3import numpy as np
4import torch
5import schedule as s
6import checkpoint as ck
7sys.path.insert(0,str(Path(__file__).resolve().parent.parent/'state-v2'))
8from records import DTYPE
9class Tests(unittest.TestCase):
10    def fixture(self):
11        m=np.zeros(26,dtype=s.META);m['source'][12:24]=1;m['kind'][24:]=[1,2];m['source'][25]=1;m['phase']=np.arange(26)%3
12        offsets=np.r_[np.arange(25),27,30].astype(np.uint64)
13        return m,offsets
14    def test_unique_counts_spacing_and_determinism(self):
15        m,o=self.fixture()
16        with tempfile.TemporaryDirectory() as t:
17            a,b=Path(t)/'a',Path(t)/'b';facts=s.build(m,o,a);s.build(m,o,b)
18            self.assertEqual(a.read_bytes(),b.read_bytes());plan=np.memmap(a,dtype=s.EXPOSURE,mode='r')
19            self.assertEqual(facts['decision_exposures'],30)
20            np.testing.assert_array_equal(s.counters(plan,len(plan),len(m)),[1]*24+[3]*2)
21            for state in (24,25):
22                positions=np.flatnonzero(plan['record']==state);self.assertTrue(np.all(np.diff(positions)>=8))
23                self.assertEqual(list(plan[positions]['ordinal']),[0,1,2]);self.assertEqual(len(set(plan[positions]['tie'])),3)
24            a0=s.DistributedPlan(plan,0,2,4);a1=s.DistributedPlan(plan,1,2,4)
25            merged=np.concatenate([a.exposures(i) for i in range(a0.steps) for a in (a0,a1)])
26            np.testing.assert_array_equal(merged,plan)
27    def test_teacher_prefix_rotates_and_resume_rejects(self):
28        m,o=self.fixture()
29        with tempfile.TemporaryDirectory() as t:
30            p=Path(t)/'s';s.build(m,o,p);plan=np.memmap(p,dtype=s.EXPOSURE,mode='r')
31            ties=np.tile([12,28,64],(30,1)).astype(np.uint16);ties[24:]=np.array([[1,18,64],[6,21,64],[11,27,64]]*2)
32            e=plan[-3:];rows=np.zeros(len(e),dtype=DTYPE);v=s.select_targets(rows,e,o,ties)
33            np.testing.assert_array_equal(v['x'][:,120:122],v['y'][:,:2])
34            saved={'format':'cm6-curriculum-checkpoint-v1','identity':{'a':1},'terminal':False,'deadline':{'absolute':3},'rank_rng':[{},{}],'exposures':ck.exposure_state(plan,13,len(m))}
35            self.assertEqual(ck.validate_resume(saved,{'a':1},plan,len(m),{'absolute':3}),13)
36            for key,value in [('terminal',True),('deadline',{'absolute':4}),('rank_rng',[{}])]:
37                bad=dict(saved,**{key:value})
38                with self.assertRaises(ValueError):ck.validate_resume(bad,{'a':1},plan,len(m),{'absolute':3})
39            bad=copy.deepcopy(saved);bad['exposures']['counts'][0]+=1
40            with self.assertRaises(ValueError):ck.validate_resume(bad,{'a':1},plan,len(m),{'absolute':3})
41    def test_cpu_optimizer_rng_resume(self):
42        torch.set_num_threads(2);torch.manual_seed(8);model=torch.nn.Linear(3,2);opt=torch.optim.AdamW(model.parameters(),lr=.01)
43        def step():
44            opt.zero_grad();x=torch.randn(4,3);loss=model(x).square().mean();loss.backward();opt.step()
45        step();state=copy.deepcopy((model.state_dict(),opt.state_dict()));rng=ck.capture_rng();step();expected=copy.deepcopy(model.state_dict())
46        model.load_state_dict(state[0]);opt.load_state_dict(state[1]);ck.restore_rng(rng);step()
47        for k,v in model.state_dict().items():torch.testing.assert_close(v,expected[k],rtol=0,atol=0)
48    def test_reject_unequal_or_unknown_labels(self):
49        m,o=self.fixture();m['source'][0]=1
50        with self.assertRaises(ValueError):s.validate(m,o)
51        m,o=self.fixture();m['kind'][0]=3
52        with self.assertRaises(ValueError):s.validate(m,o)
53if __name__=='__main__':unittest.main(verbosity=2)
54