Spaces:
Running
Running
| # Copyright 2026 The ODML Authors. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """Virtual environment manager for LiteRT-LM.""" | |
| import os | |
| import shutil | |
| import subprocess | |
| import sys | |
| import click | |
| class VenvManager: | |
| """Manages the virtual environment paths and binaries.""" | |
| def __init__(self, prefer_current_venv: bool = False): | |
| self._self_managed_venv_dir = os.path.expanduser("~/.litert-lm/.venv") | |
| if not prefer_current_venv: | |
| self.venv_dir = self._self_managed_venv_dir | |
| else: | |
| self.venv_dir = os.environ.get( | |
| "VIRTUAL_ENV", | |
| sys.prefix | |
| if sys.prefix != sys.base_prefix | |
| else self._self_managed_venv_dir, | |
| ) | |
| self.python_bin = os.path.join(self.venv_dir, "bin", "python") | |
| self.pip_bin = os.path.join(self.venv_dir, "bin", "pip") | |
| self.litert_torch_bin = os.path.join(self.venv_dir, "bin", "litert-torch") | |
| self.uv_bin = os.path.join(self.venv_dir, "bin", "uv") | |
| def ensure_venv(self): | |
| """Ensures that the virtual environment directory exists.""" | |
| if os.path.exists(self.venv_dir): | |
| return | |
| if self.venv_dir != self._self_managed_venv_dir: | |
| # Note this should never happen. | |
| raise RuntimeError( | |
| f"Virtual environment directory not found: {self.venv_dir}" | |
| ) | |
| click.echo( | |
| click.style( | |
| f"Creating virtual environment in {self.venv_dir}...", fg="cyan" | |
| ) | |
| ) | |
| os.makedirs(os.path.dirname(self.venv_dir), exist_ok=True) | |
| python_exe = sys.executable or "python3" | |
| subprocess.run([python_exe, "-m", "venv", self.venv_dir], check=True) | |
| def recreate_venv_if_self_managed(self): | |
| """Deletes and re-creates the virtual environment if it is self-managed. | |
| This ensures we are using the latest litert-torch-nightly. Since uv has | |
| local cache, if the version has been downloaded before, it will be very | |
| fast. | |
| """ | |
| if self.venv_dir != self._self_managed_venv_dir: | |
| # Only recreate if it's the default venv managed by the CLI. | |
| return | |
| if os.path.exists(self.venv_dir): | |
| click.echo( | |
| click.style( | |
| f"Deleting virtual environment in {self.venv_dir}...", fg="cyan" | |
| ) | |
| ) | |
| shutil.rmtree(self.venv_dir) | |
| self.ensure_venv() | |
| def ensure_binary(self, binary_path): | |
| """Ensures the binary exists, or installs it if using the default venv.""" | |
| if os.path.exists(binary_path): | |
| return | |
| self.ensure_venv() | |
| if binary_path == self.pip_bin: | |
| click.echo(click.style("Ensuring pip is installed...", fg="cyan")) | |
| subprocess.run( | |
| [ | |
| self.python_bin, | |
| "-m", | |
| "ensurepip", | |
| "--default-pip", | |
| ], | |
| check=True, | |
| ) | |
| elif binary_path == self.uv_bin: | |
| self.ensure_binary(self.pip_bin) | |
| click.echo( | |
| click.style( | |
| "Installing uv into the virtual environment...", fg="cyan" | |
| ) | |
| ) | |
| subprocess.run( | |
| [ | |
| self.pip_bin, | |
| "install", | |
| "uv", | |
| "-i", | |
| "https://pypi.org/simple", | |
| ], | |
| check=True, | |
| ) | |
| elif binary_path == self.litert_torch_bin: | |
| self.ensure_binary(self.uv_bin) | |
| click.echo(click.style("Installing litert-torch with uv...", fg="cyan")) | |
| subprocess.run( | |
| [ | |
| self.uv_bin, | |
| "pip", | |
| "install", | |
| "-i", | |
| "https://pypi.org/simple", | |
| "litert-torch-nightly", | |
| "--python", | |
| self.python_bin, | |
| ], | |
| check=True, | |
| ) | |