diff --git a/batchgen/server/batch_scheduler.py b/batchgen/server/batch_scheduler.py index f2400285..14b99029 100644 --- a/batchgen/server/batch_scheduler.py +++ b/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) diff --git a/docs/INSTALL.md b/docs/INSTALL.md index 10269423..d47a2a70 100644 --- a/docs/INSTALL.md +++ b/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 ` 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 ` 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 diff --git a/docs/client-api.md b/docs/client-api.md index 0b3aa3c9..de0c1e92 100644 --- a/docs/client-api.md +++ b/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 \ diff --git a/docs/manual-installation.md b/docs/manual-installation.md index 57e75d89..c30e95ce 100644 --- a/docs/manual-installation.md +++ b/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 diff --git a/scripts/install_deps.sh b/scripts/install_deps.sh index f558c7c0..6f785ee8 100755 --- a/scripts/install_deps.sh +++ b/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 . 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 ' | 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." diff --git a/setup.py b/setup.py index 26e15f47..d3259053 100644 --- a/setup.py +++ b/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", diff --git a/tests/test_pool_mode_failure_status.py b/tests/test_pool_mode_failure_status.py new file mode 100644 index 00000000..738735a9 --- /dev/null +++ b/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"]