from concurrent.futures import ThreadPoolExecutor import io import itertools import json from pathlib import Path import tempfile from threading import Event import time import unittest from unittest.mock import patch import numpy as np from tensor_pruning import prune_manifest, prune_queries from tensor_pruning.service import create_app from tensor_pruning.core import Projection, _kernel from fastapi.testclient import TestClient def reference(arrays, keep, moments): gains = [] offset = 0 for w, v in zip(arrays, moments): energies = w.astype(np.longdouble)**2 * v for head in energies: intermediate = [] for channel in head: mass = 0 for u, i in sorted((u, offset+i) for i, u in enumerate(channel)): intermediate.append((u * (2*mass+u), i)) mass += u offset += len(channel) mass = 0 for d, i in sorted(intermediate): gains.append((d * (2*mass+d), i)) mass += d keep_ids = {i for _, i in sorted(gains, reverse=True)[:keep]} offset = 0 masks = [] for w in arrays: masks.append(np.array([i in keep_ids for i in range(offset, offset+w.size)]).reshape(w.shape)) offset += w.size return masks def objective(arrays, masks, moments): return sum(np.sum(np.sum(np.sum(w.astype(np.longdouble)**2 * v * ~mask, axis=-1)**2, axis=-1)**2) for w, mask, v in zip(arrays, masks, moments)) class SelectorTest(unittest.TestCase): def test_mixed_shapes_and_ties_against_reference(self): rng = np.random.default_rng(17) for i in range(100): shapes = [(rng.integers(1, 6), rng.integers(1, 7), rng.integers(1, 14)) for _ in range(3)] weights = [rng.integers(-4, 5, size=shape).astype(dtype) for shape, dtype in zip(shapes, (np.float16, np.float32, np.float64))] moments = [rng.integers(0, 6, size=w.shape[-1]) for w in weights] total = sum(w.size for w in weights) keep = int(rng.integers(0, total+1)) actual = prune_queries(weights, keep, input_second_moments=moments, threads=2, seed=i) self.assertEqual(sum(np.count_nonzero(mask) for mask in actual), keep) for mask, expected in zip(actual, reference(weights, keep, moments)): np.testing.assert_array_equal(mask, expected) def test_exhaustive_global_optimum(self): rng = np.random.default_rng(18) for _ in range(18): weights = [rng.integers(-3, 4, size=shape).astype(np.float32) for shape in ((1, 2, 3), (2, 1, 2))] moments = [np.arange(1, w.shape[-1]+1) for w in weights] keep = int(rng.integers(0, 11)) got = prune_queries(weights, keep, input_second_moments=moments) best = float("inf") for ids in itertools.combinations(range(10), keep): flat = np.zeros(10, dtype=bool) flat[list(ids)] = True trial = [flat[:6].reshape(weights[0].shape), flat[6:].reshape(weights[1].shape)] best = min(best, objective(weights, trial, moments)) self.assertEqual(objective(weights, got, moments), best) def test_concurrent_calls_and_disk_workspace(self): rng = np.random.default_rng(19) weights = [rng.normal(size=shape).astype(dtype) for shape, dtype in [((12, 8, 65), np.float32), ((6, 5, 73), np.float64)]] moments = [rng.uniform(0.2, 2, size=w.shape[-1]) for w in weights] with tempfile.TemporaryDirectory() as scratch, ThreadPoolExecutor(3) as pool: def run(keep): return prune_queries(weights, keep, input_second_moments=moments, threads=2, scratch_dir=scratch) for keep, result in zip((2145, 3941, 6158), pool.map(run, (2145, 3941, 6158))): for mask, expected in zip(result, reference(weights, keep, moments)): np.testing.assert_array_equal(mask, expected) def test_native_bfloat16_bits(self): x = np.linspace(-6, 5, 72, dtype=np.float32).reshape(3, 4, 6) bits = (x.view(np.uint32) >> 16).astype(np.uint16) float_equivalent = (bits.astype(np.uint32) << 16).view(np.float32) moments = np.linspace(0.25, 1.25, x.shape[-1], dtype=np.float64) workspace = np.empty(x.size, dtype=np.uint32) mask = np.empty(x.shape, dtype=bool) desc = Projection(bits.ctypes.data, moments.ctypes.data, mask.ctypes.data, *x.shape, 3) self.assertEqual(_kernel()(desc, 1, workspace.ctypes.data, 29, 1, 0), 0) expected = prune_queries([float_equivalent], 29, input_second_moments=[moments])[0] np.testing.assert_array_equal(mask, expected) def test_zero_and_full_budget(self): w = np.ones((1, 2, 3), dtype=np.float16) self.assertEqual(prune_queries([w], 0)[0].sum(), 0) self.assertEqual(prune_queries([w], 6)[0].sum(), 6) with self.assertRaises(ValueError): prune_queries([w], 7) with self.assertRaises(ValueError): prune_queries([np.full_like(w, np.inf)], 0) with self.assertRaises(ValueError): prune_queries([w[..., ::-1]], 3) with self.assertRaises(ValueError): prune_queries([w], 3, input_second_moments=[[-1, 2, 3]]) with self.assertRaises(TypeError): prune_queries([np.ones((1, 2, 3), dtype=np.int16)], 3) class FileAndServiceTest(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.inputs = self.root / "data" self.inputs.mkdir() self.weights = np.arange(24, dtype=np.float32).reshape(2, 3, 4) np.save(self.inputs / "a.npy", self.weights) np.save(self.inputs / "v.npy", np.array([0.5, 1., 2., 3.])) self.spec = dict(keep=7, projections=[dict(name="q", weights="a.npy", input_second_moment="v.npy")]) (self.inputs / "job.json").write_text(json.dumps(self.spec)) def test_manifest_publishes_exact_masks(self): output = self.root / "results" result = prune_manifest(self.inputs, "job.json", output, threads=2) self.assertEqual(result["kept_slots"], 7) mask = np.load(output / "mask-0000.npy", allow_pickle=False) expected = prune_queries([self.weights], 7, input_second_moments=[[.5, 1., 2., 3.]])[0] np.testing.assert_array_equal(mask, expected) self.assertEqual(json.loads((output / "result.json").read_text()), result) with self.assertRaises(FileExistsError): prune_manifest(self.inputs, "job.json", output) with self.assertRaises(ValueError): prune_manifest(self.inputs, "job.json", self.root / "large", max_weights=12) def test_rejects_symlink_and_failed_native_run_leaves_no_output(self): external = self.root / "outside.npy" np.save(external, self.weights) (self.inputs / "escaped.npy").symlink_to(external) self.spec["projections"][0]["weights"] = "escaped.npy" (self.inputs / "escape.json").write_text(json.dumps(self.spec)) with self.assertRaises(ValueError): prune_manifest(self.inputs, "escape.json", self.root / "out") np.save(self.inputs / "a.npy", np.full_like(self.weights, np.nan)) with self.assertRaises(ValueError): prune_manifest(self.inputs, "job.json", self.root / "out") self.assertFalse((self.root / "out").exists()) def test_http_job_mask_and_auth(self): app = create_app(self.inputs, self.root / "jobs", token="test-token", threads_per_job=2) headers = {"Authorization": "Bearer test-token"} with TestClient(app) as client: self.assertEqual(client.get("/health").status_code, 200) self.assertEqual(client.post("/v1/jobs", json={"manifest": "job.json"}).status_code, 401) escape = client.post("/v1/jobs", json={"manifest": "../out.json"}, headers=headers) self.assertEqual(escape.status_code, 400) response = client.post("/v1/jobs", json={"manifest": "job.json"}, headers=headers) self.assertEqual(response.status_code, 202) endpoint = f"/v1/jobs/{response.json()['id']}" for _ in range(150): status = client.get(endpoint, headers=headers).json() if status["state"] in {"completed", "failed"}: break time.sleep(0.01) self.assertEqual(status["state"], "completed", status) binary = client.get(endpoint + "/masks/0", headers=headers) self.assertEqual(binary.status_code, 200) self.assertEqual(np.load(io.BytesIO(binary.content)).sum(), 7) self.assertEqual(client.get(endpoint + "/masks/1", headers=headers).status_code, 404) def test_queue_backpressure(self): started, proceed = Event(), Event() def slow(*args, **kwargs): started.set() if not proceed.wait(5): raise RuntimeError("test timeout") return prune_manifest(*args, **kwargs) app = create_app(self.inputs, self.root / "jobs", workers=1, queue_size=0) with patch("tensor_pruning.service.prune_manifest", slow), TestClient(app) as client: try: self.assertEqual(client.post("/v1/jobs", json={"manifest": "job.json"}).status_code, 202) self.assertTrue(started.wait(2)) response = client.post("/v1/jobs", json={"manifest": "job.json"}) self.assertEqual(response.status_code, 429) self.assertEqual(response.headers["Retry-After"], "1") finally: proceed.set() if __name__ == "__main__": unittest.main()