mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[2/3 step][ vllm ci build setup] Add vlllm buld logic and dockerfile (#160089)
# set up vllm build logic - dockerfile: please notice the dockfile introduced here is only temporary, once we migrate this file to vllm, we will fetch it directly from there - VllmBuildRunner: - implement logic to prepare and run vllm build with dockerfile - Pull Request resolved: https://github.com/pytorch/pytorch/pull/160089 Approved by: https://github.com/huydhn ghstack dependencies: #160043
This commit is contained in:
Yang Wang
committed by
PyTorch MergeBot
PyTorch MergeBot
parent
00d7d6f123
commit
639778b3ee
@ -1,11 +1,7 @@
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
from cli.lib.common.cli_helper import (
|
||||
register_target_commands_and_runner,
|
||||
RichHelp,
|
||||
TargetSpec,
|
||||
)
|
||||
from cli.lib.common.cli_helper import register_targets, RichHelp, TargetSpec
|
||||
from cli.lib.core.vllm import VllmBuildRunner
|
||||
|
||||
|
||||
@ -38,4 +34,4 @@ def register_build_commands(subparsers: argparse._SubParsersAction) -> None:
|
||||
description="Build third-party targets.\n\nAvailable targets:\n" + overview,
|
||||
formatter_class=RichHelp,
|
||||
)
|
||||
register_target_commands_and_runner(external_parser, _TARGETS)
|
||||
register_targets(external_parser, _TARGETS)
|
||||
|
||||
@ -32,14 +32,7 @@ class RichHelp(
|
||||
|
||||
|
||||
class TargetSpec(TypedDict, total=False):
|
||||
# Registry entry for a CLI subcommand.
|
||||
#
|
||||
# Keys:
|
||||
# runner – class implementing BaseRunner; will be constructed
|
||||
# with the parsed argparse Namespace.
|
||||
# help – short help text for the subparser (shows in parent help).
|
||||
# description – long help text (falls back to runner.__doc__ if missing).
|
||||
# add_arguments – optional function to add subparser-specific CLI args.
|
||||
"""CLI subcommand specification with bA."""
|
||||
|
||||
runner: Required[type[BaseRunner]]
|
||||
help: str
|
||||
@ -47,45 +40,31 @@ class TargetSpec(TypedDict, total=False):
|
||||
add_arguments: Callable[[argparse.ArgumentParser], None]
|
||||
|
||||
|
||||
def register_target_commands_and_runner(
|
||||
def register_targets(
|
||||
parser: argparse.ArgumentParser,
|
||||
target_specs: dict[str, TargetSpec],
|
||||
common_args: Callable[[argparse.ArgumentParser], None] = lambda _: None,
|
||||
placeholder_name: str = "target",
|
||||
) -> None:
|
||||
"""
|
||||
Given an argparse parser and a mapping of target names → TargetSpec,
|
||||
register each target as a subcommand and wire it to its runner class.
|
||||
- Creates a subparser for each target name.
|
||||
- description: defaults to TargetSpec['description'] or runner.__doc__.
|
||||
- description has higher priority than runner.__doc__.
|
||||
- add_arguments[Optional]: add target-specific CLI args, if 'add_arguments' exists
|
||||
- common_args[Optional]: add shared CLI args to each target parser.
|
||||
- Sets parser defaults:
|
||||
func: lambda that constructs the runner with parsed args
|
||||
and calls its .run().
|
||||
_runner_class: stored runner class for introspection/testing.
|
||||
"""
|
||||
"""Register target subcommands."""
|
||||
targets = parser.add_subparsers(
|
||||
dest=placeholder_name,
|
||||
dest="target",
|
||||
required=True,
|
||||
metavar="{" + ",".join(target_specs.keys()) + "}",
|
||||
)
|
||||
|
||||
for name, spec in target_specs.items():
|
||||
desc = (spec.get("description") or (spec["runner"].__doc__ or "")).strip()
|
||||
desc = spec.get("description") or spec["runner"].__doc__ or ""
|
||||
|
||||
p = targets.add_parser(
|
||||
name,
|
||||
help=spec.get("help", ""),
|
||||
description=desc,
|
||||
description=desc.strip(),
|
||||
formatter_class=RichHelp,
|
||||
)
|
||||
p.set_defaults(
|
||||
func=lambda args, _cls=spec["runner"]: _cls(args).run(),
|
||||
func=lambda args, cls=spec["runner"]: cls(args).run(),
|
||||
_runner_class=spec["runner"],
|
||||
)
|
||||
|
||||
if "add_arguments" in spec and callable(spec["add_arguments"]):
|
||||
spec["add_arguments"](p)
|
||||
if common_args:
|
||||
|
||||
42
.ci/lumen_cli/cli/lib/common/docker_helper.py
Normal file
42
.ci/lumen_cli/cli/lib/common/docker_helper.py
Normal file
@ -0,0 +1,42 @@
|
||||
"""
|
||||
Docker Utility helpers for CLI tasks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import docker
|
||||
from docker.errors import APIError, NotFound
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# lazy singleton so we don't reconnect every call
|
||||
_docker_client: Optional[docker.DockerClient] = None
|
||||
|
||||
|
||||
def _get_client() -> docker.DockerClient:
|
||||
global _docker_client
|
||||
if _docker_client is None:
|
||||
_docker_client = docker.from_env()
|
||||
return _docker_client
|
||||
|
||||
|
||||
def local_image_exists(
|
||||
image_name: str, client: Optional[docker.DockerClient] = None
|
||||
) -> bool:
|
||||
"""Return True if a local Docker image exists."""
|
||||
if not image_name:
|
||||
return False
|
||||
|
||||
client = client or _get_client()
|
||||
try:
|
||||
client.images.get(image_name)
|
||||
return True
|
||||
except (NotFound, APIError) as e:
|
||||
logger.error(
|
||||
"Error when checking Docker image '%s': %s",
|
||||
image_name,
|
||||
e.explanation if hasattr(e, "explanation") else str(e),
|
||||
)
|
||||
return False
|
||||
110
.ci/lumen_cli/cli/lib/common/envs_helper.py
Normal file
110
.ci/lumen_cli/cli/lib/common/envs_helper.py
Normal file
@ -0,0 +1,110 @@
|
||||
"""
|
||||
Environment Variables and Dataclasses Utility helpers for CLI tasks.
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import field, fields, is_dataclass, MISSING
|
||||
from pathlib import Path
|
||||
from textwrap import indent
|
||||
from typing import Optional, Union
|
||||
|
||||
from cli.lib.common.utils import str2bool
|
||||
|
||||
|
||||
def get_env(name: str, default: str = "") -> str:
|
||||
"""Get environment variable with default fallback."""
|
||||
return os.environ.get(name) or default
|
||||
|
||||
|
||||
def env_path_optional(
|
||||
name: str,
|
||||
default: Optional[Union[str, Path]] = None,
|
||||
resolve: bool = True,
|
||||
) -> Optional[Path]:
|
||||
"""Get environment variable as optional Path."""
|
||||
val = get_env(name) or default
|
||||
if not val:
|
||||
return None
|
||||
|
||||
path = Path(val)
|
||||
return path.resolve() if resolve else path
|
||||
|
||||
|
||||
def env_path(
|
||||
name: str,
|
||||
default: Optional[Union[str, Path]] = None,
|
||||
resolve: bool = True,
|
||||
) -> Path:
|
||||
"""Get environment variable as Path, raise if missing."""
|
||||
path = env_path_optional(name, default, resolve)
|
||||
if not path:
|
||||
raise ValueError(f"Missing path value for {name}")
|
||||
return path
|
||||
|
||||
|
||||
def env_bool(
|
||||
name: str,
|
||||
default: bool = False,
|
||||
) -> bool:
|
||||
val = get_env(name)
|
||||
if not val:
|
||||
return default
|
||||
return str2bool(val)
|
||||
|
||||
|
||||
def env_bool_field(
|
||||
name: str,
|
||||
default: bool = False,
|
||||
):
|
||||
return field(default_factory=lambda: env_bool(name, default))
|
||||
|
||||
|
||||
def env_path_field(
|
||||
name: str,
|
||||
default: Union[str, Path] = "",
|
||||
*,
|
||||
resolve: bool = True,
|
||||
) -> Path:
|
||||
return field(default_factory=lambda: env_path(name, default, resolve=resolve))
|
||||
|
||||
|
||||
def env_str_field(
|
||||
name: str,
|
||||
default: str = "",
|
||||
) -> str:
|
||||
return field(default_factory=lambda: get_env(name, default))
|
||||
|
||||
|
||||
def generate_dataclass_help(cls) -> str:
|
||||
"""Auto-generate help text for dataclass fields."""
|
||||
if not is_dataclass(cls):
|
||||
raise TypeError(f"{cls} is not a dataclass")
|
||||
|
||||
def get_value(f):
|
||||
if f.default is not MISSING:
|
||||
return f.default
|
||||
if f.default_factory is not MISSING:
|
||||
try:
|
||||
return f.default_factory()
|
||||
except Exception as e:
|
||||
return f"<error: {e}>"
|
||||
return "<required>"
|
||||
|
||||
lines = [f"{f.name:<22} = {repr(get_value(f))}" for f in fields(cls)]
|
||||
return indent("\n".join(lines), " ")
|
||||
|
||||
|
||||
def with_params_help(params_cls: type, title: str = "Parameter defaults"):
|
||||
"""
|
||||
Class decorator that appends a help table generated from another dataclass
|
||||
(e.g., VllmParameters) to the decorated class's docstring.
|
||||
"""
|
||||
if not is_dataclass(params_cls):
|
||||
raise TypeError(f"{params_cls} must be a dataclass")
|
||||
|
||||
def _decorator(cls: type) -> type:
|
||||
block = generate_dataclass_help(params_cls)
|
||||
cls.__doc__ = (cls.__doc__ or "") + f"\n\n{title}:\n{block}"
|
||||
return cls
|
||||
|
||||
return _decorator
|
||||
69
.ci/lumen_cli/cli/lib/common/git_helper.py
Normal file
69
.ci/lumen_cli/cli/lib/common/git_helper.py
Normal file
@ -0,0 +1,69 @@
|
||||
"""
|
||||
Git Utility helpers for CLI tasks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from cli.lib.common.path_helper import remove_dir
|
||||
from git import GitCommandError, RemoteProgress, Repo
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PrintProgress(RemoteProgress):
|
||||
"""Simple progress logger for git operations."""
|
||||
|
||||
def __init__(self, interval: int = 5):
|
||||
super().__init__()
|
||||
self._last_percent = -1
|
||||
self._interval = interval
|
||||
|
||||
def update(self, op_code, cur, max=None, message=""):
|
||||
msg = self._cur_line or message
|
||||
if max and cur:
|
||||
percent = int(cur / max * 100)
|
||||
if percent != self._last_percent and percent % self._interval == 0:
|
||||
self._last_percent = percent
|
||||
logger.info("Progress: %d%% - %s", percent, msg)
|
||||
elif msg:
|
||||
logger.info(msg)
|
||||
|
||||
|
||||
def clone_external_repo(target: str, repo: str, dst: str = "", update_submodules=False):
|
||||
"""Clone repository with pinned commit and optional submodules."""
|
||||
dst = dst or target
|
||||
|
||||
try:
|
||||
logger.info("Cloning %s to %s", target, dst)
|
||||
|
||||
# Clone and fetch
|
||||
remove_dir(dst)
|
||||
r = Repo.clone_from(repo, dst, progress=PrintProgress())
|
||||
r.git.fetch("--all", "--tags")
|
||||
|
||||
# Checkout pinned commit
|
||||
commit = get_post_build_pinned_commit(target)
|
||||
logger.info("Checking out pinned commit %s", commit)
|
||||
r.git.checkout(commit)
|
||||
|
||||
# Update submodules if requested
|
||||
if update_submodules and r.submodules:
|
||||
logger.info("Updating %d submodule(s)", len(r.submodules))
|
||||
for sm in r.submodules:
|
||||
sm.update(init=True, recursive=True, progress=PrintProgress())
|
||||
|
||||
logger.info("Successfully cloned %s", target)
|
||||
return r
|
||||
|
||||
except GitCommandError as e:
|
||||
logger.error("Git operation failed: %s", e)
|
||||
raise
|
||||
|
||||
|
||||
def get_post_build_pinned_commit(name: str, prefix=".github/ci_commit_pins") -> str:
|
||||
path = Path(prefix) / f"{name}.txt"
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Pin file not found: {path}")
|
||||
return path.read_text(encoding="utf-8").strip()
|
||||
@ -1,9 +1,62 @@
|
||||
"""
|
||||
File And Path Utility helpers for CLI tasks.
|
||||
"""
|
||||
"""Path utility helpers for CLI tasks."""
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
# TODO(elainewy): Add path_helper utils
|
||||
|
||||
|
||||
def get_path(path: Union[str, Path], resolve: bool = False) -> Path:
|
||||
"""Convert to Path object, optionally resolving to absolute path."""
|
||||
if not path:
|
||||
raise ValueError("Path cannot be None or empty")
|
||||
result = Path(path)
|
||||
return result.resolve() if resolve else result
|
||||
|
||||
|
||||
def ensure_dir_exists(path: Union[str, Path]) -> Path:
|
||||
"""Create directory if it doesn't exist."""
|
||||
path_obj = get_path(path)
|
||||
path_obj.mkdir(parents=True, exist_ok=True)
|
||||
return path_obj
|
||||
|
||||
|
||||
def remove_dir(path: Union[str, Path, None]) -> None:
|
||||
"""Remove directory if it exists."""
|
||||
if not path:
|
||||
return
|
||||
path_obj = get_path(path)
|
||||
if path_obj.exists():
|
||||
shutil.rmtree(path_obj)
|
||||
|
||||
|
||||
def force_create_dir(path: Union[str, Path]) -> Path:
|
||||
"""Remove directory if exists, then create fresh empty directory."""
|
||||
remove_dir(path)
|
||||
return ensure_dir_exists(path)
|
||||
|
||||
|
||||
def copy(src: Union[str, Path], dst: Union[str, Path]) -> None:
|
||||
"""Copy file or directory from src to dst."""
|
||||
src_path = get_path(src, resolve=True)
|
||||
dst_path = get_path(dst, resolve=True)
|
||||
|
||||
if not src_path.exists():
|
||||
raise FileNotFoundError(f"Source does not exist: {src_path}")
|
||||
|
||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if src_path.is_file():
|
||||
shutil.copy2(src_path, dst_path)
|
||||
elif src_path.is_dir():
|
||||
shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
|
||||
else:
|
||||
raise ValueError(f"Unsupported path type: {src_path}")
|
||||
|
||||
|
||||
def is_path_exist(path: Union[str, Path, None]) -> bool:
|
||||
"""Check if path exists."""
|
||||
return bool(path and get_path(path).exists())
|
||||
|
||||
@ -3,7 +3,77 @@ General Utility helpers for CLI tasks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
# TODO(elainewy): Add common utils
|
||||
|
||||
|
||||
def run_command(
|
||||
cmd: str,
|
||||
use_shell: bool = False,
|
||||
log_cmd: bool = True,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[dict] = None,
|
||||
check: bool = True,
|
||||
) -> int:
|
||||
"""Run a command with optional shell execution."""
|
||||
if use_shell:
|
||||
args = cmd
|
||||
log_prefix = "[shell]"
|
||||
executable = "/bin/bash"
|
||||
else:
|
||||
args = shlex.split(cmd)
|
||||
log_prefix = "[cmd]"
|
||||
executable = None
|
||||
|
||||
if log_cmd:
|
||||
display_cmd = cmd if use_shell else " ".join(args)
|
||||
logger.info("%s %s", log_prefix, display_cmd)
|
||||
|
||||
run_env = {**os.environ, **(env or {})}
|
||||
|
||||
proc = subprocess.run(
|
||||
args,
|
||||
shell=use_shell,
|
||||
executable=executable,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
cwd=cwd,
|
||||
env=run_env,
|
||||
check=False,
|
||||
)
|
||||
|
||||
if check and proc.returncode != 0:
|
||||
logger.error(
|
||||
"%s Command failed (exit %s): %s", log_prefix, proc.returncode, cmd
|
||||
)
|
||||
raise subprocess.CalledProcessError(
|
||||
proc.returncode, args if not use_shell else cmd
|
||||
)
|
||||
|
||||
return proc.returncode
|
||||
|
||||
|
||||
def str2bool(value: Optional[str]) -> bool:
|
||||
"""Convert environment variables to boolean values."""
|
||||
if not value:
|
||||
return False
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(
|
||||
f"Expected a string value for boolean conversion, got {type(value)}"
|
||||
)
|
||||
value = value.strip().lower()
|
||||
|
||||
true_value_set = {"1", "true", "t", "yes", "y", "on", "enable", "enabled", "found"}
|
||||
false_value_set = {"0", "false", "f", "no", "n", "off", "disable"}
|
||||
|
||||
if value in true_value_set:
|
||||
return True
|
||||
if value in false_value_set:
|
||||
return False
|
||||
raise ValueError(f"Invalid string value for boolean conversion: {value}")
|
||||
|
||||
@ -1,15 +1,263 @@
|
||||
import logging
|
||||
import os
|
||||
import textwrap
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from cli.lib.common.cli_helper import BaseRunner
|
||||
from cli.lib.common.docker_helper import local_image_exists
|
||||
from cli.lib.common.envs_helper import (
|
||||
env_bool_field,
|
||||
env_path_field,
|
||||
env_str_field,
|
||||
with_params_help,
|
||||
)
|
||||
from cli.lib.common.git_helper import clone_external_repo
|
||||
from cli.lib.common.path_helper import (
|
||||
copy,
|
||||
ensure_dir_exists,
|
||||
force_create_dir,
|
||||
get_path,
|
||||
is_path_exist,
|
||||
)
|
||||
from cli.lib.common.utils import run_command
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VllmBuildRunner(BaseRunner):
|
||||
# Default path for docker build artifacts
|
||||
_DEFAULT_RESULT_PATH = "./shared"
|
||||
|
||||
# Temp folder in vllm work place to cp torch whls in vllm work directory for docker build
|
||||
_VLLM_TEMP_FOLDER = "tmp"
|
||||
|
||||
|
||||
@dataclass
|
||||
class VllmBuildParameters:
|
||||
"""
|
||||
Build vllm whels in ci
|
||||
Parameters defining the vllm external input configurations.
|
||||
Combine with VllmDockerBuildArgs to define the vllm build environment
|
||||
"""
|
||||
|
||||
# USE_TORCH_WHEEL: when true, use local Torch wheels; requires TORCH_WHEELS_PATH.
|
||||
# Otherwise docker build pull torch nightly during build
|
||||
# TORCH_WHEELS_PATH: directory containing local torch wheels when use_torch_whl is True
|
||||
use_torch_whl: bool = env_bool_field("USE_TORCH_WHEEL", True)
|
||||
torch_whls_path: Path = env_path_field("TORCH_WHEELS_PATH", "./dist")
|
||||
|
||||
# USE_LOCAL_BASE_IMAGE: when true, use an existing local Docker base image; requires BASE_IMAGE
|
||||
# Otherwise, pull dockerfile's default image remotely
|
||||
# BASE_IMAGE: name:tag (only needed when use_local_base_image is True)
|
||||
use_local_base_image: bool = env_bool_field("USE_LOCAL_BASE_IMAGE", True)
|
||||
base_image: str = env_str_field("BASE_IMAGE")
|
||||
|
||||
# USE_LOCAL_DOCKERFILE: when true("1"), use a local Dockerfile; requires DOCKERFILE_PATH.
|
||||
# otherwise, use vllm's default dockerfile.torch_nightly for build
|
||||
# DOCKERFILE_PATH: path to Dockerfile used when use_local_dockerfile is True"
|
||||
use_local_dockerfile: bool = env_bool_field("USE_LOCAL_DOCKERFILE", True)
|
||||
dockerfile_path: Path = env_path_field(
|
||||
"DOCKERFILE_PATH", ".github/ci_configs/vllm/Dockerfile.tmp_vllm"
|
||||
)
|
||||
|
||||
# OUTPUT_DIR: where docker buildx (local exporter) will write artifacts
|
||||
output_dir: Path = env_path_field("OUTPUT_DIR", "shared")
|
||||
|
||||
# --- Build args ----------------------------------------------------------
|
||||
target_stage: str = env_str_field("TARGET_STAGE", "export-wheels")
|
||||
|
||||
tag_name: str = env_str_field("TAG", "vllm-wheels")
|
||||
|
||||
cuda_version: str = env_str_field("CUDA_VERSION", "12.8.1")
|
||||
|
||||
python_version: str = env_str_field("PYTHON_VERSION", "3.12")
|
||||
|
||||
max_jobs: str = env_str_field("MAX_JOBS", "64")
|
||||
|
||||
sccache_bucket: str = env_str_field("SCCACHE_BUCKET")
|
||||
|
||||
sccache_region: str = env_str_field("SCCACHE_REGION")
|
||||
|
||||
torch_cuda_arch_list: str = env_str_field("TORCH_CUDA_ARCH_LIST", "8.9")
|
||||
|
||||
def __post_init__(self):
|
||||
checks = [
|
||||
(
|
||||
self.use_torch_whl, # flag
|
||||
True, # trigger_value
|
||||
"torch_whls_path", # resource
|
||||
is_path_exist, # check_func
|
||||
"TORCH_WHEELS_PATH is not provided, but USE_TORCH_WHEEL is set to 1",
|
||||
),
|
||||
(
|
||||
self.use_local_base_image,
|
||||
True,
|
||||
"base_image",
|
||||
local_image_exists,
|
||||
f"BASE_IMAGE {self.base_image} does not found, but USE_LOCAL_BASE_IMAGE is set to 1",
|
||||
),
|
||||
(
|
||||
self.use_local_dockerfile,
|
||||
True,
|
||||
"dockerfile_path",
|
||||
is_path_exist,
|
||||
" DOCKERFILE_PATH path does not found, but USE_LOCAL_DOCKERFILE is set to 1",
|
||||
),
|
||||
]
|
||||
for flag, trigger_value, attr_name, check_func, error_msg in checks:
|
||||
value = getattr(self, attr_name)
|
||||
if flag == trigger_value:
|
||||
if not value or not check_func(value):
|
||||
raise ValueError(error_msg)
|
||||
else:
|
||||
logger.info("flag %s is not set", flag)
|
||||
if not self.output_dir:
|
||||
raise ValueError("missing required output_dir")
|
||||
|
||||
|
||||
@with_params_help(VllmBuildParameters)
|
||||
class VllmBuildRunner(BaseRunner):
|
||||
"""
|
||||
Build vLLM using docker buildx.
|
||||
|
||||
Environment variable options:
|
||||
"USE_TORCH_WHEEL": "1: use local wheels; 0: pull nightly from pypi",
|
||||
"TORCH_WHEELS_PATH": "Path to local wheels (when USE_TORCH_WHEEL=1)",
|
||||
|
||||
"USE_LOCAL_BASE_IMAGE": "1: use local base image; 0: default image",
|
||||
"BASE_IMAGE": "name:tag to indicate base image the dockerfile depends on (when USE_LOCAL_BASE_IMAGE=1)",
|
||||
|
||||
"USE_LOCAL_DOCKERFILE": "1: use local Dockerfile; 0: vllm repo default dockerfile.torch_nightly",
|
||||
"DOCKERFILE_PATH": "Path to Dockerfile (when USE_LOCAL_DOCKERFILE=1)",
|
||||
|
||||
"OUTPUT_DIR": "e.g. './shared'",
|
||||
|
||||
"TORCH_CUDA_ARCH_LIST": "e.g. '8.0' or '8.0;9.0'",
|
||||
"CUDA_VERSION": "e.g. '12.8.1'",
|
||||
"PYTHON_VERSION": "e.g. '3.12'",
|
||||
"MAX_JOBS": "e.g. '64'",
|
||||
"SCCACHE_BUCKET": "e.g. 'my-bucket'",
|
||||
"SCCACHE_REGION": "e.g. 'us-west-2'",
|
||||
"""
|
||||
|
||||
def __init__(self, args=None):
|
||||
self.work_directory = "vllm"
|
||||
|
||||
def run(self):
|
||||
logger.info("Running vllm build")
|
||||
"""
|
||||
main function to run vllm build
|
||||
1. prepare vllm build environment
|
||||
2. prepare the docker build command args
|
||||
3. run docker build
|
||||
"""
|
||||
inputs = VllmBuildParameters()
|
||||
clone_vllm()
|
||||
|
||||
self.cp_dockerfile_if_exist(inputs)
|
||||
|
||||
# cp torch wheels from root direct to vllm workspace if exist
|
||||
self.cp_torch_whls_if_exist(inputs)
|
||||
|
||||
ensure_dir_exists(inputs.output_dir)
|
||||
|
||||
cmd = self._generate_docker_build_cmd(inputs)
|
||||
logger.info("Running docker build: \n %s", cmd)
|
||||
run_command(cmd, cwd="vllm", env=os.environ.copy())
|
||||
|
||||
def cp_torch_whls_if_exist(self, inputs: VllmBuildParameters) -> str:
|
||||
if not inputs.use_torch_whl:
|
||||
return ""
|
||||
tmp_dir = f"./{self.work_directory}/{_VLLM_TEMP_FOLDER}"
|
||||
tmp_path = Path(tmp_dir)
|
||||
force_create_dir(tmp_path)
|
||||
copy(inputs.torch_whls_path, tmp_dir)
|
||||
return tmp_dir
|
||||
|
||||
def cp_dockerfile_if_exist(self, inputs: VllmBuildParameters):
|
||||
if not inputs.use_local_dockerfile:
|
||||
logger.info("using vllm default dockerfile.torch_nightly for build")
|
||||
return
|
||||
dockerfile_path = get_path(inputs.dockerfile_path, resolve=True)
|
||||
vllm_torch_dockerfile = Path(
|
||||
f"./{self.work_directory}/docker/Dockerfile.nightly_torch"
|
||||
)
|
||||
copy(dockerfile_path, vllm_torch_dockerfile)
|
||||
|
||||
def get_result_path(self, path):
|
||||
"""
|
||||
Get the absolute path of the result path
|
||||
"""
|
||||
if not path:
|
||||
path = _DEFAULT_RESULT_PATH
|
||||
abs_path = get_path(path, resolve=True)
|
||||
return abs_path
|
||||
|
||||
def _get_torch_wheel_path_arg(self, torch_whl_dir: Optional[Path]) -> str:
|
||||
if not torch_whl_dir:
|
||||
return ""
|
||||
return f"--build-arg TORCH_WHEELS_PATH={_VLLM_TEMP_FOLDER}"
|
||||
|
||||
def _get_base_image_args(self, inputs: VllmBuildParameters) -> tuple[str, str, str]:
|
||||
"""
|
||||
Returns:
|
||||
- base_image_arg: docker buildx arg string for base image
|
||||
- final_base_image_arg: docker buildx arg string for vllm-base stage
|
||||
- pull_flag: --pull=true or --pull=false depending on whether the image exists locally
|
||||
"""
|
||||
if not inputs.use_local_base_image:
|
||||
return "", "", ""
|
||||
|
||||
base_image = inputs.base_image
|
||||
|
||||
# set both base image and final base image to the same local image
|
||||
base_image_arg = f"--build-arg BUILD_BASE_IMAGE={base_image}"
|
||||
final_base_image_arg = f"--build-arg FINAL_BASE_IMAGE={base_image}"
|
||||
|
||||
if local_image_exists(base_image):
|
||||
pull_flag = "--pull=false"
|
||||
return base_image_arg, final_base_image_arg, pull_flag
|
||||
logger.info(
|
||||
"[INFO] Local image not found:%s will try to pull from remote", {base_image}
|
||||
)
|
||||
return base_image_arg, final_base_image_arg, ""
|
||||
|
||||
def _generate_docker_build_cmd(
|
||||
self,
|
||||
inputs: VllmBuildParameters,
|
||||
) -> str:
|
||||
base_image_arg, final_base_image_arg, pull_flag = self._get_base_image_args(
|
||||
inputs
|
||||
)
|
||||
torch_arg = self._get_torch_wheel_path_arg(inputs.torch_whls_path)
|
||||
|
||||
return textwrap.dedent(
|
||||
f"""
|
||||
docker buildx build \
|
||||
--output type=local,dest={inputs.output_dir} \
|
||||
-f docker/Dockerfile.nightly_torch \
|
||||
{pull_flag} \
|
||||
{torch_arg} \
|
||||
{base_image_arg} \
|
||||
{final_base_image_arg} \
|
||||
--build-arg max_jobs={inputs.max_jobs} \
|
||||
--build-arg CUDA_VERSION={inputs.cuda_version} \
|
||||
--build-arg PYTHON_VERSION={inputs.python_version} \
|
||||
--build-arg USE_SCCACHE={int(bool(inputs.sccache_bucket and inputs.sccache_region))} \
|
||||
--build-arg SCCACHE_BUCKET_NAME={inputs.sccache_bucket} \
|
||||
--build-arg SCCACHE_REGION_NAME={inputs.sccache_region} \
|
||||
--build-arg torch_cuda_arch_list='{inputs.torch_cuda_arch_list}' \
|
||||
--target {inputs.target_stage} \
|
||||
-t {inputs.tag_name} \
|
||||
--progress=plain .
|
||||
"""
|
||||
).strip()
|
||||
|
||||
|
||||
def clone_vllm():
|
||||
clone_external_repo(
|
||||
target="vllm",
|
||||
repo="https://github.com/vllm-project/vllm.git",
|
||||
dst="vllm",
|
||||
update_submodules=True,
|
||||
)
|
||||
|
||||
@ -14,10 +14,6 @@ def main():
|
||||
# Define top-level parser
|
||||
parser = argparse.ArgumentParser(description="Lumos CLI")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
# Add top-level args
|
||||
parser.add_argument(
|
||||
"--config", required=False, help="Path to config file for build and test"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-level", default="INFO", help="Log level (DEBUG, INFO, WARNING, ERROR)"
|
||||
)
|
||||
|
||||
@ -2,7 +2,10 @@
|
||||
name = "lumen-ci"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"pytest==7.3.2"
|
||||
"pyyaml==6.0.2",
|
||||
"GitPython==3.1.45",
|
||||
"docker==7.1.0",
|
||||
"pytest==7.3.2",
|
||||
]
|
||||
|
||||
[tool.setuptools]
|
||||
|
||||
@ -9,18 +9,20 @@ from cli.run import main
|
||||
|
||||
|
||||
class TestArgparseCLI(unittest.TestCase):
|
||||
def test_cli_run_build_external(self):
|
||||
@patch("cli.build_cli.register_build.VllmBuildRunner.run", return_value=None)
|
||||
@patch("cli.build_cli.register_build.VllmBuildRunner.__init__", return_value=None)
|
||||
def test_cli_run_build_external(self, mock_init, mock_run):
|
||||
from cli.run import main # import after patches if needed
|
||||
|
||||
test_args = ["cli.run", "build", "external", "vllm"]
|
||||
|
||||
with patch.object(sys, "argv", test_args):
|
||||
with self.assertLogs(level="INFO") as caplog:
|
||||
# if argparse could exit on error, wrap in try/except SystemExit if needed
|
||||
# argparse may call sys.exit on error; capture to avoid test aborts
|
||||
try:
|
||||
main()
|
||||
|
||||
# stdout print from your CLI plumbing
|
||||
# logs emitted inside your code (info/debug/error etc.)
|
||||
logs_text = "\n".join(caplog.output)
|
||||
self.assertIn("Running vllm build", logs_text)
|
||||
except SystemExit:
|
||||
pass
|
||||
mock_init.assert_called_once() # got constructed
|
||||
mock_run.assert_called_once_with() # run() called
|
||||
|
||||
def test_build_help(self):
|
||||
test_args = ["cli.run", "build", "--help"]
|
||||
|
||||
@ -4,12 +4,7 @@ import unittest
|
||||
from contextlib import redirect_stderr
|
||||
from unittest.mock import patch
|
||||
|
||||
from cli.lib.common.cli_helper import (
|
||||
BaseRunner,
|
||||
register_target_commands_and_runner,
|
||||
RichHelp,
|
||||
TargetSpec,
|
||||
)
|
||||
from cli.lib.common.cli_helper import BaseRunner, register_targets, RichHelp, TargetSpec
|
||||
|
||||
|
||||
# ---- Dummy runners for unittests----
|
||||
@ -35,7 +30,7 @@ def common_args(p: argparse.ArgumentParser) -> None:
|
||||
|
||||
def build_parser(specs: dict[str, TargetSpec]) -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(prog="app", formatter_class=RichHelp)
|
||||
register_target_commands_and_runner(
|
||||
register_targets(
|
||||
parser=parser,
|
||||
target_specs=specs,
|
||||
common_args=common_args,
|
||||
@ -89,7 +84,7 @@ class TestRegisterTargets(unittest.TestCase):
|
||||
patch.object(FooRunner, "run", return_value=None) as mock_run,
|
||||
):
|
||||
ns = parser.parse_args(["foo", "--x", "3", "--verbose"])
|
||||
ns.func(ns) # set by register_target_commands_and_runner
|
||||
ns.func(ns) # set by register_targets
|
||||
# __init__ received the Namespace
|
||||
self.assertEqual(mock_init.call_count, 1)
|
||||
(called_ns,), _ = mock_init.call_args
|
||||
|
||||
75
.ci/lumen_cli/tests/test_docker_helper.py
Normal file
75
.ci/lumen_cli/tests/test_docker_helper.py
Normal file
@ -0,0 +1,75 @@
|
||||
import unittest
|
||||
from unittest import mock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import docker.errors as derr
|
||||
from cli.lib.common.docker_helper import _get_client, local_image_exists
|
||||
|
||||
|
||||
class TestDockerImageHelpers(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Reset the singleton in the target module
|
||||
patcher = mock.patch("cli.lib.common.docker_helper._docker_client", None)
|
||||
self.addCleanup(patcher.stop)
|
||||
patcher.start()
|
||||
|
||||
def test_local_image_exists_true(self):
|
||||
# Mock a docker client whose images.get returns an object (no exception)
|
||||
mock_client = MagicMock()
|
||||
mock_client.images.get.return_value = object()
|
||||
ok = local_image_exists("repo:tag", client=mock_client)
|
||||
self.assertTrue(ok)
|
||||
|
||||
def test_local_image_exists_not_found_false(self):
|
||||
mock_client = MagicMock()
|
||||
# Raise docker.errors.NotFound
|
||||
mock_client.images.get.side_effect = derr.NotFound("nope")
|
||||
ok = local_image_exists("missing:latest", client=mock_client)
|
||||
self.assertFalse(ok)
|
||||
|
||||
def test_local_image_exists_api_error_false(self):
|
||||
mock_client = MagicMock()
|
||||
mock_client.images.get.side_effect = derr.APIError("boom", None)
|
||||
|
||||
ok = local_image_exists("broken:tag", client=mock_client)
|
||||
self.assertFalse(ok)
|
||||
|
||||
def test_local_image_exists_uses_lazy_singleton(self):
|
||||
# Patch docker.from_env used by _get_client()
|
||||
with mock.patch(
|
||||
"cli.lib.common.docker_helper.docker.from_env"
|
||||
) as mock_from_env:
|
||||
mock_docker_client = MagicMock()
|
||||
mock_from_env.return_value = mock_docker_client
|
||||
|
||||
# First call should create and cache the client
|
||||
c1 = _get_client()
|
||||
self.assertIs(c1, mock_docker_client)
|
||||
mock_from_env.assert_called_once()
|
||||
|
||||
# Second call should reuse cached client (no extra from_env calls)
|
||||
c2 = _get_client()
|
||||
self.assertIs(c2, mock_docker_client)
|
||||
mock_from_env.assert_called_once() # still once
|
||||
|
||||
def test_local_image_exists_without_client_param_calls_get_client_once(self):
|
||||
# Ensure _get_client is called and cached; local_image_exists should reuse it
|
||||
with mock.patch("cli.lib.common.docker_helper._get_client") as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# 1st call
|
||||
local_image_exists("repo:tag")
|
||||
# 2nd call
|
||||
local_image_exists("repo:tag2")
|
||||
|
||||
# local_image_exists should call _get_client each time,
|
||||
# but your _get_client itself caches docker.from_env.
|
||||
self.assertEqual(mock_get_client.call_count, 2)
|
||||
self.assertEqual(mock_client.images.get.call_count, 2)
|
||||
mock_client.images.get.assert_any_call("repo:tag")
|
||||
mock_client.images.get.assert_any_call("repo:tag2")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
149
.ci/lumen_cli/tests/test_envs_helper.py
Normal file
149
.ci/lumen_cli/tests/test_envs_helper.py
Normal file
@ -0,0 +1,149 @@
|
||||
import os
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import cli.lib.common.envs_helper as m
|
||||
|
||||
|
||||
class TestEnvHelpers(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Keep a copy of the original environment to restore later
|
||||
self._env_backup = dict(os.environ)
|
||||
|
||||
def tearDown(self):
|
||||
# Restore environment to original state
|
||||
os.environ.clear()
|
||||
os.environ.update(self._env_backup)
|
||||
|
||||
# -------- get_env --------
|
||||
def test_get_env_unset_returns_default(self):
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.assertEqual(m.get_env("FOO", "default"), "default")
|
||||
|
||||
def test_get_env_empty_returns_default(self):
|
||||
with patch.dict(os.environ, {"FOO": ""}, clear=True):
|
||||
self.assertEqual(m.get_env("FOO", "default"), "default")
|
||||
|
||||
def test_get_env_set_returns_value(self):
|
||||
with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
|
||||
self.assertEqual(m.get_env("FOO", "default"), "bar")
|
||||
|
||||
def test_get_env_not_exist_returns_default(self):
|
||||
with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
|
||||
self.assertEqual(m.get_env("TEST_NOT_EXIST", "default"), "default")
|
||||
|
||||
def test_get_env_not_exist_without_default(self):
|
||||
with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
|
||||
self.assertEqual(m.get_env("TEST_NOT_EXIST"), "")
|
||||
|
||||
# -------- env_bool --------
|
||||
def test_env_bool_uses_default_when_unset(self):
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.assertTrue(m.env_bool("FLAG", default=True))
|
||||
self.assertFalse(m.env_bool("FLAG", default=False))
|
||||
|
||||
def test_env_bool_uses_str2bool_when_set(self):
|
||||
# Patch str2bool used by env_bool so we don't depend on its exact behavior
|
||||
def fake_str2bool(s: str) -> bool:
|
||||
return s.lower() in {"1", "true", "yes", "on", "y"}
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"FLAG": "yEs"}, clear=True),
|
||||
patch.object(m, "str2bool", fake_str2bool),
|
||||
):
|
||||
self.assertTrue(m.env_bool("FLAG", default=False))
|
||||
|
||||
# -------- env_path_optional / env_path --------
|
||||
def test_env_path_optional_unset_returns_none_by_default(self):
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.assertIsNone(m.env_path_optional("P"))
|
||||
|
||||
def test_env_path_optional_unset_returns_none_when_env_var_is_empty(self):
|
||||
with patch.dict(os.environ, {"P": ""}, clear=True):
|
||||
self.assertIsNone(m.env_path_optional("P"))
|
||||
|
||||
def test_env_path_optional_unset_returns_default_str(self):
|
||||
# default as string; resolve=True by default -> absolute path
|
||||
default_str = "x/y"
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
p = m.env_path_optional("P", default=default_str)
|
||||
self.assertIsInstance(p, Path)
|
||||
self.assertIsNotNone(p)
|
||||
if p:
|
||||
self.assertTrue(p.is_absolute())
|
||||
self.assertEqual(p.parts[-2:], ("x", "y"))
|
||||
|
||||
def test_env_path_optional_unset_returns_default_path_no_resolve(self):
|
||||
d = Path("z")
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
p = m.env_path_optional("P", default=d, resolve=False)
|
||||
self.assertEqual(p, d)
|
||||
|
||||
def test_env_path_optional_respects_resolve_true(self):
|
||||
with patch.dict(os.environ, {"P": "a/b"}, clear=True):
|
||||
p = m.env_path_optional("P", resolve=True)
|
||||
self.assertIsInstance(p, Path)
|
||||
if p:
|
||||
self.assertTrue(p.is_absolute())
|
||||
|
||||
def test_env_path_optional_respects_resolve_false(self):
|
||||
with patch.dict(os.environ, {"P": "rel/dir"}, clear=True):
|
||||
p = m.env_path_optional("P", resolve=False)
|
||||
self.assertEqual(p, Path("rel/dir"))
|
||||
if p:
|
||||
self.assertFalse(p.is_absolute())
|
||||
|
||||
def test_env_path_raises_when_missing_and_default_none(self):
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with self.assertRaises(ValueError):
|
||||
m.env_path("P", None, resolve=True)
|
||||
|
||||
def test_env_path_returns_path_when_present(self):
|
||||
tmp = Path("./b").resolve()
|
||||
with patch.dict(os.environ, {"P": str(tmp)}, clear=True):
|
||||
p = m.env_path("P", None, resolve=True)
|
||||
self.assertEqual(p, tmp)
|
||||
|
||||
# -------- dataclass field helpers --------
|
||||
def test_dataclass_fields_read_env_at_instantiation(self):
|
||||
@dataclass
|
||||
class Cfg:
|
||||
flag: bool = m.env_bool_field("FLAG", default=False)
|
||||
out: Path = m.env_path_field("OUT", default="ab", resolve=True)
|
||||
name: str = m.env_str_field("NAME", default="anon")
|
||||
|
||||
# First instantiation
|
||||
with patch.dict(
|
||||
os.environ, {"FLAG": "true", "OUT": "outdir", "NAME": "alice"}, clear=True
|
||||
):
|
||||
cfg1 = Cfg()
|
||||
self.assertTrue(cfg1.flag)
|
||||
self.assertIsInstance(cfg1.out, Path)
|
||||
self.assertTrue(cfg1.out.is_absolute())
|
||||
self.assertEqual(cfg1.name, "alice")
|
||||
cfg1.name = "bob" # change instance value
|
||||
self.assertEqual(cfg1.name, "bob") # change is reflected
|
||||
|
||||
# Change env; new instance should reflect new values
|
||||
with patch.dict(os.environ, {"FLAG": "false", "NAME": ""}, clear=True):
|
||||
cfg2 = Cfg()
|
||||
self.assertFalse(cfg2.flag) # str2bool("false") -> False
|
||||
self.assertTrue("ab" in str(cfg2.out))
|
||||
self.assertIsInstance(cfg2.out, Path)
|
||||
self.assertTrue(cfg2.out.is_absolute())
|
||||
self.assertEqual(cfg2.name, "anon") # empty -> fallback to default
|
||||
|
||||
def test_dataclass_path_field_with_default_value(self):
|
||||
@dataclass
|
||||
class C2:
|
||||
out: Path = m.env_path_field("OUT", default="some/dir", resolve=False)
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
c = C2()
|
||||
self.assertEqual(c.out, Path("some/dir"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
122
.ci/lumen_cli/tests/test_path_helper.py
Normal file
122
.ci/lumen_cli/tests/test_path_helper.py
Normal file
@ -0,0 +1,122 @@
|
||||
# test_path_utils.py
|
||||
# Run: pytest -q
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
from cli.lib.common.path_helper import (
|
||||
copy,
|
||||
ensure_dir_exists,
|
||||
force_create_dir,
|
||||
get_path,
|
||||
is_path_exist,
|
||||
remove_dir,
|
||||
)
|
||||
|
||||
|
||||
class TestPathHelper(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = TemporaryDirectory()
|
||||
self.tmp_path = Path(self.tmpdir.name)
|
||||
|
||||
def tearDown(self):
|
||||
self.tmpdir.cleanup()
|
||||
|
||||
# -------- get_path --------
|
||||
def test_get_path_returns_path_for_str(self):
|
||||
# Use relative path to avoid absolute-ness
|
||||
rel_str = "sub/f.txt"
|
||||
os.chdir(self.tmp_path)
|
||||
p = get_path(rel_str, resolve=False)
|
||||
self.assertIsInstance(p, Path)
|
||||
self.assertFalse(p.is_absolute())
|
||||
self.assertEqual(str(p), rel_str)
|
||||
|
||||
def test_get_path_resolves(self):
|
||||
rel_str = "sub/f.txt"
|
||||
p = get_path(str(self.tmp_path / rel_str), resolve=True)
|
||||
self.assertTrue(p.is_absolute())
|
||||
self.assertTrue(str(p).endswith(rel_str))
|
||||
|
||||
def test_get_path_with_path_input(self):
|
||||
p_in = self.tmp_path / "sub/f.txt"
|
||||
p_out = get_path(p_in, resolve=False)
|
||||
self.assertTrue(str(p_out) == str(p_in))
|
||||
|
||||
def test_get_path_with_none_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
get_path(None) # type: ignore[arg-type]
|
||||
|
||||
def test_get_path_invalid_type_raises(self):
|
||||
with self.assertRaises(TypeError):
|
||||
get_path(123) # type: ignore[arg-type]
|
||||
|
||||
# -------- ensure_dir_exists / force_create_dir / remove_dir --------
|
||||
def test_ensure_dir_exists_creates_and_is_idempotent(self):
|
||||
d = self.tmp_path / "made"
|
||||
ensure_dir_exists(d)
|
||||
self.assertTrue(d.exists() and d.is_dir())
|
||||
ensure_dir_exists(d)
|
||||
|
||||
def test_force_create_dir_clears_existing(self):
|
||||
d = self.tmp_path / "fresh"
|
||||
(d / "inner").mkdir(parents=True)
|
||||
(d / "inner" / "f.txt").write_text("x")
|
||||
force_create_dir(d)
|
||||
self.assertTrue(d.exists())
|
||||
self.assertEqual(list(d.iterdir()), [])
|
||||
|
||||
def test_remove_dir_none_is_noop(self):
|
||||
remove_dir(None) # type: ignore[arg-type]
|
||||
|
||||
def test_remove_dir_nonexistent_is_noop(self):
|
||||
ghost = self.tmp_path / "ghost"
|
||||
remove_dir(ghost)
|
||||
|
||||
def test_remove_dir_accepts_str(self):
|
||||
d = self.tmp_path / "to_rm"
|
||||
d.mkdir()
|
||||
remove_dir(str(d))
|
||||
self.assertFalse(d.exists())
|
||||
|
||||
# -------- copy --------
|
||||
def test_copy_file_to_file(self):
|
||||
src = self.tmp_path / "src.txt"
|
||||
dst = self.tmp_path / "out" / "dst.txt"
|
||||
src.write_text("hello")
|
||||
copy(src, dst)
|
||||
self.assertEqual(dst.read_text(), "hello")
|
||||
|
||||
def test_copy_dir_to_new_dir(self):
|
||||
src = self.tmp_path / "srcdir"
|
||||
(src / "a").mkdir(parents=True)
|
||||
(src / "a" / "f.txt").write_text("content")
|
||||
dst = self.tmp_path / "destdir"
|
||||
copy(src, dst)
|
||||
self.assertEqual((dst / "a" / "f.txt").read_text(), "content")
|
||||
|
||||
def test_copy_dir_into_existing_dir_overwrite_true_merges(self):
|
||||
src = self.tmp_path / "srcdir"
|
||||
dst = self.tmp_path / "destdir"
|
||||
(src / "x").mkdir(parents=True)
|
||||
(src / "x" / "new.txt").write_text("new")
|
||||
dst.mkdir()
|
||||
(dst / "existing.txt").write_text("old")
|
||||
copy(src, dst)
|
||||
self.assertEqual((dst / "existing.txt").read_text(), "old")
|
||||
self.assertEqual((dst / "x" / "new.txt").read_text(), "new")
|
||||
|
||||
def test_is_str_path_exist(self):
|
||||
p = self.tmp_path / "x.txt"
|
||||
p.write_text("1")
|
||||
self.assertTrue(is_path_exist(str(p)))
|
||||
self.assertTrue(is_path_exist(p))
|
||||
self.assertFalse(is_path_exist(str(self.tmp_path / "missing")))
|
||||
self.assertFalse(is_path_exist(self.tmp_path / "missing"))
|
||||
self.assertFalse(is_path_exist(""))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
181
.ci/lumen_cli/tests/test_vllm.py
Normal file
181
.ci/lumen_cli/tests/test_vllm.py
Normal file
@ -0,0 +1,181 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import cli.lib.core.vllm as vllm
|
||||
|
||||
|
||||
class TestVllmBuildParameters(unittest.TestCase):
|
||||
@patch("cli.lib.core.vllm.local_image_exists", return_value=True)
|
||||
@patch("cli.lib.core.vllm.is_path_exist", return_value=True)
|
||||
@patch(
|
||||
"cli.lib.common.envs_helper.env_path_optional",
|
||||
side_effect=lambda name, default=None, resolve=True: {
|
||||
"DOCKERFILE_PATH": Path("/abs/vllm/Dockerfile"),
|
||||
"TORCH_WHEELS_PATH": Path("/abs/dist"),
|
||||
"OUTPUT_DIR": Path("/abs/shared"),
|
||||
}.get(name, Path(default) if default is not None else None),
|
||||
)
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"USE_TORCH_WHEEL": "1",
|
||||
"USE_LOCAL_BASE_IMAGE": "1",
|
||||
"USE_LOCAL_DOCKERFILE": "1",
|
||||
"BASE_IMAGE": "my/image:tag",
|
||||
"DOCKERFILE_PATH": "vllm/Dockerfile",
|
||||
"TORCH_WHEELS_PATH": "dist",
|
||||
"OUTPUT_DIR": "shared",
|
||||
},
|
||||
clear=True,
|
||||
)
|
||||
def test_params_success_normalizes_and_validates(
|
||||
self, mock_env_path, mock_is_path, mock_local_img
|
||||
):
|
||||
params = vllm.VllmBuildParameters()
|
||||
self.assertEqual(params.torch_whls_path, Path("/abs/dist"))
|
||||
self.assertEqual(params.dockerfile_path, Path("/abs/vllm/Dockerfile"))
|
||||
self.assertEqual(params.output_dir, Path("/abs/shared"))
|
||||
self.assertEqual(params.base_image, "my/image:tag")
|
||||
|
||||
@patch("cli.lib.core.vllm.is_path_exist", return_value=False)
|
||||
@patch.dict(
|
||||
os.environ, {"USE_TORCH_WHEEL": "1", "TORCH_WHEELS_PATH": "dist"}, clear=True
|
||||
)
|
||||
def test_params_missing_torch_whls_raises(self, _is_path):
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
os.chdir(td)
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
vllm.VllmBuildParameters(
|
||||
use_local_base_image=False,
|
||||
use_local_dockerfile=False,
|
||||
)
|
||||
err = cm.exception
|
||||
self.assertIn("TORCH_WHEELS_PATH", str(err))
|
||||
|
||||
@patch("cli.lib.core.vllm.local_image_exists", return_value=False)
|
||||
@patch.dict(
|
||||
os.environ, {"USE_LOCAL_BASE_IMAGE": "1", "BASE_IMAGE": "img:tag"}, clear=True
|
||||
)
|
||||
def test_params_missing_local_base_image_raises(self, _local_img):
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
os.chdir(td)
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
vllm.VllmBuildParameters(
|
||||
use_torch_whl=False,
|
||||
use_local_dockerfile=False,
|
||||
)
|
||||
err = cm.exception
|
||||
self.assertIn("BASE_IMAGE", str(err))
|
||||
|
||||
@patch("cli.lib.core.vllm.is_path_exist", return_value=False)
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"USE_LOCAL_DOCKERFILE": "1", "DOCKERFILE_PATH": "Dockerfile"},
|
||||
clear=True,
|
||||
)
|
||||
def test_params_missing_dockerfile_raises(self, _is_path):
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
os.chdir(td)
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
vllm.VllmBuildParameters(
|
||||
use_torch_whl=False,
|
||||
use_local_base_image=False,
|
||||
)
|
||||
err = cm.exception
|
||||
self.assertIn("DOCKERFILE_PATH", str(err))
|
||||
|
||||
@patch("cli.lib.core.vllm.is_path_exist", return_value=False)
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"OUTPUT_DIR": ""},
|
||||
clear=True,
|
||||
)
|
||||
def test_params_missing_output_dir(self, _is_path):
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
vllm.VllmBuildParameters()
|
||||
|
||||
|
||||
class TestBuildCmdAndRun(unittest.TestCase):
|
||||
@patch("cli.lib.core.vllm.local_image_exists", return_value=True)
|
||||
def test_generate_docker_build_cmd_includes_bits(self, _exists):
|
||||
runner = vllm.VllmBuildRunner()
|
||||
# Craft inputs that simulate a prepared build
|
||||
inputs = MagicMock()
|
||||
inputs.output_dir = Path("/abs/out")
|
||||
inputs.use_local_base_image = True
|
||||
inputs.base_image = "img:tag"
|
||||
inputs.torch_whls_path = Path("./vllm/tmp")
|
||||
inputs.max_jobs = 64
|
||||
inputs.cuda_version = "12.8.1"
|
||||
inputs.python_version = "3.12"
|
||||
inputs.sccache_bucket = "my-bucket"
|
||||
inputs.sccache_region = "us-west-2"
|
||||
inputs.torch_cuda_arch_list = "8.0;9.0"
|
||||
inputs.target_stage = "export-wheels"
|
||||
inputs.tag_name = "vllm-wheels"
|
||||
|
||||
cmd = runner._generate_docker_build_cmd(inputs)
|
||||
squashed = " ".join(cmd.split()) # normalize whitespace for matching
|
||||
|
||||
self.assertIn("--output type=local,dest=/abs/out", squashed)
|
||||
self.assertIn("-f docker/Dockerfile.nightly_torch", squashed)
|
||||
self.assertIn("--pull=false", squashed)
|
||||
self.assertIn("--build-arg TORCH_WHEELS_PATH=tmp", squashed)
|
||||
self.assertIn("--build-arg BUILD_BASE_IMAGE=img:tag", squashed)
|
||||
self.assertIn("--build-arg FINAL_BASE_IMAGE=img:tag", squashed)
|
||||
self.assertIn("--build-arg max_jobs=64", squashed)
|
||||
self.assertIn("--build-arg CUDA_VERSION=12.8.1", squashed)
|
||||
self.assertIn("--build-arg PYTHON_VERSION=3.12", squashed)
|
||||
self.assertIn("--build-arg USE_SCCACHE=1", squashed)
|
||||
self.assertIn("--build-arg SCCACHE_BUCKET_NAME=my-bucket", squashed)
|
||||
self.assertIn("--build-arg SCCACHE_REGION_NAME=us-west-2", squashed)
|
||||
self.assertIn("--build-arg torch_cuda_arch_list='8.0;9.0'", squashed)
|
||||
self.assertIn("--target export-wheels", squashed)
|
||||
self.assertIn("-t vllm-wheels", squashed)
|
||||
|
||||
@patch("cli.lib.core.vllm.run_command")
|
||||
@patch("cli.lib.core.vllm.ensure_dir_exists")
|
||||
@patch("cli.lib.core.vllm.clone_vllm")
|
||||
@patch.object(
|
||||
vllm.VllmBuildRunner,
|
||||
"_generate_docker_build_cmd",
|
||||
return_value="docker buildx ...",
|
||||
)
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
# Make __post_init__ validations pass cheaply
|
||||
"USE_TORCH_WHEEL": "0",
|
||||
"USE_LOCAL_BASE_IMAGE": "0",
|
||||
"USE_LOCAL_DOCKERFILE": "0",
|
||||
"OUTPUT_DIR": "shared",
|
||||
},
|
||||
clear=True,
|
||||
)
|
||||
def test_run_calls_clone_prepare_and_build(
|
||||
self, mock_gen, mock_clone, mock_ensure, mock_run
|
||||
):
|
||||
# Stub parameters instance so we avoid FS/Docker accesses in run()
|
||||
params = MagicMock()
|
||||
params.output_dir = Path("shared")
|
||||
params.use_local_dockerfile = False
|
||||
params.use_torch_whl = False
|
||||
|
||||
with patch("cli.lib.core.vllm.VllmBuildParameters", return_value=params):
|
||||
runner = vllm.VllmBuildRunner()
|
||||
runner.run()
|
||||
|
||||
mock_clone.assert_called_once()
|
||||
mock_ensure.assert_called_once_with(Path("shared"))
|
||||
mock_gen.assert_called_once_with(params)
|
||||
mock_run.assert_called_once()
|
||||
# ensure we run in vllm workdir
|
||||
_, kwargs = mock_run.call_args
|
||||
assert kwargs.get("cwd") == "vllm"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
414
.github/ci_configs/vllm/Dockerfile.tmp_vllm
vendored
Normal file
414
.github/ci_configs/vllm/Dockerfile.tmp_vllm
vendored
Normal file
@ -0,0 +1,414 @@
|
||||
# TODO(elainwy): remove this file after the torch nightly dockerfile is in sync in vllm repo
|
||||
# The vLLM Dockerfile is used to construct vLLM image against torch nightly and torch main that can be directly used for testing
|
||||
|
||||
ARG CUDA_VERSION=12.8.1
|
||||
ARG PYTHON_VERSION=3.12
|
||||
|
||||
# BUILD_BASE_IMAGE: used to setup python build xformers, and vllm wheels, It can be replaced with a different base image from local machine,
|
||||
# by default, it uses the torch-nightly-base stage from this docker image
|
||||
ARG BUILD_BASE_IMAGE=torch-nightly-base
|
||||
|
||||
# FINAL_BASE_IMAGE: used to set up vllm-instaled environment and build flashinfer,
|
||||
# by default, it uses devel-ubuntu22.04 official image.
|
||||
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
|
||||
|
||||
|
||||
#################### TORCH NIGHTLY BASE IMAGE ####################
|
||||
# A base image for building vLLM with devel ubuntu 22.04, this is mainly used to build vllm in vllm builtkite ci
|
||||
From nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 as torch-nightly-base
|
||||
ARG CUDA_VERSION=12.8.1
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG TARGETPLATFORM
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
|
||||
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
|
||||
|
||||
# Install Python and other dependencies if it does not existed
|
||||
RUN if ! command -v python3 >/dev/null || ! python3 --version | grep -q "${PYTHON_VERSION}"; then \
|
||||
echo "Installing Python ${PYTHON_VERSION}..." && \
|
||||
echo 'tzdata tzdata/Areas select America' | debconf-set-selections && \
|
||||
echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y ccache software-properties-common git curl sudo && \
|
||||
for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
done && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
|
||||
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && \
|
||||
ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config && \
|
||||
curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION}; \
|
||||
else \
|
||||
echo "Python ${PYTHON_VERSION} already present, skipping setup."; \
|
||||
fi \
|
||||
&& python3 --version && python3 -m pip --version
|
||||
|
||||
# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
|
||||
# as it was causing spam when compiling the CUTLASS kernels
|
||||
# Ensure gcc >= 10 to avoid CUTLASS issues (bug 92519)
|
||||
RUN current_gcc_version=$(gcc -dumpversion | cut -f1 -d.) && \
|
||||
if [ "$current_gcc_version" -lt 10 ]; then \
|
||||
echo "GCC version is $current_gcc_version, installing gcc-10..."; \
|
||||
apt-get update && \
|
||||
apt-get install -y gcc-10 g++-10 && \
|
||||
update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 100 && \
|
||||
update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-10 100; \
|
||||
else \
|
||||
echo "GCC version is $current_gcc_version, no need to install gcc-10."; \
|
||||
fi && \
|
||||
gcc --version && g++ --version
|
||||
|
||||
# install uv for faster pip installs
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
python3 -m pip install uv==0.8.4
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
|
||||
#################### TORCH NIGHTLY BASE IMAGE ####################
|
||||
|
||||
|
||||
#################### BASE BUILD IMAGE ####################
|
||||
# A base image for building vLLM with torch nightly or torch wheels
|
||||
# prepare basic build environment
|
||||
FROM ${BUILD_BASE_IMAGE} AS base
|
||||
USER root
|
||||
|
||||
# Workaround for https://github.com/openai/triton/issues/2507 and
|
||||
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
|
||||
# this won't be needed for future versions of this docker image
|
||||
# or future versions of triton.
|
||||
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
|
||||
|
||||
# Install uv for faster pip installs if not existed
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if ! python3 -m uv --version >/dev/null 2>&1; then \
|
||||
python3 -m pip install uv==0.8.4; \
|
||||
fi
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# install build and runtime dependencies
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
|
||||
# install build and runtime dependencies without stable torch version
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
# default mount file as placeholder, this just avoid the mount error
|
||||
# change to a different vllm folder if this does not exist anymore
|
||||
ARG TORCH_WHEELS_PATH="./requirements"
|
||||
ARG PINNED_TORCH_VERSION
|
||||
|
||||
# Install torch, torchaudio and torchvision based on the input
|
||||
# if TORCH_WHEELS_PATH is default "./requirements", it will pull thethe nightly versions using pip
|
||||
# otherwise, it will use the whls from TORCH_WHEELS_PATH from the host machine
|
||||
RUN --mount=type=bind,source=${TORCH_WHEELS_PATH},target=/dist \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
if [ -n "$TORCH_WHEELS_PATH" ] && [ "$TORCH_WHEELS_PATH" != "./requirements" ] && [ -d "/dist" ] && ls /dist/torch*.whl >/dev/null 2>&1; then \
|
||||
torch_whl=$(find /dist -maxdepth 1 -name 'torch-*.whl' -print -quit); \
|
||||
vision_whl=$(find /dist/vision -name 'torchvision*.whl' | head -n1 | xargs); \
|
||||
audio_whl=$(find /dist/audio -name 'torchaudio*.whl' | head -n1 | xargs); \
|
||||
uv pip install --system "${torch_whl}[opt-einsum]"; \
|
||||
uv pip install --system "${vision_whl}"; \
|
||||
uv pip install --system "${audio_whl}"; \
|
||||
elif [ -n "$PINNED_TORCH_VERSION" ]; then \
|
||||
echo "[INFO] Installing pinned torch nightly version: $PINNED_TORCH_VERSION"; \
|
||||
uv pip install --system "$PINNED_TORCH_VERSION" --index-url https://download.pytorch.org/whl/nightly/cu128; \
|
||||
else \
|
||||
echo "[INFO] Installing torch nightly with latest one"; \
|
||||
uv pip install --system torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128; \
|
||||
fi
|
||||
|
||||
# Install numba 0.61.2 for cuda environment
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system numba==0.61.2
|
||||
|
||||
# Install common dependencies from vllm common.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
|
||||
|
||||
# Must put before installing xformers, so it can install the correct version of xfomrers.
|
||||
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
ARG max_jobs=16
|
||||
ENV MAX_JOBS=${max_jobs}
|
||||
|
||||
# Build xformers with cuda and torch nightly/wheel
|
||||
# following official xformers guidance: https://github.com/facebookresearch/xformers#build
|
||||
ARG XFORMERS_COMMIT=f2de641ef670510cadab099ce6954031f52f191c
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
echo 'git clone xformers...' \
|
||||
&& git clone https://github.com/facebookresearch/xformers.git --recursive \
|
||||
&& cd xformers \
|
||||
&& git checkout ${XFORMERS_COMMIT} \
|
||||
&& git submodule update --init --recursive \
|
||||
&& echo 'finish git clone xformers...' \
|
||||
&& rm -rf build \
|
||||
&& python3 setup.py bdist_wheel --dist-dir=../xformers-dist --verbose \
|
||||
&& cd .. \
|
||||
&& rm -rf xformers
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system xformers-dist/*.whl --verbose
|
||||
|
||||
# Build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
|
||||
# track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
|
||||
RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
|
||||
RUN cat torch_build_versions.txt
|
||||
|
||||
RUN pip freeze | grep -E 'torch|xformers|torchvision|torchaudio'
|
||||
|
||||
#################### BASE BUILD IMAGE ####################
|
||||
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
# Image used to build vllm wheel
|
||||
FROM base AS build
|
||||
ARG TARGETPLATFORM
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/build.txt
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
# Max jobs used by Ninja to build extensions
|
||||
ARG max_jobs=16
|
||||
ENV MAX_JOBS=${max_jobs}
|
||||
ARG nvcc_threads=2
|
||||
ENV NVCC_THREADS=$nvcc_threads
|
||||
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
|
||||
ARG USE_SCCACHE
|
||||
ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
|
||||
ARG SCCACHE_REGION_NAME=us-west-2
|
||||
ARG SCCACHE_S3_NO_CREDENTIALS=0
|
||||
|
||||
# if USE_SCCACHE is set, use sccache to speed up compilation
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
echo "Installing sccache..." \
|
||||
&& curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
|
||||
&& tar -xzf sccache.tar.gz \
|
||||
&& sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \
|
||||
&& rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \
|
||||
&& export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
|
||||
&& export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
|
||||
&& export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
|
||||
&& export SCCACHE_IDLE_TIMEOUT=0 \
|
||||
&& export CMAKE_BUILD_TYPE=Release \
|
||||
&& sccache --show-stats \
|
||||
&& python3 setup.py bdist_wheel --dist-dir=vllm-dist --py-limited-api=cp38 \
|
||||
&& sccache --show-stats; \
|
||||
fi
|
||||
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
if [ "$USE_SCCACHE" != "1" ]; then \
|
||||
# Clean any existing CMake artifacts
|
||||
rm -rf .deps && \
|
||||
mkdir -p .deps && \
|
||||
python3 setup.py bdist_wheel --dist-dir=vllm-dist --py-limited-api=cp38; \
|
||||
fi
|
||||
|
||||
RUN echo "[DEBUG] Listing current directory:" && \
|
||||
ls -al && \
|
||||
echo "[DEBUG] Showing torch_build_versions.txt content:" && \
|
||||
cat torch_build_versions.txt
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
|
||||
|
||||
################### VLLM INSTALLED IMAGE ####################
|
||||
# Setup clean environment for vLLM for test and api server using ubuntu22.04 with AOT flashinfer
|
||||
FROM ${FINAL_BASE_IMAGE} AS vllm-base
|
||||
USER root
|
||||
# prepare for environment starts
|
||||
WORKDIR /workspace
|
||||
|
||||
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
|
||||
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
|
||||
|
||||
# Install Python and other dependencies if it does not existed
|
||||
RUN if ! command -v python3 >/dev/null || ! python3 --version | grep -q "${PYTHON_VERSION}"; then \
|
||||
echo "Installing Python ${PYTHON_VERSION}..." && \
|
||||
echo 'tzdata tzdata/Areas select America' | debconf-set-selections && \
|
||||
echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y ccache software-properties-common git curl sudo && \
|
||||
for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
done && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
|
||||
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
|
||||
update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && \
|
||||
ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config && \
|
||||
curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION}; \
|
||||
else \
|
||||
echo "Python ${PYTHON_VERSION} already present, skipping setup."; \
|
||||
fi \
|
||||
&& python3 --version && python3 -m pip --version
|
||||
|
||||
|
||||
# Get the torch versions, and whls used in previous stagtes for consistency
|
||||
COPY --from=base /workspace/torch_build_versions.txt ./torch_build_versions.txt
|
||||
COPY --from=base /workspace/xformers-dist /wheels/xformers
|
||||
COPY --from=build /workspace/vllm-dist /wheels/vllm
|
||||
RUN echo "[DEBUG] Listing current directory before torch install step:" && \
|
||||
ls -al && \
|
||||
echo "[DEBUG] Showing torch_build_versions.txt content:" && \
|
||||
cat torch_build_versions.txt
|
||||
|
||||
# Workaround for https://github.com/openai/triton/issues/2507 and
|
||||
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
|
||||
# this won't be needed for future versions of this docker image
|
||||
# or future versions of triton.
|
||||
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
|
||||
|
||||
|
||||
# Install uv for faster pip installs if not existed
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if ! python3 -m uv --version > /dev/null 2>&1; then \
|
||||
python3 -m pip install uv==0.8.4; \
|
||||
fi
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
|
||||
# Default mount file as placeholder, this just avoid the mount error
|
||||
ARG TORCH_WHEELS_PATH="./requirements"
|
||||
# Install torch, torchaudio and torchvision
|
||||
# if TORCH_WHEELS_PATH is default "./requirements", it will pull the nightly versions using pip using torch_build_versions.txt
|
||||
# otherwise, it will use the whls from TORCH_WHEELS_PATH from the host machine
|
||||
RUN --mount=type=bind,source=${TORCH_WHEELS_PATH},target=/dist \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
if [ -n "$TORCH_WHEELS_PATH" ] && [ "$TORCH_WHEELS_PATH" != "./requirements" ] && [ -d "/dist" ] && ls /dist/torch*.whl >/dev/null 2>&1; then \
|
||||
torch_whl=$(find /dist -maxdepth 1 -name 'torch-*.whl' -print -quit); \
|
||||
vision_whl=$(find /dist/vision -name 'torchvision*.whl' | head -n1 | xargs); \
|
||||
audio_whl=$(find /dist/audio -name 'torchaudio*.whl' | head -n1 | xargs); \
|
||||
echo "Found: '${torch_whl}' '${audio_whl}' '${vision_whl}'"; \
|
||||
uv pip install --system "${torch_whl}[opt-einsum]"; \
|
||||
uv pip install --system "${vision_whl}"; \
|
||||
uv pip install --system "${audio_whl}"; \
|
||||
else \
|
||||
echo "[INFO] Installing torch versions from torch_build_versions.txt"; \
|
||||
uv pip install --system $(cat torch_build_versions.txt | xargs) --index-url https://download.pytorch.org/whl/nightly/cu128; \
|
||||
fi
|
||||
|
||||
# Install the vllm wheel from previous stage
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system /wheels/vllm/*.whl --verbose
|
||||
|
||||
# Install xformers wheel from previous stage
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system /wheels/xformers/*.whl --verbose
|
||||
|
||||
|
||||
# Build flashinfer from source.
|
||||
ARG torch_cuda_arch_list='8.0;8.9;9.0a'
|
||||
# install package for build flashinfer
|
||||
# see issue: https://github.com/flashinfer-ai/flashinfer/issues/738
|
||||
|
||||
RUN pip install build==1.3.0
|
||||
RUN pip freeze | grep -E 'setuptools|packaging|build'
|
||||
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
# Build flashinfer for torch nightly from source around 10 mins
|
||||
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
|
||||
# Keep this in sync with https://github.com/vllm-project/vllm/blob/main/requirements/cuda.txt
|
||||
ARG FLASHINFER_GIT_REF="v0.2.9rc2"
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone --depth 1 --recursive --shallow-submodules \
|
||||
--branch ${FLASHINFER_GIT_REF} \
|
||||
${FLASHINFER_GIT_REPO} flashinfer \
|
||||
&& echo "Building FlashInfer with AOT for arches: ${torch_cuda_arch_list}" \
|
||||
&& cd flashinfer \
|
||||
&& python3 -m flashinfer.aot \
|
||||
&& python3 -m build --no-isolation --wheel --outdir ../wheels/flashinfer \
|
||||
&& cd .. \
|
||||
&& rm -rf flashinfer
|
||||
|
||||
# install flashinfer python
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system wheels/flashinfer/*.whl --verbose
|
||||
|
||||
# Logging to confirm the torch versions
|
||||
RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
|
||||
################### VLLM INSTALLED IMAGE ####################
|
||||
|
||||
|
||||
#################### UNITTEST IMAGE #############################
|
||||
FROM vllm-base as test
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
|
||||
COPY tests/ tests/
|
||||
COPY examples examples
|
||||
COPY benchmarks benchmarks
|
||||
COPY ./vllm/collect_env.py .
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
# Install build and runtime dependencies without stable torch version
|
||||
COPY requirements/nightly_torch_test.txt requirements/nightly_torch_test.txt
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
# install packages
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
# enable fast downloads from hf (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system hf_transfer
|
||||
ENV HF_HUB_ENABLE_HF_TRANSFER 1
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -e tests/vllm_test_utils
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/nightly_torch_test.txt
|
||||
|
||||
# Workaround for #17068
|
||||
# pinned commit for v2.2.4
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@95d8aba8a8c75aedcaa6143713b11e745e7cd0d9#egg=mamba-ssm"
|
||||
|
||||
# Logging to confirm the torch versions
|
||||
RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
|
||||
|
||||
# Logging to confirm all the packages are installed
|
||||
RUN pip freeze
|
||||
|
||||
#################### UNITTEST IMAGE #############################
|
||||
|
||||
#################### EXPORT STAGE ####################
|
||||
FROM scratch as export-wheels
|
||||
|
||||
# Just copy the wheels we prepared in previous stages
|
||||
COPY --from=base /workspace/xformers-dist /wheels/xformers
|
||||
COPY --from=build /workspace/vllm-dist /wheels/vllm
|
||||
COPY --from=vllm-base /workspace/wheels/flashinfer /wheels/flashinfer-python
|
||||
33
.github/workflows/tools-unit-tests.yml
vendored
33
.github/workflows/tools-unit-tests.yml
vendored
@ -17,7 +17,7 @@ concurrency:
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
torch-cli-unit-tests:
|
||||
lumen-cli-unit-tests-python312:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
@ -29,6 +29,11 @@ jobs:
|
||||
with:
|
||||
submodules: true
|
||||
fetch-depth: 0
|
||||
- name: Setup Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: '3.12'
|
||||
cache: pip
|
||||
|
||||
- name: Run tests
|
||||
continue-on-error: true
|
||||
@ -36,6 +41,30 @@ jobs:
|
||||
set -ex
|
||||
python3 -m venv /tmp/venv
|
||||
source /tmp/venv/bin/activate
|
||||
python -m pip install --upgrade pip
|
||||
pip install -e .ci/lumen_cli/
|
||||
pytest -v -s .ci/lumen_cli/tests/*
|
||||
|
||||
lumen-cli-compatible-python39:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
if: ${{ github.repository_owner == 'pytorch' }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout pytorch
|
||||
uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
|
||||
with:
|
||||
submodules: true
|
||||
fetch-depth: 0
|
||||
- name: Setup Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: '3.9'
|
||||
cache: 'pip'
|
||||
- name: Run tests
|
||||
continue-on-error: true
|
||||
run: |
|
||||
set -ex
|
||||
python3 -m venv /tmp/venv
|
||||
source /tmp/venv/bin/activate
|
||||
pip install -e .ci/lumen_cli/
|
||||
|
||||
Reference in New Issue
Block a user