Merge pull request #144 from statisticsnorway/dynamic-workers

Dynamic workers

Author: haraldsvik

.test.env

@@ -7,4 +7,5 @@ JOB_SERVICE_URL=http://mock.job.service
 NUMBER_OF_WORKERS=4
 SECRETS_FILE=tests/resources/secrets/secrets.json
 DOCKER_HOST_NAME=localhost
-COMMIT_ID=abc123
+COMMIT_ID=abc123
+MAX_GB_ALL_WORKERS=50

job_executor/adapter/local_storage.py

@@ -423,3 +423,14 @@ def delete_archived_input(dataset_name: str):
     archived_file: Path = INPUT_DIR / f"archive/{dataset_name}.tar"
     if archived_file.is_file():
         os.remove(archived_file)
+
+
+def get_input_tar_size_in_bytes(dataset_name: str) -> int:
+    """
+    Checks the size in bytes of the dataset.tar file.
+    returns size in bytes
+    """
+    tar_path = INPUT_DIR / f"{dataset_name}.tar"
+    if not tar_path.exists():
+        return 0
+    return os.path.getsize(tar_path)

job_executor/app.py

@@ -6,19 +6,28 @@
 from pathlib import Path
 from typing import Dict, List
 
-from job_executor.adapter import job_service
+from job_executor.adapter import job_service, local_storage
 from job_executor.config import environment
 from job_executor.config.log import setup_logging, initialize_logging_thread
 from job_executor.domain import rollback
-from job_executor.exception import RollbackException, StartupException
+from job_executor.exception import (
+    RollbackException,
+    StartupException,
+)
 from job_executor.model import Job, Datastore
 from job_executor.model.worker import Worker
-from job_executor.worker import build_dataset_worker, build_metadata_worker
+from job_executor.worker import (
+    build_dataset_worker,
+    build_metadata_worker,
+    manager_state as ManagerState,
+)
+
 
 logger = logging.getLogger()
 setup_logging()
 
 NUMBER_OF_WORKERS = int(environment.get("NUMBER_OF_WORKERS"))
+MAX_GB_ALL_WORKERS = int(environment.get("MAX_GB_ALL_WORKERS"))
 DATASTORE_DIR = environment.get("DATASTORE_DIR")
 
 datastore = None
@@ -199,7 +208,11 @@ def initialize_app():
 def main():
     initialize_app()
     logging_queue, log_thread = initialize_logging_thread()
-    workers: List[Worker] = []
+
+    manager_state = ManagerState(
+        default_max_workers=NUMBER_OF_WORKERS,
+        max_gb_all_workers=MAX_GB_ALL_WORKERS,
+    )
 
     try:
         while True:
@@ -210,12 +223,7 @@ def main():
             built_jobs = job_dict["built_jobs"]
             queued_manager_jobs = job_dict["queued_manager_jobs"]
 
-            dead_workers = [
-                worker for worker in workers if not worker.is_alive()
-            ]
-            clean_up_after_dead_workers(dead_workers)
-
-            workers = [worker for worker in workers if worker.is_alive()]
+            clean_up_after_dead_workers(manager_state)
 
             available_jobs = (
                 len(queued_worker_jobs)
@@ -229,11 +237,28 @@ def main():
                     f" (worker, built, queued manager jobs)"
                 )
             for job in queued_worker_jobs:
-                if len(workers) < NUMBER_OF_WORKERS:
-                    _handle_worker_job(job, workers, logging_queue)
+                job_size = local_storage.get_input_tar_size_in_bytes(
+                    job.dataset_name
+                )
+                if job_size == 0:
+                    logger.info(
+                        f"{job.job_id} Failed to get the size of the dataset."
+                    )
+                    job_service.update_job_status(
+                        job.job_id,
+                        "failed",
+                        log="No such dataset available for import",
+                    )
+                    continue  # skip futher processing of this job
+
+                if manager_state.can_spawn_new_worker(job_size):
+                    _handle_worker_job(
+                        job, manager_state, job_size, logging_queue
+                    )
 
             for job in built_jobs + queued_manager_jobs:
                 try:
+                    manager_state.unregister_job(job.job_id)
                     _handle_manager_job(job)
                 except Exception as exc:
                     # All exceptions that occur during the handling of a job
@@ -253,7 +278,8 @@ def main():
         log_thread.join()
 
 
-def clean_up_after_dead_workers(dead_workers: List[Worker]) -> None:
+def clean_up_after_dead_workers(manager_state) -> None:
+    dead_workers = manager_state.dead_workers
     if len(dead_workers) > 0:
         in_progress_jobs = job_service.get_jobs(ignore_completed=True)
         for dead_worker in dead_workers:
@@ -268,9 +294,12 @@ def clean_up_after_dead_workers(dead_workers: List[Worker]) -> None:
             if job and job.status not in ["queued", "built"]:
                 logger.info(f"Worker died and did not finish job {job.job_id}")
                 fix_interrupted_job(job)
+            manager_state.unregister_job(dead_worker.job_id)
 
 
-def _handle_worker_job(job: Job, workers: List[Worker], logging_queue: Queue):
+def _handle_worker_job(
+    job: Job, manager_state: ManagerState, job_size: int, logging_queue: Queue
+):
     dataset_name = job.parameters.target
     job_id = job.job_id
     operation = job.parameters.operation
@@ -285,8 +314,9 @@ def _handle_worker_job(job: Job, workers: List[Worker], logging_queue: Queue):
                 ),
             ),
             job_id=job_id,
+            job_size=job_size,
         )
-        workers.append(worker)
+        manager_state.register_job(worker, job_id, job_size)
         job_service.update_job_status(job_id, "initiated")
         worker.start()
     elif operation == "PATCH_METADATA":
@@ -301,7 +331,7 @@ def _handle_worker_job(job: Job, workers: List[Worker], logging_queue: Queue):
             ),
             job_id=job_id,
         )
-        workers.append(worker)
+        manager_state.register_job(worker)
         job_service.update_job_status(job_id, "initiated")
         worker.start()
     else:

job_executor/config/environment.py

@@ -13,6 +13,7 @@ def _initialize_environment() -> dict:
         "SECRETS_FILE": os.environ["SECRETS_FILE"],
         "DOCKER_HOST_NAME": os.environ["DOCKER_HOST_NAME"],
         "COMMIT_ID": os.environ["COMMIT_ID"],
+        "MAX_GB_ALL_WORKERS": os.environ["MAX_GB_ALL_WORKERS"],
     }
 
 

job_executor/model/worker.py

@@ -2,9 +2,13 @@
 
 
 class Worker:
-    def __init__(self, process: Process, job_id: str):
+    def __init__(self, process: Process, job_id: str, job_size: int):
         self.process = process
         self.job_id = job_id
+        self.job_size = job_size
+
+    def job_size(self) -> int:
+        return self.job_size
 
     def is_alive(self):
         return self.process.is_alive()

job_executor/worker/manager_state.py

@@ -0,0 +1,65 @@
+from job_executor.model.worker import Worker
+
+from typing import List
+
+
+class ManagerState:
+    def __init__(self, default_max_workers=4, max_gb_all_workers=50):
+        """
+        :param default_max_workers: The maximum number of workers
+        :param max_gb_all_workers: Threshold in GB (50) for when the number
+        of workers are reduced
+        """
+        self.default_max_workers = default_max_workers
+        self.max_bytes_all_workers = (
+            max_gb_all_workers * 1024**3  # Threshold in bytes
+        )
+
+        self.workers: List[Worker] = []
+
+    @property
+    def dead_workers(self) -> List[Worker]:
+        """
+        Return a list of dead workers
+        """
+        return [worker for worker in self.workers if not worker.is_alive()]
+
+    @property
+    def alive_workers(self) -> List[Worker]:
+        """
+        Return a list of alive workers
+        """
+        return [worker for worker in self.workers if worker.is_alive()]
+
+    @property
+    def current_total_size(self) -> int:
+        return sum(
+            worker.job_size for worker in self.workers if worker.is_alive()
+        )
+
+    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.default_max_workers:
+            return False
+        if (
+            self.current_total_size + new_job_size
+            >= self.max_bytes_all_workers
+        ):
+            return False
+        return True
+
+    def register_job(self, worker: Worker):
+        """
+        Called when a worker picks up a job.
+        """
+        self.workers.append(worker)
+
+    def unregister_job(self, job_id):
+        """
+        Called when a job finishes or fails.
+        """
+        self.workers = [
+            worker for worker in self.workers if worker.job_id != job_id
+        ]

tests/unit/worker/test_manager_state.py

@@ -0,0 +1,100 @@
+import time
+from job_executor.worker.manager_state import ManagerState
+from job_executor.model.worker import Worker
+from multiprocessing import Process
+
+
+def dummy():
+    time.sleep(10)
+    print("hello")
+
+
+def test_initial_state():
+    manager_state = ManagerState()
+
+    assert manager_state.current_total_size == 0
+    assert len(manager_state.workers) == 0
+
+
+def test_can_spawn_worker():
+    manager_state = ManagerState()
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=1)
+    assert can_spawn is True
+
+
+def test_cannot_spawn_worker_too_many_workers():
+    manager_state = ManagerState(default_max_workers=4)
+
+    # Register 4 jobs
+    for i in range(4):
+        worker = Worker(
+            process=Process(target=dummy),
+            job_id=f"job_{i}",
+            job_size=1024,
+        )
+        manager_state.register_job(worker)
+        worker.start()
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=1024)
+    assert can_spawn is False
+
+
+def test_cannot_spawn_worker_size_limit_reached():
+    TWENTY_GB = 20 * 1024**3
+    manager_state = ManagerState(max_gb_all_workers=20)
+
+    large_job = Worker(
+        process=Process(target=dummy),
+        job_id="job_large",
+        job_size=TWENTY_GB,
+    )
+    manager_state.register_job(large_job)
+    large_job.start()
+
+    # 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():
+    FIFTY_GB = 50 * 1024**3
+    TEN_GB = 10 * 1024**3
+    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)
+    assert can_spawn is False
+
+    # This job will be accepted
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=TEN_GB)
+    assert can_spawn is True
+    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():
+    manager_state = ManagerState(default_max_workers=4)
+
+    # Register 4 jobs
+    for i in range(4):
+        worker = Worker(
+            process=Process(target=dummy),
+            job_id=f"job_{i}",
+            job_size=1024,
+        )
+        manager_state.register_job(worker)
+        worker.start()
+
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=1024)
+    assert can_spawn is False
+
+    manager_state.unregister_job("job_1")
+    can_spawn = manager_state.can_spawn_new_worker(new_job_size=1024)
+    assert can_spawn is True