Merge pull request #314 from unslothai/fix/vlm-dataset-conversion-error-handling-local

Fix VLM training abort on URL-based dataset conversion failure

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")
 

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

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:

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:

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

studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx

@@ -340,6 +340,13 @@ export function DatasetPreviewDialog({
                 />
               </div>
 
+              {data.warning && (
+                <div className="rounded-lg border border-amber-200 bg-amber-50 px-4 py-3 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400 mb-4 flex items-start gap-2.5">
+                  <HugeiconsIcon icon={AlertCircleIcon} className="size-4 shrink-0 mt-0.5" />
+                  <span>{data.warning}</span>
+                </div>
+              )}
+
               {mappingEnabled && (
                 <DatasetMappingCard
                   mapping={manualMapping}

studio/frontend/src/features/training/types/datasets.ts

@@ -9,5 +9,6 @@ export type CheckFormatResponse = {
   total_rows?: number | null;
   is_multimodal?: boolean;
   multimodal_columns?: string[] | null;
+  warning?: string | null;
 };