[adenzhou1350]

[Integration] Release GIL during single memory registration

Description

TransferEngine.register_memory() and unregister_memory() currently keep the
Python GIL while their native registerLocalMemory() and
unregisterLocalMemory() calls execute. The batch wrappers already release the
GIL around the corresponding native calls.

This change applies the same scoped release to the single-buffer wrappers.
Python argument conversion remains before the C++ method, pybind11 reacquires
the GIL before converting the return value, and native address, length,
location, metadata, ordering, and error behavior are unchanged.

The regression test arms another Python thread before a bounded series of TCP
host-memory registration calls. It passes only when that thread can run during
both the single-register and single-unregister native spans. This is adjacent
to the GIL work in #300 and #3947 and the batch wrappers in #2191, but a current
open-PR search found no exact duplicate for these two wrappers.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Reshard (mooncake-reshard)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Historical original native comparison: base 9cfb33810629f91b07e913ad3fdaf478287f31b9, candidate
b41d8e7b1f16c803bac65b486fec65a8bdf3b9f1. The build commands and repeated control/candidate results below refer to that original comparison.

Current head 84a7d1ea: GitHub CI completed with 26 successful check runs and 3 skipped; no failures or pending checks. Both Python 3.10/3.12 wheel jobs, CTest and CI Gate succeeded. This does not establish GPU/RDMA runtime or performance.

The scoped TCP regression was also independently rerun on Linux using the CPython 3.12 wheel built from e1f87565; the native source was unchanged by this test-only follow-up.

The build and wheel commands below were run in a repository-devcontainer-derived
Linux image, with the repository mounted at /src:

Test commands:

cmake -S /src -B /build \
  -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DWITH_TE=ON \
  -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=ON \
  -DUSE_CUDA=OFF -DUSE_ETCD=OFF -DUSE_REDIS=OFF -DUSE_HTTP=ON
cmake --build /build --target engine transfer_engine_bench -j2

NON_CUDA_BUILD=1 PYTHON_VERSION=3.10 \
  OUTPUT_DIR=/evidence/dist ./scripts/build_wheel.sh

/usr/bin/python3.10 -m virtualenv /tmp/mooncake-gil-venv
source /tmp/mooncake-gil-venv/bin/activate
python -m pip install /evidence/dist/mooncake_transfer_engine_non_cuda-*.whl
python /src/mooncake-wheel/tests/test_transfer_engine_gil.py -v
python -m pip check

pre-commit run --files \
  mooncake-integration/transfer_engine/transfer_engine_py.cpp \
  mooncake-wheel/tests/test_transfer_engine_gil.py \
  scripts/test_installation.sh
git diff --check origin/main...HEAD

Test results:

• Unit tests pass
• Integration tests pass (TCP host memory)
• Manual testing done (exact unchanged-control/candidate comparison)

The candidate passed 10/10 isolated direct-module runs and 10/10 isolated runs
after fresh installation of its locally built non-CUDA wheel. The exact
unchanged control failed 3/3 because both wrappers retained the GIL. A separate
sealed probe observed both candidate wrappers release the GIL in 2/2 repeats
and both control wrappers retain it in 2/2 repeats; every registration call
returned zero.

The exact candidate configured, compiled, linked, built a repaired wheel,
installed into a fresh Python 3.10 environment, passed pip check, and had no
unresolved dynamic dependencies in the build container. PR-scoped pre-commit
passed.

No GPU or RDMA runtime was used. This PR claims only that the single-buffer
Python wrappers no longer retain the GIL during their native calls. It does not
claim lower native registration latency, RDMA throughput, GPU speedup, or
end-to-end workload improvement.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit on the files changed in this PR and all hooks pass
• I have updated the documentation (not applicable: no public API or configuration change)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (not applicable: 119 lines including tests)

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specified below)

OpenAI Codex assisted with source auditing, implementation, regression-test
design, and qualification. The human submitter remains responsible for
reviewing every changed line and defending the change before submission.

October 16 main refresh

Merged current main (5d77d34), resolving the registerMemory signature conflict by retaining the new remote_accessible parameter and forwarding it unchanged. The production delta remains the original two gil_scoped_release lines; the GIL regression and installation-test entry are retained. Generic/Ruff/format/codespell pre-commit hooks pass. The system formatter hook cannot launch /bin/bash on native Windows, so the same staged formatter was run through native Git Bash, scoped to the changed C++ file, and passed with clang-format 20.1.8. No fresh native TransferEngine/GIL integration result is claimed for this merged head; earlier validation above is historical. Current-head CI and code-owner review remain the handoff.

Current-head CTest blocker

The new CTest Python 3.12 job stops during C++ compilation, before tests run: scenario.cpp:1926 still calls BatchProbeKey(keys, TenantId), while master_service.h:264 now requires GrantLeasePolicy, keys, tenant. Current main 5d77d34 (the #4314 lease-selection change) has the same caller/signature; both files are byte-identical between this PR head and main. This is outside the two-line GIL-release production delta. Could the upstream caller/API mismatch be reconciled before the CTest rerun?

Failure log: https://github.com/kvcache-ai/Mooncake/actions/runs/38906750683/job/116994131977

No passing current-head native GIL integration result is claimed.

[adenzhou1350]

Hi @ShangmingCai — when you have a chance, could you take a focused look at this two-line wrapper change? The full non-CUDA build/wheel path and candidate/control GIL regression are documented above, and the current CI is green. Happy to adjust anything you spot.

[alogfans]

LGTM. cc @ShangmingCai

[adenzhou1350]

Thanks @alogfans for approving. The current checks are green (22 successful, 2 skipped), and GitHub still asks for a code-owner approval from @ShangmingCai, @stmatengss or @zxpdemonio. Could one of the owners take the final review when convenient?

The related integration PR #4104 is also green (22 successful, 2 skipped) and awaits @ShangmingCai or @alogfans. Both retain their separately documented scope and validation; please let me know if either needs changes before merge.

[adenzhou1350]

Updated in e1f8756. Merged current main (5d77d34), resolving the registerMemory signature conflict by retaining the new remote_accessible parameter and forwarding it unchanged. The production delta remains the original two gil_scoped_release lines; the GIL regression and installation-test entry are retained. Generic/Ruff/format/codespell pre-commit hooks pass. The system formatter hook cannot launch /bin/bash on native Windows, so the same staged formatter was run through native Git Bash, scoped to the changed C++ file, and passed with clang-format 20.1.8. No fresh native TransferEngine/GIL integration result is claimed for this merged head; earlier validation above is historical. Current-head CI and code-owner review remain the handoff.

[adenzhou1350]

The new CTest Python 3.12 job stops during C++ compilation, before tests run: scenario.cpp:1926 still calls BatchProbeKey(keys, TenantId), while master_service.h:264 now requires GrantLeasePolicy, keys, tenant. Current main 5d77d34 (the #4314 lease-selection change) has the same caller/signature; both files are byte-identical between this PR head and main. This is outside the two-line GIL-release production delta. Could the upstream caller/API mismatch be reconciled before the CTest rerun?

Failure log: https://github.com/kvcache-ai/Mooncake/actions/runs/38906750683/job/116994131977

No passing current-head native GIL integration result is claimed.

[adenzhou1350]

The wheel failure has a separate cause from the CTest build error above. Although the GIL test passes protocol="tcp", the wheel job logs auto-discovery of mana_0, selects RDMA, then fails batch registration with Cannot allocate memory / -202 before the final GIL assertions.

At this head, MC_FORCE_TCP=1 takes the supported TCP-only initialization path. I will scope that setting to this TCP host-memory test and retain the registration/GIL assertions. No fresh native pass is claimed yet: the downloaded exact-head CPython 3.12 wheel also needs CUDA 12 and other native dependencies absent from the available CPU sandbox.

[adenzhou1350]

Updated the test to force TCP while restoring the prior environment afterward. It passes on Linux with the cp312 wheel built from e1f8756; file-scoped pre-commit passes too. Native code is unchanged.

Current head 84a7d1ea: GitHub CI completed with 26 successful check runs and 3 skipped; no failures or pending checks. Both Python 3.10/3.12 wheel jobs, CTest and CI Gate succeeded. This does not establish GPU/RDMA runtime or performance.

[zxpdemonio]

LGTM.