From 80706f1f854c55ec7fbbddee4fd0ab5264992eee Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 27 Mar 2026 19:39:37 +0000 Subject: [PATCH] add Docker support: skip venv, install only missing deps --- studio/install_python_stack.py | 156 ++++++++++++++++++++++++++++++--- studio/setup.sh | 70 ++++++++------- unsloth_cli/commands/studio.py | 14 +-- 3 files changed, 189 insertions(+), 51 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 5e785c934d..a9f51dd0c2 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1421,9 +1421,9 @@ def install_python_stack() -> int: ) _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") @@ -1452,17 +1452,115 @@ def install_python_stack() -> int: ).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", + # ) # macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the # mlx-vlm / mlx-lm transformers pin -- set at module load). @@ -1586,6 +1684,11 @@ def install_python_stack() -> int: "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). @@ -1683,6 +1786,33 @@ def install_python_stack() -> int: req = REQ_ROOT / "triton-kernels.txt", constrain = False, ) + # pip_install( + # "Installing extras (no-deps)", + # "--no-deps", + # "--no-cache-dir", + # req = REQ_ROOT / "extras-no-deps.txt", + # ) + + # # 4. Overrides (torchao, transformers) -- force-reinstall + # _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, + # ) if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: _progress("flash-attn") diff --git a/studio/setup.sh b/studio/setup.sh index 3a149d465f..3090cac2ec 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -450,36 +450,7 @@ if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm fi # ── Python venv + deps ── -# UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the install root -# (mirrors install.sh). UNSLOTH_STUDIO_HOME wins when both are set. -_studio_override_var="" -_studio_override="${UNSLOTH_STUDIO_HOME:-}" -if [ -n "$_studio_override" ]; then - _studio_override_var="UNSLOTH_STUDIO_HOME" -else - _studio_override="${STUDIO_HOME:-}" - [ -n "$_studio_override" ] && _studio_override_var="STUDIO_HOME" -fi -# Strip whitespace so " " is treated as unset (matches Python .strip()). -_studio_override=$(printf '%s' "$_studio_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') -case "$_studio_override" in - "~") _studio_override="$HOME" ;; - "~/"*) _studio_override="$HOME/${_studio_override#'~/'}" ;; -esac -if [ -n "$_studio_override" ]; then - # setup.sh runs against an existing install (via 'unsloth studio update'); - # a typo in the override must fail fast instead of materializing an - # empty workspace dir. Mirrors setup.ps1 behavior. - if [ ! -d "$_studio_override" ]; then - echo "ERROR: $_studio_override_var=$_studio_override does not exist." >&2 - echo " Run install.sh to create the install root before 'unsloth studio update'." >&2 - exit 1 - fi - [ -w "$_studio_override" ] || { echo "ERROR: $_studio_override_var=$_studio_override is not writable." >&2; exit 1; } - STUDIO_HOME="$(CDPATH= cd -P -- "$_studio_override" && pwd -P)" || exit 1 -else - STUDIO_HOME="$HOME/.unsloth/studio" -fi +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" @@ -492,7 +463,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 @@ -557,6 +533,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). diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index edb8ea2cf4..c15a3d0ea5 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -626,9 +626,11 @@ def studio_default( raise typer.Exit(2) return - # Use the studio venv if it exists and we aren't 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() @@ -681,10 +683,8 @@ def studio_default( ) raise typer.Exit(rc) 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