github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-15 01:48:23 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix up some slow tests that launch python

Browse Source
This commit is contained in:
George Hotz
2025-10-14 18:42:42 +08:00
parent d3bfcd3277
commit 8a34a4e2c7
4 changed files with 48 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
test/external/external_test_dev_var.py vendored Normal file
View File

@@ -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()

27
test/test_tensor.py
View File

@@ -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)

9
test/unit/test_device.py
View File

@@ -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__":

9
tinygrad/device.py
View File

@@ -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)