From 232ed2af3f2ebb55617672adc0df8460e858f191 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 5 Dec 2023 16:17:57 -0800 Subject: [PATCH] more test cleanups (#2631) * more test cleanups * move test example back --- .pre-commit-config.yaml | 4 +- examples/mlperf/model_eval.py | 3 +- extra/helpers.py | 50 ---------------- setup.py | 1 - ...s.py => external_test_dist_collectives.py} | 0 ...t_world.py => external_test_dist_world.py} | 0 ...st_example.py => external_test_example.py} | 0 test/extra/__init__.py | 0 test/extra/test_extra_helpers.py | 57 ------------------- test/imported/__init__.py | 0 test/models/__init__.py | 0 test/unit/__init__.py | 0 12 files changed, 3 insertions(+), 112 deletions(-) delete mode 100644 extra/helpers.py rename test/external/{dist/test_collectives.py => external_test_dist_collectives.py} (100%) rename test/external/{dist/test_world.py => external_test_dist_world.py} (100%) rename test/external/{test_example.py => external_test_example.py} (100%) create mode 100644 test/extra/__init__.py delete mode 100644 test/extra/test_extra_helpers.py create mode 100644 test/imported/__init__.py create mode 100644 test/models/__init__.py create mode 100644 test/unit/__init__.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b9b3bb08..bf8e2d87 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -15,7 +15,7 @@ repos: pass_filenames: false - id: mypy name: mypy - entry: mypy tinygrad/ extra/helpers.py + entry: mypy tinygrad/ language: system always_run: true pass_filenames: false @@ -41,7 +41,7 @@ repos: pass_filenames: false - id: example name: multi device tests - entry: python3 test/external/test_example.py + entry: python3 test/external/external_test_example.py language: system always_run: true pass_filenames: false diff --git a/examples/mlperf/model_eval.py b/examples/mlperf/model_eval.py index 9dbb03c7..5febc892 100644 --- a/examples/mlperf/model_eval.py +++ b/examples/mlperf/model_eval.py @@ -26,12 +26,11 @@ def eval_resnet(): # evaluation on the mlperf classes of the validation set from imagenet from extra.datasets.imagenet import iterate - from extra.helpers import cross_process BS = 64 n,d = 0,0 st = time.perf_counter() - iterator = cross_process(lambda: iterate(BS)) + iterator = iterate(BS) x,ny = next(iterator) dat = Tensor(x) while dat is not None: diff --git a/extra/helpers.py b/extra/helpers.py deleted file mode 100644 index c580a1a6..00000000 --- a/extra/helpers.py +++ /dev/null @@ -1,50 +0,0 @@ -import multiprocessing, subprocess -import cloudpickle -from typing import Any - -def _early_exec_process(qin, qout): - while True: - path, inp = qin.get() - try: - qout.put(subprocess.check_output(path, input=inp)) - except Exception as e: - qout.put(e) - -def enable_early_exec(): - qin: multiprocessing.Queue = multiprocessing.Queue() - qout: multiprocessing.Queue = multiprocessing.Queue() - p = multiprocessing.Process(target=_early_exec_process, args=(qin, qout)) - p.daemon = True - p.start() - def early_exec(x): - qin.put(x) - ret = qout.get() - if isinstance(ret, Exception): raise ret - else: return ret - return early_exec - -def proc(itermaker, q) -> None: - try: - for x in itermaker(): q.put(x) - except Exception as e: - q.put(e) - finally: - q.put(None) - q.close() - -class _CloudpickleFunctionWrapper: - def __init__(self, fn): self.fn = fn - def __getstate__(self): return cloudpickle.dumps(self.fn) - def __setstate__(self, pfn): self.fn = cloudpickle.loads(pfn) - def __call__(self, *args, **kwargs) -> Any: return self.fn(*args, **kwargs) - -def cross_process(itermaker, maxsize=16): - q: multiprocessing.Queue = multiprocessing.Queue(maxsize) - # multiprocessing uses pickle which cannot dump lambdas, so use cloudpickle. - p = multiprocessing.Process(target=proc, args=(_CloudpickleFunctionWrapper(itermaker), q)) - p.start() - while True: - ret = q.get() - if isinstance(ret, Exception): raise ret - elif ret is None: break - else: yield ret \ No newline at end of file diff --git a/setup.py b/setup.py index 6feb3aa4..5be6ae8b 100644 --- a/setup.py +++ b/setup.py @@ -46,7 +46,6 @@ setup(name='tinygrad', "opencv-python", "tabulate", "safetensors", - "cloudpickle", "transformers", "sentencepiece", "tiktoken", diff --git a/test/external/dist/test_collectives.py b/test/external/external_test_dist_collectives.py similarity index 100% rename from test/external/dist/test_collectives.py rename to test/external/external_test_dist_collectives.py diff --git a/test/external/dist/test_world.py b/test/external/external_test_dist_world.py similarity index 100% rename from test/external/dist/test_world.py rename to test/external/external_test_dist_world.py diff --git a/test/external/test_example.py b/test/external/external_test_example.py similarity index 100% rename from test/external/test_example.py rename to test/external/external_test_example.py diff --git a/test/extra/__init__.py b/test/extra/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/extra/test_extra_helpers.py b/test/extra/test_extra_helpers.py deleted file mode 100644 index 6832b973..00000000 --- a/test/extra/test_extra_helpers.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env python -import os, cloudpickle, tempfile, unittest, subprocess -from extra.helpers import enable_early_exec, cross_process, _CloudpickleFunctionWrapper - -def normalize_line_endings(s): return s.replace(b'\r\n', b'\n') - -class TestEarlyExec(unittest.TestCase): - def setUp(self) -> None: - self.early_exec = enable_early_exec() - - def early_exec_py_file(self, file_content, exec_args): - with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp: - temp.write(file_content) - temp_path = temp.name - try: - output = self.early_exec((["python3", temp_path] + exec_args, None)) - return output - finally: - os.remove(temp_path) - - def test_enable_early_exec(self): - output = self.early_exec_py_file(b'print("Hello, world!")', []) - self.assertEqual(b"Hello, world!\n", normalize_line_endings(output)) - - def test_enable_early_exec_with_arg(self): - output = self.early_exec_py_file(b'import sys\nprint("Hello, " + sys.argv[1] + "!")', ["world"]) - self.assertEqual(b"Hello, world!\n", normalize_line_endings(output)) - - def test_enable_early_exec_process_exception(self): - with self.assertRaises(subprocess.CalledProcessError): - self.early_exec_py_file(b'raise Exception("Test exception")', []) - - def test_enable_early_exec_type_exception(self): - with self.assertRaises(TypeError): - self.early_exec((["python3"], "print('Hello, world!')")) - -class TestCrossProcess(unittest.TestCase): - - def test_cross_process(self): - def _iterate(): - for i in range(10): yield i - results = list(cross_process(_iterate)) - self.assertEqual(list(range(10)), results) - - def test_cross_process_exception(self): - def _iterate(): - for i in range(10): - if i == 5: raise ValueError("Test exception") - yield i - with self.assertRaises(ValueError): list(cross_process(_iterate)) - - def test_CloudpickleFunctionWrapper(self): - def add(x, y): return x + y - self.assertEqual(7, cloudpickle.loads(cloudpickle.dumps(_CloudpickleFunctionWrapper(add)))(3, 4)) - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/test/imported/__init__.py b/test/imported/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/models/__init__.py b/test/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/unit/__init__.py b/test/unit/__init__.py new file mode 100644 index 00000000..e69de29b