update app

app.py

@@ -38,6 +38,7 @@ CODE_MAX_TOKENS = 8192
 DATASET_PREVIEW_TIMEOUT = 60  # seconds
 RUN_CODE_URL = os.getenv("DATACHEF_RUN_URL", "http://127.0.0.1:7001/run_code")
 RUN_CODE_TIMEOUT = int(os.getenv("DATACHEF_RUN_TIMEOUT", "300"))
+DATASET_LOAD_PROBE_TIMEOUT = int(os.getenv("DATASET_LOAD_PROBE_TIMEOUT", "60"))
 RUN_CODE_KEY = os.getenv("DATACHEF_RUN_KEY")  # optional shared secret sent as x-forward-key
 _DATASET_PREVIEW_CACHE: Dict[str, Dict] = {}
 # Preset library
@@ -987,6 +988,107 @@ def _get_dataset_examples(dataset_id: str, timeout: Optional[int] = None) -> Dic
         raise RuntimeError("; ".join(errors))
 
 
+def _run_code_probe(code_text: str, timeout: int) -> Tuple[bool, str]:
+    """
+    Run lightweight code on the proxy runner to test dataset loading.
+    Returns (success, error_message).
+    """
+    if not RUN_CODE_URL:
+        return False, "RUN_CODE_URL not configured"
+
+    headers = {}
+    if RUN_CODE_KEY:
+        headers["x-forward-key"] = RUN_CODE_KEY
+
+    try:
+        resp = requests.post(
+            RUN_CODE_URL,
+            json={"code": code_text, "timeout": timeout},
+            headers=headers,
+            timeout=min(timeout, 30) + 10,
+        )
+        resp.raise_for_status()
+        data = resp.json() if resp.content else {}
+    except Exception as e:  # noqa: BLE001
+        return False, f"probe request failed: {e}"
+
+    job_id = data.get("job_id")
+    if not job_id:
+        return False, "probe runner did not return job_id"
+
+    status_base = RUN_CODE_URL.rsplit("/", 1)[0].rstrip("/")
+    status_url = f"{status_base}/status/{job_id}"
+    deadline = time.time() + timeout
+    poll_interval = 1.5
+    result_data = None
+    while time.time() < deadline:
+        try:
+            r = requests.get(status_url, headers=headers, timeout=min(timeout, 30))
+            r.raise_for_status()
+            job = r.json() if r.content else {}
+        except Exception:
+            time.sleep(poll_interval)
+            continue
+        status = job.get("status")
+        if status in {"success", "failed", "timeout", "error"}:
+            result_data = job
+            break
+        time.sleep(poll_interval)
+
+    if result_data is None:
+        return False, f"probe job {job_id} timed out after {timeout}s"
+
+    status = result_data.get("status")
+    retcode = result_data.get("returncode")
+    stderr = result_data.get("stderr") or ""
+    summary = result_data.get("summary") or ""
+    if status == "success" and retcode == 0:
+        return True, ""
+    return False, summary or stderr or f"probe failed (status={status}, returncode={retcode})"
+
+
+def _probe_dataset_load(dataset_id: str, name: Optional[str], split: str) -> Tuple[bool, str]:
+    """
+    Trigger a small load via proxy runner to weed out slow/invalid datasets.
+    Loads up to 20 samples; returns (success, error_message).
+    """
+    name_literal = "None" if not name or name == "default" else repr(name)
+    split_literal = repr(split or "train")
+    code = (
+        "from aidp import load_remote_dataset\n"
+        "from datasets import disable_caching\n"
+        "disable_caching()\n"
+        f"ds = load_remote_dataset({repr(dataset_id)}, name={name_literal}, split={split_literal}, num_samples=20, shuffle=False)\n"
+        "print('rows', len(ds))\n"
+        "print('columns', list(ds.features.keys()))\n"
+    )
+    return _run_code_probe(code, timeout=DATASET_LOAD_PROBE_TIMEOUT)
+
+
+def _filter_dataset_infos_by_load(dataset_infos: List[Dict]) -> Tuple[List[Dict], Dict[str, str]]:
+    """
+    Run load probes for each dataset. Datasets that fail or timeout are dropped.
+    Returns (kept_dataset_infos, load_errors).
+    """
+    kept: List[Dict] = []
+    errors: Dict[str, str] = {}
+    for info in dataset_infos:
+        dataset_id = info.get("dataset_id", "")
+        examples = info.get("examples") or []
+        if not examples:
+            errors[dataset_id] = "no examples to probe"
+            continue
+        first = examples[0]
+        name = first.get("name") or None
+        split = first.get("split") or "train"
+        ok, err = _probe_dataset_load(dataset_id, name, split)
+        if ok:
+            kept.append(info)
+        else:
+            errors[dataset_id] = err
+    return kept, errors
+
+
 def _fetch_json(url: str, timeout: int) -> Dict:
     try:
         with urllib.request.urlopen(url, timeout=timeout) as resp:
@@ -1171,6 +1273,12 @@ def _build_input_sample(
             dataset_infos.append(_get_dataset_examples(dataset_id, DATASET_PREVIEW_TIMEOUT))
         except Exception as e:
             errors[dataset_id] = str(e)
+
+    # Probe actual dataset loading via proxy runner to drop slow/invalid ones
+    if dataset_infos:
+        dataset_infos, load_errors = _filter_dataset_infos_by_load(dataset_infos)
+        errors.update(load_errors)
+
     if errors and not dataset_infos:
         raise ValueError(f"Dataset build error: {errors}")
     return {
@@ -1199,6 +1307,11 @@ def _preview_dataset(dataset_ids_table: List[List[str]]) -> str:
                 dataset_infos.append(_get_dataset_examples(dataset_id, DATASET_PREVIEW_TIMEOUT))
             except Exception as e:
                 errors[dataset_id] = str(e)
+
+        # Probe loadability and drop slow/failed datasets before rendering preview
+        if dataset_infos:
+            dataset_infos, load_errors = _filter_dataset_infos_by_load(dataset_infos)
+            errors.update(load_errors)
         
         return _render_preview_table_html(dataset_infos, errors if errors else None)
     except Exception as e: