Repository: livepeer/ai-runner
Pull request #454: trickle: Retry encoder if necessary
Author: login: j0sh
Created: 2025-03-13T22:48:12Z | Merged: 2025-03-17T21:56:22Z

This requires adding a last value cache for the media metadata rather than pulling the metadata from a queue, because the queue may be empty when the encoder restarts.

We also need to explicitly close the write end of any pipes that may be dangling after the exception, since pyav will no longer be closing those.


Diff:
diff --git a/runner/app/live/streamer/protocol/last_value_cache.py b/runner/app/live/streamer/protocol/last_value_cache.py
new file mode 100644
index 000000000..c5d7444d2
--- /dev/null
+++ b/runner/app/live/streamer/protocol/last_value_cache.py
@@ -0,0 +1,46 @@
+import threading
+from typing import Optional, TypeVar, Generic, Callable
+
+T = TypeVar('T')
+
+class LastValueCache(Generic[T]):
+    """
+    A thread-safe cache that stores the last value produced.
+    Consumers can wait for a value to be available with a timeout.
+    """
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._condition = threading.Condition(self._lock)
+        self._UNSET = object()
+        self._value = self._UNSET
+
+    def put(self, value: T) -> None:
+        """Store a new value and notify all waiting consumers."""
+        with self._lock:
+            self._value = value
+            self._condition.notify_all()
+
+    def get(self, timeout: float = 30.0) -> Optional[T]:
+        """
+        Get the latest value, waiting if necessary.
+
+        Args:
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The latest value, or None if timeout occurs before a value is available
+        """
+        def value_is_set():
+            return self._value is not self._UNSET
+
+        with self._condition:
+            # Wait for a value to be available
+            if not value_is_set():
+                success = self._condition.wait_for(value_is_set, timeout=timeout)
+                if not success:
+                    logging.warning(f"Timed out waiting for value (timeout={timeout}s)")
+                    return None
+
+            return self._value
+
+
diff --git a/runner/app/live/streamer/protocol/trickle.py b/runner/app/live/streamer/protocol/trickle.py
index 46a9c496d..557c94b5a 100644
--- a/runner/app/live/streamer/protocol/trickle.py
+++ b/runner/app/live/streamer/protocol/trickle.py
@@ -9,6 +9,7 @@
 from trickle import media, TricklePublisher, TrickleSubscriber, InputFrame, OutputFrame, AudioFrame, AudioOutput
 
 from .protocol import StreamProtocol
+from .last_value_cache import LastValueCache
 
 class TrickleProtocol(StreamProtocol):
     def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[str] = None, events_url: Optional[str] = None):
@@ -24,12 +25,12 @@ def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[s
         self.publish_task = None
 
     async def start(self):
-        metadata_queue = queue.Queue[dict]() # to pass video metadata from decoder to encoder
+        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_queue.put, self.emit_monitoring_event)
+            media.run_subscribe(self.subscribe_url, self.subscribe_queue.put, metadata_cache.put, self.emit_monitoring_event)
         )
         self.publish_task = asyncio.create_task(
-            media.run_publish(self.publish_url, self.publish_queue.get, metadata_queue.get, self.emit_monitoring_event)
+            media.run_publish(self.publish_url, self.publish_queue.get, metadata_cache.get, self.emit_monitoring_event)
         )
         if self.control_url and self.control_url.strip() != "":
             self.control_subscriber = TrickleSubscriber(self.control_url)
diff --git a/runner/app/live/trickle/encoder.py b/runner/app/live/trickle/encoder.py
index c3f951d6b..e800e5ccf 100644
--- a/runner/app/live/trickle/encoder.py
+++ b/runner/app/live/trickle/encoder.py
@@ -36,7 +36,7 @@ def custom_io_open(url: str, flags: int, options: dict):
         read_fd, write_fd = os.pipe()
         read_file  = os.fdopen(read_fd,  'rb', buffering=0)
         write_file = os.fdopen(write_fd, 'wb', buffering=0)
-        output_callback(read_file, url)
+        output_callback(read_file, write_file, url)
         return write_file
 
     # Open the output container in write mode
@@ -138,13 +138,13 @@ def rescale_ts(pts: int, orig_tb: Fraction, dest_tb: Fraction):
 
 def log_frame_timestamps(frame_type: str, frame: InputFrame):
     ts = frame.log_timestamps
-    
+
     def log_duration(start_key: str, end_key: str):
         if start_key in ts and end_key in ts:
             duration = ts[end_key] - ts[start_key]
             logging.debug(f"frame_type={frame_type} start_tag={start_key} end_tag={end_key} duration_s={duration}s")
-    
+
     log_duration('frame_init', 'pre_process_frame')
     log_duration('pre_process_frame', 'post_process_frame')
     log_duration('post_process_frame', 'frame_end')
-    log_duration('frame_init', 'frame_end')
\ No newline at end of file
+    log_duration('frame_init', 'frame_end')
diff --git a/runner/app/live/trickle/media.py b/runner/app/live/trickle/media.py
index c16f052e3..528305be1 100644
--- a/runner/app/live/trickle/media.py
+++ b/runner/app/live/trickle/media.py
@@ -11,6 +11,9 @@
 from .decoder import decode_av
 from .encoder import encode_av
 
+MAX_ENCODER_RETRIES = 3
+ENCODER_RETRY_RESET_SECONDS = 120 # reset retry counter after 2 minutes
+
 async def run_subscribe(subscribe_url: str, image_callback, put_metadata, monitoring_callback):
     # TODO add some pre-processing parameters, eg image size
     try:
@@ -81,6 +84,40 @@ def decode_runner():
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, decode_runner)
 
+def encode_in(task_pipes, task_lock, image_generator, sync_callback, get_metadata, **kwargs):
+    # encode_av has a tendency to crash, so restart as necessary
+    retryCount = 0
+    last_retry_time = time.time()
+    while retryCount < MAX_ENCODER_RETRIES:
+        try:
+            encode_av(image_generator, sync_callback, get_metadata, **kwargs)
+            break  # clean exit
+        except Exception as exc:
+            current_time = time.time()
+            # Reset retry counter if enough time has elapsed
+            if current_time - last_retry_time > ENCODER_RETRY_RESET_SECONDS:
+                logging.info("Resetting encoder retry count")
+                retryCount = 0
+            retryCount += 1
+            last_retry_time = current_time
+            if retryCount < MAX_ENCODER_RETRIES:
+                logging.exception(f"Error in encode_av, retrying {retryCount}/{MAX_ENCODER_RETRIES}", stack_info=True)
+            else:
+                logging.exception("Error in encode_av, maximum retries reached", stack_info=True)
+            # close leftover writer ends of any pipes to prevent hanging
+            pipe_count = 0
+            total_pipes = 0
+            with task_lock:
+                pipes = list(task_pipes)
+                total_pipes = len(pipes)
+                for p in pipes:
+                    try:
+                        p.close()
+                        pipe_count += 1
+                    except Exception as e:
+                        logging.exception("Error closing pipe on task list", stack_info=True)
+            logging.info(f"Closed pipes - {pipe_count}/{total_pipes}")
+
 async def run_publish(publish_url: str, image_generator, get_metadata, monitoring_callback):
     first_segment = True
 
@@ -111,30 +148,36 @@ async def callback(pipe_file, pipe_name):
                         }, queue_event_type="stream_trace")
                 transport.close()
 
-        def sync_callback(pipe_file, pipe_name):
+        def sync_callback(pipe_reader, pipe_writer, pipe_name):
             def do_schedule():
-                schedule_callback(callback(pipe_file, pipe_name), pipe_name)
+                schedule_callback(callback(pipe_reader, pipe_name), pipe_writer, pipe_name)
             loop.call_soon_threadsafe(do_schedule)
 
         # hold tasks since `loop.create_task` is a weak reference that gets GC'd
         # TODO use asyncio.TaskGroup once all pipelines are on Python 3.11+
+        # Also hold pipes in case we need to explicitly close them during exceptions
         live_tasks = set()
+        live_pipes = set()
         live_tasks_lock = threading.Lock()
 
-        def schedule_callback(coro, pipe_name):
+        def schedule_callback(coro, pipe_writer, pipe_name):
             task = loop.create_task(coro)
             with live_tasks_lock:
                 live_tasks.add(task)
-            def task_done(t: asyncio.Task):
+                live_pipes.add(pipe_writer)
+            def task_done2(t: asyncio.Task, p):
                 try:
                     t.result()
                 except Exception as e:
                     logging.error(f"Task {pipe_name} crashed: {e}")
                 with live_tasks_lock:
                     live_tasks.remove(t)
+                    live_pipes.remove(p)
+            def task_done(t2:asyncio.Task):
+                return task_done2(task, pipe_writer)
             task.add_done_callback(task_done)
 
-        encode_thread = threading.Thread(target=encode_av, args=(image_generator, sync_callback, get_metadata), kwargs={"audio_codec":"libopus"})
+        encode_thread = threading.Thread(target=encode_in, args=(live_pipes, live_tasks_lock, image_generator, sync_callback, get_metadata), kwargs={"audio_codec":"libopus"})
         encode_thread.start()
         logging.debug("run_publish: encoder thread started")
 


Reviews:
user: login: emranemran
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T05:17:38Z
user: login: victorges
body: LGTM mostly a question to understand a lil more how it works

state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T16:15:53Z
user: login: j0sh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T20:51:56Z
user: login: j0sh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T21:02:43Z
user: login: j0sh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T21:29:34Z

Review comments:
diff_hunk: @@ -81,6 +84,38 @@ def decode_runner():
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, decode_runner)
 
+def encode_in(task_pipes, task_lock, image_generator, sync_callback, get_metadata, **kwargs):
+    # encode_av has a tendency to crash, so restart as necessary
+    retryCount = 0
+    last_retry_time = time.time()
+    while retryCount < MAX_ENCODER_RETRIES:
+        try:
+            encode_av(image_generator, sync_callback, get_metadata, **kwargs)
+            break  # clean exit
+        except Exception as exc:
+            current_time = time.time()
+            # Reset retry counter if enough time has elapsed
+            if current_time - last_retry_time > ENCODER_RETRY_RESET_SECONDS:
+                logging.info("Resetting encoder retry count")
+                retryCount = 0
+            retryCount += 1
+            last_retry_time = current_time
+            if retryCount < MAX_ENCODER_RETRIES:
+                logging.exception(f"Error in encode_av, retrying {retryCount}/{MAX_ENCODER_RETRIES}", stack_info=True)
+            else:
+                logging.exception("Error in encode_av, maximum retries reached", stack_info=True)
+            # close leftover writer ends of any pipes to prevent hanging
+            pipe_count = 0
+            with task_lock:
+                pipes = list(task_pipes)
+                for p in pipes:
+                    try:
+                        p.close()
+                        pipe_count += 1
+                    except Exception as e:
+                        logging.exception("Error closing pipe on task list", stack_info=True)
+            logging.info(f"Closed pipes - {pipe_count}")
path: runner/app/live/trickle/media.py
user: login: emranemran
body: nit: would be good to see what the total pipe count is in the same line here so that we can easily tell how many pipes failed to close. 
created_at: 2025-03-17T05:15:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,46 @@
+import threading
+from typing import Optional, TypeVar, Generic, Callable
+
+T = TypeVar('T')
+
+class LastValueCache(Generic[T]):
+    """
+    A thread-safe cache that stores the last value produced.
+    Consumers can wait for a value to be available with a timeout.
+    """
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._condition = threading.Condition(self._lock)
+        self._UNSET = object()
+        self._value = self._UNSET
+
+    def put(self, value: T) -> None:
+        """Store a new value and notify all waiting consumers."""
+        with self._lock:
+            self._value = value
+            self._condition.notify_all()
+
+    def get(self, timeout: float = 30.0) -> Optional[T]:
+        """
+        Get the latest value, waiting if necessary.
+
+        Args:
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The latest value, or None if timeout occurs before a value is available
+        """
+        def value_is_set():
+            return self._value is not self._UNSET
+
+        with self._condition:
+            # Wait for a value to be available
+            if not value_is_set():
+                success = self._condition.wait_for(value_is_set, timeout=timeout)
+                if not success:
+                    logging.warning(f"Timed out waiting for value (timeout={timeout}s)")
path: runner/app/live/streamer/protocol/last_value_cache.py
user: login: victorges
body: Should this be a `TimeoutError`?
created_at: 2025-03-17T15:23:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -81,6 +84,38 @@ def decode_runner():
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, decode_runner)
 
+def encode_in(task_pipes, task_lock, image_generator, sync_callback, get_metadata, **kwargs):
+    # encode_av has a tendency to crash, so restart as necessary
+    retryCount = 0
+    last_retry_time = time.time()
+    while retryCount < MAX_ENCODER_RETRIES:
+        try:
+            encode_av(image_generator, sync_callback, get_metadata, **kwargs)
path: runner/app/live/trickle/media.py
user: login: victorges
body: How come (it seems like) we changed the `sync_callback` signature but it is used with the `encode_av` function in the same way? Does `encode_av` have some kind of logic to check the signature of the function and call it in a different way?
created_at: 2025-03-17T16:11:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -81,6 +84,38 @@ def decode_runner():
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, decode_runner)
 
+def encode_in(task_pipes, task_lock, image_generator, sync_callback, get_metadata, **kwargs):
+    # encode_av has a tendency to crash, so restart as necessary
+    retryCount = 0
+    last_retry_time = time.time()
+    while retryCount < MAX_ENCODER_RETRIES:
+        try:
+            encode_av(image_generator, sync_callback, get_metadata, **kwargs)
path: runner/app/live/trickle/media.py
user: login: j0sh
body: Not sure that I understand the question - we do change the callback signature, but we also change how it is called from `encode_av` [here](https://github.com/livepeer/ai-runner/pull/454/files/efa410fff18b926b05fe6f8f199fe5d5699b7af9#diff-58ccb9520686dccb238caf0d92a8d842e08674435226cf5aeaaaf023bd17573fR39)
created_at: 2025-03-17T20:51:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,46 @@
+import threading
+from typing import Optional, TypeVar, Generic, Callable
+
+T = TypeVar('T')
+
+class LastValueCache(Generic[T]):
+    """
+    A thread-safe cache that stores the last value produced.
+    Consumers can wait for a value to be available with a timeout.
+    """
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._condition = threading.Condition(self._lock)
+        self._UNSET = object()
+        self._value = self._UNSET
+
+    def put(self, value: T) -> None:
+        """Store a new value and notify all waiting consumers."""
+        with self._lock:
+            self._value = value
+            self._condition.notify_all()
+
+    def get(self, timeout: float = 30.0) -> Optional[T]:
+        """
+        Get the latest value, waiting if necessary.
+
+        Args:
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The latest value, or None if timeout occurs before a value is available
+        """
+        def value_is_set():
+            return self._value is not self._UNSET
+
+        with self._condition:
+            # Wait for a value to be available
+            if not value_is_set():
+                success = self._condition.wait_for(value_is_set, timeout=timeout)
+                if not success:
+                    logging.warning(f"Timed out waiting for value (timeout={timeout}s)")
path: runner/app/live/streamer/protocol/last_value_cache.py
user: login: j0sh
body: you mean raise an exception instead of returning None? That is probably more pythonic, but the earlier code also used None to signal EOF and I didn't want to change the behavior in that surrounding code too much right now. In any case the effect should be identical, the stream will terminate. With an exception it will probably retry a few times though - where it is not likely to succeed - which would just prolong the tear-down process.
created_at: 2025-03-17T21:02:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -81,6 +84,38 @@ def decode_runner():
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, decode_runner)
 
+def encode_in(task_pipes, task_lock, image_generator, sync_callback, get_metadata, **kwargs):
+    # encode_av has a tendency to crash, so restart as necessary
+    retryCount = 0
+    last_retry_time = time.time()
+    while retryCount < MAX_ENCODER_RETRIES:
+        try:
+            encode_av(image_generator, sync_callback, get_metadata, **kwargs)
+            break  # clean exit
+        except Exception as exc:
+            current_time = time.time()
+            # Reset retry counter if enough time has elapsed
+            if current_time - last_retry_time > ENCODER_RETRY_RESET_SECONDS:
+                logging.info("Resetting encoder retry count")
+                retryCount = 0
+            retryCount += 1
+            last_retry_time = current_time
+            if retryCount < MAX_ENCODER_RETRIES:
+                logging.exception(f"Error in encode_av, retrying {retryCount}/{MAX_ENCODER_RETRIES}", stack_info=True)
+            else:
+                logging.exception("Error in encode_av, maximum retries reached", stack_info=True)
+            # close leftover writer ends of any pipes to prevent hanging
+            pipe_count = 0
+            with task_lock:
+                pipes = list(task_pipes)
+                for p in pipes:
+                    try:
+                        p.close()
+                        pipe_count += 1
+                    except Exception as e:
+                        logging.exception("Error closing pipe on task list", stack_info=True)
+            logging.info(f"Closed pipes - {pipe_count}")
path: runner/app/live/trickle/media.py
user: login: j0sh
body: added in https://github.com/livepeer/ai-runner/pull/454/commits/36af17a2e27b888d2782ad3fb909376aba90f0a3

fwiw "failing to close" is not necessarily indicative of problem in itself; it may just have crashed after ffmpeg closed but before we had a chance to clean it up. If this happens in practice then it might offer clues about _where_ the crash is happening.
created_at: 2025-03-17T21:29:34Z
author_association: CONTRIBUTOR

PR comments:
user: login: emranemran
created_at: 2025-03-17T05:19:03Z
body: How did you test the retry logic? Were you able to trigger the same types of crash we saw on staging/prod and confirm? 
author_association: COLLABORATOR
user: login: leszko
created_at: 2025-03-17T10:51:32Z
body: Unassigning myself from reviewers, I think it's better if it's reviewed by @emranemran and @victorges 
author_association: CONTRIBUTOR
user: login: j0sh
created_at: 2025-03-17T21:31:36Z
body: > How did you test the retry logic?

by explicitly raising exceptions

> Were you able to trigger the same types of crash we saw on staging/prod and confirm?

not really - I can trigger other types of crashes with different types of input but not anything like we're seeing.
author_association: CONTRIBUTOR