mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-26 03:39:41 +08:00
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
"""Record the actual wheel-owned imports from the fresh test interpreter."""
|
|
|
|
import importlib
|
|
import importlib.metadata as metadata
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from contracts import digest, require, write_json
|
|
from resource_queue import ResourceUnavailable, reservation
|
|
|
|
|
|
def main():
|
|
require(sys.prefix != sys.base_prefix, "Test interpreter is not a venv")
|
|
for name in ("transformers", "accelerate", "sglang", "sgl-kernel", "sgl-kernel-kt"):
|
|
try:
|
|
metadata.distribution(name)
|
|
except metadata.PackageNotFoundError:
|
|
continue
|
|
raise ValueError(f"Conflicting upstream distribution: {name}")
|
|
result = {}
|
|
for package, module in {
|
|
"ktransformers": "ktransformers",
|
|
"kt-kernel": "kt_kernel",
|
|
"sglang-kt": "sglang",
|
|
"transformers-kt": "transformers",
|
|
"accelerate-kt": "accelerate",
|
|
}.items():
|
|
dist = metadata.distribution(package)
|
|
imported = importlib.import_module(module)
|
|
path = Path(imported.__file__).resolve()
|
|
require(
|
|
path.is_relative_to(Path(sys.prefix).resolve()),
|
|
"Imported code is outside the clean venv",
|
|
)
|
|
files = {Path(dist.locate_file(entry)).resolve() for entry in dist.files or []}
|
|
require(path in files, f"Import is not owned by {package}")
|
|
# Hash installed Python/native payload files as well as the entrypoint;
|
|
# tooling installation must not silently overwrite shared namespaces.
|
|
payloads = {
|
|
str(path.relative_to(sys.prefix)): digest(path)
|
|
for path in sorted(files)
|
|
if path.is_file()
|
|
and path.suffix in (".py", ".so")
|
|
and path.is_relative_to(sys.prefix)
|
|
}
|
|
result[package] = {
|
|
"version": dist.version,
|
|
"import_path": str(path),
|
|
"payload_sha256": payloads,
|
|
}
|
|
from accelerate.utils.dataclasses import KTransformersPlugin # noqa: F401
|
|
|
|
write_json(sys.argv[1], result)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Native imports can initialize CUDA too. Respect newly arrived manual
|
|
# workloads before inspecting imports, not just before loading the models.
|
|
try:
|
|
with reservation(
|
|
Path(sys.prefix).parent / "import-start.lock",
|
|
Path(sys.argv[1]).parent / "import-start-queue.jsonl",
|
|
):
|
|
main()
|
|
except ResourceUnavailable as exc:
|
|
write_json(
|
|
Path(sys.argv[1]).parent / "resource-unavailable.json", {"error": str(exc)}
|
|
)
|
|
raise
|