fix(BA-932): Unload purged docker images

changes/3923.fix.md

@@ -0,0 +1 @@
+Unload images that have been purged from the agents in the manager.

src/ai/backend/agent/docker/agent.py

@@ -1240,7 +1240,6 @@ class DockerAgent(AbstractAgent[DockerKernel, DockerKernelCreationContext]):
     monitor_docker_task: asyncio.Task
     agent_sockpath: Path
     agent_sock_task: asyncio.Task
-    scan_images_timer: asyncio.Task
     metadata_server: MetadataServer
     docker_ptask_group: aiotools.PersistentTaskGroup
     gwbridge_subnet: Optional[str]

src/ai/backend/manager/registry.py

@@ -237,6 +237,7 @@ class AgentRegistry:
     pending_waits: set[asyncio.Task[None]]
     database_ptask_group: aiotools.PersistentTaskGroup
     webhook_ptask_group: aiotools.PersistentTaskGroup
+    agent_installed_images: dict[AgentId, set[str]]
 
     def __init__(
         self,
@@ -287,6 +288,7 @@ def __init__(
             hook_plugin_ctx,
             self,
         )
+        self.agent_installed_images = {}
 
     async def init(self) -> None:
         self.heartbeat_lock = asyncio.Lock()
@@ -3149,19 +3151,21 @@ async def _update() -> None:
                 )
 
             # Update the mapping of kernel images to agents.
-            loaded_images = msgpack.unpackb(zlib.decompress(agent_info["images"]))
+            images = msgpack.unpackb(zlib.decompress(agent_info["images"]))
+            image_canonicals = set(img_info[0] for img_info in images)
+            prev_image_canonicals = self.agent_installed_images.get(agent_id, set())
 
             async def _pipe_builder(r: Redis):
                 pipe = r.pipeline()
-                for image, _ in loaded_images:
-                    try:
-                        await pipe.sadd(ImageRef.parse_image_str(image, "*").canonical, agent_id)
-                    except ValueError:
-                        # Skip opaque (non-Backend.AI) image.
-                        continue
+                for image in image_canonicals:
+                    await pipe.sadd(image, agent_id)
+
+                for image in prev_image_canonicals - image_canonicals:
+                    await pipe.srem(image, agent_id)
                 return pipe
 
             await redis_helper.execute(self.redis_image, _pipe_builder)
+            self.agent_installed_images[agent_id] = image_canonicals
 
         await self.hook_plugin_ctx.notify(
             "POST_AGENT_HEARTBEAT",