diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 5ce273f43c..f2e43f76f4 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -475,6 +475,7 @@ class UnslothTrainer: format_type=format_type, dataset_name=dataset_source, custom_format_mapping=custom_format_mapping, + progress_callback=self._update_progress, ) # Check if stopped during formatting @@ -482,6 +483,14 @@ class UnslothTrainer: print("Stopped during dataset formatting\n") return None + # Abort if dataset formatting/conversion failed + if not dataset_info.get("success", True): + errors = dataset_info.get("errors", []) + error_msg = "; ".join(errors) if errors else "Dataset formatting failed" + logger.error(f"Dataset conversion failed: {error_msg}") + self._update_progress(error=error_msg) + return None + self._update_progress(status_message=f"Dataset formatted and ready for training") print(f"Dataset formatted successfully\n") diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 81adef7577..18f6ec224b 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -34,3 +34,4 @@ class CheckFormatResponse(BaseModel): detected_text_column: Optional[str] = None preview_samples: Optional[List[Dict]] = None total_rows: Optional[int] = None + warning: Optional[str] = None diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index a53e223015..4a475ff8c2 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -211,6 +211,21 @@ def check_format( else: preview_samples = _serialize_preview_rows(preview_slice) + # Lightweight URL-based image detection for VLM datasets + warning = None + image_col = result.get("detected_image_column") + if image_col and image_col in (result.get("columns") or []): + try: + sample_val = preview_slice[0][image_col] + if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): + warning = ( + "This dataset contains image URLs instead of embedded images. " + "Images will be downloaded during training, which may be slow for large datasets." + ) + logger.info(f"URL-based image column detected: {image_col}") + except Exception: + pass + return CheckFormatResponse( requires_manual_mapping=result["requires_manual_mapping"], detected_format=result["detected_format"], @@ -222,6 +237,7 @@ def check_format( detected_text_column=result.get("detected_text_column"), preview_samples=preview_samples, total_rows=total_rows, + warning=warning, ) except HTTPException: diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 3a4d54f93f..9e1f54a75c 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -593,6 +593,7 @@ def format_and_template_dataset( aliases_for_assistant=["gpt", "assistant", "output",], batch_size=1000, num_proc=None, + progress_callback=None, ): """ Convenience function that combines format_dataset and apply_chat_template_to_dataset. @@ -638,6 +639,7 @@ def format_and_template_dataset( text_column=user_vlm_text_column, image_column=user_vlm_image_column, dataset_name=dataset_name, + progress_callback=progress_callback, ) warnings.append(f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'") @@ -734,6 +736,7 @@ def format_and_template_dataset( text_column=vlm_text_column, image_column=vlm_image_column, dataset_name=dataset_name, + progress_callback=progress_callback, ) if vlm_instruction: diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 6436e7a82a..41a9617857 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -238,24 +238,51 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None): return dataset.map(_convert, **dataset_map_kwargs) +def _format_eta(seconds): + """Format seconds into a human-readable ETA string.""" + if seconds < 60: + return f"{seconds:.0f}s" + elif seconds < 3600: + m, s = divmod(int(seconds), 60) + return f"{m}m {s}s" + else: + h, remainder = divmod(int(seconds), 3600) + m, _ = divmod(remainder, 60) + return f"{h}h {m}m" + + def convert_to_vlm_format( dataset, instruction=None, text_column="text", image_column="image", dataset_name=None, + progress_callback=None, ): """ Converts simple {image, text} format to VLM messages format. Returns a LIST, not a HuggingFace Dataset (to preserve PIL Images). + For URL-based image datasets, runs a 200-sample parallel probe first to + estimate download speed and failure rate, then reports time estimate or + warning through progress_callback before proceeding with the full conversion. + + Args: + progress_callback: Optional callable(status_message=str) to report + progress to the training overlay. + Returns: list: List of dicts with 'messages' field """ from PIL import Image from .vlm_processing import generate_smart_vlm_instruction + def _notify(msg): + """Send status update to the training overlay if callback is available.""" + if progress_callback: + progress_callback(status_message=msg) + # Generate smart instruction if not provided if instruction is None: instruction_info = generate_smart_vlm_instruction( @@ -281,12 +308,17 @@ def convert_to_vlm_format( def _convert_single_sample(sample): """Convert a single sample to VLM format.""" - # Get image (might be PIL Image or path) + # Get image (might be PIL Image, local path, or URL) image_data = sample[image_column] - # Handle image paths if isinstance(image_data, str): - image_data = Image.open(image_data).convert("RGB") + if image_data.startswith(("http://", "https://")): + import fsspec + from io import BytesIO + with fsspec.open(image_data, "rb", expand=True) as f: + image_data = Image.open(BytesIO(f.read())).convert("RGB") + else: + image_data = Image.open(image_data).convert("RGB") # Get text text_data = sample[text_column] @@ -317,11 +349,143 @@ def convert_to_vlm_format( # Return dict with messages return {"messages": messages} - # Use list comprehension and return the LIST directly - print(f"🔄 Converting {len(dataset)} samples to VLM format...") - converted_list = [_convert_single_sample(sample) for sample in dataset] + total = len(dataset) + first_image = next(iter(dataset))[image_column] + has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://")) - print(f"✅ Converted {len(converted_list)} samples") + # ── URL probe: 200 samples with parallel workers to estimate speed + failure rate ── + PROBE_SIZE = 200 + MAX_FAIL_RATE = 0.3 + + if has_urls and total > PROBE_SIZE: + import time + from concurrent.futures import ThreadPoolExecutor, as_completed + from utils.hardware import safe_num_proc + + 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...") + + probe_samples = [dataset[i] for i in range(PROBE_SIZE)] + probe_ok = 0 + probe_fail = 0 + probe_start = time.time() + + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples} + for future in as_completed(futures): + try: + future.result() + probe_ok += 1 + except Exception: + probe_fail += 1 + + probe_elapsed = time.time() - probe_start + probe_total = probe_ok + probe_fail + fail_rate = probe_fail / probe_total if probe_total > 0 else 0 + throughput = probe_total / probe_elapsed if probe_elapsed > 0 else 0 + + if fail_rate >= MAX_FAIL_RATE: + msg = ( + f"⚠️ {fail_rate:.0%} of the first {PROBE_SIZE} images failed to download " + f"({probe_fail}/{probe_total}). " + "This dataset has too many broken or unreachable image URLs. " + "Consider using a dataset with embedded images instead." + ) + print(msg) + _notify(msg) + raise ValueError(msg) + + # Estimate total time for remaining samples + remaining = total - PROBE_SIZE + estimated_seconds = remaining / throughput if throughput > 0 else 0 + eta_str = _format_eta(estimated_seconds) + + info_msg = ( + f"Downloading {total:,} images ({num_workers} workers, ~{throughput:.1f} img/s). " + f"Estimated time: ~{eta_str}" + ) + 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}") + _notify(info_msg) + + # ── Full conversion with progress ── + from tqdm import tqdm + + print(f"🔄 Converting {total} samples to VLM format...") + converted_list = [] + failed_count = 0 + + if has_urls: + # Parallel conversion for URL-based datasets + import time + from concurrent.futures import ThreadPoolExecutor, as_completed + from utils.hardware import safe_num_proc + + num_workers = safe_num_proc() + batch_size = 500 + start_time = time.time() + + for batch_start in range(0, total, batch_size): + batch_end = min(batch_start + batch_size, total) + batch_samples = [dataset[i] for i in range(batch_start, batch_end)] + + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = {executor.submit(_convert_single_sample, s): i for i, s in enumerate(batch_samples)} + batch_results = [None] * len(batch_samples) + for future in as_completed(futures): + idx = futures[future] + try: + batch_results[idx] = future.result() + except Exception: + failed_count += 1 + + converted_list.extend(r for r in batch_results if r is not None) + + # Progress update every batch + elapsed = time.time() - start_time + done = batch_end + rate = done / elapsed if elapsed > 0 else 0 + 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}") + _notify(progress_msg) + else: + # Sequential conversion for local/embedded images (fast, no I/O bottleneck) + pbar = tqdm(dataset, total=total, desc="Converting VLM samples", unit="sample") + for sample in pbar: + try: + converted_list.append(_convert_single_sample(sample)) + except Exception: + failed_count += 1 + 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") + # For datasets that skipped the probe (small URL datasets), check fail rate now + if has_urls and fail_rate >= MAX_FAIL_RATE: + msg = ( + f"⚠️ {fail_rate:.0%} of images failed to download ({failed_count}/{total}). " + "This dataset has too many broken or unreachable image URLs. " + "Consider using a dataset with embedded images instead." + ) + _notify(msg) + raise ValueError(msg) + + if len(converted_list) == 0: + raise ValueError( + f"All {total} samples failed during VLM conversion — no usable images found. " + "This dataset may contain only image URLs that are no longer accessible." + ) + + print(f"✅ Converted {len(converted_list)}/{total} samples") + _notify(f"Converted {len(converted_list):,}/{total:,} images successfully") # Return list, NOT Dataset return converted_list diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 00a738bc64..516f3612c0 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -340,6 +340,13 @@ export function DatasetPreviewDialog({ /> + {data.warning && ( +