Compare commits

...
Sign in to create a new pull request.

7 commits

Author SHA1 Message Date
Test
c8c371610b fix(test): use capsys instead of caplog for structlog output
TestLogGpuMemory tests used pytest's caplog fixture to capture log
output, but log_gpu_memory() uses structlog which writes to stdout,
not Python's standard logging module. caplog only captures standard
logging messages, so assertions always failed on empty text.

Switch to capsys which captures stdout directly.
2026-03-13 09:42:46 +00:00
Test
4b057bc84e Merge branch 'pr-4275' 2026-03-13 09:40:10 +00:00
Test
7ad2224909 Merge branch 'pr-4274' 2026-03-13 09:40:10 +00:00
Test
54afb10206 Merge branch 'pr-4268' 2026-03-13 09:40:10 +00:00
Daniel Han
9648c43065 fix(seed): disable remote code execution for seed inspect loads 2026-03-13 01:48:22 -07:00
Daniel Han
24914d359b fix: disable remote code loading for ai-assist model hint lookup 2026-03-13 01:45:24 -07:00
LeoBorcherding
d260fbb30c fix: install data-designer plugin non-editable for Colab compatibility
Editable installs (-e) work via a .pth file that is only processed at
Python startup. In Colab the kernel is already running when setup.sh
installs the plugin, so the .pth file never gets picked up and
data_designer_unstructured_seed is not importable.

Remove -e so pip copies the package files directly into site-packages,
which the live kernel can find immediately. Local venv installs are
unaffected since the venv is always created fresh before install.
2026-03-12 22:55:06 -05:00
6 changed files with 41 additions and 26 deletions

View file

@ -118,6 +118,7 @@ def _build_stream_load_kwargs(
"path": dataset_name, "path": dataset_name,
"split": split, "split": split,
"streaming": True, "streaming": True,
"trust_remote_code": False,
} }
if data_file: if data_file:
kwargs["data_files"] = [data_file] kwargs["data_files"] = [data_file]

View file

@ -0,0 +1,10 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from pathlib import Path
def test_seed_inspect_load_kwargs_disables_remote_code_execution():
seed_route = Path("studio/backend/routes/data_recipe/seed.py").read_text()
assert '"trust_remote_code": False' in seed_route

View file

@ -285,7 +285,7 @@ class TestLogGpuMemory:
def test_does_not_raise(self): def test_does_not_raise(self):
log_gpu_memory("test") log_gpu_memory("test")
def test_logs_gpu_info_when_available(self, caplog): def test_logs_gpu_info_when_available(self, capsys):
fake_info = { fake_info = {
"available": True, "available": True,
"backend": "cuda", "backend": "cuda",
@ -295,35 +295,27 @@ class TestLogGpuMemory:
"utilization_pct": 12.5, "utilization_pct": 12.5,
"free_gb": 14.0, "free_gb": 14.0,
} }
import structlog
from loggers import get_logger
with ( with patch(
patch( "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
),
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
): ):
log_gpu_memory("unit-test") log_gpu_memory("unit-test")
assert "unit-test" in caplog.text captured = capsys.readouterr().out
assert "CUDA" in caplog.text assert "unit-test" in captured
assert "FakeGPU" in caplog.text assert "CUDA" in captured
assert "FakeGPU" in captured
def test_logs_cpu_fallback_when_no_gpu(self, caplog): def test_logs_cpu_fallback_when_no_gpu(self, capsys):
fake_info = {"available": False, "backend": "cpu"} fake_info = {"available": False, "backend": "cpu"}
import structlog
from loggers import get_logger
with ( with patch(
patch( "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
),
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
): ):
log_gpu_memory("cpu-test") log_gpu_memory("cpu-test")
assert "No GPU available" in caplog.text captured = capsys.readouterr().out
assert "No GPU available" in captured
# ========== format_error_message() ========== # ========== format_error_message() ==========

View file

@ -515,7 +515,12 @@ def _run_multi_pass_advisor(
try: try:
from utils.models.model_config import load_model_config from utils.models.model_config import load_model_config
config = load_model_config(model_name, use_auth = True, token = hf_token) config = load_model_config(
model_name,
use_auth = True,
token = hf_token,
trust_remote_code = False,
)
archs = getattr(config, "architectures", []) archs = getattr(config, "architectures", [])
if archs and "Gemma3nForConditionalGeneration" in archs: if archs and "Gemma3nForConditionalGeneration" in archs:
is_gemma_3n = True is_gemma_3n = True

View file

@ -383,7 +383,10 @@ for canonical_file, model_names in MODEL_NAME_MAPPING.items():
def load_model_config( def load_model_config(
model_name: str, use_auth: bool = False, token: Optional[str] = None model_name: str,
use_auth: bool = False,
token: Optional[str] = None,
trust_remote_code: bool = True,
): ):
""" """
Load model config with optional authentication control. Load model config with optional authentication control.
@ -392,18 +395,23 @@ def load_model_config(
if token: if token:
# Explicit token provided - use it # Explicit token provided - use it
return AutoConfig.from_pretrained( return AutoConfig.from_pretrained(
model_name, trust_remote_code = True, token = token model_name, trust_remote_code = trust_remote_code, token = token
) )
if not use_auth: if not use_auth:
# Load without any authentication (for public model checks) # Load without any authentication (for public model checks)
with without_hf_auth(): with without_hf_auth():
return AutoConfig.from_pretrained( return AutoConfig.from_pretrained(
model_name, trust_remote_code = True, token = None model_name,
trust_remote_code = trust_remote_code,
token = None,
) )
# Use default authentication (cached tokens) # Use default authentication (cached tokens)
return AutoConfig.from_pretrained(model_name, trust_remote_code = True) return AutoConfig.from_pretrained(
model_name,
trust_remote_code = trust_remote_code,
)
# VLM architecture suffixes and known VLM model_type values. # VLM architecture suffixes and known VLM model_type values.

View file

@ -270,7 +270,6 @@ def install_python_stack() -> int:
"Installing local data-designer unstructured plugin", "Installing local data-designer unstructured plugin",
"--no-cache-dir", "--no-cache-dir",
"--no-deps", "--no-deps",
"-e",
str(LOCAL_DD_UNSTRUCTURED_PLUGIN), str(LOCAL_DD_UNSTRUCTURED_PLUGIN),
constrain = False, constrain = False,
) )