diff --git a/confluent_client/bin/nodeconsole b/confluent_client/bin/nodeconsole index 61943f3c..9ebf9290 100755 --- a/confluent_client/bin/nodeconsole +++ b/confluent_client/bin/nodeconsole @@ -39,6 +39,8 @@ import select import signal import socket import re +import queue +import threading import tty import termios import fcntl @@ -88,6 +90,11 @@ argparser.add_option('-i', '--interval', type='float', 'works for one node') argparser.add_option('-v', '--video', action='store_true', default=False, help='Attempt to continuously stream video from nodes that support streaming console via confluent') +argparser.add_option('-o', '--outputfile', type='string', default=None, + help='Record VNC mode output to the specified video file. When ' + 'multiple nodes are recorded, use a confluent attribute ' + 'expression (e.g. "{node}.mp4") to make the filename unique ' + 'for each node. If multiple nodes specified but not an expression, node name will be inserted.') argparser.add_option('-w','--windowed', action='store_true', default=False, help='Open terminal windows for each node. The ' 'environment variable NODECONSOLE_WINDOWED_COMMAND ' @@ -827,27 +834,34 @@ cheight = 0 imagedatabynode = {} firstnodename = None +# Terminal drawing is offloaded to a single worker thread so slow image +# encoding/output doesn't stall the asyncio loop. The lock keeps the worker +# from interleaving output with redraw() on the main thread. +draw_lock = threading.Lock() +draw_queue = queue.Queue() + def redraw(): - for node in imagedatabynode: - imgdata = imagedatabynode[node] - if node in nodepositions: - prep_node_tile(node) - cursor_save() - else: - if options.interval is not None: - if node != firstnodename: - sys.stderr.write('Multiple nodes not supported for interval') - sys.exit(1) - sticky_cursor() - sys.stdout.write('{}: '.format(node)) - # one row is used by our own name, so cheight - 1 for that allowance - draw_image(imgdata, cwidth, cheight - 1 if cheight else cheight) - if node in nodepositions: - cursor_restore() - reset_cursor(node) - else: - sys.stdout.write('\n') - sys.stdout.flush() + with draw_lock: + for node in imagedatabynode: + imgdata = imagedatabynode[node] + if node in nodepositions: + prep_node_tile(node) + cursor_save() + else: + if options.interval is not None: + if node != firstnodename: + sys.stderr.write('Multiple nodes not supported for interval') + sys.exit(1) + sticky_cursor() + sys.stdout.write('{}: '.format(node)) + # one row is used by our own name, so cheight - 1 for that allowance + draw_image(imgdata, cwidth, cheight - 1 if cheight else cheight) + if node in nodepositions: + cursor_restore() + reset_cursor(node) + else: + sys.stdout.write('\n') + sys.stdout.flush() resized = False inputwatcher = None @@ -961,9 +975,29 @@ async def grab_vncs(urlbynode): try: if streaming: directed = direct_console() + outputbynode = {} + if options.outputfile: + expr = options.outputfile + if '{' in expr: + # Let the server evaluate the expression to a unique name per node + sess = client.Command() + noderange = ','.join(urlbynode) + async for res in sess.create( + '/noderange/{}/attributes/expression'.format(noderange), + {'expression': expr}): + for node in res.get('databynode', {}): + outputbynode[node] = res['databynode'][node]['value'] + else: + for node in urlbynode: + if len(urlbynode) > 1: + base, ext = os.path.splitext(options.outputfile) + outputbynode[node] = '{}-{}{}'.format(base, node, ext) + else: + outputbynode[node] = expr for node in urlbynode: url = urlbynode[node] - tasks.append(asyncio.create_task(do_vnc(node, url))) + tasks.append(asyncio.create_task( + do_vnc(node, url, outputbynode.get(node)))) await asyncio.gather(*tasks) except Exception as e: sys.stderr.write(f"Error in grab_vncs: {e}\n") @@ -1046,13 +1080,13 @@ def toggle_focus_all(): del focused_nodes[node] -async def do_vnc(node, url): +async def do_vnc(node, url, outputfile=None): global streaming keeprunning = True retries = 5 while keeprunning: try: - async with await vnc.VNCClient.create(url) as client: + async with await vnc.VNCClient.create(url, outputfile=outputfile) as client: vncclientsbynode[node] = client while True: # Retrieve pixels as a 3D numpy array @@ -1090,6 +1124,16 @@ async def do_vnc(node, url): def draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight): + # Queue the frame for the drawing thread; frames arriving mid-draw wait here + draw_queue.put((node, imgdata, errorstr, firstnodename, cwidth, cheight)) + +def _draw_worker(): + while True: + node, imgdata, errorstr, firstnodename, cwidth, cheight = draw_queue.get() + with draw_lock: + _draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight) + +def _draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight): imagedatabynode[node] = imgdata if node in nodepositions: prep_node_tile(node) @@ -1113,6 +1157,9 @@ def draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight): sys.stdout.write('\n') sys.stdout.flush() +draw_thread = threading.Thread(target=_draw_worker, daemon=True) +draw_thread.start() + if options.screenshot or options.video: if not sys.stdout.isatty(): sys.stderr.write( diff --git a/confluent_client/confluent/vnc.py b/confluent_client/confluent/vnc.py index 08bb71d8..0b69ea4d 100644 --- a/confluent_client/confluent/vnc.py +++ b/confluent_client/confluent/vnc.py @@ -2,6 +2,9 @@ import asyncio from PIL import Image import io import numpy as np +import queue +import threading +import time import zlib # This results in an RGBA organization of pixels @@ -48,8 +51,27 @@ class VNCClient: return False @classmethod - async def create(cls, url): + async def create(cls, url, outputfile=None, fps=30): self = cls() + self.outputfile = outputfile + self.fps = fps + self.video_writer = None + self._video_size = None + self._last_frame = None + self._last_frame_time = None + self._video_queue = None + self._video_thread = None + self._cv2 = None + if outputfile: + try: + import cv2 + except ImportError: + raise ImportError("OpenCV is required for video output but is not installed.") + self._cv2 = cv2 + self._video_queue = queue.Queue() + self._video_thread = threading.Thread( + target=self._video_worker, daemon=True) + self._video_thread.start() if url.startswith('unix://'): url = url.replace('unix://', '') if url.startswith('/'): @@ -212,8 +234,52 @@ class VNCClient: for _ in range(num_rects): await self._handle_rectangle() self._updating = False + self._write_video_frame() self._request_screen_update(incremental=True) + def _write_video_frame(self): + if not self._cv2 or self.framebuffer is None: + return + # Snapshot the framebuffer now and hand it to the writer thread. Frames + # captured while a write is in progress simply queue up behind it. + frame = np.ascontiguousarray( + np.array(self.framebuffer.convert('RGB'))[:, :, ::-1]) + self._video_queue.put((frame, time.monotonic())) + + def _video_worker(self): + cv2 = self._cv2 + while True: + item = self._video_queue.get() + if item is None: + # Flush the final frame for the time it stayed on screen + if self.video_writer is not None and self._last_frame is not None: + nframes = max(1, round( + (time.monotonic() - self._last_frame_time) * self.fps)) + for _ in range(nframes): + self.video_writer.write(self._last_frame) + if self.video_writer is not None: + self.video_writer.release() + self.video_writer = None + return + frame, now = item + if self.video_writer is None: + self._video_size = (frame.shape[1], frame.shape[0]) + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + self.video_writer = cv2.VideoWriter( + self.outputfile, fourcc, self.fps, self._video_size) + if (frame.shape[1], frame.shape[0]) != self._video_size: + frame = cv2.resize(frame, self._video_size) + if self._last_frame is None: + self._last_frame = frame + self._last_frame_time = now + continue + # Hold the previous frame for the real time it was displayed + nframes = max(1, round((now - self._last_frame_time) * self.fps)) + for _ in range(nframes): + self.video_writer.write(self._last_frame) + self._last_frame = frame + self._last_frame_time = now + async def _handle_rectangle(self): if self.framebuffer is None: self.framebuffer = Image.new('RGBA', (self.width, self.height)) @@ -272,5 +338,10 @@ class VNCClient: break return length async def close(self): + if self._video_thread is not None: + # Signal the writer thread to flush and finalize the file + self._video_queue.put(None) + await asyncio.to_thread(self._video_thread.join) + self._video_thread = None self.writer.close() await self.writer.wait_closed()