zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/core/subenv.py

@@ -0,0 +1,202 @@
+"""Shared sub-environment primitive: create and drive isolated venvs.
+
+ZeroQuantz runs mutually-incompatible toolchains in dedicated virtual environments
+so they never pollute your main Python — or each other. Two consumers share this:
+
+* **Quantization** toolchains (AWQ, GPTQ, AutoRound, HQQ) — see
+  :mod:`zeroquantz.quantization.isolated`.
+* **Serving** runtimes (vLLM, SGLang, TensorRT-LLM) — see
+  :mod:`zeroquantz.deploy.runtime_env`.
+
+Each env lives under a two-level, namespaced layout::
+
+    ~/.zeroquantz/envs/
+      quant/   awq/ gptq/ autoround/ hqq/
+      serve/   vllm/ sglang/ tensorrt_llm/
+
+``uv venv`` is used when ``uv`` is on PATH (fast), otherwise ``python -m venv``.
+Builds are cached behind a ``.zeroquantz_ready`` marker and reused. Nothing is ever
+installed into your main environment.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.core.exceptions import ZeroQuantzError
+from zeroquantz.utils.logging import get_logger
+from zeroquantz.utils.paths import paths
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+log = get_logger(__name__)
+
+# PyTorch CUDA wheel index by detected CUDA version (best-effort nearest match).
+_TORCH_CUDA_INDEX = {
+    "12.8": "https://download.pytorch.org/whl/cu128",
+    "12.6": "https://download.pytorch.org/whl/cu126",
+    "12.4": "https://download.pytorch.org/whl/cu124",
+    "12.1": "https://download.pytorch.org/whl/cu121",
+    "11.8": "https://download.pytorch.org/whl/cu118",
+}
+
+
+def torch_index_for(cuda_version: str | None) -> str | None:
+    """Best-effort PyTorch CUDA wheel index URL for a detected CUDA version."""
+    if not cuda_version:
+        return None
+    for prefix, url in _TORCH_CUDA_INDEX.items():
+        if cuda_version.startswith(prefix):
+            return url
+    return "https://download.pytorch.org/whl/cu124"  # newest known
+
+
+@dataclass(frozen=True)
+class EnvSpec:
+    """A declarative description of one isolated environment."""
+
+    kind: str                       # "quant" | "serve"
+    name: str                       # family or runtime name (the leaf dir)
+    pip: tuple[str, ...] = ()        # packages to install (the toolchain/runtime)
+    torch_index: str | None = None   # install torch first from this CUDA wheel index
+    extra_index_url: str | None = None  # e.g. https://pypi.nvidia.com for TensorRT-LLM
+    python_version: str = "3.11"     # used only by `uv venv --python`
+
+
+@dataclass
+class EnvPlan:
+    """A previewable plan for building an environment (no side effects)."""
+
+    spec: EnvSpec
+    env_dir: Path
+    python: Path
+    commands: list[str] = field(default_factory=list)
+
+
+class SubEnvManager:
+    """Create and drive per-toolchain / per-runtime isolated environments."""
+
+    def __init__(self, root: Path | None = None) -> None:
+        self.root = root or (paths().ensure().home / "envs")
+
+    # ---- path resolution ----------------------------------------------------
+
+    def env_dir(self, spec: EnvSpec) -> Path:
+        return self.root / spec.kind / spec.name
+
+    @staticmethod
+    def venv_python(env_dir: Path) -> Path:
+        return env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python")
+
+    @staticmethod
+    def bin_dir(env_dir: Path) -> Path:
+        """The env's console-scripts directory (Scripts on Windows, bin elsewhere)."""
+        return env_dir / ("Scripts" if os.name == "nt" else "bin")
+
+    @staticmethod
+    def uv_available() -> bool:
+        return shutil.which("uv") is not None
+
+    def is_ready(self, spec: EnvSpec) -> bool:
+        env_dir = self.env_dir(spec)
+        return (env_dir / ".zeroquantz_ready").exists() and self.venv_python(env_dir).exists()
+
+    # ---- planning / building ------------------------------------------------
+
+    def _pip_prefix(self, py: Path, use_uv: bool) -> list[str]:
+        if use_uv:
+            return [shutil.which("uv") or "uv", "pip", "install", "--python", str(py)]
+        return [str(py), "-m", "pip", "install", "--upgrade"]
+
+    def plan(self, spec: EnvSpec) -> EnvPlan:
+        """The exact create + install commands, without running anything."""
+        env_dir = self.env_dir(spec)
+        py = self.venv_python(env_dir)
+        use_uv = self.uv_available()
+        creator = (
+            ["uv", "venv", str(env_dir), "--python", spec.python_version, "--seed"]
+            if use_uv
+            else [sys.executable, "-m", "venv", str(env_dir)]
+        )
+        pip = self._pip_prefix(py, use_uv)
+        commands = [" ".join(creator)]
+        if spec.torch_index:
+            commands.append(" ".join([*pip, "torch", "--index-url", spec.torch_index]))
+        if spec.pip:
+            extra = ["--extra-index-url", spec.extra_index_url] if spec.extra_index_url else []
+            commands.append(" ".join([*pip, *spec.pip, *extra]))
+        return EnvPlan(spec=spec, env_dir=env_dir, python=py, commands=commands)
+
+    def ensure(
+        self, spec: EnvSpec, *, progress: Callable[[str, float], None] | None = None
+    ) -> Path:
+        """Build the env if it isn't ready; return the venv's Python interpreter."""
+        env_dir = self.env_dir(spec)
+        py = self.venv_python(env_dir)
+        if self.is_ready(spec):
+            return py
+
+        use_uv = self.uv_available()
+        env_dir.parent.mkdir(parents=True, exist_ok=True)
+        if progress:
+            progress(f"creating venv for {spec.kind}/{spec.name}", 0.05)
+        if use_uv:
+            self.run([shutil.which("uv"), "venv", str(env_dir), "--python", spec.python_version, "--seed"])
+        else:
+            self.run([sys.executable, "-m", "venv", str(env_dir)])
+        if not py.exists():
+            raise ZeroQuantzError(f"venv creation failed for {spec.kind}/{spec.name} at {env_dir}.")
+
+        pip = self._pip_prefix(py, use_uv)
+        if spec.torch_index:
+            if progress:
+                progress("installing torch (CUDA)", 0.2)
+            self.run([*pip, "torch", "--index-url", spec.torch_index])
+        if spec.pip:
+            if progress:
+                progress(f"installing {spec.name} ({spec.kind})", 0.5)
+            extra = ["--extra-index-url", spec.extra_index_url] if spec.extra_index_url else []
+            self.run([*pip, *spec.pip, *extra])
+
+        (env_dir / ".zeroquantz_ready").write_text(
+            json.dumps(
+                {"pip": list(spec.pip), "torch_index": spec.torch_index,
+                 "extra_index_url": spec.extra_index_url}
+            )
+        )
+        return py
+
+    # ---- subprocess helper --------------------------------------------------
+
+    def run(
+        self, cmd: list, *, progress: Callable[[str, float], None] | None = None
+    ) -> None:
+        """Run a command in the child process, logging its output line by line."""
+        log.info("subenv: %s", " ".join(str(c) for c in cmd))
+        try:
+            proc = subprocess.Popen(
+                [str(c) for c in cmd],
+                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                text=True, encoding="utf-8", errors="replace",
+            )
+        except OSError as exc:
+            raise ZeroQuantzError(f"Failed to launch: {' '.join(map(str, cmd))}", detail=str(exc)) from exc
+        assert proc.stdout is not None
+        for line in proc.stdout:
+            line = line.rstrip()
+            if line:
+                log.debug("child: %s", line)
+        code = proc.wait()
+        if code != 0:
+            raise ZeroQuantzError(
+                f"Sub-env command failed (exit {code}): {' '.join(map(str, cmd[:3]))} …",
+                detail="See ~/.zeroquantz/logs for the child process output.",
+            )

zeroquantz/deploy/__init__.py

@@ -0,0 +1,25 @@
+"""Deployment: turn a (model, quantization format, runtime) choice into runnable
+serve commands, an OpenAI-compatible client, and container assets.
+
+vLLM / SGLang / TensorRT-LLM are each provisioned and served in their own isolated
+sub-environment (see :mod:`zeroquantz.deploy.runtime_env`)."""
+
+from __future__ import annotations
+
+from zeroquantz.deploy.assets import DeployBundle, generate_assets
+from zeroquantz.deploy.launcher import can_provision, is_runtime_available, launch
+from zeroquantz.deploy.runtime_env import RuntimeEnvRunner, runtime_spec
+from zeroquantz.deploy.targets import DeployTarget, deploy_targets_for, get_target
+
+__all__ = [
+    "DeployBundle",
+    "DeployTarget",
+    "RuntimeEnvRunner",
+    "can_provision",
+    "deploy_targets_for",
+    "generate_assets",
+    "get_target",
+    "is_runtime_available",
+    "launch",
+    "runtime_spec",
+]

zeroquantz/deploy/assets.py

@@ -0,0 +1,161 @@
+"""Generate deployment assets for a (model, format, target): serve script,
+OpenAI-compatible client, and container files (Dockerfile-free — uses the
+runtime's official image via docker-compose / docker run)."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from zeroquantz.deploy.targets import DeployTarget
+    from zeroquantz.quantization.catalog import QuantFormat
+
+
+@dataclass
+class DeployBundle:
+    directory: str
+    target: str
+    serve_command: str
+    endpoint: str | None
+    files: list[str] = field(default_factory=list)
+
+
+def generate_assets(
+    model_path: str,
+    fmt: QuantFormat,
+    target: DeployTarget,
+    out_dir: str | Path,
+    *,
+    port: int | None = None,
+) -> DeployBundle:
+    port = port or target.default_port
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    files: list[str] = []
+
+    def write(name: str, content: str, *, executable: bool = False) -> None:
+        path = out / name
+        path.write_text(content, encoding="utf-8")
+        files.append(str(path))
+
+    is_local = Path(model_path).exists()
+    local_serve = target.serve_command(model_path, fmt, port=port)
+
+    # 1) local serve script
+    write("serve.sh", f"#!/usr/bin/env bash\nset -e\n{local_serve}\n", executable=True)
+
+    endpoint = f"http://localhost:{port}/v1" if target.openai_compatible else None
+
+    # 2) container assets for server runtimes
+    if target.kind == "server" and target.container_image:
+        container_path = "/model" if is_local else model_path
+        container_serve = target.serve_command(container_path, fmt, port=port)
+        write("docker-compose.yml", _compose(target, model_path, container_serve, port, is_local))
+        write("run-docker.sh", _run_docker(target, model_path, container_serve, port, is_local),
+              executable=True)
+
+    # 3) gguf helpers
+    if target.name == "ollama":
+        write("Modelfile", f"FROM {Path(model_path).name if is_local else './model.gguf'}\n")
+
+    # 4) OpenAI-compatible client
+    if target.openai_compatible:
+        model_id = Path(model_path).name if is_local else model_path
+        write("client.py", _client_py(endpoint, model_id))
+        write("client.sh", _client_sh(endpoint, model_id), executable=True)
+
+    # 5) README
+    write("README.md", _readme(target, fmt, model_path, port, endpoint, is_local))
+
+    return DeployBundle(str(out), target.name, local_serve, endpoint, files)
+
+
+def _compose(target, model_path, container_serve, port, is_local):  # noqa: ANN001
+    volumes = f'      - "{Path(model_path).resolve()}:/model"\n' if is_local else ""
+    env = "" if is_local else "    environment:\n      - HF_TOKEN=${HF_TOKEN}\n"
+    ports = f'      - "{port}:{port}"\n' if port else ""
+    return (
+        "services:\n"
+        f"  {target.name}:\n"
+        f"    image: {target.container_image}\n"
+        "    ipc: host\n"
+        + (f"    ports:\n{ports}" if ports else "")
+        + (f"    volumes:\n{volumes}" if volumes else "")
+        + env
+        + f"    command: {container_serve}\n"
+        "    deploy:\n"
+        "      resources:\n"
+        "        reservations:\n"
+        "          devices:\n"
+        "            - driver: nvidia\n"
+        "              count: all\n"
+        "              capabilities: [gpu]\n"
+    )
+
+
+def _run_docker(target, model_path, container_serve, port, is_local):  # noqa: ANN001
+    mount = f'-v "{Path(model_path).resolve()}:/model" ' if is_local else ""
+    env = "" if is_local else '-e HF_TOKEN="$HF_TOKEN" '
+    pub = f"-p {port}:{port} " if port else ""
+    return (
+        "#!/usr/bin/env bash\nset -e\n"
+        f"docker run --gpus all --ipc=host {pub}{mount}{env}{target.container_image} \\\n"
+        f"  {container_serve}\n"
+    )
+
+
+def _client_py(endpoint: str | None, model_id: str) -> str:
+    return (
+        "from openai import OpenAI\n\n"
+        f'client = OpenAI(base_url="{endpoint}", api_key="EMPTY")\n'
+        "resp = client.chat.completions.create(\n"
+        f'    model="{model_id}",\n'
+        '    messages=[{"role": "user", "content": "Explain quantization in one sentence."}],\n'
+        ")\nprint(resp.choices[0].message.content)\n"
+    )
+
+
+def _client_sh(endpoint: str | None, model_id: str) -> str:
+    body = json.dumps({"model": model_id, "messages": [{"role": "user", "content": "Hello!"}]})
+    return (
+        "#!/usr/bin/env bash\n"
+        f"curl {endpoint}/chat/completions \\\n"
+        '  -H "Content-Type: application/json" \\\n'
+        f"  -d '{body}'\n"
+    )
+
+
+def _readme(target, fmt, model_path, port, endpoint, is_local):  # noqa: ANN001
+    lines = [
+        f"# Deploy {model_path} ({fmt.label}) on {target.label}",
+        "",
+        f"{target.note}",
+        "",
+        f"**Install:** {target.install_hint}",
+        "",
+        "## Run locally",
+        "```bash",
+        "bash serve.sh",
+        "```",
+    ]
+    if target.kind == "server" and target.container_image:
+        lines += [
+            "",
+            "## Run with Docker",
+            "```bash",
+            "bash run-docker.sh        # or: docker compose up",
+            "```",
+            "",
+            "> Note: official runtime images may define their own ENTRYPOINT; if so,",
+            "> pass only the arguments after the server executable.",
+        ]
+    if endpoint:
+        lines += ["", f"## Query it (OpenAI-compatible, {endpoint})", "```bash", "python client.py   # or: bash client.sh", "```"]
+    if not is_local:
+        lines += ["", f"> `{model_path}` is a Hub id — set `HF_TOKEN` for gated models; the server downloads it on start."]
+    if fmt.requires_precision == "fp8":
+        lines += ["", "> This format needs FP8/FP4-class hardware (Hopper/Blackwell) to run at full speed."]
+    return "\n".join(lines) + "\n"

zeroquantz/deploy/launcher.py

@@ -0,0 +1,80 @@
+"""Local launch of a deployment target.
+
+Two kinds of runtime:
+
+* **Isolated** (vLLM, SGLang, TensorRT-LLM) — provisioned in a dedicated venv under
+  ``~/.zeroquantz/envs/serve/<name>/`` and launched from it (that env's bin/Scripts
+  is prepended to PATH so ``vllm`` / ``python -m sglang…`` / ``trtllm-serve``
+  resolve to it). ``launch`` builds the env on first use.
+* **Base/system** (Transformers, llama.cpp, Ollama) — launched directly, requiring
+  the tool to be importable or on PATH.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import shutil
+import subprocess
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zeroquantz.deploy.targets import DeployTarget
+
+# How to detect a *base/system* runtime locally: a CLI on PATH or an importable module.
+_CLI = {"llamacpp": "llama-server", "ollama": "ollama"}
+_MODULE = {"transformers": "transformers"}
+
+
+def can_provision(target: DeployTarget) -> bool:
+    """True if ZeroQuantz can build an isolated env for this runtime on demand."""
+    return bool(getattr(target, "isolated_env", False) and target.pip)
+
+
+def is_runtime_available(target: DeployTarget) -> bool:
+    """True if this target can be launched *right now* (no build needed).
+
+    For isolated runtimes that means the venv is already built; for base/system
+    runtimes, that the CLI is on PATH or the module is importable.
+    """
+    if can_provision(target):
+        from zeroquantz.deploy.runtime_env import RuntimeEnvRunner
+
+        return RuntimeEnvRunner().is_ready(target)
+    cli = _CLI.get(target.name)
+    if cli and shutil.which(cli):
+        return True
+    module = _MODULE.get(target.name)
+    return bool(module and importlib.util.find_spec(module) is not None)
+
+
+def launch(
+    target: DeployTarget,
+    serve_command: str,
+    *,
+    cuda_version: str | None = None,
+    progress: Callable[[str, float], None] | None = None,
+):
+    """Start the serve command as a background process. Returns a Popen.
+
+    For an isolated runtime this **builds the venv on first use** (a multi-GB
+    install) and runs the server from it. For base/system runtimes it requires the
+    runtime to be installed. Raises RuntimeError if it can't be launched.
+    """
+    env = os.environ.copy()
+
+    if can_provision(target):
+        from zeroquantz.deploy.runtime_env import RuntimeEnvRunner
+
+        runner = RuntimeEnvRunner()
+        runner.ensure(target, progress=progress)  # build the venv if needed
+        bin_dir = str(runner.bin_dir(target))
+        env["PATH"] = bin_dir + os.pathsep + env.get("PATH", "")
+        env["VIRTUAL_ENV"] = str(runner.env_dir(target))
+        return subprocess.Popen(serve_command, shell=True, env=env)
+
+    if not is_runtime_available(target):
+        raise RuntimeError(f"{target.label} is not installed locally.")
+    return subprocess.Popen(serve_command, shell=True, env=env)

zeroquantz/deploy/runtime_env.py

@@ -0,0 +1,66 @@
+"""Build and serve each inference runtime in its own isolated sub-environment.
+
+vLLM, SGLang, and TensorRT-LLM pin mutually-incompatible torch/CUDA stacks, so
+each is provisioned in a dedicated venv under ``~/.zeroquantz/envs/serve/<runtime>/``
+(via the shared :class:`~zeroquantz.core.subenv.SubEnvManager`) and served from that
+env's interpreter — never from, and never polluting, your main environment.
+
+Runtimes that are plain system tools or run in the base env (llama.cpp, Ollama,
+Transformers) are *not* isolated and are handled directly by the launcher.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.core.subenv import EnvPlan, EnvSpec, SubEnvManager
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zeroquantz.deploy.targets import DeployTarget
+
+
+def runtime_spec(target: DeployTarget) -> EnvSpec:
+    """The EnvSpec for a runtime's isolated serving env.
+
+    We do *not* pin a torch CUDA index here: vLLM/SGLang/TensorRT-LLM wheels pull a
+    matching torch themselves, and forcing one risks a conflict.
+    """
+    return EnvSpec(
+        kind="serve",
+        name=target.name,
+        pip=tuple(target.pip),
+        torch_index=None,
+        extra_index_url=target.extra_index_url,
+        python_version=target.env_python,
+    )
+
+
+class RuntimeEnvRunner:
+    """Provision and locate per-runtime isolated serving environments."""
+
+    def __init__(self, root: Path | None = None) -> None:
+        self._envs = SubEnvManager(root=root)
+
+    def supports(self, target: DeployTarget) -> bool:
+        return bool(getattr(target, "isolated_env", False) and target.pip)
+
+    def is_ready(self, target: DeployTarget) -> bool:
+        return self.supports(target) and self._envs.is_ready(runtime_spec(target))
+
+    def env_dir(self, target: DeployTarget) -> Path:
+        return self._envs.env_dir(runtime_spec(target))
+
+    def bin_dir(self, target: DeployTarget) -> Path:
+        return SubEnvManager.bin_dir(self.env_dir(target))
+
+    def plan(self, target: DeployTarget) -> EnvPlan:
+        return self._envs.plan(runtime_spec(target))
+
+    def ensure(
+        self, target: DeployTarget, *, progress: Callable[[str, float], None] | None = None
+    ) -> Path:
+        """Build the runtime's venv if needed; return its Python interpreter."""
+        return self._envs.ensure(runtime_spec(target), progress=progress)