diff --git a/deapi/client.py b/deapi/client.py index 93c1ea9c..d3751696 100644 --- a/deapi/client.py +++ b/deapi/client.py @@ -83,6 +83,11 @@ class Client: def __init__(self): self.commandVersion = commandVersion self.read_only = False + #: How many acquisitions this client has started, and when the last one was sent + #: (time.monotonic). SimulatedGpuFrames times its frames from these: a short + #: simulated acquisition can be over before its consumer first asks. + self.acquisitions_started = 0 + self.acquisition_started_at = 0.0 def set_log_level(self, level): log = logging.getLogger("DECameraClientLib") @@ -1333,6 +1338,7 @@ def start_acquisition( log.debug(" Build Time: %.1f ms", lapsed) step_time = self.GetTime() + started = time.monotonic() response = self._sendCommand(command) if logLevel == logging.DEBUG: lapsed = (self.GetTime() - step_time) * 1000 @@ -1342,6 +1348,9 @@ def start_acquisition( if response: ret = response.acknowledge[0].error != True self.refreshProperties = True + if ret: + self.acquisition_started_at = started + self.acquisitions_started += 1 if logLevel == logging.DEBUG: lapsed = (self.GetTime() - step_time) * 1000 @@ -1982,6 +1991,52 @@ def get_movie_buffer_info(self, movieBufferInfo=None, timeoutMsec=5000): return movieBufferInfo + def get_gpu_frames(self, peer_devices=()): + """ + Zero-copy CuPy access to the server's GPU movie buffers (Windows, same computer). + + Call before start_acquisition. The returned GpuFrames holds both movie buffers as + CuPy arrays (``.buffers``) and the semaphores; iterating it yields each finished + buffer as ``batch.images`` (cupy float32, images x height x width, on GPU + ``.device``) and ``batch.info`` (frame, scanX, scanY, valid per image). The buffer + is released back to the server when the loop advances; the server's processing + waits while it is held, so close it when done (or use it in a ``with`` block). + Requires CuPy. See deapi.gpu_frames. + + ``peer_devices``: other GPUs that will work on the frames. They read the buffers + directly where the driver allows peer access (``.peer_access``), and their current + streams are synchronized before each release either way. + + Only a process on the server computer can map the buffers; from another computer + this raises with the server's reason. Against deapi's simulated server this + returns a SimulatedGpuFrames instead: the same interface, with NumPy batches of the + simulated dataset and no CuPy needed. + """ + from deapi.gpu_frames import open_gpu_frames + + return open_gpu_frames(self._attach_gpu_frames(), self, peer_devices) + + def _attach_gpu_frames(self, attach=True): + """Ask DE-Server to hand its GPU movie buffer handles to this process (or detach). + Returns what it answered: [control block handle], or ["simulated"] from the fake + server.""" + command = self._addSingleCommand( + self.GET_GPU_FRAMES, None, [str(os.getpid()), attach] + ) + response = self._sendCommand(command) + if response is False: + raise RuntimeError("DE-Server did not answer GET_GPU_FRAMES") + acknowledge = response.acknowledge[0] + if acknowledge.error: + raise RuntimeError( + "DE-Server cannot share its GPU movie buffers with this process: " + + (acknowledge.error_message or "no reason given") + ) + values = list(self.__getParameters(acknowledge)) + if attach and not values: + raise RuntimeError("DE-Server answered GET_GPU_FRAMES without a handle") + return values + def get_movie_buffer( self, movieBuffer, movieBufferSize, numFrames, timeoutMsec=5000 ): @@ -3245,6 +3300,7 @@ def ParseChangedProperties(self, changedProperties, response): GET_VIRTUAL_IMAGE_INFO = 41 GET_VIRTUAL_IMAGE = 42 GET_CLIENT_READ_ONLY = 43 + GET_GPU_FRAMES = 44 MMF_DATA_HEADER_SIZE = 24 diff --git a/deapi/gpu_frames.py b/deapi/gpu_frames.py new file mode 100644 index 00000000..3808bdd6 --- /dev/null +++ b/deapi/gpu_frames.py @@ -0,0 +1,587 @@ +"""Zero-copy CuPy access to DE-Server's GPU movie buffers (Windows, same computer). + + with client.get_gpu_frames() as frames: # before start_acquisition: maps both buffers + client.start_acquisition(1) + for batch in frames: # one finished movie buffer + batch.images # cupy float32 (images, height, width), server memory + batch.info # numpy view: frame, scanX, scanY, valid per image + # the buffer goes back to the server when the loop advances + +The server processes into two ping-pong movie buffers; while this process holds one, the +server's processing waits for it, so finish (or copy) what you need inside the loop body. +GPU work queued on the current stream of the buffers' GPU (and of each of `peer_devices`) +is synchronized before the release; synchronize work on other streams yourself. + +DEAPI hands this process a handle to a small control block once. The server duplicates the +buffer and semaphore handles into this process itself and lists them there, again whenever +it reallocates the buffers or starts a new acquisition while a buffer is held, so no DEAPI +round trip is needed during an acquisition. + +More than one GPU: `frames.device` is the GPU that holds the buffers, in this process's +numbering (read from the memory itself, so CUDA_VISIBLE_DEVICES cannot mislead it). Pass +other GPUs as `peer_devices`: where the driver allows peer access the buffers are mapped +for them too, so their kernels read server memory directly (`frames.peer_access`); +otherwise `with cp.cuda.Device(d): cp.asarray(batch.images)` copies. Either way their +current streams are synchronized before each release. + +More than one computer: only a process on the server computer can map the buffers (the +server refuses others). Run a bridge there that iterates these frames and sends each batch +on (GPUDirect RDMA, NCCL, or pinned host memory and a socket) before releasing it; the +receiving side can present the same FrameSource interface to its consumers. +""" + +import ctypes +import threading +import time +import warnings +from ctypes import c_int, c_size_t, c_uint64 as u64, c_void_p, wintypes + +import numpy as np + +VERSION, BUFFERS, MAX_IMAGES = 4, 2, 1024 # must match GpuFrameSharing/Protocol.h +H = wintypes.HANDLE +WAIT_OBJECT_0, WAIT_TIMEOUT, WAIT_FAILED = 0, 0x102, 0xFFFFFFFF +IMAGE = np.dtype( + [ + ("frame", " picked up here, including queueing), wrap and sync times. + self.profile = dict( + contextNs=0, + attachNs=0, + buffers=0, + pickupNs=0, + pickupMaxNs=0, + wrapNs=0, + wrapMaxNs=0, + syncNs=0, + syncMaxNs=0, + ) + (self._control_handle,) = attach() + address = _k32.MapViewOfFile( + H(self._control_handle), 0xF001F, 0, 0, ctypes.sizeof(_Control) + ) # all access + if not address: + error = ctypes.WinError(ctypes.get_last_error()) + _k32.CloseHandle(H(self._control_handle)) + raise error + self._control = _Control.from_address(address) + try: + if self._control.version != VERSION: + raise RuntimeError( + f"DE-Server GPU frame protocol {self._control.version}, this client expects {VERSION}" + ) + self._unshared = self._control.unshared + self._load() + except BaseException: + self._closed = True + self._unload() + self._release_control() + raise + + def _count(self, name, ns): + self.profile[name + "Ns"] += ns + self.profile[name + "MaxNs"] = max(self.profile[name + "MaxNs"], ns) + + def _load(self): + """Take the buffer and semaphore handles the server listed (again when they change). + Called with the lock held, or from __init__.""" + listed, deadline = self._control.client, time.monotonic() + 10 + while True: # 0 while the server updates the list + session = listed.session + filled, free, buffers, detach = ( + listed.filled, + list(listed.free), + list(listed.buffer), + listed.detach, + ) + if session and listed.session == session: + break + if time.monotonic() > deadline: + raise TimeoutError( + "DE-Server did not finish listing its GPU frame handles" + ) + time.sleep(0.001) + start = time.perf_counter_ns() + self._unload() + self._filled, self._free, self._detach = filled, free, detach + self.session, self.size = session, self._control.bufferBytes + allocations = [] + try: + for ( + handle + ) in buffers: # import, which leaves the duplicated handle to close + allocation = u64() + result = _cu.cuMemImportFromShareableHandle( + ctypes.byref(allocation), handle, _HANDLE_WIN32 + ) + _k32.CloseHandle(H(handle)) + _check(result) + allocations.append(allocation) + devices = set() # the GPU holding the memory, in this process's numbering + for allocation in allocations: + prop = _AllocationProp() + _check( + _cu.cuMemGetAllocationPropertiesFromHandle( + ctypes.byref(prop), allocation + ) + ) + devices.add(prop.location.id) + if len(devices) != 1: + raise RuntimeError(f"GPU movie buffers on several GPUs: {devices}") + self.device = devices.pop() + context = time.perf_counter_ns() + for device in ( + *self._peers, + self.device, + ): # contexts now, not mid-acquisition + with self._cp.cuda.Device(device): + self._cp.cuda.runtime.free(0) + context = time.perf_counter_ns() - context + self.profile["contextNs"] += context + with self._cp.cuda.Device(self.device): + self.peer_access = {d: self._can_access(d) for d in self._peers} + access = [self.device] + [ + d for d, ok in self.peer_access.items() if ok and d != self.device + ] + for allocation in allocations: + self._pointers.append(self._map(allocation, access)) + finally: + for allocation in allocations: + _cu.cuMemRelease(allocation) # a mapping keeps the memory alive + # Both movie buffers, flat; batches are views into these. + self.buffers = [ + self._cp.ndarray( + (self.size // 4,), + self._cp.float32, + self._cp.cuda.MemoryPointer( + self._cp.cuda.UnownedMemory(p, self.size, self, self.device), 0 + ), + ) + for p in self._pointers + ] + self._views, self._consumed = {}, 0 + self.profile["attachNs"] += time.perf_counter_ns() - start - context + + def _can_access(self, device): + """Whether `device`'s kernels can read memory on the buffers' GPU directly.""" + if device == self.device: + return True + can = c_int() + _check(_cu.cuDeviceCanAccessPeer(ctypes.byref(can), device, self.device)) + return bool(can.value) + + def _map(self, allocation, devices): + """Map one imported buffer for `devices` (CUDA VMM); no data is copied.""" + pointer, size = u64(), self.size + _check(_cu.cuMemAddressReserve(ctypes.byref(pointer), size, 0, 0, 0)) + try: + _check(_cu.cuMemMap(pointer, size, 0, allocation, 0)) + access = (_AccessDesc * len(devices))( + *( + _AccessDesc(_Location(_LOCATION_DEVICE, d), _ACCESS_READWRITE) + for d in devices + ) + ) + result = _cu.cuMemSetAccess(pointer, size, access, len(devices)) + if result: + _cu.cuMemUnmap(pointer, size) + _check(result) + except BaseException: + _cu.cuMemAddressFree(pointer, size) + raise + return pointer.value + + def _unload(self): + """Unmap the buffers and close the semaphore handles (lock held).""" + if self._pointers: + with self._cp.cuda.Device(self.device): + self._cp.cuda.Device(self.device).synchronize() + self.buffers, self._views = [], {} + for pointer in self._pointers: + _cu.cuMemUnmap(pointer, self.size) + _cu.cuMemAddressFree(pointer, self.size) + self._pointers, self.buffers, self._views = [], [], {} + for handle in (self._filled, *self._free, self._detach): + if handle: + _k32.CloseHandle(H(handle)) + self._filled, self._free, self._detach = None, [], None + + def server_profile(self): + """The server's counters (see PROFILE; times in ns).""" + with self._lock: + if self._control is None: + return {} + return {name: getattr(self._control, name) for name in PROFILE} + + def _check_unshared(self): + """Warn once if the server finished buffers it could not share while attached.""" + if not self._warned and self._control.unshared != self._unshared: + self._warned = True + warnings.warn( + "DE-Server finished movie buffers it could not share with this process " + f"(more than {MAX_IMAGES} images per buffer, or an event-counting mode " + "without movie images); see server_profile()['unshared']", + RuntimeWarning, + stacklevel=3, + ) + + def batches(self, idle_s=None): + """Yield each finished buffer as a Batch; with `idle_s`, also yield None after that + long without one, so the caller can check whether to stop.""" + cp, waited = self._cp, 0.0 + while True: + with self._lock: + if self._closed: + return + if self._control.client.session != self.session: # new handles listed + self._load() + result = _k32.WaitForSingleObject(H(self._filled), 100) + if result == WAIT_FAILED: + raise ctypes.WinError(ctypes.get_last_error()) + if result == WAIT_OBJECT_0: + woke = time.perf_counter_ns() + if self._control.client.session != self.session: + # A count left from a session the server has since replaced: the + # slots now describe the new session's buffers, not one we hold. + self._load() + continue + session = self.session + slot = self._control.slot[self._consumed % BUFFERS] + buffer = slot.buffer + key = (buffer, slot.count, slot.width, slot.height) + if key not in self._views: # arrays over fixed offsets, built once + self._views[key] = cp.ndarray( + (slot.count, slot.height, slot.width), + cp.float32, + self.buffers[buffer].data, + strides=(slot.imageBytes, slot.rowBytes, 4), + ) + info = np.frombuffer(slot.image, IMAGE, slot.count) # slot view + batch = Batch(self._views[key], info, self.device) + free, self._holder = self._free[buffer], threading.get_ident() + self._consumed += 1 + self._count("pickup", woke - slot.publishedNs) + self._count("wrap", time.perf_counter_ns() - woke) + self._check_unshared() + if result != WAIT_OBJECT_0: # nothing yet + waited += 0.1 + if idle_s is not None and waited >= idle_s: + waited = 0.0 + yield None + continue + waited = 0.0 + try: + yield batch + finally: + del batch + self._release(free, session) + + def _release(self, free, session): + """Give a buffer back once the GPU work queued on it is done.""" + synced = time.perf_counter_ns() + try: + for device in (self.device, *self._peers): + with self._cp.cuda.Device(device): + self._cp.cuda.get_current_stream().synchronize() + finally: + with self._lock: + # After close() or a new session the server took the buffer back itself, + # and the handle may be closed already: releasing it is no longer ours. + if not self._closed and session == self.session: + _k32.ReleaseSemaphore(H(free), 1, None) + self._holder = None + self._lock.notify_all() + self._count("sync", time.perf_counter_ns() - synced) + self.profile["buffers"] += 1 + + def close(self): + """Stop sharing: the server gets any held buffer back and stops waiting (no DEAPI + call). Safe to call more than once, and from another thread.""" + with self._lock: + if self._closed: + return + self._closed = True + me = threading.get_ident() + # Another thread still working on a batch: let it finish before unmapping. + self._lock.wait_for(lambda: self._holder in (None, me)) + if self._detach: + _k32.SetEvent(H(self._detach)) + self._unload() + self._release_control() + + def _release_control(self): + if self._control is not None: + _k32.UnmapViewOfFile(ctypes.addressof(self._control)) + _k32.CloseHandle(H(self._control_handle)) + self._control = None + + +class SimulatedGpuFrames(FrameSource): + """GpuFrames' interface over deapi's simulated server, for clients developed offline. + + Batches are NumPy float32 frames of the server's own simulated dataset (the TiltGrains + its results come from), built in this process and delivered at the server's frame rate + while it acquires. These are copies and there is no GPU: a stand-in for testing, not + the zero-copy path. `client` reads properties (`client[name]`) and `client.acquiring`, + and carries the Client's `acquisitions_started` / `acquisition_started_at`: the frames + belong to the next acquisition started after this was made, timed from its start. + """ + + BATCH = 32 + + def __init__(self, client): + self._client, self._closed = client, False + self._after = client.acquisitions_started + self.device, self.session, self.peer_access = None, 1, {} + self.profile = dict(buffers=0) + + def batches(self, idle_s=None): + """As GpuFrames.batches: each batch, and None every `idle_s` seconds without one.""" + from deapi.fake_data.grains import TiltGrains + from deapi.simulated_server.fake_server import SIMULATED_MAX_FPS + + c, waited = self._client, 0.0 + while c.acquisitions_started == self._after: # wait for the acquisition + if self._closed: + return + time.sleep(0.05) + waited += 0.05 + if idle_s is not None and waited >= idle_s: + waited = 0.0 + yield None + scanning = c["Scan - Enable"] == "On" + sx, sy = ( + (int(c["Scan - Size X"]), int(c["Scan - Size Y"])) if scanning else (1, 1) + ) + points = sx * sy + total = points * (int(c["Scan - Repeats"] or 1) if scanning else 1) + k = int(c["Sensor Size X (pixels)"]) + fps = min(float(c["Frames Per Second"]), SIMULATED_MAX_FPS) + data = TiltGrains(x_pixels=sx, y_pixels=sy, kx_pixels=k, ky_pixels=k) + signal = np.asarray(data._signal, dtype=np.float32) + start, done, checked, waited = c.acquisition_started_at, 0, 0.0, 0.0 + while done < total and not self._closed: + now = time.monotonic() + if now - checked > 0.1: # one round trip per 100 ms + checked = now + if not c.acquiring and now < start + total / fps - 0.5: + return # stopped early + due = min(total, int((now - start) * fps)) + if due <= done: + time.sleep(0.005) + waited += 0.005 + if idle_s is not None and waited >= idle_s: + waited = 0.0 + yield None + continue + n = min(self.BATCH, due - done) + frame = np.arange(done, done + n) + x, y = frame % points % sx, frame % points // sx + info = np.zeros(n, IMAGE) + info["frame"], info["scanX"], info["scanY"], info["valid"] = frame, x, y, 1 + yield Batch(signal[data.navigator[x, y]], info) + self.profile["buffers"] += 1 + done += n + waited = 0.0 + + def close(self): + self._closed = True diff --git a/deapi/simulated_server/fake_server.py b/deapi/simulated_server/fake_server.py index 56e342dc..6343a5ba 100644 --- a/deapi/simulated_server/fake_server.py +++ b/deapi/simulated_server/fake_server.py @@ -14,6 +14,11 @@ inp_file = resources.files(deapi) / "prop_dump.json" +#: The fastest this simulated camera reads out, whatever "Frames Per Second" asks for — a +#: real server clamps it to what the readout allows. SimulatedGpuFrames paces its batches +#: by the same number, so the frames last as long as the acquisition does. +SIMULATED_MAX_FPS = 2000.0 + def add_parameter(ack, value): """ @@ -265,10 +270,12 @@ def current_navigation_index(self): if self.acquisition_status == "Idle": return tuple(np.array(self.fake_data.navigator.shape) - 1) else: + # Frames so far at the rate the acquisition runs at, wrapped to the + # scan: with Scan - Repeats each pass starts again at the first position. + fps = min(float(self["Frames Per Second"]), SIMULATED_MAX_FPS) index = np.unravel_index( - int( - (time.time() - self.start_time) * float(self["Frames Per Second"]), - ), + int((time.time() - self.start_time) * fps) + % self.fake_data.navigator.size, self.fake_data.navigator.shape, ) # only works for raster scans return index @@ -333,6 +340,10 @@ def _respond_to_command(self, command=None): == self.GET_VIRTUAL_IMAGE + commandVersion * 100 ): return self._fake_get_virtual_image(command) + elif ( + command.command[0].command_id == self.GET_GPU_FRAMES + commandVersion * 100 + ): + return self._fake_get_gpu_frames(command) else: raise NotImplementedError( f"Command {command.command[0].command_id} not implemented" @@ -392,6 +403,20 @@ def _fake_set_virtual_mask(self, command): print(f"Virtual mask {mask_id} set with {mask}") return (acknowledge_return,) + def _fake_get_gpu_frames(self, command): + """There is no GPU memory to share: answer "simulated", and the client's + get_gpu_frames builds SimulatedGpuFrames over this server's dataset instead.""" + acknowledge_return = pb.DEPacket() + acknowledge_return.type = pb.DEPacket.P_ACKNOWLEDGE + ack1 = acknowledge_return.acknowledge.add() + ack1.command_id = command.command[0].command_id + add_parameter(ack1, "simulated") + return (acknowledge_return,) + + def stop(self): + """End the acquisition now (the client's UDP stop, see initialize_server).""" + self.end_time = min(self.end_time, time.time()) + def _fake_list_cameras(self, command): acknowledge_return = pb.DEPacket() acknowledge_return.type = pb.DEPacket.P_ACKNOWLEDGE @@ -420,7 +445,11 @@ def _fake_start_acquisition(self, command): acknowledge_return = pb.DEPacket() num_acq = command.command[0].parameter[0].p_int if self["Scan - Enable"] == "On": - frames = int(self["Scan - Size X"]) * int(self["Scan - Size Y"]) + frames = ( + int(self["Scan - Size X"]) + * int(self["Scan - Size Y"]) + * int(self["Scan - Repeats"] or 1) + ) self._initialize_data( int(self["Scan - Size X"]), int(self["Scan - Size Y"]), @@ -436,7 +465,7 @@ def _fake_start_acquisition(self, command): ) frames = num_acq self.number_of_frames_requested = frames - fps = float(self["Frames Per Second"]) + fps = min(float(self["Frames Per Second"]), SIMULATED_MAX_FPS) total_time = frames * num_acq / fps self.start_time = time.time() print(f"Acquisition started for {total_time} seconds") @@ -724,6 +753,13 @@ def _fake_get_result(self, command): ] image = self.fake_data.get_virtual_image(image, method=calculation_type) image = image.astype(pixel_format_dict[pixel_format]) + elif 22 <= frame_type < 26: # external images: an annular (HAADF) detector + ky, kx = self.fake_data.signal.shape[1:] + yy, xx = np.mgrid[0:ky, 0:kx] + r = np.hypot(yy - ky / 2, xx - kx / 2) + mask = np.where(r > 0.15 * min(kx, ky), 2, 1).astype(np.int8) + image = self.fake_data.get_virtual_image(mask, method="Sum") + image = image.astype(pixel_format_dict[pixel_format]) else: raise ValueError(f"Frame type {frame_type} not Supported in PythonDEServer") @@ -940,3 +976,4 @@ def _fake_get_virtual_image(self, command): GET_VIRTUAL_IMAGE_INFO = 41 GET_VIRTUAL_IMAGE = 42 GET_CLIENT_READ_ONLY = 43 + GET_GPU_FRAMES = 44 diff --git a/deapi/simulated_server/initialize_server.py b/deapi/simulated_server/initialize_server.py index 5c5751be..3624c974 100644 --- a/deapi/simulated_server/initialize_server.py +++ b/deapi/simulated_server/initialize_server.py @@ -1,4 +1,5 @@ import sys +import threading import traceback from deapi.simulated_server.fake_server import FakeServer @@ -45,6 +46,24 @@ def serve_twin(port, twin_args=()): return deapi_server.main([str(port), *twin_args]) +def _serve_stop(host, port, current): + """The client's stop_acquisition: a UDP "PyClientStopAcq" to the same port, answered + "Stopped". Without it a stop against this server would wait for a reply forever.""" + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp: + try: + udp.bind((host, port)) + except OSError as e: + sys.stderr.write( + f"stop_acquisition will not be answered: UDP {port}: {e}\n" + ) + return + while True: + message, addr = udp.recvfrom(64) + if message.startswith(b"PyClientStopAcq") and current: + current[-1].stop() + udp.sendto(b"Stopped", addr) + + # Defining main function def main(port=13240): parser = argparse.ArgumentParser() @@ -81,6 +100,10 @@ def main(port=13240): f" Port: {PORT} \n" ) sys.stderr.flush() + current = [] # the server of the connection being served, for the stop listener + threading.Thread( + target=_serve_stop, args=(HOST, PORT, current), daemon=True + ).start() while True: conn, addr = server_socket.accept() # What waits for a connection # Guard against a stalled client leaving _recv_exact blocked forever. @@ -89,6 +112,7 @@ def main(port=13240): # somehow delivers a short write. conn.settimeout(120) server = FakeServer(socket=conn) + current[:] = [server] connected = True while connected: try: diff --git a/deapi/tests/test_fake_server/test_gpu_frames.py b/deapi/tests/test_fake_server/test_gpu_frames.py new file mode 100644 index 00000000..30c1969a --- /dev/null +++ b/deapi/tests/test_fake_server/test_gpu_frames.py @@ -0,0 +1,108 @@ +"""get_gpu_frames against the simulated server: the same interface as the real, zero-copy +GpuFrames, with NumPy batches of the simulated dataset.""" + +import time + +import numpy as np +import pytest + +from deapi.buffer_protocols import pb +from deapi.gpu_frames import FrameSource, SimulatedGpuFrames + + +@pytest.fixture(autouse=True) +def _restore(client): + """The client is shared by the whole session: leave the scan as it was found.""" + before = { + name: client[name] + for name in ( + "Scan - Enable", + "Scan - Size X", + "Scan - Size Y", + "Scan - Repeats", + "Frames Per Second", + ) + } + yield + for name, value in before.items(): + client[name] = value + + +def _scan(client, size, repeats): + client["Scan - Enable"] = "On" + client["Scan - Size X"] = size + client["Scan - Size Y"] = size + client["Scan - Repeats"] = repeats + client["Frames Per Second"] = 1000 + + +def test_simulated_frames_cover_every_scan_position(client): + _scan(client, 8, 2) + frames = client.get_gpu_frames() + assert isinstance(frames, SimulatedGpuFrames) + client.start_acquisition(1) + seen, positions, deadline = 0, set(), time.monotonic() + 30 + for batch in frames.batches(idle_s=0.2): + assert time.monotonic() < deadline + if batch is None: + continue + assert batch.images.dtype == np.float32 and batch.images.ndim == 3 + assert batch.info["valid"].all() + seen += len(batch.images) + positions.update( + zip(batch.info["scanX"].tolist(), batch.info["scanY"].tolist()) + ) + frames.close() + assert seen == 8 * 8 * 2 + assert positions == {(x, y) for x in range(8) for y in range(8)} + + +def test_stop_ends_the_simulated_frames(client): + _scan(client, 8, 1000) + frames = client.get_gpu_frames() + client.start_acquisition(1) + seen, started, deadline = 0, time.monotonic(), time.monotonic() + 30 + for batch in frames.batches(idle_s=0.2): + assert time.monotonic() < deadline + if batch is not None: + seen += len(batch.images) + if seen and time.monotonic() - started > 0.5 and client.acquiring: + assert client.stop_acquisition() + frames.close() + assert 0 < seen < 8 * 8 * 1000 + assert not client.acquiring + + +def test_external_image_is_served(client): + """A HAADF-only search reads external_image1; the simulated server answers it, + including after the first pass of a repeated scan.""" + _scan(client, 8, 1000) + client.start_acquisition(1) + time.sleep(0.2) # past the first pass: 64 positions at 1000 fps + result = client.get_result("external_image1", "FLOAT32") + image = getattr(result, "image", None) + image = result[0] if image is None else image + assert np.asarray(image).size > 0 + client.stop_acquisition() + + +def test_simulated_frames_close_with_the_block(client): + with client.get_gpu_frames() as frames: + assert isinstance(frames, FrameSource) + assert frames._closed + frames.close() # again: no-op + assert list(frames.batches()) == [] + + +def test_server_refusal_is_reported(client, monkeypatch): + """From another computer (or with a consumer already attached) the server refuses, + and its reason reaches the caller.""" + refusal = pb.DEPacket() + acknowledge = refusal.acknowledge.add() + acknowledge.error = True + acknowledge.error_message = ( + "GPU frame sharing is only available to clients on the server computer" + ) + monkeypatch.setattr(client, "_sendCommand", lambda command: refusal) + with pytest.raises(RuntimeError, match="only available to clients on the server"): + client.get_gpu_frames() diff --git a/deapi/tests/test_gpu_frames_cuda.py b/deapi/tests/test_gpu_frames_cuda.py new file mode 100644 index 00000000..722a7c91 --- /dev/null +++ b/deapi/tests/test_gpu_frames_cuda.py @@ -0,0 +1,18 @@ +"""GpuFrames checks that need CuPy and a GPU but no DE-Server.""" + +import pytest + +cp = pytest.importorskip("cupy") +try: + if cp.cuda.runtime.getDeviceCount() < 1: + pytest.skip("no GPU", allow_module_level=True) +except cp.cuda.runtime.CUDARuntimeError: + pytest.skip("no CUDA driver", allow_module_level=True) + +from deapi.gpu_frames import GpuFrames + + +def test_missing_peer_device_fails_before_attaching(): + count = cp.cuda.runtime.getDeviceCount() + with pytest.raises(ValueError, match=f"sees {count} GPU"): + GpuFrames(lambda: pytest.fail("attached"), peer_devices=[count])