add Docker support: skip venv, install only missing deps

This commit is contained in:
Roland Tannous 2026-03-27 19:39:37 +00:00
parent 3869fbe1cc
commit 526c96980b
3 changed files with 229 additions and 253 deletions

View file

@ -841,216 +841,153 @@ def install_python_stack() -> int:
base_total += 3
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try to use uv for faster installs (must happen before pip upgrade
# because uv venvs don't include pip by default)
USE_UV = _bootstrap_uv()
# # 1. Try to use uv for faster installs (must happen before pip upgrade
# # because uv venvs don't include pip by default)
# USE_UV = _bootstrap_uv()
# 2. Ensure pip is available (uv venvs created by install.sh don't include pip)
_progress("pip bootstrap")
if USE_UV:
run(
"Bootstrapping pip via uv",
[
"uv",
"pip",
"install",
"--python",
sys.executable,
"pip",
],
)
else:
# pip may not exist yet (uv-created venvs omit it). Try ensurepip
# first, then upgrade. Only fall back to a direct upgrade when pip
# is already present.
_has_pip = (
subprocess.run(
[sys.executable, "-m", "pip", "--version"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
).returncode
== 0
)
# # 2. Ensure pip is available (uv venvs created by install.sh don't include pip)
# _progress("pip bootstrap")
# if USE_UV:
# run(
# "Bootstrapping pip via uv",
# [
# "uv",
# "pip",
# "install",
# "--python",
# sys.executable,
# "pip",
# ],
# )
# else:
# # pip may not exist yet (uv-created venvs omit it). Try ensurepip
# # first, then upgrade. Only fall back to a direct upgrade when pip
# # is already present.
# _has_pip = (
# subprocess.run(
# [sys.executable, "-m", "pip", "--version"],
# stdout = subprocess.DEVNULL,
# stderr = subprocess.DEVNULL,
# ).returncode
# == 0
# )
#
# if not _has_pip:
# run(
# "Bootstrapping pip via ensurepip",
# [sys.executable, "-m", "ensurepip", "--upgrade"],
# )
# else:
# run(
# "Upgrading pip",
# [sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
# )
if not _has_pip:
run(
"Bootstrapping pip via ensurepip",
[sys.executable, "-m", "ensurepip", "--upgrade"],
)
else:
run(
"Upgrading pip",
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# # 3. Core packages: unsloth-zoo + unsloth (or custom package name)
# if skip_base:
# print(_green(f"✅ {package_name} already installed — skipping base packages"))
# elif NO_TORCH:
# # No-torch update path: install unsloth + unsloth-zoo with --no-deps
# # (current PyPI metadata still declares torch as a hard dep), then
# # runtime deps with --no-deps (avoids transitive torch).
# _progress("base packages (no torch)")
# pip_install(
# f"Updating {package_name} + unsloth-zoo (no-torch mode)",
# "--no-cache-dir",
# "--no-deps",
# "--upgrade-package",
# package_name,
# "--upgrade-package",
# "unsloth-zoo",
# package_name,
# "unsloth-zoo",
# )
# pip_install(
# "Installing no-torch runtime deps",
# "--no-cache-dir",
# "--no-deps",
# req = REQ_ROOT / "no-torch-runtime.txt",
# )
# if local_repo:
# pip_install(
# "Overlaying local repo (editable)",
# "--no-cache-dir",
# "--no-deps",
# "-e",
# local_repo,
# constrain = False,
# )
# elif local_repo:
# _progress("base packages")
# pip_install(
# "Updating base packages",
# "--no-cache-dir",
# "--upgrade-package",
# "unsloth",
# "--upgrade-package",
# "unsloth-zoo",
# req = REQ_ROOT / "base.txt",
# )
# pip_install(
# "Overlaying local repo (editable)",
# "--no-cache-dir",
# "--no-deps",
# "-e",
# local_repo,
# constrain = False,
# )
# elif package_name != "unsloth":
# _progress("base packages")
# pip_install(
# f"Installing {package_name}",
# "--no-cache-dir",
# package_name,
# )
# else:
# _progress("base packages")
# pip_install(
# "Updating base packages",
# "--no-cache-dir",
# "--upgrade-package",
# "unsloth",
# "--upgrade-package",
# "unsloth-zoo",
# req = REQ_ROOT / "base.txt",
# )
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
pass
elif NO_TORCH:
# No-torch update path: install unsloth + unsloth-zoo with --no-deps
# (current PyPI metadata still declares torch as a hard dep), then
# runtime deps with --no-deps (avoids transitive torch).
_progress("base packages (no torch)")
pip_install(
f"Updating {package_name} + unsloth-zoo (no-torch mode)",
"--no-cache-dir",
"--no-deps",
"--upgrade-package",
package_name,
"--upgrade-package",
"unsloth-zoo",
package_name,
"unsloth-zoo",
)
pip_install(
"Installing no-torch runtime deps",
"--no-cache-dir",
"--no-deps",
req = REQ_ROOT / "no-torch-runtime.txt",
)
if local_repo:
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
elif local_repo:
# Local dev install: update deps from base.txt, then overlay the
# local checkout as an editable install (--no-deps so torch is
# never re-resolved).
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
elif package_name != "unsloth":
# Custom package name (e.g. roland-sloth for testing) — install directly
_progress("base packages")
pip_install(
f"Installing {package_name}",
"--no-cache-dir",
package_name,
)
else:
# Update path: upgrade only unsloth + unsloth-zoo while preserving
# existing torch/CUDA installations. Torch is pre-installed by
# install.sh / setup.ps1; --upgrade-package targets only base pkgs.
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
# pip_install(
# "Installing additional unsloth dependencies",
# "--no-cache-dir",
# req = REQ_ROOT / "extras.txt",
# )
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv received CPU-only torch (common when pip resolves torch from PyPI).
# Must come immediately after base packages so torch is present for inspection.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_ensure_rocm_torch()
# pip_install(
# "Installing extras (no-deps)",
# "--no-deps",
# "--no-cache-dir",
# req = REQ_ROOT / "extras-no-deps.txt",
# )
# Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows.
# Detect and warn so users know manual steps are needed for GPU training.
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
# Validate actual AMD GPU presence (not just tool existence)
import re as _re_win
# # 4. Overrides (torchao, transformers) -- force-reinstall
# _progress("dependency overrides")
# pip_install(
# "Installing dependency overrides",
# "--force-reinstall",
# "--no-cache-dir",
# req = REQ_ROOT / "overrides.txt",
# )
def _win_amd_smi_has_gpu(stdout: str) -> bool:
return bool(_re_win.search(r"(?im)^gpu\s*[:\[]\s*\d", stdout))
_win_amd_gpu = False
for _wcmd, _check_fn in (
(["hipinfo"], lambda out: "gcnarchname" in out.lower()),
(["amd-smi", "list"], _win_amd_smi_has_gpu),
):
_wexe = shutil.which(_wcmd[0])
if not _wexe:
continue
try:
_wr = subprocess.run(
[_wexe, *_wcmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
continue
if _wr.returncode == 0 and _check_fn(_wr.stdout):
_win_amd_gpu = True
break
if _win_amd_gpu:
_safe_print(
_dim(" Note:"),
"AMD GPU detected on Windows. ROCm-enabled PyTorch must be",
)
_safe_print(
" " * 8,
"installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
)
# 3. Extra dependencies
_progress("unsloth extras")
pip_install(
"Installing additional unsloth dependencies",
"--no-cache-dir",
req = REQ_ROOT / "extras.txt",
)
# 3b. Extra dependencies (no-deps) -- audio model support etc.
_progress("extra codecs")
pip_install(
"Installing extras (no-deps)",
"--no-deps",
"--no-cache-dir",
req = REQ_ROOT / "extras-no-deps.txt",
)
# 4. Overrides (torchao, transformers) -- force-reinstall
# Skip entirely when torch is unavailable (e.g. Intel Mac GGUF-only mode)
# because overrides.txt contains torchao which requires torch.
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
else:
_progress("dependency overrides")
pip_install(
"Installing dependency overrides",
"--force-reinstall",
"--no-cache-dir",
req = REQ_ROOT / "overrides.txt",
)
# 5. Triton kernels (no-deps, from source)
# Skip on Windows (no support) and macOS (no support).
if not IS_WINDOWS and not IS_MACOS:
_progress("triton kernels")
pip_install(
"Installing triton kernels",
"--no-deps",
"--no-cache-dir",
req = REQ_ROOT / "triton-kernels.txt",
constrain = False,
)
# # 5. Triton kernels (no-deps, from source)
# # Skip on Windows (no support) and macOS (no support).
# if not IS_WINDOWS and not IS_MACOS:
# _progress("triton kernels")
# pip_install(
# "Installing triton kernels",
# "--no-deps",
# "--no-cache-dir",
# req = REQ_ROOT / "triton-kernels.txt",
# constrain = False,
# )
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("flash-attn")

View file

@ -392,7 +392,7 @@ if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm
fi
# ── Python venv + deps ──
STUDIO_HOME="$HOME/.unsloth/studio"
STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}"
VENV_DIR="$STUDIO_HOME/unsloth_studio"
VENV_T5_530_DIR="$STUDIO_HOME/.venv_t5_530"
VENV_T5_550_DIR="$STUDIO_HOME/.venv_t5_550"
@ -405,7 +405,12 @@ VENV_T5_550_DIR="$STUDIO_HOME/.venv_t5_550"
# Note: do NOT delete $STUDIO_HOME/.venv here — install.sh handles migration
_COLAB_NO_VENV=false
if [ ! -x "$VENV_DIR/bin/python" ]; then
_DOCKER_NO_VENV=false
if [ -n "$UNSLOTH_DOCKER" ]; then
# Docker: packages already in /opt/conda — skip venv entirely.
# Only pre-install .venv_t5 for transformers 5.x switching (handled below).
_DOCKER_NO_VENV=true
elif [ ! -x "$VENV_DIR/bin/python" ]; then
if [ "$IS_COLAB" = true ]; then
# On Colab there is no Studio venv -- install backend deps into system Python.
# Strip all version constraints so pip keeps Colab's pre-installed
@ -470,6 +475,38 @@ if [ "$_COLAB_NO_VENV" = true ]; then
substep "continuing to llama.cpp install for GGUF inference support"
fi
# In Docker, packages are pre-installed in /opt/conda — only install missing
# studio/data-designer deps and pre-install .venv_t5 for transformers 5.x.
if [ "$_DOCKER_NO_VENV" = true ]; then
echo " Docker detected — skipping venv activation."
# Install branch's unsloth/unsloth_cli/studio into /opt/conda
# (overwrites PyPI version with Docker-aware code)
echo " Installing local unsloth from branch..."
pip install --force-reinstall --no-deps "$REPO_ROOT"
# Install only missing deps (studio, data-designer, plugin, metadata patch).
# Heavy packages (torch, unsloth, vllm, etc.) are already in /opt/conda.
# install_python_stack.py has steps 1-5 commented out for this branch.
python "$SCRIPT_DIR/install_python_stack.py"
# Pre-install transformers 5.x into .venv_t5
echo ""
echo " Pre-installing transformers 5.x for newer model support..."
mkdir -p "$VENV_T5_DIR"
pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0" 2>/dev/null
pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1" 2>/dev/null
pip install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2" 2>/dev/null
pip install --target "$VENV_T5_DIR" "tiktoken" 2>/dev/null
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
echo ""
echo "╔══════════════════════════════════════╗"
echo "║ Docker Studio Setup Complete! ║"
echo "╚══════════════════════════════════════╝"
exit 0
fi
# ── Check if Python deps need updating ──
# Compare installed package version against PyPI latest.
# Skip all Python dependency work if versions match (fast update path).

View file

@ -158,56 +158,58 @@ def studio_default(
if ctx.invoked_subcommand is not None:
return
# Always use the studio venv if it exists and we're not already in it
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
# In Docker, packages live in /opt/conda — skip venv re-exec entirely.
if not os.environ.get("UNSLOTH_DOCKER"):
# Always use the studio venv if it exists and we're not already in it
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
if not in_studio_venv:
studio_python = _studio_venv_python()
run_py = _find_run_py()
if studio_python and run_py:
if not silent:
typer.echo("Launching Unsloth Studio... Please wait...")
args = [
str(studio_python),
str(run_py),
"--host",
host,
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
if silent:
args.append("--silent")
# On Windows, os.execvp() spawns a child but the parent lingers,
# so Ctrl+C only kills the parent leaving the child orphaned.
# Use subprocess.run() on Windows so the parent waits for the child.
if sys.platform == "win32":
import subprocess as _sp
if not in_studio_venv:
studio_python = _studio_venv_python()
run_py = _find_run_py()
if studio_python and run_py:
if not silent:
typer.echo("Launching Unsloth Studio... Please wait...")
args = [
str(studio_python),
str(run_py),
"--host",
host,
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
if silent:
args.append("--silent")
# On Windows, os.execvp() spawns a child but the parent lingers,
# so Ctrl+C only kills the parent leaving the child orphaned.
# Use subprocess.run() on Windows so the parent waits for the child.
if sys.platform == "win32":
import subprocess as _sp
proc = _sp.Popen(args)
try:
rc = proc.wait()
except KeyboardInterrupt:
# Child has its own signal handler — let it finish
rc = proc.wait()
if rc != 0:
typer.echo(
f"\nError: Studio server exited unexpectedly (code {rc}).",
err = True,
)
typer.echo(
"Check the error above. If a package is missing, "
"re-run: unsloth studio setup",
err = True,
)
raise typer.Exit(rc)
proc = _sp.Popen(args)
try:
rc = proc.wait()
except KeyboardInterrupt:
# Child has its own signal handler — let it finish
rc = proc.wait()
if rc != 0:
typer.echo(
f"\nError: Studio server exited unexpectedly (code {rc}).",
err = True,
)
typer.echo(
"Check the error above. If a package is missing, "
"re-run: unsloth studio setup",
err = True,
)
raise typer.Exit(rc)
else:
os.execvp(str(studio_python), args)
else:
os.execvp(str(studio_python), args)
else:
typer.echo("Studio not set up. Run install.sh first.")
raise typer.Exit(1)
typer.echo("Studio not set up. Run install.sh first.")
raise typer.Exit(1)
from studio.backend.run import run_server