Simplify logic, and updated tests

Author: haraldsvik

job_executor/worker/manager_state.py

@@ -3,9 +3,6 @@
 from typing import List
 
 
-REDUCED_WORKER_NUMBER = 2
-
-
 class ManagerState:
     def __init__(self, default_max_workers=4, max_gb_all_workers=50):
         """
@@ -14,7 +11,6 @@ def __init__(self, default_max_workers=4, max_gb_all_workers=50):
         of workers are reduced
         """
         self.default_max_workers = default_max_workers
-        self.current_max_workers = default_max_workers
         self.max_bytes_all_workers = (
             max_gb_all_workers * 1024**3  # Threshold in bytes
         )
@@ -45,30 +41,20 @@ def can_spawn_new_worker(self, new_job_size: int) -> bool:
         """
         Called to check if a new worker can be spawned.
         """
-        if len(self.alive_workers) >= self.current_max_workers:
+        if len(self.alive_workers) >= self.default_max_workers:
             return False
-        if self.current_total_size + new_job_size >= self.max_bytes_all_workers:
+        if (
+            self.current_total_size + new_job_size
+            >= self.max_bytes_all_workers
+        ):
             return False
         return True
 
-    def update_worker_limit(self, new_job_size: int):
-        """
-        Check the current size beeing procces in the pipeline.
-        And changes the number of workers as needed.
-        """
-        new_total = self.current_total_size + new_job_size
-        if new_total >= self.max_bytes_all_workers:
-            self.current_max_workers = REDUCED_WORKER_NUMBER
-        else:
-            self.current_max_workers = self.default_max_workers
-
     def register_job(self, worker: Worker):
         """
         Called when a worker picks up a job.
-        When a job is register the current_max_workers are updated.
         """
         self.workers.append(worker)
-        self.update_worker_limit(worker.job_size)
 
     def unregister_job(self, job_id):
         """

tests/unit/worker/test_manager_state.py

@@ -44,7 +44,6 @@ def test_cannot_spawn_worker_size_limit_reached():
     TWENTY_GB = 20 * 1024**3
     manager_state = ManagerState(max_gb_all_workers=20)
 
-    # register large job
     large_job = Worker(
         process=Process(target=dummy),
         job_id="job_large",
@@ -53,55 +52,31 @@ def test_cannot_spawn_worker_size_limit_reached():
     manager_state.register_job(large_job)
     large_job.start()
 
-    # Max worker should now be 2, We can still spawn a second job
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TWENTY_GB)
-    assert can_spawn is True
-
-    large_job = Worker(
-        process=Process(target=dummy),
-        job_id="job_large_2",
-        job_size=TWENTY_GB,
-    )
-    manager_state.register_job(large_job)
-    large_job.start()
-
-    # We should not be able to spawn a third job
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TWENTY_GB)
+    # Only one job active but size limit is reached cannot spawn new job
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=1024)
     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
     TEN_GB = 10 * 1024**3
-    manager_state = ManagerState(max_gb_all_workers=10)
+    manager_state = ManagerState(max_gb_all_workers=20)
 
+    # This job will never be processed
     can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
-    large_job = Worker(
-        process=Process(target=dummy),
-        job_id="job_1",
-        job_size=FIFTY_GB,
-    )
-    manager_state.register_job(large_job)
-    large_job.start()
-    assert can_spawn is True
+    assert can_spawn is False
 
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
-    large_job = Worker(
-        process=Process(target=dummy),
-        job_id="job_2",
-        job_size=FIFTY_GB,
-    )
-    manager_state.register_job(large_job)
-    large_job.start()
+    # This job will be accepted
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TEN_GB)
     assert can_spawn is True
-
-    assert manager_state.max_bytes_all_workers == TEN_GB
-    assert manager_state.current_total_size == (FIFTY_GB + FIFTY_GB)
-
-    can_spawn = manager_state.can_spawn_new_worker(new_job_size=FIFTY_GB)
-    assert can_spawn is False
+    if can_spawn:
+        worker = Worker(
+            process=Process(target=dummy),
+            job_id="job_2",
+            job_size=TEN_GB,
+        )
+        manager_state.register_job(worker)
+        worker.start()
 
 
 def test_unregister_job():