[bugfix] Bound worker startup and propagate child failures

docs/flexkv_config_reference/README_en.md

@@ -206,3 +206,15 @@ Some configurations can only be set through environment variables.
 | Environment Variable | Type | Default | Description |
 |---------------------|------|---------|-------------|
 | `FLEXKV_ENABLE_COLLECTIVE_SYNC` | bool | 1 | Whether to enable cross-rank collective sync (scatter/barrier/all_reduce). This sync is primarily used for coordination between Pipeline Parallelism (PP) stages. Can be set to 0 to disable in non-PP deployments to reduce sync overhead and improve performance |
+
+## Worker startup deadlines
+
+- `FLEXKV_WORKER_SPAWN_TIMEOUT_S` (default `60`): bounds the TransferManager
+  subprocess start handshake.
+- `FLEXKV_READY_TIMEOUT_S` (default `360`): bounds the SGLang connector's wait
+  for FlexKV readiness after registration.
+
+Both values must be finite and positive. Initialization failures include the
+child traceback in the parent error. Health checks continue after readiness,
+including during completion waits. These settings do not shorten the existing
+transfer drain/shutdown timeout.

flexkv/integration/sglang/connector.py

@@ -28,12 +28,13 @@
 import atexit
 import json
 import logging
+import math
 import os
 import signal
 import socket
 import struct
 import time
-from contextlib import nullcontext
+from contextlib import nullcontext, suppress
 from dataclasses import dataclass, replace
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
@@ -108,7 +109,7 @@ class FlexKVHostReleaseShim:
     ``FlexKVConnector.shutdown()`` so unpin stays on that same path.
     """
 
-    def __init__(self, connector: "FlexKVConnector") -> None:
+    def __init__(self, connector: FlexKVConnector) -> None:
         self._connector = connector
 
     def destroy(self) -> None:
@@ -331,10 +332,8 @@ def __init__(
         #   * atexit runs kv_manager.shutdown() (cudaHostUnregister).
         #   * Stretch the parent's scheduler-exit wait so kill_process_tree
         #     does not SIGKILL mid-unpin (generic env, not FlexKV-named in TM).
-        try:
+        with suppress(Exception):
             signal.signal(signal.SIGINT, signal.SIG_IGN)
-        except Exception:  # noqa: BLE001
-            pass
         # Graceful shutdown timeout hierarchy (see FlexKV config.py):
         #   tokenizer wait (this env)                   = 1200s
         #     > TM parent-side wait (FLEXKV_TRANSFER_MANAGER_SHUTDOWN_TIMEOUT_S) = 900s
@@ -1269,40 +1268,39 @@ def check_completed_stores(self) -> List[str]:
         completed_rids: List[str] = []
         completed_by_rid: Dict[str, Any] = {}
 
-        if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
-            if self._inflight_stores:
-                fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
-                try:
-                    completed_dict = (
-                        self.kv_manager.wait(
-                            list(fk_to_rid.keys()),
-                            timeout=0.0,
-                            completely=True,
-                        )
-                        or {}
-                    )
-                except Exception as exc:  # noqa: BLE001
-                    rid = next(iter(self._inflight_stores))
-                    context = getattr(self, "_inflight_store_contexts", {}).get(rid)
-                    context = context or self._new_op_context(
-                        "store", rid, task_id=self._inflight_stores[rid]
-                    )
-                    self._log_cache_op(
-                        context,
-                        "poll",
-                        "failed",
-                        task_id=-1,
-                        flexkv_task_ids=list(fk_to_rid),
-                        error=str(exc),
+        if self._sync_ctx.is_sync_leader and self.kv_manager is not None and self._inflight_stores:
+            fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
+            try:
+                completed_dict = (
+                    self.kv_manager.wait(
+                        list(fk_to_rid.keys()),
+                        timeout=0.0,
+                        completely=True,
                     )
-                    completed_dict = {}
-                for fk_tid, response in completed_dict.items():
-                    status = _status_value(response)
-                    if not _is_terminal_status(status):
-                        continue
-                    rid = fk_to_rid[fk_tid]
-                    completed_rids.append(rid)
-                    completed_by_rid[rid] = response
+                    or {}
+                )
+            except Exception as exc:  # noqa: BLE001
+                rid = next(iter(self._inflight_stores))
+                context = getattr(self, "_inflight_store_contexts", {}).get(rid)
+                context = context or self._new_op_context(
+                    "store", rid, task_id=self._inflight_stores[rid]
+                )
+                self._log_cache_op(
+                    context,
+                    "poll",
+                    "failed",
+                    task_id=-1,
+                    flexkv_task_ids=list(fk_to_rid),
+                    error=str(exc),
+                )
+                completed_dict = {}
+            for fk_tid, response in completed_dict.items():
+                status = _status_value(response)
+                if not _is_terminal_status(status):
+                    continue
+                rid = fk_to_rid[fk_tid]
+                completed_rids.append(rid)
+                completed_by_rid[rid] = response
 
         if self._sync_ctx.needs_sync:
             completed_rids = self._sync_ctx.scatter(
@@ -2227,14 +2225,18 @@ def _store_profile_scope(self, name: str):
 
     def _wait_kv_manager_ready(self, poll_interval: float = 10.0) -> None:
         assert self.kv_manager is not None
-        wait_count = 0
+        timeout = float(os.environ.get("FLEXKV_READY_TIMEOUT_S", "360"))
+        if not math.isfinite(timeout) or timeout <= 0:
+            raise ValueError("FLEXKV_READY_TIMEOUT_S must be finite and positive")
+        start = time.monotonic()
         while not self.kv_manager.is_ready():
-            time.sleep(poll_interval)
-            wait_count += 1
+            remaining = timeout - (time.monotonic() - start)
+            if remaining <= 0:
+                raise TimeoutError(f"FlexKV did not become ready within {timeout}s {self._label}")
+            time.sleep(min(poll_interval, remaining))
             logger.info(
                 "[FlexKV] Waiting for FlexKV ready %s (waited %.0fs)",
-                self._label,
-                wait_count * poll_interval,
+                self._label, time.monotonic() - start,
             )
         logger.info("[FlexKV] FlexKV is ready %s", self._label)
 

flexkv/storage/allocator.py

@@ -16,6 +16,7 @@
 from flexkv.common.memory_handle import TensorSharedHandle
 from flexkv.common.storage import StorageHandle, AccessHandleType, KVCacheLayout, KVCacheLayoutType
 from flexkv.common.debug import flexkv_logger
+import contextlib
 
 
 class BaseStorageAllocator(ABC):
@@ -218,10 +219,8 @@ def _cleanup_hugepage_mapping(addr: int, aligned: int, fd: int,
     if fd >= 0:
         _libc.close(fd)
     if path is not None:
-        try:
+        with contextlib.suppress(FileNotFoundError):
             os.unlink(path)
-        except FileNotFoundError:
-            pass
     _live_hugepage_mappings.pop(data_ptr, None)
 
 
@@ -230,10 +229,8 @@ def _cleanup_hugepage_mmap(mm: mmap.mmap, path: str | None, data_ptr: int) -> No
         mm.close()
     finally:
         if path is not None:
-            try:
+            with contextlib.suppress(FileNotFoundError):
                 os.unlink(path)
-            except FileNotFoundError:
-                pass
         _live_hugepage_mappings.pop(data_ptr, None)
 
 
@@ -273,10 +270,8 @@ def _create_hugetlbfs_file(aligned: int) -> tuple[str, int]:
             )
     except Exception:
         os.close(fd)
-        try:
+        with contextlib.suppress(FileNotFoundError):
             os.unlink(path)
-        except FileNotFoundError:
-            pass
         raise
     return path, fd
 
@@ -429,10 +424,8 @@ def _mmap_huge(num_bytes: int, page_size_bytes: int) -> Tuple[int, int, int]:
     fd = -1
     try:
         path, fd = _create_hugetlbfs_file(aligned)
-        try:
+        with contextlib.suppress(OSError):
             os.unlink(path)
-        except OSError:
-            pass
 
         ctypes.set_errno(0)
         ret = _libc.mmap(
@@ -495,16 +488,25 @@ def alloc_hugepage_tensor(num_elements: int,
                 f"mapping_alignment={mapping_alignment}"
             )
         path, fd = _create_hugetlbfs_file(aligned)
+        mm = None
         try:
             mm = mmap.mmap(
                 fd,
                 aligned,
                 flags=mmap.MAP_SHARED,
                 prot=mmap.PROT_READ | mmap.PROT_WRITE,
             )
+            return _wrap_mmap_tensor(mm, aligned, num_elements, dtype, cleanup_path=path)
+        except Exception:
+            try:
+                if mm is not None:
+                    mm.close()
+            finally:
+                with contextlib.suppress(FileNotFoundError):
+                    os.unlink(path)
+            raise
         finally:
             os.close(fd)
-        return _wrap_mmap_tensor(mm, aligned, num_elements, dtype, cleanup_path=path)
 
     addr, aligned, fd = _mmap_huge(num_bytes, page_size_bytes)
 
@@ -695,8 +697,10 @@ def allocate(cls,
                         f"SSD allocator progress: {file_count}/{total_num_files} files created "
                         f"({file_count * 100 // total_num_files}%)"
                     )
-        flexkv_logger.info(f"SSD allocator done: {total_num_files} files in {cache_dir}, "
-                           f"each file has {real_file_size/1024/1024/1024:.2f} GB, total size {real_total_size/1024/1024/1024:.2f} GB")
+        flexkv_logger.info(
+            f"SSD allocator done: {total_num_files} files in {cache_dir}, "
+            f"each file has {real_file_size/1024/1024/1024:.2f} GB, total size {real_total_size/1024/1024/1024:.2f} GB"
+        )
         return StorageHandle(
             handle_type=AccessHandleType.FILE,
             data=ssd_files,