From 68f30a74b110784dce0da77b462e9eb2648fae3e 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 | 345 ++++++++++++++------------------- studio/setup.sh | 41 +++- unsloth_cli/commands/studio.py | 96 ++++----- 3 files changed, 229 insertions(+), 253 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2b9fd084d4..1ef626ce39 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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") diff --git a/studio/setup.sh b/studio/setup.sh index c2044b5141..d2bd9485ba 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -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). diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index a3c0840be1..ddc6216d08 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -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