"""Installation and environment-setup helpers for the MedVision benchmark.
This module bundles the routines used to provision a runtime for the benchmark:
installing the vendored ``lmms-eval`` package and the ``medvision_ds`` dataset
codebase, and configuring the environment variables that point Hugging Face and
the MedVision dataset loader at a chosen data directory. It also provides
convenience installers for CUDA, PyTorch, Flash-Attention, and vLLM used when
setting up per-model conda environments.
"""
import os
import shlex
import subprocess
import sys
from importlib.resources import files # Python 3.9+
from pathlib import Path
from huggingface_hub import snapshot_download
[docs]
def run_pip_install(requirements_path):
# Normalize input to a Path so both str and Path work.
req_path = Path(requirements_path).expanduser().resolve(strict=False)
if not req_path.exists() or not req_path.is_file():
raise FileNotFoundError(f"Requirements file not found: {requirements_path}")
# Use the current interpreter to run pip to avoid PATH/env mismatches.
cmd = [
sys.executable,
"-m",
"pip",
"install",
"--upgrade",
"--force-reinstall",
"--no-deps",
"-r",
str(req_path),
]
env = os.environ.copy()
env.setdefault("PIP_DISABLE_PIP_VERSION_CHECK", "1")
print(f"Installing packages from: {req_path}")
subprocess.run(cmd, env=env, check=True)
[docs]
def ensure_hf_hub_installed(hf_hub_version="0.35.3"):
try:
from huggingface_hub import snapshot_download # noqa: F401
except ImportError:
subprocess.run(
f"pip install huggingface_hub[cli]=={hf_hub_version}",
check=True,
shell=True,
)
def _install_lmms_eval(
lmms_eval_dir,
editable_install=False,
proj_dependency=None,
):
# compose extras text like .[extra]
extras_txt = f"[{proj_dependency}]" if proj_dependency else ""
tmp_build_lock_file = os.path.join(lmms_eval_dir, ".build.lock")
build_dir = os.path.join(lmms_eval_dir, "build")
dist_dir = os.path.join(lmms_eval_dir, "dist")
egg_info_dir = os.path.join(lmms_eval_dir, "lmms_eval.egg-info")
wheel_dir = os.path.join(lmms_eval_dir, "wheels")
os.makedirs(wheel_dir, exist_ok=True)
# Common pip flags
base_pip_flags = "--no-cache-dir --force-reinstall"
# Case A: Editable install (always from source; extras allowed)
if editable_install:
# Example: pip install -e .[qwen2_5_vl]
cmd = (
f"flock -w 600 {shlex.quote(tmp_build_lock_file)} bash -lc '"
f"python -m pip install {base_pip_flags} -e .{extras_txt}"
f"'"
)
subprocess.run(cmd, check=True, shell=True, cwd=lmms_eval_dir)
return
# Case B: Non-editable, NO extras → build wheel once, then install wheel
if proj_dependency is None:
cmd = (
f"flock -w 600 {shlex.quote(tmp_build_lock_file)} bash -lc '"
f"rm -rf {shlex.quote(build_dir)} {shlex.quote(dist_dir)} {shlex.quote(egg_info_dir)} && "
f"rm -f {shlex.quote(wheel_dir)}/*.whl && "
f"python -m pip install --upgrade build && "
f"python -m build --wheel --outdir {shlex.quote(wheel_dir)} {shlex.quote(lmms_eval_dir)} && "
f"latest_wheel=$(ls -t {shlex.quote(wheel_dir)}/lmms_eval-*.whl | head -n1) && "
f'pip install {base_pip_flags} "$latest_wheel"'
f"'"
)
subprocess.run(cmd, check=True, shell=True, cwd=lmms_eval_dir)
return
# Case C: Non-editable WITH extras → install from source with extras
# (extras on a wheel path is not supported)
cmd = (
f"flock -w 600 {shlex.quote(tmp_build_lock_file)} bash -lc '"
f"python -m pip install {base_pip_flags} .{extras_txt}"
f"'"
)
subprocess.run(cmd, check=True, shell=True, cwd=lmms_eval_dir)
def install_lmms_eval(
benchmark_dir,
lmms_eval_folder,
editable_install=False,
proj_dependency=None,
):
lmms_eval_dir = os.path.join(benchmark_dir, lmms_eval_folder)
_install_lmms_eval(
lmms_eval_dir=lmms_eval_dir,
editable_install=editable_install,
proj_dependency=proj_dependency,
)
[docs]
def install_vendored_lmms_eval(
editable_install=True,
proj_dependency=None,
):
"""Install the vendored ``lmms-eval`` package that ships inside ``medvision_bm``.
Locates the ``medvision_lmms_eval`` package shipped as package data and
installs it via the shared installer. Editable installation is used by
default because the task definition files are otherwise not discovered.
Args:
editable_install (bool): Install in editable (``pip install -e``) mode.
Defaults to ``True``.
proj_dependency (str, optional): Name of an optional-dependency extra to
install (rendered as ``.[extra]``). Defaults to ``None``.
"""
# Locate the vendored lmms-eval package, check [tool.setuptools.package-data] in pyproject.toml
lmms_eval_dir = str(files("medvision_bm").joinpath("medvision_lmms_eval"))
# NOTE: Must install the vendored lmms-eval in editable mode, otherwise tasks files won't be found.
# TODO: Check: Why editable install causes issues in some cases?
_install_lmms_eval(
lmms_eval_dir=lmms_eval_dir,
editable_install=editable_install,
proj_dependency=proj_dependency,
)
[docs]
def setup_env_hf(data_dir):
"""Point Hugging Face's cache and dataset cache at ``data_dir``.
Sets ``HF_HOME`` and ``HF_DATASETS_CACHE`` to subdirectories of ``data_dir``
so that models and datasets are cached under the benchmark's data directory.
Args:
data_dir (str): Data directory; a relative path is resolved to an
absolute path.
"""
# Safeguard data_dir: you can use relative path with this function
data_dir = os.path.abspath(data_dir)
# Set Hugging Face dataset and cache directories
os.environ["HF_DATASETS_CACHE"] = os.path.join(
data_dir, ".cache", "huggingface", "datasets"
)
os.environ["HF_HOME"] = os.path.join(data_dir, ".cache", "huggingface")
[docs]
def setup_env_medvision_ds(
data_dir,
force_install_code=True,
force_download_data=False,
):
"""Set the environment variables that configure the ``medvision_ds`` loader.
Sets ``MedVision_DATA_DIR`` to ``data_dir`` (creating the directory if
needed) and, when requested, the force-install and force-download flags read
by the dataset codebase. The flags are only set when their argument is
``True``; they are left unchanged otherwise.
Args:
data_dir (str): Data directory; a relative path is resolved to an
absolute path.
force_install_code (bool): When ``True``, set
``MedVision_FORCE_INSTALL_CODE=true`` to force reinstallation of the
dataset codebase. Defaults to ``True``.
force_download_data (bool): When ``True``, set
``MedVision_FORCE_DOWNLOAD_DATA=true`` to force re-download of the
data. Defaults to ``False``.
"""
# Safeguard data_dir: you can use relative path with this function
data_dir = os.path.abspath(data_dir)
# Set dataset directory
os.makedirs(data_dir, exist_ok=True)
os.environ["MedVision_DATA_DIR"] = data_dir
# Force install dataset codebase, default to "False"
if force_install_code:
os.environ["MedVision_FORCE_INSTALL_CODE"] = "true"
# Force download dataset, default to "False"
if force_download_data:
os.environ["MedVision_FORCE_DOWNLOAD_DATA"] = "true"
[docs]
def setup_env_hf_medvision_ds(
data_dir,
force_install_code=True,
force_download_data=False,
):
"""Configure both the ``medvision_ds`` and Hugging Face environment variables.
Calls ``setup_env_medvision_ds`` followed by ``setup_env_hf`` for the same
data directory.
Args:
data_dir (str): Data directory; a relative path is resolved to an
absolute path.
force_install_code (bool): Forwarded to ``setup_env_medvision_ds``.
Defaults to ``True``.
force_download_data (bool): Forwarded to ``setup_env_medvision_ds``.
Defaults to ``False``.
"""
# Set environment variables for medvision_ds
setup_env_medvision_ds(
data_dir=data_dir,
force_install_code=force_install_code,
force_download_data=force_download_data,
)
# Set environment variables for Hugging Face
setup_env_hf(data_dir)
[docs]
def install_medvision_ds(
data_dir,
local_dir=None,
):
"""Build and install the ``medvision_ds`` dataset codebase from source.
When ``local_dir`` is ``None``, downloads the dataset's ``src/`` directory
from the ``YongchengYAO/MedVision`` Hugging Face dataset repo into
``data_dir``; otherwise uses ``src/`` under ``local_dir``. A wheel is built
from that source (guarded by a ``flock`` build lock, with a no-lock fallback)
and installed with ``--force-reinstall``. Finally, the Hugging Face and
``medvision_ds`` environment variables are configured for ``data_dir``.
Args:
data_dir (str): Data directory used for the snapshot download and for
environment setup; a relative path is resolved to an absolute path.
local_dir (str, optional): Directory containing an existing ``src/`` tree
to install from. When provided, no download is performed. Defaults to
``None``.
"""
if local_dir is None:
# Safeguard data_dir: you can use relative path with this function
data_dir = os.path.abspath(data_dir)
os.makedirs(data_dir, exist_ok=True)
snapshot_download(
repo_id="YongchengYAO/MedVision",
allow_patterns="src/*",
repo_type="dataset",
local_dir=data_dir,
)
dir_bmvqa = os.path.abspath(os.path.join(data_dir, "src"))
else:
dir_bmvqa = os.path.abspath(os.path.join(local_dir, "src"))
tmp_build_lock_file = os.path.join(dir_bmvqa, ".build.lock")
build_dir = os.path.join(dir_bmvqa, "build")
dist_dir = os.path.join(dir_bmvqa, "dist")
egg_info_dir = os.path.join(dir_bmvqa, "medvision_ds.egg-info")
wheel_dir = os.path.join(dir_bmvqa, "wheels")
os.makedirs(wheel_dir, exist_ok=True)
cmd_w_flock = (
f"flock -w 600 {shlex.quote(tmp_build_lock_file)} bash -lc '"
f"rm -rf {shlex.quote(build_dir)} {shlex.quote(dist_dir)} {shlex.quote(egg_info_dir)} && "
f"rm -f {shlex.quote(wheel_dir)}/*.whl && "
f"python -m pip install --upgrade build && "
f"python -m build --wheel --outdir {shlex.quote(wheel_dir)} {shlex.quote(dir_bmvqa)} && "
f"latest_wheel=$(ls -t {shlex.quote(wheel_dir)}/medvision_ds-*.whl | head -n1) && "
f'pip install --no-cache-dir --force-reinstall "$latest_wheel"\''
)
# Try with flock, fallback to without flock if it fails
try:
subprocess.run(cmd_w_flock, check=True, shell=True)
except subprocess.CalledProcessError:
print("Warning: flock failed, attempting installation without file lock...")
cmd_no_flock = (
f"bash -lc '"
f"rm -rf {shlex.quote(build_dir)} {shlex.quote(dist_dir)} {shlex.quote(egg_info_dir)} && "
f"rm -f {shlex.quote(wheel_dir)}/*.whl && "
f"python -m pip install --upgrade build && "
f"python -m build --wheel --outdir {shlex.quote(wheel_dir)} {shlex.quote(dir_bmvqa)} && "
f"latest_wheel=$(ls -t {shlex.quote(wheel_dir)}/medvision_ds-*.whl | head -n1) && "
f'pip install --no-cache-dir --force-reinstall "$latest_wheel"\''
)
subprocess.run(cmd_no_flock, check=True, shell=True)
# Set environment variables for medvision_ds
setup_env_hf_medvision_ds(data_dir=data_dir)
def pip_install_medvision_ds():
try:
print(
'\n[Info] Installing medvision_ds from Hugging Face Datasets repo: pip install "git+https://huggingface.co/datasets/YongchengYAO/MedVision.git#subdirectory=src"'
)
subprocess.run(
'pip install "git+https://huggingface.co/datasets/YongchengYAO/MedVision.git#subdirectory=src"',
check=True,
shell=True,
)
print("Successfully installed medvision_ds.")
except subprocess.CalledProcessError as e:
print(f"Error installing medvision_ds: {e}", file=sys.stderr)
def pip_install_medvision_bm():
try:
print(
'\n[Info] Installing medvision_bm from GitHub repo: pip install "git+https://github.com/YongchengYAO/MedVision.git"'
)
subprocess.run(
'pip install "git+https://github.com/YongchengYAO/MedVision.git"',
check=True,
shell=True,
)
print("Successfully installed medvision_bm.")
except subprocess.CalledProcessError as e:
print(f"Error installing medvision_bm: {e}", file=sys.stderr)
[docs]
def setup_env_cuda():
print("Setting up CUDA environment...")
cuda_home = os.environ.get("CONDA_PREFIX", "")
os.environ["CUDA_HOME"] = cuda_home
os.environ["PATH"] = f"{cuda_home}/bin:{os.environ.get('PATH', '')}"
os.environ["LD_LIBRARY_PATH"] = (
f"{cuda_home}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
)
os.environ["LD_LIBRARY_PATH"] = (
f"{cuda_home}/lib:{os.environ.get('LD_LIBRARY_PATH', '')}"
)
def install_torch_cu121():
"""Install PyTorch with CUDA support."""
print("Installing PyTorch...")
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"torch==2.5.0+cu121",
"torchvision==0.20.0+cu121",
"torchaudio==2.5.0+cu121",
"--index-url",
"https://download.pytorch.org/whl/cu121",
"--force-reinstall",
],
check=True,
)
setup_env_cuda()
[docs]
def install_torch_cu124():
"""Install PyTorch with CUDA support."""
print("Installing PyTorch...")
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"torch==2.6.0+cu124",
"torchvision==0.21.0+cu124",
"torchaudio==2.6.0+cu124",
"--index-url",
"https://download.pytorch.org/whl/cu124",
"--force-reinstall",
],
check=True,
)
setup_env_cuda()
def install_flash_attention_torch_and_deps_py39():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
subprocess.run(
"pip install torch==2.6.0+cu124 torchvision==0.21.0+cu124 torchaudio==2.6.0+cu124 "
"--index-url https://download.pytorch.org/whl/cu124 --force-reinstall",
check=True,
shell=True,
)
# Install CUDA
print("Installing CUDA toolkit and components...")
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-toolkit -y",
check=True,
shell=True,
)
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-nvcc -y", check=True, shell=True
)
subprocess.run(
"conda install cudnn -y",
check=True,
shell=True,
)
subprocess.run(
"pip install --upgrade nvidia-cuda-cupti-cu12==12.4.* "
"nvidia-cuda-nvrtc-cu12==12.4.* "
"nvidia-cuda-runtime-cu12==12.4.*",
check=True,
shell=True,
)
setup_env_cuda()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
[docs]
def install_flash_attention_torch_and_deps_py39_v2():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
install_torch_cu124()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
def install_flash_attention_torch_and_deps_py310():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
subprocess.run(
"pip install torch==2.6.0+cu124 torchvision==0.21.0+cu124 torchaudio==2.6.0+cu124 "
"--index-url https://download.pytorch.org/whl/cu124 --force-reinstall",
check=True,
shell=True,
)
# Install CUDA
print("Installing CUDA toolkit and components...")
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-toolkit -y",
check=True,
shell=True,
)
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-nvcc -y", check=True, shell=True
)
subprocess.run(
"conda install cudnn -y",
check=True,
shell=True,
)
subprocess.run(
"pip install --upgrade nvidia-cuda-cupti-cu12==12.4.* "
"nvidia-cuda-nvrtc-cu12==12.4.* "
"nvidia-cuda-runtime-cu12==12.4.*",
check=True,
shell=True,
)
setup_env_cuda()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
[docs]
def install_flash_attention_torch_and_deps_py310_v2():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
install_torch_cu124()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
def install_flash_attention_torch_and_deps_py311():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
subprocess.run(
"pip install torch==2.6.0+cu124 torchvision==0.21.0+cu124 torchaudio==2.6.0+cu124 "
"--index-url https://download.pytorch.org/whl/cu124 --force-reinstall",
check=True,
shell=True,
)
# Install CUDA
print("Installing CUDA toolkit and components...")
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-toolkit -y",
check=True,
shell=True,
)
subprocess.run(
"conda install nvidia/label/cuda-12.4.0::cuda-nvcc -y", check=True, shell=True
)
subprocess.run(
"conda install cudnn -y",
check=True,
shell=True,
)
subprocess.run(
"pip install --upgrade nvidia-cuda-cupti-cu12==12.4.* "
"nvidia-cuda-nvrtc-cu12==12.4.* "
"nvidia-cuda-runtime-cu12==12.4.*",
check=True,
shell=True,
)
setup_env_cuda()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
[docs]
def install_flash_attention_torch_and_deps_py311_v2():
# Install PyTorch with CUDA support
print("Installing PyTorch with CUDA 12.4...")
install_torch_cu124()
# Install Flash Attention
print("Installing Flash Attention...")
subprocess.run(
"pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl",
check=True,
shell=True,
)
# Install numpy version 1.26.4
print("Installing numpy...")
subprocess.run("pip install numpy==1.26.4", check=True, shell=True)
# Install protobuf version 3.20
print("Installing protobuf 3.20.x")
subprocess.run("pip install protobuf==3.20", check=True, shell=True)
[docs]
def setup_env_vllm(data_dir):
# Safeguard data_dir: you can use relative path with this function
data_dir = os.path.abspath(data_dir)
# Ensure proper process spawning
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
# Set the cache directory for vllm
os.environ["XDG_CACHE_HOME"] = os.path.join(data_dir, ".cache", "vllm")
[docs]
def install_vllm(data_dir, version="0.10.0"):
# Install and setup vllm
try:
subprocess.run("pip install blobfile", check=True, shell=True)
subprocess.run(
f"pip install vllm=={version}",
check=True,
shell=True,
)
print("Successfully installed vllm")
except Exception as e:
raise RuntimeError(f"Error installing vllm: {e}")
setup_env_vllm(data_dir)
