fix: prepare v1.0.10.post5 with pool-mode failure, numa.h prereq, doc fixes

batchgen/server/batch_scheduler.py

@@ -907,9 +907,9 @@ async def _process_batch_pool_mode(
                 f"Batch with {len(entries)} requests rejected. Retry later."
             )
             logger.warning(f"[POOL] Batch {batch_id} rejected: {error_msg}")
-            self.storage.update_batch(batch_id, status="failed", error={
-                "code": "capacity_exceeded", "message": error_msg,
-            })
+            self.storage.update_batch_status(
+                batch_id, BatchStatus.FAILED, error=f"capacity_exceeded: {error_msg}"
+            )
             return
         # Store max_tokens for init message
         if not hasattr(self, '_pool_max_output_len'):
@@ -974,9 +974,9 @@ async def _wait_and_finalize_batch(
 
         if batch_failed:
             error_msg = getattr(tracker, 'error', 'timeout') if tracker else 'timeout'
-            self.storage.update_batch(batch_id, status="failed", error={
-                "code": "batch_failed", "message": str(error_msg)
-            })
+            self.storage.update_batch_status(
+                batch_id, BatchStatus.FAILED, error=f"batch_failed: {error_msg}"
+            )
             return
 
         self._finalize_batch_output(batch_id, requests, prompts)

docs/INSTALL.md

@@ -5,6 +5,9 @@
 - **CUDA**: 12.8+ toolkit installed (Blackwell/B200 requires **12.9+**, see below)
 - **Python**: 3.11+
 - **OS**: Ubuntu 22.04 (tested)
+- **libnuma**: `libnuma-dev` headers (`numa.h`) — required by the core-engine
+  JIT at first server launch (`sudo apt-get install -y libnuma-dev`). The
+  Docker image (Option A) already includes it.
 - **GitHub access**: the repository is currently **private** — anonymous
   `git clone` and raw release-asset URLs fail (404). Authenticate first
   (e.g. `gh auth login`), clone via `gh repo clone batchgen-project/batchgen`,
@@ -52,13 +55,13 @@ Total: ~40-50 min on first install.
 
 Pre-built wheels for Hopper GPUs (CUDA 12.8, PyTorch 2.9, Python 3.11) are
 available on the [GitHub Releases](https://github.com/batchgen-project/batchgen/releases) page.
-The wheel set below is pinned to `v1.0.10.post2` — the most recent release that
-ships the full dependency wheel set (later releases only ship the batchgen
-wheels). While the repository is private, plain `pip install <URL>` returns
-404; download the wheels with an authenticated client first:
+The wheel set below is pinned to `v1.0.10.post5`, which ships the full
+dependency wheel set. While the repository is private, plain
+`pip install <URL>` returns 404; download the wheels with an authenticated
+client first:
 
 ```bash
-gh release download v1.0.10.post2 -R batchgen-project/batchgen -p '*.whl' -D ./wheels
+gh release download v1.0.10.post5 -R batchgen-project/batchgen -p '*.whl' -D ./wheels
 pip install ./wheels/*.whl
 ```
 
@@ -71,13 +74,13 @@ conda activate batchgen
 pip install torch==2.9.0+cu128 --index-url https://download.pytorch.org/whl/cu128
 
 # 2. Install all wheels from the pinned release
-RELEASE_URL="https://github.com/batchgen-project/batchgen/releases/download/v1.0.10.post2"
+RELEASE_URL="https://github.com/batchgen-project/batchgen/releases/download/v1.0.10.post5"
 pip install \
   "${RELEASE_URL}/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl" \
   "${RELEASE_URL}/flash_mla-1.0.0+1408756-cp311-cp311-linux_x86_64.whl" \
   "${RELEASE_URL}/deep_gemm-2.1.1+c9f8b34-cp311-cp311-linux_x86_64.whl" \
-  "${RELEASE_URL}/batchgen_kernels-0.3.2+sm90a-cp311-cp311-linux_x86_64.whl" \
-  "${RELEASE_URL}/batchgen-1.0.10.post2-py3-none-any.whl"
+  "${RELEASE_URL}/batchgen_kernels-0.3.3+sm90a-cp311-cp311-linux_x86_64.whl" \
+  "${RELEASE_URL}/batchgen-1.0.10.post5-py3-none-any.whl"
 ```
 
 No source compilation needed — pip auto-installs all remaining Python dependencies
@@ -169,7 +172,7 @@ This does not apply to Docker (Option A), where the source is the install target
 | **Do not run from source dir** | Source tree shadows installed packages (see above) |
 | **batchgen_kernels** | Must use `--no-build-isolation` (needs installed PyTorch headers) |
 | **H20 GPUs** | Set `TORCH_CUDA_ARCH_LIST=9.0a` before building kernels |
-| **Core engine** | JIT-compiled at first server launch via ninja (automatic, ~5s) |
+| **Core engine** | JIT-compiled at first server launch via ninja (automatic, ~5s; needs `libnuma-dev`) |
 | **No JIT for compute kernels** | All 23 CUDA extensions are AOT-compiled in `batchgen_kernels` |
 
 ## Verification

docs/client-api.md

@@ -251,7 +251,7 @@ BatchGen provides test scripts that demonstrate how to use the client API. These
 ### Example: MMLU Pro Batch Test
 
 ```bash
-python test/r1_mmlu_pro_test/r1_mmlu_pro_batch_test.py \
+python tests/e2e/r1_mmlu_pro_test/r1_mmlu_pro_batch_test.py \
     --hugging_face_checkpoint /path/to/DeepSeek-R1 \
     --max_decoding_length 10240 \
     --server_host localhost \

docs/manual-installation.md

@@ -14,6 +14,8 @@ For quick installation, use the automated script instead:
 - Python 3.11+
 - CUDA 12.8+ toolkit
 - Git
+- libnuma headers (`sudo apt-get install -y libnuma-dev`) — required by the
+  core-engine JIT at first server launch
 
 Create and activate a virtual environment:
 ```bash

scripts/install_deps.sh

@@ -83,6 +83,24 @@ check_prerequisites() {
     fi
     print_success "git found"
 
+    # Check libnuma headers (numa.h) — the core engine JIT-compiles
+    # batchgen/core/Parameter_Server/posix_shm.cpp at first server launch and
+    # includes <numa.h>. Missing headers => first `launch_http_server` dies with
+    # "fatal error: numa.h: No such file or directory". Fail early with a hint.
+    if command -v cc &> /dev/null; then
+        numa_ok=$(echo '#include <numa.h>' | cc -fsyntax-only -x c - &> /dev/null && echo yes || echo no)
+    elif [[ -f /usr/include/numa.h ]]; then
+        numa_ok=yes
+    else
+        numa_ok=no
+    fi
+    if [[ "$numa_ok" != "yes" ]]; then
+        print_error "libnuma headers (numa.h) not found — required by the core-engine JIT at first launch."
+        print_error "Install them first:  sudo apt-get install -y libnuma-dev"
+        exit 1
+    fi
+    print_success "libnuma headers found"
+
     # Check CUDA
     if ! command -v nvcc &> /dev/null; then
         print_warning "nvcc not found. CUDA may not be properly configured."

setup.py

@@ -134,7 +134,7 @@ def _prepare_package_data(self):
 # install all files in the package, rather than just the egg
 setup(
     name="batchgen",
-    version=os.getenv("BATCHGEN_VERSION", "1.0.10.post4"),
+    version=os.getenv("BATCHGEN_VERSION", "1.0.10.post5"),
     packages=find_packages(
         include=[
             "batchgen",

tests/test_pool_mode_failure_status.py

@@ -0,0 +1,141 @@
+import asyncio
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_real(name, monkeypatch):
+    path = _ROOT.joinpath(*name.split(".")).with_suffix(".py")
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, name, module)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _stub(name, monkeypatch, **attrs):
+    module = types.ModuleType(name)
+    module.__path__ = []
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+    return module
+
+
+@pytest.fixture
+def scheduler_env(monkeypatch):
+    # Stub the package parents so batchgen.server.__init__ (which pulls the whole
+    # CUDA core_engine + model registry) never runs; load only the pure-Python
+    # scheduler deps for real.
+    _stub("batchgen", monkeypatch)
+    _stub("batchgen.server", monkeypatch)
+    io_struct = _load_real("batchgen.server.io_struct", monkeypatch)
+    _load_real("batchgen.server.intake_pool", monkeypatch)
+    _stub(
+        "batchgen.server.worker_manager",
+        monkeypatch,
+        WorkerManager=type("WorkerManager", (), {}),
+    )
+    _stub(
+        "batchgen.server.scheduling_pool",
+        monkeypatch,
+        SchedulingPool=type("SchedulingPool", (), {}),
+    )
+    _stub(
+        "batchgen.server.server_args",
+        monkeypatch,
+        ServerArgs=type("ServerArgs", (), {}),
+    )
+    _stub(
+        "batchgen.server.storage",
+        monkeypatch,
+        StorageManager=type("StorageManager", (), {}),
+    )
+    batch_scheduler = _load_real("batchgen.server.batch_scheduler", monkeypatch)
+    return batch_scheduler.BatchScheduler, io_struct.BatchStatus
+
+
+class _FakeStorage:
+    def __init__(self):
+        self.calls = []
+
+    def update_batch_status(self, batch_id, status, **updates):
+        self.calls.append((batch_id, status, updates))
+
+
+class _RejectingIntakePool:
+    max_capacity = 4
+
+    def size(self):
+        return 4
+
+    def submit_batch(self, batch_id, entries, priority):
+        return False
+
+
+class _EmptySchedulingPool:
+    def register_batch(self, **kwargs):
+        return None
+
+    def get_batch_tracker(self, batch_id):
+        return None
+
+
+def _bare_scheduler(BatchScheduler, storage):
+    sched = BatchScheduler.__new__(BatchScheduler)
+    sched.storage = storage
+    sched._scheduling_pool = _EmptySchedulingPool()
+    return sched
+
+
+def test_pool_capacity_rejection_marks_batch_failed(scheduler_env):
+    BatchScheduler, BatchStatus = scheduler_env
+    # Given a pool-mode scheduler whose intake pool is at capacity
+    storage = _FakeStorage()
+    sched = _bare_scheduler(BatchScheduler, storage)
+    sched._intake_pool = _RejectingIntakePool()
+    req = types.SimpleNamespace(custom_id="r0")
+    batch = types.SimpleNamespace(batchgen_debug={})
+
+    # When a batch is submitted and the intake pool rejects it
+    asyncio.run(
+        sched._process_batch_pool_mode(
+            batch_id="b0",
+            batch=batch,
+            requests=[req],
+            prompts=["hi"],
+            per_request_max_tokens=[8],
+            sampling_params=[{}],
+            incremental_kwargs={},
+        )
+    )
+
+    # Then it is marked FAILED via update_batch_status (regression guard for the missing update_batch)
+    assert len(storage.calls) == 1
+    batch_id, status, updates = storage.calls[0]
+    assert batch_id == "b0"
+    assert status is BatchStatus.FAILED
+    assert "capacity_exceeded" in updates["error"]
+
+
+def test_pool_batch_timeout_marks_batch_failed(scheduler_env):
+    BatchScheduler, BatchStatus = scheduler_env
+    # Given a pool-mode scheduler whose batch deadline is already past and no tracker exists
+    storage = _FakeStorage()
+    sched = _bare_scheduler(BatchScheduler, storage)
+    sched._batch_timeout = -1
+
+    # When the finalize task waits for a batch that never completes
+    asyncio.run(sched._wait_and_finalize_batch("b1", requests=[], prompts=[]))
+
+    # Then it is marked FAILED via update_batch_status with a batch_failed error
+    assert len(storage.calls) == 1
+    batch_id, status, updates = storage.calls[0]
+    assert batch_id == "b1"
+    assert status is BatchStatus.FAILED
+    assert "batch_failed" in updates["error"]