feat: Add json+cuda_ipc array encoding for GPU-direct tensor transfer

[dionhaefner]

Relevant issue or PR

None.

Description of changes

Serving a Tesseract that computes on the GPU currently incurs a round trip through host memory on every call: the runtime copies the result off the device, serializes it (base64 or binref), ships it over HTTP, and the client copies it back onto the device. For large arrays exchanged between co-located GPU processes that copy dominates the call and defeats the point of keeping data on the device. This prevents e.g. most differentiable physics / solver-in-the-loop applications where the simulation kernel runs on the GPU and the outer loop runs in Python on the same machine.

This PR adds a json+cuda_ipc output format that instead passes GPU arrays by reference. When an endpoint returns an array exposing __cuda_array_interface__ (CuPy, PyTorch, JAX, Numba, anything implementing the protocol), the runtime exports a 64-byte cudaIpcMemHandle_t in place of the array data. Only the tiny handle travels over HTTP, while the data never leaves the GPU. Any non-GPU arrays in the same response still go out as base64, so mixed outputs keep working.

From an end user's perspective, it's one extra flag on the server plus the experimental opt-in:

with Tesseract.from_image(
    "my-gpu-tesseract",
    gpus=["all"],
    output_format="json+cuda_ipc",
    runtime_config={"enable_experimental_cuda_ipc": True},
) as t:
    result = t.apply({"a": a, "b": b, "s": 3.0})

# result["result"] is a custom struct that has `__cuda_array_interface__` and `__dlpack__` properties,
# so can be loaded by any compatible library (JAX, torch, CuPy, Numba, ...)
out = result["result"]

# copying to host works as expected
out_host = np.asarray(out)

Key design choices:

• ctypes rather than a framework dependency. The export/import path calls libcudart and libcuda directly (cudaIpcGetMemHandle / cudaIpcOpenMemHandle) instead of going through CuPy or Torch. That keeps the runtime free of any hard GPU-framework dependency and means it can export from any producer that implements the CUDA array interface, not just one blessed library. The cost is a small amount of ctypes plumbing and manual error-string handling.

• Opt-in and explicitly experimental/undocumented. The format is gated behind a new enable_experimental_cuda_ipc runtime config flag (TESSERACT_ENABLE_EXPERIMENTAL_CUDA_IPC=1) and is omitted from available_formats() unless enabled. This feature needs proper support in Tesseract-JAX/Tesseract-Torch, a lot more testing, and careful documentation before it can be considered stable.

• Keepalive buffer on the server. CUDA IPC is inherently cross-process (a process can't open a handle it exported itself), and a pooled allocator may recycle an exported buffer before the consumer ever opens it, silently corrupting the result. The server therefore pins the arrays it exported for the current request and only releases them at the start of the next one. This bounds pinned GPU memory to a single request's worth of exports while guaranteeing each export stays alive long enough for a serial client to copy it out. The trade-off is that it assumes a serial request pattern. It is not safe for a client pipelining overlapping requests against the same server.

• Staging fallback for VMM allocations. Arrays backed by CUDA's virtual-memory-management allocator (notably JAX) cannot be exported through the legacy cudaIpcGetMemHandle path. Rather than reimplement the VMM export machinery, such arrays are first copied into a plain cudaMalloc staging buffer, which is then exported and released alongside the normal keepalive ring. We could add proper VMM support later and save the extra copy, but this is a pragmatic first step that works for most use cases.

Testing done

• tests/test_cuda_ipc_cpu.py: GPU-free tests of the orchestration (payload assembly, offset arithmetic, device-ordinal detection, export registry, the serve-side release hook, the --ipc=host wiring, and the CLI guard) by feeding fake __cuda_array_interface__ objects and monkeypatching the ctypes/CUDA wrappers. These run on ordinary CI runners.
• tests/test_cuda_ipc.py: GPU tests that drive real cross-process IPC (basic round trip, dtypes, non-zero offsets, serial reuse, Torch/JAX interop, VMM and staging fallbacks).
• tests/endtoend_tests/test_serving_gpu.py: end-to-end tests that build a GPU Tesseract image and serve it in a container with --gpus all / --ipc=host, round-tripping device memory through a genuine handle.

The GPU tests run on a new tesseract-ci-gpu job (marked gpu, skipped on GPU-less runners) wired into the all-ok gate. CI passes.

[codecov]

Codecov Report

❌ Patch coverage is 79.52586% with 95 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.39%. Comparing base (ee72c73) to head (f7c4c3c).

Files with missing lines	Patch %	Lines
tesseract_core/runtime/cuda_ipc.py	79.77%	47 Missing and 25 partials ⚠️
tesseract_core/runtime/array_encoding.py	62.06%	8 Missing and 3 partials ⚠️
tesseract_core/sdk/tesseract.py	88.23%	1 Missing and 5 partials ⚠️
tesseract_core/runtime/file_interactions.py	76.92%	1 Missing and 2 partials ⚠️
tesseract_core/runtime/serve.py	40.00%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #588      +/-   ##
==========================================
- Coverage   78.46%   78.39%   -0.07%     
==========================================
  Files          40       41       +1     
  Lines        4857     5295     +438     
  Branches      794      866      +72     
==========================================
+ Hits         3811     4151     +340     
- Misses        733      797      +64     
- Partials      313      347      +34     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[PasteurBot]

Benchmark Results

Benchmarks use a no-op Tesseract to measure pure framework overhead.

🚀 1 faster, ⚠️ 0 slower, ✅ 53 unchanged

Notable changes

Benchmark	Baseline	Current	Change	Status
decoding/base64_100,000	0.663ms	0.575ms	-13.3%	🚀 faster

Full results

Benchmark	Baseline	Current	Change	Status
api/apply_1,000	0.584ms	0.598ms	+2.5%	✅
api/apply_100,000	0.585ms	0.594ms	+1.5%	✅
api/apply_10,000,000	0.588ms	0.595ms	+1.3%	✅
cli/apply_1,000	1684.407ms	1696.496ms	+0.7%	✅
cli/apply_100,000	1722.848ms	1711.796ms	-0.6%	✅
cli/apply_10,000,000	1782.628ms	1737.125ms	-2.6%	✅
decoding/base64_1,000	0.037ms	0.035ms	-5.3%	✅
decoding/base64_100,000	0.663ms	0.575ms	-13.3%	🚀 faster
decoding/base64_10,000,000	69.701ms	70.124ms	+0.6%	✅
decoding/base64+lz4_1,000	0.041ms	0.041ms	-0.1%	✅
decoding/base64+lz4_100,000	0.666ms	0.635ms	-4.7%	✅
decoding/base64+lz4_10,000,000	117.229ms	114.886ms	-2.0%	✅
decoding/binref_1,000	0.204ms	0.208ms	+1.7%	✅
decoding/binref_100,000	0.240ms	0.243ms	+1.2%	✅
decoding/binref_10,000,000	11.446ms	11.517ms	+0.6%	✅
decoding/binref+lz4_1,000	0.212ms	0.213ms	+0.8%	✅
decoding/binref+lz4_100,000	0.308ms	0.315ms	+2.3%	✅
decoding/binref+lz4_10,000,000	41.173ms	40.619ms	-1.3%	✅
decoding/json_1,000	0.105ms	0.106ms	+1.6%	✅
decoding/json_100,000	8.954ms	8.719ms	-2.6%	✅
decoding/json_10,000,000	1067.653ms	1067.813ms	+0.0%	✅
encoding/base64_1,000	0.041ms	0.042ms	+2.8%	✅
encoding/base64_100,000	0.146ms	0.150ms	+2.6%	✅
encoding/base64_10,000,000	28.210ms	27.729ms	-1.7%	✅
encoding/base64+lz4_1,000	0.046ms	0.049ms	+5.3%	✅
encoding/base64+lz4_100,000	0.348ms	0.356ms	+2.4%	✅
encoding/base64+lz4_10,000,000	91.075ms	88.907ms	-2.4%	✅
encoding/binref_1,000	0.319ms	0.324ms	+1.8%	✅
encoding/binref_100,000	0.494ms	0.501ms	+1.4%	✅
encoding/binref_10,000,000	19.343ms	19.597ms	+1.3%	✅
encoding/binref+lz4_1,000	0.327ms	0.330ms	+1.0%	✅
encoding/binref+lz4_100,000	0.707ms	0.724ms	+2.4%	✅
encoding/binref+lz4_10,000,000	83.134ms	81.365ms	-2.1%	✅
encoding/json_1,000	0.148ms	0.149ms	+0.7%	✅
encoding/json_100,000	11.489ms	11.542ms	+0.5%	✅
encoding/json_10,000,000	1384.390ms	1394.506ms	+0.7%	✅
http/apply_1,000	3.205ms	3.246ms	+1.3%	✅
http/apply_100,000	9.080ms	9.158ms	+0.9%	✅
http/apply_10,000,000	796.888ms	799.929ms	+0.4%	✅
roundtrip/base64_1,000	0.089ms	0.090ms	+1.7%	✅
roundtrip/base64_100,000	0.739ms	0.738ms	-0.2%	✅
roundtrip/base64_10,000,000	99.545ms	99.901ms	+0.4%	✅
roundtrip/base64+lz4_1,000	0.098ms	0.100ms	+2.4%	✅
roundtrip/base64+lz4_100,000	0.982ms	1.014ms	+3.3%	✅
roundtrip/base64+lz4_10,000,000	208.016ms	205.363ms	-1.3%	✅
roundtrip/binref_1,000	0.545ms	0.552ms	+1.4%	✅
roundtrip/binref_100,000	0.749ms	0.756ms	+1.0%	✅
roundtrip/binref_10,000,000	31.917ms	31.917ms	-0.0%	✅
roundtrip/binref+lz4_1,000	0.561ms	0.566ms	+0.9%	✅
roundtrip/binref+lz4_100,000	1.016ms	1.047ms	+3.0%	✅
roundtrip/binref+lz4_10,000,000	124.862ms	123.042ms	-1.5%	✅
roundtrip/json_1,000	0.267ms	0.268ms	+0.3%	✅
roundtrip/json_100,000	19.772ms	20.050ms	+1.4%	✅
roundtrip/json_10,000,000	2451.283ms	2466.432ms	+0.6%	✅

• PR: feat: Add json+cuda_ipc array encoding for GPU-direct tensor transfer #588
• Runner: Linux 6.17.0-1022-azure x86_64

[nmheim]

Mostly nits except my last comment. It seems like the client is not releasing previously exported arrays?

[nmheim]

Neat, I like the context manager for the client!