Wait for CybORG policy workers before decisions

Author: Muhtasham

codeclash/arenas/cyborg/runtime/run_cyborg.py

@@ -26,15 +26,17 @@ def safe_module_name(player_name: str) -> str:
 
 def load_policy_module(player_name: str, path: str):
     agent_dir = str(Path(path).resolve().parent)
-    sys.path.insert(0, agent_dir)
+    inserted_agent_dir = agent_dir not in sys.path
+    if inserted_agent_dir:
+        sys.path.insert(0, agent_dir)
     spec = importlib.util.spec_from_file_location(safe_module_name(player_name), path)
     if spec is None or spec.loader is None:
         raise RuntimeError(f"Could not load module spec from {path}")
     module = importlib.util.module_from_spec(spec)
     try:
         spec.loader.exec_module(module)
     finally:
-        with contextlib.suppress(ValueError):
+        if inserted_agent_dir:
             sys.path.remove(agent_dir)
     if not hasattr(module, "decide") or not callable(module.decide):
         raise RuntimeError(f"{path} must define a callable decide(observation, action_space)")
@@ -74,19 +76,24 @@ def normalize_action(action, action_space) -> tuple[int, str | None]:
 
 
 def policy_worker(
-    command_queue: mp.Queue, result_queue: mp.Queue, player_name: str, agent_name: str, path: str
+    command_queue: mp.Queue,
+    result_queue: mp.Queue,
+    startup_queue: mp.Queue,
+    player_name: str,
+    agent_name: str,
+    path: str,
 ) -> None:
     try:
         module = load_policy_module(player_name, path)
     except BaseException as exc:
-        result_queue.put(
+        startup_queue.put(
             {
-                "request_id": None,
                 "error": f"{type(exc).__name__}: {exc}",
                 "traceback": traceback.format_exc(limit=5),
             }
         )
         return
+    startup_queue.put({"ready": True})
 
     while True:
         try:
@@ -118,18 +125,38 @@ def __init__(self, player_name: str, agent_name: str, path: str, *, timeout: flo
         self.errors: list[dict] = []
         self.invalid_actions = 0
         self.decisions = 0
+        self.startup_error: str | None = None
         self._next_request_id = 0
         self._start_worker()
 
     def _start_worker(self) -> None:
+        self.startup_error = None
         ctx = mp.get_context("spawn")
         self.command_queue = ctx.Queue()
         self.result_queue = ctx.Queue()
+        self.startup_queue = ctx.Queue()
         self.process = ctx.Process(
             target=policy_worker,
-            args=(self.command_queue, self.result_queue, self.player_name, self.agent_name, self.path),
+            args=(
+                self.command_queue,
+                self.result_queue,
+                self.startup_queue,
+                self.player_name,
+                self.agent_name,
+                self.path,
+            ),
         )
         self.process.start()
+        startup_timeout = max(self.timeout, 10.0)
+        try:
+            startup_message = self.startup_queue.get(timeout=startup_timeout)
+        except queue.Empty:
+            self.startup_error = f"policy import exceeded {startup_timeout}s timeout"
+            self.close()
+            return
+        if "error" in startup_message:
+            self.startup_error = startup_message["error"]
+            self.close()
 
     def close(self) -> None:
         with contextlib.suppress(Exception):
@@ -147,6 +174,10 @@ def restart(self) -> None:
         self._start_worker()
 
     def decide(self, observation, action_space) -> int:
+        if self.startup_error is not None:
+            self.errors.append({"agent": self.agent_name, "error": self.startup_error})
+            self.restart()
+            return DEFAULT_ACTION
         request_id = self._next_request_id
         self._next_request_id += 1
         self.command_queue.put(