Skip to content

fix(server): make native VMM registration ownership explicit - #416

Open
xiaguan wants to merge 1 commit into
feat/native-vmm-kv-clientfrom
codex/vmm-registration-resource
Open

fix(server): make native VMM registration ownership explicit#416
xiaguan wants to merge 1 commit into
feat/native-vmm-kv-clientfrom
codex/vmm-registration-resource

Conversation

@xiaguan

@xiaguan xiaguan commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

Summary

This is stacked on #414 and defines the server-side ownership contract needed by the OpenInfer consumer in openinfer-project/openinfer#740.

  • use the existing native instance_id as the fresh, non-reused process/registration-attempt identity on both the fd side channel and gRPC; no second registration ID or RPC schema change
  • keep native VMM v1 intentionally small: one GPU, one fused allocation, and multiple layer views; reject unsupported topology and mixed Python/native contexts while preserving legacy Python multi-context behavior
  • let one Arc<VmmRegistration> own the imported mapping and an operation gate; save/load tasks retain a guard until GPU work finishes, even if the RPC is cancelled
  • atomically remove a native registration before close, reject new native work, drain accepted work, remove engine raw addresses, and release the mapping on the registry actor
  • run registry-to-engine registration and cleanup as cancellation-independent tasks, with one barrier preventing registration from crossing instance/all-instance teardown
  • keep the existing v1 instance_id\0device_id SCM_RIGHTS message, but bound connections and pending fds, expire unmatched entries on activity, reject malformed/multi-fd messages without descriptor leaks, and recover only a verified stale socket inode
  • disable the fd side channel by default; native deployments must configure a path explicitly
  • auto-detect native-only CUDA devices through the CUDA driver without importing PyTorch

Scope

Native VMM v1 supports the current Qwen layout only: one GPU, one fused allocation, and multiple layer views. Multiple GPUs or allocations per native instance, including the GLM layout, remain out of scope and fail early.

A native client must generate a fresh instance_id for every process/registration attempt and must not reuse it. Native loads require wait_for_completion=true; an asynchronous native load cannot transfer the mapping lifetime into the existing detached worker contract safely.

The UDS path remains explicit configuration rather than gRPC discovery. The wire protocol and pegaflow_proto::VERSION stay on native-vmm-v1.

Validation

  • cargo check --workspace --all-targets --no-default-features --features cuda-13,rdma
  • cargo clippy --workspace --all-targets --no-default-features --features cuda-13,rdma -- -D warnings
  • cargo test -p pegaflow-server --lib --no-default-features --features cuda-13,rdma — 34 passed
  • cargo test --release -p pegaflow-server --test mock_vllm_rpc_e2e --no-default-features --features cuda-13,rdma — 8 passed
  • cargo test -p pegaflow-server --test native_vmm_rpc_e2e --no-default-features --features cuda-13,rdma -- --nocapture — 1 passed on a real GPU
  • cd python && uv run pytest — 202 passed, 11 deselected
  • native-only startup with no explicit --devices — CUDA driver auto-detected the GPU, initialized the torch-free path, served gRPC/HTTP, and shut down cleanly
  • prek run — all hooks passed, including release tests

The Python integration suite was not run because the local test environment does not include PyTorch; the native Rust GPU E2E covers the VMM registration/save/load/unregister path.

@xiaguan
xiaguan force-pushed the codex/vmm-registration-resource branch 3 times, most recently from 69a1927 to ad846ae Compare July 24, 2026 10:57
Signed-off-by: xiaguan <751080330@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant