minor refactoring

Author: haraldsvik

job_executor/worker/manager_state.py

@@ -1,3 +1,8 @@
+import logging
+
+logger = logging.getLogger()
+
+
 class ManagerState:
     def __init__(self, default_max_workers=4, max_gb_all_workers=50):
         """
@@ -27,18 +32,32 @@ def can_spawn_new_worker(self, new_job_size):
         When a job is finished and unregister the number of workers will reset
         to default_max_workers
         """
-
-        # Could tweak this to be more gradual if we want
-        if self.current_total_size >= self.max_bytes_all_workers:
-            self.current_max_workers = 2
-        else:
-            self.current_max_workers = self.default_max_workers
+        can_spawn = True
+        self.update_worker_limit(new_job_size)
 
         active_workers = len(self.datasets)
         if active_workers >= self.current_max_workers:
-            return False
+            can_spawn = False
 
-        return True
+        logger.info(
+            f"Checking can_spawn_new_worker({new_job_size}): "
+            f"active={active_workers}, dynamic_limit={self.current_max_workers}, "
+            f"current_total_size={self.current_total_size}, can_spawn={can_spawn}"
+        )
+
+        return can_spawn
+
+    def update_worker_limit(self, new_job_size):
+        """
+        Check the current size beeing procces in the pipeline.
+        And changes the number of workers as needed.
+        """
+        if (
+            self.current_total_size + new_job_size
+        ) >= self.max_bytes_all_workers:
+            self.current_max_workers = 2
+        else:
+            self.current_max_workers = self.default_max_workers
 
     def register_job(self, job_id, job_size):
         """

tests/unit/worker/test_manager_state.py

@@ -15,7 +15,7 @@ def test_can_spawn_worker():
     assert can_spawn is True
 
 
-def test_can_not_spawn_worker_to_many_workers():
+def test_cannot_spawn_worker_too_many_workers():
     manager_state = ManagerState(default_max_workers=4)
 
     # Register 4 jobs
@@ -26,20 +26,39 @@ def test_can_not_spawn_worker_to_many_workers():
     assert can_spawn is False
 
 
-def test_can_not_spawn_worker_size_limit_reached():
+def test_cannot_spawn_worker_size_limit_reached():
+    TWENTY_GB = 20 * 1024**3
     manager_state = ManagerState(max_gb_all_workers=20)
 
     # register large job
-    manager_state.register_job("job_1", 20 * 1024**3)  # 20 GB
+    manager_state.register_job("job_1", TWENTY_GB)  # 20 GB
 
     # Max worker should now be 2, We can still spawn a second job
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=10 * 1024**3)
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TWENTY_GB)
     assert can_spawn is True
 
-    manager_state.register_job("job_2", 10 * 1024**3)
+    manager_state.register_job("job_2", TWENTY_GB)
 
     # We should not be able to spawn a third job
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=10 * 1024**3)
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TWENTY_GB)
+    assert can_spawn is False
+
+
+def test_oversized_jobs():
+    # we can run jobs which exceeds the limit for the threshold,
+    # but only two of them
+    FIFTY_GB = 50 * 1024**3
+    manager_state = ManagerState(max_gb_all_workers=10)
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
+    manager_state.register_job("job_1", FIFTY_GB)
+    assert can_spawn is True
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
+    manager_state.register_job("job_2", FIFTY_GB)
+    assert can_spawn is True
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
     assert can_spawn is False