Merge pull request #37 from CAAI/dev/v1.3

fix queue ordering bug, changed autodeletion to be performed on jobs that have been finished for at least 10 minutes

Author: ChristianHinge

nodes/manager/manager.py

@@ -10,7 +10,7 @@
 from dotenv import load_dotenv
 from rhnode.version import __version__
 import time
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, validator, root_validator
 
 # load env variables from .env file if it exists
 load_dotenv()
@@ -22,9 +22,18 @@ class Job(BaseModel):
     required_gpu_mem: int
     required_memory: int
     required_threads: int
-    creation_time: float = time.time()
+    creation_time: float = None
     gpu_device_id: int = None
 
+    class Config:
+        validate_assignment = True
+
+    @root_validator
+    def number_validator(cls, values):
+        if values["creation_time"] is None:
+            values["creation_time"] = time.time()
+        return values
+
     @validator("priority")
     def check_priority(cls, v):
         if v < 1 or v > 5:
@@ -39,8 +48,8 @@ def __init__(self, job):
     def __repr__(self):
         return f"Queue item: {self.job}"
 
-    def __lt__(self, other: Job):
-        self_args = (-self.job.priority, other.job.creation_time)
+    def __lt__(self, other):
+        self_args = (-self.job.priority, self.job.creation_time)
         other_args = (-other.job.priority, other.job.creation_time)
 
         for self_arg, other_arg in zip(self_args, other_args):

rhnode/rhnode.py

@@ -76,24 +76,35 @@ def __init__(self):
         setup_frontend_routes(self)
 
     def _delete_job(self, job_id):
+        print("DELETING", job_id)
         job = self.jobs[job_id]
         job.delete_files()
         del self.jobs[job_id]
 
-    def _get_expired_job_ids(self, max_age_hours=8):
-        """Get a list of job IDs that have been in the queue for longer than max_age_hours"""
+    def _get_expired_job_ids(self, max_finished_mins=10):
+        """Get a list of job IDs that have been finished for longer than max_finished_mins"""
         expired_job_ids = []
+        current_time = time.time()
         for job_id, job in self.jobs.items():
-            if (time.time() - job.time_created) / 3600 > max_age_hours:
-                if job.status in [
-                    JobStatus.Finished,
-                    JobStatus.Error,
-                    JobStatus.Cancelled,
-                ]:
-                    expired_job_ids.append(job_id)
+            if job.time_finished is not None:
+                if (current_time - job.time_finished) / 60 > max_finished_mins:
+                    if job.status in [
+                        JobStatus.Finished,
+                        JobStatus.Error,
+                        JobStatus.Cancelled,
+                    ]:
+                        expired_job_ids.append(job_id)
 
         return expired_job_ids
 
+    def _run_cleanup(self):
+        print("Checking for expired jobs...")
+        jobs = self._get_expired_job_ids()
+
+        for job_id in jobs:
+            print("Deleting job", job_id)
+            self._delete_job(job_id)
+
     async def _delete_expired_jobs_loop(
         self, hour: int = 3, minute: int = 30, second: int = 0
     ):
@@ -113,11 +124,7 @@ async def _delete_expired_jobs_loop(
             delay = (next_run - now).total_seconds()
             print("Next check in", delay, "seconds")
             await asyncio.sleep(delay)
-            print("Checking for expired jobs...")
-            jobs = self._get_expired_job_ids()
-            for job_id in jobs:
-                print("Deleting job", job_id)
-                self._delete_job(job_id)
+            self._run_cleanup()
 
     def get_job_by_id(self, job_id: str):
         try:

rhnode/routing/backend.py

@@ -116,6 +116,14 @@ async def _parse_cli_args(cli: list):
     async def _get_job_input():
         return rhnode.help_cli_args()
 
+    @rhnode.get(rhnode._create_url("/clean_jobs"))
+    async def _clean_jobs():
+        if rhnode.rhnode_mode == "dev":
+            rhnode._run_cleanup()
+            return "ok"
+        else:
+            return "not in dev mode " + str(rhnode.rhnode_mode)
+
     @rhnode.post(rhnode._create_url("/jobs/{job_id}/upload"))
     async def _upload(
         job_id: str,

tests/atest_multijob.py

@@ -0,0 +1,37 @@
+# %%
+
+## Not to be run via pytest as of now
+import glob
+import os
+
+from rhnode import RHJob, MultiJobRunner
+import os
+import json
+import datetime
+import nibabel as nib
+from pathlib import Path
+import re
+
+ROOT = "/depict/data/quadra_fdg"
+
+jobs = []
+
+for i in range(500):
+    inputs = {
+        "scalar": 2,
+        "in_file": "/homes/hinge/Projects/rh-node/tests/data/mr.nii.gz",
+        "out_file": "out.nii.gz",
+    }
+    job = RHJob(
+        node_name="add",
+        manager_address="localhost:9050",
+        inputs=inputs,
+        check_cache=False,
+    )
+
+    jobs.append(job)
+
+job_runner = MultiJobRunner(jobs, queue_length=8)
+job_runner.start()
+
+# %%

tests/docker-compose.yaml

@@ -59,6 +59,7 @@ services:
       - "traefik.http.routers.add.rule=PathPrefix(`/add`)"
     environment:
       RH_EMAIL_ON_ERROR: [email protected]
+      RH_MODE: dev
       TZ: "Europe/Copenhagen"
 
   ## Testnode: OutputDirectory