venv.py

from __future__ import annotations

import functools, os, pathlib, shutil
import typer, yaml

from openllm.common import VENV_DIR, VERBOSE_LEVEL, BentoInfo, EnvVars, VenvSpec, output, run_command


@functools.lru_cache
def _resolve_bento_venv_spec(bento: BentoInfo, runtime_envs: EnvVars | None = None) -> VenvSpec:
    lock_file = bento.path / 'env' / 'python' / 'requirements.lock.txt'
    if not lock_file.exists():
        lock_file = bento.path / 'env' / 'python' / 'requirements.txt'

    reqs = lock_file.read_text().strip()
    bentofile = bento.path / 'bento.yaml'
    data = yaml.safe_load(bentofile.read_text())
    bento_env_list = data.get('envs', [])
    python_version = data.get('image', {})['python_version']
    bento_envs = {e['name']: e.get('value') for e in bento_env_list}
    envs = {k: runtime_envs.get(k, v) for k, v in bento_envs.items()} if runtime_envs else {}

    return VenvSpec(
        python_version=python_version,
        requirements_txt=reqs,
        name_prefix=f'{bento.tag.replace(":", "_")}-1-',
        envs=EnvVars(envs),
    )


def _ensure_venv(venv_spec: VenvSpec) -> pathlib.Path:
    venv = VENV_DIR / str(hash(venv_spec))
    if venv.exists() and not (venv / 'DONE').exists():
        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', 'uv', 'venv', venv.__fspath__(), '-p', venv_spec.python_version],
                silent=VERBOSE_LEVEL.get() < 10,
            )
            run_command(
                ['python', '-m', 'uv', 'pip', 'install', '-p', str(venv_py), 'bentoml'],
                silent=VERBOSE_LEVEL.get() < 10,
                env=venv_spec.envs,
            )
            with open(venv / 'requirements.txt', 'w') as f:
                f.write(venv_spec.normalized_requirements_txt)
            run_command(
                [
                    'python',
                    '-m',
                    'uv',
                    'pip',
                    'install',
                    '-p',
                    str(venv_py),
                    '-r',
                    (venv / 'requirements.txt').__fspath__(),
                ],
                silent=VERBOSE_LEVEL.get() < 10,
                env=venv_spec.envs,
            )
            with open(venv / 'DONE', 'w') as f:
                f.write('DONE')
        except Exception as e:
            shutil.rmtree(venv, ignore_errors=True)
            if VERBOSE_LEVEL.get() >= 10:
                output(str(e), style='red')
            output(f'Failed to install dependencies to {venv}. Cleaned up.', style='red')
            raise typer.Exit(1)
        output(f'Successfully installed dependencies to {venv}.', style='green')
        return venv
    else:
        return venv


def ensure_venv(bento: BentoInfo, runtime_envs: EnvVars | None = None) -> pathlib.Path:
    venv_spec = _resolve_bento_venv_spec(bento, runtime_envs=EnvVars(runtime_envs))
    venv = _ensure_venv(venv_spec)
    assert venv is not None
    return venv


def check_venv(bento: BentoInfo) -> bool:
    venv_spec = _resolve_bento_venv_spec(bento)
    venv = VENV_DIR / str(hash(venv_spec))
    if not venv.exists():
        return False
    if venv.exists() and not (venv / 'DONE').exists():
        return False
    return True