From 2160e8bbefc3b916c43a327d7c9fa84675ea0167 Mon Sep 17 00:00:00 2001
From: naegelyd-insitro <66270360+naegelyd-insitro@users.noreply.github.com>
Date: Tue, 5 Mar 2024 14:18:22 -0500
Subject: [PATCH] Scheduler limits bugfix (#92)

* Add regression test for keyerror in job limits util

* Don't assume all jobs share the same keys in their limits
---
 redun/scheduler.py         | 14 +++++++-------
 redun/tests/test_limits.py |  8 +++++++-
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/redun/scheduler.py b/redun/scheduler.py
index ed99cc2d..8d13e970 100644
--- a/redun/scheduler.py
+++ b/redun/scheduler.py
@@ -1311,17 +1311,17 @@ def _add_job_pending_limits(self, job: Job, eval_args: Tuple[Tuple, dict]) -> No
         """
         self._jobs_pending_limits.append((job, eval_args))
 
+    def _add_limits(self, limits1, limits2):
+        return {
+            limit_name: limits1.get(limit_name, 0) + limits2.get(limit_name, 0)
+            for limit_name in set(limits1) | set(limits2)
+        }
+
     def _check_jobs_pending_limits(self) -> None:
         """
         Checks whether Jobs pending due to limits can now satisfy limits and run.
         """
 
-        def _add_limits(limits1, limits2):
-            return {
-                limit_name: limits1[limit_name] + limits2[limit_name]
-                for limit_name in set(limits1) | set(limits2)
-            }
-
         ready_jobs: List[Tuple[Job, Tuple[tuple, dict]]] = []
         not_ready_jobs: List[Tuple[Job, Tuple[tuple, dict]]] = []
         ready_job_limits: Dict[str, int] = defaultdict(int)
@@ -1329,7 +1329,7 @@ def _add_limits(limits1, limits2):
         # Resource limits that will be consumed by ready jobs.
         for job, eval_args in self._jobs_pending_limits:
             job_limits = job.get_limits()
-            proposed_limits = _add_limits(job_limits, ready_job_limits)
+            proposed_limits = self._add_limits(job_limits, ready_job_limits)
             if self._is_job_within_limits(proposed_limits):
                 ready_jobs.append((job, eval_args))
                 ready_job_limits = proposed_limits
diff --git a/redun/tests/test_limits.py b/redun/tests/test_limits.py
index 777fc6c4..5f047b3f 100644
--- a/redun/tests/test_limits.py
+++ b/redun/tests/test_limits.py
@@ -30,7 +30,13 @@ def test_scheduler_limits_utils() -> None:
     """
     scheduler = Scheduler(config=Config(CONFIG_DICT))
 
-    job_limits = {"api": 1, "memory": MEMORY_LIMIT}
+    job_limits1 = {"api": 1, "memory": MEMORY_LIMIT}
+    job_limits2 = {"api": 1}
+
+    # Confirm we can add limits together where not all keys are present in both jobs. Regression
+    # test added for:
+    #   https://github.com/insitro/redun/issues/91
+    job_limits = scheduler._add_limits(job_limits1, job_limits2)
 
     # Check that a job that should have enough resources to run is correctly deemed within limits.
     assert scheduler._is_job_within_limits(job_limits)