follow up main code

Author: fregataa

src/ai/backend/agent/kernel_registry/loader/pickle.py

@@ -15,29 +15,44 @@
 
 
 class PickleBasedKernelRegistryLoader(AbstractKernelRegistryLoader[KernelRegistry]):
-    def __init__(self, last_registry_file_path: Path, legacy_registry_file_path: Path) -> None:
+    def __init__(
+        self,
+        last_registry_file_path: Path,
+        fallback_registry_file_path: Path,
+        legacy_registry_file_path: Path,
+    ) -> None:
         self._last_registry_file_path = last_registry_file_path
+        self._fallback_registry_file_path = fallback_registry_file_path
         self._legacy_registry_file_path = legacy_registry_file_path
 
     @override
     async def load_kernel_registry(self) -> KernelRegistry:
         legacy_registry_file = self._legacy_registry_file_path
-        last_registry_file = self._last_registry_file_path
+        fallback_registry_file = self._fallback_registry_file_path
+        final_file_path = self._last_registry_file_path
+        if not final_file_path.is_file():
+            log.warning(
+                "Registry file with name {} not found. "
+                "Falling back to path with local instance id: {}",
+                final_file_path,
+                fallback_registry_file,
+            )
+            final_file_path = fallback_registry_file
         try:
             if os.path.isfile(legacy_registry_file):
-                shutil.move(legacy_registry_file, last_registry_file)
+                shutil.move(legacy_registry_file, final_file_path)
         except Exception as e:
             log.warning(
                 "Failed to move legacy kernel registry file {} to {} (err: {})",
                 str(legacy_registry_file),
-                str(last_registry_file),
+                str(final_file_path),
                 str(e),
             )
         try:
-            with open(last_registry_file, "rb") as f:
+            with open(final_file_path, "rb") as f:
                 return pickle.load(f)
         except EOFError as e:
-            log.warning("Failed to load the last kernel registry: {}", str(last_registry_file))
+            log.warning("Failed to load the last kernel registry: {}", str(final_file_path))
             raise KernelRegistryLoadError from e
         except FileNotFoundError as e:
             raise KernelRegistryNotFound from e

src/ai/backend/agent/kernel_registry/recovery/pickle.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from ai.backend.common.types import AgentId
+
 from ...kernel import KernelRegistry
 from ..loader.pickle import PickleBasedKernelRegistryLoader
 from ..writer.pickle import PickleBasedKernelRegistryWriter
@@ -11,17 +13,20 @@
 class PickleBasedKernelRegistryRecoveryArgs:
     ipc_base_path: Path
     var_base_path: Path
+    agent_id: AgentId
     local_instance_id: str
 
 
 class PickleBasedKernelRegistryRecovery:
     def __init__(self, args: PickleBasedKernelRegistryRecoveryArgs) -> None:
-        registry_file_name = f"kernel_registry.{args.local_instance_id}.dat"
+        registry_file_name = f"kernel_registry.{args.agent_id}.dat"
+        fallback_registry_file_name = f"kernel_registry.{args.local_instance_id}.dat"
         legacy_registry_file_path = args.ipc_base_path / registry_file_name
+        fallback_registry_file_path = args.var_base_path / fallback_registry_file_name
         last_registry_file_path = args.var_base_path / registry_file_name
 
         self._loader = PickleBasedKernelRegistryLoader(
-            last_registry_file_path, legacy_registry_file_path
+            last_registry_file_path, fallback_registry_file_path, legacy_registry_file_path
         )
         self._writer = PickleBasedKernelRegistryWriter(last_registry_file_path)