diff --git a/tinygrad/runtime/ops_tinyfs.py b/tinygrad/runtime/ops_tinyfs.py index d881d0a777..95af79e3e9 100644 --- a/tinygrad/runtime/ops_tinyfs.py +++ b/tinygrad/runtime/ops_tinyfs.py @@ -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