diff --git a/test/unit/test_disk_tensor.py b/test/unit/test_disk_tensor.py index 74a8fee8be..94fa484d11 100644 --- a/test/unit/test_disk_tensor.py +++ b/test/unit/test_disk_tensor.py @@ -1,4 +1,4 @@ -import pathlib, tempfile, unittest +import os, pathlib, tempfile, unittest import numpy as np from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import DType @@ -410,5 +410,13 @@ class TestPathTensor(unittest.TestCase): self.assertEqual(t_cpu.device, "CPU") np.testing.assert_array_equal(t_cpu.numpy(), np.frombuffer(self.test_data, dtype=np.uint8)) + def test_path_tensor_disk_device_bug(self): + test_file = pathlib.Path(self.temp_dir.name) / "disk_device_bug" + with open(test_file, "wb") as f: f.write(bytes(range(10))) + os.chmod(test_file, 0o000) + with self.assertRaises(PermissionError): + Tensor(pathlib.Path(test_file)).tolist() + os.chmod(test_file, 0o644) + assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10)) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/runtime/ops_disk.py b/tinygrad/runtime/ops_disk.py index 34d499966f..1cfd5d4e6c 100644 --- a/tinygrad/runtime/ops_disk.py +++ b/tinygrad/runtime/ops_disk.py @@ -16,10 +16,11 @@ class DiskDevice(Compiled): self.fd: Optional[int] = None self.count = 0 super().__init__(device, DiskAllocator(self), None, None, None) - def _might_open(self, size): - self.count += 1 + def _might_open(self, size:int): assert self.size is None or size <= self.size, f"can't reopen Disk tensor with larger size, opened with {self.size}, tried to open with {size}" - if self.size is not None: return + if self.size is not None and hasattr(self.device, "mem"): + self.count += 1 + return filename = self.device[len("disk:"):] self.size = size @@ -34,6 +35,7 @@ class DiskDevice(Compiled): self.mem = mmap.mmap(self.fd, self.size) if hasattr(self.mem, 'madvise') and (hp := getattr(mmap, "MADV_HUGEPAGE", None)) is not None: with contextlib.suppress(OSError): self.mem.madvise(hp) # some systems have transparent_hugepage disabled + self.count += 1 def _might_close(self): self.count -= 1 if self.count == 0: