test_composite.py

1import unittest,tempfile,copy
2import numpy as np
3import torch
4from unittest.mock import patch
5import composite as c
6from records import DTYPE
7SD=np.dtype([('kind','u1'),('source','u1'),('row','<u4'),('exposure','u1')])
8class Data:
9    def __init__(self,n):
10        self.schedule=np.zeros(n,dtype=SD);self.schedule['source']=np.arange(n)%2;self.schedule['row']=np.arange(n)//2
11        self.sources={('ordinary',s):range((n+1-s)//2) for s in (0,1)};self.sources.update({('theory',s):[] for s in (0,1)})
12    def take(self,ids):
13        rows=np.zeros(len(ids),dtype=DTYPE);rows['source']=self.schedule[ids]['source'];return rows
14class Tests(unittest.TestCase):
15    def test_dummy_read_finishes_before_prefetch_initial_and_resume(self):
16        for cursor in (0,3):
17            events=[];d=Data(767);p=c.Plan(d);original=d.take
18            def take(ids):
19                self.assertNotIn('prefetch-start',events)
20                rows=original(ids);events.append('dummy-read-complete');return rows
21            d.take=take
22            def factory(plan,start,pin):
23                self.assertEqual(events,['dummy-read-complete'])
24                self.assertIs(plan,p);self.assertEqual(start,cursor);events.append('prefetch-start');return 'prefetch'
25            with patch.object(c,'Prefetch',side_effect=factory):
26                dummy,prefetch=c.start_prefetch(d,p,cursor,False)
27            self.assertEqual(len(dummy[0]),1);self.assertEqual(prefetch,'prefetch')
28            self.assertEqual(events,['dummy-read-complete','prefetch-start'])
29    def test_partial_166_and_empty_rank(self):
30        for n in (678,767,513,1):
31            d=Data(n);plans=[c.Plan(d,r,micro=64) for r in (0,1)]
32            indices=[]
33            for step in range(plans[0].steps):
34                a,b=[p.indices(step) for p in plans]
35                self.assertFalse(set(a)&set(b));indices.extend(a);indices.extend(b)
36                self.assertEqual(len(a)+len(b),plans[0].global_decisions(step))
37            self.assertEqual(sorted(indices),list(range(n)))
38    def test_prefetch_resume_and_counters(self):
39        d=Data(678);p=c.Plan(d,0,micro=64);f=c.Prefetch(p,3,False)
40        try:
41            for step in range(3,p.steps):
42                batch,tags,seconds=f.take(step);self.assertEqual(len(batch[0]),len(p.indices(step)))
43        finally:f.close()
44        ledger=c.ExposureLedger(d);v=ledger.snapshot(512);self.assertEqual(ledger.verify(v),512)
45        v['counts'][0]+=1
46        with self.assertRaises(ValueError):ledger.verify(v)
47    def test_partial_gradient_weight_matches_global(self):
48        torch.manual_seed(9);x=torch.randn(166,3);target=torch.randn(166,2);base=torch.nn.Linear(3,2)
49        expected=copy.deepcopy(base);torch.nn.functional.mse_loss(expected(x),target,reduction='sum').div(166).backward()
50        grads=[]
51        for rank in (0,1):
52            model=copy.deepcopy(base);ids=torch.arange(rank,166,2);loss=torch.nn.functional.mse_loss(model(x[ids]),target[ids],reduction='sum')*2/166;loss.backward();grads.append([p.grad for p in model.parameters()])
53        for j,p in enumerate(expected.parameters()):torch.testing.assert_close((grads[0][j]+grads[1][j])/2,p.grad,rtol=1e-5,atol=1e-6)
54    def test_invalid_and_repeated_ordinary_refused(self):
55        for field,value in (('row',1000),('kind',2),('source',2)):
56            d=Data(10);d.schedule[field][0]=value
57            with self.assertRaises(ValueError):c.ExposureLedger(d).counts(10)
58        d=Data(10);d.schedule[2]=d.schedule[0]
59        with self.assertRaisesRegex(ValueError,'overexposure'):c.ExposureLedger(d).counts(10)
60if __name__=='__main__':unittest.main(verbosity=2)
61