diff --git a/test/external/external_test_dev_var.py b/test/external/external_test_dev_var.py new file mode 100644 index 0000000000..41abbe8e79 --- /dev/null +++ b/test/external/external_test_dev_var.py @@ -0,0 +1,39 @@ +import subprocess, unittest, os, sys +from tinygrad.device import Device + +class TestTinygradSlow(unittest.TestCase): + def test_env_overwrite_default_device(self): + subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + if Device.DEFAULT != "CPU": + # setting multiple devices fail + with self.assertRaises(subprocess.CalledProcessError): + subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + # setting device via DEV + subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + with self.assertRaises(subprocess.CalledProcessError): + subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + +class TestRunAsModule(unittest.TestCase): + def test_module_runs(self): + p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE, + env={**os.environ, "DEBUG": "1"}, timeout=40,) + out = (p.stdout + p.stderr).decode() + self.assertEqual(p.returncode, 0, msg=out) + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_tensor.py b/test/test_tensor.py index 43b8202dc4..617eb242a3 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -1,4 +1,3 @@ -import subprocess import numpy as np import torch import unittest, copy, mmap, random, math, array @@ -515,32 +514,6 @@ class TestTinygrad(unittest.TestCase): print(a) print(c) - def test_env_overwrite_default_device(self): - subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - if Device.DEFAULT != "CPU": - # setting multiple devices fail - with self.assertRaises(subprocess.CalledProcessError): - subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - # setting device via DEV - subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - with self.assertRaises(subprocess.CalledProcessError): - subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - def test_no_attributeerror_after_apply_uop_exception(self): try: Tensor.arange(4).reshape(3,2) diff --git a/test/unit/test_device.py b/test/unit/test_device.py index 984feaf751..e1eaaa1314 100644 --- a/test/unit/test_device.py +++ b/test/unit/test_device.py @@ -1,7 +1,7 @@ #!/usr/bin/env python -import unittest, os, subprocess, sys +import unittest, os, subprocess from tinygrad import Tensor -from tinygrad.device import Device, Compiler +from tinygrad.device import Device, Compiler, enumerate_devices_str from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI class TestDevice(unittest.TestCase): @@ -100,10 +100,7 @@ class TestCompiler(unittest.TestCase): class TestRunAsModule(unittest.TestCase): def test_module_runs(self): - p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE, - env={**os.environ, "DEBUG": "1"}, timeout=40,) - out = (p.stdout + p.stderr).decode() - self.assertEqual(p.returncode, 0, msg=out) + out = '\n'.join(enumerate_devices_str()) self.assertIn("CPU", out) # for sanity check if __name__ == "__main__": diff --git a/tinygrad/device.py b/tinygrad/device.py index 3e1788c946..bd021ebea1 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, replace from collections import defaultdict -from typing import Any, Generic, TypeVar, Iterator, Sequence, cast +from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup @@ -357,7 +357,7 @@ if PROFILE: from tinygrad.uop.ops import launch_viz launch_viz("PROFILE", fn) -if __name__ == "__main__": +def enumerate_devices_str() -> Generator[str, None, None]: from tinygrad import Tensor, Device for device in ALL_DEVICES: @@ -376,4 +376,7 @@ if __name__ == "__main__": result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results]) except Exception as e: result = f"{colored('FAIL', 'red')} {e}" - print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}") + yield f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}" + +if __name__ == "__main__": + for s in enumerate_devices_str(): print(s)