diff --git a/tinygrad/runtime/ops_cloud.py b/tinygrad/runtime/ops_cloud.py index e1550c5eb4..df0fb6a4bf 100644 --- a/tinygrad/runtime/ops_cloud.py +++ b/tinygrad/runtime/ops_cloud.py @@ -13,7 +13,7 @@ from http.server import HTTPServer, BaseHTTPRequestHandler from tinygrad.renderer import Renderer from tinygrad.dtype import dtypes from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, Timing -from tinygrad.device import Compiled, Allocator, Compiler, Device, BufferSpec +from tinygrad.device import Compiled, Buffer, Allocator, Compiler, Device, BufferSpec # ***** API ***** @@ -75,8 +75,7 @@ class BatchRequest: @dataclass class CloudSession: programs: dict[tuple[str, str], Any] = field(default_factory=dict) - # TODO: the buffer should track this internally - buffers: dict[int, tuple[Any, int, Optional[BufferSpec]]] = field(default_factory=dict) + buffers: dict[int, Buffer] = field(default_factory=dict) class CloudHandler(BaseHTTPRequestHandler): protocol_version = 'HTTP/1.1' @@ -99,21 +98,16 @@ class CloudHandler(BaseHTTPRequestHandler): match c: case BufferAlloc(): assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated" - session.buffers[c.buffer_num] = (Device[CloudHandler.device].allocator.alloc(c.size, c.options), c.size, c.options) - case BufferFree(): - buf,sz,buffer_options = session.buffers[c.buffer_num] - Device[CloudHandler.device].allocator.free(buf,sz,buffer_options) - del session.buffers[c.buffer_num] - case CopyIn(): Device[CloudHandler.device].allocator._copyin(session.buffers[c.buffer_num][0], memoryview(bytearray(req._h[c.datahash]))) - case CopyOut(): - buf,sz,_ = session.buffers[c.buffer_num] - Device[CloudHandler.device].allocator._copyout(memoryview(ret:=bytearray(sz)), buf) + session.buffers[c.buffer_num] = Buffer(CloudHandler.device, c.size, dtypes.uint8, options=c.options, preallocate=True) + case BufferFree(): del session.buffers[c.buffer_num] + case CopyIn(): session.buffers[c.buffer_num].copyin(memoryview(bytearray(req._h[c.datahash]))) + case CopyOut(): session.buffers[c.buffer_num].copyout(memoryview(ret:=bytearray(session.buffers[c.buffer_num].nbytes))) case ProgramAlloc(): lib = Device[CloudHandler.device].compiler.compile_cached(req._h[c.datahash].decode()) session.programs[(c.name, c.datahash)] = Device[CloudHandler.device].runtime(c.name, lib) case ProgramFree(): del session.programs[(c.name, c.datahash)] case ProgramExec(): - bufs = [session.buffers[x][0] for x in c.bufs] + bufs = [session.buffers[x]._buf for x in c.bufs] extra_args = {k:v for k,v in [("global_size", c.global_size), ("local_size", c.local_size)] if v is not None} r = session.programs[(c.name, c.datahash)](*bufs, vals=c.vals, wait=c.wait, **extra_args) if r is not None: ret = str(r).encode()