System Design in Coding Questions

import torch, time, collections, math

SAFE_FRACTION = 0.8
MIN_BATCH = 1

class AdaptiveBatcher:
    def __init__(self, model, device):
        self.model = model; self.device = device
        self.buckets = collections.defaultdict(lambda: {"count":0,"mem":0.0,"time":0.0,"safe_budget":None,"backoff":0})
    def size_key(self, sample): return int(math.log2(max(1, sample.length)))  # example
    def update_stats(self, key, mem, t):
        b = self.buckets[key]; b["count"]+=1; b["mem"]+=mem; b["time"]+=t
        if b["safe_budget"] is None:
            free = torch.cuda.mem_get_info(self.device)[0]
            b["safe_budget"] = int(free * SAFE_FRACTION)
    def estimate_batch_size(self, key):
        b=self.buckets[key]
        if b["count"]==0: return MIN_BATCH
        avg = (b["mem"]/b["count"])
        return max(MIN_BATCH, int(b["safe_budget"] // avg))
    def form_batch(self, samples):
        grouped = collections.defaultdict(list)
        for s in samples: grouped[self.size_key(s)].append(s)
        batches=[]
        for k, items in grouped.items():
            bs = self.estimate_batch_size(k)
            for i in range(0, len(items), bs):
                batches.append(items[i:i+bs])
        return batches
    def run_batch(self, batch):
        try:
            torch.cuda.reset_peak_memory_stats(self.device)
            t0=time.time()
            out = self.model(batch)  # preprocess/pad as needed
            torch.cuda.synchronize()
            t=time.time()-t0
            mem = torch.cuda.max_memory_allocated(self.device)
            self.update_stats(self.size_key(batch[0]), mem, t)
            return out
        except RuntimeError as e:
            if "out of memory" in str(e):
                k=self.size_key(batch[0])
                self.buckets[k]["safe_budget"] = int(self.buckets[k]["safe_budget"] * 0.8)
                self.buckets[k]["backoff"] += 1
                # fallback: retry with smaller batch
                if len(batch) > 1:
                    return self.run_batch(batch[:len(batch)//2])
            raise

# Usage: feed incoming samples into batcher.form_batch, run_batch for each batch, monitor throughput and adapt.