Fix studio setup: venv isolation, centralized .venv_t5, uv targeting

- All platforms (including Colab) now create ~/.unsloth/studio/.venv
  with --without-pip fallback for broken ensurepip environments
- Add --python sys.executable to uv pip install in install_python_stack.py
  so uv targets the correct venv instead of system Python
- Centralize .venv_t5 bootstrap in transformers_version.py with proper
  validation (checks required packages exist, not just non-empty dir)
- Replace ~150 lines of duplicated install code across 3 worker files
  with calls to the shared _ensure_venv_t5_exists() helper
- Use uv-if-present with pip fallback; do not install uv at runtime
- Add site.addsitedir() shim in colab.py so notebook cells can import
  studio packages from the venv without system-Python double-install
- Update .venv_t5 packages: huggingface_hub 1.3.0->1.7.1, add hf_xet
- Bump transformers pin 4.57.1->4.57.6 in requirements + constraints
- Add Fast-Install helper to setup.ps1 with uv+pip fallback
- Keep Colab-specific completion banner in setup.sh
This commit is contained in:
Daniel Han 2026-03-18 07:23:41 +00:00
commit fc2d577d0d
10 changed files with 239 additions and 224 deletions

View file

@ -7,8 +7,26 @@ Uses Colab's built-in proxy - no external tunneling needed!
"""
from pathlib import Path
import site
import sys
def _bootstrap_studio_venv() -> None:
"""Expose the Studio venv's site-packages to the current interpreter.
On Colab, notebook cells run outside the venv subshell. Instead of
installing the full stack into system Python, we add the venv's
site-packages so that packages like structlog, fastapi, etc. are
importable from notebook cells.
"""
venv_lib = Path.home() / ".unsloth" / "studio" / ".venv" / "lib"
for sp in venv_lib.glob("python*/site-packages"):
if str(sp) not in sys.path:
site.addsitedir(str(sp))
_bootstrap_studio_venv()
# Add backend to path early so local modules like loggers can be imported
backend_path = str(Path(__file__).parent)
if backend_path not in sys.path:

View file

@ -40,59 +40,25 @@ def _activate_transformers_version(model_name: str) -> None:
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import needs_transformers_5, _resolve_base_model
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
else:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
r2 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
f"Failed to install transformers 5.x into {venv_t5}. "
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
)
sys.path.insert(0, venv_t5)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)

View file

@ -42,59 +42,25 @@ def _activate_transformers_version(model_name: str) -> None:
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import needs_transformers_5, _resolve_base_model
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
else:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
r2 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
f"Failed to install transformers 5.x into {venv_t5}. "
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
)
sys.path.insert(0, venv_t5)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)

View file

@ -36,59 +36,25 @@ def _activate_transformers_version(model_name: str) -> None:
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import needs_transformers_5, _resolve_base_model
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
else:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
r2 = sp.run(
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
f"Failed to install transformers 5.x into {venv_t5}. "
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
)
sys.path.insert(0, venv_t5)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)

View file

@ -11,4 +11,4 @@ git+https://github.com/meta-pytorch/OpenEnv.git
# executorch>=1.0.1 # 41.5 MB - no imports in unsloth/zoo/studio
torch-c-dlpack-ext
sentence_transformers==5.2.0
transformers==4.57.1
transformers==4.57.6

View file

@ -1,6 +1,6 @@
# Single-env pins for unsloth + studio + data-designer
# Keep compatible with unsloth transformers bounds.
transformers==4.57.1
transformers==4.57.6
trl==0.23.1
huggingface-hub==0.36.2

View file

@ -26,6 +26,7 @@ import json
import structlog
from loggers import get_logger
import os
import shutil
import subprocess
import sys
from pathlib import Path
@ -58,7 +59,7 @@ _tokenizer_class_cache: dict[str, bool] = {}
# Versions
TRANSFORMERS_5_VERSION = "5.3.0"
TRANSFORMERS_DEFAULT_VERSION = "4.57.1"
TRANSFORMERS_DEFAULT_VERSION = "4.57.6"
# Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1
_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5")
@ -216,29 +217,73 @@ def _purge_modules() -> int:
return len(to_remove)
_VENV_T5_PACKAGES = (
f"transformers=={TRANSFORMERS_5_VERSION}",
"huggingface_hub==1.7.1",
"hf_xet==1.4.2",
)
def _venv_t5_is_valid() -> bool:
"""Return True if .venv_t5/ has all required packages installed."""
if not os.path.isdir(_VENV_T5_DIR) or not os.listdir(_VENV_T5_DIR):
return False
# Check that the key package directories actually exist
for pkg_spec in _VENV_T5_PACKAGES:
pkg_name = pkg_spec.split("==")[0].replace("-", "_")
if not any(
(Path(_VENV_T5_DIR) / d).is_dir()
for d in (pkg_name, pkg_name.replace("_", "-"))
):
return False
return True
def _install_to_venv_t5(pkg: str) -> bool:
"""Install a single package into .venv_t5/, preferring uv then pip."""
# Try uv first (faster) if already on PATH -- do NOT install uv at runtime
if shutil.which("uv"):
result = subprocess.run(
[
"uv", "pip", "install",
"--python", sys.executable,
"--target", _VENV_T5_DIR,
"--no-deps", pkg,
],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
)
if result.returncode == 0:
return True
logger.warning("uv install of %s failed, falling back to pip", pkg)
# Fallback to pip
result = subprocess.run(
[
sys.executable, "-m", "pip", "install",
"--target", _VENV_T5_DIR,
"--no-deps", pkg,
],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
)
if result.returncode != 0:
logger.error("install failed:\n%s", result.stdout)
return False
return True
def _ensure_venv_t5_exists() -> bool:
"""Ensure .venv_t5/ exists. Install at runtime if missing."""
if os.path.isdir(_VENV_T5_DIR) and os.listdir(_VENV_T5_DIR):
"""Ensure .venv_t5/ exists with all required packages. Install if missing."""
if _venv_t5_is_valid():
return True
logger.warning(".venv_t5 not found at %s — installing at runtime", _VENV_T5_DIR)
logger.warning(".venv_t5 not found or incomplete at %s -- installing at runtime", _VENV_T5_DIR)
os.makedirs(_VENV_T5_DIR, exist_ok = True)
for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"):
cmd = [
sys.executable,
"-m",
"pip",
"install",
"--target",
_VENV_T5_DIR,
"--no-deps",
pkg,
]
result = subprocess.run(
cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True
)
if result.returncode != 0:
logger.error("pip install failed:\n%s", result.stdout)
for pkg in _VENV_T5_PACKAGES:
if not _install_to_venv_t5(pkg):
return False
logger.info("Installed transformers 5.x to %s", _VENV_T5_DIR)
return True