diff --git a/cli/commands/studio.py b/cli/commands/studio.py index da5198ad75..33375a89dd 100644 --- a/cli/commands/studio.py +++ b/cli/commands/studio.py @@ -12,7 +12,13 @@ import typer studio_app = typer.Typer(help = "Unsloth Studio commands.") -STUDIO_HOME = Path.home() / ".unsloth" / "studio" + +def _studio_home() -> Path: + """Studio root, overridable via UNSLOTH_STUDIO_HOME.""" + custom = os.environ.get("UNSLOTH_STUDIO_HOME") + if custom: + return Path(custom).expanduser().resolve() + return Path.home() / ".unsloth" / "studio" # __file__ is cli/commands/studio.py — two parents up is the package root # (either site-packages or the repo root for editable installs). @@ -22,9 +28,9 @@ _PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent 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" + p = _studio_home() / ".venv" / "Scripts" / "python.exe" else: - p = STUDIO_HOME / ".venv" / "bin" / "python" + p = _studio_home() / ".venv" / "bin" / "python" return p if p.is_file() else None @@ -44,7 +50,7 @@ def _find_run_py() -> Optional[Path]: "lib/python*/site-packages/studio/backend/run.py", "Lib/site-packages/studio/backend/run.py", ): - for match in (STUDIO_HOME / ".venv").glob(pattern): + for match in (_studio_home() / ".venv").glob(pattern): return match return None @@ -64,7 +70,7 @@ def _find_setup_script() -> Optional[Path]: f"lib/python*/site-packages/studio/{name}", f"Lib/site-packages/studio/{name}", ): - for match in (STUDIO_HOME / ".venv").glob(pattern): + for match in (_studio_home() / ".venv").glob(pattern): return match return None @@ -85,7 +91,7 @@ def studio_default( return # Always use the studio venv if it exists and we're not already in it - studio_venv_dir = STUDIO_HOME / ".venv" + studio_venv_dir = _studio_home() / ".venv" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) if not in_studio_venv: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 81763c3095..288d9c04a9 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -9,12 +9,15 @@ import tempfile def studio_root() -> Path: + custom = os.environ.get("UNSLOTH_STUDIO_HOME") + if custom: + return Path(custom).expanduser().resolve() return Path.home() / ".unsloth" / "studio" def cache_root() -> Path: """Central cache directory for all studio downloads (models, datasets, etc.).""" - return Path.home() / ".unsloth" / "studio" / "cache" + return studio_root() / "cache" def assets_root() -> Path: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 2420448deb..2966b2ed39 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -826,9 +826,10 @@ if (-not $PythonCmd) { Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundColor Green -# Always create a .venv for isolation -- even for pip installs. -# Created in the repo root (parent of studio/). -$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv" +# ── Studio home (configurable via UNSLOTH_STUDIO_HOME) ── +$StudioHome = if ($env:UNSLOTH_STUDIO_HOME) { $env:UNSLOTH_STUDIO_HOME } else { Join-Path $env:USERPROFILE ".unsloth\studio" } + +$VenvDir = Join-Path $StudioHome ".venv" if (-not (Test-Path $VenvDir)) { Write-Host " Creating virtual environment at $VenvDir..." -ForegroundColor Cyan & $PythonCmd -m venv $VenvDir @@ -901,7 +902,7 @@ $ErrorActionPreference = $prevEAP # The training subprocess just prepends .venv_t5/ to sys.path — instant switch. Write-Host "" Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan -$VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5" +$VenvT5Dir = Join-Path $StudioHome ".venv_t5" if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir } New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null $prevEAP_t5 = $ErrorActionPreference diff --git a/studio/setup.sh b/studio/setup.sh index 4c8a6c7dde..ce6181111a 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -227,8 +227,9 @@ if [ "$IS_COLAB" = true ]; then # Colab: install packages directly without venv install_python_stack else - # Local: create venv under ~/.unsloth/studio/ (shared location, not in repo) - STUDIO_HOME="$HOME/.unsloth/studio" + # Local: create venv under studio home (shared location, not in repo) + # Configurable via UNSLOTH_STUDIO_HOME; defaults to ~/.unsloth/studio + STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}" VENV_DIR="$STUDIO_HOME/.venv" VENV_T5_DIR="$STUDIO_HOME/.venv_t5" mkdir -p "$STUDIO_HOME"