From 817f2e8dccb82bf6158ef333aad54c15c4d266f4 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Mar 2026 12:33:16 +0000 Subject: [PATCH] feat: integrate structlog, configure workers for prod logging, and migrate print statements --- studio/backend/colab.py | 28 +- .../backend/core/data_recipe/jobs/worker.py | 15 + .../data_recipe/local_callable_validators.py | 5 +- studio/backend/core/export/export.py | 5 +- studio/backend/core/export/orchestrator.py | 5 +- studio/backend/core/export/worker.py | 16 +- studio/backend/core/inference/audio_codecs.py | 5 +- studio/backend/core/inference/inference.py | 5 +- studio/backend/core/inference/llama_cpp.py | 5 +- studio/backend/core/inference/orchestrator.py | 5 +- studio/backend/core/inference/worker.py | 24 +- studio/backend/core/training/trainer.py | 292 +++++++++--------- studio/backend/core/training/training.py | 5 +- studio/backend/core/training/worker.py | 16 +- studio/backend/loggers/__init__.py | 3 + studio/backend/loggers/config.py | 72 +++++ studio/backend/loggers/handlers.py | 96 ++++++ studio/backend/main.py | 27 +- studio/backend/requirements/studio.txt | 3 +- studio/backend/routes/datasets.py | 14 +- studio/backend/routes/export.py | 14 +- studio/backend/routes/inference.py | 14 +- studio/backend/routes/models.py | 14 +- studio/backend/routes/training.py | 14 +- studio/backend/run.py | 10 +- studio/backend/tests/test_utils.py | 6 +- studio/backend/utils/cache_cleanup.py | 5 +- .../backend/utils/datasets/chat_templates.py | 22 +- .../backend/utils/datasets/data_collators.py | 6 +- .../backend/utils/datasets/dataset_utils.py | 6 +- .../utils/datasets/format_conversion.py | 60 ++-- studio/backend/utils/hardware/hardware.py | 13 +- .../utils/inference/inference_config.py | 5 +- studio/backend/utils/models/checkpoints.py | 5 +- studio/backend/utils/models/model_config.py | 9 +- studio/backend/utils/paths/path_utils.py | 5 +- studio/backend/utils/transformers_version.py | 15 +- studio/backend/utils/utils.py | 5 +- 38 files changed, 569 insertions(+), 305 deletions(-) create mode 100644 studio/backend/loggers/config.py create mode 100644 studio/backend/loggers/handlers.py diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 69400ea673..e18cf910a2 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -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) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index fbe9d91b6d..e4dbf37031 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -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: diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index 5e2749d45d..560cccd7ff 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -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" diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 986da0def2..8d387cb726 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -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(): diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index d6a04a3c1b..e46d2ebd5c 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -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") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index a19ba429e9..9c329f3f1e 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -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"] diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index b5cdd92560..ef1c137257 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -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: diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index bd9a9c0c43..8aef1a8d5c 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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""" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index dac8d380c3..c8a26fc4f5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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: diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index acab9e2245..bd81312d0e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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") diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 21214fbbc5..0e9213664a 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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"] diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 1f2fedeb7c..10f831c51f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -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 diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5f9cb86b0c..7676346abe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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") diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 8cc6e9785f..8083bb9483 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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"] diff --git a/studio/backend/loggers/__init__.py b/studio/backend/loggers/__init__.py index e69de29bb2..81ae4948b3 100644 --- a/studio/backend/loggers/__init__.py +++ b/studio/backend/loggers/__init__.py @@ -0,0 +1,3 @@ +from .handlers import get_logger + +__all__ = ["get_logger"] diff --git a/studio/backend/loggers/config.py b/studio/backend/loggers/config.py new file mode 100644 index 0000000000..d6247a139a --- /dev/null +++ b/studio/backend/loggers/config.py @@ -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) diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py new file mode 100644 index 0000000000..08f692c256 --- /dev/null +++ b/studio/backend/loggers/handlers.py @@ -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) diff --git a/studio/backend/main.py b/studio/backend/main.py index 3a437f08ff..a5cb337d69 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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, diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 244ed4f48c..04592147b9 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -11,4 +11,5 @@ pyjwt easydict addict gradio>=4.0.0 -huggingface-hub==0.36.2 \ No newline at end of file +huggingface-hub==0.36.2 +structlog>=24.1.0 diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 25965a847e..e0029d8671 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -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 ( diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index b475f75941..ca4226c03d 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -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) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1f72f0642f..450130f2b5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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() diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 2fb0171c94..a371338703 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index e3438824f6..f2ae56c332 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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") diff --git a/studio/backend/run.py b/studio/backend/run.py index 48cc25cce0..bffcc11d5f 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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()) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 97466040cf..941149581a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -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") diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py index 9b18fb08c8..b98c67d31a 100644 --- a/studio/backend/utils/cache_cleanup.py +++ b/studio/backend/utils/cache_cleanup.py @@ -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 diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index ae80055241..3d3556b339 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -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): diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index f362e93799..8ebd2b6400 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/studio/backend/utils/datasets/data_collators.py @@ -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 diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 2b430cee14..e1f45610ab 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -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( diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 7c45a03b49..624a10ba9b 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -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 diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 5b1ef8cfa3..ebafbdeda6 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -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" ) diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 05ad01071f..08dfe32256 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -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]: diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index c601cf64db..17960f12d1 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -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]: diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index e9fea321ba..5d50654ce8 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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) ''' diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py index 1151362c4c..75ee88c31a 100644 --- a/studio/backend/utils/paths/path_utils.py +++ b/studio/backend/utils/paths/path_utils.py @@ -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: diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index d003084b85..25fa7d6585 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -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 diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index a61ce4c7e7..eac376f3ee 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -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