6f77c63229
- Workers now compute backend_path and venv_t5 locally via Path(__file__) - Moved .venv_t5 to ~/.unsloth/studio/.venv_t5 - Added ensure_studio_directories() call on server startup - Expanded CLI studio command into sub-app with setup subcommand
274 lines
9.1 KiB
Python
274 lines
9.1 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
|
|
# Copyright © 2025 Unsloth AI
|
|
|
|
import os
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import typer
|
|
|
|
studio_app = typer.Typer(help="Unsloth Studio commands.")
|
|
|
|
STUDIO_HOME = Path.home() / ".unsloth" / "studio"
|
|
|
|
|
|
def _get_repo_root() -> Optional[Path]:
|
|
"""Find the git clone repo root, or None if pure pip install."""
|
|
# Check 1: __file__ is in the repo (editable install)
|
|
candidate = Path(__file__).resolve().parent.parent.parent
|
|
if (candidate / "pyproject.toml").is_file() and (candidate / "studio" / "setup.sh").is_file():
|
|
return candidate
|
|
# Check 2: CWD is the repo (non-editable wheel, running from repo dir)
|
|
cwd = Path.cwd()
|
|
if (cwd / "pyproject.toml").is_file() and (cwd / "studio" / "setup.sh").is_file():
|
|
return cwd
|
|
return None
|
|
|
|
|
|
def _is_git_clone() -> bool:
|
|
return _get_repo_root() is not None
|
|
|
|
|
|
def _studio_venv_python() -> Optional[Path]:
|
|
"""Return the studio venv Python binary, or None if not set up."""
|
|
if platform.system() == "Windows":
|
|
p = STUDIO_HOME / ".venv" / "Scripts" / "python.exe"
|
|
else:
|
|
p = STUDIO_HOME / ".venv" / "bin" / "python"
|
|
return p if p.is_file() else None
|
|
|
|
|
|
def _find_run_py() -> Optional[Path]:
|
|
"""Find studio/backend/run.py."""
|
|
# 1. Repo root (git clone / editable)
|
|
repo = _get_repo_root()
|
|
if repo:
|
|
run_py = repo / "studio" / "backend" / "run.py"
|
|
if run_py.is_file():
|
|
return run_py
|
|
# 2. Studio venv's site-packages
|
|
for match in (STUDIO_HOME / ".venv").glob("lib/python*/site-packages/studio/backend/run.py"):
|
|
return match
|
|
# 3. Current package's site-packages
|
|
run_py = Path(__file__).resolve().parent.parent.parent / "studio" / "backend" / "run.py"
|
|
return run_py if run_py.is_file() else None
|
|
|
|
|
|
def _find_install_script() -> Optional[Path]:
|
|
"""Find studio/install_python_stack.py."""
|
|
# 1. Repo root
|
|
repo = _get_repo_root()
|
|
if repo:
|
|
s = repo / "studio" / "install_python_stack.py"
|
|
if s.is_file():
|
|
return s
|
|
# 2. Relative to __file__ (in site-packages)
|
|
s = Path(__file__).resolve().parent.parent.parent / "studio" / "install_python_stack.py"
|
|
return s if s.is_file() else None
|
|
|
|
|
|
# ── unsloth studio (server) ──────────────────────────────────────────
|
|
|
|
@studio_app.callback(invoke_without_command=True)
|
|
def studio_default(
|
|
ctx: typer.Context,
|
|
port: int = typer.Option(8000, "--port", "-p"),
|
|
host: str = typer.Option("0.0.0.0", "--host", "-H"),
|
|
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"),
|
|
silent: bool = typer.Option(False, "--silent", "-q"),
|
|
):
|
|
"""Launch the Unsloth Studio server."""
|
|
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 / ".venv"
|
|
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 with studio venv...")
|
|
args = [str(studio_python), str(run_py), "--host", host, "--port", str(port)]
|
|
if frontend:
|
|
args.extend(["--frontend", str(frontend)])
|
|
if silent:
|
|
args.append("--silent")
|
|
os.execvp(str(studio_python), args)
|
|
else:
|
|
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
|
|
raise typer.Exit(1)
|
|
|
|
from studio.backend.run import run_server
|
|
|
|
if not silent:
|
|
from studio.backend.run import _resolve_external_ip
|
|
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
|
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
|
|
|
|
run_server(
|
|
host=host,
|
|
port=port,
|
|
frontend_path=frontend,
|
|
silent=silent,
|
|
)
|
|
|
|
try:
|
|
while True:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
typer.echo("\nShutting down...")
|
|
|
|
|
|
# ── unsloth studio setup ─────────────────────────────────────────────
|
|
|
|
@studio_app.command()
|
|
def setup():
|
|
"""Run one-time Studio environment setup."""
|
|
if _is_git_clone():
|
|
_dev_setup()
|
|
else:
|
|
_pip_setup()
|
|
|
|
|
|
def _dev_setup():
|
|
"""Git-clone: run setup.sh / setup.ps1."""
|
|
repo_root = _get_repo_root()
|
|
studio_dir = repo_root / "studio"
|
|
if platform.system() == "Windows":
|
|
script = studio_dir / "setup.ps1"
|
|
subprocess.run(["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)], check=True)
|
|
else:
|
|
script = studio_dir / "setup.sh"
|
|
subprocess.run(["bash", str(script)], check=True)
|
|
|
|
|
|
def _pip_setup():
|
|
"""Pip-install: create studio venv, install all deps, build extras."""
|
|
import venv as _venv
|
|
|
|
venv_dir = STUDIO_HOME / ".venv"
|
|
venv_t5_dir = STUDIO_HOME / ".venv_t5"
|
|
|
|
if platform.system() == "Windows":
|
|
venv_python = venv_dir / "Scripts" / "python.exe"
|
|
venv_pip = venv_dir / "Scripts" / "pip.exe"
|
|
else:
|
|
venv_python = venv_dir / "bin" / "python"
|
|
venv_pip = venv_dir / "bin" / "pip"
|
|
|
|
typer.echo("Setting up Unsloth Studio...")
|
|
|
|
# 1. Create venv
|
|
if not venv_python.is_file():
|
|
typer.echo(f" Creating venv at {venv_dir}...")
|
|
STUDIO_HOME.mkdir(parents=True, exist_ok=True)
|
|
_venv.create(str(venv_dir), with_pip=True)
|
|
|
|
# 2. Install all Python deps via install_python_stack.py
|
|
install_script = _find_install_script()
|
|
if install_script:
|
|
typer.echo(" Installing Python dependencies...")
|
|
subprocess.run([str(venv_python), str(install_script)], check=True)
|
|
else:
|
|
typer.echo("Error: Could not find install_python_stack.py")
|
|
raise typer.Exit(1)
|
|
|
|
# 3. Pre-install transformers 5.x overlay
|
|
if venv_t5_dir.is_dir() and any(venv_t5_dir.iterdir()):
|
|
typer.echo(f" Transformers 5.x overlay already at {venv_t5_dir}")
|
|
else:
|
|
typer.echo(" Installing transformers 5.x overlay...")
|
|
venv_t5_dir.mkdir(parents=True, exist_ok=True)
|
|
subprocess.run(
|
|
[str(venv_pip), "install",
|
|
"--target", str(venv_t5_dir), "--no-deps", "transformers==5.2.0"],
|
|
check=True,
|
|
)
|
|
subprocess.run(
|
|
[str(venv_pip), "install",
|
|
"--target", str(venv_t5_dir), "--no-deps", "huggingface_hub==1.3.0"],
|
|
check=True,
|
|
)
|
|
typer.echo(f" Installed to {venv_t5_dir}")
|
|
|
|
# 4. Build llama.cpp
|
|
_build_llama_cpp()
|
|
|
|
typer.echo("")
|
|
typer.echo("Setup complete! Run 'unsloth studio' to start.")
|
|
|
|
|
|
def _build_llama_cpp():
|
|
"""Build llama.cpp at ~/.unsloth/llama.cpp/."""
|
|
import shutil
|
|
|
|
unsloth_home = Path.home() / ".unsloth"
|
|
llama_dir = unsloth_home / "llama.cpp"
|
|
|
|
if not shutil.which("cmake"):
|
|
typer.echo(" cmake not found — skipping llama.cpp build")
|
|
return
|
|
if not shutil.which("git"):
|
|
typer.echo(" git not found — skipping llama.cpp build")
|
|
return
|
|
|
|
typer.echo(" Building llama.cpp for GGUF inference...")
|
|
|
|
if llama_dir.exists():
|
|
shutil.rmtree(llama_dir)
|
|
unsloth_home.mkdir(parents=True, exist_ok=True)
|
|
|
|
result = subprocess.run(
|
|
["git", "clone", "--depth", "1", "https://github.com/ggml-org/llama.cpp.git", str(llama_dir)],
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(" Failed to clone llama.cpp")
|
|
return
|
|
|
|
cmake_args = []
|
|
nvcc_path = shutil.which("nvcc")
|
|
if not nvcc_path and Path("/usr/local/cuda/bin/nvcc").is_file():
|
|
nvcc_path = "/usr/local/cuda/bin/nvcc"
|
|
if nvcc_path:
|
|
typer.echo(f" Building with CUDA (nvcc: {nvcc_path})...")
|
|
cmake_args.append("-DGGML_CUDA=ON")
|
|
else:
|
|
typer.echo(" Building CPU-only...")
|
|
|
|
build_dir = llama_dir / "build"
|
|
result = subprocess.run(
|
|
["cmake", "-S", str(llama_dir), "-B", str(build_dir)] + cmake_args,
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(" cmake configure failed")
|
|
return
|
|
|
|
ncpu = str(os.cpu_count() or 4)
|
|
result = subprocess.run(
|
|
["cmake", "--build", str(build_dir), "--config", "Release",
|
|
"--target", "llama-server", f"-j{ncpu}"],
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(" llama-server build failed")
|
|
return
|
|
|
|
subprocess.run(
|
|
["cmake", "--build", str(build_dir), "--config", "Release",
|
|
"--target", "llama-quantize", f"-j{ncpu}"],
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
)
|
|
|
|
server_bin = build_dir / "bin" / "llama-server"
|
|
if server_bin.is_file():
|
|
typer.echo(f" llama-server built at {server_bin}")
|
|
else:
|
|
typer.echo(" llama-server binary not found after build")
|