diff --git a/cli/__init__.py b/cli/__init__.py index fe3321f38a..03b057362b 100644 --- a/cli/__init__.py +++ b/cli/__init__.py @@ -7,7 +7,7 @@ from cli.commands.train import train from cli.commands.inference import inference from cli.commands.export import export, list_checkpoints from cli.commands.ui import ui -from cli.commands.studio import studio +from cli.commands.studio import studio_app app = typer.Typer( help="Command-line interface for Unsloth training, inference, and export.", @@ -19,4 +19,4 @@ app.command()(inference) app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.command()(ui) -app.command()(studio) +app.add_typer(studio_app, name="studio", help="Unsloth Studio commands.") diff --git a/cli/commands/studio.py b/cli/commands/studio.py index e7be41a441..609711d53e 100644 --- a/cli/commands/studio.py +++ b/cli/commands/studio.py @@ -1,20 +1,109 @@ # 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.") -def studio( - port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."), - host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."), - frontend: Optional[Path] = typer.Option(None, "--frontend", "-f", help="Path to frontend build directory."), - silent: bool = typer.Option(False, "--silent", "-q", help="Suppress startup messages."), +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 web UI backend server.""" + """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: @@ -29,9 +118,157 @@ def studio( silent=silent, ) - # Keep running until interrupted 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") diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 13430fc067..9248a13128 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -222,12 +222,7 @@ class ExportOrchestrator: Always spawns a fresh subprocess to ensure a clean Python interpreter. """ - project_root = str( - Path(__file__).resolve().parent.parent.parent.parent.parent - ) - sub_config = { - "project_root": project_root, "checkpoint_path": checkpoint_path, "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 9c329f3f1e..0d278259ad 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -22,19 +22,20 @@ import os import sys import time import traceback +from pathlib import Path from typing import Any logger = get_logger(__name__) -def _activate_transformers_version(model_name: str, project_root: str) -> None: +def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports. If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). """ # Ensure backend is on path for utils imports - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) @@ -42,7 +43,7 @@ def _activate_transformers_version(model_name: str, project_root: str) -> None: resolved = _resolve_base_model(model_name) if needs_transformers_5(resolved): - venv_t5 = os.path.join(project_root, ".venv_t5") + 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) @@ -213,7 +214,7 @@ def run_export_process( Args: cmd_queue: mp.Queue for receiving commands from parent. resp_queue: mp.Queue for sending responses to parent. - config: Initial configuration dict with checkpoint_path and project_root. + config: Initial configuration dict with checkpoint_path. """ import queue as _queue @@ -224,18 +225,17 @@ def run_export_process( from loggers.config import LogConfig if os.getenv("ENVIRONMENT_TYPE", "production") == "production": warnings.filterwarnings("ignore") - + LogConfig.setup_logging( service_name="unsloth-studio-export-worker", env=os.getenv("ENVIRONMENT_TYPE", "production"), ) - project_root = config["project_root"] checkpoint_path = config["checkpoint_path"] # ── 1. Activate correct transformers version BEFORE any ML imports ── try: - _activate_transformers_version(checkpoint_path, project_root) + _activate_transformers_version(checkpoint_path) except Exception as exc: _send_response(resp_queue, { "type": "error", @@ -265,7 +265,7 @@ def run_export_process( "ts": time.time(), }) - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index bd81312d0e..26f3a095ea 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -277,11 +277,9 @@ class InferenceOrchestrator: try: needed_major = "5" if needs_transformers_5(model_name) else "4" - project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent) # Build config dict for subprocess sub_config = { - "project_root": project_root, "model_name": model_name, "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 0e9213664a..ed1e51d413 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -24,19 +24,20 @@ import sys import time import traceback from io import BytesIO +from pathlib import Path from typing import Any logger = get_logger(__name__) -def _activate_transformers_version(model_name: str, project_root: str) -> None: +def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports. If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). """ # Ensure backend is on path for utils imports - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) @@ -44,7 +45,7 @@ def _activate_transformers_version(model_name: str, project_root: str) -> None: resolved = _resolve_base_model(model_name) if needs_transformers_5(resolved): - venv_t5 = os.path.join(project_root, ".venv_t5") + 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) @@ -424,7 +425,7 @@ def run_inference_process( cmd_queue: mp.Queue for receiving commands from parent. resp_queue: mp.Queue for sending responses to parent. cancel_event: mp.Event shared with parent — set by parent to cancel generation. - config: Initial configuration dict with model info and project_root. + config: Initial configuration dict with model info. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports @@ -433,18 +434,17 @@ def run_inference_process( from loggers.config import LogConfig if os.getenv("ENVIRONMENT_TYPE", "production") == "production": warnings.filterwarnings("ignore") - + LogConfig.setup_logging( service_name="unsloth-studio-inference-worker", env=os.getenv("ENVIRONMENT_TYPE", "production"), ) - project_root = config["project_root"] model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── try: - _activate_transformers_version(model_name, project_root) + _activate_transformers_version(model_name) except Exception as exc: _send_response(resp_queue, { "type": "error", @@ -474,7 +474,7 @@ def run_inference_process( "ts": time.time(), }) - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 7676346abe..bbb5d00f95 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -128,12 +128,8 @@ class TrainingBackend: self.eval_enabled = False self._output_dir = None - # Resolve project root (studio/backend/core/training/ → project root) - project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent) - # Build config dict for the subprocess config = { - "project_root": project_root, "model_name": kwargs["model_name"], "training_type": kwargs.get("training_type", "LoRA/QLoRA"), "hf_token": kwargs.get("hf_token", ""), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 10c4012b51..284887e00f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -24,14 +24,14 @@ from typing import Any logger = get_logger(__name__) -def _activate_transformers_version(model_name: str, project_root: str) -> None: +def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports. If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). """ # Ensure backend is on path for utils imports - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) @@ -39,7 +39,7 @@ def _activate_transformers_version(model_name: str, project_root: str) -> None: resolved = _resolve_base_model(model_name) if needs_transformers_5(resolved): - venv_t5 = os.path.join(project_root, ".venv_t5") + 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) @@ -97,12 +97,11 @@ def run_training_process( env=os.getenv("ENVIRONMENT_TYPE", "production"), ) - project_root = config["project_root"] model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── try: - _activate_transformers_version(model_name, project_root) + _activate_transformers_version(model_name) except Exception as exc: event_queue.put({ "type": "error", @@ -128,12 +127,12 @@ def run_training_process( try: _send_status(event_queue, "Importing ML libraries...") - backend_path = os.path.join(project_root, "studio", "backend") + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) from core.training.trainer import UnslothTrainer, TrainingProgress - from utils.paths import ensure_dir, resolve_output_dir, resolve_tensorboard_dir + from utils.paths import ensure_dir, resolve_output_dir, resolve_tensorboard_dir, datasets_root import transformers logger.info("Subprocess loaded transformers %s", transformers.__version__) @@ -588,7 +587,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> all_files: list[str] = [] for dataset_file in local_datasets: file_path = dataset_file if os.path.isabs(dataset_file) else os.path.join( - project_root, "studio", "backend", "assets", "datasets", dataset_file, + str(datasets_root()), dataset_file, ) if os.path.isdir(file_path): file_path_obj = Path(file_path) diff --git a/studio/backend/run.py b/studio/backend/run.py index bffcc11d5f..ebc5eda0a9 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -92,6 +92,10 @@ def run_server( import uvicorn from main import app, setup_frontend + from utils.paths import ensure_studio_directories + + # Create all standard directories on startup + ensure_studio_directories() # Setup frontend if path provided if frontend_path: diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 234d60a006..0f24135103 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -420,10 +420,9 @@ _VLM_MODEL_TYPES = { 'internvl_chat', 'cogvlm2', 'minicpmv', } -# Pre-computed project root and .venv_t5 path for subprocess version switching. -_PROJECT_ROOT = str(Path(__file__).resolve().parent.parent.parent.parent.parent) -_VENV_T5_DIR = os.path.join(_PROJECT_ROOT, ".venv_t5") -_BACKEND_DIR = os.path.join(_PROJECT_ROOT, "studio", "backend") +# Pre-computed .venv_t5 path and backend dir for subprocess version switching. +_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5") +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) # Inline script executed in a subprocess with transformers 5.x activated. # Receives model_name and token via argv, prints JSON result to stdout. diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 25fa7d6585..b36b939548 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -54,8 +54,7 @@ TRANSFORMERS_5_VERSION = "5.2.0" TRANSFORMERS_DEFAULT_VERSION = "4.57.1" # Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1 -_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent # studio/backend/utils/ → project root -_VENV_T5_DIR = str(_PROJECT_ROOT / ".venv_t5") +_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5") def _resolve_base_model(model_name: str) -> str: