Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions weaviate_cli/commands/create.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)
Expand All @@ -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,
Expand Down Expand Up @@ -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}")
Expand Down
3 changes: 3 additions & 0 deletions weaviate_cli/defaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
88 changes: 72 additions & 16 deletions weaviate_cli/managers/data_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -30,6 +31,11 @@
UpdateDataDefaults,
DeleteDataDefaults,
)

RETRY_EXCEPTIONS = (
grpc.RpcError,
Exception,
)
from weaviate_cli.utils import pp_objects

PROPERTY_NAME_MAPPING = {
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
Loading