From 85751dc302852739ee339218b2c46ba8e3b44c7e Mon Sep 17 00:00:00 2001 From: yyj Date: Sat, 12 Sep 2026 21:14:36 +0800 Subject: [PATCH] fix(release): preserve native payload layout in candidate wheels --- .github/release/carriers.py | 44 ++++++---------- .github/release/four_main.py | 4 +- .github/release/test_carriers.py | 68 ++++++++++++++++++++++++- .github/workflows/release-four-main.yml | 4 ++ 4 files changed, 89 insertions(+), 31 deletions(-) diff --git a/.github/release/carriers.py b/.github/release/carriers.py index 0df30a894..2d21a0681 100644 --- a/.github/release/carriers.py +++ b/.github/release/carriers.py @@ -30,6 +30,7 @@ "transformers-kt": "transformers_kt_sgl_kernel_payload", "sglang-kt": "sglang_kt_sgl_kernel_payload", "ktransformers": "ktransformers_sgl_kernel_payload", + "accelerate-kt": "accelerate_kt_sgl_kernel_payload", } LARGE = ("flash_ops.abi3.so", "sm100/common_ops.abi3.so") @@ -44,7 +45,12 @@ def unpack(wheel, root): with zipfile.ZipFile(wheel) as archive: names = archive.namelist() require(len(names) == len(set(names)), "Duplicate ZIP entries") - records = [name for name in names if name.endswith(".dist-info/RECORD")] + records = [ + name + for name in names + if name.endswith(".dist-info/RECORD") + and len(PurePosixPath(name).parts) == 2 + ] require(len(records) == 1, "Require one wheel RECORD") rows = list(csv.reader(io.StringIO(archive.read(records[0]).decode()))) record = {row[0]: row[1:] for row in rows} @@ -139,6 +145,9 @@ def retag(root, python, abi): def archive_payload(sgl, output): hashes = {} + paths = [sgl / "sgl_kernel" / name for name in LARGE] + for directory in sorted(sgl.glob("*.libs")): + paths.extend(path for path in sorted(directory.rglob("*")) if path.is_file()) with ( output.open("wb") as raw, gzip.GzipFile( @@ -146,8 +155,8 @@ def archive_payload(sgl, output): ) as compressed, ): with tarfile.open(fileobj=compressed, mode="w|") as archive: - for name in LARGE: - source = sgl / "sgl_kernel" / name + for source in paths: + name = source.relative_to(sgl).as_posix() hashes[name] = sha256(source) info = archive.gettarinfo(str(source), arcname=name) info.uid = info.gid = info.mtime = 0 @@ -189,7 +198,7 @@ def binary_evidence(roots): ) kt = [entry for name, entry in evidence.items() if name.startswith("kt-kernel/")] require( - any(required <= normalize(entry["sass"]) for entry in kt), + not kt or any(required <= normalize(entry["sass"]) for entry in kt), "KT CUDA extension is missing required SASS architectures", ) return evidence @@ -201,7 +210,7 @@ def assemble(raw, output, evidence_dir): inspect_wheel(path) | {"path": path} for path in sorted(raw.glob("*.whl")) ] by_name = {entry["name"]: entry for entry in entries} - expected = set(MODULES) | {"accelerate-kt", "sgl-kernel-kt"} + expected = set(MODULES) | {"sgl-kernel-kt"} require( set(by_name) == expected and len(entries) == len(expected), "Need six fresh inputs", @@ -224,24 +233,6 @@ def assemble(raw, output, evidence_dir): unpack(entry["path"], roots[name]) save_json(evidence_dir / "cuda-binaries.json", binary_evidence(roots)) sgl = roots["sgl-kernel-kt"] - # Lazy objects are extracted into a cache, not site-packages. An - # auditwheel-renamed dependency resolved via $ORIGIN would break there. - # Torch/CUDA-runtime libraries are supplied by the pinned torch wheel; - # reject any other bundled dependency rather than emitting a broken wheel. - bundled = { - path.name - for path in sgl.rglob("*.so*") - if any(part.endswith(".libs") for part in path.parts) - } - for name in LARGE: - dynamic = subprocess.check_output( - ["readelf", "--dynamic", str(sgl / "sgl_kernel" / name)], text=True - ) - needed = set(re.findall(r"Shared library: \[([^]]+)\]", dynamic)) - require( - not (needed & bundled), - "Lazy payload depends on a relocated auditwheel library; fix the main build linkage", - ) for name in ("payload_runtime.py", "load_utils.py", "flash_attn.py"): require( (sgl / "sgl_kernel" / name).is_file(), @@ -261,6 +252,7 @@ def assemble(raw, output, evidence_dir): f"VERSION = {by_name['sgl-kernel-kt']['version']!r}\n" f"ARCHIVE_SHA256 = {sha256(archive)!r}\n" f"PAYLOAD_MODULES = {tuple(MODULES.values())!r}\nFILES = {hashes!r}\n" + f"BINARIES = {dict((name, 'sgl_kernel/' + name) for name in LARGE)!r}\n" ) kt = roots["kt-kernel"] kt_dist = next(kt.glob("*.dist-info")) @@ -355,12 +347,6 @@ def assemble(raw, output, evidence_dir): "Carrier changed dependency metadata", ) require(remaining == 0 and not source.read(1), "Incomplete payload split") - accelerate = by_name["accelerate-kt"]["path"] - require( - accelerate.stat().st_size < LIMIT, - "Accelerate wheel exceeds PyPI size limit", - ) - shutil.copyfile(accelerate, output / accelerate.name) save_json( evidence_dir / "assembly.json", { diff --git a/.github/release/four_main.py b/.github/release/four_main.py index c14cd22aa..ffcf50d69 100644 --- a/.github/release/four_main.py +++ b/.github/release/four_main.py @@ -205,7 +205,9 @@ def inspect_wheel(path: Path) -> dict: name, version, _, tags = parse_wheel_filename(path.name) with zipfile.ZipFile(path) as wheel: metadata_files = [ - p for p in wheel.namelist() if p.endswith(".dist-info/METADATA") + p + for p in wheel.namelist() + if p.endswith(".dist-info/METADATA") and len(PurePosixPath(p).parts) == 2 ] if len(metadata_files) != 1: raise ValueError(f"Expected one METADATA in {path.name}") diff --git a/.github/release/test_carriers.py b/.github/release/test_carriers.py index d69ad9f81..5d7565187 100644 --- a/.github/release/test_carriers.py +++ b/.github/release/test_carriers.py @@ -1,8 +1,10 @@ """Exercise carrier assembly with tiny synthetic wheels, not CUDA execution.""" import sys +import tarfile import zipfile from pathlib import Path +from types import SimpleNamespace import pytest @@ -14,7 +16,7 @@ def raw_wheels(tmp_path): raw = tmp_path / "raw" raw.mkdir() - for package in (*carriers.MODULES, "accelerate-kt", "sgl-kernel-kt"): + for package in (*carriers.MODULES, "sgl-kernel-kt"): root = tmp_path / package root.mkdir() dist = root / (package.replace("-", "_") + "-1.0.dist-info") @@ -73,6 +75,8 @@ def test_fresh_carriers_preserve_runtime_versions_and_sm90(tmp_path, monkeypatch assert "sgl_kernel/_payload_manifest.py" in wheel.namelist() with zipfile.ZipFile(next(output.glob("ktransformers-*.whl"))) as wheel: assert "sgl_kernel/sm90/common_ops.abi3.so" in wheel.namelist() + with zipfile.ZipFile(next(output.glob("accelerate_kt-*.whl"))) as wheel: + assert "accelerate_kt_sgl_kernel_payload/payload.part" in wheel.namelist() # Every final RECORD is independently checked, including regenerated payloads. for index, path in enumerate(output.iterdir()): carriers.unpack(path, tmp_path / f"verify-{index}") @@ -96,3 +100,65 @@ def test_size_limit_fails_without_dropping_architectures(tmp_path, monkeypatch): evidence.mkdir() with pytest.raises(ValueError, match="capacity"): carriers.assemble(raw, tmp_path / "final", evidence) + + +def test_payload_preserves_wheel_relative_libraries(tmp_path): + raw_wheels(tmp_path) + root = tmp_path / "sgl-kernel-kt" + library = root / "sgl_kernel_kt.libs/libnuma.so.1" + library.parent.mkdir() + library.write_bytes(b"dependency") + target = tmp_path / "payload.tar.gz" + hashes = carriers.archive_payload(root, target) + assert set(hashes) == {"sgl_kernel/" + name for name in carriers.LARGE} | { + "sgl_kernel_kt.libs/libnuma.so.1" + } + with tarfile.open(target) as archive: + assert set(archive.getnames()) == set(hashes) + + +def test_embedded_metadata_does_not_shadow_wheel_metadata(tmp_path): + raw_wheels(tmp_path) + root = tmp_path / "accelerate-kt" + nested = root / "accelerate_kt/vendor/example.dist-info" + nested.mkdir(parents=True) + for name in ("METADATA", "RECORD"): + (nested / name).write_text("vendored metadata") + wheel = tmp_path / "accelerate_kt-1.0-py3-none-any.whl" + carriers.pack(root, wheel) + assert inspect_wheel(wheel)["name"] == "accelerate-kt" + carriers.unpack(wheel, tmp_path / "unpacked") + + +@pytest.mark.parametrize("kt_cuda", [False, True]) +def test_cpu_only_kt_is_accepted_without_weakening_cuda_audit( + tmp_path, monkeypatch, kt_cuda +): + kt, sgl = tmp_path / "kt", tmp_path / "sgl" + kt.mkdir() + (sgl / "sgl_kernel/sm100").mkdir(parents=True) + (kt / "extension.so").touch() + (sgl / "sgl_kernel/sm100/common_ops.abi3.so").touch() + monkeypatch.setattr( + carriers.subprocess, + "check_output", + lambda args, **kw: ( + ".nv_fatbin" if kt_cuda or sgl in Path(args[-1]).parents else "" + ), + ) + monkeypatch.setattr( + carriers.subprocess, + "run", + lambda args, **kw: SimpleNamespace( + stdout=( + "sm_80 sm_86 sm_89 sm_90 sm_120" + if sgl in Path(args[-1]).parents + else "sm_80" + ) + ), + ) + if kt_cuda: + with pytest.raises(ValueError, match="KT CUDA extension"): + carriers.binary_evidence({"kt-kernel": kt, "sgl-kernel-kt": sgl}) + else: + carriers.binary_evidence({"kt-kernel": kt, "sgl-kernel-kt": sgl}) diff --git a/.github/workflows/release-four-main.yml b/.github/workflows/release-four-main.yml index bb200e291..64df428e9 100644 --- a/.github/workflows/release-four-main.yml +++ b/.github/workflows/release-four-main.yml @@ -229,6 +229,10 @@ jobs: '--exclude', 'libtorch_cuda.so', '--exclude', 'libtorch_python.so', '--exclude', 'libc10.so', '--exclude', 'libc10_cuda.so', '--exclude', 'libcudart.so.12', + '--exclude', 'libcublas.so.12', + '--exclude', 'libcublasLt.so.12', + '--exclude', 'libcurand.so.10', + '--exclude', 'libnvrtc.so.12', '-w', str(work / 'repaired'), ], check=True) else: