batches.py

1"""Exactly balanced positions, vectorized disk gather, bounded CPU pinned prefetch."""
2from concurrent.futures import ThreadPoolExecutor
3import math,time
4import numpy as np
5import torch
6
7def permutation(n,seed):
8    # Bijective affine permutation with O(1) state; no repeated records in a pass.
9    rng=np.random.default_rng(seed);a=int(rng.integers(1,n)) if n>1 else 1
10    while math.gcd(a,n)!=1:a=(a+1)%n or 1
11    return a,int(rng.integers(n))
12
13class Balanced:
14    def __init__(self,main,elite,seed,rank=0,world=2,micro=32):
15        if len(main)!=len(elite) or len(main)<1 or micro%2 or not 0<=rank<world:raise ValueError('equal source counts/even microbatch/rank required')
16        self.sources=(main,elite);self.n=len(main);self.rank=rank;self.world=world;self.micro=micro
17        self.perms=[permutation(self.n,seed+i) for i in range(2)]
18        self.steps=math.ceil(self.n/(world*(micro//2)))
19    def rows(self,step):
20        if not 0<=step<self.steps:raise IndexError('batch cursor')
21        half=self.micro//2;start=(step*self.world+self.rank)*half
22        q=np.arange(start,min(start+half,self.n),dtype=np.int64)
23        indices=[(a*q+b)%self.n for a,b in self.perms]
24        parts=[source.take(ids) for source,ids in zip(self.sources,indices)]
25        # Alternating source records preserve exactly 50/50 even on final rank.
26        out=np.empty(len(q)*2,dtype=parts[0].dtype);out[::2]=parts[0];out[1::2]=parts[1]
27        return out
28    def global_decisions(self,step):return 2*max(0,min(self.world*(self.micro//2),self.n-step*self.world*(self.micro//2)))
29
30def tensors(rows,pin=False):
31    # Three bulk allocations/conversions, never per-row tensor construction.
32    arrays=[np.ascontiguousarray(rows['x'],dtype=np.int64),np.ascontiguousarray(rows['y'],dtype=np.int64),np.ascontiguousarray(rows['attention'],dtype=np.bool_)]
33    values=[torch.from_numpy(a) for a in arrays]
34    return tuple(v.pin_memory() for v in values) if pin else tuple(values)
35
36class Prefetch:
37    def __init__(self,plan,start=0,pin=False):
38        self.plan=plan;self.pin=pin;self.pool=ThreadPoolExecutor(max_workers=1);self.futures={};self.next=start
39        for step in range(start,min(start+2,plan.steps)):self.futures[step]=self.pool.submit(self.load,step)
40    def load(self,step):
41        t=time.perf_counter();batch=tensors(self.plan.rows(step),self.pin)
42        return batch,time.perf_counter()-t
43    def take(self,step):
44        if step!=self.next:raise ValueError('ordered prefetch cursor')
45        result=self.futures.pop(step).result();self.next+=1
46        after=step+2
47        if after<self.plan.steps:self.futures[after]=self.pool.submit(self.load,after)
48        return result
49    def close(self):self.pool.shutdown(wait=True,cancel_futures=True)
50