Skip to content

Commit 74ffa7f

Browse files
bojiangpre-commit-ci[bot]
bojiang
and
pre-commit-ci[bot]
authored
chore: use uv instead of venv layers (bentoml#1054)
* chmod * chore: discard venv layers * apply black * fix * ci: auto fixes from pre-commit.ci For more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent c6e4e69 commit 74ffa7f

5 files changed

Lines changed: 58 additions & 80 deletions

File tree

‎release.sh‎

100644100755
File mode changed.

‎src/openllm/accelerator_spec.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,6 @@ def __bool__(self):
6161
'nvidia-tesla-a100': {'model': 'A100', 'memory_size': 40.0},
6262
}
6363

64-
6564
ACCELERATOR_SPECS: dict[str, Accelerator] = {key: Accelerator(**value) for key, value in ACCELERATOR_SPEC_DICT.items()}
6665

6766

‎src/openllm/clean.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import pathlib
23
import shutil
34

@@ -8,7 +9,6 @@
89

910
app = OpenLLMTyper(help='clean up and release disk space used by OpenLLM')
1011

11-
1212
HUGGINGFACE_CACHE = pathlib.Path.home() / '.cache' / 'huggingface' / 'hub'
1313

1414

@@ -30,7 +30,22 @@ def model_cache(verbose: bool = False):
3030
def venvs(verbose: bool = False):
3131
if verbose:
3232
VERBOSE_LEVEL.set(20)
33-
used_space = sum(f.stat().st_size for f in VENV_DIR.rglob('*'))
33+
34+
# Set to store paths of files to avoid double counting
35+
seen_paths = set()
36+
used_space = 0
37+
38+
for f in VENV_DIR.rglob('*'):
39+
if os.name == 'nt': # Windows system
40+
# On Windows, directly add file sizes without considering hard links
41+
used_space += f.stat().st_size
42+
else:
43+
# On non-Windows systems, use inodes to avoid double counting
44+
stat = f.stat()
45+
if stat.st_ino not in seen_paths:
46+
seen_paths.add(stat.st_ino)
47+
used_space += stat.st_size
48+
3449
sure = questionary.confirm(
3550
f'This will remove all virtual environments created by OpenLLM (~{used_space / 1024 / 1024:.2f}MB), are you sure?'
3651
).ask()

‎src/openllm/common.py‎

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
ERROR_STYLE = 'red'
2222
SUCCESS_STYLE = 'green'
2323

24-
2524
OPENLLM_HOME = pathlib.Path(os.getenv('OPENLLM_HOME', pathlib.Path.home() / '.openllm'))
2625
REPO_DIR = OPENLLM_HOME / 'repos'
2726
TEMP_DIR = OPENLLM_HOME / 'temp'
@@ -231,14 +230,33 @@ def tolist(self):
231230

232231
class VenvSpec(SimpleNamespace):
233232
python_version: str
234-
python_packages: list[str]
235-
options: list[str] = []
233+
requirements_txt: str
236234
name_prefix = ''
237235

236+
@functools.cached_property
237+
def normalized_requirements_txt(self) -> str:
238+
parameter_lines: list[str] = []
239+
dependency_lines: list[str] = []
240+
comment_lines: list[str] = []
241+
242+
for line in self.requirements_txt.splitlines():
243+
if not line.strip():
244+
continue
245+
elif line.strip().startswith('#'):
246+
comment_lines.append(line.strip())
247+
elif line.strip().startswith('-'):
248+
parameter_lines.append(line.strip())
249+
else:
250+
dependency_lines.append(line.strip())
251+
252+
parameter_lines.sort()
253+
dependency_lines.sort()
254+
return '\n'.join(parameter_lines + dependency_lines).strip()
255+
238256
def __hash__(self):
239257
return md5(
240258
# self.python_version,
241-
*sorted(self.python_packages)
259+
self.normalized_requirements_txt
242260
)
243261

244262

@@ -307,11 +325,13 @@ def run_command(cmd, cwd=None, env=None, copy_env=True, venv=None, silent=False)
307325
try:
308326
if silent:
309327
return subprocess.run( # type: ignore
310-
cmd, cwd=cwd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
328+
cmd, cwd=cwd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
311329
)
312330
else:
313-
return subprocess.run(cmd, cwd=cwd, env=env)
314-
except subprocess.CalledProcessError:
331+
return subprocess.run(cmd, cwd=cwd, env=env, check=True)
332+
except Exception as e:
333+
if VERBOSE_LEVEL.get() >= 10:
334+
output(e, style='red')
315335
output('Command failed', style='red')
316336
raise typer.Exit(1)
317337

‎src/openllm/venv.py‎

Lines changed: 14 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -2,72 +2,28 @@
22
import os
33
import pathlib
44
import shutil
5-
import typing
6-
from typing import Iterable
75

86
import typer
97

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

1210

1311
@functools.lru_cache
14-
def _resolve_packages(requirement: typing.Union[pathlib.Path, str]):
15-
from pip_requirements_parser import RequirementsFile
16-
17-
requirements_txt = RequirementsFile.from_file(str(requirement), include_nested=True)
18-
return requirements_txt
19-
20-
21-
def _filter_preheat_packages(requirements: Iterable) -> list[str]:
22-
PREHEAT_PIP_PACKAGES = ['torch', 'vllm']
23-
24-
deps: list[str] = []
25-
for req in requirements:
26-
if req.is_editable or req.is_local_path or req.is_url or req.is_wheel or not req.name or not req.specifier:
27-
continue
28-
for sp in req.specifier:
29-
if sp.operator == '==' and req.name in PREHEAT_PIP_PACKAGES:
30-
assert req.line is not None
31-
deps.append(req.line)
32-
break
33-
return deps
34-
35-
36-
@functools.lru_cache
37-
def _resolve_bento_env_specs(bento: BentoInfo):
12+
def _resolve_bento_env_spec(bento: BentoInfo):
3813
ver_file = bento.path / 'env' / 'python' / 'version.txt'
3914
assert ver_file.exists(), f'cannot find version file in {bento.path}'
4015

4116
lock_file = bento.path / 'env' / 'python' / 'requirements.lock.txt'
4217
if not lock_file.exists():
4318
lock_file = bento.path / 'env' / 'python' / 'requirements.txt'
4419

45-
reqs = _resolve_packages(lock_file)
46-
packages = reqs.requirements
47-
options = reqs.options
48-
preheat_packages = _filter_preheat_packages(packages)
4920
ver = ver_file.read_text().strip()
50-
return (
51-
VenvSpec(
52-
python_version=ver, python_packages=preheat_packages, name_prefix=f"{bento.tag.replace(':', '_')}-1-"
53-
),
54-
VenvSpec(
55-
python_version=ver,
56-
python_packages=[v.line for v in packages],
57-
options=[o.line for o in options],
58-
name_prefix=f"{bento.tag.replace(':', '_')}-2-",
59-
),
60-
)
21+
reqs = lock_file.read_text().strip()
6122

23+
return VenvSpec(python_version=ver, requirements_txt=reqs, name_prefix=f"{bento.tag.replace(':', '_')}-1-")
6224

63-
def _get_lib_dir(venv: pathlib.Path) -> pathlib.Path:
64-
if os.name == 'nt':
65-
return venv / 'Lib/site-packages'
66-
else:
67-
return next(venv.glob('lib/python*')) / 'site-packages'
6825

69-
70-
def _ensure_venv(env_spec: VenvSpec, parrent_venv: typing.Optional[pathlib.Path] = None) -> pathlib.Path:
26+
def _ensure_venv(env_spec: VenvSpec) -> pathlib.Path:
7127
venv = VENV_DIR / str(hash(env_spec))
7228
if venv.exists() and not (venv / 'DONE').exists():
7329
shutil.rmtree(venv, ignore_errors=True)
@@ -77,22 +33,18 @@ def _ensure_venv(env_spec: VenvSpec, parrent_venv: typing.Optional[pathlib.Path]
7733
venv_py = venv / 'Scripts' / 'python.exe' if os.name == 'nt' else venv / 'bin' / 'python'
7834
try:
7935
run_command(['python', '-m', 'uv', 'venv', venv], silent=VERBOSE_LEVEL.get() < 10)
80-
lib_dir = _get_lib_dir(venv)
81-
if parrent_venv is not None:
82-
parent_lib_dir = _get_lib_dir(parrent_venv)
83-
with open(lib_dir / f'{parrent_venv.name}.pth', 'w+') as f:
84-
f.write(str(parent_lib_dir))
8536
with open(venv / 'requirements.txt', 'w') as f:
86-
with open(venv / 'requirements.txt', 'w') as f:
87-
f.write('\n'.join(env_spec.options + sorted(env_spec.python_packages)))
37+
f.write(env_spec.normalized_requirements_txt)
8838
run_command(
8939
['python', '-m', 'uv', 'pip', 'install', '-p', str(venv_py), '-r', venv / 'requirements.txt'],
9040
silent=VERBOSE_LEVEL.get() < 10,
9141
)
9242
with open(venv / 'DONE', 'w') as f:
9343
f.write('DONE')
94-
except Exception:
44+
except Exception as e:
9545
shutil.rmtree(venv, ignore_errors=True)
46+
if VERBOSE_LEVEL.get() >= 10:
47+
output(e, style='red')
9648
output(f'Failed to install dependencies to {venv}. Cleaned up.', style='red')
9749
raise typer.Exit(1)
9850
output(f'Successfully installed dependencies to {venv}.', style='green')
@@ -101,26 +53,18 @@ def _ensure_venv(env_spec: VenvSpec, parrent_venv: typing.Optional[pathlib.Path]
10153
return venv
10254

10355

104-
def _ensure_venvs(env_spec_list: Iterable[VenvSpec]) -> pathlib.Path:
105-
last_venv = None
106-
for env_spec in env_spec_list:
107-
last_venv = _ensure_venv(env_spec, last_venv)
108-
assert last_venv is not None
109-
return last_venv
110-
111-
11256
def ensure_venv(bento: BentoInfo) -> pathlib.Path:
113-
return _ensure_venvs(_resolve_bento_env_specs(bento))
57+
env_spec = _resolve_bento_env_spec(bento)
58+
venv = _ensure_venv(env_spec)
59+
assert venv is not None
60+
return venv
11461

11562

116-
def _check_venv(env_spec: VenvSpec) -> bool:
63+
def check_venv(bento: BentoInfo) -> bool:
64+
env_spec = _resolve_bento_env_spec(bento)
11765
venv = VENV_DIR / str(hash(env_spec))
11866
if not venv.exists():
11967
return False
12068
if venv.exists() and not (venv / 'DONE').exists():
12169
return False
12270
return True
123-
124-
125-
def check_venv(bento: BentoInfo) -> bool:
126-
return all(_check_venv(env_spec) for env_spec in _resolve_bento_env_specs(bento))

0 commit comments

Comments
 (0)