13928b5f0e* Add configurable PyTorch mirror via UNSLOTH_PYTORCH_MIRROR env var When set, UNSLOTH_PYTORCH_MIRROR overrides the default https://download.pytorch.org/whl base URL in all four install scripts (install.sh, install.ps1, studio/setup.ps1, studio/install_python_stack.py). When unset or empty, the official URL is used. This lets users behind corporate proxies or in regions with poor connectivity to pytorch.org point at a local mirror without patching scripts. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add pytest for UNSLOTH_PYTORCH_MIRROR in install_python_stack.py Tests that _PYTORCH_WHL_BASE picks up the env var when set, falls back to the official URL when unset or empty, and preserves the value as-is (including trailing slashes). * Remove stale test assertions for missing install.sh messages * Fix GPU mocking in test_get_torch_index_url.sh Extract _has_usable_nvidia_gpu and _has_amd_rocm_gpu alongside get_torch_index_url so the GPU-presence checks work in tests. Add -L flag handling to mock nvidia-smi so it passes the GPU listing check. All 26 tests now pass on CPU-only machines. * Strip trailing slash from UNSLOTH_PYTORCH_MIRROR to avoid double-slash URLs --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for UNSLOTH_PYTORCH_MIRROR env var in install_python_stack.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# install_python_stack.py lives at repo_root/studio/install_python_stack.py
|
|
_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py"
|
|
|
|
OFFICIAL_URL = "https://download.pytorch.org/whl"
|
|
|
|
|
|
def _reload_whl_base(monkeypatch, mirror_value = None):
|
|
"""(Re-)import install_python_stack with a controlled env and return _PYTORCH_WHL_BASE."""
|
|
# Remove cached module so the module-level assignment re-executes
|
|
sys.modules.pop("install_python_stack", None)
|
|
|
|
if mirror_value is None:
|
|
monkeypatch.delenv("UNSLOTH_PYTORCH_MIRROR", raising = False)
|
|
else:
|
|
monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", mirror_value)
|
|
|
|
# Temporarily add the script's directory to sys.path for import
|
|
script_dir = str(_INSTALL_SCRIPT.parent)
|
|
monkeypatch.syspath_prepend(script_dir)
|
|
|
|
import install_python_stack
|
|
|
|
return install_python_stack._PYTORCH_WHL_BASE
|
|
|
|
|
|
class TestPyTorchMirrorEnvVar:
|
|
"""UNSLOTH_PYTORCH_MIRROR controls _PYTORCH_WHL_BASE in install_python_stack."""
|
|
|
|
def test_unset_uses_official_url(self, monkeypatch):
|
|
assert _reload_whl_base(monkeypatch) == OFFICIAL_URL
|
|
|
|
def test_empty_string_falls_back_to_official(self, monkeypatch):
|
|
assert _reload_whl_base(monkeypatch, "") == OFFICIAL_URL
|
|
|
|
def test_custom_mirror_is_used(self, monkeypatch):
|
|
mirror = "https://mirrors.nju.edu.cn/pytorch/whl"
|
|
assert _reload_whl_base(monkeypatch, mirror) == mirror
|
|
|
|
def test_trailing_slash_stripped(self, monkeypatch):
|
|
result = _reload_whl_base(monkeypatch, "https://example.com/whl/")
|
|
assert result == "https://example.com/whl"
|