Files
github-actions[bot] 54e6124925 sunnypilot v2026.05.07-4485
version: sunnypilot v2026.001.000 (dev)
date: 2026-05-07T23:07:19
master commit: c28eb958740187620f2282023b8f1997cf90f583
2026-05-07 23:07:19 +00:00

41 lines
1.1 KiB
Python

import unittest
import multiprocessing.shared_memory as shared_memory
from tinygrad.helpers import CI, WIN
from tinygrad.tensor import Tensor, Device
import numpy as np
class TestRawShmBuffer(unittest.TestCase):
@unittest.skipIf(WIN and CI, "only fails on CI windows instance")
def test_e2e(self):
t = Tensor.randn(2, 2, 2).realize()
# copy to shm
shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name
s.close()
t_shm = t.to(f"disk:shm:{shm_name}").realize()
# copy from shm
t2 = t_shm.to(Device.DEFAULT).realize()
assert np.allclose(t.numpy(), t2.numpy())
s.unlink()
@unittest.skip("big shared memory")
def test_e2e_big(self):
# bigger than this doesn't work on Linux, maybe this is a limit somewhere?
t = Tensor.randn(2048, 128, 8).realize()
# copy to shm
shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name
s.close()
t_shm = t.to(f"disk:shm:{shm_name}").realize()
# copy from shm
t2 = t_shm.to(Device.DEFAULT).realize()
assert np.allclose(t.numpy(), t2.numpy())
s.unlink()
if __name__ == "__main__":
unittest.main()