feat: configurable timeout (#13537)

This commit is contained in:
wozeparrot
2025-12-02 13:35:35 -08:00
committed by GitHub
parent 21aac568fd
commit 18640f57b2

View File

@@ -5,6 +5,7 @@ from tinygrad.helpers import DEBUG, getenv
from tinygrad import Tensor
TINYFS_ENDPOINT = getenv("TINYFS_ENDPOINT", "localhost:6767")
TINYFS_TIMEOUT = getenv("TINYFS_TIMEOUT", 60)
class TinyFSDevice(Compiled):
def __init__(self, device:str):
@@ -13,7 +14,7 @@ class TinyFSDevice(Compiled):
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.connect((TINYFS_ENDPOINT.rsplit(":", 1)[0], int(TINYFS_ENDPOINT.rsplit(":", 1)[1])))
self.sock.settimeout(10)
self.sock.settimeout(TINYFS_TIMEOUT)
self.sfile = self.sock.makefile("rwb")
# fetch node info
@@ -113,9 +114,9 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
writer.write(f"CHUNK_OUT {size}\r\n".encode())
writer.write(src.hash_buf[i*16:(i+1)*16])
await asyncio.wait_for(writer.drain(), timeout=10)
await asyncio.wait_for(writer.drain(), timeout=TINYFS_TIMEOUT)
chunk = await asyncio.wait_for(reader.readexactly(size), timeout=10)
chunk = await asyncio.wait_for(reader.readexactly(size), timeout=TINYFS_TIMEOUT)
view = dest[ptr:ptr+len(chunk)]
view[:] = chunk