diff --git a/openllm_next/common.py b/openllm_next/common.py index 8dd18c56..1d12a5aa 100644 --- a/openllm_next/common.py +++ b/openllm_next/common.py @@ -8,6 +8,7 @@ import pathlib import signal import subprocess import sys +import sysconfig import typing from contextlib import asynccontextmanager, contextmanager from types import SimpleNamespace @@ -277,6 +278,7 @@ def run_command( env = env or {} cmd = [str(c) for c in cmd] + bin_dir = "Scripts" if os.name == "nt" else "bin" if not silent: output("\n") if cwd: @@ -289,7 +291,7 @@ def run_command( output(f"$ {' '.join(cmd)}", style="bold") if venv: - py = venv / "bin" / "python" + py = venv / bin_dir / f"python{sysconfig.get_config_var('EXE')}" else: py = sys.executable diff --git a/openllm_next/venv.py b/openllm_next/venv.py index 336d3fce..a514357d 100644 --- a/openllm_next/venv.py +++ b/openllm_next/venv.py @@ -1,4 +1,5 @@ import functools +import os import pathlib import shutil import typing @@ -71,6 +72,13 @@ def _resolve_bento_env_specs(bento: BentoInfo): ) +def _get_lib_dir(venv: pathlib.Path) -> pathlib.Path: + if os.name == "nt": + return venv / "Lib/site-packages" + else: + return next(venv.glob("lib/python*")) / "site-packages" + + def _ensure_venv( env_spec: VenvSpec, parrent_venv: typing.Optional[pathlib.Path] = None, @@ -80,24 +88,33 @@ def _ensure_venv( shutil.rmtree(venv, ignore_errors=True) if not venv.exists(): output(f"Installing model dependencies({venv})...", style="green") + + venv_py = ( + venv / "Scripts" / "python.exe" + if os.name == "nt" + else venv / "bin" / "python" + ) try: run_command( - ["python", "-m", "venv", venv], + ["python", "-m", "uv", "venv", venv], silent=VERBOSE_LEVEL.get() < 10, ) - pyver = next(venv.glob("lib/python*")).name + lib_dir = _get_lib_dir(venv) if parrent_venv is not None: - with open( - venv / "lib" / pyver / "site-packages" / f"{parrent_venv.name}.pth", - "w+", - ) as f: - f.write(str(parrent_venv / "lib" / pyver / "site-packages")) + parent_lib_dir = _get_lib_dir(parrent_venv) + with open(lib_dir / f"{parrent_venv.name}.pth", "w+") as f: + f.write(str(parent_lib_dir)) with open(venv / "requirements.txt", "w") as f: f.write("\n".join(sorted(env_spec.python_packages.values()))) run_command( [ - venv / "bin" / "pip", + "python", + "-m", + "uv", + "pip", "install", + "-p", + str(venv_py), "-r", venv / "requirements.txt", "--upgrade-strategy", @@ -107,8 +124,13 @@ def _ensure_venv( ) run_command( [ - venv / "bin" / "pip", + "python", + "-m", + "uv", + "pip", "install", + "-p", + str(venv_py), "bentoml", "--upgrade-strategy", "only-if-needed", diff --git a/pyproject.toml b/pyproject.toml index 851f9219..341cd4dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "nvidia-ml-py", "dulwich", "tabulate", + "uv", ] [project.scripts]