Add adaptive metadata upload batch sizing

dagshub/common/adaptive_batching.py

@@ -0,0 +1,259 @@
+import logging
+import time
+from dataclasses import dataclass
+import itertools
+from types import SimpleNamespace
+from typing import Callable, Iterable, List, Optional, Sized, TypeVar
+
+import rich.progress
+from tenacity import wait_exponential
+
+from dagshub.common.rich_util import get_rich_progress
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+MIN_TARGET_BATCH_TIME_SECONDS = 0.01
+
+
+@dataclass(frozen=True)
+class AdaptiveBatchConfig:
+    max_batch_size: int
+    min_batch_size: int
+    initial_batch_size: int
+    target_batch_time_seconds: float
+    batch_growth_factor: int
+    retry_backoff_base_seconds: float
+    retry_backoff_max_seconds: float
+
+    @classmethod
+    def from_values(
+        cls,
+        max_batch_size: Optional[int] = None,
+        min_batch_size: Optional[int] = None,
+        initial_batch_size: Optional[int] = None,
+        target_batch_time_seconds: Optional[float] = None,
+        batch_growth_factor: Optional[int] = None,
+        retry_backoff_base_seconds: Optional[float] = None,
+        retry_backoff_max_seconds: Optional[float] = None,
+    ) -> "AdaptiveBatchConfig":
+        import dagshub.common.config as dgs_config
+
+        if max_batch_size is None:
+            max_batch_size = dgs_config.dataengine_metadata_upload_batch_size
+        if min_batch_size is None:
+            min_batch_size = dgs_config.dataengine_metadata_upload_batch_size_min
+        if initial_batch_size is None:
+            initial_batch_size = dgs_config.dataengine_metadata_upload_batch_size_initial
+        if target_batch_time_seconds is None:
+            target_batch_time_seconds = dgs_config.dataengine_metadata_upload_target_batch_time_seconds
+        if batch_growth_factor is None:
+            batch_growth_factor = dgs_config.adaptive_batch_growth_factor
+        if retry_backoff_base_seconds is None:
+            retry_backoff_base_seconds = dgs_config.adaptive_batch_retry_backoff_base_seconds
+        if retry_backoff_max_seconds is None:
+            retry_backoff_max_seconds = dgs_config.adaptive_batch_retry_backoff_max_seconds
+
+        normalized_max_batch_size = max(1, max_batch_size)
+        normalized_min_batch_size = max(1, min(min_batch_size, normalized_max_batch_size))
+        normalized_initial_batch_size = max(
+            normalized_min_batch_size,
+            min(initial_batch_size, normalized_max_batch_size),
+        )
+        normalized_target_batch_time_seconds = max(target_batch_time_seconds, MIN_TARGET_BATCH_TIME_SECONDS)
+        return cls(
+            max_batch_size=normalized_max_batch_size,
+            min_batch_size=normalized_min_batch_size,
+            initial_batch_size=normalized_initial_batch_size,
+            target_batch_time_seconds=normalized_target_batch_time_seconds,
+            batch_growth_factor=max(2, batch_growth_factor),
+            retry_backoff_base_seconds=max(0.0, retry_backoff_base_seconds),
+            retry_backoff_max_seconds=max(0.0, retry_backoff_max_seconds),
+        )
+
+
+def _clamp(value: int, lo: int, hi: int) -> int:
+    return max(lo, min(hi, value))
+
+
+def _next_batch_after_success(
+    batch_size: int,
+    config: AdaptiveBatchConfig,
+    bad_batch_size: Optional[int],
+) -> int:
+    """Pick the next batch size after a successful (fast) batch.
+
+    Strategy:
+    - If we have a known-bad upper bound, binary-search toward it.
+    - Otherwise, multiply by the growth factor.
+    - Always guarantee at least +1 progress (so we never stall).
+    """
+    if bad_batch_size is not None and batch_size < bad_batch_size:
+        # Binary search: try the midpoint between current and bad
+        candidate = (batch_size + bad_batch_size) // 2
+    else:
+        # No upper bound (or we've already passed it): grow aggressively
+        candidate = batch_size * config.batch_growth_factor
+
+    # Must advance by at least 1 to avoid stalling
+    candidate = max(candidate, batch_size + 1)
+
+    return _clamp(candidate, config.min_batch_size, config.max_batch_size)
+
+
+def _next_batch_after_retryable_failure(
+    batch_size: int,
+    config: AdaptiveBatchConfig,
+    good_batch_size: Optional[int],
+    bad_batch_size: Optional[int],
+) -> int:
+    """Pick the next batch size after a failed or slow batch.
+
+    Strategy:
+    - If we have a known-good lower bound, binary-search between it and the
+      failing size.
+    - Otherwise, halve.
+    - Must be strictly less than the current size (so we converge downward).
+    """
+    if batch_size <= config.min_batch_size:
+        return config.min_batch_size
+
+    ceiling = batch_size - 1  # must shrink
+    if bad_batch_size is not None:
+        ceiling = min(ceiling, bad_batch_size - 1)
+
+    if good_batch_size is not None and good_batch_size < ceiling:
+        # Binary search: try the midpoint between good and failing
+        candidate = (good_batch_size + ceiling) // 2
+    else:
+        # No good lower bound — probe midpoint of the valid range
+        candidate = (config.min_batch_size + ceiling) // 2
+
+    return _clamp(candidate, config.min_batch_size, ceiling)
+
+
+def _get_retry_delay_seconds(consecutive_retryable_failures: int, config: AdaptiveBatchConfig) -> float:
+    # SimpleNamespace duck-types the .attempt_number attribute that tenacity's
+    # wait strategies read, avoiding the heavier RetryCallState constructor.
+    strategy = wait_exponential(
+        multiplier=config.retry_backoff_base_seconds,
+        min=config.retry_backoff_base_seconds,
+        max=config.retry_backoff_max_seconds,
+    )
+    retry_state = SimpleNamespace(attempt_number=max(1, consecutive_retryable_failures))
+    return float(strategy(retry_state))  # type: ignore[arg-type]
+
+
+class AdaptiveBatcher:
+    """Sends items in adaptively-sized batches, growing on success and shrinking on failure."""
+
+    def __init__(
+        self,
+        is_retryable: Callable[[Exception], bool],
+        config: Optional[AdaptiveBatchConfig] = None,
+        progress_label: str = "Uploading",
+    ):
+        self._config = config if config is not None else AdaptiveBatchConfig.from_values()
+        self._is_retryable = is_retryable
+        self._progress_label = progress_label
+
+    def run(self, items: Iterable[T], operation: Callable[[List[T]], None]) -> None:
+        total: Optional[int] = len(items) if isinstance(items, Sized) else None
+        if total == 0:
+            return
+
+        config = self._config
+        current_batch_size = config.initial_batch_size
+        it = iter(items)
+        pending: List[T] = []
+
+        progress = get_rich_progress(rich.progress.MofNCompleteColumn())
+        total_task = progress.add_task(
+            f"{self._progress_label} (adaptive batch {config.min_batch_size}-{config.max_batch_size})...",
+            total=total,
+        )
+
+        last_good_batch_size: Optional[int] = None
+        last_bad_batch_size: Optional[int] = None
+        consecutive_retryable_failures = 0
+        processed = 0
+
+        with progress:
+            while True:
+                # Draw from pending (failed-batch leftovers) first, then the source iterator
+                batch = pending[:current_batch_size]
+                pending = pending[current_batch_size:]
+                if len(batch) < current_batch_size:
+                    batch.extend(itertools.islice(it, current_batch_size - len(batch)))
+                if not batch:
+                    break
+                batch_size = len(batch)
+
+                progress.update(
+                    total_task,
+                    description=f"{self._progress_label} (batch size {batch_size})...",
+                )
+                logger.debug(f"{self._progress_label}: {batch_size} entries...")
+
+                start_time = time.monotonic()
+                try:
+                    operation(batch)
+                except Exception as exc:
+                    if not self._is_retryable(exc):
+                        logger.error(
+                            f"{self._progress_label} failed with a non-retryable error; aborting.",
+                            exc_info=True,
+                        )
+                        raise
+
+                    if batch_size <= config.min_batch_size:
+                        logger.error(
+                            f"{self._progress_label} failed at minimum batch size ({batch_size}); aborting.",
+                            exc_info=True,
+                        )
+                        raise
+
+                    consecutive_retryable_failures += 1
+                    time.sleep(_get_retry_delay_seconds(consecutive_retryable_failures, config))
+
+                    last_bad_batch_size = (
+                        batch_size if last_bad_batch_size is None else min(last_bad_batch_size, batch_size)
+                    )
+                    if last_good_batch_size is not None and last_good_batch_size >= last_bad_batch_size:
+                        last_good_batch_size = None
+                    current_batch_size = _next_batch_after_retryable_failure(
+                        batch_size, config, last_good_batch_size, last_bad_batch_size
+                    )
+                    logger.warning(
+                        f"{self._progress_label} failed for batch size {batch_size} "
+                        f"({exc.__class__.__name__}: {exc}). Retrying with batch size {current_batch_size}."
+                    )
+                    # Re-queue the failed batch items for retry with smaller batch size
+                    pending = batch + pending
+                    continue
+
+                elapsed = time.monotonic() - start_time
+                consecutive_retryable_failures = 0
+                processed += batch_size
+                progress.update(total_task, advance=batch_size)
+
+                if elapsed <= config.target_batch_time_seconds:
+                    last_good_batch_size = (
+                        batch_size if last_good_batch_size is None else max(last_good_batch_size, batch_size)
+                    )
+                    # Clear stale bad bound if we succeeded fast at or above it
+                    if last_bad_batch_size is not None and batch_size >= last_bad_batch_size:
+                        last_bad_batch_size = None
+                    current_batch_size = _next_batch_after_success(batch_size, config, last_bad_batch_size)
+                else:
+                    last_bad_batch_size = (
+                        batch_size if last_bad_batch_size is None else min(last_bad_batch_size, batch_size)
+                    )
+                    if last_good_batch_size is not None and last_good_batch_size >= last_bad_batch_size:
+                        last_good_batch_size = None
+                    current_batch_size = _next_batch_after_retryable_failure(
+                        batch_size, config, last_good_batch_size, last_bad_batch_size
+                    )
+
+            progress.update(total_task, completed=processed, total=processed, refresh=True)

dagshub/common/config.py

@@ -58,7 +58,35 @@ def set_host(new_host: str):
 recommended_annotate_limit = int(os.environ.get(RECOMMENDED_ANNOTATE_LIMIT_KEY, 1e5))
 
 DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_KEY = "DAGSHUB_DE_METADATA_UPLOAD_BATCH_SIZE"
-dataengine_metadata_upload_batch_size = int(os.environ.get(DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_KEY, 15000))
+DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_MAX_KEY = "DAGSHUB_DE_METADATA_UPLOAD_BATCH_SIZE_MAX"
+dataengine_metadata_upload_batch_size = int(
+    os.environ.get(
+        DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_MAX_KEY,
+        os.environ.get(DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_KEY, 15000),
+    )
+)
+
+DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_MIN_KEY = "DAGSHUB_DE_METADATA_UPLOAD_BATCH_SIZE_MIN"
+dataengine_metadata_upload_batch_size_min = int(os.environ.get(DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_MIN_KEY, 1))
+
+DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_INITIAL_KEY = "DAGSHUB_DE_METADATA_UPLOAD_BATCH_SIZE_INITIAL"
+dataengine_metadata_upload_batch_size_initial = int(
+    os.environ.get(DATAENGINE_METADATA_UPLOAD_BATCH_SIZE_INITIAL_KEY, dataengine_metadata_upload_batch_size_min)
+)
+
+DATAENGINE_METADATA_UPLOAD_TARGET_BATCH_TIME_SECONDS_KEY = "DAGSHUB_DE_METADATA_UPLOAD_TARGET_BATCH_TIME_SECONDS"
+dataengine_metadata_upload_target_batch_time_seconds = float(
+    os.environ.get(DATAENGINE_METADATA_UPLOAD_TARGET_BATCH_TIME_SECONDS_KEY, 5.0)
+)
+
+ADAPTIVE_BATCH_GROWTH_FACTOR_KEY = "DAGSHUB_ADAPTIVE_BATCH_GROWTH_FACTOR"
+adaptive_batch_growth_factor = int(os.environ.get(ADAPTIVE_BATCH_GROWTH_FACTOR_KEY, 10))
+
+ADAPTIVE_BATCH_RETRY_BACKOFF_BASE_KEY = "DAGSHUB_ADAPTIVE_BATCH_RETRY_BACKOFF_BASE"
+adaptive_batch_retry_backoff_base_seconds = float(os.environ.get(ADAPTIVE_BATCH_RETRY_BACKOFF_BASE_KEY, 0.25))
+
+ADAPTIVE_BATCH_RETRY_BACKOFF_MAX_KEY = "DAGSHUB_ADAPTIVE_BATCH_RETRY_BACKOFF_MAX"
+adaptive_batch_retry_backoff_max_seconds = float(os.environ.get(ADAPTIVE_BATCH_RETRY_BACKOFF_MAX_KEY, 4.0))
 
 DISABLE_ANALYTICS_KEY = "DAGSHUB_DISABLE_ANALYTICS"
 disable_analytics = "DAGSHUB_DISABLE_ANALYTICS" in os.environ

dagshub/data_engine/model/datasource.py

@@ -14,7 +14,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Dict, List, Literal, Optional, Set, Tuple, Union
 
-import rich.progress
 from dataclasses_json import DataClassJsonMixin, LetterCase, config
 from pathvalidate import sanitize_filepath
 
@@ -53,6 +52,8 @@
     run_preupload_transforms,
     validate_uploading_metadata,
 )
+from dagshub.common.adaptive_batching import AdaptiveBatcher
+from dagshub.data_engine.model.metadata.util import is_retryable_metadata_upload_error
 from dagshub.data_engine.model.metadata.dtypes import DatapointMetadataUpdateEntry
 from dagshub.data_engine.model.metadata.transforms import DatasourceFieldInfo, _add_metadata
 from dagshub.data_engine.model.metadata_field_builder import MetadataFieldBuilder
@@ -753,19 +754,11 @@ def _upload_metadata(self, metadata_entries: List[DatapointMetadataUpdateEntry])
         validate_uploading_metadata(precalculated_info)
         run_preupload_transforms(self, metadata_entries, precalculated_info)
 
-        progress = get_rich_progress(rich.progress.MofNCompleteColumn())
-
-        upload_batch_size = dagshub.common.config.dataengine_metadata_upload_batch_size
-        total_entries = len(metadata_entries)
-        total_task = progress.add_task(f"Uploading metadata (batch size {upload_batch_size})...", total=total_entries)
-
-        with progress:
-            for start in range(0, total_entries, upload_batch_size):
-                entries = metadata_entries[start : start + upload_batch_size]
-                logger.debug(f"Uploading {len(entries)} metadata entries...")
-                self.source.client.update_metadata(self, entries)
-                progress.update(total_task, advance=upload_batch_size)
-            progress.update(total_task, completed=total_entries, refresh=True)
+        batcher = AdaptiveBatcher(
+            is_retryable=is_retryable_metadata_upload_error,
+            progress_label="Uploading metadata",
+        )
+        batcher.run(metadata_entries, lambda batch: self.source.client.update_metadata(self, batch))
 
         # Update the status from dagshub, so we get back the new metadata columns
         self.source.get_from_dagshub()

dagshub/data_engine/model/metadata/util.py

@@ -1,6 +1,10 @@
 import datetime
+from gql.transport.exceptions import TransportServerError, TransportConnectionFailed
+from requests import ConnectionError as RequestsConnectionError, Timeout as RequestsTimeout
 from typing import Optional
 
+from dagshub.data_engine.model.errors import DataEngineGqlError
+
 
 def _get_datetime_utc_offset(t: datetime.datetime) -> Optional[str]:
     """
@@ -19,3 +23,20 @@ def _get_datetime_utc_offset(t: datetime.datetime) -> Optional[str]:
     offset_minutes = int((offset.total_seconds() % 3600) // 60)
     offset_str = f"{offset_hours:+03d}:{offset_minutes:02d}"
     return offset_str
+
+
+def is_retryable_metadata_upload_error(exc: Exception) -> bool:
+    if isinstance(exc, DataEngineGqlError) and isinstance(exc.original_exception, Exception):
+        return is_retryable_metadata_upload_error(exc.original_exception)
+
+    return isinstance(
+        exc,
+        (
+            TransportServerError,
+            TransportConnectionFailed,
+            TimeoutError,
+            ConnectionError,
+            RequestsConnectionError,
+            RequestsTimeout,
+        ),
+    )