Add cleanup path for staged jax_callable graph captures

[JaimeFine]

Description

jax_callable() currently provides clear_jax_callable_graph_cache() for graphs cached by JaxCallableGraphMode.WARP, but I do not see an equivalent cleanup path for resources retained by the staged graph modes.

In warp/_src/jax/ffi.py, JaxCallableGraphMode.WARP_STAGED and JaxCallableGraphMode.WARP_STAGED_EX keep per-call state on FfiCallDesc, including:

• call_desc.capture
• call_desc.input_staging_arrays
• call_desc.output_staging_arrays
• call_desc.static_staging_arrays

The code comments note that there should be a way to free these retained objects.

Would the maintainers be open to adding a cleanup path for staged jax_callable() captures? One possible direction would be extending clear_jax_callable_graph_cache() so it also releases staged per-call captures and staging buffers, causing the next execution to recapture. If a separate API would fit the intended lifetime model better, I would be happy to follow that direction instead.

Context

I noticed the TODO comments in warp/_src/jax/ffi.py saying “we should have a way of freeing this,” then traced the surrounding staged jax_callable() code to understand what is being retained.

It looks like WARP_STAGED and WARP_STAGED_EX intentionally keep FfiCallDesc.capture and staging arrays alive for replay correctness, while JaxCallableGraphMode.WARP has clear_jax_callable_graph_cache() as a manual release path. I could not find an equivalent cleanup path for the staged per-call resources.

I also verified the behavior on a Jetson Orin with CUDA enabled using a small standalone repro (source):

Warp 1.17.0.dev0 initialized:
   CUDA Toolkit 12.6, Driver 12.6
   Devices:
     "cpu"      : "aarch64"
     "cuda:0"   : "Orin" (4 GiB, sm_87, mempool enabled)

For both WARP_STAGED and WARP_STAGED_EX, after first execution:

call descriptors: 1
graph cache size: 0
call_desc.capture is None: False
input_staging_arrays is None: False
output_staging_arrays is None: False
static_staging_arrays is None: False

After calling clear_jax_callable_graph_cache(jax_double), the staged resources remain retained:

graph cache size: 0
call_desc.capture is None: False
input_staging_arrays is None: False
output_staging_arrays is None: False
static_staging_arrays is None: False

Full log is here.

I’m not claiming this is necessarily a leak; the retention appears intentional for replay. The request is for a deterministic cleanup path analogous to the existing graph-cache cleanup API.

I’m opening this as a feature request/design question rather than a PR because the right API likely depends on the intended lifetime model for JAX-compiled call sites. If maintainers think this is useful, I’d be happy to help with a focused PR, tests, and docs after guidance on the preferred API shape.

Thanks!

[shi-eric]

Thanks for the detailed investigation and reproducer. WARP_STAGED and WARP_STAGED_EX keep their capture and staging allocations on the FFI call descriptor. clear_jax_callable_graph_cache() only clears the address-based cache used by WARP.

@nvlukasz left those TODOs because the cleanup choices were not obvious and needed more thought. A compiled JAX executable can continue using its call descriptor after the Python call that created it. Cleaning up the capture also means dealing with the staging buffers and memcpy state tied to it, including CUDA work that may still be running. The capture field also keeps CUDA modules loaded during JAX-managed capture, so clearing it too broadly can break that path.

We think this needs to be designed and implemented by the Warp team. An external patch would still leave most of the design and validation work with our small maintainer group, and we do not have the capacity to take that on.

We'll keep the issue open. If this is blocking your use of Warp, please tell us what you're trying to do and how the retained resources are affecting it. With a concrete use case, we can take another look at where this falls in our priorities.

[JaimeFine]

Thanks for the detailed explanation. That makes sense.

This is not currently blocking me; I found it while reading through the JAX FFI code and trying to understand possible contribution areas. I opened the issue because the TODOs looked like a real resource-lifetime question, and I wanted to avoid sending an uncoordinated PR in a design-heavy area.

I’ll leave this with the Warp team. If there is a smaller, well-scoped follow-up around tests, docs, or reproducer coverage that would be useful, I’d be happy to help.