fix(gc): check for ruuns with running state as well

ZettaAI · Nov 1, 2024 · b2e8259 · b2e8259
1 parent b375bac
commit b2e8259
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/zetta_utils/run/gc.py b/zetta_utils/run/gc.py
@@ -44,6 +44,8 @@ def _get_current_resources_and_stale_run_ids() -> (
         run_resources[str(_resource["run_id"])][_resource_id] = _resource
 
     run_ids = list(run_resources.keys())
+    runs = RUN_DB.query(column_filter={"state": ["running"]})
+    run_ids += list(runs.keys())
     stale_ids = []
     heartbeats = [x.get("heartbeat", 0) for x in RUN_DB[(run_ids, ("heartbeat",))]]
     lookback = int(os.environ["EXECUTION_HEARTBEAT_LOOKBACK"])