diff --git a/install_python_stack.py b/install_python_stack.py index e31e96232e..b9526af71b 100644 --- a/install_python_stack.py +++ b/install_python_stack.py @@ -11,25 +11,53 @@ from __future__ import annotations import os import subprocess import sys +import tempfile import urllib.request from pathlib import Path +IS_WINDOWS = sys.platform == "win32" + # ── Paths ────────────────────────────────────────────────────────────── SCRIPT_DIR = Path(__file__).resolve().parent REQ_ROOT = SCRIPT_DIR / "studio" / "backend" / "requirements" SINGLE_ENV = REQ_ROOT / "single-env" CONSTRAINTS = SINGLE_ENV / "constraints.txt" -# ── Helpers ──────────────────────────────────────────────────────────── +# ── Color support ────────────────────────────────────────────────────── + +def _enable_colors() -> bool: + """Try to enable ANSI color support. Returns True if available.""" + if not hasattr(sys.stdout, "fileno"): + return False + try: + if not os.isatty(sys.stdout.fileno()): + return False + except Exception: + return False + if IS_WINDOWS: + try: + import ctypes + kernel32 = ctypes.windll.kernel32 + # Enable ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x0004) on stdout + handle = kernel32.GetStdHandle(-11) # STD_OUTPUT_HANDLE + mode = ctypes.c_ulong() + kernel32.GetConsoleMode(handle, ctypes.byref(mode)) + kernel32.SetConsoleMode(handle, mode.value | 0x0004) + return True + except Exception: + return False + return True # Unix terminals support ANSI by default + +_HAS_COLOR = _enable_colors() def _green(msg: str) -> str: - return f"\033[92m{msg}\033[0m" + return f"\033[92m{msg}\033[0m" if _HAS_COLOR else msg def _cyan(msg: str) -> str: - return f"\033[96m{msg}\033[0m" + return f"\033[96m{msg}\033[0m" if _HAS_COLOR else msg def _red(msg: str) -> str: - return f"\033[91m{msg}\033[0m" + return f"\033[91m{msg}\033[0m" if _HAS_COLOR else msg def run(label: str, cmd: list[str], *, quiet: bool = True) -> None: @@ -47,6 +75,25 @@ def run(label: str, cmd: list[str], *, quiet: bool = True) -> None: sys.exit(result.returncode) +# Packages to skip on Windows (require special build steps) +WINDOWS_SKIP_PACKAGES = {"open_spiel"} + + +def _filter_requirements(req: Path, skip: set[str]) -> Path: + """Return a temp copy of a requirements file with certain packages removed.""" + lines = req.read_text(encoding="utf-8").splitlines(keepends=True) + filtered = [ + line for line in lines + if not any(line.strip().lower().startswith(pkg) for pkg in skip) + ] + tmp = tempfile.NamedTemporaryFile( + mode="w", suffix=".txt", delete=False, encoding="utf-8", + ) + tmp.writelines(filtered) + tmp.close() + return Path(tmp.name) + + def pip_install( label: str, *args: str, @@ -58,9 +105,20 @@ def pip_install( cmd.extend(args) if constrain and CONSTRAINTS.is_file(): cmd.extend(["-c", str(CONSTRAINTS)]) - if req is not None: - cmd.extend(["-r", str(req)]) - run(label, cmd) + actual_req = req + if req is not None and IS_WINDOWS and WINDOWS_SKIP_PACKAGES: + actual_req = _filter_requirements(req, WINDOWS_SKIP_PACKAGES) + if actual_req is not None: + cmd.extend(["-r", str(actual_req)]) + try: + run(label, cmd) + finally: + # Clean up temp file if we created one + if actual_req is not None and actual_req != req: + actual_req.unlink(missing_ok=True) + if req is not None and actual_req != req: + skipped = WINDOWS_SKIP_PACKAGES + print(_cyan(f" (Skipped on Windows: {', '.join(skipped)})")) def download_file(url: str, dest: Path) -> None: diff --git a/setup.bat b/setup.bat new file mode 100644 index 0000000000..ef16abd263 --- /dev/null +++ b/setup.bat @@ -0,0 +1,2 @@ +@echo off +powershell -ExecutionPolicy Bypass -File "%~dp0setup.ps1" %* diff --git a/setup.ps1 b/setup.ps1 index 357c9b0d98..9be3f2efb0 100644 --- a/setup.ps1 +++ b/setup.ps1 @@ -479,7 +479,7 @@ pip install --upgrade pip 2>&1 | Out-Null # The CUDA tag is chosen based on the driver's max supported CUDA version. $CuTag = Get-PytorchCudaTag Write-Host " Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan -pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" +pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" 2>&1 | Out-Null # Ordered heavy dependency installation — shared cross-platform script Write-Host " Running ordered dependency installation..." -ForegroundColor Cyan @@ -645,6 +645,45 @@ if (Test-Path $LlamaServerBin) { } } +# ============================================ +# Add shell aliases (PowerShell profile + cmd batch files) +# ============================================ +Write-Host "" +$RepoDir = $PSScriptRoot +$VenvPython = Join-Path $RepoDir ".venv\Scripts\python.exe" +$CliScript = Join-Path $RepoDir "cli.py" +$FrontendDist = Join-Path $RepoDir "studio\frontend\dist" +$AliasAdded = $false + +# --- PowerShell profile: add functions --- +$ProfileDir = Split-Path $PROFILE -Parent +if (-not (Test-Path $ProfileDir)) { New-Item -ItemType Directory -Path $ProfileDir -Force | Out-Null } +if (-not (Test-Path $PROFILE)) { New-Item -ItemType File -Path $PROFILE -Force | Out-Null } + +if (-not (Select-String -Path $PROFILE -Pattern "unsloth-studio" -Quiet -ErrorAction SilentlyContinue)) { + $block = @" + +# Unsloth Studio launcher +function unsloth-studio { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args } +function unsloth-ui { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args } +"@ + Add-Content -Path $PROFILE -Value $block + Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' added to $PROFILE" -ForegroundColor Green + $AliasAdded = $true +} else { + Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' already exist in $PROFILE" -ForegroundColor Green +} + +# --- cmd.exe: create batch files on PATH so they work from regular terminal --- +$BatDir = Join-Path $RepoDir ".venv\Scripts" +foreach ($name in @("unsloth-studio", "unsloth-ui")) { + $batPath = Join-Path $BatDir "$name.bat" + if (-not (Test-Path $batPath)) { + Set-Content -Path $batPath -Value "@echo off`r`n`"$VenvPython`" `"$CliScript`" studio -f `"$FrontendDist`" %*" + } +} +Write-Host "[OK] Batch launchers created in $BatDir (works from cmd.exe when venv is on PATH)" -ForegroundColor Green + # ============================================ # Done # ============================================ @@ -652,9 +691,15 @@ Write-Host "" Write-Host "+==============================================+" -ForegroundColor Green Write-Host "| Setup Complete! |" -ForegroundColor Green Write-Host "| |" -ForegroundColor Green -Write-Host "| Activate venv: |" -ForegroundColor Green -Write-Host "| cmd: .venv\Scripts\activate.bat |" -ForegroundColor Green -Write-Host "| PS: .\.venv\Scripts\Activate.ps1 |" -ForegroundColor Green +if ($AliasAdded) { + Write-Host "| PowerShell: run '. `$PROFILE' |" -ForegroundColor Green + Write-Host "| or open a new terminal, then: |" -ForegroundColor Green +} else { + Write-Host "| Launch with: |" -ForegroundColor Green +} Write-Host "| |" -ForegroundColor Green -Write-Host "| Then run: unsloth-roland-test studio |" -ForegroundColor Green +Write-Host "| unsloth-studio -H 0.0.0.0 -p 8000 |" -ForegroundColor Green +Write-Host "| |" -ForegroundColor Green +Write-Host "| cmd.exe: .venv\Scripts\activate.bat |" -ForegroundColor Green +Write-Host "| unsloth-studio -H 0.0.0.0 -p 8000 |" -ForegroundColor Green Write-Host "+==============================================+" -ForegroundColor Green \ No newline at end of file