create rust workspace (#32014)

2026-07-23 12:02:41 -07:00
parent d0b9689805
commit 7fe82dd02e
25 changed files with 550 additions and 260 deletions
@@ -225,16 +225,10 @@ git_describe_command = ["python3", "python/tools/get_version_tag.py"]
 # Allow editable installs even when .git metadata is not available.
 fallback_version = "0.0.0.dev0"

-[[tool.setuptools-rust.ext-modules]]
-target = "sglang.srt.grpc._core"
-path = "../rust/sglang-grpc/Cargo.toml"
-binding = "PyO3"
-
-[[tool.setuptools-rust.ext-modules]]
-target = "sglang.srt.multimodal._core"
-path = "../rust/sglang-mm/Cargo.toml"
-binding = "PyO3"
-debug = false
+# Rust extension modules are auto-discovered by setup.py from the cargo
+# workspace in ../rust ([package.metadata.sglang] python-module in each crate).
+# This CUDA pyproject builds all of them; platform variants restrict the set
+# via [tool.sglang] rust-extensions (see pyproject_other.toml).

 [tool.kernels.dependencies]
 "kernels-community/sgl-flash-attn3" = 1
@@ -192,12 +192,14 @@ dev_mps = ["sglang[all_mps]", "sglang[test]"]
 [project.scripts]
 sglang = "sglang.cli.main:main"

-# Rust-accelerated multimodal preprocessing (sglang.srt.multimodal._core).
-# grpc is intentionally omitted here (it needs proto/tonic); ROCm only builds mm.
-[[tool.setuptools-rust.ext-modules]]
-target = "sglang.srt.multimodal._core"
-path = "../rust/sglang-mm/Cargo.toml"
-binding = "PyO3"
+# Rust extension modules are auto-discovered by setup.py from the cargo
+# workspace in ../rust; this allowlist restricts which of them this platform
+# builds (case-insensitive substrings of the target module). grpc is
+# intentionally omitted here (it needs proto/tonic); ROCm only builds mm.
+# NOTE: setup.py reads this with a line regex, not a TOML parser — keep the
+# assignment on a single line.
+[tool.sglang]
+rust-extensions = ["multimodal"]

 [tool.setuptools.package-data]
 "sglang" = [
@@ -1,39 +1,152 @@
 """sglang build hooks.

-SGLANG_BUILD_RUST_EXTS controls which Rust extensions are built:
-  - unset or "all": build every declared Rust extension (the default).
-  - "none": build no Rust extensions.
-  - comma-separated names: build only extensions whose target matches one of the
-    given (case-insensitive) substrings, e.g. "grpc" matches
-    "sglang.srt.grpc._core".
+Rust extensions are auto-discovered from the cargo workspace in ../rust: every
+crate whose Cargo.toml declares

-This is a build-time environment variable, so it is read directly from
-os.environ instead of sglang.srt.environ, which is not available until after the
-package has been built.
+    [package.metadata.sglang]
+    python-module = "sglang.srt.<pkg>._core"   # import path inside the wheel
+    debug = false                              # optional RustExtension knob
+
+is built as a PyO3 extension module at that import path. Adding a new extension
+crate therefore needs no pyproject changes — declare the metadata in the crate.
+
+Two filters can narrow the discovered set:
+
+- [tool.sglang] rust-extensions in the active pyproject.toml: a list of
+  case-insensitive substrings of the target module. Platform pyprojects use
+  this to build a subset (e.g. pyproject_other.toml builds only "multimodal";
+  grpc needs proto/tonic and is intentionally CUDA-only).
+- SGLANG_BUILD_RUST_EXTS env var, applied at build time on top of the above:
+  unset or "all" builds everything, "none" builds nothing, and a
+  comma-separated list matches substrings, e.g. "grpc" matches
+  "sglang.srt.grpc._core". It is read directly from os.environ instead of
+  sglang.srt.environ, which is not importable until the package is built.
 """

+import json
 import os
+import re
+import subprocess
+from pathlib import Path

 from setuptools import setup

 try:
-    from setuptools_rust import build_rust
+    from setuptools_rust import Binding, RustExtension, build_rust
 except ModuleNotFoundError as exc:
    if exc.name != "setuptools_rust":
        raise
-    # Alternate platform pyprojects do not declare Rust extensions.
+    # Alternate platform pyprojects that build no Rust extensions do not
+    # install setuptools-rust.
    build_rust = None

 _BUILD_RUST_EXTS_ENV = "SGLANG_BUILD_RUST_EXTS"
+_PYTHON_DIR = Path(__file__).resolve().parent
+_RUST_WORKSPACE_DIR = _PYTHON_DIR.parent / "rust"
+
+
+def _cargo_workspace_metadata():
+    """The rust/ cargo workspace as JSON, straight from cargo's own parser."""
+    manifest_path = _RUST_WORKSPACE_DIR / "Cargo.toml"
+    if not manifest_path.is_file():
+        raise RuntimeError(
+            f"no cargo workspace at {manifest_path} (building outside a repo "
+            f"checkout?); set {_BUILD_RUST_EXTS_ENV}=none to build without "
+            "Rust extensions"
+        )
+    try:
+        out = subprocess.run(
+            [
+                "cargo",
+                "metadata",
+                "--format-version",
+                "1",
+                "--no-deps",
+                "--manifest-path",
+                str(manifest_path),
+            ],
+            capture_output=True,
+            check=True,
+            text=True,
+        )
+    except FileNotFoundError as exc:
+        raise RuntimeError(
+            "cargo is required to discover the Rust extension modules in "
+            f"{_RUST_WORKSPACE_DIR} (and to build them); install a Rust "
+            f"toolchain, or set {_BUILD_RUST_EXTS_ENV}=none to build without "
+            "Rust extensions"
+        ) from exc
+    except subprocess.CalledProcessError as exc:
+        raise RuntimeError(f"cargo metadata failed:\n{exc.stderr}") from exc
+    return json.loads(out.stdout)
+
+
+def _match_by_substring(declared, tokens, source):
+    """Match tokens as case-insensitive substrings of extension names."""
+    matched = set()
+    unmatched = []
+    for token in tokens:
+        hits = {ext.name for ext in declared if token in ext.name.lower()}
+        if hits:
+            matched |= hits
+        else:
+            unmatched.append(token)
+    if unmatched:
+        declared_names = sorted(ext.name for ext in declared)
+        raise ValueError(
+            f"{source} matched no discovered Rust extension for: {unmatched}; "
+            f"discovered extensions are {declared_names}"
+        )
+    return [ext for ext in declared if ext.name in matched]
+
+
+def _discovered_rust_extensions():
+    """One RustExtension per workspace crate declaring a python-module."""
+    extensions = []
+    for package in sorted(
+        _cargo_workspace_metadata()["packages"], key=lambda p: p["name"]
+    ):
+        sglang_meta = (package["metadata"] or {}).get("sglang", {})
+        if "python-module" not in sglang_meta:
+            continue
+        extensions.append(
+            RustExtension(
+                target=sglang_meta["python-module"],
+                path=package["manifest_path"],
+                binding=Binding.PyO3,
+                debug=sglang_meta.get("debug"),
+            )
+        )
+    if not extensions:
+        raise RuntimeError(
+            f"no crate under {_RUST_WORKSPACE_DIR} declares "
+            "[package.metadata.sglang] python-module; set "
+            f"{_BUILD_RUST_EXTS_ENV}=none to build without Rust extensions"
+        )
+    return extensions
+
+
+# Deliberately not a TOML parser (keeps setup.py stdlib-only): the allowlist
+# must be written as a single line, e.g. rust-extensions = ["multimodal"].
+_ALLOWLIST_RE = re.compile(r"^rust-extensions\s*=\s*\[([^\]]*)\]", re.MULTILINE)
+
+
+def _pyproject_rust_extensions(declared):
+    """Apply the active pyproject's [tool.sglang] rust-extensions allowlist."""
+    pyproject_text = (_PYTHON_DIR / "pyproject.toml").read_text(encoding="utf-8")
+    match = _ALLOWLIST_RE.search(pyproject_text)
+    if match is None:
+        return declared
+    tokens = re.findall(r'"([^"]*)"', match.group(1))
+    return _match_by_substring(
+        declared=declared,
+        tokens=[token.lower() for token in tokens],
+        source="[tool.sglang] rust-extensions",
+    )


 def _selected_rust_extensions(declared):
-    """Return the Rust extensions selected by SGLANG_BUILD_RUST_EXTS.
-
-    `ext.name` is the fully-qualified target (e.g. "sglang.srt.grpc._core") for
-    the string-target declarations in pyproject.toml, so comma-separated names
-    are matched as case-insensitive substrings of it.
-    """
+    """Apply the SGLANG_BUILD_RUST_EXTS build-time filter."""
    declared = list(declared)
    raw = os.environ.get(_BUILD_RUST_EXTS_ENV)
    if raw is None:
@@ -52,23 +165,17 @@ def _selected_rust_extensions(declared):
            f"{_BUILD_RUST_EXTS_ENV}={raw!r} has an empty item; unset it or use "
            "'all', 'none', or a comma-separated list of extension names"
        )
+    return _match_by_substring(
+        declared=declared, tokens=tokens, source=_BUILD_RUST_EXTS_ENV
+    )

-    matched = set()
-    unmatched = []
-    for token in tokens:
-        hits = {ext.name for ext in declared if token in ext.name.lower()}
-        if hits:
-            matched |= hits
-        else:
-            unmatched.append(token)
-    if unmatched:
-        declared_names = sorted(ext.name for ext in declared)
-        raise ValueError(
-            f"{_BUILD_RUST_EXTS_ENV} matched no declared Rust extension for: "
-            f"{unmatched}; declared extensions are {declared_names}"
-        )

-    return [ext for ext in declared if ext.name in matched]
+def _declared_rust_extensions():
+    # "none" short-circuits discovery so builds without a ../rust checkout
+    # (e.g. from an sdist) still work.
+    if (os.environ.get(_BUILD_RUST_EXTS_ENV) or "").strip().lower() == "none":
+        return []
+    return _pyproject_rust_extensions(_discovered_rust_extensions())


 if build_rust is not None:
@@ -84,9 +191,9 @@ if build_rust is not None:
                return
            super().run()

-    _cmdclass = {"build_rust": BuildRust}
+    setup(
+        cmdclass={"build_rust": BuildRust},
+        rust_extensions=_declared_rust_extensions(),
+    )
 else:
-    _cmdclass = {}
-
-
-setup(cmdclass=_cmdclass)
+    setup()