ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: integrate structlog, configure workers for prod logging, and migrate print statements

Browse source
This commit is contained in:
Roland Tannous 2026-03-11 12:33:16 +00:00
commit 817f2e8dcc
38 changed files with 573 additions and 309 deletions
Download patch file Download diff file Expand all files Collapse all files

28
studio/backend/colab.py
View file

@ -6,6 +6,17 @@ Colab-specific helpers for running Unsloth Studio.
Uses Colab's built-in proxy - no external tunneling needed!
"""
from pathlib import Path
import sys
# Add backend to path early so local modules like loggers can be imported
backend_path = str(Path(__file__).parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from loggers import get_logger
logger = get_logger(__name__)
def get_colab_url(port: int = 8000) -> str:
@ -19,7 +30,7 @@ def get_colab_url(port: int = 8000) -> str:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec=5)
return url if url else f"http://localhost:{port}"
except Exception as e:
print(f"Note: Could not get Colab URL ({e})")
logger.info(f"Note: Could not get Colab URL ({e})")
return f"http://localhost:{port}"
@ -61,14 +72,9 @@ def start(port: int = 8000):
"""
import sys
print("🦥 Starting Unsloth Studio...")
logger.info("🦥 Starting Unsloth Studio...")
# Add backend to path
backend_path = str(Path(__file__).parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
print(" Loading backend...")
logger.info(" Loading backend...")
from run import run_server
# Auto-detect frontend path
@ -76,14 +82,14 @@ def start(port: int = 8000):
frontend_path = repo_root / "frontend" / "dist"
if not frontend_path.exists():
print("❌ Frontend not built! Please run the setup cell first.")
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
print(" Starting server...")
logger.info(" Starting server...")
# Start server silently
run_server(host="0.0.0.0", port=port, frontend_path=frontend_path, silent=True)
print(" Server started!")
logger.info(" Server started!")
# Show the clickable link with real URL
show_link(port)

15
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -4,6 +4,8 @@
from __future__ import annotations
import json
import structlog
import loggers
import logging
import re
import shutil
@ -71,6 +73,19 @@ def run_job_process(
Subprocess entrypoint.
Sends events to `event_queue`.
"""
import os
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-data-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
)
event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()})
try:

5
studio/backend/core/data_recipe/local_callable_validators.py
View file

@ -4,7 +4,8 @@
from __future__ import annotations
import json
import logging
import structlog
from loggers import get_logger
import subprocess
from copy import deepcopy
from dataclasses import dataclass
@ -12,7 +13,7 @@ from functools import lru_cache
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
OXC_VALIDATION_FN_MARKER = "unsloth_oxc_validator"

5
studio/backend/core/export/export.py
View file

@ -7,7 +7,8 @@ Export backend - handles model exporting in various formats
"""
import glob
import json
import logging
import structlog
from loggers import get_logger
import os
import shutil
from pathlib import Path
@ -23,7 +24,7 @@ from utils.models import is_vision_model, get_base_model_from_lora
from utils.models.model_config import detect_audio_type
from core.inference import get_inference_backend
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _is_wsl():

5
studio/backend/core/export/orchestrator.py
View file

@ -14,7 +14,8 @@ the old subprocess is killed and a new one is spawned with the correct version.
Pattern follows core/inference/orchestrator.py.
"""
import atexit
import logging
import structlog
from loggers import get_logger
import multiprocessing as mp
import queue
import threading
@ -22,7 +23,7 @@ import time
from pathlib import Path
from typing import Any, List, Optional, Tuple
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")

16
studio/backend/core/export/worker.py
View file

@ -16,14 +16,15 @@ Pattern follows core/inference/worker.py and core/training/worker.py.
"""
from __future__ import annotations
import logging
import structlog
from loggers import get_logger
import os
import sys
import time
import traceback
from typing import Any
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _activate_transformers_version(model_name: str, project_root: str) -> None:
@ -217,6 +218,17 @@ def run_export_process(
import queue as _queue
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-export-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
)
project_root = config["project_root"]
checkpoint_path = config["checkpoint_path"]

5
studio/backend/core/inference/audio_codecs.py
View file

@ -8,13 +8,14 @@ Supports: SNAC (Orpheus), CSM (Sesame), BiCodec (Spark), DAC (OuteTTS)
import io
import re
import wave
import logging
import structlog
from loggers import get_logger
from typing import Optional, Tuple
import numpy as np
import torch
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _numpy_to_wav_bytes(waveform: np.ndarray, sample_rate: int) -> bytes:

5
studio/backend/core/inference/inference.py
View file

@ -20,11 +20,12 @@ from utils.utils import format_error_message
from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory
from core.inference.audio_codecs import AudioCodecManager
from io import StringIO
import logging
import structlog
from loggers import get_logger
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
class InferenceBackend:
"""Unified inference backend supporting text, vision, and LoRA models"""

5
studio/backend/core/inference/llama_cpp.py
View file

@ -9,7 +9,8 @@ through its OpenAI-compatible /v1/chat/completions endpoint.
"""
import atexit
import json
import logging
import structlog
from loggers import get_logger
import shutil
import signal
import socket
@ -21,7 +22,7 @@ from typing import Generator, Optional
import httpx
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
class LlamaCppBackend:

5
studio/backend/core/inference/orchestrator.py
View file

@ -16,7 +16,8 @@ Pattern follows core/training/training.py.
"""
import atexit
import base64
import logging
import structlog
from loggers import get_logger
import multiprocessing as mp
import queue
import threading
@ -26,7 +27,7 @@ from io import BytesIO
from pathlib import Path
from typing import Any, Generator, Optional, Tuple, Union
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")

24
studio/backend/core/inference/worker.py
View file

@ -16,7 +16,8 @@ Pattern follows core/training/worker.py.
from __future__ import annotations
import base64
import logging
import structlog
from loggers import get_logger
import os
import queue as _queue
import sys
@ -25,7 +26,7 @@ import traceback
from io import BytesIO
from typing import Any
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _activate_transformers_version(model_name: str, project_root: str) -> None:
@ -244,6 +245,8 @@ def _handle_generate(
else:
generator = backend.generate_chat_response(**gen_kwargs)
logger.info("Starting text generation for request_id=%s", request_id)
for cumulative_text in generator:
# cancel_event is an mp.Event — checked instantly, no queue polling
if cancel_event.is_set():
@ -262,6 +265,7 @@ def _handle_generate(
"request_id": request_id,
"ts": time.time(),
})
logger.info("Finished text generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Generation error: %s", exc, exc_info=True)
@ -282,6 +286,7 @@ def _handle_generate_audio(
"""Handle TTS audio generation — returns WAV bytes + sample_rate."""
request_id = cmd.get("request_id", "")
try:
logger.info("Starting audio generation for request_id=%s", request_id)
wav_bytes, sample_rate = backend.generate_audio_response(
text=cmd["text"],
temperature=cmd.get("temperature", 0.6),
@ -301,6 +306,7 @@ def _handle_generate_audio(
"sample_rate": sample_rate,
"ts": time.time(),
})
logger.info("Finished audio generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Audio generation error: %s", exc, exc_info=True)
@ -349,6 +355,8 @@ def _handle_generate_audio_input(
cancel_event=cancel_event,
)
logger.info("Starting audio input generation for request_id=%s", request_id)
for text_chunk in generator:
if cancel_event.is_set():
logger.info("Audio input generation cancelled for request %s", request_id)
@ -366,6 +374,7 @@ def _handle_generate_audio_input(
"request_id": request_id,
"ts": time.time(),
})
logger.info("Finished audio input generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Audio input generation error: %s", exc, exc_info=True)
@ -418,6 +427,17 @@ def run_inference_process(
config: Initial configuration dict with model info and project_root.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-inference-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
)
project_root = config["project_root"]
model_name = config["model_name"]

292
studio/backend/core/training/trainer.py
View file

@ -19,7 +19,8 @@ from unsloth.chat_templates import get_chat_template
import json
import threading
import math
import logging
import structlog
from loggers import get_logger
import time
from pathlib import Path
from typing import Optional, Callable
@ -32,8 +33,7 @@ from utils.datasets import format_and_template_dataset
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
from trl import SFTTrainer, SFTConfig
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
_BACKEND_ROOT = Path(__file__).resolve().parents[2]
_ASSETS_DATASETS_ROOT = _BACKEND_ROOT / "assets" / "datasets"
@ -162,7 +162,7 @@ class UnslothTrainer:
def on_step_end(self, args, state, control, **kwargs):
if trainer_ref.should_stop:
print(f"Stop detected at step {state.global_step}\n")
logger.info(f"Stop detected at step {state.global_step}\n")
control.should_training_stop = True
return control
@ -233,21 +233,21 @@ class UnslothTrainer:
self.tokenizer.save_pretrained(output_dir)
self._patch_adapter_config(output_dir)
msg = f"{label} training stopped" if label else "Training stopped"
print(f"\n{msg}. Model saved to {output_dir}\n")
logger.info(f"\n{msg}. Model saved to {output_dir}\n")
self._update_progress(
is_training=False,
status_message=f"Training stopped. Model saved to {output_dir}",
)
elif self.should_stop:
msg = f"{label} training cancelled" if label else "Training cancelled"
print(f"\n{msg}.\n")
logger.info(f"\n{msg}.\n")
self._update_progress(is_training=False, status_message="Training cancelled.")
else:
self.trainer.save_model()
self.tokenizer.save_pretrained(output_dir)
self._patch_adapter_config(output_dir)
msg = f"{label} training completed" if label else "Training completed"
print(f"\n{msg}! Model saved to {output_dir}\n")
logger.info(f"\n{msg}! Model saved to {output_dir}\n")
self._update_progress(
is_training=False,
is_completed=True,
@ -287,7 +287,7 @@ class UnslothTrainer:
del _sys.modules[key]
if removed_paths or removed_modules:
print(f"Cleaned up audio artifacts: {len(removed_paths)} paths, "
logger.info(f"Cleaned up audio artifacts: {len(removed_paths)} paths, "
f"{len(removed_modules)} modules\n")
def _resolve_audio_columns(self, dataset, custom_format_mapping: dict = None):
@ -346,7 +346,7 @@ class UnslothTrainer:
if self.trainer is not None:
del self.trainer
print("\nClearing GPU memory before training...")
logger.info("\nClearing GPU memory before training...")
clear_gpu_cache()
# Clean up sys.path and sys.modules from previous audio preprocessing
@ -413,7 +413,7 @@ class UnslothTrainer:
status_message=f"Loading {model_type_label} model... {model_display}"
)
print(f"\nLoading {model_type_label} model: {model_name}")
logger.info(f"\nLoading {model_type_label} model: {model_name}")
# Set HF token if provided
if hf_token:
@ -571,10 +571,10 @@ class UnslothTrainer:
from transformers import ProcessorMixin
tok = self.tokenizer
has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(tok, "image_processor")
print(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}")
print(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}")
print(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}")
print(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n")
logger.info(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}")
logger.info(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}")
logger.info(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}")
logger.info(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n")
else:
# Load text model - returns (model, tokenizer)
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
@ -591,7 +591,7 @@ class UnslothTrainer:
return False
self._update_progress(status_message="Model loaded successfully")
print("Model loaded successfully")
logger.info("Model loaded successfully")
return True
except OSError as e:
@ -602,7 +602,7 @@ class UnslothTrainer:
# second attempt because the failed first call's partial
# imports clean up the stale state as a side effect.
self._source_code_retried = True
print(f"\n'could not get source code' — retrying once...\n")
logger.info(f"\n'could not get source code' — retrying once...\n")
return self.load_model(model_name, max_seq_length, load_in_4bit, hf_token,
is_dataset_image, is_dataset_audio, trust_remote_code)
error_msg = str(e)
@ -656,7 +656,7 @@ class UnslothTrainer:
# Full finetuning mode - skip PEFT entirely
if not use_lora:
self._update_progress(status_message="Full finetuning mode - no LoRA adapters")
print("Full finetuning mode - training all parameters\n")
logger.info("Full finetuning mode - training all parameters\n")
return True
# LoRA/QLoRA mode - apply PEFT
@ -703,22 +703,22 @@ class UnslothTrainer:
self._update_progress(error=error_msg)
return False
print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
print(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n")
logger.info(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
logger.info(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n")
# Branch based on model type: audio, audio_vlm, vision, or text
if self._audio_type in ('csm', 'bicodec', 'dac') or self.is_audio_vlm:
# Models using FastModel.get_peft_model (codec audio + audio VLM)
from unsloth import FastModel
label = self._audio_type or 'audio_vlm'
print(f"{label} LoRA configuration:")
print(f" - Target modules: {target_modules}")
logger.info(f"{label} LoRA configuration:")
logger.info(f" - Target modules: {target_modules}")
if self.is_audio_vlm:
print(f" - Finetune vision layers: {finetune_vision_layers}")
print(f" - Finetune language layers: {finetune_language_layers}")
print(f" - Finetune attention modules: {finetune_attention_modules}")
print(f" - Finetune MLP modules: {finetune_mlp_modules}")
print()
logger.info(f" - Finetune vision layers: {finetune_vision_layers}")
logger.info(f" - Finetune language layers: {finetune_language_layers}")
logger.info(f" - Finetune attention modules: {finetune_attention_modules}")
logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}")
logger.info()
peft_kwargs = dict(
r=lora_r,
@ -745,8 +745,8 @@ class UnslothTrainer:
elif self._audio_type == 'whisper':
# Phase 2: Whisper uses FastModel.get_peft_model with task_type=None
from unsloth import FastModel
print(f"Audio model (whisper) LoRA configuration:")
print(f" - Target modules: {target_modules}\n")
logger.info(f"Audio model (whisper) LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
self.model = FastModel.get_peft_model(
self.model,
@ -764,8 +764,8 @@ class UnslothTrainer:
elif self._audio_type == 'snac':
# Orpheus uses FastLanguageModel.get_peft_model
print(f"Audio model ({self._audio_type}) LoRA configuration:")
print(f" - Target modules: {target_modules}\n")
logger.info(f"Audio model ({self._audio_type}) LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
self.model = FastLanguageModel.get_peft_model(
self.model,
@ -782,11 +782,11 @@ class UnslothTrainer:
elif self.is_vlm:
# Vision model LoRA
print(f"Vision model LoRA configuration:")
print(f" - Finetune vision layers: {finetune_vision_layers}")
print(f" - Finetune language layers: {finetune_language_layers}")
print(f" - Finetune attention modules: {finetune_attention_modules}")
print(f" - Finetune MLP modules: {finetune_mlp_modules}\n")
logger.info(f"Vision model LoRA configuration:")
logger.info(f" - Finetune vision layers: {finetune_vision_layers}")
logger.info(f" - Finetune language layers: {finetune_language_layers}")
logger.info(f" - Finetune attention modules: {finetune_attention_modules}")
logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}\n")
self.model = FastVisionModel.get_peft_model(
self.model,
@ -806,8 +806,8 @@ class UnslothTrainer:
)
else:
# Text model LoRA
print(f"Text model LoRA configuration:")
print(f" - Target modules: {target_modules}\n")
logger.info(f"Text model LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
self.model = FastLanguageModel.get_peft_model(
self.model,
@ -824,11 +824,11 @@ class UnslothTrainer:
# Check if stopped during LoRA preparation
if self.should_stop:
print("Stopped during LoRA configuration\n")
logger.info("Stopped during LoRA configuration\n")
return False
self._update_progress(status_message="LoRA adapters configured")
print("LoRA adapters configured successfully\n")
logger.info("LoRA adapters configured successfully\n")
return True
except Exception as e:
@ -838,8 +838,8 @@ class UnslothTrainer:
full_traceback = traceback.format_exc()
logger.error(f"Error preparing model: {error_details}")
logger.error(f"Full traceback:\n{full_traceback}")
print(f"\n[ERROR] Error preparing model: {error_details}", file=sys.stderr, flush=True)
print(f"[ERROR] Full traceback:\n{full_traceback}", file=sys.stderr, flush=True)
logger.info(f"\n[ERROR] Error preparing model: {error_details}")
logger.info(f"[ERROR] Full traceback:\n{full_traceback}")
self._update_progress(error=error_details)
return False
@ -995,7 +995,7 @@ class UnslothTrainer:
base_csm.forward = types.MethodType(_fixed_csm_forward, base_csm)
# Class-level: catches any path that resolves through the class dict
CsmForConditionalGeneration.forward = _fixed_csm_forward
print("Applied CSM forward fix (class + instance level)\n")
logger.info("Applied CSM forward fix (class + instance level)\n")
def _preprocess_csm_dataset(self, dataset, custom_format_mapping=None):
"""Preprocess dataset for CSM TTS training (exact notebook copy)."""
@ -1024,11 +1024,11 @@ class UnslothTrainer:
if text_col is None:
raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}")
if speaker_key is None:
print("No speaker found, adding default 'source' of 0 for all examples\n")
logger.info("No speaker found, adding default 'source' of 0 for all examples\n")
dataset = dataset.add_column("source", ["0"] * len(dataset))
speaker_key = "source"
print(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n")
logger.info(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n")
dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000))
@ -1039,7 +1039,7 @@ class UnslothTrainer:
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during CSM preprocessing\n")
logger.info("Stopped during CSM preprocessing\n")
break
example = dataset[idx]
@ -1099,7 +1099,7 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
print(f"CSM preprocessing complete: {len(result_dataset)} examples "
logger.info(f"CSM preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
return result_dataset
@ -1147,7 +1147,7 @@ class UnslothTrainer:
self._update_progress(status_message="Formatting audio VLM dataset...")
dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=safe_num_proc(4))
print(f"Audio VLM dataset formatted: {len(dataset)} examples\n")
logger.info(f"Audio VLM dataset formatted: {len(dataset)} examples\n")
return dataset
def _preprocess_snac_dataset(self, dataset, custom_format_mapping=None):
@ -1196,7 +1196,7 @@ class UnslothTrainer:
# Load SNAC codec model
self._update_progress(status_message="Loading SNAC codec model...")
print("Loading SNAC codec model...\n")
logger.info("Loading SNAC codec model...\n")
from snac import SNAC
snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME)
snac_model = snac_model.to(device).eval()
@ -1205,14 +1205,14 @@ class UnslothTrainer:
resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=SNAC_SAMPLE_RATE) if ds_sample_rate != SNAC_SAMPLE_RATE else None
self._update_progress(status_message="Encoding audio with SNAC...")
print(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
logger.info(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n")
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during SNAC preprocessing\n")
logger.info("Stopped during SNAC preprocessing\n")
break
example = dataset[idx]
@ -1300,7 +1300,7 @@ class UnslothTrainer:
)
# Free SNAC model from GPU
print("Freeing SNAC codec model from GPU...\n")
logger.info("Freeing SNAC codec model from GPU...\n")
snac_model.to("cpu")
del snac_model
import gc
@ -1314,7 +1314,7 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
print(f"SNAC preprocessing complete: {len(result_dataset)} examples "
logger.info(f"SNAC preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
return result_dataset
@ -1339,7 +1339,7 @@ class UnslothTrainer:
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
self._update_progress(status_message="Cloning Spark-TTS code repo...")
print(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n")
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n")
subprocess.run(
["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir],
check=True,
@ -1369,13 +1369,13 @@ class UnslothTrainer:
# Load BiCodec tokenizer
self._update_progress(status_message="Loading BiCodec tokenizer...")
print("Loading BiCodec tokenizer...\n")
logger.info("Loading BiCodec tokenizer...\n")
audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device)
target_sr = audio_tokenizer.config['sample_rate']
self._update_progress(status_message="Encoding audio with BiCodec...")
print(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
logger.info(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
f"has_source={has_source}, target_sr={target_sr}\n")
def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor:
@ -1407,7 +1407,7 @@ class UnslothTrainer:
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during BiCodec preprocessing\n")
logger.info("Stopped during BiCodec preprocessing\n")
break
example = dataset[idx]
@ -1492,7 +1492,7 @@ class UnslothTrainer:
)
# Free BiCodec model from GPU
print("Freeing BiCodec tokenizer from GPU...\n")
logger.info("Freeing BiCodec tokenizer from GPU...\n")
audio_tokenizer.model.cpu()
audio_tokenizer.feature_extractor.cpu()
del audio_tokenizer
@ -1507,12 +1507,12 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
print(f"BiCodec preprocessing complete: {len(result_dataset)} examples "
logger.info(f"BiCodec preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
# Debug: show first example text (truncated)
sample = result_dataset[0]["text"]
print(f"Sample text (first 200 chars): {sample[:200]}...\n")
print(f"Sample text length: {len(sample)} chars\n")
logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n")
logger.info(f"Sample text length: {len(sample)} chars\n")
return result_dataset
def _preprocess_dac_dataset(self, dataset, custom_format_mapping=None):
@ -1539,7 +1539,7 @@ class UnslothTrainer:
outetts_pkg = os.path.join(outetts_code_dir, "outetts")
if not os.path.isdir(outetts_pkg):
self._update_progress(status_message="Cloning OuteTTS code repo...")
print(f"Cloning edwko/OuteTTS to {outetts_code_dir}...\n")
logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...\n")
subprocess.run(
["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir],
check=True,
@ -1551,7 +1551,7 @@ class UnslothTrainer:
]:
if os.path.exists(fpath):
os.remove(fpath)
print(f"Removed {fpath}\n")
logger.info(f"Removed {fpath}\n")
if outetts_code_dir not in sys.path:
sys.path.insert(0, outetts_code_dir)
@ -1573,17 +1573,17 @@ class UnslothTrainer:
# Cast audio to 24kHz (notebook: dataset.cast_column("audio", Audio(sampling_rate=24000)))
from datasets import Audio
dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000))
print("Cast audio column to 24kHz\n")
logger.info("Cast audio column to 24kHz\n")
# Load Whisper for word timings
self._update_progress(status_message="Loading Whisper model for word timings...")
print("Loading Whisper model for word timings...\n")
logger.info("Loading Whisper model for word timings...\n")
import whisper
whisper_model = whisper.load_model("turbo", device=device)
# Load OuteTTS AudioProcessor + PromptProcessor
self._update_progress(status_message="Loading OuteTTS AudioProcessor...")
print("Loading OuteTTS AudioProcessor...\n")
logger.info("Loading OuteTTS AudioProcessor...\n")
model_tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B"
dummy_config = OuteTTSModelConfig(
tokenizer_path=model_tokenizer_path,
@ -1594,13 +1594,13 @@ class UnslothTrainer:
prompt_processor = PromptProcessor(model_tokenizer_path)
self._update_progress(status_message="Preprocessing audio with OuteTTS...")
print(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n")
logger.info(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n")
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during DAC preprocessing\n")
logger.info("Stopped during DAC preprocessing\n")
break
example = dataset[idx]
@ -1674,7 +1674,7 @@ class UnslothTrainer:
)
# Free Whisper from GPU (notebook: data_processor.whisper_model.to('cpu'))
print("Moving Whisper model to CPU...\n")
logger.info("Moving Whisper model to CPU...\n")
whisper_model.to('cpu')
del whisper_model
del audio_processor
@ -1690,10 +1690,10 @@ class UnslothTrainer:
)
result_dataset = HFDataset.from_list(processed_examples)
print(f"DAC preprocessing complete: {len(result_dataset)} examples "
logger.info(f"DAC preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
sample = result_dataset[0]["text"]
print(f"Sample text (first 200 chars): {sample[:200]}...\n")
logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n")
return result_dataset
def _preprocess_whisper_dataset(self, dataset, eval_split=None, custom_format_mapping=None):
@ -1726,7 +1726,7 @@ class UnslothTrainer:
eval_dataset_raw = splits["test"]
self._update_progress(status_message="Processing audio for Whisper...")
print(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
logger.info(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
f"samples={len(dataset)}\n")
def process_split(ds, split_name="train"):
@ -1734,7 +1734,7 @@ class UnslothTrainer:
skipped = 0
for idx in range(len(ds)):
if self.should_stop:
print(f"Stopped during Whisper {split_name} preprocessing\n")
logger.info(f"Stopped during Whisper {split_name} preprocessing\n")
break
example = ds[idx]
@ -1766,7 +1766,7 @@ class UnslothTrainer:
status_message=f"Processing {split_name} audio... {idx + 1}/{len(ds)}"
)
print(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n")
logger.info(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n")
return processed
train_data = process_split(dataset, "train")
@ -1857,12 +1857,12 @@ class UnslothTrainer:
# Check if stopped during dataset loading
if self.should_stop:
print("Stopped during dataset loading\n")
logger.info("Stopped during dataset loading\n")
return None
self._update_progress(status_message=f"Loaded {len(dataset)} samples from local files")
print(f"Loaded {len(dataset)} samples from local files\n")
print(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n")
logger.info(f"Loaded {len(dataset)} samples from local files\n")
logger.info(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n")
elif dataset_source:
# Load from Hugging Face
@ -1878,14 +1878,14 @@ class UnslothTrainer:
# Manual slice — stream only the rows we need instead of
# downloading the entire dataset.
rows_to_stream = dataset_slice_end + 1
print(
logger.info(
f"[dataset-slice] Manual slice specified "
f"(start={dataset_slice_start}, end={dataset_slice_end}), "
f"streaming {rows_to_stream} rows\n"
)
stream = load_dataset(**load_kwargs, streaming=True)
dataset = Dataset.from_list(list(stream.take(rows_to_stream)))
print(
logger.info(
f"[dataset-slice] Downloaded {len(dataset)} rows "
f"(requested {rows_to_stream})\n"
)
@ -1897,27 +1897,27 @@ class UnslothTrainer:
# Check if stopped during dataset loading
if self.should_stop:
print("Stopped during dataset loading\n")
logger.info("Stopped during dataset loading\n")
return None
self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
print(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n")
logger.info(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n")
# Resolve eval split from a separate HF split (explicit or auto-detected)
if eval_enabled:
effective_train = train_split or "train"
if eval_split and eval_split != effective_train:
# Explicit eval split provided - load it directly
print(f"Loading explicit eval split: '{eval_split}'\n")
logger.info(f"Loading explicit eval split: '{eval_split}'\n")
eval_load_kwargs = {"path": dataset_source, "split": eval_split}
if subset:
eval_load_kwargs["name"] = subset
eval_dataset = load_dataset(**eval_load_kwargs)
has_separate_eval_source = True
print(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n")
logger.info(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n")
elif eval_split and eval_split == effective_train:
# Same split as training — will do 80/20 split after formatting
print(f"Eval split '{eval_split}' is the same as train split — will split 80/20\n")
logger.info(f"Eval split '{eval_split}' is the same as train split — will split 80/20\n")
else:
# Auto-detect eval split from HF (returns a separate dataset, or None)
eval_dataset = self._auto_detect_eval_split_from_hf(
@ -1927,7 +1927,7 @@ class UnslothTrainer:
if eval_dataset is not None:
has_separate_eval_source = True
else:
print("Eval disabled (eval_steps <= 0), skipping eval split detection\n")
logger.info("Eval disabled (eval_steps <= 0), skipping eval split detection\n")
if dataset is None:
raise ValueError("No dataset provided")
@ -1941,12 +1941,12 @@ class UnslothTrainer:
start = max(0, min(start, total_rows - 1))
end = max(start, min(end, total_rows - 1))
dataset = dataset.select(range(start, end + 1))
print(f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n")
logger.info(f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n")
self._update_progress(status_message=f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})")
# Check if stopped before applying template
if self.should_stop:
print("Stopped before applying chat template\n")
logger.info("Stopped before applying chat template\n")
return None
# ========== AUDIO MODELS: custom preprocessing ==========
@ -1977,7 +1977,7 @@ class UnslothTrainer:
return (formatted, None)
# ========== FORMAT FIRST ==========
print(f"Formatting dataset with format_type='{format_type}'...\n")
logger.info(f"Formatting dataset with format_type='{format_type}'...\n")
dataset_info = format_and_template_dataset(
dataset,
@ -1992,7 +1992,7 @@ class UnslothTrainer:
# Check if stopped during formatting
if self.should_stop:
print("Stopped during dataset formatting\n")
logger.info("Stopped during dataset formatting\n")
return None
# Abort if dataset formatting/conversion failed
@ -2004,12 +2004,12 @@ class UnslothTrainer:
return None
self._update_progress(status_message=f"Dataset formatted and ready for training")
print(f"Dataset formatted successfully\n")
logger.info(f"Dataset formatted successfully\n")
# ========== THEN SPLIT ==========
if has_separate_eval_source and eval_dataset is not None:
# Eval came from a separate HF split — format it too
print(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n")
logger.info(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n")
eval_info = format_and_template_dataset(
eval_dataset,
model_name=self.model_name,
@ -2020,7 +2020,7 @@ class UnslothTrainer:
custom_format_mapping=custom_format_mapping,
)
eval_dataset = eval_info["dataset"]
print(f"Eval dataset formatted successfully\n")
logger.info(f"Eval dataset formatted successfully\n")
elif eval_enabled and not has_separate_eval_source:
# No separate eval source — split the already-formatted dataset
formatted_dataset = dataset_info["dataset"]
@ -2045,7 +2045,7 @@ class UnslothTrainer:
if subset:
load_kwargs["config_name"] = subset
available_splits = get_dataset_split_names(**load_kwargs)
print(f"Available splits: {available_splits}\n")
logger.info(f"Available splits: {available_splits}\n")
# Check for common eval split names
for candidate in ["eval", "validation", "valid", "val", "test"]:
@ -2055,10 +2055,10 @@ class UnslothTrainer:
eval_load_kwargs["name"] = subset
candidate_ds = load_dataset(**eval_load_kwargs)
if len(candidate_ds) >= 16:
print(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
logger.info(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
return candidate_ds
else:
print(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n")
logger.info(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n")
except Exception as e:
logger.warning(f"Could not check dataset splits: {e}")
@ -2077,16 +2077,16 @@ class UnslothTrainer:
n = len(dataset)
if n < MIN_TOTAL_ROWS:
print(f"Dataset too small ({n} rows) for eval split, skipping eval\n")
logger.info(f"Dataset too small ({n} rows) for eval split, skipping eval\n")
return None
eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * n)))
# Ensure we don't take more than half the dataset
eval_size = min(eval_size, n // 2)
print(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
logger.info(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
split_result = dataset.train_test_split(test_size=eval_size, seed=3407)
print(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n")
logger.info(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n")
return (split_result['train'], split_result['test'])
def start_training(self,
@ -2225,7 +2225,7 @@ class UnslothTrainer:
training_args.get('max_steps', 0),
)
self._update_progress(total_steps=total, status_message="Starting CSM training...")
print(f"CSM training config: {config}\n")
logger.info(f"CSM training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "CSM")
return
@ -2254,7 +2254,7 @@ class UnslothTrainer:
training_args.get('max_steps', 0),
)
self._update_progress(total_steps=total, status_message="Starting SNAC training...")
print(f"SNAC training config: {config}\n")
logger.info(f"SNAC training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "SNAC")
return
@ -2293,7 +2293,7 @@ class UnslothTrainer:
training_args.get('max_steps', 0),
)
self._update_progress(total_steps=total, status_message="Starting Whisper training...")
print(f"Whisper training config: {config}\n")
logger.info(f"Whisper training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "Whisper")
return
@ -2307,12 +2307,12 @@ class UnslothTrainer:
model_name_lower = self.model_name.lower()
is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower
print("Configuring data collator...\n")
logger.info("Configuring data collator...\n")
data_collator = None # Default to built-in data collator
if is_deepseek_ocr:
# Special DeepSeek OCR collator - auto-install if needed
print("Detected DeepSeek OCR model\n")
logger.info("Detected DeepSeek OCR model\n")
# Ensure DeepSeek OCR module is installed
if not _ensure_deepseek_ocr_installed():
error_msg = (
@ -2328,7 +2328,7 @@ class UnslothTrainer:
try:
from backend.data_utils import DeepSeekOCRDataCollator
print("Configuring DeepSeek OCR data collator...\n")
logger.info("Configuring DeepSeek OCR data collator...\n")
FastVisionModel.for_training(self.model)
data_collator = DeepSeekOCRDataCollator(
tokenizer=self.tokenizer,
@ -2338,7 +2338,7 @@ class UnslothTrainer:
crop_mode=True,
train_on_responses_only=training_args.get('train_on_completions', False),
)
print("DeepSeek OCR data collator configured successfully\n")
logger.info("DeepSeek OCR data collator configured successfully\n")
except Exception as e:
logger.error(f"Failed to configure DeepSeek OCR collator: {e}")
@ -2349,7 +2349,7 @@ class UnslothTrainer:
elif self.is_audio_vlm:
# Audio VLM collator (e.g. Gemma 3N with audio data)
# Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook
print("Configuring audio VLM data collator...\n")
logger.info("Configuring audio VLM data collator...\n")
processor = self.tokenizer # FastModel returns processor as tokenizer
audio_col_name = getattr(self, '_audio_vlm_audio_col', 'audio')
@ -2379,16 +2379,16 @@ class UnslothTrainer:
return batch
data_collator = audio_vlm_collate_fn
print("Audio VLM data collator configured\n")
logger.info("Audio VLM data collator configured\n")
elif self.is_vlm:
# Standard VLM collator (images)
print("Using UnslothVisionDataCollator for vision model\n")
logger.info("Using UnslothVisionDataCollator for vision model\n")
from unsloth.trainer import UnslothVisionDataCollator
FastVisionModel.for_training(self.model)
data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
print("Vision data collator configured\n")
logger.info("Vision data collator configured\n")
# ========== TRAINING CONFIGURATION ==========
# Handle warmup_steps vs warmup_ratio
@ -2396,7 +2396,7 @@ class UnslothTrainer:
warmup_ratio_val = training_args.get('warmup_ratio', None)
lr_value = training_args.get('learning_rate', 2e-4)
print(f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n")
logger.info(f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n")
config_args = {
"per_device_train_batch_size": training_args.get('batch_size', 2),
@ -2414,7 +2414,7 @@ class UnslothTrainer:
"dataset_num_proc": 1 if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) else safe_num_proc(max(1, os.cpu_count() // 4)),
"max_seq_length": training_args.get('max_seq_length', 2048),
}
print(f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})")
logger.info(f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})")
# On Windows with transformers 5.x, disable DataLoader multiprocessing
# to avoid issues with modified sys.path (.venv_t5) in spawned workers.
@ -2426,14 +2426,14 @@ class UnslothTrainer:
# Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
if warmup_ratio_val is not None:
config_args["warmup_ratio"] = warmup_ratio_val
print(f"Using warmup_ratio: {warmup_ratio_val}\n")
logger.info(f"Using warmup_ratio: {warmup_ratio_val}\n")
elif warmup_steps_val is not None:
config_args["warmup_steps"] = warmup_steps_val
print(f"Using warmup_steps: {warmup_steps_val}\n")
logger.info(f"Using warmup_steps: {warmup_steps_val}\n")
else:
# Default to warmup_steps if neither provided
config_args["warmup_steps"] = 5
print(f"Using default warmup_steps: 5\n")
logger.info(f"Using default warmup_steps: 5\n")
# Add save_steps if specified
save_steps_val = training_args.get('save_steps', 0)
@ -2446,9 +2446,9 @@ class UnslothTrainer:
if max_steps_val and max_steps_val > 0:
del config_args["num_train_epochs"] # Remove epochs
config_args["max_steps"] = max_steps_val # Use steps instead
print(f"Training for {max_steps_val} steps\n")
logger.info(f"Training for {max_steps_val} steps\n")
else:
print(f"Training for {config_args['num_train_epochs']} epochs\n")
logger.info(f"Training for {config_args['num_train_epochs']} epochs\n")
# ========== EVAL CONFIGURATION ==========
eval_dataset = training_args.get('eval_dataset', None)
@ -2457,13 +2457,13 @@ class UnslothTrainer:
if eval_steps_val > 0:
config_args["eval_strategy"] = "steps"
config_args["eval_steps"] = eval_steps_val
print(f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
print(f"Eval dataset: {len(eval_dataset)} rows\n")
logger.info(f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
logger.info(f"Eval dataset: {len(eval_dataset)} rows\n")
else:
print(f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n")
print("To enable evaluation, set eval_steps > 0.0\n")
logger.info(f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n")
logger.info("To enable evaluation, set eval_steps > 0.0\n")
else:
print("No eval dataset — evaluation disabled\n")
logger.info("No eval dataset — evaluation disabled\n")
# Add model-specific parameters
# Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults
@ -2473,7 +2473,7 @@ class UnslothTrainer:
if self.is_vlm or self.is_audio_vlm:
# Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns)
label = "audio VLM" if self.is_audio_vlm else "vision"
print(f"Configuring {label} model training parameters\n")
logger.info(f"Configuring {label} model training parameters\n")
# Use provided values or defaults for vision models
optim_value = training_args.get('optim', "adamw_torch_fused")
lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine")
@ -2489,7 +2489,7 @@ class UnslothTrainer:
"max_length": training_args.get('max_seq_length', 2048),
})
else:
print("Configuring text model training parameters\n")
logger.info("Configuring text model training parameters\n")
config_args.update({
"optim": optim_value,
"lr_scheduler_type": lr_scheduler_type_value,
@ -2500,19 +2500,19 @@ class UnslothTrainer:
if not is_deepseek_ocr:
packing_enabled = training_args.get('packing', False)
config_args["packing"] = packing_enabled
print(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n")
logger.info(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n")
# Audio codec overrides — BiCodec/DAC use the text SFTTrainer path
if self._audio_type == 'bicodec':
config_args["packing"] = False
print("Applied BiCodec overrides: packing=False\n")
logger.info("Applied BiCodec overrides: packing=False\n")
elif self._audio_type == 'dac':
config_args["packing"] = False
print("Applied DAC overrides: packing=False\n")
logger.info("Applied DAC overrides: packing=False\n")
print(f"The configuration is: {config_args}")
logger.info(f"The configuration is: {config_args}")
print("Training configuration prepared\n")
logger.info("Training configuration prepared\n")
# ========== TRAINER INITIALIZATION ==========
if self.is_audio_vlm:
# Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset
@ -2551,7 +2551,7 @@ class UnslothTrainer:
from transformers import ProcessorMixin
sft_tokenizer = self.tokenizer
if isinstance(self.tokenizer, ProcessorMixin) and hasattr(self.tokenizer, 'tokenizer'):
print(f" ⚠️ Unwrapping Processor → raw tokenizer for text-only SFTTrainer")
logger.info(f" ⚠️ Unwrapping Processor → raw tokenizer for text-only SFTTrainer")
sft_tokenizer = self.tokenizer.tokenizer
trainer_kwargs = {
@ -2568,7 +2568,7 @@ class UnslothTrainer:
# saves include preprocessor_config.json (needed for GGUF export).
if sft_tokenizer is not self.tokenizer:
self.trainer.processing_class = self.tokenizer
print("Trainer initialized\n")
logger.info("Trainer initialized\n")
# ========== TRAIN ON RESPONSES ONLY ==========
# Determine if we should train on responses only
@ -2580,26 +2580,26 @@ class UnslothTrainer:
# Audio VLM handles label masking in its collator, so skip
if train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
try:
print("Configuring train on responses only...\n")
logger.info("Configuring train on responses only...\n")
# Get the template mapping for this model
model_name_lower = self.model_name.lower()
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
print(f"Detected template: {template_name}\n")
logger.info(f"Detected template: {template_name}\n")
if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["instruction"]
response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
print(f"Instruction marker: {instruction_part[:50]}...\n")
print(f"Response marker: {response_part[:50]}...\n")
logger.info(f"Instruction marker: {instruction_part[:50]}...\n")
logger.info(f"Response marker: {response_part[:50]}...\n")
else:
print(f"No response mapping found for template: {template_name}\n")
logger.info(f"No response mapping found for template: {template_name}\n")
train_on_responses_enabled = False
else:
print(f"No template mapping found for model: {self.model_name}\n")
logger.info(f"No template mapping found for model: {self.model_name}\n")
train_on_responses_enabled = False
except Exception as e:
@ -2617,7 +2617,7 @@ class UnslothTrainer:
response_part=response_part,
num_proc=config_args["dataset_num_proc"],
)
print("Train on responses only configured successfully\n")
logger.info("Train on responses only configured successfully\n")
# ── Safety net: check if all samples were filtered out ──
# Unsloth's train_on_responses_only masks non-response
@ -2646,21 +2646,21 @@ class UnslothTrainer:
return
if dropped > 0:
print(
logger.info(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
print(f"Post-filter dataset size: {filtered_len} samples\n")
logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
train_on_responses_enabled = False
else:
if train_on_responses_enabled and is_deepseek_ocr:
print("Train on responses handled by DeepSeek OCR collator\n")
logger.info("Train on responses handled by DeepSeek OCR collator\n")
else:
print("Training on full sequences (including prompts)\n")
logger.info("Training on full sequences (including prompts)\n")
# ========== PROGRESS TRACKING ==========
self.trainer.add_callback(self._create_progress_callback())
@ -2677,7 +2677,7 @@ class UnslothTrainer:
# ========== START TRAINING ==========
self._update_progress(status_message="Starting training...")
print("Starting training...\n")
logger.info("Starting training...\n")
self.trainer.train()
# ========== SAVE MODEL ==========
@ -2724,7 +2724,7 @@ class UnslothTrainer:
def stop_training(self, save: bool = True):
"""Stop ongoing training"""
print(f"\nStopping training (save={save})...")
logger.info(f"\nStopping training (save={save})...")
self.should_stop = True
self.save_on_stop = save
stop_msg = (
@ -2738,7 +2738,7 @@ class UnslothTrainer:
if self.trainer:
try:
# The callback will catch should_stop flag and stop the training loop
print("Training will stop at next step...\n")
logger.info("Training will stop at next step...\n")
except Exception as e:
logger.error(f"Error stopping trainer: {e}")
@ -2778,7 +2778,7 @@ def _ensure_deepseek_ocr_installed():
try:
logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...")
print("\n Downloading DeepSeek OCR module from HuggingFace...\n")
logger.info("\n Downloading DeepSeek OCR module from HuggingFace...\n")
from huggingface_hub import snapshot_download
import sys
@ -2805,12 +2805,12 @@ def _ensure_deepseek_ocr_installed():
from deepseek_ocr.modeling_deepseekocr import format_messages
logger.info("DeepSeek OCR module installed successfully")
print("DeepSeek OCR module installed successfully!\n")
logger.info("DeepSeek OCR module installed successfully!\n")
return True
except Exception as e:
logger.error(f"Failed to install DeepSeek OCR module: {e}")
print(f"\n❌ Failed to install DeepSeek OCR module: {e}\n")
logger.info(f"\n❌ Failed to install DeepSeek OCR module: {e}\n")
return False
# Global trainer instance

5
studio/backend/core/training/training.py
View file

@ -18,14 +18,15 @@ import multiprocessing as mp
import queue
import threading
import time
import logging
import structlog
from loggers import get_logger
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Tuple, Any
import matplotlib.pyplot as plt
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")

16
studio/backend/core/training/worker.py
View file

@ -12,7 +12,8 @@ Pattern follows core/data_recipe/jobs/worker.py.
"""
from __future__ import annotations
import logging
import structlog
from loggers import get_logger
import os
import sys
import time
@ -20,7 +21,7 @@ import traceback
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _activate_transformers_version(model_name: str, project_root: str) -> None:
@ -84,6 +85,17 @@ def run_training_process(
config: Training configuration dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-training-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
)
project_root = config["project_root"]
model_name = config["model_name"]

3
studio/backend/loggers/__init__.py
View file

@ -0,0 +1,3 @@
from .handlers import get_logger
__all__ = ["get_logger"]

72
studio/backend/loggers/config.py Normal file
View file

@ -0,0 +1,72 @@
"""Logging configuration for structured logging with structlog.
This module provides centralized logging configuration with environment-specific
formats and processors. Supports both development and production environments
with consistent structured logging.
Key Features:
- Environment-specific formatting (JSON for production, console for development)
- Timestamp standardization (ISO format)
- Context variable integration
- Log level filtering
- Logger caching for performance
"""
import logging
import os
import sys
from typing import Optional
import structlog
class LogConfig:
"""Structured logging configuration for the application.
Provides static method to configure structlog with environment-specific
formatting and processors for consistent structured logging.
"""
@staticmethod
def setup_logging(
service_name: str = "unsloth-studio-backend", env: Optional[str] = None
) -> structlog.BoundLogger:
"""Configure structured logging for the application.
Args:
service_name: Name of the service for logging identification
env: Environment (development/production), affects logging format
"""
# Determine log level from environment
log_level_name = os.getenv("LOG_LEVEL", "INFO").upper()
# Fallback to INFO if an invalid level is provided
log_level = getattr(logging, log_level_name, logging.INFO)
structlog.configure(
processors=[
# Reorder processors to control field order
structlog.processors.TimeStamper(fmt="iso"), # timestamp first
structlog.processors.add_log_level, # level second
structlog.contextvars.merge_contextvars,
# Custom processor to flatten the extra field
lambda logger, method_name, event_dict: {
"timestamp": event_dict.get("timestamp"),
"level": event_dict.get("level"),
"event": event_dict.get("event"),
**(event_dict.get("extra", {})), # Flatten extra into main dict
**{
k: v
for k, v in event_dict.items()
if k not in ["timestamp", "level", "event", "extra"]
},
},
(
structlog.processors.JSONRenderer(sort_keys=False) # Preserve order
if env == "production"
else structlog.dev.ConsoleRenderer()
),
],
wrapper_class=structlog.make_filtering_bound_logger(log_level),
logger_factory=structlog.PrintLoggerFactory(file=sys.stdout),
cache_logger_on_first_use=True,
)
return structlog.get_logger(service_name)

96
studio/backend/loggers/handlers.py Normal file
View file

@ -0,0 +1,96 @@
"""Logging handlers and middleware for structured logging.
This module provides FastAPI middleware and structlog processors for:
- Request/response logging with timing
- Sensitive data filtering in logs
- Structured logging configuration
- Error handling with detailed context
Key Components:
- LoggingMiddleware: FastAPI middleware for request/response logging
- filter_sensitive_data: Structlog processor for data sanitization
- get_logger: Factory function for structured loggers
"""
import time
from typing import Callable
import structlog
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
logger = structlog.get_logger(__name__)
class LoggingMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: Callable) -> Response:
start_time = time.time()
try:
response = await call_next(request)
# Log response
process_time = (time.time() - start_time) * 1000
EXCLUDED_PATHS = {
"/api/train/status",
"/api/train/metrics",
"/api/train/hardware",
"/api/system"
}
is_excluded = (
request.url.path in EXCLUDED_PATHS
or request.url.path.startswith("/assets/")
or request.url.path.endswith((".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf"))
)
if not is_excluded:
logger.info(
"request_completed",
method=request.method,
path=request.url.path,
status_code=response.status_code,
process_time_ms=round(process_time, 2),
)
return response
except Exception as e:
logger.error(
"request_failed",
path=request.url.path,
method=request.method,
error=str(e),
exc_info=True,
)
raise
def filter_sensitive_data(logger, method_name, event_dict):
"""Structlog processor to filter out base64 data from logs."""
def filter_value(value):
if (
isinstance(value, str)
and len(value) > 100
and ("," in value or "/" in value)
):
# Likely base64 data, truncate it
return value[:20] + "..."
elif isinstance(value, dict):
return {k: filter_value(v) for k, v in value.items()}
elif isinstance(value, list):
return [filter_value(item) for item in value]
return value
return {k: filter_value(v) for k, v in event_dict.items()}
def get_logger(name: str) -> structlog.BoundLogger:
"""Get a logger instance for a specific module.
Args:
name: Usually __name__ of the module
Returns:
A bound structured logger
"""
return structlog.get_logger(name)

27
studio/backend/main.py
View file

@ -5,10 +5,21 @@
Main FastAPI application for Unsloth UI Backend
"""
import os
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
import secrets
import shutil
import warnings
from contextlib import asynccontextmanager
# Suppress annoying dependency warnings in production
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
# Alternatively, you can be more specific:
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", module="triton.*")
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
@ -55,8 +66,9 @@ async def lifespan(app: FastAPI):
sm_version = props.major * 10 + props.minor
if sm_version >= 120:
os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
import logging
logging.getLogger(__name__).info(
import structlog
from loggers import get_logger
get_logger(__name__).info(
f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0"
)
@ -83,6 +95,17 @@ app = FastAPI(
lifespan=lifespan,
)
# Initialize structured logging
from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
logger = LogConfig.setup_logging(
service_name="unsloth-studio-backend",
env=os.getenv("ENVIRONMENT_TYPE", "production")
)
app.add_middleware(LoggingMiddleware)
# CORS middleware
app.add_middleware(
CORSMiddleware,

3
studio/backend/requirements/studio.txt
View file

@ -11,4 +11,5 @@ pyjwt
easydict
addict
gradio>=4.0.0
huggingface-hub==0.36.2
huggingface-hub==0.36.2
structlog>=24.1.0

14
studio/backend/routes/datasets.py
View file

@ -11,7 +11,8 @@ import sys
from pathlib import Path
from uuid import uuid4
from fastapi import APIRouter, Depends, HTTPException, UploadFile
import logging
import structlog
from loggers import get_logger
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
@ -23,16 +24,9 @@ from utils.datasets import check_dataset_format
from auth.authentication import get_current_subject
router = APIRouter()
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Configure logger
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
from models.datasets import (

14
studio/backend/routes/export.py
View file

@ -8,7 +8,8 @@ Export API routes: checkpoint discovery and model export operations.
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
import logging
import structlog
from loggers import get_logger
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
@ -39,16 +40,9 @@ from models import (
)
router = APIRouter()
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Configure logger
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)

14
studio/backend/routes/inference.py
View file

@ -12,7 +12,8 @@ from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from typing import Optional
import json
import logging
import structlog
from loggers import get_logger
import asyncio
import threading
@ -65,16 +66,9 @@ import base64
import numpy as np
router = APIRouter()
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Configure logger
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
# GGUF inference backend (llama-server)
_llama_cpp_backend = LlamaCppBackend()

14
studio/backend/routes/models.py
View file

@ -8,7 +8,8 @@ import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
from typing import List, Optional
import logging
import structlog
from loggers import get_logger
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
@ -66,7 +67,7 @@ from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType
from models.responses import LoRABaseModelResponse, VisionCheckResponse, EmbeddingCheckResponse
router = APIRouter()
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding: bool = False) -> ModelType:
@ -79,14 +80,7 @@ def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding:
return "vision"
return "text"
# Configure logger
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
def _resolve_hf_cache_dir() -> Path:

14
studio/backend/routes/training.py
View file

@ -9,7 +9,8 @@ from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from typing import Dict, Optional, Any
import logging
import structlog
from loggers import get_logger
import asyncio
from datetime import datetime
@ -48,16 +49,9 @@ class TrainingStopRequest(PydanticBaseModel):
save: bool = True
router = APIRouter()
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Configure logger
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
@router.get("/hardware")

10
studio/backend/run.py
View file

@ -5,7 +5,12 @@
Run script for Unsloth UI Backend.
Works independently and can be moved to any directory.
"""
import os
import sys
# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
os.environ["PYTHONWARNINGS"] = "ignore"
from pathlib import Path
# Add the backend directory to Python path
@ -13,6 +18,9 @@ backend_dir = Path(__file__).parent
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
from loggers import get_logger
logger = get_logger(__name__)
def _resolve_external_ip() -> str:
"""
@ -96,7 +104,7 @@ def run_server(
# Run server
def _run():
config = uvicorn.Config(app, host=host, port=port, log_level="info")
config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=False)
server = uvicorn.Server(config)
asyncio.run(server.serve())

6
studio/backend/tests/test_utils.py
View file

@ -273,7 +273,8 @@ class TestLogGpuMemory:
"utilization_pct": 12.5,
"free_gb": 14.0,
}
import logging
import structlog
from loggers import get_logger
with patch("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")
@ -284,7 +285,8 @@ class TestLogGpuMemory:
def test_logs_cpu_fallback_when_no_gpu(self, caplog):
fake_info = {"available": False, "backend": "cpu"}
import logging
import structlog
from loggers import get_logger
with patch("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")

5
studio/backend/utils/cache_cleanup.py
View file

@ -9,10 +9,11 @@ FastModel.from_pretrained() and contains model-type-specific compiled Python
files. It should be cleared between model loads to avoid stale artefacts.
"""
import shutil
import logging
import structlog
from loggers import get_logger
from pathlib import Path
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Possible locations where unsloth_compiled_cache may appear
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend

22
studio/backend/utils/datasets/chat_templates.py
View file

@ -12,6 +12,10 @@ from torch.utils.data import IterableDataset
from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic
from .model_mappings import MODEL_TO_TEMPLATE_MAPPER
from loggers import get_logger
logger = get_logger(__name__)
DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@ -53,15 +57,15 @@ def get_tokenizer_chat_template(tokenizer, model_name):
# Direct match in MODEL_TO_TEMPLATE_MAPPER
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
matched_template = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
print(f"📝 Applying Unsloth chat template: {matched_template}")
logger.info(f"📝 Applying Unsloth chat template: {matched_template}")
try:
tokenizer = get_chat_template(
tokenizer,
chat_template=matched_template,
)
except Exception as e:
print(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
print(f" Falling back to tokenizer's default chat template")
logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
logger.info(f" Falling back to tokenizer's default chat template")
else:
# Check if tokenizer actually has a chat_template set
has_chat_template = (
@ -69,18 +73,18 @@ def get_tokenizer_chat_template(tokenizer, model_name):
and tokenizer.chat_template is not None
)
if has_chat_template:
print(f"📝 Using tokenizer's own chat template (no Unsloth template match)")
logger.info(f"📝 Using tokenizer's own chat template (no Unsloth template match)")
else:
# Base model with no chat template — apply default ChatML
print(f"📝 No chat template found — applying default ChatML template (base model)")
logger.info(f"📝 No chat template found — applying default ChatML template (base model)")
try:
tokenizer = get_chat_template(
tokenizer,
chat_template="chatml",
)
except Exception as e:
print(f"⚠️ Failed to apply default ChatML template: {e}")
print(f" Falling back to tokenizer as-is")
logger.info(f"⚠️ Failed to apply default ChatML template: {e}")
logger.info(f" Falling back to tokenizer as-is")
return tokenizer
@ -253,9 +257,9 @@ def apply_chat_template_to_dataset(
try:
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
print(f"📝 Set alpaca chat template on tokenizer for model saving")
logger.info(f"📝 Set alpaca chat template on tokenizer for model saving")
except Exception as e:
print(f"⚠️ Could not set alpaca template on tokenizer: {e}")
logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")
# Use custom template if provided
def _format_alpaca_custom(examples):

6
studio/backend/utils/datasets/data_collators.py
View file

@ -11,6 +11,10 @@ particularly for VLM/OCR processing.
import torch
from dataclasses import dataclass
from typing import Any, List, Optional, Union
from loggers import get_logger
logger = get_logger(__name__)
@dataclass
@ -116,7 +120,7 @@ class DeepSeekOCRDataCollator:
return inputs
except Exception as e:
print(f"⚠️ DeepSeekOCRDataCollator error: {e}")
logger.info(f"⚠️ DeepSeekOCRDataCollator error: {e}")
raise

6
studio/backend/utils/datasets/dataset_utils.py
View file

@ -42,6 +42,10 @@ from .chat_templates import (
from .vlm_processing import generate_smart_vlm_instruction
from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
from loggers import get_logger
logger = get_logger(__name__)
def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
@ -696,7 +700,7 @@ def format_and_template_dataset(
f"text='{user_vlm_text_column}') failed: {e} — "
f"falling back to auto-detection"
)
print(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
custom_format_mapping = None # clear so auto-detection runs below
else:
errors.append(

60
studio/backend/utils/datasets/format_conversion.py
View file

@ -11,6 +11,10 @@ This module contains functions for converting between dataset formats
import os
from datasets import IterableDataset
from loggers import get_logger
logger = get_logger(__name__)
def standardize_chat_format(
@ -301,12 +305,12 @@ def convert_to_vlm_format(
instruction_column = instruction_info.get("instruction_column")
uses_dynamic = instruction_info["uses_dynamic_instruction"]
print(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}")
print(f"📝 Confidence: {instruction_info['confidence']:.2f}")
logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}")
logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}")
if not uses_dynamic:
print(f"📝 Using instruction: '{instruction}'")
logger.info(f"📝 Using instruction: '{instruction}'")
else:
print(f"📝 Using dynamic instructions from column: '{instruction_column}'")
logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'")
else:
instruction_column = None
uses_dynamic = False
@ -382,7 +386,7 @@ def convert_to_vlm_format(
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
print(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
_image_lookup = {
os.path.basename(f): f
@ -390,12 +394,12 @@ def convert_to_vlm_format(
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS)
}
if first_image in _image_lookup:
print(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
else:
print(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
_image_lookup = None
except Exception as e:
print(f"⚠️ Failed to build HF repo image lookup: {e}")
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
_image_lookup = None
# ── URL probe: 200 samples with parallel workers to estimate speed + failure rate ──
@ -409,7 +413,7 @@ def convert_to_vlm_format(
num_workers = safe_num_proc()
_notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...")
print(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
probe_samples = [dataset[i] for i in range(PROBE_SIZE)]
probe_ok = 0
@ -437,7 +441,7 @@ def convert_to_vlm_format(
"This dataset has too many broken or unreachable image URLs. "
"Consider using a dataset with embedded images instead."
)
print(msg)
logger.info(msg)
_notify(msg)
raise ValueError(msg)
@ -453,14 +457,14 @@ def convert_to_vlm_format(
if probe_fail > 0:
info_msg += f" | {fail_rate:.0%} broken URLs will be skipped"
print(f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s")
print(f"⏱️ Estimated time for {total:,} samples: ~{eta_str}")
logger.info(f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s")
logger.info(f"⏱️ Estimated time for {total:,} samples: ~{eta_str}")
_notify(info_msg)
# ── Full conversion with progress ──
from tqdm import tqdm
print(f"🔄 Converting {total} samples to VLM format...")
logger.info(f"🔄 Converting {total} samples to VLM format...")
converted_list = []
failed_count = 0
@ -488,7 +492,7 @@ def convert_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
converted_list.extend(r for r in batch_results if r is not None)
@ -499,7 +503,7 @@ def convert_to_vlm_format(
remaining_time = (total - done) / rate if rate > 0 else 0
eta_str = _format_eta(remaining_time)
progress_msg = f"Downloading images: {done:,}/{total:,} ({done*100//total}%) | ~{eta_str} remaining | {failed_count} skipped"
print(f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}")
logger.info(f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}")
_notify(progress_msg)
else:
# Sequential conversion for local/embedded images (fast, no I/O bottleneck)
@ -511,13 +515,13 @@ def convert_to_vlm_format(
failed_count += 1
if failed_count == 1:
# Log the first failure to aid debugging
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
pbar.close()
if failed_count > 0:
fail_rate = failed_count / total
print(f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
logger.info(f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
# For datasets that skipped the probe (small URL datasets), check fail rate now
if has_urls and fail_rate >= MAX_FAIL_RATE:
msg = (
@ -534,7 +538,7 @@ def convert_to_vlm_format(
"This dataset may contain only image URLs that are no longer accessible."
)
print(f"✅ Converted {len(converted_list)}/{total} samples")
logger.info(f"✅ Converted {len(converted_list)}/{total} samples")
_notify(f"Converted {len(converted_list):,}/{total:,} images successfully")
# Return list, NOT Dataset
@ -592,7 +596,7 @@ def convert_sharegpt_with_images_to_vlm_format(
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
print(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
_image_lookup = {
os.path.basename(f): f
@ -604,12 +608,12 @@ def convert_sharegpt_with_images_to_vlm_format(
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS):
_image_lookup[f] = f
if first_image in _image_lookup:
print(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
else:
print(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
_image_lookup = None
except Exception as e:
print(f"⚠️ Failed to build HF repo image lookup: {e}")
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
_image_lookup = None
def _resolve_image(image_data):
@ -670,7 +674,7 @@ def convert_sharegpt_with_images_to_vlm_format(
return {"messages": new_messages}
# ── Full conversion with progress ──
print(f"🔄 Converting {total} samples from ShareGPT+image format...")
logger.info(f"🔄 Converting {total} samples from ShareGPT+image format...")
converted_list = []
failed_count = 0
@ -681,12 +685,12 @@ def convert_sharegpt_with_images_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
print(f"⚠️ First conversion failure: {type(e).__name__}: {e}")
logger.info(f"⚠️ First conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
pbar.close()
if failed_count > 0:
print(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
if len(converted_list) == 0:
raise ValueError(
@ -694,7 +698,7 @@ def convert_sharegpt_with_images_to_vlm_format(
"no usable samples found."
)
print(f"✅ Converted {len(converted_list)}/{total} samples")
logger.info(f"✅ Converted {len(converted_list)}/{total} samples")
_notify(f"Converted {len(converted_list):,}/{total:,} samples successfully")
return converted_list
@ -712,7 +716,7 @@ def convert_llava_to_vlm_format(dataset):
"""
from PIL import Image
print(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
def _convert_single_sample(sample):
"""Convert a single llava sample to standard VLM format."""
@ -768,5 +772,5 @@ def convert_llava_to_vlm_format(dataset):
# Convert using list comprehension
converted_list = [_convert_single_sample(sample) for sample in dataset]
print(f"✅ Converted {len(converted_list)} samples")
logger.info(f"✅ Converted {len(converted_list)} samples")
return converted_list

13
studio/backend/utils/hardware/hardware.py
View file

@ -16,11 +16,12 @@ Usage:
...
"""
import platform
import logging
import structlog
from loggers import get_logger
from enum import Enum
from typing import Optional, Dict, Any
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# ========== Device Enum ==========
@ -82,19 +83,19 @@ def detect_hardware() -> DeviceType:
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
device_name = torch.cuda.get_device_properties(0).name
logger.info(f"Hardware detected: CUDA — {device_name}")
print(f"Hardware detected: CUDA — {device_name}")
return DEVICE
# --- MLX: Apple Silicon ---
if is_apple_silicon() and _has_mlx():
DEVICE = DeviceType.MLX
chip = platform.processor() or platform.machine()
logger.info(f"Hardware detected: MLX — Apple Silicon ({chip})")
print(f"Hardware detected: MLX — Apple Silicon ({chip})")
return DEVICE
# --- Fallback ---
DEVICE = DeviceType.CPU
logger.info("Hardware detected: CPU (no GPU backend available)")
print("Hardware detected: CPU (no GPU backend available)")
return DEVICE
@ -458,7 +459,7 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
if get_physical_gpu_count() > 1:
capped = min(4, desired)
print(
logger.info(
f"⚙️ Multi-GPU detected ({get_physical_gpu_count()} GPUs) — "
f"capping num_proc {desired} → {capped} to avoid fork deadlocks"
)

5
studio/backend/utils/inference/inference_config.py
View file

@ -10,11 +10,12 @@ from model YAML configuration files, with fallback to default.yaml.
from pathlib import Path
from typing import Dict, Any
import yaml
import logging
import structlog
from loggers import get_logger
from utils.models.model_config import load_model_defaults
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def load_inference_config(model_identifier: str) -> Dict[str, Any]:

5
studio/backend/utils/models/checkpoints.py
View file

@ -5,11 +5,12 @@
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
import json
import logging
import structlog
from loggers import get_logger
from pathlib import Path
from typing import List, Optional, Tuple
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:

9
studio/backend/utils/models/model_config.py
View file

@ -9,7 +9,8 @@ from dataclasses import dataclass
from typing import Optional, Dict, Any
from utils.paths import normalize_path, is_local_path, is_model_cached
from utils.utils import without_hf_auth
import logging
import structlog
from loggers import get_logger
import os
import subprocess
import sys
@ -19,7 +20,7 @@ import json
import yaml
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Model name mapping: maps all equivalent model names to their canonical YAML config file
# Format: "canonical_model_name.yaml": [list of all equivalent model names]
@ -459,10 +460,10 @@ try:
model_type = getattr(config, "model_type", "unknown")
archs = getattr(config, "architectures", [])
print(json.dumps({"is_vision": is_vlm, "model_type": model_type,
logger.info(json.dumps({"is_vision": is_vlm, "model_type": model_type,
"architectures": archs}))
except Exception as exc:
print(json.dumps({"error": str(exc)}))
logger.info(json.dumps({"error": str(exc)}))
sys.exit(1)
'''

5
studio/backend/utils/paths/path_utils.py
View file

@ -7,9 +7,10 @@ Path utilities for model and dataset handling
import os
from pathlib import Path
from typing import Optional
import logging
import structlog
from loggers import get_logger
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
def normalize_path(path: str) -> str:

15
studio/backend/utils/transformers_version.py
View file

@ -23,23 +23,16 @@ Strategy:
import importlib
import json
import logging
import structlog
from loggers import get_logger
import os
import subprocess
import sys
from pathlib import Path
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
# Ensure our logger is visible even if root logger isn't configured for INFO.
if not logger.handlers:
_handler = logging.StreamHandler()
_handler.setLevel(logging.INFO)
_handler.setFormatter(
logging.Formatter("[%(name)s|%(levelname)s]%(message)s")
)
logger.addHandler(_handler)
logger.setLevel(logging.INFO)
# ---------------------------------------------------------------------------
# Detection

5
studio/backend/utils/utils.py
View file

@ -5,14 +5,15 @@
Shared backend utilities
"""
import os
import logging
import structlog
from loggers import get_logger
from contextlib import contextmanager
from pathlib import Path
import shutil
import tempfile
logger = logging.getLogger(__name__)
logger = get_logger(__name__)
@contextmanager