Merge branch 'main' into feat-merge-inference

Author: anna-grim

src/deep_neurographs/utils/ml_util.py

@@ -20,6 +20,73 @@
 GNN_DEPTH = 2
 
 
+# --- GPU Scheduler ---
+import torch
+from multiprocessing import Process, Queue
+from queue import Empty
+
+class GPUScheduler:
+    def __init__(self, model_path, num_gpus):
+        self.model_path = model_path
+        self.num_gpus = num_gpus
+        self.job_queues = []
+        self.return_queues = []
+        self.processes = []
+        self._init_workers()
+
+    def _init_workers(self):
+        for gpu_id in range(self.num_gpus):
+            job_q = Queue()
+            ret_q = Queue()
+            p = Process(
+                target=self._gpu_worker,
+                args=(gpu_id, self.model_path, job_q, ret_q),
+            )
+            p.start()
+            self.job_queues.append(job_q)
+            self.return_queues.append(ret_q)
+            self.processes.append(p)
+
+    def _gpu_worker(self, gpu_id, model_path, job_queue, return_queue):
+        device = torch.device(f"cuda:{gpu_id}")
+        model = torch.load(model_path, map_location=device)
+        model.eval()
+        while True:
+            job = job_queue.get()
+            if job is None:
+                break  # Sentinel to exit
+            batch, job_id = job
+            with torch.no_grad():
+                batch = batch.to(device)
+                preds = model(batch)
+                return_queue.put((job_id, preds.cpu()))
+
+    def submit(self, batch, job_id):
+        """Submit a batch to the next GPU in round-robin fashion."""
+        gpu_id = job_id % self.num_gpus
+        self.job_queues[gpu_id].put((batch, job_id))
+
+    def get_result(self, job_id, timeout=None):
+        """Retrieve results from the return queues."""
+        for q in self.return_queues:
+            try:
+                result_job_id, preds = q.get(timeout=timeout)
+                if result_job_id == job_id:
+                    return preds
+                else:
+                    # Re-enqueue if it's not the one we're looking for
+                    q.put((result_job_id, preds))
+            except Empty:
+                continue
+        return None  # or raise TimeoutError
+
+    def shutdown(self):
+        """Stop all GPU worker processes cleanly."""
+        for q in self.job_queues:
+            q.put(None)
+        for p in self.processes:
+            p.join()
+
 # --- Batch Generation ---
 def get_batch(graph, proposals, batch_size, flagged_proposals=set()):
     """