Merge pull request #29 from CAAI/dev/cleanup_rhprocess

small cleanup. Manager Load estimation. The manager will check which server has the fewest queued jobs and assign the job to that server

Author: ChristianHinge

nodes/manager/manager.py

@@ -154,6 +154,46 @@ def get_active_jobs_info(self):
             )
         return active_jobs_info
 
+    def get_queued_jobs_info(self):
+        return {
+            "queued_jobs": [
+                {
+                    "priority": job[0] * -1,
+                    "job_id": job[1],
+                    "required_gpu_mem": job[2],
+                    "required_threads": job[3],
+                    "required_memory": job[4],
+                }
+                for job in self.job_queue
+            ]
+        }
+
+    def get_load(self):
+        active_jobs_info = self.get_active_jobs_info()
+        queued_jobs_info = self.get_queued_jobs_info()["queued_jobs"]
+        sum_gpu_mem = (
+            sum(self.gpu_devices_mem_available)
+            if isinstance(self.gpu_devices_mem_available, list)
+            else self.gpu_devices_mem_available
+        )
+        tot_required_gpu_mem = sum(
+            x["required_gpu_mem"] for x in active_jobs_info + queued_jobs_info
+        )
+        tot_required_mem = sum(
+            x["required_memory"] for x in active_jobs_info + queued_jobs_info
+        )
+        tot_required_threads = sum(
+            x["required_threads"] for x in active_jobs_info + queued_jobs_info
+        )
+
+        load = max(
+            tot_required_gpu_mem / sum_gpu_mem,
+            tot_required_mem / self.memory_max,
+            tot_required_threads / self.threads_max,
+        )
+
+        return load
+
 
 templates = Jinja2Templates(
     directory=os.path.dirname(__file__) + "/resources/templates"
@@ -194,15 +234,28 @@ def get_addr_to_run_node(self, node_name):
         if node_name in self.nodes.keys():
             return "localhost:8000"
 
+        lowest_load = None
+        addr_with_lowest_load = None
         # query other hosts for available gpu
         for addr in self.other_addrs:
             url = f"http://{addr}/manager/dispatcher/has_node/{node_name}"
             response = requests.get(url)
             response.raise_for_status()
             has_node = response.json()
+            assert isinstance(has_node, bool)
             if has_node:
-                return addr
-        raise Exception("No servers were found with the node")
+                url = f"http://{addr}/manager/get_load"
+                response = requests.get(url)
+                response.raise_for_status()
+                node_load = response.json()
+                if lowest_load == None or node_load < lowest_load:
+                    lowest_load = node_load
+                    addr_with_lowest_load = addr
+
+        if addr_with_lowest_load is not None:
+            return addr_with_lowest_load
+        else:
+            raise Exception("No servers were found with the node")
 
     def setup_routes(self):
         @self.post("/manager/register_node")
@@ -251,18 +304,11 @@ async def get_active_jobs():
 
         @self.get("/manager/get_queued_jobs")
         async def get_queued_jobs():
-            return {
-                "queued_jobs": [
-                    {
-                        "priority": job[0] * -1,
-                        "job_id": job[1],
-                        "required_gpu_mem": job[2],
-                        "required_threads": job[3],
-                        "required_memory": job[4],
-                    }
-                    for job in self.queue.job_queue
-                ]
-            }
+            self.queue.get_queued_jobs_info()
+
+        @self.get("/manager/get_load")
+        async def get_load():
+            return self.queue.get_load()
 
         @self.get("/manager/get_resource_info")
         async def get_resource_info():

rhnode/common.py

@@ -134,3 +134,22 @@ def validate_input_output_spec(input_spec, output_spec):
         raise ValueError(
             "Input spec and output spec must not have overlapping key names"
         )
+
+
+def convert_string_to_type(text, type_):
+    """Convert a string to a type. This is used to parse the inputs to the process function."""
+    if type_ == str or type_ == FilePath:
+        return text
+    elif type_ == bool:
+        return text_to_bool(text)
+    else:
+        return type_(text)
+
+
+def text_to_bool(value):
+    if value.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif value.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise ValueError(f"Could not convert {value} to bool")

rhnode/rhjob.py

@@ -8,7 +8,6 @@
 from requests.exceptions import HTTPError
 
 
-
 class RHJob:
     def __init__(
         self,
@@ -269,7 +268,9 @@ def wait_for_finish(self):
         if self.strict_output_dir:
             output_path = self.output_directory
         else:
-            output_path = _create_output_directory_name(self.output_directory, self.node_identifier)
+            output_path = _create_output_directory_name(
+                self.output_directory, self.node_identifier
+            )
 
         output = None
         while output is None:

rhnode/rhnode.py

@@ -10,7 +10,8 @@
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi import FastAPI, File, Form, UploadFile, BackgroundTasks
 from .rhprocess import RHProcess
-from .frontend import setup_frontend_routes
+from .routing.frontend import setup_frontend_routes
+from .routing.backend import setup_api_routes
 import traceback
 from fastapi import Response
 from fastapi import HTTPException
@@ -70,11 +71,10 @@ def __init__(self):
         else:
             self.email_sender = None
 
-        self.setup_api_routes()
+        setup_api_routes(self)
         setup_frontend_routes(self)
 
     def _delete_job(self, job_id):
-        pass
         job = self.jobs[job_id]
         job.delete_files()
         del self.jobs[job_id]
@@ -207,7 +207,9 @@ def _get_output_with_download_links(self, job_id):
                 else val
                 for key, val in job.output.dict(exclude_unset=True).items()
                 # Only return links for non FilePath and FilePath with values that are not None unless required
-                if not self.output_spec.__fields__[key].type_ == FilePath or val is not None or (val is None and self.output_spec.__fields__[key].required)
+                if not self.output_spec.__fields__[key].type_ == FilePath
+                or val is not None
+                or (val is None and self.output_spec.__fields__[key].required)
             }
         )
 
@@ -225,153 +227,6 @@ def _ensure_job_status(self, status, valid_statuses):
                 ),
             )
 
-    def setup_api_routes(self):
-        """Setup the API routes for the node. See frontend.py for the frontend routes."""
-
-        @self.post(self._create_url("/jobs"))
-        async def _post_new_job(inputs: self.input_spec_no_file) -> str:
-            job_id = self.CREATE_JOB(inputs)
-            return job_id
-
-        @self.post(self._create_url("/jobs/{job_id}/start"))
-        async def START_JOB(
-            job_id: str, job_meta_data: JobMetaData, background_tasks: BackgroundTasks
-        ) -> Response:
-            job_obj = self.get_job_by_id(job_id)
-            self._ensure_job_status(job_obj.status, JobStatus.Preparing)
-            if not job_obj.is_ready_to_run():
-                raise HTTPException(
-                    status_code=400,
-                    detail="Job is not ready to run. Some files are likely missing.",
-                )
-            background_tasks.add_task(job_obj.run, job_meta_data)
-            return Response(status_code=204)
-
-        @self.get(self._create_url("/jobs/{job_id}/status"))
-        async def _get_job_status(job_id: str) -> JobStatus:
-            return self.get_job_by_id(job_id).status
-
-        @self.get(self._create_url("/jobs/{job_id}/data"))
-        async def _get_job_data_download_urls(job_id: str) -> self.output_spec_url:
-            return self._get_output_with_download_links(job_id)
-
-        @self.get(self._create_url("/jobs/{job_id}/error"))
-        def _get_job_error(job_id: str):
-            job = self.get_job_by_id(job_id)
-            self._ensure_job_status(job.status, [JobStatus.Cancelled, JobStatus.Error])
-            return job.error
-
-        @self.post(self._create_url("/jobs/{job_id}/stop"))
-        def _stop_task(job_id: str):
-            job = self.get_job_by_id(job_id)
-            if job.status not in [
-                JobStatus.Finished,
-                JobStatus.Cancelled,
-                JobStatus.Cancelling,
-                JobStatus.Error,
-            ]:
-                job.stop()
-                return Response(status_code=204)
-
-            return Response(status_code=200, content="Job is already in terminal phase")
-
-        @self.post(self._create_url("/jobs/{job_id}/delete"))
-        async def _delete_job(job_id: str):
-            """Delete a job from the node."""
-            job = self.get_job_by_id(job_id)
-            self._delete_job(job_id)
-
-        @self.get(self._create_url("/jobs/{job_id}/download/{filename}"))
-        def _get_file(job_id, filename):
-            job = self.get_job_by_id(job_id)
-            self._ensure_job_status(job.status, JobStatus.Finished)
-            if not filename in self.output_spec.__fields__:
-                raise HTTPException(
-                    status_code=404,
-                    detail="The requested file key {} is invalid.".format(filename),
-                )
-            try:
-                fname = job.output.dict()[filename]
-            except KeyError:
-                raise HTTPException(
-                    status_code=404,
-                    detail="The requested file corresponding to key {} could not be found.".format(
-                        filename
-                    ),
-                )           
-            return FileResponse(
-                fname, filename=create_file_name_from_key(filename, fname)
-            )
-
-        @self.get(self._create_url("/filename_keys"))
-        async def _get_file_keys():
-            return self.input_file_keys
-
-        @self.get(self._create_url("/keys"))
-        async def _get_keys():
-            return {
-                "output_keys": list(self.output_spec.__fields__.keys()),
-                "input_keys": list(self.input_spec.__fields__.keys()),
-            }
-
-        @self.post(self._create_url("/cli/parse"))
-        async def _parse_cli_args(cli: list):
-            try:
-                return self.parse_cli_args(cli)
-            except Exception as e:
-                raise HTTPException(status_code=400, detail=str(e))
-
-        @self.get(self._create_url("/cli/help"))
-        async def _get_job_input():
-            return self.help_cli_args()
-
-        @self.post(self._create_url("/jobs/{job_id}/upload"))
-        async def _upload(
-            job_id: str,
-            file: UploadFile = File(...),
-            key: str = Form(...),
-        ):
-            job = self.get_job_by_id(job_id)
-            self._ensure_job_status(job.status, JobStatus.Preparing)
-
-            """Upload an input file to a job. The key must be one of the file keys."""
-
-            if not key in self.input_file_keys:
-                raise HTTPException(
-                    status_code=404,
-                    detail="The requested file key {} is invalid.".format(key),
-                )
-            # The outer with statement is to ensure that the file is validated after upload
-            with job.upload_file(key, file.filename) as fpath:
-                with open(fpath, "wb") as f:
-                    f.write(await file.read())
-
-            return Response(status_code=204)
-
-        ### OTHER
-        @self.on_event("startup")
-        async def register_on_startup():
-            # print(multiprocessing.get_start_method())
-            """Looks for manager node at startup and initializes multiprocessing module"""
-            # if not os.environ.get("PYTEST", False):
-            multiprocessing.set_start_method("spawn")
-            asyncio.create_task(self._register_with_manager())
-
-        @self.exception_handler(500)
-        async def internal_exception_handler(request: Request, exc: Exception):
-            if self.email_sender:
-                self.email_sender.send_email_exception(
-                    self.name, self.host_name, datetime.datetime.now()
-                )
-
-            return JSONResponse(
-                status_code=500, content={"code": 500, "msg": "Internal Server Error"}
-            )
-
-        @self.on_event("startup")
-        async def start_cleaning_loop():
-            asyncio.create_task(self._delete_expired_jobs_loop())
-
     @classmethod
     def process_wrapper(cls, inputs, job, result_queue):
         """Wrapper for the process function. It has two purposes: catching errors and packing the output of the process function into the "queue" object."""
@@ -423,22 +278,3 @@ def help_cli_args(self):
         for key, val in self.output_spec.__fields__.items():
             help_string += f"\t{key}: {val.type_.__name__}\n"
         return help_string
-
-
-def convert_string_to_type(text, type_):
-    """Convert a string to a type. This is used to parse the inputs to the process function."""
-    if type_ == str or type_ == FilePath:
-        return text
-    elif type_ == bool:
-        return text_to_bool(text)
-    else:
-        return type_(text)
-
-
-def text_to_bool(value):
-    if value.lower() in ("yes", "true", "t", "y", "1"):
-        return True
-    elif value.lower() in ("no", "false", "f", "n", "0"):
-        return False
-    else:
-        raise ValueError(f"Could not convert {value} to bool")