From eee0a9b1c1abff31fd0cb524721ea87ed3e4e745 Mon Sep 17 00:00:00 2001 From: fastagents Date: Mon, 29 Jun 2026 11:03:12 -0600 Subject: [PATCH] Add live stream memory cleanup hook --- live/streamdiffusion/pipeline/pipeline.py | 44 ++++++++ runner/src/runner/live/pipelines/interface.py | 9 ++ runner/src/runner/live/process/process.py | 63 +++++++++-- .../runner/live/process/process_guardian.py | 5 +- .../runner/live/streamer/protocol/trickle.py | 63 +++++++++-- runner/src/runner/live/trickle/encoder.py | 91 +++++++++------- .../tests/live/test_stream_cleanup_source.py | 103 ++++++++++++++++++ 7 files changed, 318 insertions(+), 60 deletions(-) create mode 100644 runner/tests/live/test_stream_cleanup_source.py diff --git a/live/streamdiffusion/pipeline/pipeline.py b/live/streamdiffusion/pipeline/pipeline.py index 77e5a5e47..271870741 100644 --- a/live/streamdiffusion/pipeline/pipeline.py +++ b/live/streamdiffusion/pipeline/pipeline.py @@ -2,6 +2,8 @@ import logging import asyncio import base64 +import ctypes +import gc import re from pathlib import Path from typing import Dict, List, Optional, Any, cast @@ -30,6 +32,10 @@ ENGINES_DIR = Path("./engines") # this one is used only for realesrgan_trt which has ./models hardcoded LOCAL_MODELS_DIR = Path("./models") +STREAM_CLEANUP_GC_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_GC" +STREAM_CLEANUP_MALLOC_TRIM_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_MALLOC_TRIM" +STREAM_CLEANUP_CUDA_CACHE_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_CUDA_CACHE" +STREAM_CLEANUP_STYLE_CACHE_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_STYLE_CACHE" class StreamDiffusion(Pipeline): def __init__(self): @@ -280,6 +286,21 @@ async def stop(self): self.params = None self.frame_queue = asyncio.Queue() + async def cleanup_stream(self): + async with self._pipeline_lock: + _clear_asyncio_queue(self.frame_queue) + if _env_enabled(STREAM_CLEANUP_STYLE_CACHE_ENV, default=False): + self._cached_style_image_tensor = None + self._cached_style_image_url = None + + if _env_enabled(STREAM_CLEANUP_GC_ENV, default=True): + await asyncio.to_thread(gc.collect) + if _env_enabled(STREAM_CLEANUP_MALLOC_TRIM_ENV, default=True): + await asyncio.to_thread(_malloc_trim) + if _env_enabled(STREAM_CLEANUP_CUDA_CACHE_ENV, default=False) and torch.cuda.is_available(): + torch.cuda.empty_cache() + logging.info("StreamDiffusion stream cleanup complete") + @classmethod def prepare_models(cls): from .prepare import prepare_streamdiffusion_models @@ -287,6 +308,29 @@ def prepare_models(cls): prepare_streamdiffusion_models() +def _env_enabled(name: str, *, default: bool) -> bool: + raw = os.getenv(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +def _malloc_trim() -> None: + try: + libc = ctypes.CDLL("libc.so.6") + libc.malloc_trim(0) + except Exception: + logging.debug("malloc_trim unavailable", exc_info=True) + + +def _clear_asyncio_queue(queue: asyncio.Queue) -> None: + while True: + try: + queue.get_nowait() + except asyncio.QueueEmpty: + return + + def _prepare_controlnet_configs(params: StreamDiffusionParams) -> Optional[List[Dict[str, Any]]]: """Prepare ControlNet configurations for wrapper""" if not params.controlnets: diff --git a/runner/src/runner/live/pipelines/interface.py b/runner/src/runner/live/pipelines/interface.py index 1cdc245a2..fc3d24dd5 100644 --- a/runner/src/runner/live/pipelines/interface.py +++ b/runner/src/runner/live/pipelines/interface.py @@ -108,6 +108,15 @@ async def stop(self): """ pass + async def cleanup_stream(self): + """Release per-stream state while keeping the warmed pipeline alive. + + Called after a stream goes offline. Implementations should clear + frame queues, per-stream caches, and allocator state that does not need + to survive into the next stream, without unloading model weights. + """ + pass + @classmethod @abstractmethod def prepare_models(cls): diff --git a/runner/src/runner/live/process/process.py b/runner/src/runner/live/process/process.py index 07b11edeb..e19e6af79 100644 --- a/runner/src/runner/live/process/process.py +++ b/runner/src/runner/live/process/process.py @@ -28,6 +28,8 @@ from .loading_overlay import LoadingOverlayRenderer +STREAM_CLEANUP_CONTROL_KEY = "__stream_cleanup__" + class PipelineProcess: @staticmethod @@ -139,20 +141,28 @@ def _set_pipeline_ready(self, ready: bool): def update_params(self, params: dict): self.param_update_queue.put(params) - def reset_stream(self, request_id: str, manifest_id: str, stream_id: str): + def reset_stream( + self, + request_id: str, + manifest_id: str, + stream_id: str, + *, + cleanup_stream: bool = False, + ): # Clear queues to avoid using frames from previous sessions clear_queue(self.input_queue) clear_queue(self.output_queue) clear_queue(self.param_update_queue) clear_queue(self.error_queue) clear_queue(self.log_queue) - self.param_update_queue.put( - { - "request_id": request_id, - "manifest_id": manifest_id, - "stream_id": stream_id, - } - ) + message = { + "request_id": request_id, + "manifest_id": manifest_id, + "stream_id": stream_id, + } + if cleanup_stream: + message[STREAM_CLEANUP_CONTROL_KEY] = True + self.param_update_queue.put(message) # TODO: Once audio is implemented, combined send_input with input_loop # We don't need additional queueing as comfystream already maintains a queue @@ -330,7 +340,7 @@ async def _param_update_loop(self, pipeline: Pipeline, overlay: LoadingOverlayRe while not self.is_done(): reload_task = None try: - params = await self._get_latest_params(timeout=0.1) + params = await self._get_latest_params(timeout=0.1, pipeline=pipeline, overlay=overlay) if params is None: continue @@ -382,7 +392,12 @@ async def _param_update_loop(self, pipeline: Pipeline, overlay: LoadingOverlayRe except Exception: logging.warning("Failed to prewarm loading overlay caches", exc_info=True) - async def _get_latest_params(self, timeout: float) -> dict | None: + async def _get_latest_params( + self, + timeout: float, + pipeline: Pipeline, + overlay: LoadingOverlayRenderer, + ) -> dict | None: """ Get the latest params from the param_update_queue, skipping stale entries before the latest. Already filters and processes params that are only logging updates. Waits for timeout seconds for a new params entry, or returns @@ -396,20 +411,44 @@ async def _get_latest_params(self, timeout: float) -> dict | None: except queue.Empty: return None - if self._handle_logging_params(params): + if await self._handle_process_message(params, pipeline, overlay): params = None # Drain the params queue to get the latest params while not self.param_update_queue.empty(): try: new_params = self.param_update_queue.get_nowait() - if not self._handle_logging_params(new_params): + if not await self._handle_process_message(new_params, pipeline, overlay): params = new_params except queue.Empty: break return params + async def _handle_process_message( + self, + params: dict, + pipeline: Pipeline, + overlay: LoadingOverlayRenderer, + ) -> bool: + consumed = self._handle_logging_params(params) + if isinstance(params, dict) and params.get(STREAM_CLEANUP_CONTROL_KEY): + await self._cleanup_stream_state(pipeline, overlay) + consumed = True + return consumed + + async def _cleanup_stream_state(self, pipeline: Pipeline, overlay: LoadingOverlayRenderer): + logging.info("PipelineProcess: stream cleanup requested") + clear_queue(self.input_queue) + clear_queue(self.output_queue) + overlay.end_reload() + try: + await pipeline.cleanup_stream() + except Exception as e: + self._report_error("Error cleaning up pipeline stream state", e) + else: + logging.info("PipelineProcess: stream cleanup complete") + def _report_error(self, msg: str, error: Exception | None = None, silent=False): if not silent: logging.error(msg, exc_info=error) diff --git a/runner/src/runner/live/process/process_guardian.py b/runner/src/runner/live/process/process_guardian.py index 50659bd66..bd4a0693d 100644 --- a/runner/src/runner/live/process/process_guardian.py +++ b/runner/src/runner/live/process/process_guardian.py @@ -301,8 +301,9 @@ async def _monitor_loop(self): continue if state == PipelineState.OFFLINE: - # Revert to initial params when the stream stops - self.process.reset_stream("", "", "") + # Revert to initial params and release stream-scoped state + # while keeping the warmed pipeline alive. + self.process.reset_stream("", "", "", cleanup_stream=True) self.process.update_params(self.pipeline_spec.initial_params) if state != PipelineState.ERROR: diff --git a/runner/src/runner/live/streamer/protocol/trickle.py b/runner/src/runner/live/streamer/protocol/trickle.py index c9c4c11c6..139f704c2 100644 --- a/runner/src/runner/live/streamer/protocol/trickle.py +++ b/runner/src/runner/live/streamer/protocol/trickle.py @@ -1,5 +1,6 @@ import asyncio import logging +import os import queue import json from typing import AsyncGenerator, Optional @@ -11,6 +12,40 @@ from .protocol import StreamProtocol from .last_value_cache import LastValueCache +DEFAULT_PUBLISH_QUEUE_MAXSIZE = 1 + + +def _publish_queue_maxsize() -> int: + raw = os.environ.get("LIVE_TRICKLE_PUBLISH_QUEUE_MAXSIZE") + if raw is None or raw.strip() == "": + return DEFAULT_PUBLISH_QUEUE_MAXSIZE + try: + value = int(raw) + except ValueError: + logging.warning( + "Invalid LIVE_TRICKLE_PUBLISH_QUEUE_MAXSIZE=%r, using %s", + raw, + DEFAULT_PUBLISH_QUEUE_MAXSIZE, + ) + return DEFAULT_PUBLISH_QUEUE_MAXSIZE + return max(1, value) + + +def _put_drop_oldest(_queue: queue.Queue, item) -> int: + dropped = 0 + while True: + try: + _queue.put_nowait(item) + return dropped + except queue.Full: + try: + old = _queue.get_nowait() + except queue.Empty: + continue + dropped += 1 + del old + + class TrickleProtocol(StreamProtocol): def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[str] = None, events_url: Optional[str] = None, input_width: Optional[int] = DEFAULT_WIDTH, input_height: Optional[int] = DEFAULT_HEIGHT, output_width: Optional[int] = None, output_height: Optional[int] = None): self.subscribe_url = subscribe_url @@ -18,7 +53,8 @@ def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[s self.control_url = control_url self.events_url = events_url self.subscribe_queue = queue.Queue[InputFrame]() - self.publish_queue = queue.Queue[OutputFrame]() + self.publish_queue = queue.Queue[OutputFrame](maxsize=_publish_queue_maxsize()) + self.publish_queue_dropped = 0 self.control_subscriber = None self.events_publisher = None self.subscribe_task = None @@ -30,7 +66,8 @@ def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[s async def start(self): self.subscribe_queue = queue.Queue[InputFrame]() - self.publish_queue = queue.Queue[OutputFrame]() + self.publish_queue = queue.Queue[OutputFrame](maxsize=_publish_queue_maxsize()) + self.publish_queue_dropped = 0 metadata_cache = LastValueCache[dict]() # to pass video metadata from decoder to encoder self.subscribe_task = asyncio.create_task( media.run_subscribe(self.subscribe_url, self.subscribe_queue.put, metadata_cache.put, self.emit_monitoring_event, self.input_width, self.input_height) @@ -49,7 +86,7 @@ async def stop(self): # send sentinel None values to stop the trickle tasks gracefully self.subscribe_queue.put(None) - self.publish_queue.put(None) + _put_drop_oldest(self.publish_queue, None) if self.control_subscriber: await self.control_subscriber.close() @@ -86,17 +123,30 @@ def dequeue_frame(): # TEMP: Put audio immediately into the publish queue # TODO: Remove once there is ComfyUI audio support if isinstance(image, AudioFrame): - publish_queue.put(AudioOutput([image])) + self.publish_queue_dropped += _put_drop_oldest( + publish_queue, + AudioOutput([image]), + ) continue yield image async def egress_loop(self, output_frames: AsyncGenerator[OutputFrame, None]): publish_queue = self.publish_queue def enqueue_bytes(frame: OutputFrame): - publish_queue.put(frame) + return _put_drop_oldest(publish_queue, frame) async for frame in output_frames: - await asyncio.to_thread(enqueue_bytes, frame) + try: + dropped = await asyncio.to_thread(enqueue_bytes, frame) + self.publish_queue_dropped += dropped + if dropped and self.publish_queue_dropped % 100 == 1: + logging.warning( + "Dropped stale publish frames dropped_total=%s queue_maxsize=%s", + self.publish_queue_dropped, + publish_queue.maxsize, + ) + finally: + del frame async def emit_monitoring_event(self, event: dict, queue_event_type: str = "ai_stream_events"): if not self.events_publisher: @@ -134,4 +184,3 @@ async def control_loop(self, done: asyncio.Event) -> AsyncGenerator[dict, None]: except Exception: logging.error(f"Error in control loop", exc_info=True) continue - diff --git a/runner/src/runner/live/trickle/encoder.py b/runner/src/runner/live/trickle/encoder.py index fb3f8a885..30182b6fd 100644 --- a/runner/src/runner/live/trickle/encoder.py +++ b/runner/src/runner/live/trickle/encoder.py @@ -99,47 +99,60 @@ def custom_io_open(url: str, flags: int, options: dict): break if isinstance(avframe, VideoOutput): - if not output_video_stream: - # received video but no video output, so drop - continue - avframe.log_timestamps["frame_end"] = time.time() - log_frame_timestamps("Video", avframe.frame) - - tensor = avframe.tensor.squeeze(0) - image_np = (tensor * 255).byte().cpu().numpy() - image = Image.fromarray(image_np) - - frame = av.video.frame.VideoFrame.from_image(image) - frame.pts = rescale_ts(avframe.timestamp, avframe.time_base, output_video_stream.codec_context.time_base) - frame.time_base = output_video_stream.codec_context.time_base - current = avframe.timestamp * avframe.time_base - - # if there is pending audio, check whether video is too far behind - if len(audio_buffer) > 0: - audio_ts = audio_buffer[0].timestamp * audio_buffer[0].time_base - delta = audio_ts - current - if delta > AUDIO_BUFFER_SECS: - # video is too far behind audio, drop the frame to try and catch up - dropped_video += 1 - if dropped_video > MAX_DROPPED_VIDEO: - # dropped too many frames, may be unlikely to catch up so stop the stream - logging.error(f"A/V is out of sync, exiting from video audio_ts={float(audio_ts)} video_ts={float(current)} delta={float(delta)}") - break - # drop the frame by skipping the rest of the following code + tensor = None + image_np = None + image = None + frame = None + encoded_packets = None + try: + if not output_video_stream: + # received video but no video output, so drop continue + avframe.log_timestamps["frame_end"] = time.time() + log_frame_timestamps("Video", avframe.frame) + + tensor = avframe.tensor.squeeze(0) + image_np = (tensor * 255).byte().cpu().numpy() + image = Image.fromarray(image_np) + + frame = av.video.frame.VideoFrame.from_image(image) + frame.pts = rescale_ts(avframe.timestamp, avframe.time_base, output_video_stream.codec_context.time_base) + frame.time_base = output_video_stream.codec_context.time_base + current = avframe.timestamp * avframe.time_base + + # if there is pending audio, check whether video is too far behind + if len(audio_buffer) > 0: + audio_ts = audio_buffer[0].timestamp * audio_buffer[0].time_base + delta = audio_ts - current + if delta > AUDIO_BUFFER_SECS: + # video is too far behind audio, drop the frame to try and catch up + dropped_video += 1 + if dropped_video > MAX_DROPPED_VIDEO: + # dropped too many frames, may be unlikely to catch up so stop the stream + logging.error(f"A/V is out of sync, exiting from video audio_ts={float(audio_ts)} video_ts={float(current)} delta={float(delta)}") + break + # drop the frame by skipping the rest of the following code + continue - if not last_kf or float(current - last_kf) >= GOP_SECS: - frame.pict_type = av.video.frame.PictureType.I - last_kf = current - if first_video_ts is None: - first_video_ts = current - encoded_packets = output_video_stream.encode(frame) - for ep in encoded_packets: - if last_video_ts and frame.pts <= last_video_ts: - logging.info(f"Timestamps out of order; dropping video last={last_video_ts} current={frame.pts}") - continue - last_video_ts = frame.pts - output_container.mux(ep) + if not last_kf or float(current - last_kf) >= GOP_SECS: + frame.pict_type = av.video.frame.PictureType.I + last_kf = current + if first_video_ts is None: + first_video_ts = current + encoded_packets = output_video_stream.encode(frame) + for ep in encoded_packets: + if last_video_ts and frame.pts <= last_video_ts: + logging.info(f"Timestamps out of order; dropping video last={last_video_ts} current={frame.pts}") + continue + last_video_ts = frame.pts + output_container.mux(ep) + finally: + del encoded_packets + del frame + del image + del image_np + del tensor + del avframe continue if isinstance(avframe, AudioOutput): diff --git a/runner/tests/live/test_stream_cleanup_source.py b/runner/tests/live/test_stream_cleanup_source.py new file mode 100644 index 000000000..8b94f9695 --- /dev/null +++ b/runner/tests/live/test_stream_cleanup_source.py @@ -0,0 +1,103 @@ +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[3] +PIPELINE_PATH = REPO_ROOT / "live" / "streamdiffusion" / "pipeline" / "pipeline.py" +PROCESS_PATH = REPO_ROOT / "runner" / "src" / "runner" / "live" / "process" / "process.py" +GUARDIAN_PATH = ( + REPO_ROOT / "runner" / "src" / "runner" / "live" / "process" / "process_guardian.py" +) +TRICKLE_PATH = ( + REPO_ROOT + / "runner" + / "src" + / "runner" + / "live" + / "streamer" + / "protocol" + / "trickle.py" +) +ENCODER_PATH = REPO_ROOT / "runner" / "src" / "runner" / "live" / "trickle" / "encoder.py" +INTERFACE_PATH = ( + REPO_ROOT / "runner" / "src" / "runner" / "live" / "pipelines" / "interface.py" +) + + +def _function_source(text: str, signature: str, next_marker: str) -> str: + start = text.index(signature) + end = text.index(next_marker, start) + return text[start:end] + + +def test_streamdiffusion_cleanup_stream_drains_queue_in_place(): + text = PIPELINE_PATH.read_text() + cleanup_source = _function_source( + text, + " async def cleanup_stream(self):", + " @classmethod", + ) + + assert "self.frame_queue =" not in cleanup_source + assert "_clear_asyncio_queue(self.frame_queue)" in cleanup_source + assert "gc.collect" in text + assert "malloc_trim" in text + assert "LIVE_PIPELINE_STREAM_CLEANUP_CUDA_CACHE" in text + + +def test_process_cleanup_message_reaches_child_pipeline(): + text = PROCESS_PATH.read_text() + + assert "STREAM_CLEANUP_CONTROL_KEY" in text + assert "cleanup_stream: bool = False" in text + assert "message[STREAM_CLEANUP_CONTROL_KEY] = True" in text + assert "await pipeline.cleanup_stream()" in text + assert "PipelineProcess: stream cleanup requested" in text + + +def test_guardian_requests_cleanup_when_stream_goes_offline(): + text = GUARDIAN_PATH.read_text() + + assert 'self.process.reset_stream("", "", "", cleanup_stream=True)' in text + assert "self.process.update_params(self.pipeline_spec.initial_params)" in text + + +def test_pipeline_interface_exposes_cleanup_hook(): + text = INTERFACE_PATH.read_text() + + assert "async def cleanup_stream(self):" in text + assert "while keeping the warmed pipeline alive" in text + + +def test_publish_queue_is_bounded_and_drops_stale_frames(): + text = TRICKLE_PATH.read_text() + + assert "DEFAULT_PUBLISH_QUEUE_MAXSIZE = 1" in text + assert "LIVE_TRICKLE_PUBLISH_QUEUE_MAXSIZE" in text + assert "queue.Queue[OutputFrame](maxsize=_publish_queue_maxsize())" in text + assert "def _put_drop_oldest" in text + assert "old = _queue.get_nowait()" in text + assert "del old" in text + assert "self.publish_queue_dropped" in text + + +def test_trickle_egress_releases_output_frame_reference(): + text = TRICKLE_PATH.read_text() + egress_source = _function_source( + text, + " async def egress_loop(self, output_frames: AsyncGenerator[OutputFrame, None]):", + " async def emit_monitoring_event", + ) + + assert "_put_drop_oldest(publish_queue, frame)" in egress_source + assert "finally:" in egress_source + assert "del frame" in egress_source + + +def test_encoder_releases_video_frame_conversion_locals(): + text = ENCODER_PATH.read_text() + + assert "tensor = None" in text + assert "image_np = None" in text + assert "encoded_packets = None" in text + assert "finally:" in text + assert "del avframe" in text