Spaces:
Sleeping
Sleeping
| 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() | |