feat(BA-695): Configurable setup for kernel initialization polling

changes/3657.feature.md

@@ -0,0 +1 @@
+Add configurable setup for kernel initialization polling

src/ai/backend/agent/agent.py

@@ -2227,10 +2227,20 @@ async def create_kernel(
                 current_task = asyncio.current_task()
                 assert current_task is not None
                 self._pending_creation_tasks[kernel_id].add(current_task)
+                kernel_init_polling_attempt = cast(
+                    int, self.local_config["agent"]["kernel-lifecycles"]["init-polling-attempt"]
+                )
+                kernel_init_polling_timeout = cast(
+                    float,
+                    self.local_config["agent"]["kernel-lifecycles"]["init-polling-timeout-sec"],
+                )
                 try:
                     async for attempt in AsyncRetrying(
                         wait=wait_fixed(0.3),
-                        stop=(stop_after_attempt(10) | stop_after_delay(60)),
+                        stop=(
+                            stop_after_attempt(kernel_init_polling_attempt)
+                            | stop_after_delay(kernel_init_polling_timeout)
+                        ),
                         retry=retry_if_exception_type(zmq.error.ZMQError),
                     ):
                         with attempt:

src/ai/backend/agent/config.py

@@ -171,8 +171,15 @@
 
 DEFAULT_PULL_TIMEOUT = 2 * 60 * 60  # 2 hours
 DEFAULT_PUSH_TIMEOUT = None  # Infinite
+DEFAULT_KERNEL_INIT_POLLING_ATTEMPT = 10
+DEFAULT_KERNEL_INIT_POLLING_TIMEOUT = 60.0  # 60 seconds
 
 default_api_config = {"pull-timeout": DEFAULT_PULL_TIMEOUT, "push-timeout": DEFAULT_PUSH_TIMEOUT}
+default_kernel_lifecycles = {
+    "init-polling-attempt": DEFAULT_KERNEL_INIT_POLLING_ATTEMPT,
+    "init-polling-timeout-sec": DEFAULT_KERNEL_INIT_POLLING_TIMEOUT,
+}
+
 
 agent_etcd_config_iv = t.Dict({
     t.Key("container-logs", default=default_container_logs_config): t.Dict({
@@ -187,6 +194,16 @@
             0:
         ],  # Set the image push timeout in seconds. 'None' indicates an infinite timeout.
     }).allow_extra("*"),
+    t.Key("kernel-lifecycles", default=default_kernel_lifecycles): t.Dict({
+        t.Key(
+            "init-polling-attempt",
+            default=default_kernel_lifecycles["init-polling-attempt"],
+        ): t.ToInt,
+        t.Key(
+            "init-polling-timeout-sec",
+            default=default_kernel_lifecycles["init-polling-timeout-sec"],
+        ): t.ToFloat,
+    }),
 }).allow_extra("*")
 
 container_etcd_config_iv = t.Dict({