Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 49 additions & 15 deletions studio/install_python_stack.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,20 +22,20 @@

IS_WINDOWS = sys.platform == "win32"

# ── Verbosity control ──────────────────────────────────────────────────────────
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal progress bar (one line, in-place).
# Set UNSLOTH_VERBOSE=1 in the environment to restore full per-step output:
# Linux/Mac: UNSLOTH_VERBOSE=1 ./studio/setup.sh
# Windows: $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1
VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"

# Progress bar state updated by _progress() as each install step runs.
# Progress bar state -- updated by _progress() as each install step runs.
# _TOTAL counts: pip-upgrade + 7 shared steps + triton (non-Windows) + local-plugin + finalize
# Update _TOTAL here if you add or remove install steps in install_python_stack().
_STEP: int = 0
_TOTAL: int = 0 # set at runtime in install_python_stack() based on platform

# ── Paths ──────────────────────────────────────────────────────────────
# -- Paths --------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
REQ_ROOT = SCRIPT_DIR / "backend" / "requirements"
SINGLE_ENV = REQ_ROOT / "single-env"
Expand All @@ -44,7 +44,39 @@
SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed"
)

# ── Color support ──────────────────────────────────────────────────────
# -- Unicode-safe printing ---------------------------------------------
# On Windows the default console encoding can be a legacy code page
# (e.g. CP1252) that cannot represent Unicode glyphs such as ✅ or ❌.
# _safe_print() gracefully degrades to ASCII equivalents so the
# installer never crashes just because of a status glyph.

_UNICODE_TO_ASCII: dict[str, str] = {
"\u2705": "[OK]", # ✅
"\u274c": "[FAIL]", # ❌
"\u26a0\ufe0f": "[!]", # ⚠️ (warning + variation selector)
"\u26a0": "[!]", # ⚠ (warning without variation selector)
}


def _safe_print(*args: object, **kwargs: object) -> None:
"""Drop-in print() replacement that survives non-UTF-8 consoles."""
try:
print(*args, **kwargs)
except UnicodeEncodeError:
# Stringify, then swap emoji for ASCII equivalents
text = " ".join(str(a) for a in args)
for uni, ascii_alt in _UNICODE_TO_ASCII.items():
text = text.replace(uni, ascii_alt)
# Final fallback: replace any remaining unencodable chars
print(
text.encode(sys.stdout.encoding or "ascii", errors = "replace").decode(
sys.stdout.encoding or "ascii", errors = "replace"
Comment on lines +61 to +73
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To make _safe_print a more robust "drop-in replacement" for print() as the docstring states, I suggest a small improvement. The current implementation in the except block doesn't handle keyword arguments like sep correctly. It joins all arguments with a space and then passes the original kwargs to a print call with a single argument, where sep has no effect.

A more robust approach is to process each argument individually to make it safe, and then pass the safe arguments to print. This preserves the behavior of all keyword arguments like sep, end, etc.

Suggested change
def _safe_print(*args: object, **kwargs: object) -> None:
"""Drop-in print() replacement that survives non-UTF-8 consoles."""
try:
print(*args, **kwargs)
except UnicodeEncodeError:
# Stringify, then swap emoji for ASCII equivalents
text = " ".join(str(a) for a in args)
for uni, ascii_alt in _UNICODE_TO_ASCII.items():
text = text.replace(uni, ascii_alt)
# Final fallback: replace any remaining unencodable chars
print(
text.encode(sys.stdout.encoding or "ascii", errors = "replace").decode(
sys.stdout.encoding or "ascii", errors = "replace"
def _safe_print(*args: object, **kwargs: object) -> None:
"""Drop-in print() replacement that survives non-UTF-8 consoles."""
try:
print(*args, **kwargs)
except UnicodeEncodeError:
# Sanitize each argument individually to preserve kwargs like 'sep'
safe_args = []
for arg in args:
s = str(arg)
for uni, ascii_alt in _UNICODE_TO_ASCII.items():
s = s.replace(uni, ascii_alt)
safe_args.append(
s.encode(sys.stdout.encoding or "ascii", errors="replace").decode(
sys.stdout.encoding or "ascii", errors="replace"
)
)
print(*safe_args, **kwargs)

),
**kwargs,
)


# -- Color support ------------------------------------------------------


def _enable_colors() -> bool:
Expand Down Expand Up @@ -72,7 +104,7 @@ def _enable_colors() -> bool:
return True # Unix terminals support ANSI by default


# Colors disabled Colab and most CI runners render ANSI fine, but plain output
# Colors disabled -- Colab and most CI runners render ANSI fine, but plain output
# is cleaner in the notebook cell. Re-enable by setting _HAS_COLOR = _enable_colors()
_HAS_COLOR = False

Expand All @@ -92,7 +124,7 @@ def _red(msg: str) -> str:
def _progress(label: str) -> None:
"""Print an in-place progress bar for the current install step.

Uses only stdlib (sys.stdout) no extra packages required.
Uses only stdlib (sys.stdout) -- no extra packages required.
In VERBOSE mode this is a no-op; per-step labels are printed by run() instead.
"""
global _STEP
Expand All @@ -119,7 +151,7 @@ def run(
stderr = subprocess.STDOUT if quiet else None,
)
if result.returncode != 0:
print(_red(f"❌ {label} failed (exit code {result.returncode}):"))
_safe_print(_red(f"❌ {label} failed (exit code {result.returncode}):"))
if result.stdout:
print(result.stdout.decode(errors = "replace"))
sys.exit(result.returncode)
Expand All @@ -129,7 +161,7 @@ def run(
# Packages to skip on Windows (require special build steps)
WINDOWS_SKIP_PACKAGES = {"open_spiel", "triton_kernels"}

# ── uv bootstrap ──────────────────────────────────────────────────────
# -- uv bootstrap ------------------------------------------------------

USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack()
UV_NEEDS_SYSTEM = False # Set by _bootstrap_uv() via probe
Expand Down Expand Up @@ -267,7 +299,9 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
text = True,
)
if result.returncode != 0:
print(_red(f" ⚠️ Could not find package {package_name}, skipping patch"))
_safe_print(
_red(f" ⚠️ Could not find package {package_name}, skipping patch")
)
return

location = None
Expand All @@ -277,15 +311,15 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
break

if not location:
print(_red(f" ⚠️ Could not determine location of {package_name}"))
_safe_print(_red(f" ⚠️ Could not determine location of {package_name}"))
return

dest = Path(location) / relative_path
print(_cyan(f" Patching {dest.name} in {package_name}..."))
download_file(url, dest)


# ── Main install sequence ─────────────────────────────────────────────
# -- Main install sequence ---------------------------------------------


def install_python_stack() -> int:
Expand Down Expand Up @@ -316,7 +350,7 @@ def install_python_stack() -> int:
req = REQ_ROOT / "extras.txt",
)

# 3b. Extra dependencies (no-deps) audio model support etc.
# 3b. Extra dependencies (no-deps) -- audio model support etc.
_progress("extra codecs")
pip_install(
"Installing extras (no-deps)",
Expand All @@ -325,7 +359,7 @@ def install_python_stack() -> int:
req = REQ_ROOT / "extras-no-deps.txt",
)

# 4. Overrides (torchao, transformers) force-reinstall
# 4. Overrides (torchao, transformers) -- force-reinstall
_progress("dependency overrides")
pip_install(
"Installing dependency overrides",
Expand Down Expand Up @@ -393,7 +427,7 @@ def install_python_stack() -> int:

# 11. Local Data Designer seed plugin
if not LOCAL_DD_UNSTRUCTURED_PLUGIN.is_dir():
print(
_safe_print(
_red(
f"❌ Missing local plugin directory: {LOCAL_DD_UNSTRUCTURED_PLUGIN}",
),
Expand Down Expand Up @@ -422,7 +456,7 @@ def install_python_stack() -> int:
stderr = subprocess.DEVNULL,
)

print(_green("✅ Python dependencies installed"))
_safe_print(_green("✅ Python dependencies installed"))
return 0


Expand Down