diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f41e80a..515f255 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,6 +23,15 @@ jobs: with: python-version: "3.12" cache: pip + # Installing the package compiles the extension, which needs an ONNX + # Runtime SDK matching the wheel. CMake downloads it here and reuses it on + # later runs. + - name: Cache the ONNX Runtime SDK + uses: actions/cache@v4 + with: + path: ~/.cache/anytime-inference-planner + key: ort-sdk-${{ runner.os }}-${{ hashFiles('pyproject.toml') }} + restore-keys: ort-sdk-${{ runner.os }}- - name: Install run: pip install -e ".[dev]" - name: ruff check @@ -45,9 +54,28 @@ jobs: with: python-version: ${{ matrix.python-version }} cache: pip + - name: Cache the ONNX Runtime SDK + uses: actions/cache@v4 + with: + path: ~/.cache/anytime-inference-planner + key: ort-sdk-${{ runner.os }}-${{ hashFiles('pyproject.toml') }} + restore-keys: ort-sdk-${{ runner.os }}- - name: Install run: pip install -e ".[bench]" pytest + # The extension is the runtime, so an install that silently fell back to + # the Python backend would hide a broken build behind passing tests. + - name: Assert the extension was built + run: | + python -c " + from anytime_serving.serving import extension_available, load_extension + import onnxruntime + assert extension_available(), 'anytime_runtime did not build or import' + print('anytime_runtime links ONNX Runtime', load_extension().onnxruntime_version()) + print('onnxruntime wheel reports ', onnxruntime.__version__) + " - name: Run tests + env: + ANYTIME_REQUIRE_BACKENDS: extension,python run: pytest -q test-minimal: @@ -59,6 +87,12 @@ jobs: with: python-version: "3.12" cache: pip + - name: Cache the ONNX Runtime SDK + uses: actions/cache@v4 + with: + path: ~/.cache/anytime-inference-planner + key: ort-sdk-${{ runner.os }}-${{ hashFiles('pyproject.toml') }} + restore-keys: ort-sdk-${{ runner.os }}- # The serving control plane must run without the research stack. Installing # only the base dependencies enforces that: a module-scope torch or pandas # import anywhere on the serving path fails this job. @@ -75,50 +109,47 @@ jobs: run: python -c "import anytime_serving.serving, anytime_serving.models.cascade" # The remaining suites need pandas/scipy and are covered by the `test` job. - name: Run serving tests and import-boundary guards + env: + ANYTIME_REQUIRE_BACKENDS: extension,python run: | pytest -q \ tests/test_admission.py \ tests/test_selector.py \ tests/test_serving.py \ tests/test_load_monitor.py \ + tests/test_runtime_engine.py \ tests/test_import_boundaries.py - runtime: - name: C++ runtime build + engine: + name: Engine parity against the reference backend runs-on: ubuntu-latest - env: - # Must match the onnxruntime wheel the Python side installs. The version is - # not cosmetic: building the worker against 1.20.1 while profiling with the - # 1.26.0 wheel measured DistilBERT at 98.9 ms in the worker versus 13.0 ms - # in the profiler, a 7.6x gap that made every admission decision wrong. - ONNXRUNTIME_VERSION: 1.26.0 steps: - uses: actions/checkout@v4 - - name: Fetch ONNX Runtime - run: | - curl -fsSL -o onnxruntime.tgz \ - "https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz" - mkdir -p "$RUNNER_TEMP/onnxruntime" - # Release tarball layouts differ between versions: some nest a - # versioned directory, some do not. Locate the one holding include/. - tar -xzf onnxruntime.tgz -C "$RUNNER_TEMP/onnxruntime" --strip-components=1 - root="$RUNNER_TEMP/onnxruntime" - if [ ! -d "$root/include" ]; then - root=$(find "$RUNNER_TEMP/onnxruntime" -maxdepth 2 -type d -name include \ - -exec dirname {} \; | head -1) - fi - echo "ONNXRUNTIME_ROOT_PATH=$root" >> "$GITHUB_ENV" - - name: Configure - run: | - cmake -S runtime_cpp -B runtime_cpp/build \ - -DCMAKE_BUILD_TYPE=Release \ - -DONNXRUNTIME_ROOT_PATH="$ONNXRUNTIME_ROOT_PATH" - - name: Build - run: cmake --build runtime_cpp/build -j"$(nproc)" - - name: Verify the binary starts and reports its usage + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + - name: Cache the ONNX Runtime SDK + uses: actions/cache@v4 + with: + path: ~/.cache/anytime-inference-planner + key: ort-sdk-${{ runner.os }}-${{ hashFiles('pyproject.toml') }} + restore-keys: ort-sdk-${{ runner.os }}- + - name: Install + run: pip install -e . pytest + # Both copies of ONNX Runtime in this process must be the same version. The + # extension reports what it linked; the wheel reports what it will import. + - name: Report the resolved ONNX Runtime versions run: | - # No --model argument is an error; the binary must exit non-zero with a - # message rather than hang waiting on stdin. - if runtime_cpp/build/anytime_runtime and unpack it. Release tarball -layouts differ between versions: some nest a versioned directory. Point -`ONNXRUNTIME_ROOT_PATH` at whichever directory contains `include/` and `lib/`. - -The binary lands at `runtime_cpp/build/anytime_runtime`. The Python client finds -it via `ANYTIME_RUNTIME_BIN`, or by searching upward from the package and the -working directory. Set the environment variable explicitly for a non-editable -install, where the package no longer sits inside the source tree. +That is the whole thing. `scikit-build-core` runs CMake, and CMake resolves an ONNX +Runtime SDK matching the installed `onnxruntime` wheel, downloading it once into +`~/.cache/anytime-inference-planner/` if it is not already there. ## Match the ONNX Runtime version to the Python wheel -This is the single most important build detail. Building the worker against -1.20.1 while the Python side used the 1.26.0 wheel measured DistilBERT at -**98.9 ms inside the worker against 13.0 ms in the profiler**, a 7.6x gap. The -timing is taken around `session->Run()`, so this was not transport overhead; the -older release simply lacks the optimisations the newer one applies to this graph -on arm64. - -Because the profiler and the serving path disagreed, every service time the -planner used was wrong by almost an order of magnitude and no request met its -deadline. With matched versions the worker measures 14.05 ms, within 8% of the -in-process session, and the JSON transport costs 0.27 ms per request. - -`CMakeLists.txt` guards a related trap: `find_library` caches its result, so -reconfiguring with a different `ONNXRUNTIME_ROOT_PATH` would keep linking the -previously found library while compiling against the new headers. That mismatch -surfaces only at runtime, as `The requested API version [N] is not available`. -The cache entry is now invalidated when the root path changes, and the resolved -library is printed at configure time. - -## Protocol - -Handshake: the worker prints `ready` on stdout once every model has loaded. - -Request, one JSON object per line: - -```json -{"request_id": "r-1", - "variant": "distilbert_fp32", - "inputs": {"input_ids": {"shape": [1,128], "dtype": "int64", "data": ""}, - "attention_mask": {"shape": [1,128], "dtype": "int64", "data": ""}}} -``` - -Response: - -```json -{"request_id": "r-1", - "logits": {"shape": [1,2], "dtype": "float32", "data": ""}, - "latency_ms": 14.05} -``` - -Error response, for anything recoverable: - -```json -{"request_id": "r-1", "error": "unknown variant"} -``` - -The Python client raises `RuntimeError` on an error response. +This is the single most important build detail, and it is now enforced rather than +documented. + +Building the Stage 1 worker against 1.20.1 while the Python side used the 1.26.0 +wheel measured DistilBERT at **98.9 ms inside the worker against 13.0 ms in the +profiler**, a 7.6x gap. The timing was taken around `session->Run()`, so this was +not transport overhead; the older release simply lacks the optimisations the newer +one applies to this graph on arm64. Because the profiler and the serving path +disagreed, every service time the planner used was wrong by almost an order of +magnitude and no request met its deadline. Nothing crashed and no test failed. The +numbers were just false. + +With the extension, both copies of ONNX Runtime are loaded into the same process, +so this stopped being a performance question and became a correctness one. The +version is therefore never written down twice. It is read from the wheel the target +interpreter will import, and checked in three places: + +1. **Configure time.** The SDK version must equal the wheel version, or CMake + fails. +2. **Compile time.** `ORT_API_VERSION` is read out of the resolved headers and + asserted in `src/tensor.cpp`, so a header from a different include path breaks + the build rather than the run. +3. **Import time.** `load_extension()` compares + `anytime_runtime.onnxruntime_version()` against `onnxruntime.__version__`, + which covers an extension carried into an environment with a different wheel. + +A related consequence: CI no longer pins a version anywhere. It had been building +the C++ side against 1.26.0 while `pip` installed 1.28.0, since the dependency +floor is `>=1.16`. That was harmless only because the two never shared a process. + +## Cost of the transport it replaced + +Measured on DistilBERT at batch size one, 60 requests after warm-up: + +| Backend | Inference p50 | Wall p50 | Transport | +| --- | --- | --- | --- | +| `python` | 13.667 ms | 13.673 ms | 0.006 ms | +| `extension` | 13.609 ms | 13.628 ms | 0.018 ms | +| `subprocess` | 13.629 ms | 13.926 ms | 0.296 ms | + +All three agreed bitwise on the logits, and their inference times agreed within +0.4%. That agreement is what a matched version looks like, and it is the check that +was missing in Stage 1. + +Transport fell from 0.296 ms to 0.018 ms. Against 14 ms of inference that saving is +not why the change was made: a subprocess boundary rules out batching across +requests and sharing a KV cache between them, which is what the rest of Stage 2 +needs. The `subprocess` row is from the Stage 1 worker, which was removed once it +had served as the reference the extension was validated against. ## Behaviour -- **One inference at a time.** Run a pool from the Python side for concurrency. - Intra-op threads are set to 1 so N workers behave as N independent servers, - which is what makes the M/M/c queueing model in - [`planner.md`](planner.md) valid. -- **Per-request error isolation.** A malformed request, an unknown variant, or an - ONNX Runtime failure produces an error response and the worker continues. An - earlier version wrapped the whole request loop in one `try`, so a single bad - request terminated the worker and failed every subsequent request routed to it. -- **Input filtering.** Variants can declare different inputs. The worker keeps the - subset its graph declares and drops the rest, and errors if a declared input is - missing rather than running on a partial feed. -- **Supported dtypes.** Inputs `float32` and `int64`; outputs `float32`. -- **No external dependencies.** JSON parsing and base64 are implemented in - `main.cpp` so the worker needs only ONNX Runtime. +- **Tensors are borrowed, not copied.** Inputs point into the numpy buffer; + outputs are numpy views over ONNX Runtime's buffers, kept alive by a capsule + owning the `Ort::Value`. A response therefore holds runtime memory until it is + dropped, so copy `logits` if it needs to outlive the request. +- **The GIL is released during inference.** Without that a pool would serialise on + the interpreter lock, and the M/M/c model in [`planner.md`](planner.md) would + describe a machine that does not exist. +- **One engine per worker.** Intra-op threads are set to 1 so that N workers behave + as N independent servers, which is what makes that queueing model valid, and + matches how the service times in `configs/serving.yaml` were measured. +- **Input filtering.** Variants can declare different inputs, so callers pass the + union and each graph takes the subset it declares. A declared input that is + missing is an error rather than a run on a partial feed. +- **Error contract.** An unknown variant or a missing declared input raises + `RuntimeError`; a dtype the engine does not accept raises `ValueError`. Both + backends behave identically, which is what makes them comparable. +- **Supported dtypes.** Inputs and outputs: float32, float64, int32, int64, bool. + +## Backends + +`RuntimeClient` picks the extension when it imports and matches the wheel, and +otherwise warns and falls back to ONNX Runtime's own Python API. The fallback exists +so the control plane runs where the extension has not been built; it is not the +serving path, and a measurement taken through it does not describe one. + +`backend="extension"` or `backend="python"` pins the choice. Tests pin it so a +parity failure cannot hide behind a silent fallback. ## Tests -`tests/test_runtime_subprocess.py` exercises the real binary: round-trip -fidelity, the base64 codec across all three padding cases, variant routing, -concurrent dispatch through a pool, and error isolation. Marked `needs_runtime` -and skipped when the binary is absent. - ```bash -pytest -q tests/test_runtime_subprocess.py +pytest -q tests/test_runtime_engine.py +ANYTIME_REQUIRE_BACKENDS=extension,python pytest -q tests/test_runtime_engine.py ``` -## Planned replacement - -The subprocess boundary costs a base64 encode, a JSON parse, and a copy in each -direction. At 0.27 ms against 14 ms of inference that is acceptable, but it rules -out batching across requests and sharing a KV cache between them. The intended -replacement is a pybind11 extension linking the same ONNX Runtime as the Python -wheel, exposing zero-copy numpy tensors and a C++ scheduler. +`tests/test_runtime_engine.py` asserts the backends agree bitwise on a graph with +real arithmetic in it, that outputs are views rather than copies, that a strided +input is made contiguous rather than misread, and that an unsupported dtype is +refused. `ANYTIME_REQUIRE_BACKENDS` turns a missing backend from a skip into a +failure, so the comparison cannot decay into one backend checked against itself. diff --git a/pyproject.toml b/pyproject.toml index 9990213..ea2cd14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,19 @@ [build-system] -requires = ["hatchling>=1.21"] -build-backend = "hatchling.build" +# The runtime is a compiled extension rather than a subprocess, so building the +# package means running CMake. +# +# onnxruntime appears here as a fallback, not as the authority. The extension has +# to link the version the *target* environment imports at runtime, so +# runtime/cmake/ResolveOnnxRuntime.cmake reads that environment first and only +# falls back to this copy when installing into a clean environment, where no wheel +# exists yet at build time. The specifier matches the runtime dependency so the two +# resolve alike, and the equality is re-checked when the extension is imported. +requires = [ + "scikit-build-core>=0.10", + "pybind11>=3.0", + "onnxruntime>=1.16", +] +build-backend = "scikit_build_core.build" [project] name = "anytime-serving" @@ -55,8 +68,15 @@ dev = [ "ruff>=0.4", ] -[tool.hatch.build.targets.wheel] -packages = ["src/anytime_serving"] +[tool.scikit-build] +cmake.source-dir = "runtime" +cmake.version = ">=3.20" +wheel.packages = ["src/anytime_serving"] +# Keep the configure and build output visible. A version mismatch against the +# onnxruntime wheel is reported as a configure-time error, and that message is +# worth reading rather than burying. +build.verbose = true +logging.level = "INFO" [tool.ruff] line-length = 100 diff --git a/pytest.ini b/pytest.ini index 49e44b4..f5a73b2 100644 --- a/pytest.ini +++ b/pytest.ini @@ -8,7 +8,7 @@ addopts = markers = slow: takes more than a few seconds; deselect with -m "not slow" needs_torch: requires torch/torchvision/transformers - needs_runtime: requires the compiled C++ runtime binary + needs_runtime: requires a compiled native runtime (the anytime_runtime extension or the subprocess worker) filterwarnings = error::DeprecationWarning:anytime_serving.* default diff --git a/runtime/CMakeLists.txt b/runtime/CMakeLists.txt new file mode 100644 index 0000000..1428355 --- /dev/null +++ b/runtime/CMakeLists.txt @@ -0,0 +1,56 @@ +cmake_minimum_required(VERSION 3.20) +project(anytime_runtime CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE Release) +endif() + +list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") + +find_package(Python REQUIRED COMPONENTS Interpreter Development.Module) +find_package(pybind11 CONFIG REQUIRED) + +# Resolves an SDK whose version equals the onnxruntime wheel's. See the module +# for why that equality is enforced at configure time rather than documented. +include(ResolveOnnxRuntime) +anytime_resolve_onnxruntime() + +add_library(anytime_core STATIC + src/tensor.cpp + src/engine.cpp) + +target_include_directories(anytime_core + PUBLIC + "${CMAKE_CURRENT_SOURCE_DIR}/include" + "${ANYTIME_ORT_ROOT}/include") + +# Asserted in tensor.cpp against the value ORT_API_VERSION actually has, so a +# header arriving from some other include path fails the build rather than the +# run. +target_compile_definitions(anytime_core + PUBLIC ANYTIME_EXPECTED_ORT_API_VERSION=${ANYTIME_ORT_API_VERSION}) + +target_link_libraries(anytime_core PUBLIC "${ANYTIME_ORT_LIBRARY}") + +pybind11_add_module(anytime_runtime MODULE bindings/module.cpp) +target_link_libraries(anytime_runtime PRIVATE anytime_core) + +# The extension loads libonnxruntime from the resolved SDK. The path is absolute +# because the SDK lives outside the wheel; ONNXRUNTIME_ROOT_PATH is how a +# relocatable build points somewhere else. +cmake_path(GET ANYTIME_ORT_LIBRARY PARENT_PATH ANYTIME_ORT_LIBDIR) +if(APPLE) + set_target_properties(anytime_runtime PROPERTIES + BUILD_RPATH "${ANYTIME_ORT_LIBDIR}" + INSTALL_RPATH "${ANYTIME_ORT_LIBDIR}") +elseif(UNIX) + set_target_properties(anytime_runtime PROPERTIES + BUILD_RPATH "${ANYTIME_ORT_LIBDIR}" + INSTALL_RPATH "${ANYTIME_ORT_LIBDIR}") +endif() + +install(TARGETS anytime_runtime DESTINATION .) diff --git a/runtime/README.md b/runtime/README.md new file mode 100644 index 0000000..6ad9ea6 --- /dev/null +++ b/runtime/README.md @@ -0,0 +1,99 @@ +# anytime_runtime (C++ extension) + +In-process ONNX Runtime engine, exposed to Python through pybind11. Replaces the +Stage 1 subprocess worker: the transport is a function call, so there is no JSON +framing, no base64, and no copy in either direction. + +## Build + +The extension is built by installing the package. `scikit-build-core` drives +CMake, so there is no separate build step: + +```bash +pip install -e . +``` + +`onnxruntime` must already be installed in the target environment, because the +extension is linked against the version that environment will import. + +## The version is derived, never written down twice + +`cmake/ResolveOnnxRuntime.cmake` reads `onnxruntime.__version__` from the target +interpreter and resolves an SDK to match: + +- `ONNXRUNTIME_ROOT_PATH`, if set, must contain a `VERSION_NUMBER` equal to the + wheel version. Configuration fails otherwise. +- Otherwise the matching release archive is downloaded once into + `~/.cache/anytime-inference-planner/` and reused. `ANYTIME_ORT_CACHE` overrides + the location. + +This matters more than it looks. Stage 1 built the worker against 1.20.1 while the +Python side used the 1.26.0 wheel and measured DistilBERT at **98.9 ms against +13.0 ms** inside `session->Run()`, a 7.6x gap that made every admission decision +wrong. Nothing failed; the numbers were just false. With the extension both copies +of the library are loaded into one process, so the equality is enforced in three +places: + +1. **Configure time.** The SDK version must equal the wheel version. +2. **Compile time.** `ORT_API_VERSION` is read out of the resolved headers and + asserted in `src/tensor.cpp`, so a header from another include path fails the + build rather than the run. +3. **Import time.** `load_extension()` in `serving/onnx_runtime.py` compares + `anytime_runtime.onnxruntime_version()` against `onnxruntime.__version__` and + raises on a mismatch, which covers an extension carried into a different + environment. + +The probe deliberately clears `PYTHONPATH`. pip builds in an isolated environment +layered onto the target interpreter through that variable, and reading its +onnxruntime instead of the target's would reintroduce the mismatch through the +build system. + +## Layout + +```text +include/anytime/tensor.hpp element types, borrowed tensor views +include/anytime/engine.hpp sessions, input filtering, one run +src/ implementations +bindings/module.cpp the pybind11 module +cmake/ ONNX Runtime resolution +``` + +Headers arrive as the work that needs them lands, rather than as empty +placeholders: the scheduler, batch assembly, and KV cache are not here yet. + +## Behaviour + +- **Tensors are borrowed, not copied.** Inputs point into the numpy buffer, whose + references the bindings hold for the whole call. Outputs are numpy views over + ONNX Runtime's own buffers, kept alive by a capsule owning the `Ort::Value`. A + response therefore holds runtime memory until it is dropped. +- **A non-contiguous input is made contiguous.** ONNX Runtime reads the buffer + directly, so a strided array would otherwise be misread. +- **The GIL is released around inference.** Without that a pool would serialise on + the interpreter lock, and the M/M/c model the admission controller uses would + describe a machine that does not exist. +- **One session set per engine.** The Python pool holds one engine per slot, so N + workers stay N independent single-threaded servers. Intra-op and inter-op + threads default to one each, matching how the service times in + `configs/serving.yaml` were measured. +- **Input filtering.** Variants of one task can declare different inputs, so + callers pass the union and each graph takes the subset it declares. A declared + input that is missing raises instead, since running on a partial feed would + silently produce wrong output. +- **Error contract.** Anything meaning "the runtime could not serve this request" + (an unknown variant, a missing input) raises `RuntimeError`. A malformed + argument (a dtype the engine does not accept) raises `ValueError`. +- **Supported input dtypes.** float32, float64, int32, int64, bool. Outputs are + mapped back from the same set. + +## Tests + +`tests/test_runtime_engine.py` compares the extension against the two backends it +replaces on a graph with real arithmetic in it, and asserts they agree bitwise. +`ANYTIME_REQUIRE_BACKENDS` turns a missing backend from a skip into a failure, so +the comparison cannot decay into a backend checked against itself. + +```bash +pytest -q tests/test_runtime_engine.py +ANYTIME_REQUIRE_BACKENDS=extension,python,subprocess pytest -q tests/test_runtime_engine.py +``` diff --git a/runtime/bindings/module.cpp b/runtime/bindings/module.cpp new file mode 100644 index 0000000..8a5262d --- /dev/null +++ b/runtime/bindings/module.cpp @@ -0,0 +1,179 @@ +// pybind11 bindings for the in-process engine. +// +// Two things here are load-bearing beyond wiring: +// +// Tensors cross the boundary without copying. Inputs are borrowed straight from +// the numpy buffer, so the references are held on the stack for the whole call. +// Outputs are numpy views over the buffers ONNX Runtime allocated, kept alive by +// a capsule that owns the Ort::Value. +// +// The GIL is released around Session::Run. Without that, a pool of workers would +// serialise on the interpreter lock and the M/M/c model the admission controller +// relies on would describe a machine that does not exist. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "anytime/engine.hpp" +#include "anytime/tensor.hpp" + +namespace py = pybind11; + +namespace { + +anytime::DType dtype_from_numpy(const py::array& array, const std::string& name) { + const py::dtype dtype = array.dtype(); + if (dtype.is(py::dtype::of())) return anytime::DType::Float32; + if (dtype.is(py::dtype::of())) return anytime::DType::Float64; + if (dtype.is(py::dtype::of())) return anytime::DType::Int32; + if (dtype.is(py::dtype::of())) return anytime::DType::Int64; + if (dtype.is(py::dtype::of())) return anytime::DType::Bool; + throw std::invalid_argument( + "feed '" + name + "' has dtype " + + py::str(dtype).cast() + + ", which the engine does not accept. Supported: float32, float64, int32, " + "int64, bool. Cast before submitting rather than relying on an implicit " + "conversion."); +} + +py::dtype numpy_dtype_from_ort(ONNXTensorElementDataType type) { + switch (type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return py::dtype::of(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: return py::dtype::of(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: return py::dtype::of(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: return py::dtype::of(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: return py::dtype::of(); + default: + throw std::runtime_error( + "graph produced an output of type " + anytime::ort_type_name(type) + + ", which the bindings cannot map to numpy"); + } +} + +// Hands ONNX Runtime's output buffer to numpy without copying it. The capsule +// owns the Ort::Value, so the buffer outlives this call for exactly as long as +// the returned array is referenced. +py::array wrap_output(Ort::Value&& value) { + const auto info = value.GetTensorTypeAndShapeInfo(); + const ONNXTensorElementDataType element_type = info.GetElementType(); + const py::dtype dtype = numpy_dtype_from_ort(element_type); + const std::vector raw_shape = info.GetShape(); + + std::vector shape(raw_shape.begin(), raw_shape.end()); + std::vector strides(shape.size()); + py::ssize_t stride = static_cast(dtype.itemsize()); + for (std::size_t i = shape.size(); i-- > 0;) { + strides[i] = stride; + stride *= shape[i]; + } + + if (info.GetElementCount() == 0) { + // Nothing to share. An owning empty array avoids handing numpy a base + // object for a buffer that may not exist. + return py::array(dtype, shape, strides); + } + + auto* owned = new Ort::Value(std::move(value)); + py::capsule owner(owned, [](void* pointer) { + delete static_cast(pointer); + }); + return py::array(dtype, shape, strides, owned->GetTensorMutableRawData(), owner); +} + +py::tuple run_engine(anytime::Engine& engine, const std::string& variant, + const py::dict& feeds) { + // Built while the GIL is held. `retained` keeps every numpy buffer alive + // across the release below; the views point into those buffers. + std::vector retained; + std::vector> views; + retained.reserve(feeds.size()); + views.reserve(feeds.size()); + + for (const auto item : feeds) { + const auto name = py::cast(item.first); + // ONNX Runtime reads the buffer directly, so it has to be C-contiguous. + // ensure() returns the same object when it already is, and converts only + // when it is not, rather than misreading a strided array. + py::array array = py::array::ensure(item.second, py::array::c_style); + if (!array) { + throw std::invalid_argument("feed '" + name + "' is not array-like"); + } + + anytime::TensorView view; + view.dtype = dtype_from_numpy(array, name); + view.shape.assign(array.shape(), array.shape() + array.ndim()); + view.data = array.data(); + + retained.push_back(std::move(array)); + views.emplace_back(name, std::move(view)); + } + + anytime::RunResult result; + { + py::gil_scoped_release release; + result = engine.run(variant, views); + } + + py::list outputs; + for (auto& value : result.outputs) { + outputs.append(wrap_output(std::move(value))); + } + return py::make_tuple(std::move(outputs), result.latency_ms); +} + +} // namespace + +PYBIND11_MODULE(anytime_runtime, module) { + module.doc() = + "In-process ONNX Runtime engine for the Anytime Inference Planner.\n\n" + "Tensors cross the boundary as borrowed buffers rather than copies, and " + "the GIL is released around inference so a pool of workers runs " + "concurrently."; + + module.def("onnxruntime_version", []() { + return std::string(OrtGetApiBase()->GetVersionString()); + }, "Version of the ONNX Runtime this extension is linked against.\n\n" + "Must equal onnxruntime.__version__: both libraries are loaded into this " + "process, and Stage 1 measured a 7.6x difference in inference time from a " + "mismatch between them."); + + module.def("ort_api_version", []() { return ORT_API_VERSION; }, + "ORT_API_VERSION of the headers this extension was compiled " + "against."); + + py::class_(module, "Engine", + "Holds one ONNX Runtime session per named variant.\n\n" + "Thread counts default to one of each so that N pooled engines behave as " + "N independent single-threaded servers, which is what makes the queueing " + "model in the admission controller valid.") + .def(py::init([](const std::vector>& models, + int intra_op_threads, int inter_op_threads) { + return std::make_unique(models, intra_op_threads, + inter_op_threads); + }), + py::arg("models"), py::arg("intra_op_threads") = 1, + py::arg("inter_op_threads") = 1, + "models is a sequence of (variant, onnx_path) pairs.") + .def("run", &run_engine, py::arg("variant"), py::arg("feeds"), + "Run one variant.\n\n" + "Returns (outputs, latency_ms), where outputs are numpy views over " + "ONNX Runtime's buffers in the order the graph declares them, and " + "latency_ms is the time spent inside Session::Run.\n\n" + "Feeds the graph does not declare are dropped, since variants of one " + "task can declare different inputs and callers pass the union. A " + "declared input that is missing raises instead: running on a partial " + "feed would silently produce wrong output.") + .def_property_readonly("variants", &anytime::Engine::variants, + "Loaded variant names.") + .def("input_names", &anytime::Engine::input_names, py::arg("variant"), + "Input names the variant's graph declares.") + .def("output_names", &anytime::Engine::output_names, py::arg("variant"), + "Output names the variant's graph declares, in run order."); +} diff --git a/runtime/cmake/ResolveOnnxRuntime.cmake b/runtime/cmake/ResolveOnnxRuntime.cmake new file mode 100644 index 0000000..be2bcb1 --- /dev/null +++ b/runtime/cmake/ResolveOnnxRuntime.cmake @@ -0,0 +1,264 @@ +# Resolve an ONNX Runtime SDK that matches the installed onnxruntime wheel. +# +# The extension and the Python wheel run in one process, so their ONNX Runtime +# versions cannot be allowed to differ. Stage 1 built the worker against 1.20.1 +# while the wheel was 1.26.0 and measured DistilBERT at 98.9 ms against 13.0 ms +# inside session->Run(); every service time the planner used was wrong by almost +# an order of magnitude. Under pybind11 a mismatch is worse than a performance +# cliff, because both copies of the library are loaded at once. +# +# The version is therefore never written down twice. It is read from the wheel +# that the target interpreter will import at runtime, and everything else is +# derived from it: +# +# 1. ONNXRUNTIME_ROOT_PATH, if set, must contain a VERSION_NUMBER equal to the +# wheel version. Configuration fails otherwise. +# 2. Otherwise the matching release archive is downloaded once into a cache +# directory and reused. +# +# Sets, in the parent scope: +# ANYTIME_ORT_ROOT directory holding include/ and lib/ +# ANYTIME_ORT_VERSION version string, as reported by the wheel +# ANYTIME_ORT_LIBRARY full path to the shared library +# ANYTIME_ORT_API_VERSION ORT_API_VERSION as written in the resolved headers + +include_guard(GLOBAL) + +function(_anytime_ort_wheel_version out_var) + if(NOT Python_EXECUTABLE) + message(FATAL_ERROR + "Python_EXECUTABLE is not set, so the onnxruntime wheel version cannot " + "be read. Build through pip (which supplies it via scikit-build-core) " + "or pass -DPython_EXECUTABLE=/path/to/python.") + endif() + + # Probed twice, in order of authority. + # + # First with PYTHONPATH cleared. pip builds in an isolated environment layered + # onto the target interpreter through that variable, and the overlay carries + # its own onnxruntime, which can be a different version from the one the target + # environment imports at runtime. Reading the overlay when the target already + # has a wheel would reintroduce the mismatch this module exists to prevent, + # through the build system instead of through the developer. + execute_process( + COMMAND "${CMAKE_COMMAND}" -E env --unset=PYTHONPATH + "${Python_EXECUTABLE}" -c + "import onnxruntime; print(onnxruntime.__version__)" + OUTPUT_VARIABLE version + ERROR_VARIABLE target_error + RESULT_VARIABLE probe_status + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(probe_status EQUAL 0) + message(STATUS "onnxruntime resolved from the target environment") + set(${out_var} "${version}" PARENT_SCOPE) + return() + endif() + + # Then with the build environment left intact. On a first install into a clean + # environment there is no wheel to read yet: pip installs build requirements, + # builds, and only then installs runtime dependencies. The build requirement + # pins the same specifier as the runtime dependency, so the version resolved + # here is the one pip is about to install. load_extension() re-checks the + # equality at import time, which is what catches the case where it is not. + execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import onnxruntime; print(onnxruntime.__version__)" + OUTPUT_VARIABLE version + ERROR_VARIABLE build_error + RESULT_VARIABLE probe_status + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(probe_status EQUAL 0) + message(STATUS + "onnxruntime is not installed in the target environment yet; resolved " + "${version} from the build environment instead. The version is " + "re-checked when anytime_runtime is imported.") + set(${out_var} "${version}" PARENT_SCOPE) + return() + endif() + + message(FATAL_ERROR + "Could not import onnxruntime with ${Python_EXECUTABLE}, with or without " + "the build environment on the path. The extension is linked against the " + "same ONNX Runtime that interpreter imports at runtime, so a wheel has to " + "be resolvable:\n" + " ${Python_EXECUTABLE} -m pip install onnxruntime\n" + "target environment:\n${target_error}\n" + "build environment:\n${build_error}") +endfunction() + +# Release archives are published per platform under a stable naming scheme. +function(_anytime_ort_archive_name version out_name out_extension) + if(APPLE) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64|aarch64") + set(platform "osx-arm64") + else() + set(platform "osx-x86_64") + endif() + set(extension "tgz") + elseif(UNIX) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64") + set(platform "linux-aarch64") + else() + set(platform "linux-x64") + endif() + set(extension "tgz") + else() + # Windows archives are zips with a different internal layout. Rather than + # guess, ask for an explicit SDK. + message(FATAL_ERROR + "Automatic ONNX Runtime download is implemented for macOS and Linux " + "only. Download the ${version} release for this platform from " + "https://github.com/microsoft/onnxruntime/releases and pass " + "-DONNXRUNTIME_ROOT_PATH=/path/to/onnxruntime.") + endif() + + set(${out_name} "onnxruntime-${platform}-${version}" PARENT_SCOPE) + set(${out_extension} "${extension}" PARENT_SCOPE) +endfunction() + +# The unpacked tree is the directory holding include/. Some releases nest a +# versioned directory inside the archive and some do not, so locate it rather +# than assuming either layout. +function(_anytime_ort_find_root search_dir out_var) + if(EXISTS "${search_dir}/include/onnxruntime_c_api.h") + set(${out_var} "${search_dir}" PARENT_SCOPE) + return() + endif() + file(GLOB candidates LIST_DIRECTORIES true "${search_dir}/*") + foreach(candidate IN LISTS candidates) + if(IS_DIRECTORY "${candidate}" AND + EXISTS "${candidate}/include/onnxruntime_c_api.h") + set(${out_var} "${candidate}" PARENT_SCOPE) + return() + endif() + endforeach() + set(${out_var} "" PARENT_SCOPE) +endfunction() + +function(_anytime_ort_download version out_root) + _anytime_ort_archive_name("${version}" archive_name archive_extension) + + # A durable cache, so the SDK is fetched once per version per machine rather + # than once per build tree. Stage 1 kept it in a temporary directory and lost + # it between sessions. + if(DEFINED ENV{ANYTIME_ORT_CACHE}) + set(cache_root "$ENV{ANYTIME_ORT_CACHE}") + elseif(DEFINED ENV{XDG_CACHE_HOME}) + set(cache_root "$ENV{XDG_CACHE_HOME}/anytime-inference-planner") + else() + set(cache_root "$ENV{HOME}/.cache/anytime-inference-planner") + endif() + set(unpacked "${cache_root}/${archive_name}") + + _anytime_ort_find_root("${unpacked}" cached_root) + if(cached_root) + message(STATUS "ONNX Runtime ${version}: using cached SDK at ${cached_root}") + set(${out_root} "${cached_root}" PARENT_SCOPE) + return() + endif() + + set(url "https://github.com/microsoft/onnxruntime/releases/download/v${version}/${archive_name}.${archive_extension}") + set(archive "${cache_root}/${archive_name}.${archive_extension}") + message(STATUS "ONNX Runtime ${version}: downloading ${url}") + + file(MAKE_DIRECTORY "${cache_root}") + file(DOWNLOAD "${url}" "${archive}" STATUS download_status SHOW_PROGRESS) + list(GET download_status 0 download_code) + if(NOT download_code EQUAL 0) + list(GET download_status 1 download_message) + file(REMOVE "${archive}") + message(FATAL_ERROR + "Failed to download ONNX Runtime ${version} from ${url}: " + "${download_message}\nDownload it manually and pass " + "-DONNXRUNTIME_ROOT_PATH=/path/to/onnxruntime instead.") + endif() + + file(MAKE_DIRECTORY "${unpacked}") + file(ARCHIVE_EXTRACT INPUT "${archive}" DESTINATION "${unpacked}") + file(REMOVE "${archive}") + + _anytime_ort_find_root("${unpacked}" downloaded_root) + if(NOT downloaded_root) + message(FATAL_ERROR + "Downloaded ONNX Runtime ${version} but found no include/ directory " + "under ${unpacked}.") + endif() + set(${out_root} "${downloaded_root}" PARENT_SCOPE) +endfunction() + +function(anytime_resolve_onnxruntime) + _anytime_ort_wheel_version(wheel_version) + message(STATUS "onnxruntime wheel reports version ${wheel_version}") + + if(NOT ONNXRUNTIME_ROOT_PATH AND DEFINED ENV{ONNXRUNTIME_ROOT_PATH}) + set(ONNXRUNTIME_ROOT_PATH "$ENV{ONNXRUNTIME_ROOT_PATH}") + endif() + + if(ONNXRUNTIME_ROOT_PATH) + _anytime_ort_find_root("${ONNXRUNTIME_ROOT_PATH}" root) + if(NOT root) + message(FATAL_ERROR + "ONNXRUNTIME_ROOT_PATH=${ONNXRUNTIME_ROOT_PATH} contains no " + "include/onnxruntime_c_api.h.") + endif() + + if(NOT EXISTS "${root}/VERSION_NUMBER") + message(FATAL_ERROR + "${root} has no VERSION_NUMBER, so its version cannot be checked " + "against the ${wheel_version} wheel. Point ONNXRUNTIME_ROOT_PATH at " + "an unpacked release archive, or unset it to download one.") + endif() + file(READ "${root}/VERSION_NUMBER" sdk_version) + string(STRIP "${sdk_version}" sdk_version) + + if(NOT sdk_version VERSION_EQUAL wheel_version) + message(FATAL_ERROR + "ONNX Runtime version mismatch. The SDK at ${root} is " + "${sdk_version}; the onnxruntime wheel that will be imported in the " + "same process is ${wheel_version}. These must be equal: the two " + "libraries are loaded together, and Stage 1 measured a 7.6x " + "difference in session->Run() from exactly this mismatch. Unset " + "ONNXRUNTIME_ROOT_PATH to download the matching SDK.") + endif() + message(STATUS "ONNX Runtime ${sdk_version}: using SDK at ${root}") + else() + _anytime_ort_download("${wheel_version}" root) + endif() + + # find_library caches its result, so a changed root would otherwise keep + # linking the previously resolved library against the new headers. That + # mismatch surfaces only at runtime, as "The requested API version [N] is not + # available". + if(DEFINED CACHE{ANYTIME_ORT_LIBRARY} AND + NOT ANYTIME_ORT_LIBRARY MATCHES "^${root}/") + message(STATUS "ONNX Runtime root changed; re-resolving the library") + unset(ANYTIME_ORT_LIBRARY CACHE) + endif() + find_library(ANYTIME_ORT_LIBRARY + NAMES onnxruntime + PATHS "${root}/lib" + NO_DEFAULT_PATH + REQUIRED) + + # Read ORT_API_VERSION out of the headers that were just resolved and assert + # it in the translation units. This catches a stray onnxruntime_c_api.h + # arriving from some other include path, without assuming anything about how + # release versions map onto API versions. + file(STRINGS "${root}/include/onnxruntime_c_api.h" api_version_line + REGEX "^#define ORT_API_VERSION[ \t]+[0-9]+") + if(NOT api_version_line) + message(FATAL_ERROR + "Could not read ORT_API_VERSION from " + "${root}/include/onnxruntime_c_api.h.") + endif() + string(REGEX MATCH "[0-9]+" api_version "${api_version_line}") + + message(STATUS "ONNX Runtime library: ${ANYTIME_ORT_LIBRARY}") + message(STATUS "ONNX Runtime API version: ${api_version}") + + set(ANYTIME_ORT_ROOT "${root}" PARENT_SCOPE) + set(ANYTIME_ORT_VERSION "${wheel_version}" PARENT_SCOPE) + set(ANYTIME_ORT_API_VERSION "${api_version}" PARENT_SCOPE) +endfunction() diff --git a/runtime/include/anytime/engine.hpp b/runtime/include/anytime/engine.hpp new file mode 100644 index 0000000..3258281 --- /dev/null +++ b/runtime/include/anytime/engine.hpp @@ -0,0 +1,88 @@ +// In-process inference engine. +// +// Owns one ONNX Runtime session per named variant and runs them synchronously. +// This replaces the line-delimited JSON worker from Stage 1: the transport is now +// a function call, so there is no framing, no base64, and no copy in either +// direction. +// +// Threading in this phase mirrors what the subprocess pool did. One Engine holds +// its own sessions, and the Python pool holds one Engine per slot, so N workers +// remain N independent single-threaded servers and the M/M/c model the admission +// controller uses stays valid. Sharing one session across the pool would save +// memory and is what the scheduler will do once it lands, but it would change the +// concurrency model in the same commit that introduces the engine. + +#ifndef ANYTIME_ENGINE_HPP +#define ANYTIME_ENGINE_HPP + +#include + +#include +#include +#include +#include +#include + +#include "anytime/tensor.hpp" + +namespace anytime { + +// One loaded graph, with the input and output names it declares. +class Model { +public: + Model(Ort::Env& env, const std::string& path, int intra_op_threads, + int inter_op_threads); + + const std::vector& input_names() const { return input_names_; } + const std::vector& output_names() const { return output_names_; } + bool declares_input(const std::string& name) const; + Ort::Session& session() { return *session_; } + +private: + std::unique_ptr session_; + std::vector input_names_; + std::vector output_names_; +}; + +// Outputs of one run, in the order the graph declares them, plus the time spent +// inside Session::Run. The values own their buffers; the bindings hand those +// buffers to numpy without copying. +struct RunResult { + std::vector outputs; + double latency_ms = 0.0; +}; + +class Engine { +public: + // Each entry maps a variant name to an ONNX graph path. Threads default to + // one of each so a pooled engine behaves as a single server, matching how + // the service times in configs/serving.yaml were measured. + explicit Engine(const std::vector>& models, + int intra_op_threads = 1, int inter_op_threads = 1); + + // Runs `variant` over `feeds`. + // + // Feeds a graph does not declare are dropped rather than rejected: variants + // of one task can declare different inputs (a DistilBERT graph takes + // input_ids and attention_mask, a BERT graph also takes token_type_ids), so + // callers pass the union. A declared input that is missing is an error, since + // running on a partial feed would silently produce wrong logits. + RunResult run(const std::string& variant, + const std::vector>& feeds); + + std::vector variants() const; + const std::vector& input_names(const std::string& variant) const; + const std::vector& output_names(const std::string& variant) const; + +private: + Model& lookup(const std::string& variant); + const Model& lookup(const std::string& variant) const; + + Ort::Env env_; + Ort::MemoryInfo memory_; + std::map> models_; +}; + +} // namespace anytime + +#endif // ANYTIME_ENGINE_HPP diff --git a/runtime/include/anytime/tensor.hpp b/runtime/include/anytime/tensor.hpp new file mode 100644 index 0000000..2f4c774 --- /dev/null +++ b/runtime/include/anytime/tensor.hpp @@ -0,0 +1,58 @@ +// Element types and borrowed tensor views. +// +// The engine never owns input buffers. Callers hand it a pointer, a shape, and a +// dtype; ONNX Runtime reads through that pointer during Run. This is what makes +// the numpy boundary free of copies, and it means the caller has to keep the +// buffer alive for the duration of the call. The bindings do that by holding the +// numpy references on the stack across the whole run. + +#ifndef ANYTIME_TENSOR_HPP +#define ANYTIME_TENSOR_HPP + +#include + +#include +#include +#include +#include + +namespace anytime { + +// Element types accepted on graph inputs. Deliberately narrow: exported graphs +// declare int64 token ids and masks and float32 activations, and an unsupported +// dtype has to fail loudly rather than be reinterpreted. Outputs are mapped back +// from whatever the graph produces, which is a wider set. +enum class DType { + Float32, + Float64, + Int32, + Int64, + Bool, +}; + +std::size_t dtype_size(DType dtype); +const char* dtype_name(DType dtype); +ONNXTensorElementDataType dtype_to_ort(DType dtype); +DType dtype_from_ort(ONNXTensorElementDataType type); + +// Name of an ONNX element type, for error messages about types the engine does +// not accept. +std::string ort_type_name(ONNXTensorElementDataType type); + +// A borrowed, C-contiguous buffer. Non-owning by design; see the file comment. +struct TensorView { + const void* data = nullptr; + std::vector shape; + DType dtype = DType::Float32; + + std::size_t element_count() const; + std::size_t byte_count() const; +}; + +// Wraps a borrowed buffer as an Ort::Value. No copy: the returned value points +// into view.data, which must outlive it. +Ort::Value borrow_as_tensor(const TensorView& view, const OrtMemoryInfo* memory); + +} // namespace anytime + +#endif // ANYTIME_TENSOR_HPP diff --git a/runtime/src/engine.cpp b/runtime/src/engine.cpp new file mode 100644 index 0000000..dfd7f03 --- /dev/null +++ b/runtime/src/engine.cpp @@ -0,0 +1,158 @@ +#include "anytime/engine.hpp" + +#include +#include +#include +#include + +namespace anytime { +namespace { + +std::string join(const std::vector& parts) { + std::ostringstream out; + for (std::size_t i = 0; i < parts.size(); ++i) { + if (i) out << ", "; + out << parts[i]; + } + return out.str(); +} + +} // namespace + +Model::Model(Ort::Env& env, const std::string& path, int intra_op_threads, + int inter_op_threads) { + Ort::SessionOptions options; + options.SetIntraOpNumThreads(intra_op_threads); + options.SetInterOpNumThreads(inter_op_threads); + options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); + + session_ = std::make_unique(env, path.c_str(), options); + + Ort::AllocatorWithDefaultOptions allocator; + const std::size_t input_count = session_->GetInputCount(); + input_names_.reserve(input_count); + for (std::size_t i = 0; i < input_count; ++i) { + auto name = session_->GetInputNameAllocated(i, allocator); + input_names_.emplace_back(name.get()); + } + const std::size_t output_count = session_->GetOutputCount(); + output_names_.reserve(output_count); + for (std::size_t i = 0; i < output_count; ++i) { + auto name = session_->GetOutputNameAllocated(i, allocator); + output_names_.emplace_back(name.get()); + } +} + +bool Model::declares_input(const std::string& name) const { + return std::find(input_names_.begin(), input_names_.end(), name) != + input_names_.end(); +} + +Engine::Engine(const std::vector>& models, + int intra_op_threads, int inter_op_threads) + : env_(ORT_LOGGING_LEVEL_WARNING, "anytime_runtime"), + memory_(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault)) { + if (models.empty()) { + throw std::invalid_argument("Engine needs at least one model"); + } + for (const auto& [variant, path] : models) { + if (models_.count(variant)) { + throw std::invalid_argument("duplicate variant: " + variant); + } + models_[variant] = std::make_unique(env_, path, intra_op_threads, + inter_op_threads); + } +} + +Model& Engine::lookup(const std::string& variant) { + const auto it = models_.find(variant); + if (it == models_.end()) { + // runtime_error rather than out_of_range: pybind11 maps it to + // RuntimeError, which is the contract every backend shares for "the + // runtime could not serve this request". Malformed arguments (a dtype the + // engine does not accept) stay invalid_argument and surface as ValueError. + throw std::runtime_error("unknown variant: " + variant + " (loaded: " + + join(variants()) + ")"); + } + return *it->second; +} + +const Model& Engine::lookup(const std::string& variant) const { + const auto it = models_.find(variant); + if (it == models_.end()) { + // runtime_error rather than out_of_range: pybind11 maps it to + // RuntimeError, which is the contract every backend shares for "the + // runtime could not serve this request". Malformed arguments (a dtype the + // engine does not accept) stay invalid_argument and surface as ValueError. + throw std::runtime_error("unknown variant: " + variant + " (loaded: " + + join(variants()) + ")"); + } + return *it->second; +} + +std::vector Engine::variants() const { + std::vector names; + names.reserve(models_.size()); + for (const auto& [variant, model] : models_) { + names.push_back(variant); + } + return names; +} + +const std::vector& Engine::input_names(const std::string& variant) const { + return lookup(variant).input_names(); +} + +const std::vector& Engine::output_names(const std::string& variant) const { + return lookup(variant).output_names(); +} + +RunResult Engine::run(const std::string& variant, + const std::vector>& feeds) { + Model& model = lookup(variant); + + std::vector input_names; + std::vector inputs; + input_names.reserve(feeds.size()); + inputs.reserve(feeds.size()); + + for (const auto& [name, view] : feeds) { + if (!model.declares_input(name)) continue; + input_names.push_back(name.c_str()); + inputs.push_back(borrow_as_tensor(view, memory_)); + } + + if (inputs.size() != model.input_names().size()) { + std::vector missing; + for (const auto& declared : model.input_names()) { + const bool supplied = std::any_of( + feeds.begin(), feeds.end(), + [&declared](const auto& feed) { return feed.first == declared; }); + if (!supplied) missing.push_back(declared); + } + throw std::runtime_error("variant " + variant + " is missing input(s): " + + join(missing) + " (declares: " + + join(model.input_names()) + ")"); + } + + std::vector output_names; + output_names.reserve(model.output_names().size()); + for (const auto& name : model.output_names()) { + output_names.push_back(name.c_str()); + } + + const auto start = std::chrono::steady_clock::now(); + std::vector outputs = model.session().Run( + Ort::RunOptions{nullptr}, + input_names.data(), inputs.data(), inputs.size(), + output_names.data(), output_names.size()); + const auto end = std::chrono::steady_clock::now(); + + RunResult result; + result.outputs = std::move(outputs); + result.latency_ms = + std::chrono::duration(end - start).count(); + return result; +} + +} // namespace anytime diff --git a/runtime/src/tensor.cpp b/runtime/src/tensor.cpp new file mode 100644 index 0000000..e7fad2c --- /dev/null +++ b/runtime/src/tensor.cpp @@ -0,0 +1,115 @@ +#include "anytime/tensor.hpp" + +#include + +// The headers this compiles against must be the ones CMake resolved from the SDK +// that matches the onnxruntime wheel. A stray onnxruntime_c_api.h arriving from +// another include path would compile cleanly and then disagree with the linked +// library at runtime, which is the failure mode this whole build path exists to +// prevent. +#ifndef ANYTIME_EXPECTED_ORT_API_VERSION +#error "ANYTIME_EXPECTED_ORT_API_VERSION is not defined; configure through runtime/CMakeLists.txt" +#endif +static_assert(ORT_API_VERSION == ANYTIME_EXPECTED_ORT_API_VERSION, + "onnxruntime_c_api.h does not match the SDK CMake resolved. The " + "include path is picking up a different ONNX Runtime."); + +namespace anytime { + +std::size_t dtype_size(DType dtype) { + switch (dtype) { + case DType::Float32: return sizeof(float); + case DType::Float64: return sizeof(double); + case DType::Int32: return sizeof(int32_t); + case DType::Int64: return sizeof(int64_t); + case DType::Bool: return sizeof(bool); + } + throw std::runtime_error("unhandled dtype"); +} + +const char* dtype_name(DType dtype) { + switch (dtype) { + case DType::Float32: return "float32"; + case DType::Float64: return "float64"; + case DType::Int32: return "int32"; + case DType::Int64: return "int64"; + case DType::Bool: return "bool"; + } + throw std::runtime_error("unhandled dtype"); +} + +ONNXTensorElementDataType dtype_to_ort(DType dtype) { + switch (dtype) { + case DType::Float32: return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; + case DType::Float64: return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; + case DType::Int32: return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; + case DType::Int64: return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; + case DType::Bool: return ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; + } + throw std::runtime_error("unhandled dtype"); +} + +DType dtype_from_ort(ONNXTensorElementDataType type) { + switch (type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return DType::Float32; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: return DType::Float64; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: return DType::Int32; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: return DType::Int64; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: return DType::Bool; + default: + throw std::runtime_error("unsupported ONNX element type: " + + ort_type_name(type)); + } +} + +std::string ort_type_name(ONNXTensorElementDataType type) { + switch (type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return "float32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: return "float64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: return "float16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: return "bfloat16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: return "int8"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: return "uint8"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: return "int16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: return "uint16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: return "int32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: return "uint32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: return "int64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: return "uint64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: return "bool"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: return "string"; + default: return "unknown"; + } +} + +std::size_t TensorView::element_count() const { + std::size_t count = 1; + for (const int64_t dimension : shape) { + if (dimension < 0) { + throw std::runtime_error("tensor shape has a negative dimension"); + } + count *= static_cast(dimension); + } + return count; +} + +std::size_t TensorView::byte_count() const { + return element_count() * dtype_size(dtype); +} + +Ort::Value borrow_as_tensor(const TensorView& view, const OrtMemoryInfo* memory) { + if (view.data == nullptr) { + throw std::runtime_error("tensor view has no data pointer"); + } + // CreateTensor over caller memory: ONNX Runtime reads through the pointer and + // does not take ownership, so nothing is copied on the way in. + return Ort::Value::CreateTensor( + memory, + const_cast(view.data), + view.byte_count(), + view.shape.data(), + view.shape.size(), + dtype_to_ort(view.dtype)); +} + +} // namespace anytime diff --git a/runtime_cpp/CMakeLists.txt b/runtime_cpp/CMakeLists.txt deleted file mode 100644 index 8438b24..0000000 --- a/runtime_cpp/CMakeLists.txt +++ /dev/null @@ -1,56 +0,0 @@ -cmake_minimum_required(VERSION 3.14) -project(anytime_runtime CXX) - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_POSITION_INDEPENDENT_CODE ON) - -if(NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE Release) -endif() - -# Locate ONNX Runtime. Either set ONNXRUNTIME_ROOT_PATH to the unpacked release -# tarball (which contains include/ and lib/) or install ONNX Runtime such that -# CMake can find it via find_package. -if(NOT DEFINED ONNXRUNTIME_ROOT_PATH AND DEFINED ENV{ONNXRUNTIME_ROOT_PATH}) - set(ONNXRUNTIME_ROOT_PATH $ENV{ONNXRUNTIME_ROOT_PATH}) -endif() - -if(NOT ONNXRUNTIME_ROOT_PATH) - message(FATAL_ERROR - "ONNXRUNTIME_ROOT_PATH is not set. Download ONNX Runtime from " - "https://github.com/microsoft/onnxruntime/releases and pass " - "-DONNXRUNTIME_ROOT_PATH=/path/to/onnxruntime to cmake.") -endif() - -add_executable(anytime_runtime src/main.cpp) - -target_include_directories(anytime_runtime PRIVATE ${ONNXRUNTIME_ROOT_PATH}/include) - -# find_library caches its result, so reconfiguring with a different -# ONNXRUNTIME_ROOT_PATH would otherwise keep linking the previously found -# library while compiling against the new headers. That mismatch surfaces only at -# runtime, as "The requested API version [N] is not available". -if(DEFINED CACHE{ONNXRUNTIME_LIB} AND NOT ONNXRUNTIME_LIB MATCHES "^${ONNXRUNTIME_ROOT_PATH}/") - message(STATUS "ONNXRUNTIME_ROOT_PATH changed; re-resolving ONNXRUNTIME_LIB") - unset(ONNXRUNTIME_LIB CACHE) -endif() - -find_library(ONNXRUNTIME_LIB - NAMES onnxruntime - PATHS ${ONNXRUNTIME_ROOT_PATH}/lib - NO_DEFAULT_PATH - REQUIRED) - -message(STATUS "Linking ONNX Runtime: ${ONNXRUNTIME_LIB}") - -target_link_libraries(anytime_runtime PRIVATE ${ONNXRUNTIME_LIB}) - -if(APPLE) - set_target_properties(anytime_runtime PROPERTIES - BUILD_RPATH "${ONNXRUNTIME_ROOT_PATH}/lib" - INSTALL_RPATH "${ONNXRUNTIME_ROOT_PATH}/lib") -elseif(UNIX) - target_link_options(anytime_runtime PRIVATE - "-Wl,-rpath,${ONNXRUNTIME_ROOT_PATH}/lib") -endif() diff --git a/runtime_cpp/README.md b/runtime_cpp/README.md deleted file mode 100644 index e51e2f0..0000000 --- a/runtime_cpp/README.md +++ /dev/null @@ -1,57 +0,0 @@ -# anytime_runtime (C++) - -Single-threaded ONNX Runtime worker. Loads one ONNX graph per variant and answers -line-delimited JSON requests on stdin. `RuntimePool` in -`src/anytime_serving/serving/onnx_runtime.py` spawns one process per pool slot and -dispatches across them. - -## Build - -```bash -cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \ - -DONNXRUNTIME_ROOT_PATH=/path/to/onnxruntime -cmake --build build -j -``` - -Download the release matching your platform from -. Point -`ONNXRUNTIME_ROOT_PATH` at the directory containing `include/` and `lib/`. - -**The version must match the `onnxruntime` Python wheel.** A mismatch is a silent -performance cliff, not a link error. See -[`../docs/runtime.md`](../docs/runtime.md). - -The binary lands at `build/anytime_runtime`. The Python client finds it via -`ANYTIME_RUNTIME_BIN` or by searching upward from the package and the working -directory. - -## Protocol - -Handshake: the worker prints `ready` once every model has loaded. - -```json -{"request_id":"r-1","variant":"distilbert_fp32", - "inputs":{"input_ids":{"shape":[1,128],"dtype":"int64","data":""}, - "attention_mask":{"shape":[1,128],"dtype":"int64","data":""}}} -``` - -```json -{"request_id":"r-1", - "logits":{"shape":[1,2],"dtype":"float32","data":""}, - "latency_ms":14.05} -``` - -Recoverable failures return `{"request_id":"r-1","error":"..."}` and the worker -continues serving. - -## Notes - -- One inference per process. Run a pool from the Python side for concurrency. -- Input dtypes `float32` and `int64`; output dtype `float32`. -- Inputs a graph does not declare are dropped; a missing declared input is an - error. -- No dependencies beyond ONNX Runtime. JSON and base64 are implemented in - `src/main.cpp`. - -Full details, including the version-matching rationale and the CMake cache -pitfall: [`../docs/runtime.md`](../docs/runtime.md). diff --git a/runtime_cpp/src/main.cpp b/runtime_cpp/src/main.cpp deleted file mode 100644 index 7b6e34e..0000000 --- a/runtime_cpp/src/main.cpp +++ /dev/null @@ -1,494 +0,0 @@ -// anytime_runtime: a minimal C++ inference worker. -// -// Loads one or more ONNX models (one per "variant" name) and accepts -// line-delimited JSON requests on stdin. Each request specifies the variant -// and the input tensors; the worker runs ONNX Runtime and writes a response -// line containing the output tensor and the measured per-request latency. -// -// Protocol (one JSON object per line): -// request : {"request_id": "...", "variant": "fp32", -// "inputs": {"": {"shape": [...], "dtype": "...", -// "data": ""}}} -// response: {"request_id": "...", "logits": {"shape": [...], -// "dtype": "float32", "data": ""}, "latency_ms": 12.3} -// -// On startup, after every model is loaded, the worker prints the literal -// string "ready" followed by a newline. The Python client uses this as a -// handshake. - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -std::string base64_encode(const uint8_t* data, std::size_t length) { - static constexpr char kAlphabet[] = - "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; - std::string out; - out.reserve(((length + 2) / 3) * 4); - std::size_t i = 0; - for (; i + 3 <= length; i += 3) { - const uint32_t value = (static_cast(data[i]) << 16) | - (static_cast(data[i + 1]) << 8) | - static_cast(data[i + 2]); - out.push_back(kAlphabet[(value >> 18) & 0x3F]); - out.push_back(kAlphabet[(value >> 12) & 0x3F]); - out.push_back(kAlphabet[(value >> 6) & 0x3F]); - out.push_back(kAlphabet[value & 0x3F]); - } - if (i < length) { - const std::size_t remaining = length - i; - uint32_t value = static_cast(data[i]) << 16; - if (remaining == 2) { - value |= static_cast(data[i + 1]) << 8; - } - out.push_back(kAlphabet[(value >> 18) & 0x3F]); - out.push_back(kAlphabet[(value >> 12) & 0x3F]); - out.push_back(remaining == 2 ? kAlphabet[(value >> 6) & 0x3F] : '='); - out.push_back('='); - } - return out; -} - -std::vector base64_decode(const std::string& input) { - static int8_t table[256]; - static bool initialised = false; - if (!initialised) { - std::fill(std::begin(table), std::end(table), int8_t{-1}); - static const char* alphabet = - "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; - for (int i = 0; i < 64; ++i) { - table[static_cast(alphabet[i])] = static_cast(i); - } - initialised = true; - } - std::vector out; - out.reserve((input.size() / 4) * 3); - uint32_t buffer = 0; - int bits = 0; - for (char c : input) { - if (c == '=' || c == '\n' || c == '\r' || c == ' ') continue; - const int8_t value = table[static_cast(c)]; - if (value < 0) throw std::runtime_error("invalid base64 character"); - buffer = (buffer << 6) | static_cast(value); - bits += 6; - if (bits >= 8) { - bits -= 8; - out.push_back(static_cast((buffer >> bits) & 0xFF)); - } - } - return out; -} - -// Tiny JSON reader sufficient for the request envelope produced by the Python -// client. Strings are unescaped only for the characters the client emits -// (\\ and \"). Numbers are parsed as doubles. Arrays of numbers are returned -// as std::vector. Anything more exotic raises std::runtime_error. -class JsonReader { -public: - explicit JsonReader(const std::string& source) : source_(source), pos_(0) {} - - struct Value { - enum class Kind { Null, Bool, Number, String, Array, Object }; - Kind kind = Kind::Null; - bool boolean = false; - double number = 0.0; - std::string string; - std::vector array; - std::vector> object; - - const Value& find(const std::string& key) const { - for (const auto& [k, v] : object) { - if (k == key) return v; - } - throw std::runtime_error("missing JSON key: " + key); - } - std::vector as_int_array() const { - std::vector out; - out.reserve(array.size()); - for (const auto& v : array) { - out.push_back(static_cast(v.number)); - } - return out; - } - }; - - Value parse() { - skip_ws(); - Value v = parse_value(); - skip_ws(); - if (pos_ != source_.size()) { - throw std::runtime_error("trailing characters in JSON"); - } - return v; - } - -private: - void skip_ws() { - while (pos_ < source_.size() && std::isspace(static_cast(source_[pos_]))) { - ++pos_; - } - } - - Value parse_value() { - skip_ws(); - if (pos_ >= source_.size()) throw std::runtime_error("unexpected end of JSON"); - char c = source_[pos_]; - if (c == '{') return parse_object(); - if (c == '[') return parse_array(); - if (c == '"') return parse_string(); - if (c == 't' || c == 'f') return parse_bool(); - if (c == 'n') { expect_literal("null"); Value v; v.kind = Value::Kind::Null; return v; } - return parse_number(); - } - - Value parse_object() { - Value v; - v.kind = Value::Kind::Object; - ++pos_; // consume '{' - skip_ws(); - if (pos_ < source_.size() && source_[pos_] == '}') { ++pos_; return v; } - while (true) { - skip_ws(); - Value key = parse_string(); - skip_ws(); - if (pos_ >= source_.size() || source_[pos_] != ':') { - throw std::runtime_error("expected ':' in JSON object"); - } - ++pos_; - Value value = parse_value(); - v.object.emplace_back(std::move(key.string), std::move(value)); - skip_ws(); - if (pos_ < source_.size() && source_[pos_] == ',') { ++pos_; continue; } - if (pos_ < source_.size() && source_[pos_] == '}') { ++pos_; break; } - throw std::runtime_error("expected ',' or '}' in JSON object"); - } - return v; - } - - Value parse_array() { - Value v; - v.kind = Value::Kind::Array; - ++pos_; // consume '[' - skip_ws(); - if (pos_ < source_.size() && source_[pos_] == ']') { ++pos_; return v; } - while (true) { - v.array.push_back(parse_value()); - skip_ws(); - if (pos_ < source_.size() && source_[pos_] == ',') { ++pos_; continue; } - if (pos_ < source_.size() && source_[pos_] == ']') { ++pos_; break; } - throw std::runtime_error("expected ',' or ']' in JSON array"); - } - return v; - } - - Value parse_string() { - Value v; - v.kind = Value::Kind::String; - if (source_[pos_] != '"') throw std::runtime_error("expected '\"' in JSON string"); - ++pos_; - std::string out; - while (pos_ < source_.size() && source_[pos_] != '"') { - char c = source_[pos_++]; - if (c == '\\' && pos_ < source_.size()) { - char escape = source_[pos_++]; - switch (escape) { - case '"': out.push_back('"'); break; - case '\\': out.push_back('\\'); break; - case '/': out.push_back('/'); break; - case 'n': out.push_back('\n'); break; - case 't': out.push_back('\t'); break; - case 'r': out.push_back('\r'); break; - default: throw std::runtime_error("unsupported JSON escape"); - } - } else { - out.push_back(c); - } - } - if (pos_ >= source_.size()) throw std::runtime_error("unterminated JSON string"); - ++pos_; // closing quote - v.string = std::move(out); - return v; - } - - Value parse_bool() { - Value v; - v.kind = Value::Kind::Bool; - if (source_[pos_] == 't') { - expect_literal("true"); - v.boolean = true; - } else { - expect_literal("false"); - v.boolean = false; - } - return v; - } - - Value parse_number() { - std::size_t start = pos_; - if (source_[pos_] == '-') ++pos_; - while (pos_ < source_.size() && (std::isdigit(static_cast(source_[pos_])) || - source_[pos_] == '.' || source_[pos_] == 'e' || - source_[pos_] == 'E' || source_[pos_] == '+' || - source_[pos_] == '-')) { - ++pos_; - } - Value v; - v.kind = Value::Kind::Number; - v.number = std::stod(source_.substr(start, pos_ - start)); - return v; - } - - void expect_literal(const char* literal) { - const std::size_t n = std::strlen(literal); - if (pos_ + n > source_.size() || source_.compare(pos_, n, literal) != 0) { - throw std::runtime_error(std::string("expected literal ") + literal); - } - pos_ += n; - } - - const std::string& source_; - std::size_t pos_; -}; - -std::string escape_json(const std::string& s) { - std::string out; - out.reserve(s.size()); - for (char c : s) { - switch (c) { - case '"': out += "\\\""; break; - case '\\': out += "\\\\"; break; - case '\n': out += "\\n"; break; - case '\t': out += "\\t"; break; - case '\r': out += "\\r"; break; - default: out.push_back(c); - } - } - return out; -} - -struct LoadedModel { - std::unique_ptr session; - std::vector input_names; - std::vector output_names; -}; - -LoadedModel load_model(Ort::Env& env, const std::string& path) { - Ort::SessionOptions options; - options.SetIntraOpNumThreads(1); - options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); - - LoadedModel model; - model.session = std::make_unique(env, path.c_str(), options); - - Ort::AllocatorWithDefaultOptions allocator; - const std::size_t num_inputs = model.session->GetInputCount(); - for (std::size_t i = 0; i < num_inputs; ++i) { - auto name = model.session->GetInputNameAllocated(i, allocator); - model.input_names.emplace_back(name.get()); - } - const std::size_t num_outputs = model.session->GetOutputCount(); - for (std::size_t i = 0; i < num_outputs; ++i) { - auto name = model.session->GetOutputNameAllocated(i, allocator); - model.output_names.emplace_back(name.get()); - } - return model; -} - -Ort::Value make_tensor(const JsonReader::Value& payload, - Ort::MemoryInfo& memory, - std::vector& raw_storage, - std::vector& shape_storage, - std::vector& float_storage, - std::vector& int64_storage) { - const std::string dtype = payload.find("dtype").string; - shape_storage = payload.find("shape").as_int_array(); - raw_storage = base64_decode(payload.find("data").string); - - if (dtype == "float32") { - const std::size_t count = raw_storage.size() / sizeof(float); - float_storage.resize(count); - std::memcpy(float_storage.data(), raw_storage.data(), raw_storage.size()); - return Ort::Value::CreateTensor(memory, float_storage.data(), count, - shape_storage.data(), shape_storage.size()); - } - if (dtype == "int64") { - const std::size_t count = raw_storage.size() / sizeof(int64_t); - int64_storage.resize(count); - std::memcpy(int64_storage.data(), raw_storage.data(), raw_storage.size()); - return Ort::Value::CreateTensor(memory, int64_storage.data(), count, - shape_storage.data(), shape_storage.size()); - } - throw std::runtime_error("unsupported dtype: " + dtype); -} - -std::string serialise_output(const Ort::Value& tensor) { - auto info = tensor.GetTensorTypeAndShapeInfo(); - const std::vector shape = info.GetShape(); - const std::size_t element_count = info.GetElementCount(); - const auto element_type = info.GetElementType(); - if (element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { - throw std::runtime_error("only float32 output tensors are supported"); - } - const float* data = tensor.GetTensorData(); - const std::string encoded = base64_encode( - reinterpret_cast(data), element_count * sizeof(float)); - - std::ostringstream shape_json; - shape_json << '['; - for (std::size_t i = 0; i < shape.size(); ++i) { - if (i) shape_json << ','; - shape_json << shape[i]; - } - shape_json << ']'; - - std::ostringstream out; - out << "{\"shape\":" << shape_json.str() - << ",\"dtype\":\"float32\"" - << ",\"data\":\"" << encoded << "\"}"; - return out.str(); -} - -struct ParsedArgs { - std::vector> models; -}; - -ParsedArgs parse_args(int argc, char** argv) { - ParsedArgs out; - for (int i = 1; i < argc; ++i) { - std::string flag = argv[i]; - if (flag == "--model" && i + 1 < argc) { - std::string value = argv[++i]; - const std::size_t eq = value.find('='); - if (eq == std::string::npos) { - throw std::runtime_error("--model expects ="); - } - out.models.emplace_back(value.substr(0, eq), value.substr(eq + 1)); - } else { - throw std::runtime_error("unknown argument: " + flag); - } - } - if (out.models.empty()) { - throw std::runtime_error("at least one --model must be specified"); - } - return out; -} - -} // namespace - - -int main(int argc, char** argv) { - try { - const ParsedArgs args = parse_args(argc, argv); - - Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "anytime_runtime"); - std::unordered_map models; - for (const auto& [variant, path] : args.models) { - models[variant] = load_model(env, path); - } - - std::cout << "ready" << std::endl; - - Ort::MemoryInfo memory = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); - - std::string line; - while (std::getline(std::cin, line)) { - if (line.empty()) continue; - - // Per-request failures must not take the worker down: a pool shares - // one process per worker, so exiting here would fail every in-flight - // and subsequent request on this worker. Anything recoverable is - // reported as an error response and the loop continues. - std::string request_id; - try { - JsonReader reader(line); - JsonReader::Value request = reader.parse(); - - request_id = request.find("request_id").string; - const std::string variant = request.find("variant").string; - const JsonReader::Value& inputs = request.find("inputs"); - - const auto it = models.find(variant); - if (it == models.end()) { - std::cerr << "unknown variant: " << variant << std::endl; - std::cout << "{\"request_id\":\"" << escape_json(request_id) - << "\",\"error\":\"unknown variant\"}" << std::endl; - continue; - } - LoadedModel& model = it->second; - - std::vector> raw_storage(inputs.object.size()); - std::vector> shape_storage(inputs.object.size()); - std::vector> float_storage(inputs.object.size()); - std::vector> int64_storage(inputs.object.size()); - std::vector tensors; - tensors.reserve(inputs.object.size()); - std::vector input_name_ptrs; - input_name_ptrs.reserve(inputs.object.size()); - - // Variants of the same task can declare different inputs: a - // DistilBERT graph takes input_ids and attention_mask, a BERT - // graph additionally takes token_type_ids. The caller sends the - // union, so drop anything this graph does not declare rather than - // failing the request. - for (std::size_t i = 0; i < inputs.object.size(); ++i) { - const auto& [name, payload] = inputs.object[i]; - const bool declared = std::find(model.input_names.begin(), - model.input_names.end(), - name) != model.input_names.end(); - if (!declared) continue; - input_name_ptrs.push_back(name.c_str()); - tensors.push_back(make_tensor(payload, memory, - raw_storage[i], shape_storage[i], - float_storage[i], int64_storage[i])); - } - if (tensors.size() != model.input_names.size()) { - throw std::runtime_error( - "request supplies " + std::to_string(tensors.size()) + - " of the " + std::to_string(model.input_names.size()) + - " inputs this graph declares"); - } - - std::vector output_name_ptrs; - output_name_ptrs.reserve(model.output_names.size()); - for (const auto& n : model.output_names) output_name_ptrs.push_back(n.c_str()); - - const auto start = std::chrono::steady_clock::now(); - auto outputs = model.session->Run( - Ort::RunOptions{nullptr}, - input_name_ptrs.data(), tensors.data(), tensors.size(), - output_name_ptrs.data(), output_name_ptrs.size()); - const auto end = std::chrono::steady_clock::now(); - const double latency_ms = - std::chrono::duration(end - start).count(); - - std::string logits_json = serialise_output(outputs.front()); - std::cout << "{\"request_id\":\"" << escape_json(request_id) - << "\",\"logits\":" << logits_json - << ",\"latency_ms\":" << latency_ms << "}" << std::endl; - } catch (const std::exception& exc) { - std::cerr << "request error: " << exc.what() << std::endl; - std::cout << "{\"request_id\":\"" << escape_json(request_id) - << "\",\"error\":\"" << escape_json(exc.what()) << "\"}" - << std::endl; - } - } - return 0; - } catch (const std::exception& exc) { - std::cerr << "fatal: " << exc.what() << std::endl; - return 1; - } -} diff --git a/scripts/demo_serving.py b/scripts/demo_serving.py index c7e2a06..462930f 100644 --- a/scripts/demo_serving.py +++ b/scripts/demo_serving.py @@ -22,7 +22,7 @@ from onnx import TensorProto, helper from anytime_serving.serving.load_monitor import LoadMonitor -from anytime_serving.serving.onnx_runtime import InferenceRequest, RuntimePool, find_runtime_binary +from anytime_serving.serving.onnx_runtime import InferenceRequest, RuntimePool from anytime_serving.serving.selector import AdaptiveSelector, VariantProfile from anytime_serving.serving.server import AdaptiveServer, drive_workload, poisson_arrivals from anytime_serving.utils.logger import get_logger @@ -71,7 +71,7 @@ def _run_policy( monitor.start() rng = np.random.default_rng(seed) try: - with RuntimePool(workers, model_paths, binary=find_runtime_binary()) as pool: + with RuntimePool(workers, model_paths) as pool: server = AdaptiveServer(pool, selector, monitor) try: diff --git a/scripts/profile_variants.py b/scripts/profile_variants.py index 248aedc..f779298 100644 --- a/scripts/profile_variants.py +++ b/scripts/profile_variants.py @@ -5,11 +5,22 @@ carrying a number over from a different host, thread count, or batch size makes every downstream admission decision wrong. -This script measures every exported candidate under the same ONNX Runtime -configuration the serving worker uses (one intra-op thread per worker, batch size -one), computes the Pareto frontier over (service time, accuracy), and writes: - - results/variant_profiles.json all measurements plus host metadata +Measurement goes **through the serving path**, not beside it. Every latency and +every logit here comes from the same `RuntimeClient` the server dispatches to. +Stage 1 profiled through a separate ONNX Runtime session instead, and that is what +let a 7.6x discrepancy hide: the C++ worker was built against ONNX Runtime 1.20.1 +while the profiler used the 1.26.0 wheel, so DistilBERT measured 98.9 ms in the +worker against 13.0 ms in the profiler. Nothing failed. Every service time the +planner used was simply false. + +To keep that from recurring in a form the engine cannot see, each variant is also +measured through a separate ONNX Runtime session and the two are required to agree. +A divergence beyond --agreement-tolerance fails the run rather than being written +to disk. This is the check whose absence made Stage 1's numbers wrong. + +Writes: + + results/variant_profiles.json all measurements, agreement, host metadata configs/serving.yaml the frontier, as the serving harness reads it Which variants end up on the frontier is a property of the hardware, not an @@ -40,6 +51,12 @@ import numpy as np import onnxruntime as ort +from anytime_serving.serving.onnx_runtime import ( + InferenceRequest, + RuntimeClient, + extension_available, + load_extension, +) from anytime_serving.utils.logger import get_logger LOGGER = get_logger("scripts.profile_variants") @@ -50,11 +67,20 @@ QUICK_WARMUP_ITERATIONS = 5 QUICK_MEASURE_ITERATIONS = 40 QUICK_ACCURACY_SAMPLES = 128 +# Independent measurement passes per variant. One pass is not enough: the same +# 200-request p50 measured eight times ranged over 6.5% for DistilBERT and 3.6% +# for MiniLM on this host, driven by thermal state rather than by anything the +# code does. +DEFAULT_REPEATS = 3 +# Matched versions measured within 0.4% on this host, and Stage 1 saw 8% across a +# process boundary. 15% is loose enough not to fire on scheduler noise and tight +# enough that anything structural, let alone a 7.6x version mismatch, trips it. +DEFAULT_AGREEMENT_TOLERANCE = 0.15 @dataclass class VariantMeasurement: - """Measured cost and quality for one variant.""" + """Measured cost and quality for one variant, as served.""" name: str model: str @@ -68,19 +94,30 @@ class VariantMeasurement: latency_mean_ms: float latency_stdev_ms: float iterations: int + passes: int + # Range of the per-pass p50s, as a percentage. This is the honest uncertainty + # on service_time_ms; the within-pass stdev above is much smaller and does not + # describe how far the number moves between runs. + service_time_spread_pct: float accuracy: float accuracy_samples: int + # Time inside inference as the engine reports it, against the same graph run + # through a separate session. Agreement is what says the number is real. + engine_inference_p50_ms: float + direct_session_p50_ms: float + agreement_ratio: float compute_cost_per_request: float = 1.0 on_pareto_frontier: bool = False dominated_by: str | None = None def _session(graph: Path) -> ort.InferenceSession: - """Build a session matching the serving worker's configuration. + """Build a cross-check session matching the engine's configuration. - runtime_cpp/src/main.cpp sets intra-op threads to one so that N pooled - workers behave as N independent single-threaded servers. Profiling with a - different thread count would measure a machine the planner never runs on. + runtime/src/engine.cpp sets intra-op and inter-op threads to one so that N + pooled workers behave as N independent single-threaded servers. This session + exists only to be compared against the engine, so it has to be configured the + same way or the comparison measures the configuration difference instead. """ options = ort.SessionOptions() options.intra_op_num_threads = 1 @@ -101,7 +138,12 @@ def _find_graph(directory: Path) -> Path: return graphs[0] -def _feeds(tokenizer, session: ort.InferenceSession, text: str) -> dict[str, np.ndarray]: +def _tokenize(tokenizer, text: str) -> dict[str, np.ndarray]: + """Encode one string. Not filtered to any graph's declared inputs. + + The engine drops inputs a graph does not declare, so the union is what the + server sends and therefore what should be measured. + """ encoded = tokenizer( text, return_tensors="np", @@ -109,41 +151,120 @@ def _feeds(tokenizer, session: ort.InferenceSession, text: str) -> dict[str, np. truncation=True, max_length=SEQUENCE_LENGTH, ) - expected = {i.name for i in session.get_inputs()} - return {name: value.astype(np.int64) for name, value in encoded.items() if name in expected} + return {name: value.astype(np.int64) for name, value in encoded.items()} + + +def _declared(session: ort.InferenceSession, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + expected = {spec.name for spec in session.get_inputs()} + return {name: value for name, value in feeds.items() if name in expected} -def _measure_latency( +def _percentiles(samples: list[float], key: str) -> dict[str, float]: + array = np.asarray(samples) + return { + f"{key}_p50_ms": float(np.percentile(array, 50)), + f"{key}_p95_ms": float(np.percentile(array, 95)), + f"{key}_p99_ms": float(np.percentile(array, 99)), + f"{key}_mean_ms": statistics.fmean(samples), + f"{key}_stdev_ms": statistics.stdev(samples) if len(samples) > 1 else 0.0, + } + + +def _spread_pct(values: list[float]) -> float: + """Range of a set of measurements as a percentage of the smallest.""" + if len(values) < 2 or min(values) <= 0.0: + return 0.0 + return (max(values) - min(values)) / min(values) * 100.0 + + +def _measure_through_engine( + client: RuntimeClient, + variant: str, + feeds: dict[str, np.ndarray], + *, + warmup: int, + iterations: int, + repeats: int, +) -> dict[str, float]: + """Latency as the server experiences it, through the runtime client. + + Repeated in independent passes, each with its own warm-up, and reported as the + median of the per-pass medians. One pass is not enough: the same 200-request + p50 measured eight times on this host ranged over 6.5% for DistilBERT, because + thermal state and scheduler placement drift with whatever ran before. The + within-pass standard deviation is around 0.45 ms and says nothing about that, + so quoting a single p50 with a within-pass spread overstates its precision. + """ + request = InferenceRequest(variant=variant, inputs=feeds) + pass_wall_p50: list[float] = [] + pass_inference_p50: list[float] = [] + wall: list[float] = [] + inference: list[float] = [] + + for _ in range(repeats): + for _ in range(warmup): + client.infer(request) + pass_wall: list[float] = [] + pass_inference: list[float] = [] + for _ in range(iterations): + response = client.infer(request) + pass_wall.append(response.wall_latency_ms) + pass_inference.append(response.runtime_latency_ms) + pass_wall_p50.append(float(np.percentile(pass_wall, 50))) + pass_inference_p50.append(float(np.percentile(pass_inference, 50))) + wall.extend(pass_wall) + inference.extend(pass_inference) + + # Tails come from the pooled samples, so they reflect variation across passes + # as well as within them. The headline median does not, by construction. + stats = _percentiles(wall, "wall") + stats.update(_percentiles(inference, "inference")) + stats["wall_p50_ms"] = statistics.median(pass_wall_p50) + stats["inference_p50_ms"] = statistics.median(pass_inference_p50) + stats["wall_p50_spread_pct"] = _spread_pct(pass_wall_p50) + stats["iterations"] = float(iterations) + stats["passes"] = float(repeats) + return stats + + +def _measure_direct_session( session: ort.InferenceSession, feeds: dict[str, np.ndarray], *, warmup: int, iterations: int, + repeats: int, ) -> dict[str, float]: - for _ in range(warmup): - session.run(None, feeds) + """The same graph through a separate session, for the agreement check. + Repeated the same number of times as the engine measurement, so the two are + compared like for like rather than one median against a noisier one. + """ + fed = _declared(session, feeds) + pass_p50: list[float] = [] samples: list[float] = [] - for _ in range(iterations): - start = time.perf_counter() - session.run(None, feeds) - samples.append((time.perf_counter() - start) * 1000.0) - array = np.asarray(samples) - return { - "latency_p50_ms": float(np.percentile(array, 50)), - "latency_p95_ms": float(np.percentile(array, 95)), - "latency_p99_ms": float(np.percentile(array, 99)), - "latency_mean_ms": statistics.fmean(samples), - "latency_stdev_ms": statistics.stdev(samples) if len(samples) > 1 else 0.0, - "iterations": float(len(samples)), - } + for _ in range(repeats): + for _ in range(warmup): + session.run(None, fed) + pass_samples: list[float] = [] + for _ in range(iterations): + start = time.perf_counter() + session.run(None, fed) + pass_samples.append((time.perf_counter() - start) * 1000.0) + pass_p50.append(float(np.percentile(pass_samples, 50))) + samples.extend(pass_samples) + + stats = _percentiles(samples, "direct") + stats["direct_p50_ms"] = statistics.median(pass_p50) + stats["direct_p50_spread_pct"] = _spread_pct(pass_p50) + return stats def _measure_accuracy( - session: ort.InferenceSession, tokenizer, samples: int | None + client: RuntimeClient, variant: str, tokenizer, samples: int | None ) -> tuple[float, int]: - """Accuracy on the SST-2 validation split.""" + """Accuracy on the SST-2 validation split, scored through the engine.""" from datasets import load_dataset dataset = load_dataset("glue", "sst2", split="validation") @@ -152,7 +273,8 @@ def _measure_accuracy( correct = 0 for row in dataset: - logits = session.run(None, _feeds(tokenizer, session, row["sentence"]))[0] + feeds = _tokenize(tokenizer, row["sentence"]) + logits = client.infer(InferenceRequest(variant=variant, inputs=feeds)).logits if int(np.argmax(logits, axis=-1)[0]) == int(row["label"]): correct += 1 return correct / len(dataset), len(dataset) @@ -184,68 +306,114 @@ def _mark_pareto_frontier(measurements: list[VariantMeasurement]) -> None: candidate.dominated_by = dominator.name if dominator else None -def profile(model_dir: Path, quick: bool) -> list[VariantMeasurement]: - from transformers import AutoTokenizer - - warmup = QUICK_WARMUP_ITERATIONS if quick else WARMUP_ITERATIONS - iterations = QUICK_MEASURE_ITERATIONS if quick else MEASURE_ITERATIONS - accuracy_samples = QUICK_ACCURACY_SAMPLES if quick else None - +def _discover(model_dir: Path) -> list[tuple[str, str, str, Path]]: + """Find exported variants as (name, model, precision, graph).""" directories = sorted(model_dir.glob("text_*_*")) if not directories: raise SystemExit( f"no exported variants under {model_dir}. Run: " f"python scripts/export_onnx.py --task text --output-dir {model_dir}" ) - - measurements: list[VariantMeasurement] = [] + found = [] for directory in directories: if not directory.is_dir(): continue # Directory names are text__. _, model_name, precision = directory.name.split("_", 2) - graph = _find_graph(directory) - tokenizer = AutoTokenizer.from_pretrained(directory) - session = _session(graph) - - name = f"{model_name}_{precision}" - LOGGER.info("Profiling %s (%s)", name, graph.name) - latency = _measure_latency( - session, - _feeds(tokenizer, session, "a genuinely measured service time"), - warmup=warmup, - iterations=iterations, - ) - accuracy, n_samples = _measure_accuracy(session, tokenizer, accuracy_samples) - - measurements.append( - VariantMeasurement( - name=name, - model=model_name, - precision=precision, - graph=graph.name, - size_mb=round(graph.stat().st_size / 1e6, 1), - # The planner admits on expected service time, so use the median: - # it is not skewed by occasional scheduler noise. - service_time_ms=round(latency["latency_p50_ms"], 3), - latency_p50_ms=round(latency["latency_p50_ms"], 3), - latency_p95_ms=round(latency["latency_p95_ms"], 3), - latency_p99_ms=round(latency["latency_p99_ms"], 3), - latency_mean_ms=round(latency["latency_mean_ms"], 3), - latency_stdev_ms=round(latency["latency_stdev_ms"], 3), - iterations=int(latency["iterations"]), - accuracy=round(accuracy, 4), - accuracy_samples=n_samples, + found.append((f"{model_name}_{precision}", model_name, precision, _find_graph(directory))) + return found + + +def profile( + model_dir: Path, quick: bool, tolerance: float, repeats: int +) -> tuple[list[VariantMeasurement], str]: + from transformers import AutoTokenizer + + warmup = QUICK_WARMUP_ITERATIONS if quick else WARMUP_ITERATIONS + iterations = QUICK_MEASURE_ITERATIONS if quick else MEASURE_ITERATIONS + accuracy_samples = QUICK_ACCURACY_SAMPLES if quick else None + + variants = _discover(model_dir) + # One client holding every variant, which is how the server loads them. + client = RuntimeClient({name: graph for name, _, _, graph in variants}) + LOGGER.info("Profiling through the %s backend", client.backend_name) + + measurements: list[VariantMeasurement] = [] + disagreements: list[str] = [] + try: + for name, model_name, precision, graph in variants: + tokenizer = AutoTokenizer.from_pretrained(graph.parent) + feeds = _tokenize(tokenizer, "a genuinely measured service time") + + LOGGER.info("Profiling %s (%s)", name, graph.name) + engine = _measure_through_engine( + client, name, feeds, warmup=warmup, iterations=iterations, repeats=repeats ) - ) - LOGGER.info( - " p50=%.2fms p95=%.2fms p99=%.2fms accuracy=%.4f (n=%d) size=%.1fMB", - latency["latency_p50_ms"], - latency["latency_p95_ms"], - latency["latency_p99_ms"], - accuracy, - n_samples, - graph.stat().st_size / 1e6, + direct = _measure_direct_session( + _session(graph), feeds, warmup=warmup, iterations=iterations, repeats=repeats + ) + accuracy, n_samples = _measure_accuracy(client, name, tokenizer, accuracy_samples) + + ratio = engine["inference_p50_ms"] / direct["direct_p50_ms"] + if abs(ratio - 1.0) > tolerance: + disagreements.append( + f" {name}: engine {engine['inference_p50_ms']:.3f} ms vs separate " + f"session {direct['direct_p50_ms']:.3f} ms ({ratio:.3f}x)" + ) + + measurements.append( + VariantMeasurement( + name=name, + model=model_name, + precision=precision, + graph=graph.name, + size_mb=round(graph.stat().st_size / 1e6, 1), + # The planner admits on expected service time, so use the + # median of what a request actually costs the pool: the wall + # time through the client, not just time inside inference. + service_time_ms=round(engine["wall_p50_ms"], 3), + latency_p50_ms=round(engine["wall_p50_ms"], 3), + latency_p95_ms=round(engine["wall_p95_ms"], 3), + latency_p99_ms=round(engine["wall_p99_ms"], 3), + latency_mean_ms=round(engine["wall_mean_ms"], 3), + latency_stdev_ms=round(engine["wall_stdev_ms"], 3), + iterations=int(engine["iterations"]), + passes=int(engine["passes"]), + service_time_spread_pct=round(engine["wall_p50_spread_pct"], 2), + accuracy=round(accuracy, 4), + accuracy_samples=n_samples, + engine_inference_p50_ms=round(engine["inference_p50_ms"], 3), + direct_session_p50_ms=round(direct["direct_p50_ms"], 3), + agreement_ratio=round(ratio, 4), + ) + ) + LOGGER.info( + " p50=%.2fms (+/-%.1f%% over %d passes) p95=%.2fms accuracy=%.4f (n=%d) " + "size=%.1fMB agreement=%.3fx", + engine["wall_p50_ms"], + engine["wall_p50_spread_pct"], + repeats, + engine["wall_p95_ms"], + accuracy, + n_samples, + graph.stat().st_size / 1e6, + ratio, + ) + finally: + client.close() + + if disagreements: + raise SystemExit( + "The engine and a separate ONNX Runtime session disagree by more than " + f"the {tolerance} relative bound:\n" + + "\n".join(disagreements) + + "\n\nThis is the check that Stage 1 lacked. A 7.6x version mismatch " + "between the C++ worker and the onnxruntime wheel went unnoticed for " + "exactly this reason, and every service time the planner used was wrong. " + "Do not write these numbers to disk: confirm the extension and the wheel " + "report the same ONNX Runtime version, then re-run. Pass " + "--agreement-tolerance to widen the bound only if the difference is " + "understood." ) _mark_pareto_frontier(measurements) @@ -257,21 +425,25 @@ def profile(model_dir: Path, quick: bool) -> list[VariantMeasurement]: measurement.compute_cost_per_request = round( measurement.service_time_ms / reference.service_time_ms, 4 ) - return measurements + return measurements, client.backend_name -def _host_metadata() -> dict[str, object]: - return { +def _host_metadata(backend: str) -> dict[str, object]: + metadata: dict[str, object] = { "platform": platform.platform(), "processor": platform.processor() or platform.machine(), "machine": platform.machine(), "python": platform.python_version(), "onnxruntime": ort.__version__, "providers": ort.get_available_providers(), + "backend": backend, "intra_op_num_threads": 1, "batch_size": 1, "sequence_length": SEQUENCE_LENGTH, } + if extension_available(): + metadata["extension_onnxruntime"] = load_extension().onnxruntime_version() + return metadata def _write_serving_config(path: Path, frontier: list[VariantMeasurement], workers: int) -> None: @@ -306,9 +478,10 @@ def _write_serving_config(path: Path, frontier: list[VariantMeasurement], worker existing.setdefault("admission", {"safety_factor": 1.0}) path.write_text( - "# Generated by scripts/profile_variants.py. Values are measured on the\n" - "# host recorded in results/variant_profiles.json; re-run after changing\n" - "# hardware, thread count, or model variants.\n" + yaml.safe_dump(existing, sort_keys=False) + "# Generated by scripts/profile_variants.py. Values are measured through the\n" + "# serving path on the host recorded in results/variant_profiles.json;\n" + "# re-run after changing hardware, thread count, or model variants.\n" + + yaml.safe_dump(existing, sort_keys=False) ) @@ -321,18 +494,61 @@ def main() -> int: action="store_true", help="Fewer iterations and a subset of the validation split", ) + parser.add_argument( + "--repeats", + type=int, + default=DEFAULT_REPEATS, + help=( + "Independent measurement passes per variant; the reported service time " + f"is the median of their medians (default {DEFAULT_REPEATS})" + ), + ) + parser.add_argument( + "--agreement-tolerance", + type=float, + default=DEFAULT_AGREEMENT_TOLERANCE, + help=( + "Maximum relative difference between inference time through the engine " + "and through a separate session before the run fails, as a fraction " + f"(default {DEFAULT_AGREEMENT_TOLERANCE})" + ), + ) + parser.add_argument( + "--allow-fallback-backend", + action="store_true", + help=( + "Profile even when the anytime_runtime extension is unavailable. The " + "numbers then describe the Python fallback rather than the serving path" + ), + ) parser.add_argument("--output", type=Path, default=Path("results/variant_profiles.json")) parser.add_argument("--serving-config", type=Path, default=Path("configs/serving.yaml")) args = parser.parse_args() - measurements = profile(args.model_dir, args.quick) + if not extension_available() and not args.allow_fallback_backend: + raise SystemExit( + "anytime_runtime is not available, so profiling would measure the Python " + "fallback rather than the path that serves traffic. That divergence is " + "what made Stage 1's service times wrong. Build the extension with:\n" + " pip install -e .\n" + "or pass --allow-fallback-backend to record fallback numbers deliberately." + ) + + if args.repeats < 1: + raise SystemExit("--repeats must be at least 1") + + measurements, backend = profile( + args.model_dir, args.quick, args.agreement_tolerance, args.repeats + ) frontier = [m for m in measurements if m.on_pareto_frontier] payload = { - "host": _host_metadata(), + "host": _host_metadata(backend), "quick": args.quick, "task": "text", "accuracy_dataset": "glue/sst2 validation", + "agreement_tolerance": args.agreement_tolerance, + "measurement_passes": args.repeats, "variants": [asdict(m) for m in measurements], "pareto_frontier": [m.name for m in frontier], } @@ -344,15 +560,26 @@ def main() -> int: LOGGER.info("Wrote %s", args.serving_config) LOGGER.info("") - LOGGER.info("%-18s %10s %10s %9s %s", "variant", "p50 (ms)", "accuracy", "size(MB)", "verdict") + LOGGER.info( + "%-18s %10s %8s %10s %9s %10s %s", + "variant", + "p50 (ms)", + "spread", + "accuracy", + "size(MB)", + "agreement", + "verdict", + ) for m in sorted(measurements, key=lambda m: m.service_time_ms): verdict = "frontier" if m.on_pareto_frontier else f"dominated by {m.dominated_by}" LOGGER.info( - "%-18s %10.2f %10.4f %9.1f %s", + "%-18s %10.2f %7.1f%% %10.4f %9.1f %9.3fx %s", m.name, m.service_time_ms, + m.service_time_spread_pct, m.accuracy, m.size_mb, + m.agreement_ratio, verdict, ) diff --git a/scripts/run_load_sweep.py b/scripts/run_load_sweep.py index d6886bf..f70262d 100644 --- a/scripts/run_load_sweep.py +++ b/scripts/run_load_sweep.py @@ -34,7 +34,7 @@ import yaml from anytime_serving.serving.load_monitor import LoadMonitor -from anytime_serving.serving.onnx_runtime import InferenceRequest, RuntimePool, find_runtime_binary +from anytime_serving.serving.onnx_runtime import InferenceRequest, RuntimePool from anytime_serving.serving.selector import AdaptiveSelector, VariantProfile from anytime_serving.serving.server import AdaptiveServer, drive_workload, poisson_arrivals from anytime_serving.utils.logger import get_logger @@ -133,7 +133,7 @@ def _run_policy( monitor = LoadMonitor(interval_s=0.1, alpha=0.3) monitor.start() try: - with RuntimePool(workers, model_paths, binary=find_runtime_binary()) as pool: + with RuntimePool(workers, model_paths) as pool: server = AdaptiveServer(pool, selector, monitor) try: diff --git a/src/anytime_serving/serving/__init__.py b/src/anytime_serving/serving/__init__.py index 4ea7620..cf06b89 100644 --- a/src/anytime_serving/serving/__init__.py +++ b/src/anytime_serving/serving/__init__.py @@ -3,11 +3,13 @@ from .admission import AdmissionDecision, MMcAdmissionController, erlang_b, erlang_c from .load_monitor import LoadMonitor, LoadSample from .onnx_runtime import ( + BACKENDS, InferenceRequest, InferenceResponse, RuntimeClient, RuntimePool, - find_runtime_binary, + extension_available, + load_extension, ) from .selector import AdaptiveSelector, SelectionDecision, VariantProfile from .server import AdaptiveServer, ServedRequest, ServerStats, drive_workload, poisson_arrivals @@ -22,11 +24,13 @@ "AdaptiveSelector", "SelectionDecision", "VariantProfile", + "BACKENDS", "RuntimeClient", "RuntimePool", "InferenceRequest", "InferenceResponse", - "find_runtime_binary", + "extension_available", + "load_extension", "AdaptiveServer", "ServedRequest", "ServerStats", diff --git a/src/anytime_serving/serving/onnx_runtime.py b/src/anytime_serving/serving/onnx_runtime.py index 02acb0d..3ff8d96 100644 --- a/src/anytime_serving/serving/onnx_runtime.py +++ b/src/anytime_serving/serving/onnx_runtime.py @@ -1,28 +1,30 @@ -"""Client for the C++ ONNX runtime worker. - -The C++ binary (``runtime_cpp/``) accepts line-delimited JSON requests on -stdin and writes line-delimited JSON responses to stdout. This module wraps -that protocol and provides a worker pool for concurrent dispatch. - -A pure-Python fallback worker is also provided. It runs ONNX Runtime in-process -and is used when the C++ binary has not been built (e.g. on CI without a C++ -toolchain). The protocol it implements is identical, so unit tests exercise -the same code path the production server uses. +"""Client for the inference runtime. + +Two backends implement one interface, both taking and returning numpy arrays: + +``extension`` + The ``anytime_runtime`` pybind11 module, which runs ONNX Runtime in this + process. Tensors are borrowed rather than copied and the GIL is released + around inference, so a pool of workers runs concurrently. This is the serving + path. +``python`` + ONNX Runtime through its own Python wheel. The reference implementation the + tests compare the extension against, and the fallback where the extension has + not been built. + +The backend is chosen automatically unless ``backend=`` names one. Tests pin it +explicitly so a parity failure cannot be hidden by a silent fallback. """ from __future__ import annotations -import base64 -import json -import os import queue -import subprocess import threading import time import uuid from dataclasses import dataclass, field from pathlib import Path -from typing import IO +from typing import Any import numpy as np @@ -30,6 +32,51 @@ LOGGER = get_logger("serving.onnx_runtime") +BACKENDS = ("extension", "python") + +_extension: Any | None = None + + +def load_extension() -> Any: + """Import ``anytime_runtime``, refusing a version mismatch. + + The extension and the ``onnxruntime`` wheel are two independent copies of + ONNX Runtime loaded into one process. Stage 1 built the C++ worker against + 1.20.1 while profiling against the 1.26.0 wheel and measured DistilBERT at + 98.9 ms versus 13.0 ms inside ``session->Run()``; every service time the + planner used was wrong by almost an order of magnitude, and nothing failed. + The build enforces this equality at configure time, and this is the gate that + still applies when a built extension is carried into an environment with a + different wheel. + """ + global _extension + if _extension is not None: + return _extension + + import anytime_runtime + import onnxruntime + + linked = anytime_runtime.onnxruntime_version() + installed = onnxruntime.__version__ + if linked != installed: + raise RuntimeError( + f"anytime_runtime links ONNX Runtime {linked} but the installed wheel is " + f"{installed}. Both are loaded into this process, and a mismatch measured " + f"a 7.6x difference in inference time during Stage 1. Rebuild the " + f"extension against the current wheel: pip install -e . --no-cache-dir" + ) + _extension = anytime_runtime + return _extension + + +def extension_available() -> bool: + """Whether the in-process extension can be imported and matches the wheel.""" + try: + load_extension() + except (ImportError, RuntimeError): + return False + return True + @dataclass class InferenceRequest: @@ -60,117 +107,76 @@ def feed(self, default_input_name: str) -> dict[str, np.ndarray]: @dataclass class InferenceResponse: + """Result of one inference. + + Under the extension backend ``logits`` is a view over the buffer ONNX Runtime + allocated, not a copy, so holding the response holds that buffer. Copy it if + it needs to outlive the request. + """ + request_id: str logits: np.ndarray runtime_latency_ms: float wall_latency_ms: float -def _encode(array: np.ndarray) -> dict: - contiguous = np.ascontiguousarray(array) - return { - "shape": list(contiguous.shape), - "dtype": str(contiguous.dtype), - "data": base64.b64encode(contiguous.tobytes()).decode("ascii"), - } +class _RuntimeBackend: + """Runs one variant over a name -> tensor mapping. + Implementations return the first graph output and the time spent inside + inference, excluding anything the client adds around it. -def _decode(payload: dict) -> np.ndarray: - raw = base64.b64decode(payload["data"]) - dtype = np.dtype(payload["dtype"]) - return np.frombuffer(raw, dtype=dtype).reshape(payload["shape"]) + Anything meaning "the runtime could not serve this request" -- an unknown + variant, a missing declared input -- raises ``RuntimeError``. A malformed + argument, such as a dtype the engine does not accept, raises ``ValueError``. + """ + name = "abstract" -class _RuntimeBackend: - """Abstract backend that takes a request dict and returns a response dict.""" - - def submit(self, request: dict) -> dict: + def infer(self, variant: str, feeds: dict[str, np.ndarray]) -> tuple[np.ndarray, float]: raise NotImplementedError def close(self) -> None: raise NotImplementedError -class _SubprocessBackend(_RuntimeBackend): - """Backend backed by the C++ binary, communicating via stdin/stdout.""" - - def __init__(self, binary: Path, model_paths: dict[str, Path]) -> None: - args = [str(binary)] - for variant, path in model_paths.items(): - args.extend(["--model", f"{variant}={path}"]) - self._process = subprocess.Popen( - args, - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - bufsize=1, - text=True, - ) - # Popen types the pipes as Optional. All three were requested above, so - # bind them once to non-optional handles rather than narrowing at every - # use site. - assert self._process.stdin is not None - assert self._process.stdout is not None - assert self._process.stderr is not None - self._stdin: IO[str] = self._process.stdin - self._stdout: IO[str] = self._process.stdout - self._stderr: IO[str] = self._process.stderr - - # Drain handshake line (the C++ binary prints "ready" once both models - # have loaded). - handshake = self._stdout.readline().strip() - if handshake != "ready": - stderr = self._stderr.read() - raise RuntimeError(f"C++ runtime failed to start: {handshake!r} ({stderr})") - - def submit(self, request: dict) -> dict: - line = json.dumps(request) + "\n" - self._stdin.write(line) - self._stdin.flush() - response_line = self._stdout.readline() - if not response_line: - stderr = self._stderr.read() - raise RuntimeError(f"C++ runtime closed unexpectedly: {stderr}") - response = json.loads(response_line) - # The worker reports a recoverable problem (an unknown variant, say) as an - # error field and stays alive for the next request. - if "error" in response: - raise RuntimeError( - f"C++ runtime rejected request {response.get('request_id', '?')}: " - f"{response['error']}" - ) - return response +class _ExtensionBackend(_RuntimeBackend): + """In-process ONNX Runtime through the ``anytime_runtime`` extension.""" + + name = "extension" + + def __init__(self, model_paths: dict[str, Path]) -> None: + extension = load_extension() + self._engine = extension.Engine([(v, str(p)) for v, p in model_paths.items()]) + self._variants = frozenset(self._engine.variants) + + def infer(self, variant: str, feeds: dict[str, np.ndarray]) -> tuple[np.ndarray, float]: + if variant not in self._variants: + raise RuntimeError(f"unknown variant {variant!r}; loaded: {sorted(self._variants)}") + outputs, latency_ms = self._engine.run(variant, feeds) + return outputs[0], float(latency_ms) def close(self) -> None: - if self._process.poll() is None: - try: - self._stdin.close() - except OSError: - pass - self._process.terminate() - try: - self._process.wait(timeout=2.0) - except subprocess.TimeoutExpired: - self._process.kill() - # Popen leaves stdout/stderr open after the child exits; close them so a - # long-lived pool does not accumulate file descriptors. - for stream in (self._stdin, self._stdout, self._stderr): - try: - stream.close() - except OSError: - pass + # Dropping the engine releases the sessions and their arenas. + self._engine = None class _PythonBackend(_RuntimeBackend): - """Pure-Python ONNX Runtime backend. Used when the C++ binary is absent.""" + """ONNX Runtime through its Python wheel. The reference implementation.""" + + name = "python" def __init__(self, model_paths: dict[str, Path]) -> None: import onnxruntime as ort # local import so the dep is optional at import time options = ort.SessionOptions() options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL self._sessions = { - variant: ort.InferenceSession(str(path), sess_options=options) + variant: ort.InferenceSession( + str(path), sess_options=options, providers=["CPUExecutionProvider"] + ) for variant, path in model_paths.items() } self._declared_inputs = { @@ -178,72 +184,76 @@ def __init__(self, model_paths: dict[str, Path]) -> None: for variant, session in self._sessions.items() } - def submit(self, request: dict) -> dict: - variant = request["variant"] - session = self._sessions[variant] - # Variants of the same task can declare different inputs, so the caller - # sends the union and each graph takes the subset it declares. Mirrors the - # filtering in runtime_cpp/src/main.cpp. + def infer(self, variant: str, feeds: dict[str, np.ndarray]) -> tuple[np.ndarray, float]: + if variant not in self._sessions: + raise RuntimeError(f"unknown variant {variant!r}; loaded: {sorted(self._sessions)}") + # Variants of one task can declare different inputs, so the caller sends + # the union and each graph takes the subset it declares. Mirrors the + # filtering in runtime/src/engine.cpp. declared = self._declared_inputs[variant] - feeds = { - name: _decode(payload) - for name, payload in request["inputs"].items() - if name in declared - } - missing = declared - feeds.keys() + fed = {name: tensor for name, tensor in feeds.items() if name in declared} + missing = declared - fed.keys() if missing: - raise RuntimeError(f"request is missing inputs {sorted(missing)} for {variant!r}") + raise RuntimeError(f"variant {variant!r} is missing input(s): {sorted(missing)}") start = time.perf_counter() - outputs = session.run(None, feeds) + outputs = self._sessions[variant].run(None, fed) elapsed_ms = (time.perf_counter() - start) * 1000.0 - return { - "request_id": request["request_id"], - "logits": _encode(outputs[0]), - "latency_ms": elapsed_ms, - } + return outputs[0], elapsed_ms def close(self) -> None: self._sessions.clear() +def _make_backend(model_paths: dict[str, Path], requested: str | None) -> _RuntimeBackend: + """Build the requested backend, or pick one.""" + if requested is not None: + if requested not in BACKENDS: + raise ValueError(f"backend must be one of {BACKENDS}, got {requested!r}") + if requested == "extension": + return _ExtensionBackend(model_paths) + return _PythonBackend(model_paths) + + if extension_available(): + return _ExtensionBackend(model_paths) + LOGGER.warning( + "anytime_runtime is unavailable; falling back to the Python backend. Measured " + "service times will not reflect the serving path. Build the extension with: " + "pip install -e ." + ) + return _PythonBackend(model_paths) + + class RuntimeClient: - """Single inference worker, backed by either the C++ binary or ONNX Runtime.""" + """Single inference worker.""" def __init__( self, model_paths: dict[str, Path], *, - binary: Path | None = None, + backend: str | None = None, input_name: str = "input", ) -> None: if not model_paths: raise ValueError("model_paths must be non-empty") self._input_name = input_name - if binary is not None and binary.exists(): - LOGGER.info("Using C++ runtime binary at %s", binary) - self._backend: _RuntimeBackend = _SubprocessBackend(binary, model_paths) - else: - if binary is not None: - LOGGER.warning("C++ binary %s not found; falling back to Python backend", binary) - self._backend = _PythonBackend(model_paths) + self._backend = _make_backend(model_paths, backend) self._lock = threading.Lock() + @property + def backend_name(self) -> str: + """Which backend is serving. Recorded alongside every measurement.""" + return self._backend.name + def infer(self, request: InferenceRequest) -> InferenceResponse: - payload = { - "request_id": request.request_id, - "variant": request.variant, - "inputs": { - name: _encode(tensor) for name, tensor in request.feed(self._input_name).items() - }, - } + feeds = request.feed(self._input_name) start = time.perf_counter() with self._lock: - response = self._backend.submit(payload) + logits, runtime_latency_ms = self._backend.infer(request.variant, feeds) wall_ms = (time.perf_counter() - start) * 1000.0 return InferenceResponse( - request_id=response["request_id"], - logits=_decode(response["logits"]), - runtime_latency_ms=float(response["latency_ms"]), + request_id=request.request_id, + logits=logits, + runtime_latency_ms=runtime_latency_ms, wall_latency_ms=wall_ms, ) @@ -265,13 +275,13 @@ def __init__( size: int, model_paths: dict[str, Path], *, - binary: Path | None = None, + backend: str | None = None, input_name: str = "input", ) -> None: if size <= 0: raise ValueError("size must be positive") self._clients: list[RuntimeClient] = [ - RuntimeClient(model_paths, binary=binary, input_name=input_name) for _ in range(size) + RuntimeClient(model_paths, backend=backend, input_name=input_name) for _ in range(size) ] self._free: queue.Queue[RuntimeClient] = queue.Queue() for client in self._clients: @@ -285,6 +295,11 @@ def size(self) -> int: """ return len(self._clients) + @property + def backend_name(self) -> str: + """Which backend the workers use.""" + return self._clients[0].backend_name if self._clients else "none" + def infer(self, request: InferenceRequest) -> InferenceResponse: client = self._free.get() try: @@ -302,45 +317,3 @@ def __enter__(self) -> RuntimePool: def __exit__(self, exc_type, exc, tb) -> None: self.close() - - -def _build_dir_candidates(root: Path) -> tuple[Path, ...]: - build = root / "runtime_cpp" / "build" - return ( - build / "anytime_runtime", - build / "Release" / "anytime_runtime.exe", - ) - - -def find_runtime_binary() -> Path | None: - """Locate the compiled C++ runtime, if it exists. - - ``ANYTIME_RUNTIME_BIN`` takes precedence and is the only reliable option for - a non-editable install, where the package no longer lives inside the source - tree. Otherwise search upward from this file (covering an editable install) - and from the working directory (covering a plain checkout). - """ - env_override = os.environ.get("ANYTIME_RUNTIME_BIN") - if env_override: - candidate = Path(env_override) - return candidate if candidate.exists() else None - - roots: list[Path] = [] - # Walk up from this module: src/anytime_serving/serving/ -> repo root. Done by - # search rather than a fixed parent index so moving the package does not - # silently break resolution. - here = Path(__file__).resolve() - roots.extend(here.parents) - cwd = Path.cwd().resolve() - roots.append(cwd) - roots.extend(cwd.parents) - - seen: set[Path] = set() - for root in roots: - if root in seen: - continue - seen.add(root) - for candidate in _build_dir_candidates(root): - if candidate.exists(): - return candidate - return None diff --git a/tests/test_runtime_engine.py b/tests/test_runtime_engine.py new file mode 100644 index 0000000..8e14db3 --- /dev/null +++ b/tests/test_runtime_engine.py @@ -0,0 +1,371 @@ +"""Tests for the in-process engine, against the reference implementation. + +The point of most of this file is cross-backend agreement. The extension runs +ONNX Runtime in this process over borrowed buffers; the ``python`` backend runs the +same graph through the onnxruntime wheel. Both must produce bitwise identical +output, because they are the same library doing the same work and any difference +means the tensor path is corrupting something. + +The Stage 1 subprocess worker was the third backend here until it was removed. It +agreed bitwise with the extension on DistilBERT, which is what justified deleting +it: a replacement should be checked against the thing it replaces before that +thing goes away. + +Tests needing the compiled extension are skipped when it is absent. Build it with: + + pip install -e . +""" + +import os +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +onnx = pytest.importorskip("onnx", reason="onnx is required to build the fixture graphs") +from onnx import TensorProto, helper, numpy_helper # noqa: E402 + +from anytime_serving.serving.onnx_runtime import ( # noqa: E402 + BACKENDS, + InferenceRequest, + RuntimeClient, + RuntimePool, + extension_available, + load_extension, +) + +requires_extension = pytest.mark.skipif( + not extension_available(), reason="anytime_runtime is not built; see module docstring" +) + +# Fixed weights, so every backend is asked to compute exactly the same thing and +# the comparison is over arithmetic rather than over random inputs. +WEIGHT = np.arange(16, dtype=np.float32).reshape(4, 4) / 8.0 - 1.0 +BIAS = np.array([0.25, -0.5, 0.125, 1.0], dtype=np.float32) + + +def _build_graph(path: Path) -> None: + """A graph with real arithmetic in it. + + An Identity graph would round-trip a payload without exercising a kernel, so + it cannot show that two backends compute the same result. MatMul, Add, and + Relu do. + """ + inp = helper.make_tensor_value_info("input", TensorProto.FLOAT, [None, 4]) + out = helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, 4]) + nodes = [ + helper.make_node("MatMul", ["input", "weight"], ["projected"]), + helper.make_node("Add", ["projected", "bias"], ["biased"]), + helper.make_node("Relu", ["biased"], ["logits"]), + ] + graph = helper.make_graph( + nodes, + "affine_relu", + [inp], + [out], + initializer=[ + numpy_helper.from_array(WEIGHT, "weight"), + numpy_helper.from_array(BIAS, "bias"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 14)], ir_version=8) + onnx.save(model, str(path)) + + +def _build_two_input_graph(path: Path) -> None: + """Declares two inputs, so input filtering can be tested.""" + left = helper.make_tensor_value_info("input", TensorProto.FLOAT, [None, 4]) + right = helper.make_tensor_value_info("scale", TensorProto.FLOAT, [None, 4]) + out = helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, 4]) + graph = helper.make_graph( + [helper.make_node("Mul", ["input", "scale"], ["logits"])], + "scaled", + [left, right], + [out], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 14)], ir_version=8) + onnx.save(model, str(path)) + + +@pytest.fixture(scope="module") +def graphs(): + with tempfile.TemporaryDirectory() as tmp: + affine = Path(tmp) / "affine.onnx" + two_input = Path(tmp) / "two_input.onnx" + _build_graph(affine) + _build_two_input_graph(two_input) + yield {"affine": affine, "two_input": two_input} + + +@pytest.fixture(scope="module") +def model_paths(graphs): + return {"fp32": graphs["affine"], "int8": graphs["affine"]} + + +def _available_backends() -> list[str]: + available = ["python"] + if extension_available(): + available.append("extension") + return available + + +def _client(model_paths, backend): + return RuntimeClient(model_paths, backend=backend) + + +def _expected(data: np.ndarray) -> np.ndarray: + return np.maximum(data @ WEIGHT + BIAS, 0.0) + + +# --- the premise the comparisons rest on ------------------------------------ + + +def test_required_backends_are_actually_present(): + """Fail rather than skip where a backend is supposed to exist. + + Most tests here skip a backend that is not built, which is right locally but + would let a CI job quietly compare one backend against itself and report + success. ``ANYTIME_REQUIRE_BACKENDS`` names the backends an environment + promises to provide, and CI sets it so the parity comparison cannot decay into + a no-op without anyone noticing. + """ + required = os.environ.get("ANYTIME_REQUIRE_BACKENDS", "").strip() + if not required: + pytest.skip("ANYTIME_REQUIRE_BACKENDS is not set") + + requested = [name.strip() for name in required.split(",") if name.strip()] + unknown = sorted(set(requested) - set(BACKENDS)) + assert not unknown, f"ANYTIME_REQUIRE_BACKENDS names unknown backend(s): {unknown}" + + available = _available_backends() + missing = sorted(set(requested) - set(available)) + assert not missing, ( + f"ANYTIME_REQUIRE_BACKENDS demands {requested} but only {available} are " + f"available. The cross-backend comparison would have skipped instead of " + f"failing." + ) + + +# --- version equality ------------------------------------------------------- + + +@requires_extension +def test_extension_links_the_installed_onnxruntime(): + """The two ONNX Runtimes in this process must be the same version. + + Enforced at CMake configure time as well. This is the gate that still applies + when a built extension is carried into an environment with a different wheel, + which is how the Stage 1 mismatch would have been caught immediately. + """ + import onnxruntime + + extension = load_extension() + assert extension.onnxruntime_version() == onnxruntime.__version__ + + +@requires_extension +def test_extension_reports_its_compiled_api_version(): + extension = load_extension() + assert extension.ort_api_version() > 0 + + +# --- cross-backend agreement ------------------------------------------------ + + +def test_every_backend_computes_the_same_result(model_paths): + """Bitwise agreement across every backend built in this environment.""" + backends = _available_backends() + if len(backends) < 2: + pytest.skip(f"only the {backends[0]!r} backend is available; nothing to compare") + + rng = np.random.default_rng(0) + data = rng.standard_normal((3, 4)).astype(np.float32) + request = InferenceRequest(variant="fp32", data=data) + + results = {} + for backend in backends: + with _client(model_paths, backend) as client: + assert client.backend_name == backend + results[backend] = client.infer(request).logits + + # Against the arithmetic, then against each other, so a shared error in two + # backends cannot pass as agreement. + reference = _expected(data) + for backend, logits in results.items(): + np.testing.assert_allclose(logits, reference, rtol=1e-6, atol=1e-6) + for backend in backends[1:]: + np.testing.assert_array_equal( + results[backend], + results[backends[0]], + err_msg=f"{backend} disagrees with {backends[0]} bitwise", + ) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_backend_round_trips_a_batch(model_paths, backend): + if backend not in _available_backends(): + pytest.skip(f"the {backend!r} backend is not available here") + rng = np.random.default_rng(1) + data = rng.standard_normal((5, 4)).astype(np.float32) + with _client(model_paths, backend) as client: + response = client.infer(InferenceRequest(variant="fp32", data=data)) + np.testing.assert_allclose(response.logits, _expected(data), rtol=1e-6, atol=1e-6) + assert response.runtime_latency_ms >= 0.0 + assert response.wall_latency_ms >= response.runtime_latency_ms - 1e-6 + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_backend_routes_between_variants(model_paths, backend): + if backend not in _available_backends(): + pytest.skip(f"the {backend!r} backend is not available here") + data = np.full((1, 4), 0.5, dtype=np.float32) + with _client(model_paths, backend) as client: + for variant in ("fp32", "int8"): + response = client.infer(InferenceRequest(variant=variant, data=data)) + np.testing.assert_allclose(response.logits, _expected(data), rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_backend_reports_unknown_variant_and_stays_usable(model_paths, backend): + """An unknown variant raises without taking the worker down.""" + if backend not in _available_backends(): + pytest.skip(f"the {backend!r} backend is not available here") + data = np.ones((1, 4), dtype=np.float32) + with _client(model_paths, backend) as client: + with pytest.raises(RuntimeError, match="unknown variant"): + client.infer(InferenceRequest(variant="does_not_exist", data=data)) + response = client.infer(InferenceRequest(variant="fp32", data=data)) + np.testing.assert_allclose(response.logits, _expected(data), rtol=1e-6, atol=1e-6) + + +# --- input handling --------------------------------------------------------- + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_backend_drops_undeclared_inputs(graphs, backend): + """Callers send the union of every variant's inputs; extras are dropped. + + Variants of one task declare different inputs, so a graph must take the subset + it declares rather than failing on the rest. + """ + if backend not in _available_backends(): + pytest.skip(f"the {backend!r} backend is not available here") + paths = {"affine": graphs["affine"]} + data = np.full((1, 4), 2.0, dtype=np.float32) + with _client(paths, backend) as client: + response = client.infer( + InferenceRequest( + variant="affine", + inputs={"input": data, "token_type_ids": np.zeros((1, 4), dtype=np.int64)}, + ) + ) + np.testing.assert_allclose(response.logits, _expected(data), rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_backend_rejects_a_missing_declared_input(graphs, backend): + """A partial feed is an error, not a run on whatever was supplied.""" + if backend not in _available_backends(): + pytest.skip(f"the {backend!r} backend is not available here") + paths = {"scaled": graphs["two_input"]} + with _client(paths, backend) as client: + with pytest.raises(RuntimeError, match="missing input"): + client.infer( + InferenceRequest( + variant="scaled", inputs={"input": np.ones((1, 4), dtype=np.float32)} + ) + ) + + +@requires_extension +def test_extension_accepts_a_non_contiguous_input(model_paths): + """A strided array is made contiguous rather than misread. + + ONNX Runtime reads the buffer directly, so handing it a transposed view + without converting would silently transpose the arithmetic. + """ + rng = np.random.default_rng(2) + base = rng.standard_normal((4, 4)).astype(np.float32) + strided = base.T + assert not strided.flags.c_contiguous + + with _client(model_paths, "extension") as client: + from_strided = client.infer(InferenceRequest(variant="fp32", data=strided)).logits + from_copy = client.infer( + InferenceRequest(variant="fp32", data=np.ascontiguousarray(strided)) + ).logits + np.testing.assert_array_equal(from_strided, from_copy) + np.testing.assert_allclose(from_strided, _expected(strided), rtol=1e-6, atol=1e-6) + + +@requires_extension +def test_extension_rejects_an_unsupported_dtype(model_paths): + """An unaccepted dtype fails loudly rather than being reinterpreted.""" + with _client(model_paths, "extension") as client: + with pytest.raises(ValueError, match="float16|does not accept"): + client.infer(InferenceRequest(variant="fp32", data=np.ones((1, 4), dtype=np.float16))) + + +# --- the properties the extension exists for -------------------------------- + + +@requires_extension +def test_extension_output_is_a_view_not_a_copy(model_paths): + """Outputs borrow ONNX Runtime's buffer. + + Copying every output would undo the reason for running in-process, so assert + the array does not own its memory and is kept alive by its base. + """ + data = np.ones((2, 4), dtype=np.float32) + with _client(model_paths, "extension") as client: + logits = client.infer(InferenceRequest(variant="fp32", data=data)).logits + assert not logits.flags.owndata + assert logits.base is not None + # The buffer has to stay valid once the client is closed, since the response + # outlives the call that produced it. + np.testing.assert_allclose(logits, _expected(data), rtol=1e-6, atol=1e-6) + + +@requires_extension +def test_extension_releases_the_gil_during_inference(model_paths): + """Concurrent pool workers must actually overlap. + + Holding the GIL through inference would serialise the pool, and the M/M/c + model the admission controller uses would describe a machine that does not + exist. Correctness under concurrency is asserted here; the throughput claim + itself belongs to the benchmarks, not to a timing-sensitive unit test. + """ + from concurrent.futures import ThreadPoolExecutor + + rng = np.random.default_rng(3) + requests = [ + InferenceRequest( + variant="int8" if i % 2 else "fp32", + data=rng.standard_normal((1, 4)).astype(np.float32), + ) + for i in range(16) + ] + with RuntimePool(size=4, model_paths=model_paths, backend="extension") as pool: + with ThreadPoolExecutor(max_workers=4) as executor: + responses = list(executor.map(pool.infer, requests)) + + for request, response in zip(requests, responses, strict=True): + assert response.request_id == request.request_id + assert request.data is not None + np.testing.assert_allclose(response.logits, _expected(request.data), rtol=1e-6, atol=1e-6) + + +# --- backend selection ------------------------------------------------------ + + +@requires_extension +def test_the_extension_is_preferred_when_available(model_paths): + """Automatic selection picks the in-process path, not the fallback.""" + with RuntimeClient(model_paths) as client: + assert client.backend_name == "extension" + + +def test_an_unknown_backend_name_is_rejected(model_paths): + with pytest.raises(ValueError, match="backend must be one of"): + RuntimeClient(model_paths, backend="does_not_exist") diff --git a/tests/test_runtime_subprocess.py b/tests/test_runtime_subprocess.py deleted file mode 100644 index b20cc24..0000000 --- a/tests/test_runtime_subprocess.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Protocol tests against the compiled C++ runtime. - -The Python fallback backend and the C++ subprocess backend implement the same -request/response protocol, but only the fallback is exercised by the rest of the -suite. These tests run the real binary so the framing, base64 codec, and -handshake are covered too. - -Skipped when the binary has not been built. Build it with: - - cmake -S runtime_cpp -B runtime_cpp/build -DCMAKE_BUILD_TYPE=Release \\ - -DONNXRUNTIME_ROOT_PATH=/path/to/onnxruntime - cmake --build runtime_cpp/build -j -""" - -import tempfile -from pathlib import Path - -import numpy as np -import pytest - -onnx = pytest.importorskip("onnx", reason="onnx is required to build the fixture model") -from onnx import TensorProto, helper # noqa: E402 - -from anytime_serving.serving.onnx_runtime import ( # noqa: E402 - InferenceRequest, - RuntimeClient, - RuntimePool, - find_runtime_binary, -) - -pytestmark = pytest.mark.needs_runtime - -BINARY = find_runtime_binary() -requires_binary = pytest.mark.skipif( - BINARY is None, reason="C++ runtime binary not built; see module docstring" -) - - -def _build_identity_model(path: Path, dim: int = 4) -> None: - inp = helper.make_tensor_value_info("input", TensorProto.FLOAT, [None, dim]) - out = helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, dim]) - graph = helper.make_graph( - [helper.make_node("Identity", ["input"], ["logits"])], "identity", [inp], [out] - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 14)], ir_version=8) - onnx.save(model, str(path)) - - -@pytest.fixture(scope="module") -def model_paths(): - with tempfile.TemporaryDirectory() as tmp: - fp32 = Path(tmp) / "fp32.onnx" - int8 = Path(tmp) / "int8.onnx" - _build_identity_model(fp32) - _build_identity_model(int8) - yield {"fp32": fp32, "int8": int8} - - -@requires_binary -def test_subprocess_round_trip_preserves_payload(model_paths): - with RuntimeClient(model_paths, binary=BINARY) as client: - data = np.arange(8, dtype=np.float32).reshape(2, 4) - response = client.infer(InferenceRequest(variant="fp32", data=data)) - np.testing.assert_array_almost_equal(response.logits, data) - assert response.runtime_latency_ms >= 0.0 - assert response.wall_latency_ms >= response.runtime_latency_ms - 1e-6 - - -@requires_binary -def test_subprocess_base64_codec_survives_all_byte_values(model_paths): - """Exercise the hand-rolled base64 codec across padding lengths. - - Payload sizes of 4, 5, and 6 floats give byte counts with remainders 0, 1, - and 2 mod 3, covering every padding branch in the encoder. - """ - try: - for count in (4, 5, 6, 8, 16): - data = np.linspace(-1e6, 1e6, count, dtype=np.float32).reshape(1, count) - # The graph pins its input width, so rebuild it per payload size. - _build_identity_model(model_paths["fp32"], dim=count) - with RuntimeClient({"fp32": model_paths["fp32"]}, binary=BINARY) as sized: - response = sized.infer(InferenceRequest(variant="fp32", data=data)) - np.testing.assert_array_equal(response.logits, data) - finally: - # Restore the shared module-scoped fixture for the other tests. - _build_identity_model(model_paths["fp32"]) - - -@requires_binary -def test_subprocess_routes_between_variants(model_paths): - with RuntimeClient(model_paths, binary=BINARY) as client: - for variant in ("fp32", "int8"): - data = np.full((1, 4), 3.5, dtype=np.float32) - response = client.infer(InferenceRequest(variant=variant, data=data)) - np.testing.assert_array_almost_equal(response.logits, data) - - -@requires_binary -def test_subprocess_pool_serves_concurrent_requests(model_paths): - from concurrent.futures import ThreadPoolExecutor - - with RuntimePool(size=2, model_paths=model_paths, binary=BINARY) as pool: - rng = np.random.default_rng(0) - requests = [ - InferenceRequest( - variant="int8" if i % 2 else "fp32", - data=rng.standard_normal((1, 4)).astype(np.float32), - ) - for i in range(8) - ] - with ThreadPoolExecutor(max_workers=4) as executor: - responses = list(executor.map(pool.infer, requests)) - for request, response in zip(requests, responses, strict=True): - np.testing.assert_array_almost_equal(response.logits, request.data) - assert response.request_id == request.request_id - - -@requires_binary -def test_subprocess_reports_unknown_variant_and_stays_usable(model_paths): - """An unknown variant raises a descriptive error without killing the worker.""" - with RuntimeClient(model_paths, binary=BINARY) as client: - with pytest.raises(RuntimeError, match="unknown variant"): - client.infer( - InferenceRequest(variant="does_not_exist", data=np.zeros((1, 4), np.float32)) - ) - # The worker reports the problem and continues serving. - data = np.ones((1, 4), dtype=np.float32) - response = client.infer(InferenceRequest(variant="fp32", data=data)) - np.testing.assert_array_almost_equal(response.logits, data) - - -def test_find_runtime_binary_honours_env_override(monkeypatch, tmp_path): - """The env override wins, and a missing path resolves to None.""" - monkeypatch.setenv("ANYTIME_RUNTIME_BIN", str(tmp_path / "absent")) - assert find_runtime_binary() is None - - present = tmp_path / "anytime_runtime" - present.write_text("") - monkeypatch.setenv("ANYTIME_RUNTIME_BIN", str(present)) - assert find_runtime_binary() == present