diff --git a/weaviate_cli/commands/create.py b/weaviate_cli/commands/create.py index 9400576..a38f909 100644 --- a/weaviate_cli/commands/create.py +++ b/weaviate_cli/commands/create.py @@ -541,6 +541,24 @@ def create_backup_cli( type=click.IntRange(min=1), help=f"Number of tenants to process in parallel (default: {CreateDataDefaults.parallel_workers}). Set to 1 to disable parallelism.", ) +@click.option( + "--max_retries", + default=CreateDataDefaults.max_retries, + type=int, + help=f"Maximum number of retries for transient errors like shard not found (default: {CreateDataDefaults.max_retries}).", +) +@click.option( + "--retry_initial_delay", + default=CreateDataDefaults.retry_initial_delay, + type=float, + help=f"Initial delay between retries in seconds (default: {CreateDataDefaults.retry_initial_delay}).", +) +@click.option( + "--retry_max_delay", + default=CreateDataDefaults.retry_max_delay, + type=float, + help=f"Maximum delay between retries in seconds (default: {CreateDataDefaults.retry_max_delay}).", +) @click.option( "--json", "json_output", is_flag=True, default=False, help="Output in JSON format." ) @@ -557,6 +575,9 @@ def create_data_cli( tenant_suffix, vector_dimensions, uuid, + max_retries, + retry_initial_delay, + retry_max_delay, wait_for_indexing, verbose, multi_vector, @@ -616,6 +637,9 @@ def create_data_cli( concurrent_requests=concurrent_requests, parallel_workers=parallel_workers, json_output=json_output, + max_retries=max_retries, + retry_initial_delay=retry_initial_delay, + retry_max_delay=retry_max_delay, ) except Exception as e: click.echo(f"Error: {e}") diff --git a/weaviate_cli/defaults.py b/weaviate_cli/defaults.py index fbd6c71..0fd93dc 100644 --- a/weaviate_cli/defaults.py +++ b/weaviate_cli/defaults.py @@ -125,6 +125,9 @@ class CreateDataDefaults: batch_size: int = 1000 dynamic_batch: bool = False parallel_workers: int = MAX_WORKERS + max_retries: int = 3 + retry_initial_delay: float = 1.0 + retry_max_delay: float = 30.0 @dataclass diff --git a/weaviate_cli/managers/data_manager.py b/weaviate_cli/managers/data_manager.py index d2fcc94..458c87d 100644 --- a/weaviate_cli/managers/data_manager.py +++ b/weaviate_cli/managers/data_manager.py @@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta from typing import Dict, List, Optional, Union, Any, Tuple +import grpc import click import numpy as np @@ -30,6 +31,11 @@ UpdateDataDefaults, DeleteDataDefaults, ) + +RETRY_EXCEPTIONS = ( + grpc.RpcError, + Exception, +) from weaviate_cli.utils import pp_objects PROPERTY_NAME_MAPPING = { @@ -659,13 +665,46 @@ def __ingest_data( batch_size: int = 1000, concurrent_requests: int = MAX_WORKERS, json_output: bool = False, + max_retries: int = CreateDataDefaults.max_retries, + retry_initial_delay: float = CreateDataDefaults.retry_initial_delay, + retry_max_delay: float = CreateDataDefaults.retry_max_delay, ) -> Collection: + def _is_retryable_error(e: Exception) -> bool: + error_str = str(e).lower() + retryable_patterns = [ + "shard not found", + "shard unavailable", + "temporarily unavailable", + "service unavailable", + "connection reset", + "connection refused", + "deadline exceeded", + " UNAVAILABLE", + " CANCELLED", + ] + return any(pattern.lower() in error_str for pattern in retryable_patterns) + + def _do_ingest() -> Tuple[int, List, _ErrorTracker]: + return self.__producer_consumer_ingest( + collection=cl_collection, + num_objects=num_objects, + vectorizer=vectorizer, + vector_dimensions=vector_dimensions or 1536, + named_vectors=named_vectors, + uuid=uuid, + dynamic_batch=dynamic_batch, + batch_size=batch_size, + concurrent_requests=concurrent_requests, + multi_vector=multi_vector, + skip_seed=skip_seed, + verbose=verbose, + ) + if randomize: if not json_output: click.echo(f"Generating and ingesting {num_objects} objects") start_time = time.time() - # Determine vector dimensions based on vectorizer config = collection.config.get() if not config.vectorizer and config.vector_config: @@ -682,21 +721,29 @@ def __ingest_data( cl_collection = collection.with_consistency_level(cl) - # Single consumer that feeds the batcher; batcher does its own HTTP parallelism - counter, failed_objects, error_tracker = self.__producer_consumer_ingest( - collection=cl_collection, - num_objects=num_objects, - vectorizer=vectorizer, - vector_dimensions=vector_dimensions or 1536, - named_vectors=named_vectors, - uuid=uuid, - dynamic_batch=dynamic_batch, - batch_size=batch_size, - concurrent_requests=concurrent_requests, - multi_vector=multi_vector, - skip_seed=skip_seed, - verbose=verbose, - ) + delay = retry_initial_delay + last_error = None + for attempt in range(max_retries + 1): + try: + counter, _, error_tracker = _do_ingest() + break + except Exception as e: + last_error = e + if attempt < max_retries and _is_retryable_error(e): + wait_time = min(delay, retry_max_delay) + if verbose: + click.echo( + f"Retryable error encountered (attempt {attempt + 1}/{max_retries + 1}): {e}. " + f"Retrying in {wait_time:.1f}s...", + err=True, + ) + time.sleep(wait_time) + delay *= 2 + else: + raise + + if last_error and "counter" not in locals(): + raise last_error # New compact failure summary (replaces old block) if error_tracker.total > 0: @@ -828,6 +875,9 @@ def create_data( concurrent_requests: int = MAX_WORKERS, parallel_workers: int = CreateDataDefaults.parallel_workers, json_output: bool = False, + max_retries: int = CreateDataDefaults.max_retries, + retry_initial_delay: float = CreateDataDefaults.retry_initial_delay, + retry_max_delay: float = CreateDataDefaults.retry_max_delay, ) -> Collection: if not self.client.collections.exists(collection): @@ -901,6 +951,9 @@ def _ingest_one_tenant(tenant: str): batch_size=batch_size, concurrent_requests=effective_concurrent, json_output=json_output, + max_retries=max_retries, + retry_initial_delay=retry_initial_delay, + retry_max_delay=retry_max_delay, ) _after = len(col) else: @@ -937,6 +990,9 @@ def _ingest_one_tenant(tenant: str): batch_size=batch_size, concurrent_requests=effective_concurrent, json_output=json_output, + max_retries=max_retries, + retry_initial_delay=retry_initial_delay, + retry_max_delay=retry_max_delay, ) _after = len(col.with_tenant(tenant)) if wait_for_indexing: