Wire up Daft execution/S3 tuning and fix Ray parallelism collapse

README.md

@@ -105,6 +105,40 @@ USE_RAY=0               # set to 1 to use Ray for Daft-backed steps
 RAY_ADDRESS=            # leave empty to start/join local Ray
 ```
 
+### Tuning for large Ray runs
+
+All knobs below are applied through `configure_daft_runner()` / `daft_io_config()`
+at workflow startup (Daft itself only reads a handful of `DAFT_*` env vars natively).
+
+```sh
+# Parallelism: analyze stages split the manifest into this many partitions
+# before download. Unset/auto = 2 × cluster CPUs on Ray, no split on native.
+ANALYZE_NUM_PARTITIONS=auto
+
+# Memory safety: rows per morsel. Keep small — rows carry image/audio bytes.
+DAFT_DEFAULT_MORSEL_SIZE=32
+
+# LLM concurrency: async requests per parallel task. Total in-flight requests
+# = DEEPSEEK_CONCURRENCY × parallel tasks; budget against provider rate limits.
+DEEPSEEK_CONCURRENCY=8
+
+# S3/MinIO: connections per IO thread, retries, timeouts (now actually applied).
+S3_MAX_CONNECTIONS=8
+S3_NUM_TRIES=5
+S3_READ_TIMEOUT_MS=60000
+
+# Actor startup: raise if many UDF actors cold-start and download models at once.
+DAFT_ACTOR_UDF_READY_TIMEOUT=600
+
+# Long jobs: enable Daft's event log for post-mortem debugging (read natively by Daft).
+DAFT_EVENT_LOG_ENABLED=1
+DAFT_EVENT_LOG_DIR=/tmp/daft-events
+```
+
+Output layout: JSONL writes coalesce partitions (≈ `ANALYZE_NUM_PARTITIONS / 8` files)
+and Lance writes honor `LANCE_MAX_ROWS_PER_FILE` / `LANCE_MAX_BYTES_PER_FILE`, so high
+analysis parallelism does not produce a flood of tiny output files or fragments.
+
 ## Usage — audio workflow pipeline
 
 The manifest must be parquet, jsonl, or csv with at minimum `doc_id` and `s3_url` columns.  

multimodal_toolkit/audio/config.py

@@ -12,6 +12,9 @@
 DEEPSEEK_BASE_URL = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
 DEEPSEEK_MODEL = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "")
+# llm_prompt 的 async 并发协程数（每个并行 task 内）。总并发 =
+# 该值 × 并行 task 数，需按 API 供应商限流预算调整。
+DEEPSEEK_CONCURRENCY = _shared_config.env_int("DEEPSEEK_CONCURRENCY", 8)
 
 MIN_DURATION_S = float(os.getenv("MIN_DURATION_S", "0"))
 MAX_DURATION_S = float(os.getenv("MAX_DURATION_S", "1800"))

multimodal_toolkit/audio/udfs.py

@@ -46,7 +46,7 @@ def prompt_udf(transcripts, acoustic_emotions):
     ]
 
 
-@daft.cls(cpus=1)
+@daft.cls(cpus=1, max_retries=2)
 class AsrUDF:
     def __init__(self) -> None:
         from multimodal_toolkit.audio.asr import SenseVoiceASR

multimodal_toolkit/audio/workflow/analyze.py

@@ -29,7 +29,14 @@
     duration_udf,
     prompt_udf,
 )
-from ...storage.io import configure_daft_runner, daft_io_config, read_manifest
+from ...storage.io import (
+    coalesce_for_write,
+    configure_daft_runner,
+    daft_io_config,
+    read_manifest,
+    spread_partitions,
+)
+from ... import config as shared_config
 
 # PII patterns (Chinese mainland): 18-digit resident ID and 11-digit mobile
 # number. Redacted from transcripts before they are sent to the LLM or stored.
@@ -53,7 +60,7 @@
 ]
 
 
-@daft.cls(cpus=1)
+@daft.cls(cpus=1, max_retries=2)
 class _EmbedUDF:
     """Turns raw audio into a fixed-length float vector (the "embedding").
 
@@ -80,7 +87,9 @@ def __call__(self, audio_bytes_col):
 
 def _build_analysis_df(manifest: str, io_config) -> daft.DataFrame:
     """Download audio from S3 and run the full analysis pipeline."""
+    # 下载前先切分区：manifest 文件太小，不切的话 Ray 上整条链路只有一个 task。
     df = read_manifest(manifest)
+    df = spread_partitions(df)
     df = df.with_column(
         "audio_bytes", download(col("s3_url"), on_error="null", io_config=io_config)
     )
@@ -131,6 +140,8 @@ def _build_analysis_df(manifest: str, io_config) -> daft.DataFrame:
                 use_chat_completions=True,
                 response_format={"type": "json_object"},
                 temperature=0,
+                concurrency=config.DEEPSEEK_CONCURRENCY,
+                on_error="ignore",
             ),
         )
     else:
@@ -163,10 +174,16 @@ def run(manifest: str, out_path: str, embed: bool = False) -> None:
         embed_udf = _EmbedUDF()
         df = df.with_column("audio_embedding", embed_udf(col("audio_bytes")))
         output = df.select(*_BASE_OUTPUT_COLS, "audio_embedding")
-        output.write_lance(out_path, mode="overwrite", io_config=io_config)
+        output.write_lance(
+            out_path,
+            mode="overwrite",
+            io_config=io_config,
+            max_rows_per_file=shared_config.LANCE_MAX_ROWS_PER_FILE,
+            max_bytes_per_file=shared_config.LANCE_MAX_BYTES_PER_FILE,
+        )
         print(f"[ok] wrote analysis+embedding lance staging table: {out_path}")
     else:
-        output = df.select(*_BASE_OUTPUT_COLS)
+        output = coalesce_for_write(df.select(*_BASE_OUTPUT_COLS))
         output.write_json(out_path, write_mode="overwrite", io_config=io_config)
         print(f"[ok] wrote analysis JSONL: {out_path}")
 

multimodal_toolkit/audio/workflow/ingest.py

@@ -13,6 +13,7 @@
 from daft import col, lit
 from daft.functions import download
 
+from ... import config
 from ...storage.blob import validate_blob_v2
 from ...storage.io import configure_daft_runner, daft_io_config, lance_write_mode, read_analysis_output
 
@@ -35,7 +36,14 @@ def run(analysis_path: str, lance_uri: str) -> None:
     )
 
     mode = lance_write_mode(lance_uri)
-    df.write_lance(lance_uri, mode=mode, io_config=io_config, blob_columns=["audio_blob"])
+    df.write_lance(
+        lance_uri,
+        mode=mode,
+        io_config=io_config,
+        blob_columns=["audio_blob"],
+        max_rows_per_file=config.LANCE_MAX_ROWS_PER_FILE,
+        max_bytes_per_file=config.LANCE_MAX_BYTES_PER_FILE,
+    )
     validate_blob_v2(lance_uri, "audio_blob")
 
     result = daft.read_lance(lance_uri, io_config=io_config)

multimodal_toolkit/config.py

@@ -2,6 +2,7 @@
 
 import os
 from pathlib import Path
+from typing import TypeVar
 
 try:
     from dotenv import load_dotenv
@@ -19,11 +20,98 @@ def load_env() -> None:
 
 load_env()
 
+_T = TypeVar("_T", int, float)
+
+
+def env_bool(name: str, default: bool) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    normalized = raw.strip().lower()
+    if normalized in {"1", "true", "yes", "on"}:
+        return True
+    if normalized in {"0", "false", "no", "off"}:
+        return False
+    raise ValueError(f"{name} must be a boolean, got {raw!r}")
+
+
+def _env_number(name: str, default: _T, parser: type[_T], *, minimum: _T) -> _T:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    try:
+        value = parser(raw)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be a {parser.__name__}, got {raw!r}") from exc
+    if value < minimum:
+        raise ValueError(f"{name} must be >= {minimum}, got {value}")
+    return value
+
+
+def env_int(name: str, default: int, *, minimum: int = 1) -> int:
+    return _env_number(name, default, int, minimum=minimum)
+
+
+def env_float(name: str, default: float, *, minimum: float = 0.0) -> float:
+    return _env_number(name, default, float, minimum=minimum)
+
+
+def env_optional_int(name: str) -> int | None:
+    raw = os.getenv(name)
+    if raw is None or not raw.strip() or raw.strip().lower() in {"auto", "none"}:
+        return None
+    return env_int(name, 1)
+
+
+def env_choice(name: str, default: str, choices: set[str]) -> str:
+    value = os.getenv(name, default).strip().lower()
+    if value not in choices:
+        allowed = ", ".join(sorted(choices))
+        raise ValueError(f"{name} must be one of {allowed}, got {value!r}")
+    return value
+
 S3_ENDPOINT = os.getenv("MINIO_ENDPOINT", "http://127.0.0.1:9000")
 S3_KEY = os.getenv("MINIO_ROOT_USER", "minioadmin")
 S3_SECRET = os.getenv("MINIO_ROOT_PASSWORD", "minioadmin")
 S3_REGION = os.getenv("MINIO_REGION", "us-east-1")
 S3_USE_SSL = S3_ENDPOINT.startswith("https")
+S3_MAX_CONNECTIONS = env_int("S3_MAX_CONNECTIONS", 8)
+S3_RETRY_INITIAL_BACKOFF_MS = env_int("S3_RETRY_INITIAL_BACKOFF_MS", 1000)
+S3_CONNECT_TIMEOUT_MS = env_int("S3_CONNECT_TIMEOUT_MS", 10_000)
+S3_READ_TIMEOUT_MS = env_int("S3_READ_TIMEOUT_MS", 60_000)
+S3_NUM_TRIES = env_int("S3_NUM_TRIES", 5)
+S3_RETRY_MODE = env_choice("S3_RETRY_MODE", "adaptive", {"adaptive", "standard"})
 
-USE_RAY = os.getenv("USE_RAY", "0").lower() in ("1", "true", "yes")
+USE_RAY = env_bool("USE_RAY", False)
 RAY_ADDRESS = os.getenv("RAY_ADDRESS") or None  # None = start/join local Ray
+
+# Keep USE_RAY as a compatibility fallback. DAFT_RUNNER is the canonical switch.
+DAFT_RUNNER = env_choice(
+    "DAFT_RUNNER",
+    "ray" if USE_RAY else "native",
+    {"native", "ray"},
+)
+DAFT_NATIVE_NUM_THREADS = env_int("DAFT_NATIVE_NUM_THREADS", 4)
+DAFT_DEFAULT_MORSEL_SIZE = env_int("DAFT_DEFAULT_MORSEL_SIZE", 32)
+DAFT_ENABLE_DYNAMIC_BATCHING = env_bool("DAFT_ENABLE_DYNAMIC_BATCHING", True)
+DAFT_ENABLE_SCAN_TASK_SPLIT_AND_MERGE = env_bool(
+    "DAFT_ENABLE_SCAN_TASK_SPLIT_AND_MERGE", True
+)
+DAFT_SCAN_TASKS_MIN_SIZE_BYTES = env_int("DAFT_SCAN_TASKS_MIN_SIZE_BYTES", 32 * 1024 * 1024)
+DAFT_SCAN_TASKS_MAX_SIZE_BYTES = env_int("DAFT_SCAN_TASKS_MAX_SIZE_BYTES", 128 * 1024 * 1024)
+if DAFT_SCAN_TASKS_MAX_SIZE_BYTES < DAFT_SCAN_TASKS_MIN_SIZE_BYTES:
+    raise ValueError("DAFT_SCAN_TASKS_MAX_SIZE_BYTES must be >= DAFT_SCAN_TASKS_MIN_SIZE_BYTES")
+DAFT_MAX_SOURCES_PER_SCAN_TASK = env_int("DAFT_MAX_SOURCES_PER_SCAN_TASK", 16)
+DAFT_SCANTASK_MAX_PARALLEL = env_int("DAFT_SCANTASK_MAX_PARALLEL", 4)
+DAFT_JSON_TARGET_FILESIZE = env_int("DAFT_JSON_TARGET_FILESIZE", 128 * 1024 * 1024)
+DAFT_MAINTAIN_ORDER = env_bool("DAFT_MAINTAIN_ORDER", False)
+# 600s：Ray 上几十个 actor 同时冷启动 + 首次下载模型时，300s 会集体超时。
+DAFT_ACTOR_UDF_READY_TIMEOUT = env_int("DAFT_ACTOR_UDF_READY_TIMEOUT", 600)
+
+# analyze 阶段的分区数。manifest 只有几 MB，scan task 按字节切分不起作用，
+# 不显式切分的话整条 download+分析链路在 Ray 上只有一个 task。
+# 未设置（或 auto）时：Ray 模式按集群 CPU 数推导，native 模式不切分。
+ANALYZE_NUM_PARTITIONS = env_optional_int("ANALYZE_NUM_PARTITIONS")
+
+LANCE_MAX_ROWS_PER_FILE = env_int("LANCE_MAX_ROWS_PER_FILE", 100_000)
+LANCE_MAX_BYTES_PER_FILE = env_int("LANCE_MAX_BYTES_PER_FILE", 512 * 1024 * 1024)

multimodal_toolkit/image/udfs.py

@@ -49,7 +49,7 @@
 }
 
 
-@daft.cls(cpus=1)
+@daft.cls(cpus=1, max_retries=2)
 class ImageQualityUDF:
     def __init__(self) -> None:
         # 每个 worker 进程只加载一次模型，而不是每行加载——模型初始化

multimodal_toolkit/image/workflow/analyze.py

@@ -22,7 +22,14 @@
 from daft.functions import download, when
 
 from .. import config
-from ...storage.io import configure_daft_runner, daft_io_config, read_manifest
+from ...storage.io import (
+    coalesce_for_write,
+    configure_daft_runner,
+    daft_io_config,
+    read_manifest,
+    spread_partitions,
+)
+from ... import config as shared_config
 from ..rules import add_rule_columns
 from ..udfs import ImageQualityUDF
 
@@ -56,7 +63,7 @@
 ]
 
 
-@daft.cls(cpus=1)
+@daft.cls(cpus=1, max_retries=2)
 class _ImageEmbedUDF:
     def __init__(self) -> None:
         from multimodal_toolkit.image.embedding import get_embedder
@@ -87,7 +94,9 @@ def run(manifest: str, out_path: str, embed: bool = False) -> None:
 
     # manifest 只有 doc_id + s3_url 两列；按 s3_url 下载图片字节，
     # 失败的行 image_bytes 为 null（on_error="null"），保留不丢。
+    # 下载前先切分区：manifest 文件太小，不切的话 Ray 上整条链路只有一个 task。
     df = read_manifest(manifest)
+    df = spread_partitions(df)
     df = df.with_column(
         "image_bytes", download(col("s3_url"), on_error="null", io_config=io_config)
     )
@@ -119,10 +128,16 @@ def run(manifest: str, out_path: str, embed: bool = False) -> None:
         embed_udf = _ImageEmbedUDF()
         df = df.with_column("image_embedding", embed_udf(col("image_bytes"), col("status")))
         output = df.select(*_OUTPUT_COLS, "image_embedding")
-        output.write_lance(out_path, mode="overwrite", io_config=io_config)
+        output.write_lance(
+            out_path,
+            mode="overwrite",
+            io_config=io_config,
+            max_rows_per_file=shared_config.LANCE_MAX_ROWS_PER_FILE,
+            max_bytes_per_file=shared_config.LANCE_MAX_BYTES_PER_FILE,
+        )
         print(f"[ok] wrote image analysis+embedding lance staging table: {out_path}")
     else:
-        output = df.select(*_OUTPUT_COLS)
+        output = coalesce_for_write(df.select(*_OUTPUT_COLS))
         output.write_json(out_path, write_mode="overwrite", io_config=io_config)
         print(f"[ok] wrote image analysis JSONL: {out_path}")
 

multimodal_toolkit/image/workflow/ingest.py

@@ -12,6 +12,7 @@
 from daft import col, lit
 from daft.functions import download, when
 
+from ... import config
 from ...storage.blob import validate_blob_v2
 from ...storage.io import (
     configure_daft_runner,
@@ -107,7 +108,14 @@ def run(analysis_path: str, lance_uri: str) -> None:
         lit(now).cast(daft.DataType.timestamp("us", "UTC")),
     )
 
-    df.write_lance(lance_uri, mode=mode, io_config=io_config, blob_columns=["image_blob"])
+    df.write_lance(
+        lance_uri,
+        mode=mode,
+        io_config=io_config,
+        blob_columns=["image_blob"],
+        max_rows_per_file=config.LANCE_MAX_ROWS_PER_FILE,
+        max_bytes_per_file=config.LANCE_MAX_BYTES_PER_FILE,
+    )
     # 校验 image_blob 确实以 lance blob v2 编码落盘（而不是被静默降级成
     # 普通 large_binary），库版本升级时这是最容易出问题的地方。
     validate_blob_v2(lance_uri, "image_blob")