Skip to content

Packaging: which artifacts should ship as pip wheels #62

Description

@jiejingzhangamd

pyproject.toml already declares amd-infera with three console scripts, but nothing publishes a wheel and there is no release workflow. Everything reaches a machine today by being baked into a ~40-80 GB engine image, which is why an upstream base bump means rebuilding the world.

Building the overlay payload (#61) forced a full audit of what we ship, and the conclusion is stronger than expected: nothing we add is tied to a specific vendor image. What looked like hard image-level coupling is coupling to two coarse, stable interfaces — the CPython minor and the ROCm major.

Audit

Artifact Form Actually binds Ships as
infera (server, router policies, engine wrappers, gaie) pure Python nothing py3-none-any wheel
infera.kvd pure Python (python -m infera.kvd; Dockerfile.kvd is just python:3.10-slim + pip install) nothing same wheel
infera-router (Rust) ~5 MB binary; links only libc/libstdc++/libgcc_s/libmno ROCm, no torch glibc manylinux platform wheel, or package data
gaie protobuf/gRPC stubs generated .py, checked in grpcio version part of the [gaie] extra
Mooncake compiled extension + ~40 system libs CPython minor + ROCm major a ROCm-built wheel, auditwheel-repaired
hipFile libhipfile.so + ais-check (a Python script) + a small module ROCm major same wheel, or its own
aiter kernels JIT-compiled per GPU arch at run time GPU arch not ours to ship — cache ~/.aiter/build
infera-operator Go binary + Helm chart container image + chart (already has one)

The correction

An earlier version of this issue said Mooncake and hipFile "must be built against their base" and could not be wheels. That was wrong, and #61 proves it on hardware: Mooncake built in our image, with its transitive libs bundled next to the .so, loads and instantiates TransferEngine inside an unmodified vllm/vllm-openai-rocm:kimi-k3; libhipfile.so loads and ais-check runs there too. Bundling those libs is exactly what auditwheel repair automates.

So the wheel boundary is not "pure Python only" — it is everything, split by (CPython minor, ROCm major).

Why the upstream Mooncake wheel does not help

mooncake-transfer-engine on PyPI ships cp310–cp313 manylinux wheels, but they are CUDA builds: engine.so lists libcudart.so.12 and libcuda.so.1 as NEEDED and contains zero HIP symbols. A ROCm variant has to be built and published by us.

Good news on the maintenance side: all three of our Mooncake patches are now obsolete. B.1 (USE_HIP_DMABUF propagated to rdma_transport) and B.3 (auto-chunk MRs over max_mr_size, upstream PR #2644) are merged; B.2 (gating the HIP transport) was solved upstream more precisely with isHipReachableTarget(), which skips HIP buffers for cross-host targets instead of disabling intra-node P2P wholesale. A ROCm wheel would therefore be stock upstream source with ROCm build flags — no fork to carry.

Questions to settle

  • How many wheels? amd-infera (pure Python, py3-none-any) plus a native wheel per (rocm, python)? Or fold the native bits into extras, e.g. amd-infera[rocm7]?
  • How does infera-router ship? (a) manylinux_2_28_x86_64 wheel built with maturin/cibuildwheel; (b) package data, which costs the py3-none-any tag; (c) stays in the payload image only.
  • Who builds the ROCm Mooncake wheel, and against which ROCm? Needs a ROCm toolchain in release CI, auditwheel repair for the ~40 transitive libs, and a naming scheme that encodes the ROCm major (PyPI has no ROCm platform tag — this may have to live on an internal index).
  • Publish where? PyPI for the pure-Python wheel, internal index for the ROCm-tagged ones? setuptools_scm already derives the version from git tags, so git tag vX.Y.Z is the natural trigger.
  • Minimum Python. requires-python = ">=3.10" matches the vendor split (SGLang 3.10, vLLM 3.12). Worth asserting in CI so a 3.11+ feature does not slip in unnoticed.

Why it is worth doing

With #61 the payload is assembled by pip install --target from a repo checkout plus harvesting binaries out of an existing engine image. With wheels it becomes pip install --target /payload amd-infera==X.Y.Z amd-infera-rocm7==X.Y.Z and needs no checkout and no source image — and a user can pip install straight into whatever upstream image they like. The thing we rebuild on every source change drops from tens of gigabytes to a few megabytes.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions