From 004313cf09f1d826759933e4116b4a8138abc417 Mon Sep 17 00:00:00 2001 From: Tim Blakely Date: Fri, 26 Sep 2025 13:53:28 -0700 Subject: [PATCH] Fix a weird bug where using FFN inference interactively in Colab caused a kernel crash (!) when subsequent inference calls are made. PiperOrigin-RevId: 811944130 --- ffn/inference/executor.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/ffn/inference/executor.py b/ffn/inference/executor.py index c0461c1..df78185 100644 --- a/ffn/inference/executor.py +++ b/ffn/inference/executor.py @@ -294,9 +294,14 @@ def _run_executor(self): self.active_clients += 1 logging.info('client %d starting', client_id) else: - self.registered_clients.remove(-client_id - 1) - logging.info('client %d terminating', -client_id - 1) - self.active_clients -= 1 + try: + self.registered_clients.remove(-client_id - 1) + logging.info('client %d terminating', -client_id - 1) + self.active_clients -= 1 + except KeyError: + logging.warning( + 'client %d not known or already terminated', -client_id - 1 + ) else: client_id, seed, image, fetches = data l = len(ready)