diff --git a/.github/workflows/backend-tests-postgres.yml b/.github/workflows/backend-tests-postgres.yml index 222c61f8..3ab31e7a 100644 --- a/.github/workflows/backend-tests-postgres.yml +++ b/.github/workflows/backend-tests-postgres.yml @@ -59,3 +59,63 @@ jobs: - name: Run backend test suite on Postgres run: make test + + db-sharding: + runs-on: ubuntu-latest + timeout-minutes: 20 + + services: + postgres: + image: postgres:16 + env: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: efficientai_test + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U postgres -d efficientai_test" + --health-interval 10s + --health-timeout 5s + --health-retries 10 + + env: + PGHOST: localhost + PGPORT: 5432 + PGUSER: postgres + PGPASSWORD: postgres + DATABASE_URL: postgresql://postgres:postgres@localhost:5432/efficientai_catalog + TEST_DATABASE_URL: postgresql://postgres:postgres@localhost:5432/efficientai_test + CATALOG_DATABASE_URL: postgresql://postgres:postgres@localhost:5432/efficientai_catalog + SHARD_DATABASE_URL_01: postgresql://postgres:postgres@localhost:5432/efficientai_data_01 + SHARD_DATABASE_URL_02: postgresql://postgres:postgres@localhost:5432/efficientai_data_02 + SHARDING_INTEGRATION_TEST: "1" + REDIS_URL: redis://localhost:6379/0 + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: "pip" + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[dev]" + + - name: Create catalog and row-shard databases + run: | + for db in efficientai_catalog efficientai_data_01 efficientai_data_02; do + psql -d postgres -v ON_ERROR_STOP=1 -c "DROP DATABASE IF EXISTS ${db};" + psql -d postgres -v ON_ERROR_STOP=1 -c "CREATE DATABASE ${db};" + done + + - name: Run db_sharding unit tests + run: make test-sharding + + - name: Run sharding Postgres integration tests + run: make test-sharding-integration diff --git a/Makefile b/Makefile index 1bf7d429..c071e76e 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install-dev check-pytest test test-docker-db test-unit test-integration test-phase1 test-file test-k +.PHONY: help install-dev check-pytest test test-docker-db test-unit test-integration test-phase1 test-file test-k test-sharding test-sharding-integration PYTHON ?= python PYTEST ?= $(PYTHON) -m pytest @@ -19,7 +19,8 @@ help: ## Show available make targets @echo " make test-integration - run integration tests (marker: integration)" @echo " make test-phase1 - run current Phase 1 suites" @echo " make test-file FILE=...- run a specific test file/path" - @echo " make test-k K=... - run tests matching expression" + @echo " make test-sharding - run call-import db_sharding tests (unit; no integration marker)" + @echo " make test-sharding-integration - run 2-shard Postgres integration tests (CI)" install-dev: ## Install project and dev dependencies $(PYTHON) -m pip install -e ".[dev]" @@ -57,3 +58,10 @@ test-file: check-pytest ## Run one test module/file; usage: make test-file FILE= test-k: check-pytest ## Run tests by keyword expression; usage: make test-k K=password @if [ -z "$(K)" ]; then echo "K is required. Example: make test-k K=password"; exit 1; fi $(PYTEST) tests -k "$(K)" $(PYTEST_FLAGS) $(PYTEST_ARGS) + +test-sharding: check-pytest ## Run db_sharding unit tests (excludes integration marker) + $(PYTEST) tests/test_db_sharding -m "not integration" $(PYTEST_FLAGS) $(PYTEST_ARGS) + +test-sharding-integration: check-pytest ## Run 2-shard Postgres sharding integration tests + @if [ -z "$$SHARDING_INTEGRATION_TEST" ]; then export SHARDING_INTEGRATION_TEST=1; fi + $(PYTEST) tests/test_db_sharding/test_sharding_postgres_integration.py -m integration $(PYTEST_FLAGS) $(PYTEST_ARGS) diff --git a/app/api/v1/routes/call_import_evaluations.py b/app/api/v1/routes/call_import_evaluations.py index cf212589..ef316f2e 100644 --- a/app/api/v1/routes/call_import_evaluations.py +++ b/app/api/v1/routes/call_import_evaluations.py @@ -486,6 +486,15 @@ def _serialize_eval( db, row.organization_id, metric_ids_for_lookup ) + from app.services.call_imports.progress_counters import merge_eval_counters_for_ui + + ui_completed_raw, ui_failed_raw = merge_eval_counters_for_ui(row) + total = int(row.total_rows or 0) + ui_completed = ( + min(ui_completed_raw, total) if total else ui_completed_raw + ) + ui_failed = min(ui_failed_raw, total) if total else ui_failed_raw + return CallImportEvaluationResponse( id=row.id, call_import_id=row.call_import_id, @@ -513,8 +522,8 @@ def _serialize_eval( ], status=row.status, total_rows=row.total_rows, - completed_rows=row.completed_rows, - failed_rows=row.failed_rows, + completed_rows=ui_completed, + failed_rows=ui_failed, error_message=row.error_message, llm_provider=row.llm_provider, llm_model=row.llm_model, @@ -1505,6 +1514,7 @@ def _apply_direction(column_expr): sort_recognized = False sort_by_clean = (sort_by or "").strip() primary_sort = None + metric_uuid: Optional[UUID] = None if sort_by_clean == "row_index": sort_recognized = True # Falls through to the default ``order_by`` below with @@ -1570,8 +1580,83 @@ def _apply_direction(column_expr): # so a typo'd / stale ``sort_by`` doesn't quietly invert the # default order. query = query.order_by(CallImportRow.row_index.asc()) - total = query.count() - rows = query.offset((page - 1) * page_size).limit(page_size).all() + from app.db_sharding.eval_rows import fetch_evaluation_row_pairs_page + from app.db_sharding.sessions import is_sharding_enabled + + def _pair_row_index( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> int: + return int(pair[1].row_index or 0) + + def _directed_string(value: Optional[str], desc: bool) -> Tuple[int, ...]: + text = value or "" + if not desc: + return (0, *text.encode("utf-8")) + return (1, *(-byte for byte in text.encode("utf-8"))) + + if sort_by_clean == "conversation_id": + def _pair_sort_key( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> Tuple[Any, ...]: + return ( + _directed_string(pair[1].conversation_id, direction_desc), + _pair_row_index(pair), + ) + elif sort_by_clean == "status": + def _pair_sort_key( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> Tuple[Any, ...]: + return ( + _directed_string(pair[0].status, direction_desc), + _pair_row_index(pair), + ) + elif sort_by_clean.startswith("metric:") and metric_uuid is not None: + metric_id_str = str(metric_uuid) + + def _pair_sort_key( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> Tuple[Any, ...]: + scores = pair[0].metric_scores or {} + entry = scores.get(metric_id_str, {}) + raw_value = entry.get("value") if isinstance(entry, dict) else None + null_rank = 1 if raw_value is None else 0 + return ( + null_rank, + _directed_string( + str(raw_value) if raw_value is not None else None, + direction_desc, + ), + _pair_row_index(pair), + ) + elif sort_recognized and sort_by_clean == "row_index": + def _pair_sort_key( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> Tuple[int, ...]: + idx = _pair_row_index(pair) + return (-idx,) if direction_desc else (idx,) + else: + def _pair_sort_key( + pair: Tuple[CallImportEvaluationRow, CallImportRow], + ) -> Tuple[int, ...]: + return (_pair_row_index(pair),) + + if is_sharding_enabled(): + def _build_query(session: Session): + return query.with_session(session) + + total, rows = fetch_evaluation_row_pairs_page( + db, + _build_query, + page=page, + page_size=page_size, + sort_key=_pair_sort_key, + bounded_shard_fetch=( + not sort_recognized or sort_by_clean == "row_index" + ), + ) + else: + total = query.count() + rows = query.offset((page - 1) * page_size).limit(page_size).all() # Row detail shows the diarised transcript that normal metrics score. items: List[CallImportEvaluationRowResponse] = [ @@ -1782,13 +1867,25 @@ def _add_metric_column(metric: Metric) -> None: *metric_headers, ] - rows = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join(CallImportRow, CallImportRow.id == CallImportEvaluationRow.call_import_row_id) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .order_by(CallImportRow.row_index.asc()) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + rows = sorted( + load_evaluation_row_pairs(db, eval_id), + key=lambda pair: int(pair[1].row_index or 0), + ) + else: + rows = ( + db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.evaluation_id == eval_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) def _project_rows() -> Iterator[Dict[str, str]]: for eval_row, source_row in rows: @@ -3258,13 +3355,25 @@ async def generate_call_import_evaluation_pdf_report( raise HTTPException(status_code=404, detail="Call import evaluation not found") is_internal = payload.report_type == "internal" - rows = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join(CallImportRow, CallImportRow.id == CallImportEvaluationRow.call_import_row_id) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .order_by(CallImportRow.row_index.asc()) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + rows = sorted( + load_evaluation_row_pairs(db, eval_id), + key=lambda pair: int(pair[1].row_index or 0), + ) + else: + rows = ( + db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.evaluation_id == eval_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) report_config = payload.report_config if isinstance(payload.report_config, dict) else {} metrics = _display_metrics_for_pdf_report(db, organization_id, evaluation) configured_quality_ids = { @@ -3727,6 +3836,8 @@ async def cancel_call_import_evaluation( ) _claim_evaluation_bulk_operation(eval_id, "abort") + evaluation.status = "cancelled" + db.commit() from app.workers.tasks.call_import_bulk_ops import ( cancel_call_import_evaluation_task, @@ -3848,6 +3959,34 @@ async def cancel_call_import_evaluation_row( _require_no_evaluation_bulk_operation(eval_id) + from app.db_sharding.eval_rows import evaluation_row_session + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + try: + with evaluation_row_session(eval_row_id) as ( + row_db, + _catalog_db, + eval_row, + source_row, + _shard_id, + ): + if eval_row.evaluation_id != eval_id: + raise HTTPException( + status_code=404, + detail="Evaluation row not found in this run", + ) + _apply_evaluation_cancel([eval_row]) + row_db.commit() + _rollup_evaluation_status(evaluation, db) + db.commit() + row_db.refresh(eval_row) + return _to_evaluation_row_response(eval_row, source_row) + except LookupError as exc: + raise HTTPException( + status_code=404, detail="Evaluation row not found in this run" + ) from exc + eval_row = ( db.query(CallImportEvaluationRow) .filter( @@ -4337,11 +4476,7 @@ async def get_call_import_evaluation_aggregate( status_code=404, detail="Call import evaluation not found" ) - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) metrics = _compute_metric_aggregates(db, evaluation, eval_rows) @@ -4349,15 +4484,9 @@ async def get_call_import_evaluation_aggregate( resolved_baseline_id: Optional[UUID] = None if baseline_evaluation_id is not None: call_import = _require_import(db, call_import_id, organization_id) - rows = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join( - CallImportRow, - CallImportRow.id == CallImportEvaluationRow.call_import_row_id, - ) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + + rows = load_evaluation_row_pairs(db, eval_id) period_start, _, _, _ = _report_period_from_rows(rows) baseline_evaluation = _resolve_baseline_evaluation( db, @@ -4737,11 +4866,10 @@ def _generate_and_persist_tldr_summary( ) -> EvaluationTldrSummary: """LLM TLDR generation used by the imports-queue Celery worker.""" eval_id = evaluation.id - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + + pairs = load_evaluation_row_pairs(db, eval_id) + eval_rows = [eval_row for eval_row, _ in pairs] aggregate = _compute_metric_aggregates(db, evaluation, eval_rows) if not aggregate: raise HTTPException( @@ -4790,7 +4918,11 @@ def _generate_and_persist_tldr_summary( ) from e summary = _parse_insights_response(llm_result.get("text", "")) - summary.generated_at_completed_rows = evaluation.completed_rows + total = int(evaluation.total_rows or 0) + ui_completed = min(int(evaluation.completed_rows or 0), total) if total else int( + evaluation.completed_rows or 0 + ) + summary.generated_at_completed_rows = ui_completed summary.provider = provider_enum.value summary.model = model_str summary.is_stale = False @@ -5008,12 +5140,7 @@ def _enqueue_user_insights_job( llm_budget = normalize_max_llm_calls(max_llm_calls) completed_count = ( - db.query(CallImportEvaluationRow) - .filter( - CallImportEvaluationRow.evaluation_id == evaluation.id, - CallImportEvaluationRow.status == "completed", - ) - .count() + _count_completed_eval_rows(db, evaluation.id) if db is not None else evaluation.completed_rows ) @@ -5118,11 +5245,7 @@ async def generate_call_import_evaluation_user_insights( if cached is not None and cached.status in {"running", "completed"}: return cached - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) if not any(row.status == "completed" for row in eval_rows): raise HTTPException( status_code=400, @@ -5268,19 +5391,27 @@ def _enqueue_prompt_improvements_job( db.commit() +def _load_eval_rows(db: Session, evaluation_id: UUID) -> List[CallImportEvaluationRow]: + from app.db_sharding.eval_rows import load_evaluation_rows_for_run + + return load_evaluation_rows_for_run(db, evaluation_id) + + +def _count_completed_eval_rows(db: Session, evaluation_id: UUID) -> int: + from app.db_sharding.eval_rows import count_evaluation_rows_for_run + + return count_evaluation_rows_for_run( + db, evaluation_id, statuses=["completed"] + ) + + def _completed_row_pairs_for_evaluation( db: Session, evaluation_id: UUID, ) -> List[Tuple[CallImportEvaluationRow, CallImportRow]]: - row_pairs = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join( - CallImportRow, - CallImportRow.id == CallImportEvaluationRow.call_import_row_id, - ) - .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + + row_pairs = load_evaluation_row_pairs(db, evaluation_id) return [ (eval_row, source_row) for eval_row, source_row in row_pairs @@ -5293,6 +5424,8 @@ def _resolve_metric_cluster_row_selection( evaluation: CallImportEvaluation, eval_rows: List[CallImportEvaluationRow], evaluation_row_ids: Optional[List[UUID]], + *, + row_limit: Optional[int] = None, policies: Optional[Dict[str, MetricFailurePolicy]] = None, ) -> Tuple[List[Tuple[CallImportEvaluationRow, CallImportRow]], List[str]]: """Return filtered completed row pairs and the selected row id strings.""" @@ -5318,10 +5451,19 @@ def _resolve_metric_cluster_row_selection( eligible = list_eligible_cluster_rows( evaluation, completed_pairs, metrics, policies ) - eligible_id_set = {str(item["evaluation_row_id"]) for item in eligible} + eligible_ordered_ids = [str(item["evaluation_row_id"]) for item in eligible] + eligible_id_set = set(eligible_ordered_ids) + + if evaluation_row_ids is None and row_limit is not None: + selected_ids = eligible_ordered_ids[:row_limit] + filtered = filter_completed_row_pairs( + completed_pairs, + [UUID(rid) for rid in selected_ids], + ) + return filtered, selected_ids if evaluation_row_ids is None: - selected_ids = sorted(eligible_id_set) + selected_ids = eligible_ordered_ids filtered = filter_completed_row_pairs( completed_pairs, [UUID(rid) for rid in selected_ids], @@ -5377,11 +5519,7 @@ def _enqueue_metric_clusters_job( total_calls = 1 row_ids_for_task: Optional[List[str]] = None if db is not None: - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == evaluation.id) - .all() - ) + eval_rows = _load_eval_rows(db, evaluation.id) if selected_evaluation_row_ids is None: _, selected_evaluation_row_ids = _resolve_metric_cluster_row_selection( db, @@ -5552,11 +5690,7 @@ async def get_call_import_evaluation_metric_cluster_failure_policies( status_code=404, detail="Call import evaluation not found" ) - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) metrics, aggregates, policies, source, child_names_by_parent = _clustering_context( db, evaluation, eval_rows ) @@ -5613,11 +5747,7 @@ async def save_call_import_evaluation_metric_cluster_failure_policies( status_code=404, detail="Call import evaluation not found" ) - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) metrics, aggregates, _existing, _source, child_names_by_parent = _clustering_context( db, evaluation, eval_rows ) @@ -5674,6 +5804,8 @@ async def save_call_import_evaluation_metric_cluster_failure_policies( async def list_call_import_evaluation_metric_cluster_eligible_rows( call_import_id: UUID, eval_id: UUID, + limit: Optional[int] = Query(default=None, ge=1), + count_only: bool = Query(default=False), api_key: str = Depends(get_api_key), organization_id: UUID = Depends(get_organization_id), db: Session = Depends(get_db), @@ -5696,20 +5828,20 @@ async def list_call_import_evaluation_metric_cluster_eligible_rows( status_code=404, detail="Call import evaluation not found" ) - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) completed_pairs = _completed_row_pairs_for_evaluation(db, eval_id) metrics, _aggregates, policies, _source, _child_map = _clustering_context( db, evaluation, eval_rows ) - raw_items = list_eligible_cluster_rows( + all_eligible = list_eligible_cluster_rows( evaluation, completed_pairs, metrics, policies ) + total = len(all_eligible) + if count_only: + return MetricClusterEligibleRowsResponse(items=[], total=total) + raw_items = all_eligible if limit is None else all_eligible[:limit] items = [MetricClusterEligibleRow.model_validate(item) for item in raw_items] - return MetricClusterEligibleRowsResponse(items=items, total=len(items)) + return MetricClusterEligibleRowsResponse(items=items, total=total) @router.get( @@ -5780,11 +5912,7 @@ async def generate_call_import_evaluation_metric_clusters( if cached is not None and cached.status in {"running", "completed"}: return cached - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) if not any(row.status == "completed" for row in eval_rows): raise HTTPException( status_code=400, @@ -5794,6 +5922,12 @@ async def generate_call_import_evaluation_metric_clusters( ), ) + if body.evaluation_row_ids and body.row_limit is not None: + raise HTTPException( + status_code=400, + detail="Specify either evaluation_row_ids or row_limit, not both.", + ) + if body.evaluation_row_ids: completed_pairs = _completed_row_pairs_for_evaluation(db, evaluation.id) completed_id_set = {str(eval_row.id) for eval_row, _ in completed_pairs} @@ -5847,6 +5981,7 @@ async def generate_call_import_evaluation_metric_clusters( evaluation, eval_rows, body.evaluation_row_ids, + row_limit=body.row_limit, policies=merged_policies, ) if not selected_row_ids: @@ -6379,15 +6514,14 @@ def _get_running_discovered_labels( """ parent_id_str = str(parent_metric_id) - rows = ( - db.query(CallImportEvaluationRow.metric_scores) - .filter( - CallImportEvaluationRow.evaluation_id == eval_id, - CallImportEvaluationRow.status - == CallImportRowStatus.COMPLETED.value, - ) - .all() - ) + from app.db_sharding.eval_rows import load_evaluation_rows_for_run + + eval_rows = load_evaluation_rows_for_run(db, eval_id) + rows = [ + (row.metric_scores,) + for row in eval_rows + if row.status == CallImportRowStatus.COMPLETED.value + ] # Suppress slugs that have either: # * been promoted to a real child of the parent (so the panel doesn't @@ -6499,15 +6633,14 @@ def _get_running_discovered_metrics( "count": 12} """ - rows = ( - db.query(CallImportEvaluationRow.metric_scores) - .filter( - CallImportEvaluationRow.evaluation_id == eval_id, - CallImportEvaluationRow.status - == CallImportRowStatus.COMPLETED.value, - ) - .all() - ) + from app.db_sharding.eval_rows import load_evaluation_rows_for_run + + eval_rows = load_evaluation_rows_for_run(db, eval_id) + rows = [ + (row.metric_scores,) + for row in eval_rows + if row.status == CallImportRowStatus.COMPLETED.value + ] promoted_slugs: set[str] = set() if organization_id is not None: @@ -6935,11 +7068,7 @@ async def get_call_import_evaluation_flow( ) extra_children = [c for c in all_children if c.id not in existing_ids] - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + eval_rows = _load_eval_rows(db, eval_id) alias_map = _alias_map_for_parent(evaluation, parent.id) return _build_flow_graph( @@ -7101,28 +7230,21 @@ async def merge_call_import_evaluation_discovered_labels( ) parent_id_str = str(parent.id) - rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.eval_rows import foreach_evaluation_row_mutating - for row in rows: + def _merge_discovered_label_row(row: CallImportEvaluationRow) -> bool: scores = ( row.metric_scores if isinstance(row.metric_scores, dict) else None ) if not scores: - continue + return False parent_entry = scores.get(parent_id_str) if not isinstance(parent_entry, dict): - continue + return False mutated = False - # 1. Rewrite the discovered_labels list. If the row already has - # an entry for to_key we keep that (it carries the user-chosen - # name + sample rationale) and just drop the from_key entry. discovered = parent_entry.get("discovered_labels") if isinstance(discovered, list): kept: List[Dict[str, Any]] = [] @@ -7153,9 +7275,6 @@ async def merge_call_import_evaluation_discovered_labels( if mutated: parent_entry["discovered_labels"] = kept - # 2. Rewrite any discovered slugs inside the sequence array so - # the flow chart stays consistent. Dedupe so we don't end up - # with adjacent identical entries. seq = parent_entry.get("sequence") if isinstance(seq, list): new_seq: List[str] = [] @@ -7178,15 +7297,10 @@ async def merge_call_import_evaluation_discovered_labels( mutated = True if mutated: - # ``JSON`` columns aren't auto-tracked when the same dict is - # mutated in place — ``scores`` IS ``row.metric_scores``, so - # the in-place edits above already updated SQLAlchemy's - # cached "committed" snapshot to the post-edit dict. Without - # ``flag_modified`` the subsequent ``dict(scores)`` reassign - # compares equal to that snapshot and SQLAlchemy skips the - # UPDATE, leaving the per-row payload stale on disk. row.metric_scores = dict(scores) - flag_modified(row, "metric_scores") + return mutated + + foreach_evaluation_row_mutating(db, eval_id, _merge_discovered_label_row) # Persist the merge at the evaluation level too. This is what makes # the merge survive future scoring: rows that finish AFTER this @@ -7301,27 +7415,21 @@ async def delete_call_import_evaluation_discovered_label( ) parent_id_str = str(parent.id) - rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.eval_rows import foreach_evaluation_row_mutating - for row in rows: + def _delete_discovered_label_row(row: CallImportEvaluationRow) -> bool: scores = ( row.metric_scores if isinstance(row.metric_scores, dict) else None ) if not scores: - continue + return False parent_entry = scores.get(parent_id_str) if not isinstance(parent_entry, dict): - continue + return False mutated = False - - # 1. Strip the slug from discovered_labels. discovered = parent_entry.get("discovered_labels") if isinstance(discovered, list): kept = [ @@ -7337,8 +7445,6 @@ async def delete_call_import_evaluation_discovered_label( parent_entry["discovered_labels"] = kept mutated = True - # 2. Strip the slug from the sequence array, collapsing adjacent - # duplicates that the deletion exposes. seq = parent_entry.get("sequence") if isinstance(seq, list): new_seq: List[str] = [] @@ -7360,12 +7466,10 @@ async def delete_call_import_evaluation_discovered_label( mutated = True if mutated: - # See merge endpoint above: in-place edits to ``scores`` / - # ``parent_entry`` already mutated SQLAlchemy's committed - # snapshot, so the reassign alone wouldn't trigger an - # UPDATE. Flagging the column forces it. row.metric_scores = dict(scores) - flag_modified(row, "metric_scores") + return mutated + + foreach_evaluation_row_mutating(db, eval_id, _delete_discovered_label_row) # 3. Persist the tombstone on the evaluation so workers that finish # later don't re-surface the deleted slug. We also retarget any @@ -7535,23 +7639,19 @@ async def merge_call_import_evaluation_discovered_metrics( items=[DiscoveredMetricItem(**item) for item in items_raw], ) - rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.eval_rows import foreach_evaluation_row_mutating - for row in rows: + def _merge_discovered_metric_row(row: CallImportEvaluationRow) -> bool: scores = ( row.metric_scores if isinstance(row.metric_scores, dict) else None ) if not scores: - continue + return False discovered = scores.get(DISCOVERED_METRICS_KEY) if not isinstance(discovered, list): - continue + return False kept: List[Dict[str, Any]] = [] mutated = False @@ -7582,7 +7682,9 @@ async def merge_call_import_evaluation_discovered_metrics( if mutated: scores[DISCOVERED_METRICS_KEY] = kept row.metric_scores = dict(scores) - flag_modified(row, "metric_scores") + return mutated + + foreach_evaluation_row_mutating(db, eval_id, _merge_discovered_metric_row) raw_aliases = ( evaluation.discovered_metric_aliases @@ -7650,23 +7752,19 @@ async def delete_call_import_evaluation_discovered_metric( detail="key must be a non-empty slug.", ) - rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == eval_id) - .all() - ) + from app.db_sharding.eval_rows import foreach_evaluation_row_mutating - for row in rows: + def _delete_discovered_metric_row(row: CallImportEvaluationRow) -> bool: scores = ( row.metric_scores if isinstance(row.metric_scores, dict) else None ) if not scores: - continue + return False discovered = scores.get(DISCOVERED_METRICS_KEY) if not isinstance(discovered, list): - continue + return False kept = [ e for e in discovered @@ -7676,13 +7774,16 @@ async def delete_call_import_evaluation_discovered_metric( == target_key ) ] - if len(kept) != len(discovered): - if kept: - scores[DISCOVERED_METRICS_KEY] = kept - else: - scores.pop(DISCOVERED_METRICS_KEY, None) - row.metric_scores = dict(scores) - flag_modified(row, "metric_scores") + if len(kept) == len(discovered): + return False + if kept: + scores[DISCOVERED_METRICS_KEY] = kept + else: + scores.pop(DISCOVERED_METRICS_KEY, None) + row.metric_scores = dict(scores) + return True + + foreach_evaluation_row_mutating(db, eval_id, _delete_discovered_metric_row) raw_aliases = ( evaluation.discovered_metric_aliases @@ -7746,6 +7847,19 @@ async def delete_call_import_evaluation_row( if not evaluation: raise HTTPException(status_code=404, detail="Call import evaluation not found") + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + from app.db_sharding.eval_rows import delete_evaluation_row_on_shards + + if not delete_evaluation_row_on_shards(eval_row_id, eval_id): + raise HTTPException( + status_code=404, detail="Evaluation row not found in this run" + ) + _rollup_evaluation_status(evaluation, db) + db.commit() + return Response(status_code=status.HTTP_204_NO_CONTENT) + eval_row = ( db.query(CallImportEvaluationRow) .filter( @@ -7903,22 +8017,29 @@ def _enqueue_eval_rows_with_optional_transcribe( eval_only_count = 0 transcribe_count = 0 - for eval_row, source_row in eval_rows_with_source: - if _needs_transcribe_for_eval( - evaluation, - source_row, - transcribe_overwrite=transcribe_overwrite, - ): - transcribe_count += 1 - else: - eval_only_count += 1 + if eval_rows_with_source: + for eval_row, source_row in eval_rows_with_source: + if _needs_transcribe_for_eval( + evaluation, + source_row, + transcribe_overwrite=transcribe_overwrite, + ): + transcribe_count += 1 + else: + eval_only_count += 1 - restricted_metric_ids_str: Optional[List[str]] = ( - [str(mid) for mid in restricted_metric_ids] if restricted_metric_ids else None - ) - if restricted_metric_ids_str: - for eval_row, _ in eval_rows_with_source: - store_row_restricted_metrics(eval_row.id, restricted_metric_ids_str) + restricted_metric_ids_str: Optional[List[str]] = ( + [str(mid) for mid in restricted_metric_ids] + if restricted_metric_ids + else None + ) + if restricted_metric_ids_str: + for eval_row, _ in eval_rows_with_source: + store_row_restricted_metrics(eval_row.id, restricted_metric_ids_str) + else: + restricted_metric_ids_str = ( + [str(mid) for mid in restricted_metric_ids] if restricted_metric_ids else None + ) store_evaluation_transcribe_overwrite( evaluation.id, overwrite=transcribe_overwrite, @@ -8197,6 +8318,18 @@ def _gather_retry_targets( in flight; ``include_completed`` controls whether previously- successful rows are eligible. """ + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + from app.db_sharding.eval_rows import gather_retry_targets_sharded + + return gather_retry_targets_sharded( + db, + evaluation, + requested_ids, + include_completed=include_completed, + ) + eval_rows_query = db.query(CallImportEvaluationRow).filter( CallImportEvaluationRow.evaluation_id == evaluation.id ) @@ -8450,20 +8583,31 @@ async def retry_call_import_evaluation( skipped: List[CallImportEvaluationRetrySkippedItem] = [] if requested_ids is None: - from sqlalchemy import func + from app.db_sharding.eval_rows import count_evaluation_rows_for_run + from app.db_sharding.sessions import is_sharding_enabled - count_query = db.query(func.count(CallImportEvaluationRow.id)).filter( - CallImportEvaluationRow.evaluation_id == eval_id - ) - if include_completed: - count_query = count_query.filter( - CallImportEvaluationRow.status.in_(["failed", "completed"]) + if is_sharding_enabled(): + statuses = ( + ["failed", "completed"] if include_completed else ["failed"] + ) + target_count = count_evaluation_rows_for_run( + db, eval_id, statuses=statuses ) else: - count_query = count_query.filter( - CallImportEvaluationRow.status == "failed" + from sqlalchemy import func + + count_query = db.query(func.count(CallImportEvaluationRow.id)).filter( + CallImportEvaluationRow.evaluation_id == eval_id ) - target_count = int(count_query.scalar() or 0) + if include_completed: + count_query = count_query.filter( + CallImportEvaluationRow.status.in_(["failed", "completed"]) + ) + else: + count_query = count_query.filter( + CallImportEvaluationRow.status == "failed" + ) + target_count = int(count_query.scalar() or 0) if target_count == 0: return CallImportEvaluationRetryResponse( requeued=0, @@ -8570,15 +8714,14 @@ async def retry_call_import_evaluation_row( _require_no_evaluation_bulk_operation(eval_id) - eval_row = ( - db.query(CallImportEvaluationRow) - .filter( - CallImportEvaluationRow.id == eval_row_id, - CallImportEvaluationRow.evaluation_id == eval_id, - ) - .first() + from app.db_sharding.eval_rows import ( + evaluation_row_session, + find_evaluation_row_in_run, ) - if not eval_row: + from app.db_sharding.sessions import is_sharding_enabled + + eval_row, _source_stub = find_evaluation_row_in_run(db, eval_id, eval_row_id) + if eval_row is None: raise HTTPException( status_code=404, detail="Evaluation row not found in this run" ) @@ -8594,8 +8737,6 @@ async def retry_call_import_evaluation_row( targets, _ = _gather_retry_targets(db, evaluation, [eval_row.id]) if not targets: - # Status was ``completed`` (or source row vanished) — surface a - # 409 instead of silently no-op'ing so the UI can show why. raise HTTPException( status_code=409, detail=( @@ -8604,9 +8745,22 @@ async def retry_call_import_evaluation_row( ), ) - for er, source_row in targets: - _prepare_source_row_for_retry(source_row, transcribe_overwrite=False) - _reset_eval_row_for_retry(er) + if is_sharding_enabled(): + with evaluation_row_session(eval_row_id) as ( + row_db, + _catalog_db, + eval_row, + source_row, + _shard_id, + ): + _prepare_source_row_for_retry(source_row, transcribe_overwrite=False) + _reset_eval_row_for_retry(eval_row) + row_db.commit() + targets = [(eval_row, source_row)] + else: + for er, source_row in targets: + _prepare_source_row_for_retry(source_row, transcribe_overwrite=False) + _reset_eval_row_for_retry(er) evaluation.error_message = None evaluation.finished_at = None @@ -8625,8 +8779,20 @@ async def retry_call_import_evaluation_row( logger.exception( "Failed to re-enqueue retry for evaluation row {}", eval_row_id ) - eval_row.status = "failed" - eval_row.error_message = f"Failed to re-enqueue retry: {exc}" + if is_sharding_enabled(): + with evaluation_row_session(eval_row_id) as ( + row_db, + _catalog_db, + eval_row, + _source_row, + _shard_id, + ): + eval_row.status = "failed" + eval_row.error_message = f"Failed to re-enqueue retry: {exc}" + row_db.commit() + else: + eval_row.status = "failed" + eval_row.error_message = f"Failed to re-enqueue retry: {exc}" _rollup_evaluation_status(evaluation, db) db.commit() raise HTTPException( @@ -8634,6 +8800,16 @@ async def retry_call_import_evaluation_row( detail=f"Failed to re-enqueue retry: {exc}", ) + if is_sharding_enabled(): + with evaluation_row_session(eval_row_id) as ( + _row_db, + _catalog_db, + eval_row, + source_row, + _shard_id, + ): + return _to_evaluation_row_response(eval_row, source_row) + db.refresh(eval_row) source_row = targets[0][1] return _to_evaluation_row_response(eval_row, source_row) diff --git a/app/api/v1/routes/call_imports.py b/app/api/v1/routes/call_imports.py index dd8d3f9c..6e59f992 100644 --- a/app/api/v1/routes/call_imports.py +++ b/app/api/v1/routes/call_imports.py @@ -16,7 +16,7 @@ import re from datetime import date, datetime, time, timedelta from typing import Any, Dict, Iterable, List, Optional, Tuple -from uuid import UUID +from uuid import UUID, uuid4 from fastapi import APIRouter, Body, BackgroundTasks, Depends, File, Form, HTTPException, Query, Response, UploadFile, status from loguru import logger @@ -26,6 +26,7 @@ from app.config import settings from app.core.auth.rbac import require_admin from app.database import get_db +from app.db_sharding.sessions import is_sharding_enabled from app.dependencies import ( get_api_key, get_organization_id, @@ -94,6 +95,37 @@ def _normalize_dataset(raw: Optional[str]) -> Optional[str]: return cleaned or None +def _serialize_call_import(db: Session, call_import: CallImport) -> CallImportResponse: + """Catalog parent fields; counters come from SQL rollup (not Redis merge).""" + from app.services.call_imports.bulk_ops import rollup_call_import_batch_status + from app.services.call_imports.progress_counters import ( + clear_import_progress_redis, + read_import_progress, + ) + + redis_completed, redis_failed = read_import_progress(call_import.id) + if ( + redis_completed + or redis_failed + or int(call_import.completed_rows or 0) > int(call_import.total_rows or 0) + or int(call_import.failed_rows or 0) > int(call_import.total_rows or 0) + ): + rollup_call_import_batch_status(db, call_import) + db.flush() + + clear_import_progress_redis(call_import.id) + db.refresh(call_import) + total = int(call_import.total_rows or 0) + completed = min(int(call_import.completed_rows or 0), total) if total else int( + call_import.completed_rows or 0 + ) + failed = min(int(call_import.failed_rows or 0), total) if total else int( + call_import.failed_rows or 0 + ) + base = CallImportResponse.model_validate(call_import) + return base.model_copy(update={"completed_rows": completed, "failed_rows": failed}) + + def _resolve_tags( db: Session, organization_id: UUID, tag_ids: Optional[List[UUID]] ) -> List[CallImportTag]: @@ -1142,6 +1174,7 @@ def _materialize_rows( row_model = CallImportRow( call_import_id=call_import.id, organization_id=organization_id, + workspace_id=call_import.workspace_id, row_index=idx, conversation_id=row["conversation_id"], recording_date=( @@ -1458,7 +1491,7 @@ async def create_call_import( raise db.refresh(call_import) - return CallImportResponse.model_validate(call_import) + return _serialize_call_import(db, call_import) @router.patch( @@ -1581,7 +1614,7 @@ async def update_call_import_mapping( call_import.status = CallImportStatus.MAPPED db.commit() db.refresh(call_import) - return CallImportResponse.model_validate(call_import) + return _serialize_call_import(db, call_import) @router.post( @@ -2073,25 +2106,13 @@ async def upload_call_import_audio( # intentionally have no telephony provider. call_import.provider = None + row_mappings: List[Dict[str, Any]] = [] for idx, item in enumerate(prepared): - row = CallImportRow( - call_import_id=call_import.id, - organization_id=organization_id, - row_index=idx, - conversation_id=item["conversation_id"], - recording_url=None, - transcript=None, - transcript_source=None, - raw_columns={"conversation_id": item["conversation_id"]}, - status=CallImportRowStatus.COMPLETED, - ) - db.add(row) - db.flush() - + row_id = uuid4() key = _audio_s3_key( organization_id, call_import.id, - row.id, + row_id, item["extension"], ) s3_service.upload_file_by_key( @@ -2101,9 +2122,36 @@ async def upload_call_import_audio( ) uploaded_keys.append(key) - row.recording_s3_key = key - row.recording_content_type = item["content_type"] - row.recording_size_bytes = len(item["contents"]) + row_mappings.append( + { + "id": row_id, + "call_import_id": call_import.id, + "organization_id": organization_id, + "workspace_id": workspace_id, + "row_index": idx, + "conversation_id": item["conversation_id"], + "recording_url": None, + "transcript": None, + "transcript_source": None, + "raw_columns": {"conversation_id": item["conversation_id"]}, + "status": CallImportRowStatus.COMPLETED, + "recording_s3_key": key, + "recording_content_type": item["content_type"], + "recording_size_bytes": len(item["contents"]), + } + ) + + if is_sharding_enabled(): + from app.db_sharding.row_ops import ( + bulk_insert_mappings_on_shards, + register_shard_slices, + ) + + bulk_insert_mappings_on_shards(db, call_import.id, row_mappings) + register_shard_slices(db, call_import.id, len(row_mappings)) + else: + for mapping in row_mappings: + db.add(CallImportRow(**mapping)) db.commit() except Exception as exc: @@ -2231,7 +2279,7 @@ async def list_call_imports( ) return CallImportListResponse( - items=[CallImportResponse.model_validate(item) for item in items], + items=[_serialize_call_import(db, item) for item in items], total=total, page=page, page_size=page_size, @@ -2416,7 +2464,7 @@ async def update_call_import( db.commit() db.refresh(call_import) - return CallImportResponse.model_validate(call_import) + return _serialize_call_import(db, call_import) @router.get( @@ -2483,29 +2531,71 @@ async def get_call_import_detail( db.commit() db.refresh(call_import) - rows_query = db.query(CallImportRow).filter( - CallImportRow.call_import_id == call_import.id - ) - search_term = (q or "").strip() diarised_status_filter = (diarised_status or "").strip() or None filtered_total_rows: Optional[int] = None - if search_term: - rows_query = rows_query.filter( - CallImportRow.conversation_id.ilike(f"%{search_term}%") - ) - if diarised_status_filter: - rows_query = rows_query.filter( - CallImportRow.diarised_transcript_status == diarised_status_filter - ) - # Surface the post-filter total whenever any filter is active so - # the UI can paginate against the slice it's actually displaying. - if search_term or diarised_status_filter: - filtered_total_rows = rows_query.count() + has_row_filters = bool(search_term or diarised_status_filter) + + if has_row_filters: + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import count_call_import_rows_filtered + + filtered_total_rows = count_call_import_rows_filtered( + db, + call_import.id, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + else: + rows_query = db.query(CallImportRow).filter( + CallImportRow.call_import_id == call_import.id + ) + if search_term: + rows_query = rows_query.filter( + CallImportRow.conversation_id.ilike(f"%{search_term}%") + ) + if diarised_status_filter: + rows_query = rows_query.filter( + CallImportRow.diarised_transcript_status == diarised_status_filter + ) + filtered_total_rows = rows_query.count() if row_limit == 0: rows: List[CallImportRow] = [] + elif is_sharding_enabled(): + from app.db_sharding.scatter_gather import ( + fetch_call_import_rows_filtered_page, + fetch_call_import_rows_page, + ) + + if has_row_filters: + rows = fetch_call_import_rows_filtered_page( + db, + call_import.id, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + offset=row_offset, + limit=row_limit, + ) + else: + rows = fetch_call_import_rows_page( + db, + call_import.id, + offset=row_offset, + limit=row_limit, + ) else: + rows_query = db.query(CallImportRow).filter( + CallImportRow.call_import_id == call_import.id + ) + if search_term: + rows_query = rows_query.filter( + CallImportRow.conversation_id.ilike(f"%{search_term}%") + ) + if diarised_status_filter: + rows_query = rows_query.filter( + CallImportRow.diarised_transcript_status == diarised_status_filter + ) rows = ( rows_query.order_by(CallImportRow.row_index) .offset(row_offset) @@ -2517,17 +2607,26 @@ async def get_call_import_detail( # across the whole batch — much cheaper than paging through every # row to recount on the client and lets the UI render a # transcribe/diarise progress bar without a separate roundtrip. - diarised_status_counts: Dict[str, int] = {} - for status_value, count in ( - db.query(CallImportRow.diarised_transcript_status, func.count()) - .filter(CallImportRow.call_import_id == call_import.id) - .group_by(CallImportRow.diarised_transcript_status) - .all() - ): - if isinstance(status_value, str): - diarised_status_counts[status_value] = int(count or 0) + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import aggregate_diarised_transcript_counts - detail = CallImportDetailResponse.model_validate(call_import) + diarised_status_counts = aggregate_diarised_transcript_counts( + db, call_import.id + ) + else: + diarised_status_counts: Dict[str, int] = {} + for status_value, count in ( + db.query(CallImportRow.diarised_transcript_status, func.count()) + .filter(CallImportRow.call_import_id == call_import.id) + .group_by(CallImportRow.diarised_transcript_status) + .all() + ): + if isinstance(status_value, str): + diarised_status_counts[status_value] = int(count or 0) + + detail = CallImportDetailResponse.model_validate( + _serialize_call_import(db, call_import).model_dump() + ) detail.rows = [CallImportRowResponse.model_validate(r) for r in rows] detail.filtered_total_rows = filtered_total_rows detail.diarised_pending_rows = diarised_status_counts.get("pending", 0) @@ -2588,15 +2687,27 @@ async def list_call_import_row_ids( detail="Call import not found", ) + search_term = (q or "").strip() + status_filter = (diarised_status or "").strip() or None + + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import list_call_import_row_ids_filtered + + ids = list_call_import_row_ids_filtered( + db, + call_import.id, + search_term=search_term, + diarised_status_filter=status_filter, + ) + return CallImportRowIdsResponse(ids=ids, total=len(ids)) + rows_query = db.query(CallImportRow.id).filter( CallImportRow.call_import_id == call_import.id ) - search_term = (q or "").strip() if search_term: rows_query = rows_query.filter( CallImportRow.conversation_id.ilike(f"%{search_term}%") ) - status_filter = (diarised_status or "").strip() or None if status_filter: rows_query = rows_query.filter( CallImportRow.diarised_transcript_status == status_filter diff --git a/app/cli.py b/app/cli.py index 74acf517..59a669bd 100644 --- a/app/cli.py +++ b/app/cli.py @@ -466,11 +466,12 @@ def worker(config: str, loglevel: str, queues: Optional[str], concurrency: Optio ) @click.option( "--imports-worker-concurrency", - default=32, + default=12, type=int, help=( "Concurrency for the imports+diarization+evaluations worker " - "(default: 32; thread pool; consumes imports, then diarization, then evaluations)" + "(default: 12; thread pool). Use lower values (8–12) when DB sharding " + "is enabled; 32 threads can exhaust per-shard SQLAlchemy pools." ), ) def start_all( @@ -636,6 +637,8 @@ def _stream(): ) if imports_worker: + from app.workers.config import IMPORTS_WORKER_QUEUES + worker_imports_process = _spawn_worker( [ "celery", @@ -644,14 +647,14 @@ def _stream(): "worker", f"--loglevel={worker_loglevel}", "-Q", - "imports,diarization,evaluations", + IMPORTS_WORKER_QUEUES, "-P", "threads", "-c", str(imports_worker_concurrency), ], label=( - "Celery worker (imports+diarization+evaluations queues, " + f"Celery worker ({IMPORTS_WORKER_QUEUES} queues, " f"pool=threads, concurrency={imports_worker_concurrency})" ), prefix="[WORKER-IMPORTS]", @@ -687,8 +690,10 @@ def _stream(): if watch_frontend: click.echo(f" Frontend watcher: Active (rebuilding on file changes)") if imports_worker: + from app.workers.config import IMPORTS_WORKER_QUEUES + click.echo( - " Workers: default queue + imports,diarization,evaluations queues " + f" Workers: default queue + {IMPORTS_WORKER_QUEUES} " f"(concurrency={imports_worker_concurrency}; imports preferred)" ) else: @@ -797,6 +802,168 @@ def init_config(output: str): sys.exit(1) +def _bootstrap_sharding_config(config_path: str) -> None: + from app.config import load_config_from_file + from app.db_sharding.pool_manager import db_pool_manager + + load_config_from_file(config_path) + db_pool_manager.reset() + + +@click.group() +def sharding(): + """Call-import data-plane sharding admin (rebalance / registry).""" + pass + + +main.add_command(sharding) + + +@sharding.command("list-slices") +@click.option( + "--config", + "-c", + type=click.Path(exists=True, readable=True), + default="config.yml", + help="Path to configuration YAML file", +) +@click.option( + "--call-import-id", + required=True, + help="Call import UUID whose slice registry to inspect", +) +def sharding_list_slices(config: str, call_import_id: str): + """Show catalog slice registry rows for a call import.""" + from uuid import UUID + + from app.db_sharding.pool_manager import open_catalog_session + from app.db_sharding.rebalance import RebalanceError, list_shard_slices, require_sharding_enabled + + try: + _bootstrap_sharding_config(config) + require_sharding_enabled() + cid = UUID(call_import_id) + except (ValueError, RebalanceError) as exc: + click.echo(f"❌ {exc}", err=True) + sys.exit(1) + + catalog = open_catalog_session() + try: + slices = list_shard_slices(catalog, cid) + if not slices: + click.echo(f"No registry slices for call_import {cid}") + return + click.echo(f"call_import_id: {cid}") + click.echo(f"slices: {len(slices)}") + for item in slices: + click.echo( + f" slice {item.slice_id}: shard={item.shard_id} " + f"rows [{item.row_index_min}, {item.row_index_max}] " + f"(count={item.row_count})" + ) + finally: + catalog.close() + + +@sharding.command("rebalance-slices") +@click.option( + "--config", + "-c", + type=click.Path(exists=True, readable=True), + default="config.yml", + help="Path to configuration YAML file", +) +@click.option("--call-import-id", required=True, help="Call import UUID to rebalance") +@click.option("--from-shard", required=True, help="Source shard id (e.g. data-shard-01)") +@click.option("--to-shard", required=True, help="Target shard id (e.g. data-shard-02)") +@click.option( + "--slice-id", + "slice_ids", + multiple=True, + type=int, + help="Move only these slice ids (default: all slices on --from-shard)", +) +@click.option( + "--dry-run", + is_flag=True, + help="Plan only: print row counts without copying or updating registry", +) +@click.option( + "--force", + is_flag=True, + help="Skip quiescence checks (use only after pausing workers)", +) +def sharding_rebalance_slices( + config: str, + call_import_id: str, + from_shard: str, + to_shard: str, + slice_ids: tuple[int, ...], + dry_run: bool, + force: bool, +): + """Copy call-import slice rows (and eval rows) from one shard to another.""" + from uuid import UUID + + from app.db_sharding.pool_manager import open_catalog_session + from app.db_sharding.rebalance import ( + RebalanceError, + require_sharding_enabled, + build_rebalance_plan, + execute_rebalance_slices, + ) + + try: + _bootstrap_sharding_config(config) + require_sharding_enabled() + cid = UUID(call_import_id) + except (ValueError, RebalanceError) as exc: + click.echo(f"❌ {exc}", err=True) + sys.exit(1) + + catalog = open_catalog_session() + try: + plan = build_rebalance_plan( + catalog, + cid, + from_shard_id=from_shard, + to_shard_id=to_shard, + slice_ids=slice_ids or None, + ) + click.echo("Rebalance plan:") + click.echo(f" call_import_id: {plan.call_import_id}") + click.echo(f" from_shard: {plan.from_shard_id}") + click.echo(f" to_shard: {plan.to_shard_id}") + click.echo(f" slices: {len(plan.slices)}") + for item in plan.slices: + click.echo( + f" slice {item.slice_id}: rows [{item.row_index_min}, {item.row_index_max}]" + ) + click.echo(f" import_rows: {plan.import_row_count}") + click.echo(f" eval_rows: {plan.eval_row_count}") + + result = execute_rebalance_slices( + catalog, + plan, + dry_run=dry_run, + force=force, + ) + if result.dry_run: + click.echo("✅ Dry run complete (no changes made)") + else: + click.echo( + "✅ Rebalance complete: " + f"slices={result.slices_moved}, " + f"import_rows={result.import_rows_moved}, " + f"eval_rows={result.eval_rows_moved}" + ) + except RebalanceError as exc: + click.echo(f"❌ {exc}", err=True) + sys.exit(1) + finally: + catalog.close() + + if __name__ == "__main__": main() diff --git a/app/config.py b/app/config.py index a6e7356f..eacc00cd 100644 --- a/app/config.py +++ b/app/config.py @@ -3,7 +3,7 @@ import json import yaml from pathlib import Path -from typing import List, Optional, Union +from typing import Any, Dict, List, Optional, Union from pydantic import field_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -30,6 +30,14 @@ class Settings(BaseSettings): POSTGRES_PORT: int = 5432 POSTGRES_DB: str = "efficientai" + # Database sharding (call-import row shards; default off) + DB_SHARDING_ENABLED: bool = False + DB_CATALOG_URL: Optional[str] = None + DB_SHARD_ROW_CHUNK_SIZE: int = 500 + DB_POOL_SIZE: int = 10 + DB_MAX_OVERFLOW: int = 20 + DB_SHARD_ENTRIES: List[dict] = [] + # Redis REDIS_URL: Optional[str] = None REDIS_HOST: str = "localhost" @@ -416,7 +424,21 @@ def load_config_from_file(config_path: str) -> None: f"postgresql://{settings.POSTGRES_USER}:{settings.POSTGRES_PASSWORD}" f"@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}" ) - + if "pool_size" in db_config: + settings.DB_POOL_SIZE = int(db_config["pool_size"]) + if "max_overflow" in db_config: + settings.DB_MAX_OVERFLOW = int(db_config["max_overflow"]) + if "catalog_url" in db_config: + settings.DB_CATALOG_URL = db_config["catalog_url"] + sharding_cfg = db_config.get("sharding") or {} + if isinstance(sharding_cfg, dict): + if "enabled" in sharding_cfg: + settings.DB_SHARDING_ENABLED = bool(sharding_cfg["enabled"]) + if "row_chunk_size" in sharding_cfg: + settings.DB_SHARD_ROW_CHUNK_SIZE = int(sharding_cfg["row_chunk_size"]) + if "shards" in db_config and isinstance(db_config["shards"], list): + settings.DB_SHARD_ENTRIES = list(db_config["shards"]) + if "redis" in config_data: redis_config = config_data["redis"] if "url" in redis_config: diff --git a/app/core/migrations.py b/app/core/migrations.py index 531e49aa..4df87b83 100644 --- a/app/core/migrations.py +++ b/app/core/migrations.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import List, Optional, Tuple from sqlalchemy import text, inspect -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.exc import ProgrammingError from app.database import engine, SessionLocal, Base import logging @@ -19,11 +19,64 @@ MIGRATIONS_DIR = Path(__file__).parent.parent / "migrations" +def _migration_scope_for_file(migration_file: Path) -> str: + import importlib.util + + version = migration_file.stem + spec = importlib.util.spec_from_file_location(f"migrations_scope_{version}", migration_file) + if spec is None or spec.loader is None: + return "all" + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return str(getattr(module, "MIGRATION_SCOPE", "all") or "all") + + +def _migration_applies_to_engine(migration_file: Path, engine_role: str) -> bool: + if engine_role in ("all", "legacy"): + return True + scope = _migration_scope_for_file(migration_file) + if scope == "all": + return True + return scope == engine_role + + +def _canonical_db_url_key(url: str) -> tuple: + """Compare DB URLs ignoring driver suffix normalisation (postgresql vs postgresql+psycopg2).""" + from sqlalchemy.engine import make_url + + parsed = make_url(url) + driver = (parsed.drivername or "").split("+", 1)[0] + return ( + driver, + parsed.username or "", + parsed.password or "", + parsed.host or "", + parsed.port, + parsed.database or "", + ) + + +def _engine_role_for_url(engine_url: str) -> str: + from app.config import settings + + if not getattr(settings, "DB_SHARDING_ENABLED", False): + return "legacy" + catalog_url = getattr(settings, "DB_CATALOG_URL", None) or settings.DATABASE_URL + try: + if _canonical_db_url_key(engine_url) == _canonical_db_url_key(catalog_url): + return "catalog" + except Exception: + if str(engine_url) == str(catalog_url): + return "catalog" + return "shard" + + class MigrationRunner: """Handles running database migrations in order.""" - def __init__(self, db: Session): + def __init__(self, db: Session, *, engine_role: str = "legacy"): self.db = db + self.engine_role = engine_role self.ensure_migrations_table() def ensure_migrations_table(self): @@ -81,6 +134,8 @@ def get_pending_migrations(self) -> List[Path]: for migration_file in migration_files: version = migration_file.stem # filename without .py if version not in applied and not version.startswith("__"): + if not _migration_applies_to_engine(migration_file, self.engine_role): + continue pending.append(migration_file) return pending @@ -163,85 +218,77 @@ def run_migrations(): """ Run all pending database migrations. This should be called on application startup. - + + When sharding is enabled, applies the same migration set to each unique + engine URL (catalog + row shards) until catalog/shard split (Phase 8). + Raises: RuntimeError: If migrations fail, preventing application startup """ + from app.db_sharding.pool_manager import db_pool_manager + # Ensure logging is configured at INFO level logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', force=True ) - + logger.info("=" * 60) logger.info("🔄 Starting database migrations...") logger.info("=" * 60) - - db = SessionLocal() - try: - runner = MigrationRunner(db) - - # Check current migration status - applied = runner.get_applied_migrations() - pending = runner.get_pending_migrations() - - if applied: - logger.info(f"📊 Currently applied migrations: {len(applied)}") - for version in applied[-5:]: # Show last 5 - logger.info(f" ✓ {version}") - if len(applied) > 5: - logger.info(f" ... and {len(applied) - 5} more") - - if not pending: - logger.info("✅ Database is up to date - no migrations needed") - return - - # Run migrations - success = runner.run_all() - if not success: - logger.error("") - logger.error("=" * 60) - logger.error("❌ MIGRATION FAILED - Application cannot start!") - logger.error("=" * 60) - logger.error("The application will not start until migrations succeed.") - logger.error("") - logger.error("To diagnose the issue:") - logger.error(" 1. Run: eai migrate --verbose") - logger.error(" 2. Check the error messages above") - logger.error(" 3. Fix any database schema issues") - logger.error("=" * 60) - raise RuntimeError("Database migrations failed") - - # Verify migrations were applied - logger.info("") - logger.info("🔍 Verifying migrations were applied...") - final_pending = runner.get_pending_migrations() - if final_pending: - logger.warning(f"⚠️ Warning: {len(final_pending)} migration(s) still pending after run:") - for migration_file in final_pending: - logger.warning(f" - {migration_file.name}") - logger.warning("This may indicate a migration tracking issue.") - else: - logger.info("✅ Verification complete - all migrations applied successfully") - - logger.info("=" * 60) - - except RuntimeError: - # Re-raise RuntimeError (migration failures) - raise - except Exception as e: - logger.error("") - logger.error("=" * 60) - logger.error("❌ UNEXPECTED ERROR during migrations!") - logger.error("=" * 60) - logger.error(f"Error: {e}") - import traceback - logger.error(f"Traceback:\n{traceback.format_exc()}") - logger.error("=" * 60) - raise - finally: - db.close() + + engines = db_pool_manager.all_engines_for_migrations() + for idx, eng in enumerate(engines): + url_hint = str(eng.url).split("@")[-1] if eng.url else f"engine-{idx}" + logger.info("Migration target: %s", url_hint) + factory = sessionmaker(autocommit=False, autoflush=False, bind=eng) + db = factory() + try: + engine_role = _engine_role_for_url(str(eng.url)) + runner = MigrationRunner(db, engine_role=engine_role) + + applied = runner.get_applied_migrations() + pending = runner.get_pending_migrations() + + if applied: + logger.info(f"📊 Currently applied migrations: {len(applied)}") + for version in applied[-5:]: + logger.info(f" ✓ {version}") + if len(applied) > 5: + logger.info(f" ... and {len(applied) - 5} more") + + if not pending: + logger.info("✅ Database is up to date - no migrations needed (%s)", url_hint) + continue + + success = runner.run_all() + if not success: + logger.error("") + logger.error("=" * 60) + logger.error("❌ MIGRATION FAILED - Application cannot start!") + logger.error("=" * 60) + logger.error("Target: %s", url_hint) + raise RuntimeError("Database migrations failed") + + final_pending = runner.get_pending_migrations() + if final_pending: + logger.warning( + f"⚠️ Warning: {len(final_pending)} migration(s) still pending after run on {url_hint}" + ) + else: + logger.info("✅ Verification complete - all migrations applied (%s)", url_hint) + except RuntimeError: + raise + except Exception as e: + logger.error("❌ UNEXPECTED ERROR during migrations on %s: %s", url_hint, e) + import traceback + logger.error(f"Traceback:\n{traceback.format_exc()}") + raise + finally: + db.close() + + logger.info("=" * 60) def check_migrations_status() -> Tuple[bool, List[str]]: diff --git a/app/database.py b/app/database.py index 8fda0d0b..7d1001db 100644 --- a/app/database.py +++ b/app/database.py @@ -1,20 +1,32 @@ """Database connection and session management.""" -from sqlalchemy import create_engine -from sqlalchemy.orm import declarative_base, sessionmaker -from app.config import settings - -# Create database engine -engine = create_engine( - settings.DATABASE_URL, - pool_pre_ping=True, - pool_size=10, - max_overflow=20, - connect_args={"options": "-c timezone=UTC"}, -) - -# Create session factory -SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) +from sqlalchemy.orm import declarative_base + +from app.db_sharding.pool_manager import db_pool_manager + +class _EngineProxy: + """Proxy so `from app.database import engine` works after lazy init.""" + + def __getattr__(self, name): + return getattr(db_pool_manager.catalog_engine, name) + + def __repr__(self): + return repr(db_pool_manager.catalog_engine) + + +engine = _EngineProxy() + + +class _SessionLocalFactory: + def __call__(self): + factory = db_pool_manager.catalog_session_factory() + return factory() + + def __getattr__(self, name): + return getattr(db_pool_manager.catalog_session_factory(), name) + + +SessionLocal = _SessionLocalFactory() # Base class for models Base = declarative_base() @@ -41,16 +53,15 @@ def init_db(): # on a fresh database. import app.models.database # noqa: F401 - Base.metadata.create_all(bind=engine) + for eng in db_pool_manager.all_engines_for_migrations(): + Base.metadata.create_all(bind=eng) _run_column_migrations() def _run_column_migrations(): """Add new columns to existing tables if they don't exist yet.""" - from sqlalchemy import text, inspect - - inspector = inspect(engine) - + from sqlalchemy import inspect, text + migrations = [ ("evaluators", "name", "ALTER TABLE evaluators ADD COLUMN name VARCHAR"), ("evaluators", "custom_prompt", "ALTER TABLE evaluators ADD COLUMN custom_prompt TEXT"), @@ -60,23 +71,28 @@ def _run_column_migrations(): ("evaluators", "llm_provider", "ALTER TABLE evaluators ADD COLUMN llm_provider VARCHAR"), ("evaluators", "llm_model", "ALTER TABLE evaluators ADD COLUMN llm_model VARCHAR"), ] - - with engine.begin() as conn: - existing_cols = {c["name"] for c in inspector.get_columns("evaluators")} - - for table, column, ddl in migrations: - if ddl and column not in existing_cols: - conn.execute(text(ddl)) - - nullable_changes = [ - ("evaluators", "agent_id"), - ("evaluators", "persona_id"), - ("evaluators", "scenario_id"), - ("evaluator_results", "agent_id"), - ] - for table, column in nullable_changes: - try: - conn.execute(text(f"ALTER TABLE {table} ALTER COLUMN {column} DROP NOT NULL")) - except Exception: - pass + for eng in db_pool_manager.all_engines_for_migrations(): + inspector = inspect(eng) + if not inspector.has_table("evaluators"): + continue + with eng.begin() as conn: + existing_cols = {c["name"] for c in inspector.get_columns("evaluators")} + + for table, column, ddl in migrations: + if ddl and column not in existing_cols: + conn.execute(text(ddl)) + + nullable_changes = [ + ("evaluators", "agent_id"), + ("evaluators", "persona_id"), + ("evaluators", "scenario_id"), + ("evaluator_results", "agent_id"), + ] + for table, column in nullable_changes: + if not inspector.has_table(table): + continue + try: + conn.execute(text(f"ALTER TABLE {table} ALTER COLUMN {column} DROP NOT NULL")) + except Exception: + pass diff --git a/app/db_sharding/__init__.py b/app/db_sharding/__init__.py new file mode 100644 index 00000000..3788c55c --- /dev/null +++ b/app/db_sharding/__init__.py @@ -0,0 +1,13 @@ +"""Platform data-plane sharding: pools, routing, sessions (call-import rows today).""" + +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.router import ShardRouter +from app.db_sharding.sessions import catalog_session, is_sharding_enabled, row_shard_session + +__all__ = [ + "ShardRouter", + "catalog_session", + "db_pool_manager", + "is_sharding_enabled", + "row_shard_session", +] diff --git a/app/db_sharding/eval_rows.py b/app/db_sharding/eval_rows.py new file mode 100644 index 00000000..f60a9c50 --- /dev/null +++ b/app/db_sharding/eval_rows.py @@ -0,0 +1,317 @@ +"""Evaluation row lookups and fan-out helpers for sharded data plane.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, Callable, List, Optional, Sequence, Tuple, TypeVar +from uuid import UUID + +from sqlalchemy.orm import Query, Session + +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.row_ops import ( + close_row_sessions, + locate_call_import_evaluation_row, +) +from app.db_sharding.scatter_gather import load_evaluation_row_pairs +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import CallImportEvaluation, CallImportEvaluationRow, CallImportRow + +T = TypeVar("T") + + +def load_evaluation_rows_for_run( + catalog_db: Session, + evaluation_id: UUID, +) -> List[CallImportEvaluationRow]: + """All evaluation rows for a run (scatter-gather when sharded).""" + return [eval_row for eval_row, _ in load_evaluation_row_pairs(catalog_db, evaluation_id)] + + +def find_evaluation_row_in_run( + catalog_db: Session, + evaluation_id: UUID, + eval_row_id: UUID, +) -> Tuple[Optional[CallImportEvaluationRow], Optional[CallImportRow]]: + for eval_row, source_row in load_evaluation_row_pairs(catalog_db, evaluation_id): + if eval_row.id == eval_row_id: + return eval_row, source_row + return None, None + + +def count_evaluation_rows_for_run( + catalog_db: Session, + evaluation_id: UUID, + *, + statuses: Optional[Sequence[str]] = None, +) -> int: + rows = load_evaluation_rows_for_run(catalog_db, evaluation_id) + if not statuses: + return len(rows) + allowed = set(statuses) + return sum(1 for row in rows if (row.status or "") in allowed) + + +def foreach_evaluation_row_mutating( + catalog_db: Session, + evaluation_id: UUID, + mutate: Callable[[CallImportEvaluationRow], bool], +) -> int: + """Run ``mutate(row)`` on every eval row; commit per shard when sharded.""" + from sqlalchemy.orm.attributes import flag_modified + + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .all() + ) + changed = 0 + for row in rows: + if mutate(row): + flag_modified(row, "metric_scores") + changed += 1 + return changed + + router = db_pool_manager.router + assert router is not None + total = 0 + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + rows = ( + shard_db.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .all() + ) + shard_changed = False + for row in rows: + if mutate(row): + flag_modified(row, "metric_scores") + shard_changed = True + total += 1 + if shard_changed: + shard_db.commit() + except Exception: + shard_db.rollback() + raise + finally: + shard_db.close() + return total + + +def scatter_gather_eval_query( + catalog_db: Session, + build_query: Callable[[Session], Query], +) -> List[T]: + """Run the same ORM query on each row shard and concatenate results.""" + if not is_sharding_enabled(): + return list(build_query(catalog_db).all()) + + merged: List[T] = [] + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + merged.extend(build_query(shard_db).all()) + finally: + shard_db.close() + return merged + + +def scatter_gather_eval_query_count( + catalog_db: Session, + build_query: Callable[[Session], Query], +) -> int: + if not is_sharding_enabled(): + return int(build_query(catalog_db).count()) + total = 0 + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + total += int(build_query(shard_db).count()) + finally: + shard_db.close() + return total + + +def paginate_pairs( + pairs: Sequence[Tuple[CallImportEvaluationRow, CallImportRow]], + *, + page: int, + page_size: int, +) -> Tuple[int, List[Tuple[CallImportEvaluationRow, CallImportRow]]]: + ordered = sorted(pairs, key=lambda pair: int(pair[1].row_index or 0)) + total = len(ordered) + start = max(0, (page - 1) * page_size) + end = start + page_size + return total, list(ordered[start:end]) + + +PairSortKey = Callable[[Tuple[CallImportEvaluationRow, CallImportRow]], Any] + + +def fetch_evaluation_row_pairs_page( + catalog_db: Session, + build_query: Callable[[Session], Query], + *, + page: int, + page_size: int, + sort_key: PairSortKey, + sort_desc: bool = False, + bounded_shard_fetch: bool = True, +) -> Tuple[int, List[Tuple[CallImportEvaluationRow, CallImportRow]]]: + """Return ``(total, page_slice)`` without loading every eval row pair. + + When sharding is enabled and ``bounded_shard_fetch`` is true (``row_index`` + sort only), each shard returns at most ``page * page_size`` rows; results + are merged in Python and sliced to the requested page. + + For other sort keys, every shard returns all matching rows so globally + correct ordering is preserved across shards. + """ + page = max(1, page) + page_size = max(1, page_size) + total = scatter_gather_eval_query_count(catalog_db, build_query) + if total == 0: + return 0, [] + + if not is_sharding_enabled(): + rows = ( + build_query(catalog_db) + .offset((page - 1) * page_size) + .limit(page_size) + .all() + ) + return total, list(rows) + + merged: List[Tuple[CallImportEvaluationRow, CallImportRow]] = [] + router = db_pool_manager.router + assert router is not None + over_fetch = page * page_size + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + shard_query = build_query(shard_db) + if bounded_shard_fetch: + merged.extend(shard_query.limit(over_fetch).all()) + else: + merged.extend(shard_query.all()) + finally: + shard_db.close() + + merged.sort(key=sort_key, reverse=sort_desc) + start = (page - 1) * page_size + return total, list(merged[start : start + page_size]) + + +def gather_retry_targets_sharded( + catalog_db: Session, + evaluation: CallImportEvaluation, + requested_ids: Optional[List[UUID]], + *, + include_completed: bool, +) -> Tuple[List[Tuple[CallImportEvaluationRow, CallImportRow]], List]: + """Return (targets, skipped_as_dict_list) for bulk retry when sharding is on.""" + from app.models.schemas import CallImportEvaluationRetrySkippedItem + + pairs = load_evaluation_row_pairs(catalog_db, evaluation.id) + eval_by_id = {er.id: (er, sr) for er, sr in pairs} + + targets: List[Tuple[CallImportEvaluationRow, CallImportRow]] = [] + skipped: List[CallImportEvaluationRetrySkippedItem] = [] + + if requested_ids is None: + if include_completed: + candidate_ids = [ + er.id + for er, _ in pairs + if er.status in ("failed", "completed") + ] + else: + candidate_ids = [er.id for er, _ in pairs if er.status == "failed"] + else: + requested_set = set(requested_ids) + candidate_ids = [eid for eid in requested_set if eid in eval_by_id] + for missing in requested_set - set(candidate_ids): + skipped.append( + CallImportEvaluationRetrySkippedItem( + eval_row_id=missing, + reason="unknown", + ) + ) + + for eid in candidate_ids: + eval_row, source_row = eval_by_id[eid] + if eval_row.status in {"pending", "running"}: + skipped.append( + CallImportEvaluationRetrySkippedItem( + eval_row_id=eval_row.id, + reason="in_progress", + ) + ) + continue + if eval_row.status == "completed" and not include_completed: + skipped.append( + CallImportEvaluationRetrySkippedItem( + eval_row_id=eval_row.id, + reason="completed", + ) + ) + continue + targets.append((eval_row, source_row)) + + return targets, skipped + + +@contextmanager +def evaluation_row_session(eval_row_id: UUID | str): + """Yield (row_db, catalog_db, eval_row, source_row, shard_id).""" + row_db, catalog_db, eval_row, source_row, shard_id = ( + locate_call_import_evaluation_row(eval_row_id) + ) + try: + yield row_db, catalog_db, eval_row, source_row, shard_id + finally: + close_row_sessions(row_db, catalog_db) + + +def delete_evaluation_row_on_shards( + eval_row_id: UUID, + evaluation_id: UUID, +) -> bool: + """Delete eval row from the shard that holds it. Returns True if deleted.""" + if not is_sharding_enabled(): + return False + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + eval_row = ( + shard_db.query(CallImportEvaluationRow) + .filter( + CallImportEvaluationRow.id == eval_row_id, + CallImportEvaluationRow.evaluation_id == evaluation_id, + ) + .first() + ) + if eval_row is None: + continue + shard_db.delete(eval_row) + shard_db.commit() + return True + except Exception: + shard_db.rollback() + raise + finally: + shard_db.close() + return False diff --git a/app/db_sharding/import_dispatch.py b/app/db_sharding/import_dispatch.py new file mode 100644 index 00000000..3b85c0dd --- /dev/null +++ b/app/db_sharding/import_dispatch.py @@ -0,0 +1,325 @@ +"""Per-shard import and diarization dispatch helpers.""" + +from __future__ import annotations + +from typing import List, Optional, Tuple +from uuid import UUID + +from sqlalchemy.orm import Session + +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import CallImport, CallImportRow +from app.models.enums import CallImportRowStatus, CallImportStatus + + +def _import_blocked_by_rebalance(call_import_id: UUID) -> bool: + from app.db_sharding.rebalance import is_import_rebalance_locked + + return is_import_rebalance_locked(call_import_id) + + +def pending_import_workspaces(catalog_db: Session) -> List[UUID]: + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImport.workspace_id) + .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) + .filter( + CallImport.status != CallImportStatus.DELETING, + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + import_ids: set[UUID] = set() + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + rows = ( + shard_db.query(CallImportRow.call_import_id) + .filter( + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + import_ids.update(row[0] for row in rows if row[0] is not None) + finally: + shard_db.close() + if not import_ids: + return [] + from app.db_sharding.rebalance import filter_unlocked_call_import_ids + + import_ids = set(filter_unlocked_call_import_ids(import_ids)) + if not import_ids: + return [] + rows = ( + catalog_db.query(CallImport.workspace_id) + .filter( + CallImport.id.in_(import_ids), + CallImport.status != CallImportStatus.DELETING, + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + +def call_imports_with_pending_rows( + catalog_db: Session, + workspace_id: UUID, +) -> List[UUID]: + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImport.id) + .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) + .filter( + CallImport.workspace_id == workspace_id, + CallImport.status != CallImportStatus.DELETING, + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + candidates = [ + row[0] + for row in catalog_db.query(CallImport.id) + .filter( + CallImport.workspace_id == workspace_id, + CallImport.status != CallImportStatus.DELETING, + ) + .all() + if row[0] is not None + ] + pending: set[UUID] = set() + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + rows = ( + shard_db.query(CallImportRow.call_import_id) + .filter( + CallImportRow.call_import_id.in_(candidates), + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + pending.update(row[0] for row in rows if row[0] is not None) + finally: + shard_db.close() + from app.db_sharding.rebalance import filter_unlocked_call_import_ids + + return sorted(filter_unlocked_call_import_ids(pending)) + + +def pending_import_row_for_call_import( + catalog_db: Session, + call_import_id: UUID, +) -> Optional[Tuple[CallImportRow, CallImport]]: + call_import = ( + catalog_db.query(CallImport) + .filter( + CallImport.id == call_import_id, + CallImport.status != CallImportStatus.DELETING, + ) + .first() + ) + if call_import is None: + return None + if _import_blocked_by_rebalance(call_import_id): + return None + + if not is_sharding_enabled(): + row = ( + catalog_db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .order_by(CallImportRow.created_at.asc()) + .first() + ) + return (row, call_import) if row is not None else None + + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = ( + shard_db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.PENDING, + CallImportRow.celery_task_id.is_(None), + ) + .order_by(CallImportRow.created_at.asc()) + .first() + ) + if row is not None: + return row, call_import + finally: + shard_db.close() + return None + + +def pending_diarization_workspaces(catalog_db: Session) -> List[UUID]: + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImport.workspace_id) + .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) + .filter( + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + import_ids: set[UUID] = set() + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + rows = ( + shard_db.query(CallImportRow.call_import_id) + .filter( + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + import_ids.update(row[0] for row in rows if row[0] is not None) + finally: + shard_db.close() + if not import_ids: + return [] + from app.db_sharding.rebalance import filter_unlocked_call_import_ids + + import_ids = set(filter_unlocked_call_import_ids(import_ids)) + if not import_ids: + return [] + rows = ( + catalog_db.query(CallImport.workspace_id) + .filter(CallImport.id.in_(import_ids)) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + +def call_imports_with_pending_diarization( + catalog_db: Session, + workspace_id: UUID, +) -> List[UUID]: + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImport.id) + .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) + .filter( + CallImport.workspace_id == workspace_id, + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + candidates = [ + row[0] + for row in catalog_db.query(CallImport.id) + .filter(CallImport.workspace_id == workspace_id) + .all() + if row[0] is not None + ] + pending: set[UUID] = set() + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + rows = ( + shard_db.query(CallImportRow.call_import_id) + .filter( + CallImportRow.call_import_id.in_(candidates), + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + pending.update(row[0] for row in rows if row[0] is not None) + finally: + shard_db.close() + from app.db_sharding.rebalance import filter_unlocked_call_import_ids + + return sorted(filter_unlocked_call_import_ids(pending)) + + +def pending_diarization_row_for_call_import( + catalog_db: Session, + call_import_id: UUID, +) -> Optional[Tuple[CallImportRow, CallImport]]: + call_import = ( + catalog_db.query(CallImport).filter(CallImport.id == call_import_id).first() + ) + if call_import is None: + return None + if _import_blocked_by_rebalance(call_import_id): + return None + if not is_sharding_enabled(): + row = ( + catalog_db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .order_by(CallImportRow.created_at.asc()) + .first() + ) + return (row, call_import) if row is not None else None + + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = ( + shard_db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.diarised_transcript_status == "pending", + CallImportRow.celery_task_id.is_(None), + ) + .order_by(CallImportRow.created_at.asc()) + .first() + ) + if row is not None: + return row, call_import + finally: + shard_db.close() + return None diff --git a/app/db_sharding/pool_manager.py b/app/db_sharding/pool_manager.py new file mode 100644 index 00000000..12bbb143 --- /dev/null +++ b/app/db_sharding/pool_manager.py @@ -0,0 +1,216 @@ +"""SQLAlchemy engines and session factories for catalog + row shards.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List, Optional + +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine, make_url +from sqlalchemy.orm import Session, sessionmaker + +from app.db_sharding.router import ShardRouter + + +@dataclass(frozen=True) +class ShardEntry: + id: str + url: str + + +class DatabasePoolManager: + """Lazy-init pools from application settings.""" + + def __init__(self) -> None: + self._initialized = False + self._sharding_enabled = False + self._catalog_engine: Optional[Engine] = None + self._legacy_engine: Optional[Engine] = None + self._shard_engines: Dict[str, Engine] = {} + self._catalog_session_factory: Optional[sessionmaker] = None + self._legacy_session_factory: Optional[sessionmaker] = None + self._shard_session_factories: Dict[str, sessionmaker] = {} + self._router: Optional[ShardRouter] = None + self._shard_entries: List[ShardEntry] = [] + + @property + def sharding_enabled(self) -> bool: + self._ensure_initialized() + return self._sharding_enabled + + @property + def router(self) -> Optional[ShardRouter]: + self._ensure_initialized() + return self._router + + @property + def catalog_engine(self) -> Engine: + self._ensure_initialized() + if self._sharding_enabled: + assert self._catalog_engine is not None + return self._catalog_engine + assert self._legacy_engine is not None + return self._legacy_engine + + def shard_engine(self, shard_id: str) -> Engine: + self._ensure_initialized() + if not self._sharding_enabled: + return self.catalog_engine + try: + return self._shard_engines[shard_id] + except KeyError as exc: + raise KeyError(f"unknown shard id: {shard_id}") from exc + + def catalog_session_factory(self) -> sessionmaker: + self._ensure_initialized() + if self._sharding_enabled: + assert self._catalog_session_factory is not None + return self._catalog_session_factory + assert self._legacy_session_factory is not None + return self._legacy_session_factory + + def shard_session_factory(self, shard_id: str) -> sessionmaker: + self._ensure_initialized() + if not self._sharding_enabled: + return self.catalog_session_factory() + return self._shard_session_factories[shard_id] + + def all_engines_for_migrations(self) -> List[Engine]: + """Engines that should receive schema migrations (unique URLs).""" + self._ensure_initialized() + seen: set[str] = set() + engines: List[Engine] = [] + for eng in [self.catalog_engine, *self._shard_engines.values()]: + url = str(eng.url) + if url in seen: + continue + seen.add(url) + engines.append(eng) + return engines + + def reset(self) -> None: + """Dispose engines (tests).""" + for eng in list(self._shard_engines.values()): + eng.dispose() + if self._catalog_engine is not None: + self._catalog_engine.dispose() + if self._legacy_engine is not None: + self._legacy_engine.dispose() + self._initialized = False + self._catalog_engine = None + self._legacy_engine = None + self._shard_engines.clear() + self._catalog_session_factory = None + self._legacy_session_factory = None + self._shard_session_factories.clear() + self._router = None + self._shard_entries: List[ShardEntry] = [] + self._sharding_enabled = False + + def _ensure_initialized(self) -> None: + if not self._initialized: + from app.config import settings + + self._configure_from_settings(settings) + + def _configure_from_settings(self, settings) -> None: + pool_size = int(getattr(settings, "DB_POOL_SIZE", 10)) + max_overflow = int(getattr(settings, "DB_MAX_OVERFLOW", 20)) + + def make_engine(url: str) -> Engine: + return create_engine(url, **_create_engine_kwargs(url, pool_size, max_overflow)) + + enabled = bool(getattr(settings, "DB_SHARDING_ENABLED", False)) + database_url = settings.DATABASE_URL + if not database_url: + raise RuntimeError("DATABASE_URL is not configured") + + if not enabled: + self._legacy_engine = make_engine(database_url) + self._legacy_session_factory = sessionmaker( + autocommit=False, autoflush=False, bind=self._legacy_engine + ) + self._sharding_enabled = False + self._initialized = True + return + + catalog_url = getattr(settings, "DB_CATALOG_URL", None) or database_url + shard_entries = _parse_shard_entries(settings, fallback_url=database_url) + chunk_size = int(getattr(settings, "DB_SHARD_ROW_CHUNK_SIZE", 500)) + shard_pool_size = pool_size + shard_max_overflow = max_overflow + if len(shard_entries) > 1: + per_shard = max(4, pool_size // len(shard_entries)) + shard_pool_size = max(8, per_shard) + shard_max_overflow = max(8, max_overflow // len(shard_entries)) + + self._catalog_engine = make_engine(catalog_url) + self._catalog_session_factory = sessionmaker( + autocommit=False, autoflush=False, bind=self._catalog_engine + ) + self._shard_entries = shard_entries + for entry in shard_entries: + eng = create_engine( + entry.url, + **_create_engine_kwargs(entry.url, shard_pool_size, shard_max_overflow), + ) + self._shard_engines[entry.id] = eng + self._shard_session_factories[entry.id] = sessionmaker( + autocommit=False, autoflush=False, bind=eng + ) + self._router = ShardRouter( + [e.id for e in shard_entries], + row_chunk_size=chunk_size, + ) + self._sharding_enabled = True + self._initialized = True + + +def _create_engine_kwargs(url: str, pool_size: int, max_overflow: int) -> dict: + """Dialect-appropriate kwargs (SQLite tests use SingletonThreadPool).""" + dialect_name = make_url(url).get_backend_name() + if dialect_name == "sqlite": + return {"pool_pre_ping": True} + return { + "pool_pre_ping": True, + "pool_size": pool_size, + "max_overflow": max_overflow, + "connect_args": {"options": "-c timezone=UTC"}, + } + + +def _parse_shard_entries(settings, *, fallback_url: str) -> List[ShardEntry]: + raw = getattr(settings, "DB_SHARD_ENTRIES", None) or [] + entries: List[ShardEntry] = [] + for item in raw: + if not isinstance(item, dict): + continue + shard_id = str(item.get("id") or "").strip() + url = str(item.get("url") or "").strip() + if shard_id and url: + entries.append(ShardEntry(id=shard_id, url=url)) + if not entries: + entries.append(ShardEntry(id="data-shard-01", url=fallback_url)) + return entries + + +db_pool_manager = DatabasePoolManager() + + +def open_catalog_session() -> Session: + factory = db_pool_manager.catalog_session_factory() + return factory() + + +def open_row_shard_session( + call_import_id, + row_index: int, +) -> tuple[Session, str]: + manager = db_pool_manager + if not manager.sharding_enabled: + return open_catalog_session(), "legacy" + router = manager.router + assert router is not None + shard_id = router.shard_id_for_row(call_import_id, row_index) + session = manager.shard_session_factory(shard_id)() + return session, shard_id diff --git a/app/db_sharding/rebalance.py b/app/db_sharding/rebalance.py new file mode 100644 index 00000000..ba57b812 --- /dev/null +++ b/app/db_sharding/rebalance.py @@ -0,0 +1,581 @@ +"""Backfill / rebalance tooling: move call-import slice rows between shards.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, List, Optional, Sequence, Set +from uuid import UUID + +import redis +from loguru import logger +from sqlalchemy import and_, inspect, or_, select +from sqlalchemy.orm import Session + +from app.config import settings +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.row_ops import _reset_shard_write_role, _shard_write_without_catalog_fks +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import ( + CallImport, + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, + CallImportShardSlice, +) +from app.models.enums import CallImportRowStatus, CallImportStatus + +_REBALANCE_LOCK_KEY_PREFIX = "rebalance:import:" +_REBALANCE_LOCK_TTL_SECONDS = 3600 + +_redis_client: redis.Redis | None = None + +_IN_FLIGHT_IMPORT_ROW_STATUSES = frozenset( + {CallImportRowStatus.PENDING.value, CallImportRowStatus.PROCESSING.value} +) +_IN_FLIGHT_EVAL_ROW_STATUSES = frozenset({"pending", "running"}) + + +class RebalanceError(Exception): + """Operator-facing rebalance validation or execution failure.""" + + +@dataclass(frozen=True) +class SliceInfo: + slice_id: int + shard_id: str + row_index_min: int + row_index_max: int + + @property + def row_count(self) -> int: + return self.row_index_max - self.row_index_min + 1 + + +@dataclass(frozen=True) +class RebalancePlan: + call_import_id: UUID + from_shard_id: str + to_shard_id: str + slices: tuple[SliceInfo, ...] + import_row_count: int + eval_row_count: int + + +@dataclass(frozen=True) +class RebalanceResult: + dry_run: bool + call_import_id: UUID + from_shard_id: str + to_shard_id: str + slices_moved: int + import_rows_moved: int + eval_rows_moved: int + + +def _redis() -> redis.Redis: + global _redis_client + if _redis_client is None: + _redis_client = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis_client + + +def rebalance_lock_key(call_import_id: UUID | str) -> str: + return f"{_REBALANCE_LOCK_KEY_PREFIX}{call_import_id}" + + +def is_import_rebalance_locked(call_import_id: UUID | str) -> bool: + try: + return bool(_redis().get(rebalance_lock_key(call_import_id))) + except redis.RedisError: + return False + + +def acquire_rebalance_lock(call_import_id: UUID | str) -> bool: + """Best-effort pause: block fair dispatch for this import while rebalancing.""" + try: + return bool( + _redis().set( + rebalance_lock_key(call_import_id), + "1", + nx=True, + ex=_REBALANCE_LOCK_TTL_SECONDS, + ) + ) + except redis.RedisError as exc: + logger.warning("Could not acquire rebalance lock for {}: {}", call_import_id, exc) + return False + + +def release_rebalance_lock(call_import_id: UUID | str) -> None: + try: + _redis().delete(rebalance_lock_key(call_import_id)) + except redis.RedisError as exc: + logger.warning("Could not release rebalance lock for {}: {}", call_import_id, exc) + + +def filter_unlocked_call_import_ids( + call_import_ids: Iterable[UUID], +) -> List[UUID]: + """Drop imports that are mid-rebalance (fair dispatch should skip them).""" + return [cid for cid in call_import_ids if not is_import_rebalance_locked(cid)] + + +def filter_evaluations_not_rebalancing( + catalog_db: Session, + evaluation_ids: Iterable[UUID], +) -> List[UUID]: + """Drop evaluations whose parent import is mid-rebalance.""" + ids = list(evaluation_ids) + if not ids: + return [] + rows = catalog_db.execute( + select(CallImportEvaluation.id, CallImportEvaluation.call_import_id).where( + CallImportEvaluation.id.in_(ids) + ) + ).all() + return [ + evaluation_id + for evaluation_id, call_import_id in rows + if not is_import_rebalance_locked(call_import_id) + ] + + +def _require_sharding() -> None: + if not is_sharding_enabled(): + raise RebalanceError( + "database.sharding.enabled must be true (set catalog_url + shards in config)" + ) + + +def require_sharding_enabled() -> None: + _require_sharding() + + +def _configured_shard_ids() -> Set[str]: + router = db_pool_manager.router + if router is None: + return set() + return set(router.shard_ids) + + +def _validate_shard_id(shard_id: str, *, label: str) -> None: + configured = _configured_shard_ids() + if shard_id not in configured: + raise RebalanceError( + f"{label} shard {shard_id!r} is not configured " + f"(available: {sorted(configured)})" + ) + + +def list_shard_slices(catalog_db: Session, call_import_id: UUID) -> List[SliceInfo]: + rows = catalog_db.execute( + select( + CallImportShardSlice.slice_id, + CallImportShardSlice.shard_id, + CallImportShardSlice.row_index_min, + CallImportShardSlice.row_index_max, + ) + .where(CallImportShardSlice.call_import_id == call_import_id) + .order_by(CallImportShardSlice.slice_id.asc()) + ).all() + return [ + SliceInfo( + slice_id=int(slice_id), + shard_id=str(shard_id), + row_index_min=int(row_index_min), + row_index_max=int(row_index_max), + ) + for slice_id, shard_id, row_index_min, row_index_max in rows + ] + + +def _row_index_filter(model, slices: Sequence[SliceInfo]): + return or_( + *[ + and_( + model.row_index >= slice_info.row_index_min, + model.row_index <= slice_info.row_index_max, + ) + for slice_info in slices + ] + ) + + +def _orm_mapping(instance) -> dict: + return { + column.key: getattr(instance, column.key) + for column in inspect(instance).mapper.column_attrs + } + + +def _evaluation_ids_for_import(catalog_db: Session, call_import_id: UUID) -> List[UUID]: + rows = catalog_db.execute( + select(CallImportEvaluation.id).where( + CallImportEvaluation.call_import_id == call_import_id + ) + ).all() + return [row[0] for row in rows] + + +def _count_rows_on_shard( + shard_db: Session, + *, + call_import_id: UUID, + slices: Sequence[SliceInfo], + evaluation_ids: Sequence[UUID], +) -> tuple[int, int]: + import_count = ( + shard_db.query(CallImportRow.id) + .filter( + CallImportRow.call_import_id == call_import_id, + _row_index_filter(CallImportRow, slices), + ) + .count() + ) + eval_count = 0 + if evaluation_ids: + row_ids = [ + row_id + for (row_id,) in shard_db.query(CallImportRow.id) + .filter( + CallImportRow.call_import_id == call_import_id, + _row_index_filter(CallImportRow, slices), + ) + .all() + ] + if row_ids: + eval_count = ( + shard_db.query(CallImportEvaluationRow.id) + .filter(CallImportEvaluationRow.call_import_row_id.in_(row_ids)) + .count() + ) + return import_count, eval_count + + +def assert_import_rebalance_ready( + catalog_db: Session, + call_import_id: UUID, + *, + force: bool = False, +) -> None: + call_import = ( + catalog_db.query(CallImport) + .filter(CallImport.id == call_import_id) + .first() + ) + if call_import is None: + raise RebalanceError(f"call_import {call_import_id} not found on catalog") + + if call_import.status == CallImportStatus.DELETING: + raise RebalanceError("import is being deleted; rebalance is not allowed") + + if force: + return + + if call_import.status in { + CallImportStatus.PROCESSING, + CallImportStatus.PENDING, + }: + raise RebalanceError( + f"import status is {call_import.status.value!r}; " + "wait for terminal status or pass --force after pausing workers" + ) + + evaluation_ids = _evaluation_ids_for_import(catalog_db, call_import_id) + router = db_pool_manager.router + assert router is not None + + in_flight_import = 0 + in_flight_eval = 0 + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + in_flight_import += ( + shard_db.query(CallImportRow.id) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status.in_(_IN_FLIGHT_IMPORT_ROW_STATUSES), + ) + .count() + ) + if evaluation_ids: + in_flight_eval += ( + shard_db.query(CallImportEvaluationRow.id) + .filter( + CallImportEvaluationRow.evaluation_id.in_(evaluation_ids), + CallImportEvaluationRow.status.in_(_IN_FLIGHT_EVAL_ROW_STATUSES), + ) + .count() + ) + finally: + shard_db.close() + + if in_flight_import or in_flight_eval: + raise RebalanceError( + "import still has in-flight rows " + f"(import_rows={in_flight_import}, eval_rows={in_flight_eval}); " + "pause workers or pass --force" + ) + + +def build_rebalance_plan( + catalog_db: Session, + call_import_id: UUID, + *, + from_shard_id: str, + to_shard_id: str, + slice_ids: Optional[Iterable[int]] = None, +) -> RebalancePlan: + _require_sharding() + _validate_shard_id(from_shard_id, label="source") + _validate_shard_id(to_shard_id, label="target") + if from_shard_id == to_shard_id: + raise RebalanceError("source and target shard must differ") + + all_slices = list_shard_slices(catalog_db, call_import_id) + if not all_slices: + raise RebalanceError( + f"no registry slices for call_import {call_import_id}; " + "run materialize/register_shard_slices first" + ) + + selected = [s for s in all_slices if s.shard_id == from_shard_id] + if slice_ids is not None: + wanted = {int(value) for value in slice_ids} + selected = [s for s in selected if s.slice_id in wanted] + missing = wanted - {s.slice_id for s in selected} + if missing: + raise RebalanceError( + f"slice id(s) {sorted(missing)} not registered on shard {from_shard_id!r}" + ) + + if not selected: + raise RebalanceError( + f"no slices on shard {from_shard_id!r} for call_import {call_import_id}" + ) + + evaluation_ids = _evaluation_ids_for_import(catalog_db, call_import_id) + factory = db_pool_manager.shard_session_factory(from_shard_id) + shard_db = factory() + try: + import_count, eval_count = _count_rows_on_shard( + shard_db, + call_import_id=call_import_id, + slices=selected, + evaluation_ids=evaluation_ids, + ) + finally: + shard_db.close() + + return RebalancePlan( + call_import_id=call_import_id, + from_shard_id=from_shard_id, + to_shard_id=to_shard_id, + slices=tuple(selected), + import_row_count=import_count, + eval_row_count=eval_count, + ) + + +def _copy_rows_between_shards( + *, + call_import_id: UUID, + slices: Sequence[SliceInfo], + from_shard_id: str, + to_shard_id: str, +) -> tuple[int, int, List[UUID], List[UUID]]: + """Copy rows to the target shard; source rows are deleted after catalog commit.""" + source_db = db_pool_manager.shard_session_factory(from_shard_id)() + target_db = db_pool_manager.shard_session_factory(to_shard_id)() + import_moved = 0 + eval_moved = 0 + row_ids: List[UUID] = [] + eval_row_ids: List[UUID] = [] + try: + import_rows = ( + source_db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + _row_index_filter(CallImportRow, slices), + ) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + if not import_rows: + return 0, 0, [], [] + + row_ids = [row.id for row in import_rows] + eval_rows = ( + source_db.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.call_import_row_id.in_(row_ids)) + .all() + ) + eval_row_ids = [row.id for row in eval_rows] + + existing_import_ids = { + row_id + for (row_id,) in target_db.query(CallImportRow.id) + .filter(CallImportRow.id.in_(row_ids)) + .all() + } + existing_eval_ids = { + row_id + for (row_id,) in target_db.query(CallImportEvaluationRow.id) + .filter(CallImportEvaluationRow.id.in_(eval_row_ids)) + .all() + } + + import_to_insert = [ + _orm_mapping(row) for row in import_rows if row.id not in existing_import_ids + ] + eval_to_insert = [ + _orm_mapping(row) for row in eval_rows if row.id not in existing_eval_ids + ] + + if import_to_insert or eval_to_insert: + _shard_write_without_catalog_fks(target_db) + try: + if import_to_insert: + target_db.bulk_insert_mappings(CallImportRow, import_to_insert) + if eval_to_insert: + target_db.bulk_insert_mappings(CallImportEvaluationRow, eval_to_insert) + target_db.commit() + except Exception: + target_db.rollback() + raise + finally: + try: + _reset_shard_write_role(target_db) + except Exception: + pass + + import_moved = len(import_rows) + eval_moved = len(eval_rows) + finally: + source_db.close() + target_db.close() + + return import_moved, eval_moved, row_ids, eval_row_ids + + +def _delete_rows_from_source_shard( + *, + from_shard_id: str, + row_ids: Sequence[UUID], + eval_row_ids: Sequence[UUID], +) -> None: + """Remove copied rows from the source shard after the catalog registry is updated.""" + if not row_ids: + return + source_db = db_pool_manager.shard_session_factory(from_shard_id)() + try: + if eval_row_ids: + source_db.query(CallImportEvaluationRow).filter( + CallImportEvaluationRow.id.in_(list(eval_row_ids)) + ).delete(synchronize_session=False) + source_db.query(CallImportRow).filter(CallImportRow.id.in_(list(row_ids))).delete( + synchronize_session=False + ) + source_db.commit() + except Exception: + source_db.rollback() + raise + finally: + source_db.close() + + +def _update_slice_registry( + catalog_db: Session, + call_import_id: UUID, + slices: Sequence[SliceInfo], + *, + to_shard_id: str, +) -> None: + slice_id_set = {slice_info.slice_id for slice_info in slices} + for slice_info in slices: + catalog_db.merge( + CallImportShardSlice( + call_import_id=call_import_id, + slice_id=slice_info.slice_id, + shard_id=to_shard_id, + row_index_min=slice_info.row_index_min, + row_index_max=slice_info.row_index_max, + ) + ) + catalog_db.flush() + logger.info( + "Updated registry for call_import {} slices {} -> shard {}", + call_import_id, + sorted(slice_id_set), + to_shard_id, + ) + + +def execute_rebalance_slices( + catalog_db: Session, + plan: RebalancePlan, + *, + dry_run: bool = False, + force: bool = False, +) -> RebalanceResult: + _require_sharding() + assert_import_rebalance_ready(catalog_db, plan.call_import_id, force=force) + + if dry_run: + return RebalanceResult( + dry_run=True, + call_import_id=plan.call_import_id, + from_shard_id=plan.from_shard_id, + to_shard_id=plan.to_shard_id, + slices_moved=len(plan.slices), + import_rows_moved=plan.import_row_count, + eval_rows_moved=plan.eval_row_count, + ) + + lock_acquired = acquire_rebalance_lock(plan.call_import_id) + if not lock_acquired: + raise RebalanceError( + f"another rebalance lock is active for import {plan.call_import_id}" + ) + + try: + import_moved, eval_moved, row_ids, eval_row_ids = _copy_rows_between_shards( + call_import_id=plan.call_import_id, + slices=plan.slices, + from_shard_id=plan.from_shard_id, + to_shard_id=plan.to_shard_id, + ) + if plan.import_row_count > 0 and import_moved == 0: + raise RebalanceError( + f"expected {plan.import_row_count} import row(s) on " + f"{plan.from_shard_id!r} but found none; " + "registry may not match shard data" + ) + _update_slice_registry( + catalog_db, + plan.call_import_id, + plan.slices, + to_shard_id=plan.to_shard_id, + ) + catalog_db.commit() + _delete_rows_from_source_shard( + from_shard_id=plan.from_shard_id, + row_ids=row_ids, + eval_row_ids=eval_row_ids, + ) + except Exception: + catalog_db.rollback() + raise + finally: + release_rebalance_lock(plan.call_import_id) + + return RebalanceResult( + dry_run=False, + call_import_id=plan.call_import_id, + from_shard_id=plan.from_shard_id, + to_shard_id=plan.to_shard_id, + slices_moved=len(plan.slices), + import_rows_moved=import_moved, + eval_rows_moved=eval_moved, + ) diff --git a/app/db_sharding/registry.py b/app/db_sharding/registry.py new file mode 100644 index 00000000..88ecd2cd --- /dev/null +++ b/app/db_sharding/registry.py @@ -0,0 +1,30 @@ +"""Load shard slice registry from catalog DB for router overrides.""" + +from __future__ import annotations + +import uuid +from typing import Dict, Tuple + +from sqlalchemy import select +from sqlalchemy.orm import Session + + +def load_slice_registry_for_import( + db: Session, + call_import_id: uuid.UUID | str, +) -> Dict[Tuple[str, int], str]: + """ + Returns mapping (call_import_id str, slice_id) -> shard_id from + ``call_import_shard_slices``. Empty when table missing or no rows. + """ + from app.models.database import CallImportShardSlice + + cid = call_import_id if isinstance(call_import_id, uuid.UUID) else uuid.UUID(str(call_import_id)) + rows = db.execute( + select( + CallImportShardSlice.slice_id, + CallImportShardSlice.shard_id, + ).where(CallImportShardSlice.call_import_id == cid) + ).all() + key_prefix = str(cid) + return {(key_prefix, int(slice_id)): str(shard_id) for slice_id, shard_id in rows} diff --git a/app/db_sharding/router.py b/app/db_sharding/router.py new file mode 100644 index 00000000..d32454a7 --- /dev/null +++ b/app/db_sharding/router.py @@ -0,0 +1,75 @@ +"""Consistent-hash routing for call-import row shards.""" + +from __future__ import annotations + +import hashlib +import uuid +from typing import List, Sequence + + +def _normalize_import_id(call_import_id: uuid.UUID | str) -> str: + if isinstance(call_import_id, uuid.UUID): + return str(call_import_id) + return str(call_import_id) + + +class ShardRouter: + """ + Maps (call_import_id, row_index) to a configured shard id. + + slice_id = row_index // chunk_size; physical shard = + hash(call_import_id, slice_id) mod N. With one shard configured, + every route hits that shard (single-node enterprise mode). + """ + + def __init__( + self, + shard_ids: Sequence[str], + *, + row_chunk_size: int = 500, + ) -> None: + if row_chunk_size < 1: + raise ValueError("row_chunk_size must be >= 1") + if not shard_ids: + raise ValueError("at least one shard id is required when sharding is enabled") + self._shard_ids: List[str] = list(shard_ids) + self.row_chunk_size = row_chunk_size + + @property + def shard_ids(self) -> List[str]: + return list(self._shard_ids) + + @property + def shard_count(self) -> int: + return len(self._shard_ids) + + def slice_id_for_row_index(self, row_index: int) -> int: + if row_index < 0: + raise ValueError("row_index must be >= 0") + return row_index // self.row_chunk_size + + def shard_id_for_row( + self, + call_import_id: uuid.UUID | str, + row_index: int, + *, + slice_registry: dict[tuple[str, int], str] | None = None, + ) -> str: + slice_id = self.slice_id_for_row_index(row_index) + if slice_registry: + key = (_normalize_import_id(call_import_id), slice_id) + override = slice_registry.get(key) + if override is not None: + return override + return self.shard_id_for_slice(call_import_id, slice_id) + + def shard_id_for_slice( + self, + call_import_id: uuid.UUID | str, + slice_id: int, + ) -> str: + key = f"{_normalize_import_id(call_import_id)}:{slice_id}" + digest = hashlib.sha256(key.encode("utf-8")).digest() + bucket = int.from_bytes(digest[:8], "big") + idx = bucket % len(self._shard_ids) + return self._shard_ids[idx] diff --git a/app/db_sharding/row_ops.py b/app/db_sharding/row_ops.py new file mode 100644 index 00000000..8549e07c --- /dev/null +++ b/app/db_sharding/row_ops.py @@ -0,0 +1,560 @@ +"""Row placement, lookup, and bulk insert across shards.""" + +from __future__ import annotations + +from collections import defaultdict +from contextlib import contextmanager +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple +from uuid import UUID, uuid4 + +from sqlalchemy.orm import Session +from sqlalchemy import text + +from app.db_sharding.pool_manager import db_pool_manager, open_catalog_session +from app.db_sharding.registry import load_slice_registry_for_import +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import CallImport, CallImportRow, CallImportShardSlice +from app.models.database import CallImportEvaluationRow + + +def router_for_import(catalog_db: Session, call_import_id: UUID) -> Tuple[Any, Optional[dict]]: + """Return (router, slice_registry) for routing rows of this import.""" + manager = db_pool_manager + router = manager.router + if router is None: + return None, None + registry = load_slice_registry_for_import(catalog_db, call_import_id) + return router, registry or None + + +def shard_id_for_row( + catalog_db: Session, + call_import_id: UUID, + row_index: int, +) -> str: + if not is_sharding_enabled(): + return "legacy" + router, registry = router_for_import(catalog_db, call_import_id) + assert router is not None + return router.shard_id_for_row( + call_import_id, + row_index, + slice_registry=registry, + ) + + +def register_shard_slices( + catalog_db: Session, + call_import_id: UUID, + total_rows: int, +) -> None: + """Persist slice → shard assignments on the catalog for an import.""" + if not is_sharding_enabled() or total_rows <= 0: + return + router, _ = router_for_import(catalog_db, call_import_id) + assert router is not None + chunk = router.row_chunk_size + slice_meta: Dict[int, Dict[str, Any]] = {} + num_slices = (total_rows + chunk - 1) // chunk + for slice_id in range(num_slices): + row_index_min = slice_id * chunk + row_index_max = min((slice_id + 1) * chunk - 1, total_rows - 1) + shard_id = router.shard_id_for_row(call_import_id, row_index_min) + slice_meta[slice_id] = { + "shard_id": shard_id, + "row_index_min": row_index_min, + "row_index_max": row_index_max, + } + for slice_id, meta in slice_meta.items(): + catalog_db.merge( + CallImportShardSlice( + call_import_id=call_import_id, + slice_id=slice_id, + shard_id=meta["shard_id"], + row_index_min=meta["row_index_min"], + row_index_max=meta["row_index_max"], + ) + ) + catalog_db.flush() + + +def partition_mappings_by_shard( + catalog_db: Session, + call_import_id: UUID, + mappings: Iterable[dict], +) -> Dict[str, List[dict]]: + buckets: Dict[str, List[dict]] = defaultdict(list) + if not is_sharding_enabled(): + buckets["legacy"] = list(mappings) + return buckets + router, registry = router_for_import(catalog_db, call_import_id) + assert router is not None + for mapping in mappings: + row_index = int(mapping["row_index"]) + shard_id = router.shard_id_for_row( + call_import_id, + row_index, + slice_registry=registry, + ) + buckets[shard_id].append(mapping) + return buckets + + +def _shard_db_is_postgresql(shard_db: Session) -> bool: + bind = shard_db.get_bind() + return bind is not None and bind.dialect.name == "postgresql" + + +def _shard_write_without_catalog_fks(shard_db: Session) -> None: + """Allow row inserts on shards when parent rows live on catalog only.""" + if not _shard_db_is_postgresql(shard_db): + return + shard_db.execute(text("SET session_replication_role = replica")) + + +def _reset_shard_write_role(shard_db: Session) -> None: + if not _shard_db_is_postgresql(shard_db): + return + shard_db.execute(text("SET session_replication_role = DEFAULT")) + + +@contextmanager +def shard_row_write_context(shard_db: Session) -> Iterator[None]: + """Bypass catalog-only FK parents while mutating shard row tables.""" + if not is_sharding_enabled(): + yield + return + _shard_write_without_catalog_fks(shard_db) + try: + yield + finally: + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + + +def flush_shard_row_session(shard_db: Session) -> None: + with shard_row_write_context(shard_db): + shard_db.flush() + + +def commit_shard_row_session(shard_db: Session) -> None: + with shard_row_write_context(shard_db): + shard_db.commit() + + +def commit_pending_shard_sessions(sessions: List[Session]) -> None: + """Commit staged shard writes after the catalog transaction succeeds.""" + for shard_db in sessions: + try: + commit_shard_row_session(shard_db) + finally: + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + shard_db.close() + + +def rollback_pending_shard_sessions(sessions: List[Session]) -> None: + """Discard staged shard writes when the catalog transaction fails.""" + for shard_db in sessions: + try: + shard_db.rollback() + finally: + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + shard_db.close() + + +def bulk_insert_mappings_on_shards( + catalog_db: Session, + call_import_id: UUID, + mappings: List[dict], + *, + orm_class=CallImportRow, + defer_commit: bool = False, +) -> tuple[int, List[Session]]: + """Insert row mappings on the correct shard sessions; catalog_db unused when legacy. + + When ``defer_commit`` is true, rows are flushed on each shard but not committed + until ``commit_pending_shard_sessions`` runs after the catalog commit succeeds. + """ + if not mappings: + return 0, [] + if not is_sharding_enabled(): + catalog_db.bulk_insert_mappings(orm_class, mappings) + catalog_db.flush() + return len(mappings), [] + + from app.models.database import CallImportRow as RowModel + + buckets = partition_mappings_by_shard(catalog_db, call_import_id, mappings) + inserted = 0 + pending: List[Session] = [] + for shard_id, shard_mappings in buckets.items(): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + _shard_write_without_catalog_fks(shard_db) + shard_db.bulk_insert_mappings(RowModel, shard_mappings) + if defer_commit: + flush_shard_row_session(shard_db) + pending.append(shard_db) + else: + shard_db.commit() + inserted += len(shard_mappings) + except Exception: + shard_db.rollback() + raise + finally: + if not defer_commit: + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + shard_db.close() + return inserted, pending + + +def partition_eval_mappings_by_shard( + catalog_db: Session, + call_import_id: UUID, + mappings: Iterable[dict], + *, + index_by_source_id: dict[UUID, int], +) -> Dict[str, List[dict]]: + buckets: Dict[str, List[dict]] = defaultdict(list) + if not is_sharding_enabled(): + buckets["legacy"] = list(mappings) + return buckets + router, registry = router_for_import(catalog_db, call_import_id) + assert router is not None + for mapping in mappings: + source_id = mapping["call_import_row_id"] + row_index = index_by_source_id.get(source_id) + if row_index is None: + raise ValueError(f"unknown call_import_row_id for sharded eval insert: {source_id}") + shard_id = router.shard_id_for_row( + call_import_id, + row_index, + slice_registry=registry, + ) + buckets[shard_id].append(mapping) + return buckets + + +def bulk_insert_evaluation_rows_on_shards( + catalog_db: Session, + call_import_id: UUID, + evaluation_id: UUID, + source_row_ids: List[UUID], + *, + workspace_id: UUID, + index_by_source_id: dict[UUID, int], + defer_commit: bool = False, +) -> tuple[int, List[Session]]: + """Insert eval-row stubs on the same shard as each source import row.""" + from app.models.database import CallImportEvaluationRow as EvalRowModel + + if not source_row_ids: + return 0, [] + + def _mapping(source_row_id: UUID) -> dict: + return { + "id": uuid4(), + "evaluation_id": evaluation_id, + "call_import_row_id": source_row_id, + "workspace_id": workspace_id, + "status": "pending", + "metric_scores": {}, + } + + if not is_sharding_enabled(): + mappings = [_mapping(source_row_id) for source_row_id in source_row_ids] + catalog_db.bulk_insert_mappings(EvalRowModel, mappings) + catalog_db.flush() + return len(mappings), [] + + inserted = 0 + pending: List[Session] = [] + for start in range(0, len(source_row_ids), 500): + chunk = source_row_ids[start : start + 500] + mappings = [_mapping(source_row_id) for source_row_id in chunk] + buckets = partition_eval_mappings_by_shard( + catalog_db, + call_import_id, + mappings, + index_by_source_id=index_by_source_id, + ) + for shard_id, shard_mappings in buckets.items(): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + _shard_write_without_catalog_fks(shard_db) + shard_db.bulk_insert_mappings(EvalRowModel, shard_mappings) + if defer_commit: + flush_shard_row_session(shard_db) + pending.append(shard_db) + else: + shard_db.commit() + inserted += len(shard_mappings) + except Exception: + shard_db.rollback() + raise + finally: + if not defer_commit: + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + shard_db.close() + return inserted, pending + + +def locate_call_import_row( + row_id: UUID | str, +) -> Tuple[Session, Optional[Session], CallImportRow, str]: + """ + Find a call import row. Returns (row_db, catalog_db, row, shard_id). + + When sharding is off, row_db and catalog_db are the same session. + Caller must close session(s): if catalog_db is not row_db, close both. + """ + from app.database import SessionLocal + + rid = row_id if isinstance(row_id, UUID) else UUID(str(row_id)) + if not is_sharding_enabled(): + db = SessionLocal() + row = db.query(CallImportRow).filter(CallImportRow.id == rid).first() + if row is None: + db.close() + raise LookupError(f"call_import_row {rid} not found") + return db, db, row, "legacy" + + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = shard_db.query(CallImportRow).filter(CallImportRow.id == rid).first() + if row is not None: + catalog_db = open_catalog_session() + return shard_db, catalog_db, row, shard_id + except Exception: + shard_db.close() + raise + shard_db.close() + raise LookupError(f"call_import_row {rid} not found on any shard") + + +def locate_call_import_evaluation_row( + eval_row_id: UUID | str, +) -> Tuple[Session, Optional[Session], CallImportEvaluationRow, CallImportRow, str]: + """ + Find eval + source rows. Returns + (row_db, catalog_db, eval_row, source_row, shard_id). + """ + from app.database import SessionLocal + from app.models.database import CallImportEvaluationRow + + eid = eval_row_id if isinstance(eval_row_id, UUID) else UUID(str(eval_row_id)) + if not is_sharding_enabled(): + db = SessionLocal() + row = ( + db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.id == eid) + .first() + ) + if row is None: + db.close() + raise LookupError(f"call_import_evaluation_row {eid} not found") + eval_row, source_row = row + return db, db, eval_row, source_row, "legacy" + + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = ( + shard_db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.id == eid) + .first() + ) + if row is not None: + eval_row, source_row = row + catalog_db = open_catalog_session() + return shard_db, catalog_db, eval_row, source_row, shard_id + except Exception: + shard_db.close() + raise + shard_db.close() + raise LookupError(f"call_import_evaluation_row {eid} not found on any shard") + + +def close_row_sessions(row_db: Session, catalog_db: Optional[Session]) -> None: + row_db.close() + if catalog_db is not None and catalog_db is not row_db: + catalog_db.close() + + +def update_call_import_rows_on_shards( + catalog_db: Session, + call_import_id: UUID, + updates: List[dict], +) -> None: + """Apply field updates to import rows on the correct DB session(s).""" + if not updates: + return + + allowed_fields = { + "diarised_transcript_status", + "diarised_transcript_error", + "celery_task_id", + } + + if not is_sharding_enabled(): + by_id = {item["id"]: item for item in updates} + rows = ( + catalog_db.query(CallImportRow) + .filter(CallImportRow.id.in_(list(by_id.keys()))) + .all() + ) + for row in rows: + patch = by_id.get(row.id) + if patch is None: + continue + for field in allowed_fields: + if field in patch: + setattr(row, field, patch[field]) + catalog_db.commit() + return + + by_shard: Dict[str, List[dict]] = defaultdict(list) + for item in updates: + row_index = int(item.get("row_index", 0)) + shard_id = shard_id_for_row(catalog_db, call_import_id, row_index) + by_shard[shard_id].append(item) + + for shard_id, shard_updates in by_shard.items(): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + by_id = {item["id"]: item for item in shard_updates} + rows = ( + shard_db.query(CallImportRow) + .filter(CallImportRow.id.in_(list(by_id.keys()))) + .all() + ) + for row in rows: + patch = by_id.get(row.id) + if patch is None: + continue + for field in allowed_fields: + if field in patch: + setattr(row, field, patch[field]) + commit_shard_row_session(shard_db) + except Exception: + shard_db.rollback() + raise + finally: + shard_db.close() + + +def delete_call_import_rows_on_shards( + catalog_db: Session, + call_import_id: UUID, + rows: List[CallImportRow], +) -> int: + """Delete import rows from catalog or row shards.""" + if not rows: + return 0 + + row_ids = [row.id for row in rows if row.id is not None] + if not row_ids: + return 0 + + if not is_sharding_enabled(): + deleted = ( + catalog_db.query(CallImportRow) + .filter( + CallImportRow.id.in_(row_ids), + CallImportRow.call_import_id == call_import_id, + ) + .delete(synchronize_session=False) + ) + catalog_db.flush() + return int(deleted or 0) + + by_shard: Dict[str, List[UUID]] = defaultdict(list) + for row in rows: + if row.id is None: + continue + shard_id = shard_id_for_row( + catalog_db, + call_import_id, + int(row.row_index or 0), + ) + by_shard[shard_id].append(row.id) + + deleted_total = 0 + for shard_id, shard_row_ids in by_shard.items(): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + deleted = ( + shard_db.query(CallImportRow) + .filter( + CallImportRow.id.in_(shard_row_ids), + CallImportRow.call_import_id == call_import_id, + ) + .delete(synchronize_session=False) + ) + commit_shard_row_session(shard_db) + deleted_total += int(deleted or 0) + except Exception: + shard_db.rollback() + raise + finally: + shard_db.close() + return deleted_total + + +def new_row_mapping( + *, + call_import: CallImport, + organization_id: UUID, + workspace_id: UUID, + row_index: int, + row: dict, + transcript_source: Optional[str], +) -> dict: + csv_transcript = row["transcript"] + return { + "id": uuid4(), + "call_import_id": call_import.id, + "organization_id": organization_id, + "workspace_id": workspace_id, + "row_index": row_index, + "conversation_id": row["conversation_id"], + "recording_date": row.get("recording_date"), + "recording_url": row["recording_url"], + "transcript": csv_transcript, + "transcript_source": transcript_source, + "raw_columns": row.get("parameter_values") or None, + "status": row.get("status"), + } diff --git a/app/db_sharding/scatter_gather.py b/app/db_sharding/scatter_gather.py new file mode 100644 index 00000000..917acc31 --- /dev/null +++ b/app/db_sharding/scatter_gather.py @@ -0,0 +1,1086 @@ +"""Scatter-gather reads and fan-out writes across row shards.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, TYPE_CHECKING +from uuid import UUID + +from sqlalchemy import case, func +from sqlalchemy.orm import Session, load_only + +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import CallImportEvaluation, CallImportEvaluationRow, CallImportRow + +if TYPE_CHECKING: + from app.db_sharding.session_cache import ShardSessionCache + +T = TypeVar("T") + + +def _open_shard_session( + shard_id: str, + shard_cache: Optional["ShardSessionCache"], +) -> tuple[Session, bool]: + """Return (session, should_close). Uses cache when provided.""" + if shard_cache is not None: + return shard_cache.session_for(shard_id), False + factory = db_pool_manager.shard_session_factory(shard_id) + return factory(), True + + +def shard_ids_for_import(catalog_db: Session, call_import_id: UUID) -> List[str]: + """Shard ids that may hold rows for this import (from registry or all shards).""" + if not is_sharding_enabled(): + return ["legacy"] + from app.db_sharding.registry import load_slice_registry_for_import + + registry = load_slice_registry_for_import(catalog_db, call_import_id) + if registry: + return sorted({v for v in registry.values()}) + router = db_pool_manager.router + assert router is not None + return list(router.shard_ids) + + +def scatter_gather_on_shards( + shard_ids: Sequence[str], + fn: Callable[[Session, str], T], + *, + max_workers: Optional[int] = None, +) -> List[T]: + if not is_sharding_enabled() or len(shard_ids) <= 1: + sid = shard_ids[0] if shard_ids else "legacy" + if sid == "legacy": + from app.database import SessionLocal + + db = SessionLocal() + try: + return [fn(db, sid)] + finally: + db.close() + factory = db_pool_manager.shard_session_factory(sid) + db = factory() + try: + return [fn(db, sid)] + finally: + db.close() + + workers = max_workers or min(len(shard_ids), 6) + results: List[T] = [] + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = {} + for shard_id in shard_ids: + futures[pool.submit(_run_on_shard, shard_id, fn)] = shard_id + for future in as_completed(futures): + results.append(future.result()) + return results + + +def _run_on_shard(shard_id: str, fn: Callable[[Session, str], T]) -> T: + factory = db_pool_manager.shard_session_factory(shard_id) + db = factory() + try: + return fn(db, shard_id) + finally: + db.close() + + +def merge_rows_by_index(rows: Iterable[CallImportRow]) -> List[CallImportRow]: + return sorted(rows, key=lambda r: int(r.row_index or 0)) + + +def _apply_call_import_row_filters( + query, + *, + search_term: str, + diarised_status_filter: Optional[str], +): + if search_term: + query = query.filter( + CallImportRow.conversation_id.ilike(f"%{search_term}%") + ) + if diarised_status_filter: + query = query.filter( + CallImportRow.diarised_transcript_status == diarised_status_filter + ) + return query + + +def _merged_call_import_rows_for_import( + catalog_db: Session, + call_import_id: UUID, + *, + search_term: str = "", + diarised_status_filter: Optional[str] = None, +) -> List[CallImportRow]: + """All matching rows merged by ``row_index`` (scatter-gather when sharded).""" + if not is_sharding_enabled(): + query = catalog_db.query(CallImportRow).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return query.order_by(CallImportRow.row_index.asc()).all() + + def load_shard(db: Session, _shard_id: str) -> List[CallImportRow]: + query = db.query(CallImportRow).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return query.order_by(CallImportRow.row_index.asc()).all() + + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + merged: List[CallImportRow] = [] + for part in scatter_gather_on_shards(shard_ids, load_shard): + merged.extend(part) + merged = merge_rows_by_index(merged) + if merged: + return merged + + query = catalog_db.query(CallImportRow).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return query.order_by(CallImportRow.row_index.asc()).all() + + +def count_call_import_rows_filtered( + catalog_db: Session, + call_import_id: UUID, + *, + search_term: str = "", + diarised_status_filter: Optional[str] = None, +) -> int: + if not is_sharding_enabled(): + query = catalog_db.query(func.count(CallImportRow.id)).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return int(query.scalar() or 0) + + total = 0 + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def count_shard(db: Session, _shard_id: str) -> int: + query = db.query(func.count(CallImportRow.id)).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return int(query.scalar() or 0) + + for part in scatter_gather_on_shards(shard_ids, count_shard): + total += int(part) + if total == 0: + query = catalog_db.query(func.count(CallImportRow.id)).filter( + CallImportRow.call_import_id == call_import_id + ) + query = _apply_call_import_row_filters( + query, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + total = int(query.scalar() or 0) + return total + + +def fetch_call_import_rows_filtered_page( + catalog_db: Session, + call_import_id: UUID, + *, + search_term: str = "", + diarised_status_filter: Optional[str] = None, + offset: int, + limit: int, +) -> List[CallImportRow]: + merged = _merged_call_import_rows_for_import( + catalog_db, + call_import_id, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return merged[offset : offset + limit] + + +def list_call_import_row_ids_filtered( + catalog_db: Session, + call_import_id: UUID, + *, + search_term: str = "", + diarised_status_filter: Optional[str] = None, +) -> List[UUID]: + merged = _merged_call_import_rows_for_import( + catalog_db, + call_import_id, + search_term=search_term, + diarised_status_filter=diarised_status_filter, + ) + return [row.id for row in merged if row.id is not None] + + +def pending_eval_workspaces( + catalog_db: Session, + *, + shard_cache: Optional["ShardSessionCache"] = None, +) -> List[UUID]: + """Workspaces with pending eval rows (catalog + shard queries when sharded).""" + if not is_sharding_enabled(): + return _pending_eval_workspaces_mono(catalog_db) + + evaluation_ids: set[UUID] = set() + + def collect_pending(db: Session, _shard_id: str) -> List[UUID]: + rows = ( + db.query(CallImportEvaluationRow.evaluation_id) + .filter( + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return [row[0] for row in rows if row[0] is not None] + + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + shard_db, close_shard = _open_shard_session(shard_id, shard_cache) + try: + for eid in collect_pending(shard_db, shard_id): + evaluation_ids.add(eid) + finally: + if close_shard: + shard_db.close() + + if not evaluation_ids: + return [] + + from app.db_sharding.rebalance import filter_evaluations_not_rebalancing + + evaluation_ids = set(filter_evaluations_not_rebalancing(catalog_db, evaluation_ids)) + if not evaluation_ids: + return [] + + rows = ( + catalog_db.query(CallImportEvaluation.workspace_id) + .filter( + CallImportEvaluation.id.in_(evaluation_ids), + CallImportEvaluation.status != "cancelled", + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + +def _pending_eval_workspaces_mono(db: Session) -> List[UUID]: + from app.models.database import CallImportEvaluation + + rows = ( + db.query(CallImportEvaluation.workspace_id) + .join( + CallImportEvaluationRow, + CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + ) + .filter( + CallImportEvaluation.status != "cancelled", + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + +def evaluations_with_pending_rows( + catalog_db: Session, + workspace_id: UUID, + *, + shard_cache: Optional["ShardSessionCache"] = None, +) -> List[UUID]: + if not is_sharding_enabled(): + return _evaluations_with_pending_mono(catalog_db, workspace_id) + + from app.models.database import CallImportEvaluation + + eval_ids = [ + row[0] + for row in catalog_db.query(CallImportEvaluation.id) + .filter( + CallImportEvaluation.workspace_id == workspace_id, + CallImportEvaluation.status != "cancelled", + ) + .all() + if row[0] is not None + ] + if not eval_ids: + return [] + + pending: set[UUID] = set() + router = db_pool_manager.router + assert router is not None + + for shard_id in router.shard_ids: + shard_db, close_shard = _open_shard_session(shard_id, shard_cache) + try: + rows = ( + shard_db.query(CallImportEvaluationRow.evaluation_id) + .filter( + CallImportEvaluationRow.evaluation_id.in_(eval_ids), + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + pending.update(row[0] for row in rows if row[0] is not None) + finally: + if close_shard: + shard_db.close() + from app.db_sharding.rebalance import filter_evaluations_not_rebalancing + + return sorted(filter_evaluations_not_rebalancing(catalog_db, pending)) + + +def _evaluations_with_pending_mono(db: Session, workspace_id: UUID) -> List[UUID]: + from app.models.database import CallImportEvaluation + + rows = ( + db.query(CallImportEvaluation.id) + .join( + CallImportEvaluationRow, + CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + ) + .filter( + CallImportEvaluation.workspace_id == workspace_id, + CallImportEvaluation.status != "cancelled", + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .distinct() + .all() + ) + return sorted({row[0] for row in rows if row[0] is not None}) + + +def pending_eval_row_triples( + catalog_db: Session, + evaluation_id: UUID, + *, + limit: int, + shard_cache: Optional["ShardSessionCache"] = None, +) -> List[tuple]: + """Return (eval_row, source_row, evaluation) tuples up to limit.""" + from app.models.database import CallImportEvaluation + + evaluation = ( + catalog_db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == evaluation_id) + .first() + ) + if evaluation is None: + return [] + from app.db_sharding.rebalance import is_import_rebalance_locked + + if is_import_rebalance_locked(evaluation.call_import_id): + return [] + + if not is_sharding_enabled(): + return ( + catalog_db.query(CallImportEvaluationRow, CallImportRow, CallImportEvaluation) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .join( + CallImportEvaluation, + CallImportEvaluation.id == CallImportEvaluationRow.evaluation_id, + ) + .filter( + CallImportEvaluation.id == evaluation_id, + CallImportEvaluation.status != "cancelled", + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .order_by(CallImportEvaluationRow.created_at.asc()) + .limit(limit) + .all() + ) + + collected: List[tuple] = [] + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + if len(collected) >= limit: + break + shard_db, close_shard = _open_shard_session(shard_id, shard_cache) + try: + batch = ( + shard_db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter( + CallImportEvaluationRow.evaluation_id == evaluation_id, + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .order_by(CallImportEvaluationRow.created_at.asc()) + .limit(max(1, limit - len(collected))) + .all() + ) + for eval_row, source_row in batch: + collected.append((eval_row, source_row, evaluation)) + finally: + if close_shard: + shard_db.close() + return collected + + +def aggregate_evaluation_row_counts( + catalog_db: Session, + evaluation_id: UUID, +) -> tuple[int, int, int]: + """Return (total, completed, failed) across all shards.""" + from sqlalchemy import case, func + + if not is_sharding_enabled(): + row = ( + catalog_db.query( + func.count(CallImportEvaluationRow.id), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "completed", 1), + else_=0, + ) + ), + 0, + ), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "failed", 1), + else_=0, + ) + ), + 0, + ), + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .one() + ) + return int(row[0] or 0), int(row[1] or 0), int(row[2] or 0) + + total = completed = failed = 0 + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = ( + shard_db.query( + func.count(CallImportEvaluationRow.id), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "completed", 1), + else_=0, + ) + ), + 0, + ), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "failed", 1), + else_=0, + ) + ), + 0, + ), + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .one() + ) + total += int(row[0] or 0) + completed += int(row[1] or 0) + failed += int(row[2] or 0) + finally: + shard_db.close() + return total, completed, failed + + +def count_eval_rows_in_progress(catalog_db: Session, evaluation_id: UUID) -> int: + from sqlalchemy import func + + if not is_sharding_enabled(): + return int( + catalog_db.query(func.count()) + .filter( + CallImportEvaluationRow.evaluation_id == evaluation_id, + CallImportEvaluationRow.status.in_(["pending", "running"]), + ) + .scalar() + or 0 + ) + + total = 0 + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + total += int( + shard_db.query(func.count()) + .filter( + CallImportEvaluationRow.evaluation_id == evaluation_id, + CallImportEvaluationRow.status.in_(["pending", "running"]), + ) + .scalar() + or 0 + ) + finally: + shard_db.close() + return total + + +def fetch_call_import_rows_page( + catalog_db: Session, + call_import_id: UUID, + *, + offset: int, + limit: int, +) -> List[CallImportRow]: + """Scatter-gather row list merged by row_index.""" + if not is_sharding_enabled(): + return ( + catalog_db.query(CallImportRow) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .offset(offset) + .limit(limit) + .all() + ) + + merged = _merged_call_import_rows_for_import(catalog_db, call_import_id) + return merged[offset : offset + limit] + + +def _load_call_import_rows_scatter( + catalog_db: Session, + call_import_id: UUID, + *, + columns: tuple, + requested_row_ids: Optional[List[UUID]] = None, +) -> List[CallImportRow]: + """Load import rows with ``load_only`` columns (catalog or scatter-gather).""" + if not is_sharding_enabled(): + query = ( + catalog_db.query(CallImportRow) + .options(load_only(*columns)) + .filter(CallImportRow.call_import_id == call_import_id) + ) + if requested_row_ids: + query = query.filter(CallImportRow.id.in_(requested_row_ids)) + return query.order_by(CallImportRow.row_index.asc()).all() + + def load_shard(db: Session, _shard_id: str) -> List[CallImportRow]: + query = ( + db.query(CallImportRow) + .options(load_only(*columns)) + .filter(CallImportRow.call_import_id == call_import_id) + ) + if requested_row_ids: + query = query.filter(CallImportRow.id.in_(requested_row_ids)) + return query.order_by(CallImportRow.row_index.asc()).all() + + merged: List[CallImportRow] = [] + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + for part in scatter_gather_on_shards(shard_ids, load_shard): + merged.extend(part) + merged = merge_rows_by_index(merged) + if not merged: + query = ( + catalog_db.query(CallImportRow) + .options(load_only(*columns)) + .filter(CallImportRow.call_import_id == call_import_id) + ) + if requested_row_ids: + query = query.filter(CallImportRow.id.in_(requested_row_ids)) + merged = query.order_by(CallImportRow.row_index.asc()).all() + elif requested_row_ids: + requested = set(requested_row_ids) + merged = [row for row in merged if row.id in requested] + return merged + + +def load_call_import_rows_for_transcription( + catalog_db: Session, + call_import_id: UUID, + *, + requested_row_ids: Optional[List[UUID]] = None, +) -> List[CallImportRow]: + """Rows needed to decide diarization enqueue (shard-aware when enabled).""" + columns = ( + CallImportRow.id, + CallImportRow.row_index, + CallImportRow.recording_s3_key, + CallImportRow.diarised_transcript, + CallImportRow.diarised_transcript_status, + CallImportRow.diarised_transcript_error, + CallImportRow.celery_task_id, + ) + return _load_call_import_rows_scatter( + catalog_db, + call_import_id, + columns=columns, + requested_row_ids=requested_row_ids, + ) + + +def load_call_import_rows_for_delete( + catalog_db: Session, + call_import_id: UUID, + row_ids: List[UUID], +) -> List[CallImportRow]: + """Rows targeted for bulk delete (shard-aware when enabled).""" + columns = ( + CallImportRow.id, + CallImportRow.row_index, + CallImportRow.recording_s3_key, + CallImportRow.celery_task_id, + CallImportRow.status, + ) + return _load_call_import_rows_scatter( + catalog_db, + call_import_id, + columns=columns, + requested_row_ids=row_ids, + ) + + +def _legacy_catalog_rows( + catalog_db: Session, + call_import_id: UUID, +) -> List[CallImportRow]: + """Pre-sharding rows still stored on the catalog DB (same DB as headers).""" + return ( + catalog_db.query(CallImportRow) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + + +def aggregate_diarised_transcript_counts( + catalog_db: Session, + call_import_id: UUID, +) -> Dict[str, int]: + """Batch-wide diarisation status counts (scatter-gather when sharded).""" + if not is_sharding_enabled(): + counts: Dict[str, int] = {} + for status_value, count in ( + catalog_db.query(CallImportRow.diarised_transcript_status, func.count()) + .filter(CallImportRow.call_import_id == call_import_id) + .group_by(CallImportRow.diarised_transcript_status) + .all() + ): + if isinstance(status_value, str): + counts[status_value] = int(count or 0) + return counts + + merged: Dict[str, int] = {} + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def count_shard(db: Session, _shard_id: str) -> Dict[str, int]: + part: Dict[str, int] = {} + for status_value, count in ( + db.query(CallImportRow.diarised_transcript_status, func.count()) + .filter(CallImportRow.call_import_id == call_import_id) + .group_by(CallImportRow.diarised_transcript_status) + .all() + ): + if isinstance(status_value, str): + part[status_value] = int(count or 0) + return part + + for part in scatter_gather_on_shards(shard_ids, count_shard): + for key, value in part.items(): + merged[key] = merged.get(key, 0) + int(value) + if not merged: + for status_value, count in ( + catalog_db.query(CallImportRow.diarised_transcript_status, func.count()) + .filter(CallImportRow.call_import_id == call_import_id) + .group_by(CallImportRow.diarised_transcript_status) + .all() + ): + if isinstance(status_value, str): + merged[status_value] = int(count or 0) + return merged + + +def count_call_import_rows( + catalog_db: Session, + call_import_id: UUID, +) -> int: + """Total rows for an import (shards + optional legacy catalog copy).""" + if not is_sharding_enabled(): + return int( + catalog_db.query(func.count(CallImportRow.id)) + .filter(CallImportRow.call_import_id == call_import_id) + .scalar() + or 0 + ) + total = 0 + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def count_shard(db: Session, _shard_id: str) -> int: + return int( + db.query(func.count(CallImportRow.id)) + .filter(CallImportRow.call_import_id == call_import_id) + .scalar() + or 0 + ) + + for part in scatter_gather_on_shards(shard_ids, count_shard): + total += int(part) + if total == 0: + total = int( + catalog_db.query(func.count(CallImportRow.id)) + .filter(CallImportRow.call_import_id == call_import_id) + .scalar() + or 0 + ) + return total + + +def count_completed_call_import_rows( + catalog_db: Session, + call_import_id: UUID, +) -> int: + """Completed import rows (scatter-gather when sharded).""" + from app.models.enums import CallImportRowStatus + + if not is_sharding_enabled(): + return int( + catalog_db.query(func.count(CallImportRow.id)) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .scalar() + or 0 + ) + + total = 0 + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def count_shard(db: Session, _shard_id: str) -> int: + return int( + db.query(func.count(CallImportRow.id)) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .scalar() + or 0 + ) + + for part in scatter_gather_on_shards(shard_ids, count_shard): + total += int(part) + if total == 0: + total = int( + catalog_db.query(func.count(CallImportRow.id)) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .scalar() + or 0 + ) + return total + + +def list_completed_source_row_ids_ordered( + catalog_db: Session, + call_import_id: UUID, +) -> List[UUID]: + """Completed source row ids ordered by row_index (shard-aware).""" + from app.models.enums import CallImportRowStatus + + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImportRow.id) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + return [row_id for (row_id,) in rows if row_id is not None] + + merged: List[CallImportRow] = [] + + def load_shard(db: Session, _shard_id: str) -> List[CallImportRow]: + return ( + db.query(CallImportRow) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + for part in scatter_gather_on_shards(shard_ids, load_shard): + merged.extend(part) + merged = merge_rows_by_index(merged) + if not merged: + rows = ( + catalog_db.query(CallImportRow.id) + .filter( + CallImportRow.call_import_id == call_import_id, + CallImportRow.status == CallImportRowStatus.COMPLETED, + ) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + return [row_id for (row_id,) in rows if row_id is not None] + return [row.id for row in merged if row.id is not None] + + +def list_source_row_ids_ordered( + catalog_db: Session, + call_import_id: UUID, +) -> List[UUID]: + """All source row ids for an import, ordered by row_index.""" + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImportRow.id) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + return [row[0] for row in rows if row[0] is not None] + + merged: List[Tuple[UUID, int]] = [ + (row_id, int(row_index or 0)) + for row_id, row_index in catalog_db.query( + CallImportRow.id, + CallImportRow.row_index, + ) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + if row_id is not None + ] + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def load_shard(db: Session, _shard_id: str) -> List[Tuple[UUID, int]]: + return [ + (row_id, int(row_index or 0)) + for row_id, row_index in db.query( + CallImportRow.id, + CallImportRow.row_index, + ) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .all() + if row_id is not None + ] + + for part in scatter_gather_on_shards(shard_ids, load_shard): + merged.extend(part) + merged.sort(key=lambda pair: pair[1]) + if merged: + return [row_id for row_id, _ in merged] + + rows = ( + catalog_db.query(CallImportRow.id) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + return [row[0] for row in rows if row[0] is not None] + + +def source_row_index_map( + catalog_db: Session, + call_import_id: UUID, +) -> dict[UUID, int]: + """Map call_import_row id -> row_index (for shard routing eval rows).""" + if not is_sharding_enabled(): + rows = ( + catalog_db.query(CallImportRow.id, CallImportRow.row_index) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + ) + return {row_id: int(row_index or 0) for row_id, row_index in rows} + + merged: List[Tuple[UUID, int]] = [ + (row_id, int(row_index or 0)) + for row_id, row_index in catalog_db.query( + CallImportRow.id, + CallImportRow.row_index, + ) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + if row_id is not None + ] + shard_ids = shard_ids_for_import(catalog_db, call_import_id) + + def load_shard(db: Session, _shard_id: str) -> List[Tuple[UUID, int]]: + return [ + (row_id, int(row_index or 0)) + for row_id, row_index in db.query( + CallImportRow.id, + CallImportRow.row_index, + ) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + if row_id is not None + ] + + for part in scatter_gather_on_shards(shard_ids, load_shard): + merged.extend(part) + merged.sort(key=lambda pair: pair[1]) + if merged: + return {row_id: row_index for row_id, row_index in merged} + + rows = ( + catalog_db.query(CallImportRow.id, CallImportRow.row_index) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + ) + return {row_id: int(row_index or 0) for row_id, row_index in rows} + + +def count_evaluation_cancel_targets_sharded( + catalog_db: Session, + evaluation_id: UUID, + *, + pending_only: bool, + in_progress_only: bool, +) -> int: + """Count eval rows eligible for cancel across row shards.""" + from sqlalchemy import func + + if not is_sharding_enabled(): + query = catalog_db.query(func.count(CallImportEvaluationRow.id)).filter( + CallImportEvaluationRow.evaluation_id == evaluation_id + ) + if in_progress_only: + query = query.filter( + CallImportEvaluationRow.status.in_(("pending", "running")) + ) + elif pending_only: + query = query.filter(CallImportEvaluationRow.status == "pending") + return int(query.scalar() or 0) + + total = 0 + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + query = shard_db.query(func.count(CallImportEvaluationRow.id)).filter( + CallImportEvaluationRow.evaluation_id == evaluation_id + ) + if in_progress_only: + query = query.filter( + CallImportEvaluationRow.status.in_(("pending", "running")) + ) + elif pending_only: + query = query.filter(CallImportEvaluationRow.status == "pending") + total += int(query.scalar() or 0) + finally: + shard_db.close() + return total + + +def _evaluation_row_pair_sort_key( + pair: tuple[CallImportEvaluationRow, CallImportRow], +) -> tuple[int, str]: + _, source_row = pair + return (int(source_row.row_index or 0), str(pair[0].id)) + + +def _sort_evaluation_row_pairs( + pairs: List[tuple[CallImportEvaluationRow, CallImportRow]], +) -> List[tuple[CallImportEvaluationRow, CallImportRow]]: + return sorted(pairs, key=_evaluation_row_pair_sort_key) + + +def load_evaluation_row_pairs( + catalog_db: Session, + evaluation_id: UUID, +) -> List[tuple[CallImportEvaluationRow, CallImportRow]]: + """All eval/source row pairs for insights and PDF aggregation.""" + if not is_sharding_enabled(): + pairs = ( + catalog_db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .all() + ) + return _sort_evaluation_row_pairs(pairs) + + pairs: List[tuple[CallImportEvaluationRow, CallImportRow]] = [] + router = db_pool_manager.router + assert router is not None + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + pairs.extend( + shard_db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .all() + ) + finally: + shard_db.close() + if not pairs: + pairs = ( + catalog_db.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + .all() + ) + return _sort_evaluation_row_pairs(pairs) diff --git a/app/db_sharding/session_cache.py b/app/db_sharding/session_cache.py new file mode 100644 index 00000000..ea67b0aa --- /dev/null +++ b/app/db_sharding/session_cache.py @@ -0,0 +1,28 @@ +"""Reuse SQLAlchemy shard sessions within a worker task (fair dispatch, scatter reads).""" + +from __future__ import annotations + +from sqlalchemy.orm import Session + + +class ShardSessionCache: + """One open session per shard id for the lifetime of a dispatch/scatter pass.""" + + def __init__(self) -> None: + self._by_shard: dict[str, Session] = {} + + def session_for(self, shard_id: str) -> Session: + from app.db_sharding.pool_manager import db_pool_manager + + existing = self._by_shard.get(shard_id) + if existing is not None: + return existing + factory = db_pool_manager.shard_session_factory(shard_id) + session = factory() + self._by_shard[shard_id] = session + return session + + def close_all(self) -> None: + for session in self._by_shard.values(): + session.close() + self._by_shard.clear() diff --git a/app/db_sharding/sessions.py b/app/db_sharding/sessions.py new file mode 100644 index 00000000..45a44c64 --- /dev/null +++ b/app/db_sharding/sessions.py @@ -0,0 +1,37 @@ +"""Context managers for catalog and row-shard sessions.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Iterator, Tuple +from uuid import UUID + +from sqlalchemy.orm import Session + +from app.db_sharding.pool_manager import db_pool_manager, open_row_shard_session +from app.db_sharding.pool_manager import open_catalog_session + + +@contextmanager +def catalog_session() -> Iterator[Session]: + db = open_catalog_session() + try: + yield db + finally: + db.close() + + +@contextmanager +def row_shard_session( + call_import_id: UUID | str, + row_index: int, +) -> Iterator[Tuple[Session, str]]: + db, shard_id = open_row_shard_session(call_import_id, row_index) + try: + yield db, shard_id + finally: + db.close() + + +def is_sharding_enabled() -> bool: + return db_pool_manager.sharding_enabled diff --git a/app/migrations/054_call_import_sharding.py b/app/migrations/054_call_import_sharding.py new file mode 100644 index 00000000..272cb425 --- /dev/null +++ b/app/migrations/054_call_import_sharding.py @@ -0,0 +1,230 @@ +""" +Migration: Call-import sharding registry, workspace denorm, dispatch indexes. + +Catalog table ``call_import_shard_slices`` records which shard owns each +slice of rows for an import. Denormalized ``workspace_id`` on row tables +supports shard-local queries without joining catalog parents. +""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Add call_import_shard_slices registry, workspace_id on row tables, " + "and dispatch indexes for evaluation and diarization workers." +) + +# ``all`` until catalog/shard migration split (Phase 8 hardening). +MIGRATION_SCOPE = "all" + + +def _column_exists(db: Session, table_name: str, column_name: str) -> bool: + row = db.execute( + text( + """ + SELECT 1 + FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table_name, "column_name": column_name}, + ).first() + return row is not None + + +def _table_exists(db: Session, table_name: str) -> bool: + row = db.execute( + text( + """ + SELECT 1 + FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table_name}, + ).first() + return row is not None + + +def _index_exists(db: Session, index_name: str) -> bool: + row = db.execute( + text( + """ + SELECT 1 FROM pg_indexes WHERE indexname = :index_name + """ + ), + {"index_name": index_name}, + ).first() + return row is not None + + +def upgrade(db: Session): + if not _table_exists(db, "call_import_shard_slices"): + db.execute( + text( + """ + CREATE TABLE call_import_shard_slices ( + call_import_id UUID NOT NULL + REFERENCES call_imports(id) ON DELETE CASCADE, + slice_id INTEGER NOT NULL, + shard_id VARCHAR(64) NOT NULL, + row_index_min INTEGER NOT NULL, + row_index_max INTEGER NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (call_import_id, slice_id) + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_call_import_shard_slices_shard + ON call_import_shard_slices (shard_id, call_import_id) + """ + ) + ) + print("Created call_import_shard_slices") + + if _table_exists(db, "call_import_rows") and not _column_exists( + db, "call_import_rows", "workspace_id" + ): + db.execute( + text( + """ + ALTER TABLE call_import_rows + ADD COLUMN workspace_id UUID NULL + """ + ) + ) + db.execute( + text( + """ + UPDATE call_import_rows r + SET workspace_id = c.workspace_id + FROM call_imports c + WHERE r.call_import_id = c.id AND r.workspace_id IS NULL + """ + ) + ) + db.execute( + text( + """ + ALTER TABLE call_import_rows + ALTER COLUMN workspace_id SET NOT NULL + """ + ) + ) + db.execute( + text( + """ + ALTER TABLE call_import_rows + ADD CONSTRAINT fk_call_import_rows_workspace + FOREIGN KEY (workspace_id) REFERENCES workspaces(id) ON DELETE RESTRICT + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_call_import_rows_workspace_id + ON call_import_rows (workspace_id) + """ + ) + ) + print("Added call_import_rows.workspace_id") + + if _table_exists(db, "call_import_evaluation_rows") and not _column_exists( + db, "call_import_evaluation_rows", "workspace_id" + ): + db.execute( + text( + """ + ALTER TABLE call_import_evaluation_rows + ADD COLUMN workspace_id UUID NULL + """ + ) + ) + db.execute( + text( + """ + UPDATE call_import_evaluation_rows er + SET workspace_id = c.workspace_id + FROM call_import_evaluations e + JOIN call_imports c ON c.id = e.call_import_id + WHERE er.evaluation_id = e.id AND er.workspace_id IS NULL + """ + ) + ) + db.execute( + text( + """ + ALTER TABLE call_import_evaluation_rows + ALTER COLUMN workspace_id SET NOT NULL + """ + ) + ) + db.execute( + text( + """ + ALTER TABLE call_import_evaluation_rows + ADD CONSTRAINT fk_call_import_eval_rows_workspace + FOREIGN KEY (workspace_id) REFERENCES workspaces(id) ON DELETE RESTRICT + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_call_import_evaluation_rows_workspace_id + ON call_import_evaluation_rows (workspace_id) + """ + ) + ) + print("Added call_import_evaluation_rows.workspace_id") + + if _table_exists(db, "call_import_evaluation_rows") and not _index_exists( + db, "ix_cier_eval_status_task" + ): + db.execute( + text( + """ + CREATE INDEX ix_cier_eval_status_task + ON call_import_evaluation_rows (evaluation_id, status, celery_task_id) + """ + ) + ) + print("Created ix_cier_eval_status_task") + + if _table_exists(db, "call_import_rows") and not _index_exists( + db, "ix_cir_import_diarise_task" + ): + db.execute( + text( + """ + CREATE INDEX ix_cir_import_diarise_task + ON call_import_rows (call_import_id, diarised_transcript_status, celery_task_id) + """ + ) + ) + print("Created ix_cir_import_diarise_task") + + db.commit() + + +def downgrade(db: Session): + if _index_exists(db, "ix_cir_import_diarise_task"): + db.execute(text("DROP INDEX ix_cir_import_diarise_task")) + if _index_exists(db, "ix_cier_eval_status_task"): + db.execute(text("DROP INDEX ix_cier_eval_status_task")) + if _column_exists(db, "call_import_evaluation_rows", "workspace_id"): + db.execute( + text( + "ALTER TABLE call_import_evaluation_rows DROP COLUMN workspace_id" + ) + ) + if _column_exists(db, "call_import_rows", "workspace_id"): + db.execute(text("ALTER TABLE call_import_rows DROP COLUMN workspace_id")) + if _table_exists(db, "call_import_shard_slices"): + db.execute(text("DROP TABLE call_import_shard_slices")) + db.commit() diff --git a/app/migrations/catalog/README.md b/app/migrations/catalog/README.md new file mode 100644 index 00000000..ff207759 --- /dev/null +++ b/app/migrations/catalog/README.md @@ -0,0 +1,4 @@ +# Catalog-only migrations (headers, registry, integrations). +# Postgres database name is typically ``efficientai_catalog`` (see config.yml.example). +# Set ``MIGRATION_SCOPE = "catalog"`` on new migration modules in this directory +# once the hard split is enforced (Phase 8). diff --git a/app/migrations/shard/README.md b/app/migrations/shard/README.md new file mode 100644 index 00000000..de41e3ea --- /dev/null +++ b/app/migrations/shard/README.md @@ -0,0 +1,3 @@ +# Data-plane shard migrations (row-heavy tables; call-import rows today). +# Postgres databases are typically ``efficientai_data_01``, ``efficientai_data_02``, … +# Set ``MIGRATION_SCOPE = "shard"`` on new migration modules here after split. diff --git a/app/models/database.py b/app/models/database.py index bdb17647..7e3466d5 100644 --- a/app/models/database.py +++ b/app/models/database.py @@ -16,6 +16,7 @@ String, Text, UniqueConstraint, + select, text, ) from sqlalchemy.dialects.postgresql import UUID @@ -1879,6 +1880,23 @@ class CallImport(Base): ) +class CallImportShardSlice(Base): + """Registry row: which shard stores a slice of rows for an import.""" + + __tablename__ = "call_import_shard_slices" + + call_import_id = Column( + UUID(as_uuid=True), + ForeignKey("call_imports.id", ondelete="CASCADE"), + primary_key=True, + ) + slice_id = Column(Integer, primary_key=True) + shard_id = Column(String(64), nullable=False, index=True) + row_index_min = Column(Integer, nullable=False) + row_index_max = Column(Integer, nullable=False) + created_at = Column(DateTime(timezone=True), server_default=func.now()) + + class CallImportRow(Base): """A single row within a CallImport batch (one CSV line / one external call).""" @@ -1895,6 +1913,12 @@ class CallImportRow(Base): index=True, ) organization_id = Column(UUID(as_uuid=True), ForeignKey("organizations.id"), nullable=False, index=True) + workspace_id = Column( + UUID(as_uuid=True), + ForeignKey("workspaces.id", ondelete="RESTRICT"), + nullable=False, + index=True, + ) row_index = Column(Integer, nullable=False) # Was historically named ``external_call_id``; renamed to @@ -2038,6 +2062,20 @@ class CallImportRow(Base): call_import = relationship("CallImport", back_populates="rows") +@event.listens_for(CallImportRow, "before_insert") +def _call_import_row_fill_workspace_id(_mapper, connection, target): + """Denormalize workspace_id from the parent import when omitted.""" + if target.workspace_id is not None or target.call_import_id is None: + return + workspace_id = connection.execute( + select(CallImport.workspace_id).where( + CallImport.id == target.call_import_id + ) + ).scalar_one_or_none() + if workspace_id is not None: + target.workspace_id = workspace_id + + class CallImportTag(Base): """User-defined tag that can be attached to one or more call imports. @@ -2315,6 +2353,12 @@ class CallImportEvaluationRow(Base): nullable=False, index=True, ) + workspace_id = Column( + UUID(as_uuid=True), + ForeignKey("workspaces.id", ondelete="RESTRICT"), + nullable=False, + index=True, + ) status = Column(String(20), nullable=False, default="pending", index=True) # Same shape as EvaluatorResult.metric_scores: {metric_id_str: {value, type, metric_name, ...}} @@ -2333,6 +2377,20 @@ class CallImportEvaluationRow(Base): source_row = relationship("CallImportRow") +@event.listens_for(CallImportEvaluationRow, "before_insert") +def _call_import_evaluation_row_fill_workspace_id(_mapper, connection, target): + """Denormalize workspace_id from the parent evaluation when omitted.""" + if target.workspace_id is not None or target.evaluation_id is None: + return + workspace_id = connection.execute( + select(CallImportEvaluation.workspace_id).where( + CallImportEvaluation.id == target.evaluation_id + ) + ).scalar_one_or_none() + if workspace_id is not None: + target.workspace_id = workspace_id + + class CallImportEvaluationReportSnapshot(Base): """Persisted PDF-report aggregate used for period-over-period deltas.""" diff --git a/app/models/schemas.py b/app/models/schemas.py index bcbb1635..5f1dc4c2 100644 --- a/app/models/schemas.py +++ b/app/models/schemas.py @@ -4376,6 +4376,14 @@ class EvaluationMetricClustersRequest(BaseModel): "all completed rows with at least one flagged quality metric are used." ), ) + row_limit: Optional[int] = Field( + default=None, + ge=1, + description=( + "Use the first N eligible rows (by row order). Mutually exclusive " + "with evaluation_row_ids." + ), + ) failure_policies: Optional[Dict[str, MetricFailurePolicy]] = Field( default=None, description="Per-metric failure policies confirmed in the cluster modal.", diff --git a/app/services/call_imports/bulk_ops.py b/app/services/call_imports/bulk_ops.py index 14c903fd..96be5883 100644 --- a/app/services/call_imports/bulk_ops.py +++ b/app/services/call_imports/bulk_ops.py @@ -13,15 +13,17 @@ from loguru import logger from sqlalchemy.orm import Session, load_only +from sqlalchemy import func + +from app.db_sharding.sessions import is_sharding_enabled from app.models.database import ( CallImport, CallImportEvaluation, CallImportEvaluationRow, CallImportRow, - CallImportRowStatus, ) -from app.models.enums import CallImportStatus +from app.models.enums import CallImportRowStatus, CallImportStatus from app.models.schemas import CallImportTranscribeRequest _BULK_INSERT_CHUNK = 1000 @@ -68,14 +70,13 @@ def select_rows_for_transcription( requested_row_ids: Optional[List[UUID]] = None, ) -> tuple[List[CallImportRow], Dict[str, int]]: """Pick rows to diarise, loading only columns needed for the decision.""" - query = ( - db.query(CallImportRow) - .options(load_only(*_ROW_TRANSCRIBE_COLUMNS)) - .filter(CallImportRow.call_import_id == call_import.id) + from app.db_sharding.scatter_gather import load_call_import_rows_for_transcription + + rows = load_call_import_rows_for_transcription( + db, + call_import.id, + requested_row_ids=requested_row_ids, ) - if requested_row_ids: - query = query.filter(CallImportRow.id.in_(requested_row_ids)) - rows = query.order_by(CallImportRow.row_index.asc()).all() if requested_row_ids: found_ids = {r.id for r in rows} @@ -152,19 +153,34 @@ def execute_bulk_diarization( stored_set = set(stored_ids) failed_set = set(failed_ids) + updates: List[dict] = [] queued = 0 for row in rows: if row.id in stored_set: - row.diarised_transcript_status = "pending" - row.diarised_transcript_error = None - row.celery_task_id = None + updates.append( + { + "id": row.id, + "row_index": row.row_index, + "diarised_transcript_status": "pending", + "diarised_transcript_error": None, + "celery_task_id": None, + } + ) queued += 1 elif row.id in failed_set: - row.diarised_transcript_status = "failed" - row.diarised_transcript_error = _REDIS_PARAMS_STORE_ERROR + updates.append( + { + "id": row.id, + "row_index": row.row_index, + "diarised_transcript_status": "failed", + "diarised_transcript_error": _REDIS_PARAMS_STORE_ERROR, + } + ) - if stored_set or failed_set: - db.commit() + if updates: + from app.db_sharding.row_ops import update_call_import_rows_on_shards + + update_call_import_rows_on_shards(db, call_import.id, updates) if queued > 0: schedule_fair_diarization_dispatch(max_workspace_turns=999) @@ -177,35 +193,22 @@ def execute_bulk_diarization( def count_completed_source_rows(db: Session, call_import_id: UUID) -> int: - from sqlalchemy import func + from app.db_sharding.scatter_gather import count_completed_call_import_rows - return int( - db.query(func.count(CallImportRow.id)) - .filter( - CallImportRow.call_import_id == call_import_id, - CallImportRow.status == CallImportRowStatus.COMPLETED, - ) - .scalar() - or 0 - ) + return count_completed_call_import_rows(db, call_import_id) def _completed_source_row_ids(db: Session, call_import_id: UUID) -> List[UUID]: - return [ - row_id - for (row_id,) in ( - db.query(CallImportRow.id) - .filter( - CallImportRow.call_import_id == call_import_id, - CallImportRow.status == CallImportRowStatus.COMPLETED, - ) - .order_by(CallImportRow.row_index.asc()) - .all() - ) - ] + from app.db_sharding.scatter_gather import list_completed_source_row_ids_ordered + + return list_completed_source_row_ids_ordered(db, call_import_id) def count_all_source_rows(db: Session, call_import_id: UUID) -> int: + from app.db_sharding.scatter_gather import count_call_import_rows + + if is_sharding_enabled(): + return count_call_import_rows(db, call_import_id) from sqlalchemy import func return int( @@ -217,6 +220,10 @@ def count_all_source_rows(db: Session, call_import_id: UUID) -> int: def _all_source_row_ids(db: Session, call_import_id: UUID) -> List[UUID]: + from app.db_sharding.scatter_gather import list_source_row_ids_ordered + + if is_sharding_enabled(): + return list_source_row_ids_ordered(db, call_import_id) return [ row_id for (row_id,) in ( @@ -230,12 +237,31 @@ def _all_source_row_ids(db: Session, call_import_id: UUID) -> List[UUID]: def bulk_insert_evaluation_rows( db: Session, + call_import_id: UUID, evaluation_id: UUID, source_row_ids: List[UUID], -) -> None: + *, + workspace_id: UUID, +) -> List[Session]: """Insert eval-row stubs in chunks without per-row ORM overhead.""" if not source_row_ids: - return + return [] + + if is_sharding_enabled(): + from app.db_sharding.row_ops import bulk_insert_evaluation_rows_on_shards + from app.db_sharding.scatter_gather import source_row_index_map + + index_by_id = source_row_index_map(db, call_import_id) + _inserted, pending = bulk_insert_evaluation_rows_on_shards( + db, + call_import_id, + evaluation_id, + source_row_ids, + workspace_id=workspace_id, + index_by_source_id=index_by_id, + defer_commit=True, + ) + return pending for start in range(0, len(source_row_ids), _BULK_INSERT_CHUNK): chunk = source_row_ids[start : start + _BULK_INSERT_CHUNK] @@ -244,6 +270,7 @@ def bulk_insert_evaluation_rows( "id": uuid4(), "evaluation_id": evaluation_id, "call_import_row_id": source_row_id, + "workspace_id": workspace_id, "status": "pending", "metric_scores": {}, } @@ -251,6 +278,7 @@ def bulk_insert_evaluation_rows( ] db.bulk_insert_mappings(CallImportEvaluationRow, mappings) db.flush() + return [] def materialize_and_enqueue_evaluation( @@ -284,32 +312,32 @@ def materialize_and_enqueue_evaluation( db.commit() return - bulk_insert_evaluation_rows(db, evaluation_id, source_row_ids) - db.commit() + pending_shard_sessions: List[Session] = [] + try: + pending_shard_sessions = bulk_insert_evaluation_rows( + db, + evaluation.call_import_id, + evaluation_id, + source_row_ids, + workspace_id=evaluation.workspace_id, + ) + db.commit() + from app.db_sharding.row_ops import commit_pending_shard_sessions - eval_rows = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) - .all() - ) - source_rows = ( - db.query(CallImportRow) - .options(load_only(*_EVAL_SOURCE_COLUMNS)) - .filter(CallImportRow.id.in_(source_row_ids)) - .all() - ) - source_by_id = {row.id: row for row in source_rows} - bucket: List[Tuple[CallImportEvaluationRow, CallImportRow]] = [] - for eval_row in eval_rows: - source_row = source_by_id.get(eval_row.call_import_row_id) - if source_row is not None: - bucket.append((eval_row, source_row)) + commit_pending_shard_sessions(pending_shard_sessions) + pending_shard_sessions = [] + except Exception: + db.rollback() + from app.db_sharding.row_ops import rollback_pending_shard_sessions + + rollback_pending_shard_sessions(pending_shard_sessions) + raise try: _enqueue_eval_rows_with_optional_transcribe( db, evaluation, - bucket, + [], transcribe_overwrite=transcribe_overwrite, ) evaluation.status = "running" @@ -339,31 +367,17 @@ def execute_bulk_row_delete( if not row_ids: return 0 - rows = ( - db.query(CallImportRow) - .options(load_only(*_ROW_DELETE_COLUMNS)) - .filter( - CallImportRow.id.in_(row_ids), - CallImportRow.call_import_id == call_import.id, - ) - .all() - ) + from app.db_sharding.row_ops import delete_call_import_rows_on_shards + from app.db_sharding.scatter_gather import load_call_import_rows_for_delete + + rows = load_call_import_rows_for_delete(db, call_import.id, row_ids) if not rows: return 0 _revoke_pending_tasks(rows) _delete_s3_objects(organization_id, call_import.id, rows) - deleted_ids = [row.id for row in rows] - deleted = ( - db.query(CallImportRow) - .filter( - CallImportRow.id.in_(deleted_ids), - CallImportRow.call_import_id == call_import.id, - ) - .delete(synchronize_session=False) - ) - db.flush() + deleted = delete_call_import_rows_on_shards(db, call_import.id, rows) _recompute_call_import_counters(db, call_import) db.commit() @@ -514,13 +528,16 @@ def bulk_materialize_call_import_rows( call_import: CallImport, parsed_rows: List[Dict[str, Any]], organization_id: UUID, -) -> int: + *, + defer_shard_commit: bool = False, +) -> tuple[int, List[Session]]: """Insert import rows in chunks without per-row ORM overhead.""" from app.api.v1.routes.call_imports import _parse_recording_date_cell if not parsed_rows: - return 0 + return 0, [] + pending: List[Session] = [] for start in range(0, len(parsed_rows), _BULK_INSERT_CHUNK): chunk = parsed_rows[start : start + _BULK_INSERT_CHUNK] mappings = [] @@ -532,6 +549,7 @@ def bulk_materialize_call_import_rows( "id": uuid4(), "call_import_id": call_import.id, "organization_id": organization_id, + "workspace_id": call_import.workspace_id, "row_index": idx, "conversation_id": row["conversation_id"], "recording_date": ( @@ -548,9 +566,17 @@ def bulk_materialize_call_import_rows( "status": CallImportRowStatus.PENDING, } ) - db.bulk_insert_mappings(CallImportRow, mappings) + from app.db_sharding.row_ops import bulk_insert_mappings_on_shards + + _inserted, shard_pending = bulk_insert_mappings_on_shards( + db, + call_import.id, + mappings, + defer_commit=defer_shard_commit, + ) + pending.extend(shard_pending) db.flush() - return len(parsed_rows) + return len(parsed_rows), pending def execute_call_import_materialization( @@ -606,8 +632,12 @@ def execute_call_import_materialization( .filter(CallImportRow.call_import_id == call_import.id) .limit(1) .first() + if not is_sharding_enabled() + else None ) - if existing_rows is not None: + if existing_rows is not None or ( + is_sharding_enabled() and int(call_import.total_rows or 0) > 0 + ): logger.info( "execute_call_import_materialization: import {} already has rows", call_import_id, @@ -663,20 +693,41 @@ def execute_call_import_materialization( call_import.completed_rows = 0 call_import.failed_rows = 0 + row_count = len(parsed_rows) + pending_shard_sessions: List[Session] = [] try: - row_count = bulk_materialize_call_import_rows( + if is_sharding_enabled(): + from app.db_sharding.row_ops import register_shard_slices + + register_shard_slices(db, call_import.id, row_count) + + row_count, pending_shard_sessions = bulk_materialize_call_import_rows( db, call_import, parsed_rows, organization_id, + defer_shard_commit=is_sharding_enabled(), ) + if not is_sharding_enabled(): + from app.db_sharding.row_ops import register_shard_slices + + register_shard_slices(db, call_import.id, row_count) db.commit() + if is_sharding_enabled(): + from app.db_sharding.row_ops import commit_pending_shard_sessions + + commit_pending_shard_sessions(pending_shard_sessions) + pending_shard_sessions = [] except Exception as exc: # noqa: BLE001 + db.rollback() + if is_sharding_enabled(): + from app.db_sharding.row_ops import rollback_pending_shard_sessions + + rollback_pending_shard_sessions(pending_shard_sessions) logger.exception( "Failed to materialize rows for call import {}", call_import_id, ) - db.rollback() call_import.status = CallImportStatus.FAILED call_import.error_message = f"Failed to materialize import rows: {exc}" db.commit() @@ -702,6 +753,101 @@ def execute_call_import_materialization( return {"total_rows": row_count, "status": "processing"} +def _aggregate_import_row_status_counts( + db: Session, + call_import_id: UUID, +) -> tuple[int, int, int, int, int]: + """Return total, completed, failed, pending, processing via SQL aggregate.""" + if is_sharding_enabled(): + from app.db_sharding.pool_manager import db_pool_manager + from app.db_sharding.scatter_gather import shard_ids_for_import + + total = completed = failed = pending = processing = 0 + for shard_id in shard_ids_for_import(db, call_import_id): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + row = ( + shard_db.query( + func.count(CallImportRow.id), + func.count().filter( + CallImportRow.status == CallImportRowStatus.COMPLETED + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.FAILED + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.PENDING + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.PROCESSING + ), + ) + .filter(CallImportRow.call_import_id == call_import_id) + .one() + ) + total += int(row[0] or 0) + completed += int(row[1] or 0) + failed += int(row[2] or 0) + pending += int(row[3] or 0) + processing += int(row[4] or 0) + finally: + shard_db.close() + if total == 0: + row = ( + db.query( + func.count(CallImportRow.id), + func.count().filter( + CallImportRow.status == CallImportRowStatus.COMPLETED + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.FAILED + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.PENDING + ), + func.count().filter( + CallImportRow.status == CallImportRowStatus.PROCESSING + ), + ) + .filter(CallImportRow.call_import_id == call_import_id) + .one() + ) + total = int(row[0] or 0) + if total > 0: + return ( + total, + int(row[1] or 0), + int(row[2] or 0), + int(row[3] or 0), + int(row[4] or 0), + ) + return total, completed, failed, pending, processing + + row = ( + db.query( + func.count(CallImportRow.id), + func.count().filter( + CallImportRow.status == CallImportRowStatus.COMPLETED + ), + func.count().filter(CallImportRow.status == CallImportRowStatus.FAILED), + func.count().filter(CallImportRow.status == CallImportRowStatus.PENDING), + func.count().filter( + CallImportRow.status == CallImportRowStatus.PROCESSING + ), + ) + .filter(CallImportRow.call_import_id == call_import_id) + .one() + ) + return ( + int(row[0] or 0), + int(row[1] or 0), + int(row[2] or 0), + int(row[3] or 0), + int(row[4] or 0), + ) + + def rollup_call_import_batch_status(db: Session, call_import: CallImport) -> None: """Recompute batch counters and terminal status on the parent import. @@ -711,17 +857,8 @@ def rollup_call_import_batch_status(db: Session, call_import: CallImport) -> Non parent stuck in ``processing`` once every evaluation run is terminal. """ - counts = ( - db.query(CallImportRow.status) - .filter(CallImportRow.call_import_id == call_import.id) - .all() - ) - total = len(counts) - completed = sum(1 for (status,) in counts if status == CallImportRowStatus.COMPLETED) - failed = sum(1 for (status,) in counts if status == CallImportRowStatus.FAILED) - pending = sum(1 for (status,) in counts if status == CallImportRowStatus.PENDING) - processing = sum( - 1 for (status,) in counts if status == CallImportRowStatus.PROCESSING + total, completed, failed, pending, processing = _aggregate_import_row_status_counts( + db, call_import.id ) pending_or_processing = pending + processing @@ -773,6 +910,10 @@ def rollup_call_import_batch_status(db: Session, call_import: CallImport) -> Non else: call_import.status = CallImportStatus.PARTIAL + from app.services.call_imports.progress_counters import clear_import_progress_redis + + clear_import_progress_redis(call_import.id) + _EVAL_CANCEL_COLUMNS = ( CallImportEvaluationRow.id, @@ -788,6 +929,15 @@ def count_evaluation_cancel_targets( mode: Literal["abort", "force_fail_pending"], ) -> int: """Count rows eligible for bulk cancel without loading full ORM objects.""" + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import count_evaluation_cancel_targets_sharded + + return count_evaluation_cancel_targets_sharded( + db, + evaluation_id, + pending_only=(mode == "force_fail_pending"), + in_progress_only=(mode == "abort"), + ) from sqlalchemy import func query = db.query(func.count(CallImportEvaluationRow.id)).filter( @@ -837,10 +987,7 @@ def execute_evaluation_cancel( mode: Literal["abort", "force_fail_pending"], ) -> dict: """Cancel eval rows in chunks off the API thread.""" - from app.api.v1.routes.call_import_evaluations import ( - EVAL_CANCELLED_BY_USER_ERROR, - _rollup_evaluation_status, - ) + from app.api.v1.routes.call_import_evaluations import _rollup_evaluation_status evaluation = ( db.query(CallImportEvaluation) @@ -856,41 +1003,17 @@ def execute_evaluation_cancel( try: cancelled_total = 0 - while True: - query = ( - db.query(CallImportEvaluationRow) - .options(load_only(*_EVAL_CANCEL_COLUMNS)) - .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + if is_sharding_enabled(): + cancelled_total = _cancel_evaluation_rows_all_shards( + evaluation_id, + mode=mode, ) - if mode == "abort": - query = query.filter( - CallImportEvaluationRow.status.in_(("pending", "running")) - ) - else: - query = query.filter(CallImportEvaluationRow.status == "pending") - - rows = ( - query.order_by(CallImportEvaluationRow.id.asc()) - .limit(_BULK_INSERT_CHUNK) - .all() + else: + cancelled_total = _cancel_evaluation_rows_on_session( + db, + evaluation_id, + mode=mode, ) - if not rows: - break - - task_ids: List[str] = [] - now = datetime.now(timezone.utc) - for row in rows: - task_id = (row.celery_task_id or "").strip() - if task_id: - task_ids.append(task_id) - row.status = "failed" - row.error_message = EVAL_CANCELLED_BY_USER_ERROR - row.finished_at = now - row.celery_task_id = None - cancelled_total += 1 - - _batch_revoke_celery_task_ids(task_ids, terminate=True) - db.commit() db.refresh(evaluation) _rollup_evaluation_status(evaluation, db) @@ -909,6 +1032,172 @@ def execute_evaluation_cancel( clear_evaluation_bulk_operation(evaluation_id) +def _cancel_evaluation_rows_on_session( + db: Session, + evaluation_id: UUID, + *, + mode: Literal["abort", "force_fail_pending"], +) -> int: + from app.api.v1.routes.call_import_evaluations import EVAL_CANCELLED_BY_USER_ERROR + + cancelled_total = 0 + while True: + query = ( + db.query(CallImportEvaluationRow) + .options(load_only(*_EVAL_CANCEL_COLUMNS)) + .filter(CallImportEvaluationRow.evaluation_id == evaluation_id) + ) + if mode == "abort": + query = query.filter( + CallImportEvaluationRow.status.in_(("pending", "running")) + ) + else: + query = query.filter(CallImportEvaluationRow.status == "pending") + + rows = ( + query.order_by(CallImportEvaluationRow.id.asc()) + .limit(_BULK_INSERT_CHUNK) + .all() + ) + if not rows: + break + + task_ids: List[str] = [] + now = datetime.now(timezone.utc) + for row in rows: + task_id = (row.celery_task_id or "").strip() + if task_id: + task_ids.append(task_id) + row.status = "failed" + row.error_message = EVAL_CANCELLED_BY_USER_ERROR + row.finished_at = now + row.celery_task_id = None + cancelled_total += 1 + + _batch_revoke_celery_task_ids(task_ids, terminate=True) + db.commit() + return cancelled_total + + +def _cancel_evaluation_rows_all_shards( + evaluation_id: UUID, + *, + mode: Literal["abort", "force_fail_pending"], +) -> int: + from app.db_sharding.pool_manager import db_pool_manager + + router = db_pool_manager.router + assert router is not None + total = 0 + for shard_id in router.shard_ids: + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + total += _cancel_evaluation_rows_on_session( + shard_db, + evaluation_id, + mode=mode, + ) + finally: + shard_db.close() + return total + + +def _persist_evaluation_retry_targets( + catalog_db: Session, + evaluation: CallImportEvaluation, + targets: List[Tuple[CallImportEvaluationRow, CallImportRow]], + *, + metric_ids: Optional[List[UUID]] = None, + transcribe_overwrite: bool = False, +) -> None: + """Apply retry resets on the correct DB session (per shard when sharding).""" + from collections import defaultdict + + from sqlalchemy.orm.attributes import flag_modified + + from app.api.v1.routes.call_import_evaluations import ( + _prepare_source_row_for_retry, + _reset_eval_row_for_retry, + ) + from app.db_sharding.pool_manager import db_pool_manager + from app.db_sharding.row_ops import shard_id_for_row + + task_ids: List[str] = [] + for eval_row, _source_row in targets: + if eval_row.celery_task_id and eval_row.status in {"pending", "running"}: + task_id = (eval_row.celery_task_id or "").strip() + if task_id: + task_ids.append(task_id) + _batch_revoke_celery_task_ids(task_ids, terminate=False) + + if not is_sharding_enabled(): + for eval_row, source_row in targets: + _prepare_source_row_for_retry( + source_row, + transcribe_overwrite=transcribe_overwrite, + ) + _reset_eval_row_for_retry( + eval_row, + metric_ids=metric_ids, + skip_revoke=True, + ) + catalog_db.commit() + return + + by_shard: dict[str, List[Tuple[UUID, UUID]]] = defaultdict(list) + for eval_row, source_row in targets: + shard_id = shard_id_for_row( + catalog_db, + evaluation.call_import_id, + int(source_row.row_index or 0), + ) + by_shard[shard_id].append((eval_row.id, source_row.id)) + + router = db_pool_manager.router + assert router is not None + for shard_id, id_pairs in by_shard.items(): + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + eval_row_ids = [pair[0] for pair in id_pairs] + source_row_ids = [pair[1] for pair in id_pairs] + eval_by_id = { + row.id: row + for row in shard_db.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.id.in_(eval_row_ids)) + .all() + } + source_by_id = { + row.id: row + for row in shard_db.query(CallImportRow) + .filter(CallImportRow.id.in_(source_row_ids)) + .all() + } + for eval_row_id, source_row_id in id_pairs: + bound_eval = eval_by_id.get(eval_row_id) + bound_source = source_by_id.get(source_row_id) + if bound_eval is None or bound_source is None: + continue + _prepare_source_row_for_retry( + bound_source, + transcribe_overwrite=transcribe_overwrite, + ) + _reset_eval_row_for_retry( + bound_eval, + metric_ids=metric_ids, + skip_revoke=True, + ) + if metric_ids: + flag_modified(bound_eval, "metric_scores") + shard_db.commit() + except Exception: + shard_db.rollback() + raise + finally: + shard_db.close() + + def execute_evaluation_retry( db: Session, evaluation_id: UUID, @@ -955,26 +1244,13 @@ def execute_evaluation_retry( for start in range(0, len(targets), _BULK_INSERT_CHUNK): chunk = targets[start : start + _BULK_INSERT_CHUNK] - task_ids: List[str] = [] - for eval_row, source_row in chunk: - if eval_row.celery_task_id and eval_row.status in { - "pending", - "running", - }: - task_id = (eval_row.celery_task_id or "").strip() - if task_id: - task_ids.append(task_id) - _prepare_source_row_for_retry( - source_row, - transcribe_overwrite=transcribe_overwrite, - ) - _reset_eval_row_for_retry( - eval_row, - metric_ids=metric_ids, - skip_revoke=True, - ) - _batch_revoke_celery_task_ids(task_ids, terminate=False) - db.commit() + _persist_evaluation_retry_targets( + db, + evaluation, + chunk, + metric_ids=metric_ids, + transcribe_overwrite=transcribe_overwrite, + ) try: evaluate_only, transcribe_chain = ( @@ -991,9 +1267,25 @@ def execute_evaluation_retry( "Failed to re-enqueue evaluation {} after retry reset", evaluation_id, ) - for eval_row, _ in targets: - eval_row.status = "failed" - eval_row.error_message = f"Failed to re-enqueue retry: {exc}" + err_msg = f"Failed to re-enqueue retry: {exc}" + target_ids = {eval_row.id for eval_row, _ in targets} + if is_sharding_enabled(): + from app.db_sharding.eval_rows import foreach_evaluation_row_mutating + + def _mark_failed(row: CallImportEvaluationRow) -> bool: + if row.id not in target_ids: + return False + row.status = "failed" + row.error_message = err_msg + return True + + foreach_evaluation_row_mutating(db, evaluation_id, _mark_failed) + else: + for eval_row, _ in targets: + if eval_row.id in target_ids: + eval_row.status = "failed" + eval_row.error_message = err_msg + db.commit() _rollup_evaluation_status(evaluation, db) db.commit() return { @@ -1002,6 +1294,8 @@ def execute_evaluation_retry( "error": str(exc), } + evaluation.error_message = None + evaluation.finished_at = None _rollup_evaluation_status(evaluation, db) db.commit() return { diff --git a/app/services/call_imports/progress_counters.py b/app/services/call_imports/progress_counters.py new file mode 100644 index 00000000..e50c29fe --- /dev/null +++ b/app/services/call_imports/progress_counters.py @@ -0,0 +1,268 @@ +"""Redis-backed completion counters with debounced catalog flush.""" + +from __future__ import annotations + +from typing import Optional +from uuid import UUID + +import redis +from loguru import logger +from sqlalchemy.orm import Session + +from app.config import settings + +_redis: redis.Redis | None = None + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def _eval_completed_key(evaluation_id: UUID | str) -> str: + return f"eval:{evaluation_id}:completed" + + +def _eval_failed_key(evaluation_id: UUID | str) -> str: + return f"eval:{evaluation_id}:failed" + + +def _import_completed_key(call_import_id: UUID | str) -> str: + return f"import:{call_import_id}:completed" + + +def _import_failed_key(call_import_id: UUID | str) -> str: + return f"import:{call_import_id}:failed" + + +def record_eval_row_terminal( + evaluation_id: UUID | str, + *, + completed_delta: int = 0, + failed_delta: int = 0, +) -> None: + try: + client = _client() + if completed_delta: + client.hincrby("eval:progress", _eval_completed_key(evaluation_id), completed_delta) + if failed_delta: + client.hincrby("eval:progress", _eval_failed_key(evaluation_id), failed_delta) + except redis.RedisError as exc: + logger.debug("eval progress counter skipped: {}", exc) + + +def record_import_row_terminal( + call_import_id: UUID | str, + *, + completed_delta: int = 0, + failed_delta: int = 0, +) -> None: + try: + client = _client() + if completed_delta: + client.hincrby( + "import:progress", + _import_completed_key(call_import_id), + completed_delta, + ) + if failed_delta: + client.hincrby( + "import:progress", + _import_failed_key(call_import_id), + failed_delta, + ) + except redis.RedisError as exc: + logger.debug("import progress counter skipped: {}", exc) + + +def read_eval_progress(evaluation_id: UUID | str) -> tuple[int, int]: + try: + client = _client() + completed = int(client.hget("eval:progress", _eval_completed_key(evaluation_id)) or 0) + failed = int(client.hget("eval:progress", _eval_failed_key(evaluation_id)) or 0) + return completed, failed + except redis.RedisError: + return 0, 0 + + +def flush_eval_progress_to_catalog(db: Session, evaluation_id: UUID) -> None: + """Merge Redis deltas into catalog parent columns (best-effort).""" + from app.models.database import CallImportEvaluation + + completed_delta, failed_delta = read_eval_progress(evaluation_id) + if not completed_delta and not failed_delta: + return + + # Decrement Redis before the catalog write. If Redis fails, skip the flush + # entirely so merge_eval_counters_for_ui never double-counts stale deltas. + try: + client = _client() + if completed_delta: + client.hincrby( + "eval:progress", + _eval_completed_key(evaluation_id), + -completed_delta, + ) + if failed_delta: + client.hincrby( + "eval:progress", + _eval_failed_key(evaluation_id), + -failed_delta, + ) + except redis.RedisError as exc: + logger.warning("eval progress flush skipped (redis decrement failed): {}", exc) + return + + try: + evaluation = ( + db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == evaluation_id) + .first() + ) + if evaluation is None: + raise LookupError(f"evaluation {evaluation_id} not found") + evaluation.completed_rows = int(evaluation.completed_rows or 0) + completed_delta + evaluation.failed_rows = int(evaluation.failed_rows or 0) + failed_delta + db.flush() + except Exception as exc: + logger.warning( + "eval progress catalog flush failed, restoring redis deltas: {}", + exc, + ) + record_eval_row_terminal( + evaluation_id, + completed_delta=completed_delta, + failed_delta=failed_delta, + ) + + +def clear_import_progress_redis(call_import_id: UUID | str) -> None: + """Drop stale Redis deltas after catalog counters were reconciled from row data.""" + try: + client = _client() + client.hdel( + "import:progress", + _import_completed_key(call_import_id), + _import_failed_key(call_import_id), + ) + except redis.RedisError as exc: + logger.debug("import progress redis clear skipped: {}", exc) + + +def clear_eval_progress_redis(evaluation_id: UUID | str) -> None: + """Drop stale Redis deltas after catalog counters were reconciled from row data.""" + try: + client = _client() + client.hdel( + "eval:progress", + _eval_completed_key(evaluation_id), + _eval_failed_key(evaluation_id), + ) + except redis.RedisError as exc: + logger.debug("eval progress redis clear skipped: {}", exc) + + +def merge_eval_counters_for_ui( + evaluation, +) -> tuple[int, int]: + """Catalog counters plus unflushed Redis deltas for progress display.""" + completed = int(getattr(evaluation, "completed_rows", 0) or 0) + failed = int(getattr(evaluation, "failed_rows", 0) or 0) + rc, rf = read_eval_progress(evaluation.id) + return completed + rc, failed + rf + + +def read_import_progress(call_import_id: UUID | str) -> tuple[int, int]: + try: + client = _client() + completed = int( + client.hget("import:progress", _import_completed_key(call_import_id)) or 0 + ) + failed = int( + client.hget("import:progress", _import_failed_key(call_import_id)) or 0 + ) + return completed, failed + except redis.RedisError: + return 0, 0 + + +def merge_import_counters_for_ui(call_import) -> tuple[int, int]: + completed = int(getattr(call_import, "completed_rows", 0) or 0) + failed = int(getattr(call_import, "failed_rows", 0) or 0) + rc, rf = read_import_progress(call_import.id) + return completed + rc, failed + rf + + +def record_import_row_status_transition( + call_import_id: UUID | str, + *, + previous_status: str, + new_status: str, +) -> None: + from app.workers.tasks.evaluate_call_import_row_core import ( + counter_deltas_for_status_transition, + ) + + def _norm(value) -> str: + raw = getattr(value, "value", value) + return str(raw or "").strip().lower() + + completed_delta, failed_delta = counter_deltas_for_status_transition( + _norm(previous_status), + _norm(new_status), + ) + if completed_delta or failed_delta: + record_import_row_terminal( + call_import_id, + completed_delta=completed_delta, + failed_delta=failed_delta, + ) + + +def flush_import_progress_to_catalog(db: Session, call_import_id: UUID) -> None: + """Merge Redis import deltas into catalog parent (best-effort).""" + from app.models.database import CallImport + + completed_delta, failed_delta = read_import_progress(call_import_id) + if not completed_delta and not failed_delta: + return + + try: + client = _client() + if completed_delta: + client.hincrby( + "import:progress", + _import_completed_key(call_import_id), + -completed_delta, + ) + if failed_delta: + client.hincrby( + "import:progress", + _import_failed_key(call_import_id), + -failed_delta, + ) + except redis.RedisError as exc: + logger.warning("import progress flush skipped (redis decrement failed): {}", exc) + return + + try: + call_import = ( + db.query(CallImport).filter(CallImport.id == call_import_id).first() + ) + if call_import is None: + raise LookupError(f"call_import {call_import_id} not found") + call_import.completed_rows = int(call_import.completed_rows or 0) + completed_delta + call_import.failed_rows = int(call_import.failed_rows or 0) + failed_delta + db.flush() + except Exception as exc: + logger.warning( + "import progress catalog flush failed, restoring redis deltas: {}", + exc, + ) + record_import_row_terminal( + call_import_id, + completed_delta=completed_delta, + failed_delta=failed_delta, + ) diff --git a/app/services/telephony/exotel_client.py b/app/services/telephony/exotel_client.py index 9d5f42f3..8ba30879 100644 --- a/app/services/telephony/exotel_client.py +++ b/app/services/telephony/exotel_client.py @@ -163,8 +163,7 @@ def get_call_recording_url(self, call_sid: str) -> str: f"Exotel server error fetching call detail (HTTP {resp.status_code})" ) if resp.status_code == 400: - self._penalize_if_fingerprinted() - raise CredentialedRecordingThrottledError( + raise ExotelTransientError( f"Unexpected HTTP 400 fetching call detail: {resp.text[:200]}" ) if resp.status_code >= 400: diff --git a/app/services/telephony/recording_download.py b/app/services/telephony/recording_download.py index 322db1a6..e4d40ec8 100644 --- a/app/services/telephony/recording_download.py +++ b/app/services/telephony/recording_download.py @@ -219,8 +219,8 @@ def _validate_redirect(request: httpx.Request) -> None: f"Rate limited fetching recording (HTTP 429)", retry_after_seconds=retry_after, ) - raise ExotelInvalidContentError( - f"Unexpected HTTP 429 fetching recording: {resp.text[:200]}" + raise ExotelTransientError( + f"Rate limited fetching recording (HTTP 429): {resp.text[:200]}" ) if 500 <= resp.status_code < 600: raise ExotelTransientError( @@ -228,13 +228,11 @@ def _validate_redirect(request: httpx.Request) -> None: ) if resp.status_code == 400: if auth is not None: - retry_after = _raise_credentialed_throttle( - fingerprint=credential_fingerprint, - message=f"Unexpected HTTP 400 fetching recording: {resp.text[:200]}", - ) - raise CredentialedRecordingThrottledError( - f"Unexpected HTTP 400 fetching recording: {resp.text[:200]}", - retry_after_seconds=retry_after, + # Exotel occasionally returns transient 400s on valid credentialed fetches; + # retry the row but do not penalize the shared credential (that blocks + # every other worker for 60s). + raise ExotelTransientError( + f"Unexpected HTTP 400 fetching recording: {resp.text[:200]}" ) raise ExotelInvalidContentError( f"Unexpected HTTP 400 fetching recording: {resp.text[:200]}" diff --git a/app/workers/catalog_cache.py b/app/workers/catalog_cache.py new file mode 100644 index 00000000..41b84ccb --- /dev/null +++ b/app/workers/catalog_cache.py @@ -0,0 +1,30 @@ +"""Short-lived in-process cache for catalog reads in bulk workers.""" + +from __future__ import annotations + +import time +from typing import Any, Callable, Dict, Tuple + +_TTL_SECONDS = 120.0 +_cache: Dict[Tuple[str, str], Tuple[float, Any]] = {} + + +def cached_catalog_fetch( + namespace: str, + key: str, + loader: Callable[[], Any], + *, + ttl_seconds: float = _TTL_SECONDS, +) -> Any: + cache_key = (namespace, key) + now = time.monotonic() + hit = _cache.get(cache_key) + if hit is not None and now - hit[0] < ttl_seconds: + return hit[1] + value = loader() + _cache[cache_key] = (now, value) + return value + + +def clear_worker_catalog_cache() -> None: + _cache.clear() diff --git a/app/workers/concurrency/dispatch_reconcile.py b/app/workers/concurrency/dispatch_reconcile.py new file mode 100644 index 00000000..2c48a2f4 --- /dev/null +++ b/app/workers/concurrency/dispatch_reconcile.py @@ -0,0 +1,152 @@ +"""Best-effort cleanup of orphaned dispatch locks after worker restarts.""" + +from __future__ import annotations + +from typing import Optional, TYPE_CHECKING + +from loguru import logger +from sqlalchemy.orm import Session + +from app.models.database import CallImportEvaluationRow, CallImportRow +from app.models.enums import CallImportRowStatus + +if TYPE_CHECKING: + from app.db_sharding.session_cache import ShardSessionCache + +_RECONCILE_BATCH = 500 + + +def _celery_task_is_active(task_id: str) -> bool: + """Return True when a Celery task id is still queued or running.""" + cleaned = (task_id or "").strip() + if not cleaned: + return False + try: + from app.workers.celery_app import celery_app + + result = celery_app.AsyncResult(cleaned) + state = result.state + if state in {"STARTED", "RETRY"}: + return True + if state in {"SUCCESS", "FAILURE", "REVOKED"}: + return False + + inspect = celery_app.control.inspect(timeout=0.5) + if inspect is None: + return False + for method_name in ("active", "reserved", "scheduled"): + method = getattr(inspect, method_name, None) + if method is None: + continue + payload = method() or {} + for tasks in payload.values(): + for task in tasks: + if task.get("id") == cleaned: + return True + return False + except Exception as exc: + logger.debug("Celery task inspect failed for {}: {}", task_id, exc) + return True + + +def _reconcile_import_rows_on_session(db: Session) -> int: + cleared = 0 + rows = ( + db.query(CallImportRow) + .filter( + CallImportRow.celery_task_id.isnot(None), + CallImportRow.status.in_( + (CallImportRowStatus.PENDING, CallImportRowStatus.PROCESSING) + ), + ) + .limit(_RECONCILE_BATCH) + .all() + ) + for row in rows: + task_id = (row.celery_task_id or "").strip() + if not task_id or _celery_task_is_active(task_id): + continue + row.celery_task_id = None + if ( + row.status == CallImportRowStatus.PROCESSING + and not (row.recording_s3_key or "").strip() + ): + row.status = CallImportRowStatus.PENDING + cleared += 1 + if cleared: + db.commit() + return cleared + + +def reconcile_orphaned_import_dispatch_locks(catalog_db: Session) -> int: + """Clear stale import-row task ids so fair import dispatch can resume.""" + from app.db_sharding.sessions import is_sharding_enabled + + if not is_sharding_enabled(): + return _reconcile_import_rows_on_session(catalog_db) + + from app.db_sharding.pool_manager import db_pool_manager + + router = db_pool_manager.router + assert router is not None + total = 0 + for shard_id in router.shard_ids: + shard_db = db_pool_manager.shard_session_factory(shard_id)() + try: + total += _reconcile_import_rows_on_session(shard_db) + finally: + shard_db.close() + return total + + +def _reconcile_eval_rows_on_session(db: Session) -> int: + cleared = 0 + rows = ( + db.query(CallImportEvaluationRow) + .filter( + CallImportEvaluationRow.celery_task_id.isnot(None), + CallImportEvaluationRow.status == "pending", + ) + .limit(_RECONCILE_BATCH) + .all() + ) + for row in rows: + task_id = (row.celery_task_id or "").strip() + if not task_id or _celery_task_is_active(task_id): + continue + row.celery_task_id = None + cleared += 1 + if cleared: + db.commit() + return cleared + + +def reconcile_orphaned_eval_dispatch_locks( + catalog_db: Session, + *, + shard_cache: Optional["ShardSessionCache"] = None, +) -> int: + """Clear stale eval-row task ids so fair eval dispatch can resume.""" + from app.db_sharding.sessions import is_sharding_enabled + + if not is_sharding_enabled(): + return _reconcile_eval_rows_on_session(catalog_db) + + from app.db_sharding.pool_manager import db_pool_manager + + router = db_pool_manager.router + assert router is not None + total = 0 + for shard_id in router.shard_ids: + if shard_cache is not None: + shard_db = shard_cache.session_for(shard_id) + owns_session = False + else: + shard_db = db_pool_manager.shard_session_factory(shard_id)() + owns_session = True + try: + total += _reconcile_eval_rows_on_session(shard_db) + finally: + if owns_session: + shard_db.close() + return total diff --git a/app/workers/concurrency/eval_dispatch.py b/app/workers/concurrency/eval_dispatch.py index 83f40a46..7afbd735 100644 --- a/app/workers/concurrency/eval_dispatch.py +++ b/app/workers/concurrency/eval_dispatch.py @@ -10,7 +10,9 @@ from sqlalchemy.orm import Session from app.database import SessionLocal +from app.db_sharding.session_cache import ShardSessionCache from app.models.database import ( + CallImport, CallImportEvaluation, CallImportEvaluationRow, CallImportRow, @@ -26,9 +28,10 @@ DIARIZATION_QUEUE = "diarization" EVALUATIONS_QUEUE = "evaluations" AUDIO_METRICS_QUEUE = "audio-metrics" -# Lightweight scheduler tasks — must not sit behind ``imports`` / ``evaluations`` -# fan-out on the call-import worker (see docker-compose queue order). -DISPATCH_QUEUE = "celery" +# Lightweight scheduler tasks — runs on ``evaluations`` so sandbox/dev can use +# ``worker-imports`` (imports,diarization,evaluations) without a separate celery worker. +# Row scoring also uses ``evaluations``; dispatch tasks are short DB/Redis work. +DISPATCH_QUEUE = "evaluations" DispatchSingleRowResult = Literal[ "dispatched", "skip", "at_capacity", "credential_throttled" @@ -105,16 +108,52 @@ def _fail_eval_row_for_import( db: Session, eval_row: CallImportEvaluationRow, source_row: CallImportRow, + *, + catalog_db: Session | None = None, + evaluation: CallImportEvaluation | None = None, ) -> None: from datetime import datetime, timezone + from app.db_sharding.row_ops import commit_shard_row_session + from app.workers.tasks.evaluate_call_import_row_core import ( + commit_terminal_row_and_rollup, + ) + + previous_status = eval_row.status or "pending" eval_row.status = "failed" eval_row.error_message = ( source_row.error_message or "Recording fetch failed" ) eval_row.finished_at = datetime.now(timezone.utc) eval_row.celery_task_id = None - db.commit() + if evaluation is not None: + commit_terminal_row_and_rollup( + db, + evaluation, + eval_row, + previous_row_status=previous_status, + catalog_db=( + catalog_db if catalog_db is not None and catalog_db is not db else None + ), + ) + else: + commit_shard_row_session(db) + + +def source_row_import_blocks_eval(source_row: CallImportRow) -> bool: + """True when import failure prevents the eval pipeline from continuing.""" + if (source_row.recording_s3_key or "").strip(): + return False + return source_row.status == CallImportRowStatus.FAILED + + +def recover_eval_row_for_eval_chain(eval_row: CallImportEvaluationRow) -> None: + """Undo a premature eval-row failure so the eval chain can continue.""" + if eval_row.status != "failed": + return + eval_row.status = "pending" + eval_row.error_message = None + eval_row.finished_at = None def build_eval_chain_transcribe_apply_async( @@ -173,22 +212,27 @@ def enqueue_eval_chain_transcribe_after_import( transcribe_overwrite: bool = False, ) -> bool: """Directly chain diarisation after a successful eval-chain recording fetch.""" - source_row.diarised_transcript_status = "pending" - source_row.diarised_transcript_error = None - source_row.celery_task_id = slot_task_id - eval_row.celery_task_id = None - db.flush() + from app.db_sharding.row_ops import shard_row_write_context - async_result = build_eval_chain_transcribe_apply_async( - evaluation=evaluation, - eval_row=eval_row, - source_row=source_row, - reserved_task_id=slot_task_id, - restricted_metric_ids=restricted_metric_ids, - transcribe_overwrite=transcribe_overwrite, - ) - eval_row.celery_task_id = async_result.id - db.commit() + recover_eval_row_for_eval_chain(eval_row) + + with shard_row_write_context(db): + source_row.diarised_transcript_status = "pending" + source_row.diarised_transcript_error = None + source_row.celery_task_id = slot_task_id + eval_row.celery_task_id = None + db.flush() + + async_result = build_eval_chain_transcribe_apply_async( + evaluation=evaluation, + eval_row=eval_row, + source_row=source_row, + reserved_task_id=slot_task_id, + restricted_metric_ids=restricted_metric_ids, + transcribe_overwrite=transcribe_overwrite, + ) + eval_row.celery_task_id = async_result.id + db.commit() return True @@ -209,19 +253,83 @@ def _reserve_slot_and_enqueue( ): return False + from app.db_sharding.row_ops import shard_row_write_context + try: - async_result = enqueue_fn(reserved_task_id) + with shard_row_write_context(db): + async_result = enqueue_fn(reserved_task_id) + eval_row.celery_task_id = async_result.id + db.commit() except Exception: release_eval_slot_for_celery_task(reserved_task_id) raise + return True + +def _attach_sharded_eval_dispatch_rows( + catalog_db: Session, + evaluation: CallImportEvaluation, + eval_row: CallImportEvaluationRow, + source_row: CallImportRow, + *, + shard_cache: ShardSessionCache | None = None, +) -> tuple[Session, CallImportEvaluationRow, CallImportRow, bool] | None: + """Bind eval/source rows on a shard session without extra catalog connections.""" + from app.db_sharding.pool_manager import db_pool_manager + from app.db_sharding.row_ops import shard_id_for_row + from app.db_sharding.sessions import is_sharding_enabled + + if not is_sharding_enabled(): + return catalog_db, eval_row, source_row, False + + shard_id = shard_id_for_row( + catalog_db, + evaluation.call_import_id, + int(source_row.row_index or 0), + ) + owns_session = shard_cache is None + if shard_cache is not None: + shard_db = shard_cache.session_for(shard_id) + else: + shard_db = db_pool_manager.shard_session_factory(shard_id)() try: - eval_row.celery_task_id = async_result.id - db.commit() + bound_eval = ( + shard_db.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.id == eval_row.id) + .first() + ) + bound_source = ( + shard_db.query(CallImportRow) + .filter(CallImportRow.id == source_row.id) + .first() + ) + if bound_eval is None or bound_source is None: + if owns_session: + shard_db.close() + return None + if (bound_eval.status or "") != "pending" or bound_eval.celery_task_id: + if owns_session: + shard_db.close() + return None + return shard_db, bound_eval, bound_source, owns_session except Exception: - release_eval_slot_for_celery_task(reserved_task_id) + if owns_session: + shard_db.close() raise - return True + + +def _load_call_import_for_eval_row( + catalog_db: Session, + source_row: CallImportRow, +) -> CallImport | None: + call_import_id = source_row.call_import_id + if call_import_id is None: + return None + return ( + catalog_db.query(CallImport) + .filter(CallImport.id == call_import_id) + .first() + ) def _try_dispatch_single_row( @@ -233,6 +341,8 @@ def _try_dispatch_single_row( restricted_metric_ids: Optional[List[str]] = None, transcribe_overwrite: bool = False, auto_transcribe: bool = True, + call_import: CallImport | None = None, + shard_cache: ShardSessionCache | None = None, ) -> EvalDispatchOutcome: """Dispatch one eval row (transcribe or evaluate). @@ -242,6 +352,7 @@ def _try_dispatch_single_row( * ``at_capacity`` — inflight cap reached; stop the workspace batch * ``credential_throttled`` — shared telephony credential is backing off """ + from app.db_sharding.sessions import is_sharding_enabled from app.workers.tasks.evaluate_call_import_row import ( evaluate_call_import_row_task, ) @@ -255,147 +366,180 @@ def _try_dispatch_single_row( process_call_import_row_task, ) - if (evaluation.status or "").strip().lower() == "cancelled": + attached = _attach_sharded_eval_dispatch_rows( + db, + evaluation, + eval_row, + source_row, + shard_cache=shard_cache, + ) + if attached is None: return EvalDispatchOutcome("skip") - if source_row.status == CallImportRowStatus.FAILED: - _fail_eval_row_for_import(db, eval_row, source_row) - return EvalDispatchOutcome("skip") + mutate_db, eval_row, source_row, owns_shard_session = attached + catalog_db = db - if source_row.status == CallImportRowStatus.PROCESSING: - if not (source_row.recording_s3_key or "").strip(): + try: + if (evaluation.status or "").strip().lower() == "cancelled": return EvalDispatchOutcome("skip") - if _needs_import_for_eval(source_row): - from app.workers.concurrency.import_dispatch import ( - _peek_authenticated_import_credit, + from app.services.call_imports.evaluation_bulk_op import ( + get_evaluation_bulk_operation, ) - call_import = source_row.call_import - throttled = _peek_authenticated_import_credit( - db=db, - call_import=call_import, - ) - if throttled is not None: - return EvalDispatchOutcome( - "credential_throttled", - wait_seconds=throttled.wait_seconds, - ) + if get_evaluation_bulk_operation(evaluation.id): + return EvalDispatchOutcome("skip") - def _enqueue_import(reserved_task_id: str): - source_row.celery_task_id = reserved_task_id - db.flush() - return process_call_import_row_task.apply_async( - args=(str(source_row.id),), - kwargs={ - "_eval_slot_task_id": reserved_task_id, - "run_eval_row_id": str(eval_row.id), - }, - queue=IMPORTS_QUEUE, - task_id=reserved_task_id, + if source_row.status == CallImportRowStatus.FAILED: + if source_row_import_blocks_eval(source_row): + _fail_eval_row_for_import( + mutate_db, + eval_row, + source_row, + catalog_db=catalog_db, + evaluation=evaluation, + ) + return EvalDispatchOutcome("skip") + + if source_row.status == CallImportRowStatus.PROCESSING: + if not (source_row.recording_s3_key or "").strip(): + return EvalDispatchOutcome("skip") + + if _needs_import_for_eval(source_row): + from app.workers.concurrency.import_dispatch import ( + _peek_authenticated_import_credit, ) - if _reserve_slot_and_enqueue( - evaluation=evaluation, - eval_row=eval_row, - db=db, - enqueue_fn=_enqueue_import, - ): - return EvalDispatchOutcome("dispatched") - return EvalDispatchOutcome("at_capacity") - - transcribe_mode = ( - getattr(evaluation, "transcribe_mode", None) or "stt_llm" - ).strip().lower() - - if _needs_transcribe_for_eval( - evaluation, - source_row, - transcribe_overwrite=transcribe_overwrite, - auto_transcribe=auto_transcribe, - ): - if _diarisation_in_flight(source_row): - return EvalDispatchOutcome("skip") - - dia_status = (source_row.diarised_transcript_status or "").strip().lower() - if ( - dia_status == "failed" - and not transcribe_overwrite - and not (source_row.diarised_transcript or "").strip() + call_import = call_import or _load_call_import_for_eval_row( + catalog_db, source_row + ) + if call_import is None: + return EvalDispatchOutcome("skip") + throttled = _peek_authenticated_import_credit( + db=catalog_db, + call_import=call_import, + ) + if throttled is not None: + return EvalDispatchOutcome( + "credential_throttled", + wait_seconds=throttled.wait_seconds, + ) + + def _enqueue_import(reserved_task_id: str): + source_row.celery_task_id = reserved_task_id + mutate_db.flush() + return process_call_import_row_task.apply_async( + args=(str(source_row.id),), + kwargs={ + "_eval_slot_task_id": reserved_task_id, + "run_eval_row_id": str(eval_row.id), + }, + queue=IMPORTS_QUEUE, + task_id=reserved_task_id, + ) + + if _reserve_slot_and_enqueue( + evaluation=evaluation, + eval_row=eval_row, + db=mutate_db, + enqueue_fn=_enqueue_import, + ): + return EvalDispatchOutcome("dispatched") + return EvalDispatchOutcome("at_capacity") + + if _needs_transcribe_for_eval( + evaluation, + source_row, + transcribe_overwrite=transcribe_overwrite, + auto_transcribe=auto_transcribe, ): - return EvalDispatchOutcome("skip") - - def _enqueue_transcribe(reserved_task_id: str): - source_row.diarised_transcript_status = "pending" - source_row.diarised_transcript_error = None - if transcribe_overwrite: - source_row.diarised_transcript = None - source_row.celery_task_id = reserved_task_id - db.flush() - return build_eval_chain_transcribe_apply_async( + if _diarisation_in_flight(source_row): + return EvalDispatchOutcome("skip") + + dia_status = ( + source_row.diarised_transcript_status or "" + ).strip().lower() + if ( + dia_status == "failed" + and not transcribe_overwrite + and not (source_row.diarised_transcript or "").strip() + ): + return EvalDispatchOutcome("skip") + + def _enqueue_transcribe(reserved_task_id: str): + source_row.diarised_transcript_status = "pending" + source_row.diarised_transcript_error = None + if transcribe_overwrite: + source_row.diarised_transcript = None + source_row.celery_task_id = reserved_task_id + mutate_db.flush() + return build_eval_chain_transcribe_apply_async( + evaluation=evaluation, + eval_row=eval_row, + source_row=source_row, + reserved_task_id=reserved_task_id, + restricted_metric_ids=restricted_metric_ids, + transcribe_overwrite=transcribe_overwrite, + ) + + if _reserve_slot_and_enqueue( evaluation=evaluation, eval_row=eval_row, - source_row=source_row, - reserved_task_id=reserved_task_id, - restricted_metric_ids=restricted_metric_ids, - transcribe_overwrite=transcribe_overwrite, - ) - - if _reserve_slot_and_enqueue( - evaluation=evaluation, - eval_row=eval_row, - db=db, - enqueue_fn=_enqueue_transcribe, + db=mutate_db, + enqueue_fn=_enqueue_transcribe, + ): + return EvalDispatchOutcome("dispatched") + return EvalDispatchOutcome("at_capacity") + + if row_needs_audio_phase( + catalog_db, + evaluation, + source_row, + restricted_metric_ids=restricted_metric_ids, ): - return EvalDispatchOutcome("dispatched") - return EvalDispatchOutcome("at_capacity") - if row_needs_audio_phase( - db, - evaluation, - source_row, - restricted_metric_ids=restricted_metric_ids, - ): + def _enqueue_audio(reserved_task_id: str): + kwargs = {"_eval_slot_task_id": reserved_task_id} + if restricted_metric_ids: + kwargs["restricted_metric_ids"] = restricted_metric_ids + return evaluate_call_import_row_audio_task.apply_async( + args=(str(eval_row.id),), + kwargs=kwargs, + queue=AUDIO_METRICS_QUEUE, + task_id=reserved_task_id, + ) + + if _reserve_slot_and_enqueue( + evaluation=evaluation, + eval_row=eval_row, + db=mutate_db, + enqueue_fn=_enqueue_audio, + ): + return EvalDispatchOutcome("dispatched") + return EvalDispatchOutcome("at_capacity") - def _enqueue_audio(reserved_task_id: str): + def _enqueue_eval(reserved_task_id: str): kwargs = {"_eval_slot_task_id": reserved_task_id} if restricted_metric_ids: kwargs["restricted_metric_ids"] = restricted_metric_ids - return evaluate_call_import_row_audio_task.apply_async( + return evaluate_call_import_row_task.apply_async( args=(str(eval_row.id),), kwargs=kwargs, - queue=AUDIO_METRICS_QUEUE, + queue=EVALUATIONS_QUEUE, task_id=reserved_task_id, ) if _reserve_slot_and_enqueue( evaluation=evaluation, eval_row=eval_row, - db=db, - enqueue_fn=_enqueue_audio, + db=mutate_db, + enqueue_fn=_enqueue_eval, ): return EvalDispatchOutcome("dispatched") return EvalDispatchOutcome("at_capacity") - - def _enqueue_eval(reserved_task_id: str): - kwargs = {"_eval_slot_task_id": reserved_task_id} - if restricted_metric_ids: - kwargs["restricted_metric_ids"] = restricted_metric_ids - return evaluate_call_import_row_task.apply_async( - args=(str(eval_row.id),), - kwargs=kwargs, - queue=EVALUATIONS_QUEUE, - task_id=reserved_task_id, - ) - - if _reserve_slot_and_enqueue( - evaluation=evaluation, - eval_row=eval_row, - db=db, - enqueue_fn=_enqueue_eval, - ): - return EvalDispatchOutcome("dispatched") - return EvalDispatchOutcome("at_capacity") + finally: + if owns_shard_session and is_sharding_enabled() and mutate_db is not db: + mutate_db.close() @celery_app.task(name="dispatch_evaluation_rows", queue=EVALUATIONS_QUEUE) diff --git a/app/workers/concurrency/fair_diarization_dispatch.py b/app/workers/concurrency/fair_diarization_dispatch.py index 91659af4..2499ce9c 100644 --- a/app/workers/concurrency/fair_diarization_dispatch.py +++ b/app/workers/concurrency/fair_diarization_dispatch.py @@ -20,6 +20,12 @@ ) from app.workers.concurrency.eval_dispatch import DIARIZATION_QUEUE from app.workers.config import celery_app +from app.db_sharding.import_dispatch import ( + call_imports_with_pending_diarization as sharded_call_imports_diarize, + pending_diarization_row_for_call_import as sharded_pending_diarize_row, + pending_diarization_workspaces as sharded_pending_diarize_workspaces, +) +from app.db_sharding.sessions import is_sharding_enabled _RR_CURSOR_KEY = "diarisation:fair:rr_cursor" _WS_CALL_IMPORT_RR_CURSOR_KEY_PREFIX = "diarisation:fair:rr_cursor:ws:" @@ -76,6 +82,8 @@ def _set_workspace_call_import_rr_cursor(workspace_id: UUID, cursor: int) -> Non def _workspaces_with_pending_diarization(db: Session) -> List[UUID]: + if is_sharding_enabled(): + return sharded_pending_diarize_workspaces(db) rows = ( db.query(CallImport.workspace_id) .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) @@ -93,6 +101,8 @@ def _call_imports_with_pending_diarization( db: Session, workspace_id: UUID, ) -> List[UUID]: + if is_sharding_enabled(): + return sharded_call_imports_diarize(db, workspace_id) rows = ( db.query(CallImport.id) .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) @@ -111,6 +121,9 @@ def _pending_row_for_call_import( db: Session, call_import_id: UUID, ) -> tuple[CallImportRow, CallImport] | None: + if is_sharding_enabled(): + pending = sharded_pending_diarize_row(db, call_import_id) + return pending row = ( db.query(CallImportRow, CallImport) .join(CallImport, CallImport.id == CallImportRow.call_import_id) diff --git a/app/workers/concurrency/fair_dispatch.py b/app/workers/concurrency/fair_dispatch.py index 380e8b7d..913734ea 100644 --- a/app/workers/concurrency/fair_dispatch.py +++ b/app/workers/concurrency/fair_dispatch.py @@ -13,6 +13,7 @@ from app.config import settings from app.database import SessionLocal from app.models.database import ( + CallImport, CallImportEvaluation, CallImportEvaluationRow, CallImportRow, @@ -23,8 +24,17 @@ _try_dispatch_single_row, ) from app.workers.config import celery_app +from app.db_sharding.scatter_gather import ( + evaluations_with_pending_rows as sharded_evaluations_with_pending_rows, + pending_eval_row_triples as sharded_pending_eval_row_triples, + pending_eval_workspaces as sharded_pending_eval_workspaces, +) +from app.db_sharding.sessions import is_sharding_enabled +from app.db_sharding.session_cache import ShardSessionCache _RR_CURSOR_KEY = "eval:fair:rr_cursor" +_DISPATCH_LOCK_KEY = "eval:fair:dispatch_lock" +_DISPATCH_LOCK_TTL_SECONDS = 90 _WS_EVAL_RR_CURSOR_KEY_PREFIX = "eval:fair:rr_cursor:ws:" _RESTRICTED_ROW_KEY_PREFIX = "eval:restricted:row:" _TRANSCRIBE_OVERWRITE_KEY_PREFIX = "eval:transcribe_overwrite:" @@ -171,7 +181,35 @@ def _set_workspace_eval_rr_cursor(workspace_id: UUID, cursor: int) -> None: ) -def _workspaces_with_pending_rows(db: Session) -> List[UUID]: +def _try_acquire_dispatch_lock() -> bool: + try: + return bool( + _get_redis().set( + _DISPATCH_LOCK_KEY, + "1", + nx=True, + ex=_DISPATCH_LOCK_TTL_SECONDS, + ) + ) + except redis.RedisError as exc: + logger.warning("Fair dispatch lock acquire failed: {}", exc) + return True + + +def _release_dispatch_lock() -> None: + try: + _get_redis().delete(_DISPATCH_LOCK_KEY) + except redis.RedisError as exc: + logger.warning("Fair dispatch lock release failed: {}", exc) + + +def _workspaces_with_pending_rows( + db: Session, + *, + shard_cache: ShardSessionCache | None = None, +) -> List[UUID]: + if is_sharding_enabled(): + return sharded_pending_eval_workspaces(db, shard_cache=shard_cache) rows = ( db.query(CallImportEvaluation.workspace_id) .join( @@ -196,23 +234,39 @@ def _workspaces_with_pending_rows(db: Session) -> List[UUID]: def _evaluations_with_pending_rows( db: Session, workspace_id: UUID, + *, + shard_cache: ShardSessionCache | None = None, ) -> List[UUID]: - rows = ( - db.query(CallImportEvaluation.id) - .join( - CallImportEvaluationRow, - CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + from app.services.call_imports.evaluation_bulk_op import ( + get_evaluation_bulk_operation, + ) + + if is_sharding_enabled(): + evaluation_ids = sharded_evaluations_with_pending_rows( + db, workspace_id, shard_cache=shard_cache ) - .filter( - CallImportEvaluation.workspace_id == workspace_id, - CallImportEvaluation.status != "cancelled", - CallImportEvaluationRow.status == "pending", - CallImportEvaluationRow.celery_task_id.is_(None), + else: + rows = ( + db.query(CallImportEvaluation.id) + .join( + CallImportEvaluationRow, + CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + ) + .filter( + CallImportEvaluation.workspace_id == workspace_id, + CallImportEvaluation.status != "cancelled", + CallImportEvaluationRow.status == "pending", + CallImportEvaluationRow.celery_task_id.is_(None), + ) + .distinct() + .all() ) - .distinct() - .all() - ) - return sorted({row[0] for row in rows if row[0] is not None}) + evaluation_ids = sorted({row[0] for row in rows if row[0] is not None}) + return [ + evaluation_id + for evaluation_id in evaluation_ids + if not get_evaluation_bulk_operation(evaluation_id) + ] def _pending_rows_for_evaluation( @@ -220,7 +274,15 @@ def _pending_rows_for_evaluation( evaluation_id: UUID, *, limit: int, + shard_cache: ShardSessionCache | None = None, ) -> List[tuple[CallImportEvaluationRow, CallImportRow, CallImportEvaluation]]: + if is_sharding_enabled(): + return sharded_pending_eval_row_triples( + db, + evaluation_id, + limit=limit, + shard_cache=shard_cache, + ) return ( db.query(CallImportEvaluationRow, CallImportRow, CallImportEvaluation) .join( @@ -248,12 +310,15 @@ def _dispatch_batch_for_workspace( workspace_id: UUID, *, batch_size: int, + shard_cache: ShardSessionCache | None = None, ) -> tuple[int, bool, int]: """Dispatch up to ``batch_size`` pending rows for one workspace turn. Returns ``(dispatched_count, hit_capacity, backoff_seconds)``. """ - evaluations = _evaluations_with_pending_rows(db, workspace_id) + evaluations = _evaluations_with_pending_rows( + db, workspace_id, shard_cache=shard_cache + ) if not evaluations: return 0, False, 0 @@ -269,12 +334,20 @@ def _dispatch_batch_for_workspace( db, evaluation_id, limit=max(1, batch_size - dispatched), + shard_cache=shard_cache, ) if not pending: skips += 1 cursor = (cursor + 1) % len(evaluations) continue + evaluation_header = pending[0][2] + call_import = ( + db.query(CallImport) + .filter(CallImport.id == evaluation_header.call_import_id) + .first() + ) + evaluation_dispatched = False for eval_row, source_row, evaluation in pending: restricted_metric_ids = get_row_restricted_metrics(eval_row.id) @@ -287,6 +360,8 @@ def _dispatch_batch_for_workspace( restricted_metric_ids=restricted_metric_ids, transcribe_overwrite=transcribe_overwrite, auto_transcribe=True, + call_import=call_import, + shard_cache=shard_cache, ) if outcome.result == "dispatched": clear_row_restricted_metrics(eval_row.id) @@ -373,11 +448,28 @@ def schedule_fair_dispatch( @celery_app.task(name="dispatch_fair_eval_rows", queue=DISPATCH_QUEUE) def dispatch_fair_eval_rows_task(max_workspace_turns: int = 1) -> dict: """Round-robin pending eval rows across workspaces (batch K per turn).""" + if not _try_acquire_dispatch_lock(): + schedule_fair_dispatch(max_workspace_turns=1, countdown=5) + return {"status": "deferred", "reason": "dispatch_lock_held"} + db = SessionLocal() + shard_cache = ShardSessionCache() if is_sharding_enabled() else None total_dispatched = 0 hit_capacity = False try: - workspaces = _workspaces_with_pending_rows(db) + from app.workers.concurrency.dispatch_reconcile import ( + reconcile_orphaned_eval_dispatch_locks, + ) + + reconciled = reconcile_orphaned_eval_dispatch_locks( + db, shard_cache=shard_cache + ) + if reconciled: + logger.info( + "Reconciled {} orphaned eval-row dispatch lock(s) after restart", + reconciled, + ) + workspaces = _workspaces_with_pending_rows(db, shard_cache=shard_cache) if not workspaces: return {"status": "ok", "dispatched": 0, "workspaces": 0} @@ -396,6 +488,7 @@ def dispatch_fair_eval_rows_task(max_workspace_turns: int = 1) -> dict: db, workspace_id, batch_size=batch_size, + shard_cache=shard_cache, ) ) hit_capacity = hit_capacity or workspace_at_capacity @@ -411,7 +504,9 @@ def dispatch_fair_eval_rows_task(max_workspace_turns: int = 1) -> dict: skips += 1 cursor = (cursor + 1) % len(workspaces) - if hit_capacity and _workspaces_with_pending_rows(db): + if hit_capacity and _workspaces_with_pending_rows( + db, shard_cache=shard_cache + ): _schedule_dispatch_deduped( max_workspace_turns=1, countdown=backoff_seconds or _DISPATCH_AT_CAPACITY_BACKOFF_SECONDS, @@ -429,7 +524,10 @@ def dispatch_fair_eval_rows_task(max_workspace_turns: int = 1) -> dict: logger.exception("dispatch_fair_eval_rows failed") raise finally: + if shard_cache is not None: + shard_cache.close_all() db.close() + _release_dispatch_lock() def finish_eval_work_and_redispatch( diff --git a/app/workers/concurrency/fair_import_dispatch.py b/app/workers/concurrency/fair_import_dispatch.py index 868d2150..04d8f804 100644 --- a/app/workers/concurrency/fair_import_dispatch.py +++ b/app/workers/concurrency/fair_import_dispatch.py @@ -16,6 +16,12 @@ from app.workers.concurrency.eval_dispatch import IMPORTS_QUEUE from app.workers.concurrency.import_dispatch import _try_dispatch_single_import_row from app.workers.config import celery_app +from app.db_sharding.import_dispatch import ( + call_imports_with_pending_rows as sharded_call_imports_with_pending_rows, + pending_import_row_for_call_import as sharded_pending_import_row, + pending_import_workspaces as sharded_pending_import_workspaces, +) +from app.db_sharding.sessions import is_sharding_enabled _RR_CURSOR_KEY = "import:fair:rr_cursor" _WS_CALL_IMPORT_RR_CURSOR_KEY_PREFIX = "import:fair:rr_cursor:ws:" @@ -74,6 +80,8 @@ def _set_workspace_call_import_rr_cursor(workspace_id: UUID, cursor: int) -> Non def _workspaces_with_pending_imports(db: Session) -> List[UUID]: + if is_sharding_enabled(): + return sharded_pending_import_workspaces(db) rows = ( db.query(CallImport.workspace_id) .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) @@ -92,6 +100,8 @@ def _call_imports_with_pending_rows( db: Session, workspace_id: UUID, ) -> List[UUID]: + if is_sharding_enabled(): + return sharded_call_imports_with_pending_rows(db, workspace_id) rows = ( db.query(CallImport.id) .join(CallImportRow, CallImportRow.call_import_id == CallImport.id) @@ -111,6 +121,9 @@ def _pending_row_for_call_import( db: Session, call_import_id: UUID, ) -> tuple[CallImportRow, CallImport] | None: + if is_sharding_enabled(): + pending = sharded_pending_import_row(db, call_import_id) + return pending row = ( db.query(CallImportRow, CallImport) .join(CallImport, CallImport.id == CallImportRow.call_import_id) @@ -221,8 +234,18 @@ def schedule_fair_import_dispatch( @celery_app.task(name="dispatch_fair_import_rows", queue=IMPORTS_QUEUE) def dispatch_fair_import_rows_task(max_workspace_turns: int = 1) -> dict: """Round-robin pending import rows across workspaces (batch K per turn).""" + from app.workers.concurrency.dispatch_reconcile import ( + reconcile_orphaned_import_dispatch_locks, + ) + db = SessionLocal() try: + reconciled = reconcile_orphaned_import_dispatch_locks(db) + if reconciled: + logger.info( + "Reconciled {} orphaned import-row dispatch lock(s) after restart", + reconciled, + ) workspaces = _workspaces_with_pending_imports(db) if not workspaces: return {"status": "ok", "dispatched": 0, "workspaces": 0} diff --git a/app/workers/config.py b/app/workers/config.py index 347874bf..387c65a3 100644 --- a/app/workers/config.py +++ b/app/workers/config.py @@ -71,6 +71,10 @@ log_startup_status(component="celery-worker") +# Queues consumed by the dedicated call-import / evaluation worker. +IMPORTS_WORKER_QUEUES = "imports,diarization,eval-control,evaluations" +EVAL_CONTROL_QUEUE = "eval-control" + # Create Celery app celery_app = Celery( "efficientai", @@ -105,24 +109,26 @@ ) # Route call-import recording fetch to the imports queue (preferred by workers -# that consume ``imports,diarization,evaluations``). Diarisation work (manual -# bulk transcribe and eval-chain transcribe) goes to the diarization queue. -# Eval fair dispatch, LLM scoring, and post-eval LLM jobs go to the -# evaluations queue so large scoring fan-outs do not head-of-line block -# recording fetch for other workspaces. +# that consume ``imports,diarization,eval-control,evaluations``). Diarisation +# work (manual bulk transcribe and eval-chain transcribe) goes to the +# diarization queue. Bulk eval control (materialize, cancel, retry) goes to +# ``eval-control`` so operator actions are not head-of-line blocked by scoring. +# Eval fair dispatch and LLM scoring go to the evaluations queue. celery_app.conf.task_routes = { "process_call_import_row": {"queue": "imports"}, "bulk_diarize_call_import": {"queue": "imports"}, "bulk_delete_call_import_rows": {"queue": "imports"}, "materialize_call_import_rows": {"queue": "imports"}, "delete_call_import": {"queue": "imports"}, - "materialize_call_import_evaluation": {"queue": "celery"}, - "materialize_mapped_call_import_evaluation": {"queue": "celery"}, + "materialize_call_import_evaluation": {"queue": EVAL_CONTROL_QUEUE}, + "materialize_mapped_call_import_evaluation": {"queue": EVAL_CONTROL_QUEUE}, + "retry_call_import_evaluation": {"queue": EVAL_CONTROL_QUEUE}, + "cancel_call_import_evaluation": {"queue": EVAL_CONTROL_QUEUE}, "evaluate_call_import_row": {"queue": "evaluations"}, "evaluate_call_import_row_audio": {"queue": "audio-metrics"}, "transcribe_call_import_row": {"queue": "diarization"}, "dispatch_evaluation_rows": {"queue": "evaluations"}, - "dispatch_fair_eval_rows": {"queue": "celery"}, + "dispatch_fair_eval_rows": {"queue": "evaluations"}, "dispatch_fair_diarization_rows": {"queue": "diarization"}, "dispatch_fair_import_rows": {"queue": "imports"}, "generate_evaluation_tldr_insights": {"queue": "evaluations"}, diff --git a/app/workers/tasks/evaluate_call_import_row.py b/app/workers/tasks/evaluate_call_import_row.py index 2cea1d65..9ccb9dbb 100644 --- a/app/workers/tasks/evaluate_call_import_row.py +++ b/app/workers/tasks/evaluate_call_import_row.py @@ -17,6 +17,7 @@ from uuid import UUID from loguru import logger +from sqlalchemy.orm import Session from app.database import SessionLocal from app.models.database import ( @@ -59,6 +60,30 @@ _commit_terminal_row_and_rollup = commit_terminal_row_and_rollup +def _rollup_terminal( + row_db: Session, + catalog_db, + evaluation: CallImportEvaluation, + eval_row: CallImportEvaluationRow, + *, + previous_row_status: str, +) -> None: + cat = catalog_db if catalog_db is not row_db else None + _commit_terminal_row_and_rollup( + row_db, + evaluation, + eval_row, + previous_row_status=previous_row_status, + catalog_db=cat, + ) + + +def _persist_eval_sessions(row_db, catalog_db) -> None: + row_db.commit() + if catalog_db is not row_db: + catalog_db.commit() + + def _run_llm_scoring( *, eval_row_id: UUID, @@ -72,6 +97,7 @@ def _run_llm_scoring( ai_providers: list, llm_provider: str | None, llm_model: str | None, + llm_credential_id: str | None, llm_config: dict | None, metric_llm_overrides: dict, discover_new_metrics: bool, @@ -94,22 +120,26 @@ def _run_llm_scoring( else None ) + run_provider = (llm_provider or "").strip() or None + run_model = (llm_model or "").strip() or None + run_llm_config = llm_config if isinstance(llm_config, dict) else None + run_credential_id = (llm_credential_id or "").strip() or None + overrides = metric_llm_overrides if isinstance(metric_llm_overrides, dict) else {} + if comparison_metrics: - run_provider = (llm_provider or "").strip() or None - run_model = (llm_model or "").strip() or None - run_llm_config = llm_config if isinstance(llm_config, dict) else None - overrides = metric_llm_overrides if isinstance(metric_llm_overrides, dict) else {} for cmp_metric in comparison_metrics: override = overrides.get(str(cmp_metric.id)) or {} provider = override.get("provider") or run_provider or None model = override.get("model") or run_model or None llm_cfg = override.get("llm_config") or run_llm_config + credential_id = override.get("credential_id") or run_credential_id evaluator_obj = None if provider and model: evaluator_obj = SimpleNamespace( llm_provider=provider, llm_model=model, llm_config=llm_cfg, + llm_credential_id=credential_id, custom_prompt=None, ) try: @@ -160,36 +190,32 @@ def _llm_config_key(cfg: dict | None) -> str | None: def _resolve_pm( metric: Metric, - ) -> tuple[str | None, str | None, dict | None]: + ) -> tuple[str | None, str | None, dict | None, str | None]: override = overrides.get(str(metric.id)) or {} provider = override.get("provider") or run_provider or None model = override.get("model") or run_model or None llm_cfg = override.get("llm_config") or run_llm_config - return provider, model, llm_cfg - - run_provider = (llm_provider or "").strip() or None - run_model = (llm_model or "").strip() or None - run_llm_config = llm_config if isinstance(llm_config, dict) else None - overrides = metric_llm_overrides if isinstance(metric_llm_overrides, dict) else {} + credential_id = override.get("credential_id") or run_credential_id + return provider, model, llm_cfg, credential_id - BucketKey = tuple[tuple[str | None, str | None, str | None], UUID | None] + BucketKey = tuple[tuple[str | None, str | None, str | None, str | None], UUID | None] groups: dict[BucketKey, list[Metric]] = {} for metric in standalone_metrics: - provider, model, llm_cfg = _resolve_pm(metric) + provider, model, llm_cfg, credential_id = _resolve_pm(metric) groups.setdefault( - ((provider, model, _llm_config_key(llm_cfg)), None), + ((provider, model, _llm_config_key(llm_cfg), credential_id), None), [], ).append(metric) for parent_id, children in children_by_parent.items(): - provider, model, llm_cfg = _resolve_pm(children[0]) + provider, model, llm_cfg, credential_id = _resolve_pm(children[0]) groups.setdefault( - ((provider, model, _llm_config_key(llm_cfg)), parent_id), + ((provider, model, _llm_config_key(llm_cfg), credential_id), parent_id), [], ).extend(children) metric_discovery_emitted = False for (config, parent_id), bucket in groups.items(): - provider, model, llm_config_key = config + provider, model, llm_config_key, credential_id = config llm_cfg = json.loads(llm_config_key) if llm_config_key else None evaluator_obj = None if provider and model: @@ -197,6 +223,7 @@ def _resolve_pm( llm_provider=provider, llm_model=model, llm_config=llm_cfg, + llm_credential_id=credential_id, custom_prompt=None, ) parent_metric = parents_by_id.get(parent_id) if parent_id else None @@ -306,20 +333,24 @@ def evaluate_call_import_row_task( scoring_inputs: dict[str, Any] | None = None restricted_metric_uuids: list[UUID] | None = None try: - db = SessionLocal() + from app.db_sharding.row_ops import ( + close_row_sessions, + locate_call_import_evaluation_row, + ) + + row_db = catalog_db = None try: row_uuid = UUID(eval_row_id) - eval_row = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == row_uuid) - .first() - ) - if not eval_row: + try: + row_db, catalog_db, eval_row, source_row, _shard_id = ( + locate_call_import_evaluation_row(row_uuid) + ) + except LookupError: logger.warning("CallImportEvaluationRow {} not found", eval_row_id) return {"status": "skipped", "reason": "row_not_found"} evaluation = ( - db.query(CallImportEvaluation) + catalog_db.query(CallImportEvaluation) .filter(CallImportEvaluation.id == eval_row.evaluation_id) .first() ) @@ -329,26 +360,10 @@ def evaluate_call_import_row_task( ) eval_row.status = "failed" eval_row.error_message = "Evaluation parent not found" - db.commit() + row_db.commit() return {"status": "failed", "reason": "evaluation_missing"} previous_row_status = eval_row.status - source_row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == eval_row.call_import_row_id) - .first() - ) - if not source_row: - eval_row.status = "failed" - eval_row.error_message = "Source call import row not found" - eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, - evaluation, - eval_row, - previous_row_status=previous_row_status, - ) - return {"status": "failed", "reason": "source_row_missing"} eval_row.status = "running" eval_row.celery_task_id = self.request.id @@ -357,7 +372,7 @@ def evaluate_call_import_row_task( if evaluation.status == "pending": evaluation.status = "running" evaluation.started_at = evaluation.started_at or _now() - db.commit() + _persist_eval_sessions(row_db, catalog_db) previous_row_status = "running" production_transcript = (source_row.transcript or "").strip() @@ -383,8 +398,9 @@ def evaluate_call_import_row_task( eval_row.status = "completed" eval_row.error_message = None eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status=previous_row_status, @@ -395,7 +411,7 @@ def evaluate_call_import_row_task( } metrics = load_enabled_metrics( - db, evaluation, restricted_metric_ids=restricted_metric_ids + catalog_db, evaluation, restricted_metric_ids=restricted_metric_ids ) if not metrics: eval_row.status = "failed" @@ -403,8 +419,9 @@ def evaluate_call_import_row_task( "No enabled metrics selected for this evaluation" ) eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status=previous_row_status, @@ -417,6 +434,14 @@ def evaluate_call_import_row_task( else {} ) parent_import = getattr(source_row, "call_import", None) + if parent_import is None and source_row.call_import_id: + from app.models.database import CallImport + + parent_import = ( + catalog_db.query(CallImport) + .filter(CallImport.id == source_row.call_import_id) + .first() + ) custom_column_mapping = ( parent_import.custom_column_mapping if parent_import is not None @@ -472,8 +497,9 @@ def evaluate_call_import_row_task( ) eval_row.metric_scores = _as_json_dict(metric_scores) eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status=previous_row_status, @@ -503,8 +529,9 @@ def evaluate_call_import_row_task( ) eval_row.metric_scores = _as_json_dict(metric_scores) eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status=previous_row_status, @@ -526,7 +553,7 @@ def evaluate_call_import_row_task( transcript_metrics = [] ai_providers = ( - db.query(AIProvider) + catalog_db.query(AIProvider) .filter( AIProvider.organization_id == evaluation.organization_id, AIProvider.is_active.is_(True), @@ -541,7 +568,7 @@ def evaluate_call_import_row_task( running_discovered_by_parent: dict[UUID, list] = {} if transcript_metrics and transcript: parents_by_id, children_by_parent, standalone_metrics = ( - _build_parent_groups(db, transcript_metrics) + _build_parent_groups(catalog_db, transcript_metrics) ) if bool(getattr(evaluation, "discover_new_metrics", False)): from app.api.v1.routes.call_import_evaluations import ( @@ -554,7 +581,7 @@ def evaluate_call_import_row_task( else {} ) running_discovered_metrics = _get_running_discovered_metrics( - db, + catalog_db, evaluation.id, organization_id=evaluation.organization_id, alias_map=alias_map_metrics, @@ -577,7 +604,7 @@ def evaluate_call_import_row_task( continue running_discovered_by_parent[parent_id] = ( _get_running_discovered_labels( - db, + catalog_db, evaluation.id, parent_metric.id, organization_id=evaluation.organization_id, @@ -599,6 +626,11 @@ def evaluate_call_import_row_task( "ai_providers": ai_providers, "llm_provider": evaluation.llm_provider, "llm_model": evaluation.llm_model, + "llm_credential_id": ( + str(evaluation.llm_credential_id) + if evaluation.llm_credential_id + else None + ), "llm_config": evaluation.llm_config, "metric_llm_overrides": evaluation.metric_llm_overrides, "discover_new_metrics": bool( @@ -614,7 +646,8 @@ def evaluate_call_import_row_task( "pre_llm_metric_scores": dict(metric_scores), } finally: - db.close() + if row_db is not None: + close_row_sessions(row_db, catalog_db) pre_llm_metric_scores = scoring_inputs.pop("pre_llm_metric_scores") llm_result = _run_llm_scoring(**scoring_inputs) @@ -625,39 +658,47 @@ def evaluate_call_import_row_task( evaluation_failed = llm_result["evaluation_failed"] primary_error_message = llm_result["primary_error_message"] - db = SessionLocal() + row_db = catalog_db = None try: - eval_row = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == UUID(eval_row_id)) - .first() + from app.db_sharding.row_ops import ( + close_row_sessions, + locate_call_import_evaluation_row, ) - if not eval_row: + + try: + row_db, catalog_db, eval_row, _source_row, _ = ( + locate_call_import_evaluation_row(UUID(eval_row_id)) + ) + except LookupError: + logger.warning( + "[CallImportEval {}] Row not found on any shard — skipping", + eval_row_id, + ) return {"status": "skipped", "reason": "row_not_found"} evaluation = ( - db.query(CallImportEvaluation) + catalog_db.query(CallImportEvaluation) .filter(CallImportEvaluation.id == eval_row.evaluation_id) .first() ) if not evaluation: return {"status": "failed", "reason": "evaluation_missing"} - if _was_cancelled_externally(db, eval_row): + if _was_cancelled_externally(row_db, eval_row): logger.info( "[CallImportEval {}] Skipping terminal write; " "row was cancelled by user", eval_row.id, ) try: - rollup_parent( - db, + _rollup_parent( + catalog_db, evaluation, previous_row_status="running", new_row_status="failed", ) - db.commit() + catalog_db.commit() except Exception: # noqa: BLE001 — rollup is best-effort here - db.rollback() + catalog_db.rollback() return { "status": "cancelled", "eval_row_id": eval_row_id, @@ -678,7 +719,7 @@ def evaluate_call_import_row_task( ) normalize_scores_with_aliases( - metric_scores, evaluation, db, evaluation.organization_id + metric_scores, evaluation, catalog_db, evaluation.organization_id ) new_scores = _as_json_dict(metric_scores) @@ -695,8 +736,9 @@ def evaluate_call_import_row_task( else: eval_row.metric_scores = new_scores eval_row.finished_at = _now() - _commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status="running", @@ -709,7 +751,8 @@ def evaluate_call_import_row_task( "subset_retry": bool(restricted_metric_uuids), } finally: - db.close() + if row_db is not None: + close_row_sessions(row_db, catalog_db) finally: from app.workers.concurrency.fair_dispatch import ( finish_eval_work_and_redispatch, diff --git a/app/workers/tasks/evaluate_call_import_row_audio.py b/app/workers/tasks/evaluate_call_import_row_audio.py index 3281ac9f..fb746b59 100644 --- a/app/workers/tasks/evaluate_call_import_row_audio.py +++ b/app/workers/tasks/evaluate_call_import_row_audio.py @@ -11,12 +11,7 @@ from loguru import logger -from app.database import SessionLocal -from app.models.database import ( - CallImportEvaluation, - CallImportEvaluationRow, - CallImportRow, -) +from app.models.database import CallImportEvaluation from app.workers.config import celery_app from app.workers.concurrency.eval_dispatch import EVALUATIONS_QUEUE from app.workers.tasks.evaluate_call_import_row_core import ( @@ -32,6 +27,23 @@ ) +def _rollup_terminal(row_db, catalog_db, evaluation, eval_row, *, previous_row_status): + cat = catalog_db if catalog_db is not row_db else None + commit_terminal_row_and_rollup( + row_db, + evaluation, + eval_row, + previous_row_status=previous_row_status, + catalog_db=cat, + ) + + +def _persist_sessions(row_db, catalog_db) -> None: + row_db.commit() + if catalog_db is not row_db: + catalog_db.commit() + + @celery_app.task( name="evaluate_call_import_row_audio", bind=True, @@ -46,222 +58,211 @@ def evaluate_call_import_row_audio_task( _eval_slot_task_id: Optional[str] = None, ): """Score audio-only metrics; chain LLM phase or finalize when done.""" - db = SessionLocal() + from app.db_sharding.eval_rows import evaluation_row_session + slot_task_id = _eval_slot_task_id or self.request.id chain_llm = False try: - row_uuid = UUID(eval_row_id) - eval_row = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == row_uuid) - .first() - ) - if not eval_row: - logger.warning("CallImportEvaluationRow {} not found", eval_row_id) - return {"status": "skipped", "reason": "row_not_found"} - - evaluation = ( - db.query(CallImportEvaluation) - .filter(CallImportEvaluation.id == eval_row.evaluation_id) - .first() - ) - if not evaluation: - eval_row.status = "failed" - eval_row.error_message = "Evaluation parent not found" - db.commit() - return {"status": "failed", "reason": "evaluation_missing"} + with evaluation_row_session(eval_row_id) as ( + row_db, + catalog_db, + eval_row, + source_row, + _shard_id, + ): + evaluation = ( + catalog_db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == eval_row.evaluation_id) + .first() + ) + if not evaluation: + eval_row.status = "failed" + eval_row.error_message = "Evaluation parent not found" + row_db.commit() + return {"status": "failed", "reason": "evaluation_missing"} - source_row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == eval_row.call_import_row_id) - .first() - ) - if not source_row: previous_row_status = eval_row.status - eval_row.status = "failed" - eval_row.error_message = "Source call import row not found" - eval_row.finished_at = now_utc() - commit_terminal_row_and_rollup( - db, - evaluation, - eval_row, - previous_row_status=previous_row_status, - ) - return {"status": "failed", "reason": "source_row_missing"} + eval_row.status = "running" + eval_row.celery_task_id = self.request.id + eval_row.error_message = None + eval_row.started_at = eval_row.started_at or now_utc() + if evaluation.status == "pending": + evaluation.status = "running" + evaluation.started_at = evaluation.started_at or now_utc() + _persist_sessions(row_db, catalog_db) + previous_row_status = "running" - previous_row_status = eval_row.status - eval_row.status = "running" - eval_row.celery_task_id = self.request.id - eval_row.error_message = None - eval_row.started_at = eval_row.started_at or now_utc() - if evaluation.status == "pending": - evaluation.status = "running" - evaluation.started_at = evaluation.started_at or now_utc() - db.commit() - previous_row_status = "running" + restricted_uuids = parse_restricted_metric_uuids(restricted_metric_ids) + if restricted_metric_ids is not None and restricted_uuids is not None: + selected_raw = {str(x) for x in (evaluation.selected_metric_ids or [])} + if restricted_uuids and not any( + str(mid) in selected_raw for mid in restricted_uuids + ): + eval_row.status = "completed" + eval_row.error_message = None + eval_row.finished_at = now_utc() + _rollup_terminal( + row_db, + catalog_db, + evaluation, + eval_row, + previous_row_status=previous_row_status, + ) + return { + "status": "skipped", + "reason": "restricted_metric_ids_no_match", + } - restricted_uuids = parse_restricted_metric_uuids(restricted_metric_ids) - if restricted_metric_ids is not None and restricted_uuids is not None: - selected_raw = {str(x) for x in (evaluation.selected_metric_ids or [])} - if restricted_uuids and not any( - str(mid) in selected_raw for mid in restricted_uuids - ): - eval_row.status = "completed" - eval_row.error_message = None + metrics = load_enabled_metrics( + catalog_db, evaluation, restricted_metric_ids=restricted_metric_ids + ) + if not metrics: + eval_row.status = "failed" + eval_row.error_message = "No enabled metrics selected for this evaluation" eval_row.finished_at = now_utc() - commit_terminal_row_and_rollup( - db, + _rollup_terminal( + row_db, + catalog_db, evaluation, eval_row, previous_row_status=previous_row_status, ) - return { - "status": "skipped", - "reason": "restricted_metric_ids_no_match", - } - - metrics = load_enabled_metrics( - db, evaluation, restricted_metric_ids=restricted_metric_ids - ) - if not metrics: - eval_row.status = "failed" - eval_row.error_message = "No enabled metrics selected for this evaluation" - eval_row.finished_at = now_utc() - commit_terminal_row_and_rollup( - db, - evaluation, - eval_row, - previous_row_status=previous_row_status, - ) - return {"status": "failed", "reason": "no_metrics"} - - ( - _transcript_metrics, - audio_metrics, - _comparison_metrics, - metric_scores, - ) = categorize_row_metrics(db, evaluation, source_row, metrics) + return {"status": "failed", "reason": "no_metrics"} - recording_s3_key = (source_row.recording_s3_key or "").strip() or None - result_id = f"call-import-eval:{eval_row.id}" - audio_failed = False + ( + _transcript_metrics, + audio_metrics, + _comparison_metrics, + metric_scores, + ) = categorize_row_metrics(catalog_db, evaluation, source_row, metrics) - if audio_metrics and recording_s3_key: - from app.workers.tasks.helpers.audio_evaluation import ( - evaluate_audio_metrics, - handle_audio_evaluation_error, - ) + recording_s3_key = (source_row.recording_s3_key or "").strip() or None + result_id = f"call-import-eval:{eval_row.id}" + audio_failed = False - try: - audio_scores = evaluate_audio_metrics( - audio_s3_key=recording_s3_key, - audio_metrics=audio_metrics, - result_id=result_id, - ) - metric_scores.update(audio_scores) - except Exception as audio_err: # noqa: BLE001 - logger.exception( - "[CallImportEval {}] Audio analysis failed", eval_row.id + if audio_metrics and recording_s3_key: + from app.workers.tasks.helpers.audio_evaluation import ( + evaluate_audio_metrics, + handle_audio_evaluation_error, ) - metric_scores.update( - handle_audio_evaluation_error(audio_metrics, audio_err) - ) - audio_failed = True - if was_cancelled_externally(db, eval_row): - try: - rollup_parent( - db, - evaluation, - previous_row_status="running", - new_row_status="failed", - ) - db.commit() - except Exception: # noqa: BLE001 - db.rollback() - return {"status": "cancelled", "eval_row_id": eval_row_id} + try: + audio_scores = evaluate_audio_metrics( + audio_s3_key=recording_s3_key, + audio_metrics=audio_metrics, + result_id=result_id, + ) + metric_scores.update(audio_scores) + except Exception as audio_err: # noqa: BLE001 + logger.exception( + "[CallImportEval {}] Audio analysis failed", eval_row.id + ) + metric_scores.update( + handle_audio_evaluation_error(audio_metrics, audio_err) + ) + audio_failed = True - existing = ( - eval_row.metric_scores if isinstance(eval_row.metric_scores, dict) else {} - ) - merged = dict(existing) - for key, value in as_json_dict(metric_scores).items(): - merged[key] = value - eval_row.metric_scores = merged - db.commit() + if was_cancelled_externally(row_db, eval_row): + try: + rollup_parent( + catalog_db, + evaluation, + previous_row_status="running", + new_row_status="failed", + ) + catalog_db.commit() + except Exception: # noqa: BLE001 + catalog_db.rollback() + return {"status": "cancelled", "eval_row_id": eval_row_id} - needs_llm = row_needs_llm_phase( - db, - evaluation, - source_row, - restricted_metric_ids=restricted_metric_ids, - ) + existing = ( + eval_row.metric_scores if isinstance(eval_row.metric_scores, dict) else {} + ) + merged = dict(existing) + for key, value in as_json_dict(metric_scores).items(): + merged[key] = value + eval_row.metric_scores = merged + row_db.commit() - if needs_llm: - from app.workers.tasks.evaluate_call_import_row import ( - evaluate_call_import_row_task, + needs_llm = row_needs_llm_phase( + catalog_db, + evaluation, + source_row, + restricted_metric_ids=restricted_metric_ids, ) - chain_kwargs: dict = { - "_skip_audio": True, - "_eval_slot_task_id": slot_task_id, - } - if restricted_metric_ids: - chain_kwargs["restricted_metric_ids"] = restricted_metric_ids - try: - evaluate_call_import_row_task.apply_async( - args=(str(eval_row.id),), - kwargs=chain_kwargs, - queue=EVALUATIONS_QUEUE, - ) - except Exception: - logger.exception( - "[CallImportEval {}] Failed to enqueue LLM phase", - eval_row.id, - ) - eval_row.status = "failed" - eval_row.error_message = "Failed to enqueue LLM evaluation phase" - eval_row.finished_at = now_utc() - commit_terminal_row_and_rollup( - db, - evaluation, - eval_row, - previous_row_status=previous_row_status, + if needs_llm: + from app.workers.tasks.evaluate_call_import_row import ( + evaluate_call_import_row_task, ) + + chain_kwargs: dict = { + "_skip_audio": True, + "_eval_slot_task_id": slot_task_id, + } + if restricted_metric_ids: + chain_kwargs["restricted_metric_ids"] = restricted_metric_ids + try: + evaluate_call_import_row_task.apply_async( + args=(str(eval_row.id),), + kwargs=chain_kwargs, + queue=EVALUATIONS_QUEUE, + ) + except Exception: + logger.exception( + "[CallImportEval {}] Failed to enqueue LLM phase", + eval_row.id, + ) + eval_row.status = "failed" + eval_row.error_message = "Failed to enqueue LLM evaluation phase" + eval_row.finished_at = now_utc() + _rollup_terminal( + row_db, + catalog_db, + evaluation, + eval_row, + previous_row_status=previous_row_status, + ) + return { + "status": "failed", + "eval_row_id": eval_row_id, + "reason": "chain_enqueue_failed", + } + chain_llm = True return { - "status": "failed", + "status": "chained", "eval_row_id": eval_row_id, - "reason": "chain_enqueue_failed", + "next": "llm_phase", } - chain_llm = True + + if audio_failed: + eval_row.status = "failed" + eval_row.error_message = "Evaluation failed for one or more audio metrics" + else: + eval_row.status = "completed" + eval_row.error_message = None + + eval_row.finished_at = now_utc() + _rollup_terminal( + row_db, + catalog_db, + evaluation, + eval_row, + previous_row_status=previous_row_status, + ) + return { - "status": "chained", + "status": eval_row.status, "eval_row_id": eval_row_id, - "next": "llm_phase", + "phase": "audio_only", } - - if audio_failed: - eval_row.status = "failed" - eval_row.error_message = "Evaluation failed for one or more audio metrics" - else: - eval_row.status = "completed" - eval_row.error_message = None - - eval_row.finished_at = now_utc() - commit_terminal_row_and_rollup( - db, - evaluation, - eval_row, - previous_row_status=previous_row_status, + except LookupError: + logger.warning( + "[CallImportEval audio {}] Row not found on any shard — skipping", + eval_row_id, ) - - return { - "status": eval_row.status, - "eval_row_id": eval_row_id, - "phase": "audio_only", - } + return {"status": "skipped", "reason": "row_not_found"} finally: - db.close() if not chain_llm: from app.workers.concurrency.fair_dispatch import ( finish_eval_work_and_redispatch, diff --git a/app/workers/tasks/evaluate_call_import_row_core.py b/app/workers/tasks/evaluate_call_import_row_core.py index 4dccdde2..81f71199 100644 --- a/app/workers/tasks/evaluate_call_import_row_core.py +++ b/app/workers/tasks/evaluate_call_import_row_core.py @@ -243,39 +243,68 @@ def reconcile_evaluation_counters( db: Session, evaluation: CallImportEvaluation, ) -> None: - """Sync parent counters from child rows using one aggregate query.""" - counts = ( - db.query( - func.count().label("total"), - func.coalesce( - func.sum( - case( - (CallImportEvaluationRow.status == "completed", 1), - else_=0, - ) - ), - 0, - ).label("completed"), - func.coalesce( - func.sum( - case( - (CallImportEvaluationRow.status == "failed", 1), - else_=0, - ) - ), - 0, - ).label("failed"), + """Sync parent counters from child rows using aggregate queries.""" + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import aggregate_evaluation_row_counts + + total, completed, failed = aggregate_evaluation_row_counts(db, evaluation.id) + else: + counts = ( + db.query( + func.count().label("total"), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "completed", 1), + else_=0, + ) + ), + 0, + ).label("completed"), + func.coalesce( + func.sum( + case( + (CallImportEvaluationRow.status == "failed", 1), + else_=0, + ) + ), + 0, + ).label("failed"), + ) + .filter(CallImportEvaluationRow.evaluation_id == evaluation.id) + .one() + ) + total = int(counts.total or 0) + completed = int(counts.completed or 0) + failed = int(counts.failed or 0) + + db.execute( + update(CallImportEvaluation) + .where(CallImportEvaluation.id == evaluation.id) + .values( + total_rows=total, + completed_rows=completed, + failed_rows=failed, ) - .filter(CallImportEvaluationRow.evaluation_id == evaluation.id) - .one() ) - evaluation.total_rows = int(counts.total or 0) - evaluation.completed_rows = int(counts.completed or 0) - evaluation.failed_rows = int(counts.failed or 0) + db.flush() + db.refresh(evaluation) + from app.services.call_imports.progress_counters import clear_eval_progress_redis + + clear_eval_progress_redis(evaluation.id) def _count_in_progress_rows(db: Session, evaluation_id: UUID) -> int: """Count rows still pending or running (cheap indexed query).""" + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled(): + from app.db_sharding.scatter_gather import count_eval_rows_in_progress + + return count_eval_rows_in_progress(db, evaluation_id) + return int( db.query(func.count()) .filter( @@ -313,20 +342,30 @@ def _apply_parent_status_from_counters( def commit_terminal_row_and_rollup( - db: Session, + row_db: Session, evaluation: CallImportEvaluation, eval_row: CallImportEvaluationRow, *, previous_row_status: str, + catalog_db: Session | None = None, ) -> None: - db.commit() + parent_db = catalog_db if catalog_db is not None and catalog_db is not row_db else row_db + row_db.commit() + if parent_db is not row_db: + evaluation = ( + parent_db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == evaluation.id) + .first() + ) + if evaluation is None: + return rollup_parent( - db, + parent_db, evaluation, previous_row_status=previous_row_status, new_row_status=eval_row.status, ) - db.commit() + parent_db.commit() def rollup_parent( @@ -363,12 +402,11 @@ def rollup_parent( else: reconcile_evaluation_counters(db, evaluation) - evaluation = ( - db.query(CallImportEvaluation) - .filter(CallImportEvaluation.id == evaluation_id) - .with_for_update() - .one() - ) + from app.services.call_imports.progress_counters import clear_eval_progress_redis + + clear_eval_progress_redis(evaluation_id) + + db.refresh(evaluation) if previous_row_status is not None and new_row_status is not None: expected_in_progress = ( int(evaluation.total_rows or 0) @@ -377,19 +415,18 @@ def rollup_parent( ) if expected_in_progress != _count_in_progress_rows(db, evaluation_id): reconcile_evaluation_counters(db, evaluation) - db.flush() - evaluation = ( - db.query(CallImportEvaluation) - .filter(CallImportEvaluation.id == evaluation_id) - .with_for_update() - .one() - ) - _apply_parent_status_from_counters(evaluation) + db.refresh(evaluation) completed = int(evaluation.completed_rows or 0) already_billed = int(getattr(evaluation, "billed_completed_rows", 0) or 0) delta = completed - already_billed if delta > 0: + evaluation = ( + db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == evaluation_id) + .with_for_update() + .one() + ) from app.services.billing.flexprice_service import ( record_call_import_evaluation_completed, ) @@ -407,6 +444,8 @@ def rollup_parent( if billing_accepted: evaluation.billed_completed_rows = completed + _apply_parent_status_from_counters(evaluation) + def parse_restricted_metric_uuids( restricted_metric_ids: Optional[List[str]], diff --git a/app/workers/tasks/generate_evaluation_metric_clusters.py b/app/workers/tasks/generate_evaluation_metric_clusters.py index 1906976e..e91dc108 100644 --- a/app/workers/tasks/generate_evaluation_metric_clusters.py +++ b/app/workers/tasks/generate_evaluation_metric_clusters.py @@ -69,15 +69,9 @@ def generate_evaluation_metric_clusters_task( UUID(credential_id) if credential_id else None, ) - rows = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join( - CallImportRow, - CallImportRow.id == CallImportEvaluationRow.call_import_row_id, - ) - .filter(CallImportEvaluationRow.evaluation_id == evaluation.id) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + + rows = load_evaluation_row_pairs(db, evaluation.id) completed_pairs = [ (eval_row, source_row) for eval_row, source_row in rows diff --git a/app/workers/tasks/generate_evaluation_user_insights.py b/app/workers/tasks/generate_evaluation_user_insights.py index 6c07e9f7..6029ce4a 100644 --- a/app/workers/tasks/generate_evaluation_user_insights.py +++ b/app/workers/tasks/generate_evaluation_user_insights.py @@ -55,15 +55,9 @@ def generate_evaluation_user_insights_task( model, ) - rows = ( - db.query(CallImportEvaluationRow, CallImportRow) - .join( - CallImportRow, - CallImportRow.id == CallImportEvaluationRow.call_import_row_id, - ) - .filter(CallImportEvaluationRow.evaluation_id == evaluation.id) - .all() - ) + from app.db_sharding.scatter_gather import load_evaluation_row_pairs + + rows = load_evaluation_row_pairs(db, evaluation.id) completed_pairs = [ (eval_row, source_row) for eval_row, source_row in rows diff --git a/app/workers/tasks/helpers/llm_evaluation.py b/app/workers/tasks/helpers/llm_evaluation.py index 01d13fc3..66f78c14 100644 --- a/app/workers/tasks/helpers/llm_evaluation.py +++ b/app/workers/tasks/helpers/llm_evaluation.py @@ -1193,6 +1193,13 @@ def evaluate_with_llm( evaluation_start_time = time.time() evaluator_llm_config = getattr(evaluator, "llm_config", None) if evaluator else None + evaluator_credential_id = getattr(evaluator, "llm_credential_id", None) if evaluator else None + parsed_credential_id = None + if evaluator_credential_id: + try: + parsed_credential_id = UUID(str(evaluator_credential_id)) + except (TypeError, ValueError): + parsed_credential_id = None llm_result = llm_service.generate_response( messages=messages, llm_provider=llm_provider, @@ -1201,6 +1208,7 @@ def evaluate_with_llm( db=db, llm_config=evaluator_llm_config, task_defaults={"temperature": 0.3, "max_tokens": dynamic_max_tokens}, + credential_id=parsed_credential_id, ) evaluation_time = time.time() - evaluation_start_time diff --git a/app/workers/tasks/process_call_import_row.py b/app/workers/tasks/process_call_import_row.py index 06e0ec7a..5138cc3a 100644 --- a/app/workers/tasks/process_call_import_row.py +++ b/app/workers/tasks/process_call_import_row.py @@ -27,7 +27,6 @@ from loguru import logger from sqlalchemy.orm.exc import StaleDataError -from app.database import SessionLocal from app.workers.config import celery_app @@ -82,8 +81,14 @@ def _safe_commit( context: str = "update", ) -> bool: """Commit pending ORM changes; return False when the row/import was deleted.""" + from app.db_sharding.row_ops import commit_shard_row_session + from app.db_sharding.sessions import is_sharding_enabled + try: - db.commit() + if is_sharding_enabled(): + commit_shard_row_session(db) + else: + db.commit() return True except StaleDataError: db.rollback() @@ -95,10 +100,42 @@ def _safe_commit( return False -def _row_or_import_gone(db, row_id: UUID) -> Optional[str]: +def _row_or_import_gone( + db, + row_id: UUID, + *, + catalog_db=None, +) -> Optional[str]: """Return a skip reason when the row or its parent import no longer exists.""" from app.models.database import CallImport, CallImportRow from app.models.enums import CallImportStatus + from app.db_sharding.sessions import is_sharding_enabled + + if is_sharding_enabled() and catalog_db is not None and catalog_db is not db: + row_pk = ( + db.query(CallImportRow.id) + .filter(CallImportRow.id == row_id) + .first() + ) + if row_pk is None: + return "row_deleted" + call_import_id = ( + db.query(CallImportRow.call_import_id) + .filter(CallImportRow.id == row_id) + .scalar() + ) + if call_import_id is None: + return "row_deleted" + import_status = ( + catalog_db.query(CallImport.status) + .filter(CallImport.id == call_import_id) + .scalar() + ) + if import_status is None: + return "row_deleted" + if import_status == CallImportStatus.DELETING: + return "import_deleting" + return None db.expire_all() hit = ( @@ -159,6 +196,15 @@ def _rollup_parent_status(db, call_import) -> None: rollup_call_import_batch_status(db, call_import) +def _rollup_parent_on_catalog(catalog_db, row_db, call_import) -> None: + from app.db_sharding.sessions import is_sharding_enabled + + target = catalog_db if is_sharding_enabled() and catalog_db is not row_db else row_db + _rollup_parent_status(target, call_import) + if target is not row_db: + target.commit() + + @celery_app.task(name="process_call_import_row", bind=True, max_retries=3) def process_call_import_row_task( self, @@ -194,15 +240,31 @@ def process_call_import_row_task( slot_task_id = _eval_slot_task_id or self.request.id eval_chain_chained_transcribe = False - db = SessionLocal() + row_db = catalog_db = None + from app.db_sharding.row_ops import close_row_sessions, locate_call_import_row + from app.models.database import CallImport + try: - row_uuid = UUID(row_id) - row = db.query(CallImportRow).filter(CallImportRow.id == row_uuid).first() - if row is None: + try: + row_db, catalog_db, row, shard_id = locate_call_import_row(row_id) + except LookupError: logger.warning("CallImportRow {} not found, skipping", row_id) return {"status": "skipped", "reason": "row_not_found"} - call_import = row.call_import + db = row_db + row_uuid = row.id + + if catalog_db is not row_db: + call_import = ( + catalog_db.query(CallImport) + .filter(CallImport.id == row.call_import_id) + .first() + ) + logger.bind(shard_id=shard_id).debug( + "process_call_import_row on shard {}", shard_id + ) + else: + call_import = row.call_import from app.models.enums import CallImportStatus @@ -236,7 +298,7 @@ def process_call_import_row_task( # for legacy rows imported before the column existed. client = telephony_service.get_provider_client( row.organization_id, - db, + catalog_db if catalog_db is not row_db else db, provider=call_import.provider, credential_id=call_import.telephony_integration_id, ) @@ -246,7 +308,7 @@ def process_call_import_row_task( row.error_message = f"Provider client error: {exc}" if not _safe_commit(db, row_id=row_id, context="provider_client_error"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_provider_client_error"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "provider_client_error"} @@ -272,7 +334,7 @@ def process_call_import_row_task( row.error_message = msg if not _safe_commit(db, row_id=row_id, context="direct_url_no_source"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_direct_url_no_source"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "no_recording_source"} @@ -320,7 +382,7 @@ def process_call_import_row_task( ) if not _safe_commit(db, row_id=row_id, context="direct_url_failed"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_direct_url_failed"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "non_retryable_provider_error"} @@ -353,7 +415,7 @@ def process_call_import_row_task( row.error_message = msg if not _safe_commit(db, row_id=row_id, context="no_recording_source"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_no_recording_source"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "no_recording_source"} @@ -419,7 +481,7 @@ def process_call_import_row_task( ) if not _safe_commit(db, row_id=row_id, context="non_retryable_provider_error"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_non_retryable_provider_error"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "non_retryable_provider_error"} @@ -434,7 +496,7 @@ def process_call_import_row_task( row.error_message = f"Cloud blob storage unavailable: {err}" if not _safe_commit(db, row_id=row_id, context="s3_unavailable"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_s3_unavailable"): return {"status": "skipped", "reason": "row_deleted"} return {"status": "failed", "reason": "s3_unavailable"} @@ -456,7 +518,11 @@ def process_call_import_row_task( return {"status": "skipped", "reason": "row_deleted"} raise self.retry(exc=exc, countdown=_RETRYABLE_COUNTDOWN_SECONDS) - skip_reason = _row_or_import_gone(db, row_uuid) + skip_reason = _row_or_import_gone( + db, + row_uuid, + catalog_db=catalog_db if catalog_db is not row_db else None, + ) if skip_reason: db.rollback() logger.info( @@ -469,12 +535,13 @@ def process_call_import_row_task( row.recording_s3_key = key row.recording_content_type = content_type row.recording_size_bytes = len(audio_bytes) + previous_import_status = row.status row.status = CallImportRowStatus.COMPLETED row.error_message = None if not _safe_commit(db, row_id=row_id, context="mark_completed"): return {"status": "skipped", "reason": "row_deleted"} - _rollup_parent_status(db, call_import) + _rollup_parent_on_catalog(catalog_db, row_db, call_import) if not _safe_commit(db, row_id=row_id, context="rollup_completed"): return {"status": "skipped", "reason": "row_deleted"} @@ -490,8 +557,14 @@ def process_call_import_row_task( .first() ) if eval_row is not None: + # Evaluation headers live on the catalog when sharding is on. + eval_header_db = ( + catalog_db + if catalog_db is not None and catalog_db is not row_db + else db + ) evaluation = ( - db.query(CallImportEvaluation) + eval_header_db.query(CallImportEvaluation) .filter(CallImportEvaluation.id == eval_row.evaluation_id) .first() ) @@ -519,41 +592,52 @@ def process_call_import_row_task( ) return {"status": "skipped", "reason": "row_deleted"} finally: - db.close() + if row_db is not None: + close_row_sessions(row_db, catalog_db if catalog_db is not row_db else None) from app.workers.concurrency.limits import slot_registered_for_task if slot_registered_for_task(slot_task_id): if run_eval_row_id: if not eval_chain_chained_transcribe: - cleanup_db = SessionLocal() - try: - from app.models.database import CallImportEvaluationRow, CallImportRow - from app.models.enums import CallImportRowStatus - from app.workers.concurrency.eval_dispatch import ( - _fail_eval_row_for_import, - ) + from app.db_sharding.row_ops import ( + close_row_sessions as close_eval_sessions, + locate_call_import_evaluation_row, + ) + from app.workers.concurrency.eval_dispatch import ( + _fail_eval_row_for_import, + recover_eval_row_for_eval_chain, + source_row_import_blocks_eval, + ) - eval_row = ( - cleanup_db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == UUID(run_eval_row_id)) - .first() - ) - source_row = ( - cleanup_db.query(CallImportRow) - .filter(CallImportRow.id == UUID(row_id)) - .first() + try: + cleanup_row_db, cleanup_catalog_db, eval_row, source_row, _ = ( + locate_call_import_evaluation_row( + UUID(run_eval_row_id) + ) ) - if ( - eval_row is not None - and source_row is not None - and source_row.status != CallImportRowStatus.COMPLETED - and eval_row.status == "pending" - ): - _fail_eval_row_for_import( - cleanup_db, eval_row, source_row + except LookupError: + pass + else: + try: + if source_row_import_blocks_eval(source_row): + if eval_row.status == "pending": + _fail_eval_row_for_import( + cleanup_row_db, eval_row, source_row + ) + else: + recover_eval_row_for_eval_chain(eval_row) + if eval_row.status == "pending": + from app.db_sharding.row_ops import ( + commit_shard_row_session, + ) + + eval_row.celery_task_id = None + source_row.celery_task_id = None + commit_shard_row_session(cleanup_row_db) + finally: + close_eval_sessions( + cleanup_row_db, cleanup_catalog_db ) - finally: - cleanup_db.close() from app.workers.concurrency.fair_dispatch import ( finish_eval_work_and_redispatch, diff --git a/app/workers/tasks/transcribe_call_import_row.py b/app/workers/tasks/transcribe_call_import_row.py index ebdd3420..c6e89198 100644 --- a/app/workers/tasks/transcribe_call_import_row.py +++ b/app/workers/tasks/transcribe_call_import_row.py @@ -289,53 +289,77 @@ def _summarize_exc(exc: BaseException, *, max_chars: int = 240) -> str: return _compact_diarisation_error(text, max_chars=max_chars) -def _apply_eval_chain_transcribe_cleanup(db, run_eval_row_id: str) -> None: +def _apply_eval_chain_transcribe_cleanup(run_eval_row_id: str) -> None: """Clear eval-row dispatch state after eval-chain transcribe completes.""" - from app.models.database import CallImportEvaluationRow, CallImportRow - - eval_row = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == UUID(run_eval_row_id)) - .first() + from app.db_sharding.row_ops import ( + close_row_sessions, + locate_call_import_evaluation_row, ) - if eval_row is None: - return - - eval_row.celery_task_id = None - source_row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == eval_row.call_import_row_id) - .first() + from app.models.database import CallImportEvaluation + from app.workers.tasks.evaluate_call_import_row_core import ( + commit_terminal_row_and_rollup, ) - if ( - source_row is not None - and (source_row.diarised_transcript_status or "").lower() == "failed" - and eval_row.status == "pending" - ): - eval_row.status = "failed" - eval_row.error_message = ( - source_row.diarised_transcript_error or "Diarisation failed" + + try: + row_db, catalog_db, eval_row, source_row, _ = ( + locate_call_import_evaluation_row(UUID(run_eval_row_id)) ) - eval_row.finished_at = _now() - db.commit() + except LookupError: + return + try: + eval_row.celery_task_id = None + previous_status = eval_row.status or "pending" + marked_failed = False + if ( + (source_row.diarised_transcript_status or "").lower() == "failed" + and eval_row.status == "pending" + ): + eval_row.status = "failed" + eval_row.error_message = ( + source_row.diarised_transcript_error or "Diarisation failed" + ) + eval_row.finished_at = _now() + marked_failed = True + + if marked_failed: + parent_db = catalog_db if catalog_db is not None else row_db + evaluation = ( + parent_db.query(CallImportEvaluation) + .filter(CallImportEvaluation.id == eval_row.evaluation_id) + .first() + ) + if evaluation is not None: + commit_terminal_row_and_rollup( + row_db, + evaluation, + eval_row, + previous_row_status=previous_status, + catalog_db=( + catalog_db + if catalog_db is not None and catalog_db is not row_db + else None + ), + ) + else: + row_db.commit() + else: + row_db.commit() + finally: + close_row_sessions(row_db, catalog_db) def _was_cancelled_by_row_id(row_id: str | UUID) -> bool: """Re-read the row in a short-lived session during slow I/O.""" - from app.models.database import CallImportRow + from app.db_sharding.row_ops import close_row_sessions, locate_call_import_row - db = SessionLocal() try: - row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == UUID(str(row_id))) - .first() - ) - if row is None: - return False - return _was_cancelled_externally(db, row) + row_db, catalog_db, row, _ = locate_call_import_row(row_id) + except LookupError: + return False + try: + return _was_cancelled_externally(row_db, row) finally: - db.close() + close_row_sessions(row_db, catalog_db) def _persist_diarization_failure( @@ -347,23 +371,21 @@ def _persist_diarization_failure( """Write a terminal diarisation failure without holding a long session.""" from app.models.database import CallImportRow - db = SessionLocal() + from app.db_sharding.row_ops import close_row_sessions, locate_call_import_row + try: - row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == UUID(str(row_id))) - .first() - ) - if row is None: - return {"status": "skipped", "reason": "row_not_found"} - if _was_cancelled_externally(db, row): + row_db, catalog_db, row, _ = locate_call_import_row(row_id) + except LookupError: + return {"status": "skipped", "reason": "row_not_found"} + try: + if _was_cancelled_externally(row_db, row): return {"status": "cancelled", "reason": "cancelled_by_user"} row.diarised_transcript_status = "failed" row.diarised_transcript_error = error_message - db.commit() + row_db.commit() return {"status": "failed", "reason": reason} finally: - db.close() + close_row_sessions(row_db, catalog_db) def _run_diarization_pipeline(ctx: dict[str, Any]) -> dict[str, Any]: @@ -577,18 +599,15 @@ def _finalize_diarization_row( if pipeline_ctx.get("stt_provider"): provider_enum = ModelProvider(pipeline_ctx["stt_provider"]) - db = SessionLocal() - try: - row = ( - db.query(CallImportRow) - .filter(CallImportRow.id == UUID(str(row_id))) - .first() - ) - if row is None: - return {"status": "skipped", "reason": "row_not_found"} + from app.db_sharding.row_ops import close_row_sessions, locate_call_import_row + try: + row_db, catalog_db, row, _ = locate_call_import_row(row_id) + except LookupError: + return {"status": "skipped", "reason": "row_not_found"} + try: if work_result.get("reason") == "no_speech_detected": - if _was_cancelled_externally(db, row): + if _was_cancelled_externally(row_db, row): return {"status": "cancelled", "reason": "cancelled_by_user"} row.diarised_transcript = "" row.diarised_segments = [] @@ -601,7 +620,7 @@ def _finalize_diarization_row( ) row.diarised_prompt = effective_prompt row.diarised_at = _now() - db.commit() + row_db.commit() return { "status": "completed", "row_id": str(row_id), @@ -633,7 +652,7 @@ def _finalize_diarization_row( "storable transcript; marking completed with empty transcript.", row_id, ) - if _was_cancelled_externally(db, row): + if _was_cancelled_externally(row_db, row): return { "status": "cancelled", "reason": "cancelled_by_user", @@ -648,7 +667,7 @@ def _finalize_diarization_row( ) row.diarised_prompt = effective_prompt row.diarised_at = _now() - db.commit() + row_db.commit() return { "status": "completed", "row_id": str(row_id), @@ -658,7 +677,7 @@ def _finalize_diarization_row( "characters": 0, } - if _was_cancelled_externally(db, row): + if _was_cancelled_externally(row_db, row): logger.info( "Row {} was cancelled by the user mid-flight; " "skipping success write and preserving cancelled state.", @@ -677,7 +696,7 @@ def _finalize_diarization_row( row.diarised_transcript_error = None row.diarised_prompt = effective_prompt row.diarised_at = _now() - db.commit() + row_db.commit() return { "status": "completed", @@ -691,7 +710,7 @@ def _finalize_diarization_row( "turn_count": len(turns) if has_real_diarisation else 0, } finally: - db.close() + close_row_sessions(row_db, catalog_db) @celery_app.task( @@ -754,22 +773,28 @@ def transcribe_call_import_row_task( slot_task_id = _eval_slot_task_id or self.request.id pipeline_ctx: dict[str, Any] | None = None try: - db = SessionLocal() + from app.db_sharding.row_ops import ( + close_row_sessions, + locate_call_import_evaluation_row, + locate_call_import_row, + ) + + row_db = catalog_db = None try: if run_eval_row_id: - from app.models.database import CallImportEvaluationRow - - eval_row_for_chain = ( - db.query(CallImportEvaluationRow) - .filter(CallImportEvaluationRow.id == UUID(run_eval_row_id)) - .first() - ) - if eval_row_for_chain is not None: + try: + _er_db, _cat, eval_row_for_chain, _, _ = ( + locate_call_import_evaluation_row(UUID(run_eval_row_id)) + ) evaluation_id_for_dispatch = str( eval_row_for_chain.evaluation_id ) - row_uuid = UUID(row_id) - row = db.query(CallImportRow).filter(CallImportRow.id == row_uuid).first() + close_row_sessions(_er_db, _cat) + except LookupError: + pass + + row_db, catalog_db, row, _shard_id = locate_call_import_row(row_id) + row_uuid = row.id if row is None: logger.warning( "transcribe_call_import_row: row {} not found, skipping", @@ -791,7 +816,7 @@ def transcribe_call_import_row_task( existing_diarised = (row.diarised_transcript or "").strip() if existing_diarised and not overwrite_existing: row.diarised_transcript_status = "completed" - db.commit() + row_db.commit() return { "status": "skipped", "reason": "transcript_present", @@ -804,7 +829,7 @@ def transcribe_call_import_row_task( row.diarised_transcript_error = ( "No recording available for this row; cannot diarise." ) - db.commit() + row_db.commit() return {"status": "skipped", "reason": "no_recording"} normalised_mode = (mode or "stt_llm").strip().lower() @@ -814,7 +839,7 @@ def transcribe_call_import_row_task( f"Unknown diarisation mode '{mode}'. Expected " "'stt_llm' or 'llm_only'." ) - db.commit() + row_db.commit() return {"status": "failed", "reason": "unknown_mode"} provider_enum: Optional[ModelProvider] = None @@ -825,7 +850,7 @@ def transcribe_call_import_row_task( "STT provider/model not configured. Pick an STT " "model in the Diarise modal." ) - db.commit() + row_db.commit() return {"status": "failed", "reason": "missing_stt"} try: provider_enum = ModelProvider(stt_provider.lower()) @@ -834,7 +859,7 @@ def transcribe_call_import_row_task( row.diarised_transcript_error = ( f"Unknown STT provider '{stt_provider}'." ) - db.commit() + row_db.commit() return {"status": "failed", "reason": "unknown_provider"} llm_provider_value = (diarization_llm_provider or "").strip() @@ -845,7 +870,7 @@ def transcribe_call_import_row_task( "Diarisation LLM provider/model not configured. Pick " "a chat model in the Diarise modal." ) - db.commit() + row_db.commit() return {"status": "failed", "reason": "missing_llm_diariser"} row.diarised_transcript_status = "running" @@ -879,7 +904,7 @@ def transcribe_call_import_row_task( except (TypeError, ValueError): credential_uuid = None - db.commit() + row_db.commit() pipeline_ctx = { "row_id": row_id, @@ -896,7 +921,8 @@ def transcribe_call_import_row_task( "effective_prompt": effective_prompt, } finally: - db.close() + if row_db is not None: + close_row_sessions(row_db, catalog_db) if pipeline_ctx is None: return {"status": "skipped", "reason": "setup_incomplete"} @@ -911,58 +937,43 @@ def transcribe_call_import_row_task( work_result=work_result, ) except Exception as exc: # noqa: BLE001 — terminal row state + no retry loop + from app.db_sharding.row_ops import close_row_sessions, locate_call_import_row + logger.exception( "transcribe_call_import_row crashed for row {}", row_id ) - fail_db = SessionLocal() try: - row = ( - fail_db.query(CallImportRow) - .filter(CallImportRow.id == UUID(row_id)) - .first() - ) - if row is not None: - if _was_cancelled_externally(fail_db, row): - return {"status": "cancelled", "reason": "cancelled_by_user"} - if (row.diarised_transcript_status or "").lower() == "running": - row.diarised_transcript_status = "failed" - row.diarised_transcript_error = _summarize_exc(exc) - fail_db.commit() + row_db, catalog_db, row, _ = locate_call_import_row(row_id) + except LookupError: + return {"status": "failed", "reason": "unexpected_error"} + try: + if _was_cancelled_externally(row_db, row): + return {"status": "cancelled", "reason": "cancelled_by_user"} + if (row.diarised_transcript_status or "").lower() == "running": + row.diarised_transcript_status = "failed" + row.diarised_transcript_error = _summarize_exc(exc) + row_db.commit() except Exception: logger.exception( "Failed to persist unexpected-error state for row {}", row_id, ) try: - fail_db.rollback() + row_db.rollback() except Exception: pass finally: - fail_db.close() + close_row_sessions(row_db, catalog_db) return {"status": "failed", "reason": "unexpected_error"} finally: - cleanup_db = SessionLocal() - try: - if run_eval_row_id: - try: - _apply_eval_chain_transcribe_cleanup(cleanup_db, run_eval_row_id) - except Exception: - logger.exception( - "Failed to finalize eval-chain transcribe cleanup for row {}", - run_eval_row_id, - ) - try: - cleanup_db.rollback() - _apply_eval_chain_transcribe_cleanup( - cleanup_db, run_eval_row_id - ) - except Exception: - logger.exception( - "Failed to clear stale celery_task_id for eval row {}", - run_eval_row_id, - ) - finally: - cleanup_db.close() + if run_eval_row_id: + try: + _apply_eval_chain_transcribe_cleanup(run_eval_row_id) + except Exception: + logger.exception( + "Failed to finalize eval-chain transcribe cleanup for row {}", + run_eval_row_id, + ) if run_eval_row_id: from app.workers.concurrency.fair_dispatch import ( finish_eval_work_and_redispatch, diff --git a/config.docker.sharding.example.yml b/config.docker.sharding.example.yml new file mode 100644 index 00000000..d36561f2 --- /dev/null +++ b/config.docker.sharding.example.yml @@ -0,0 +1,22 @@ +# Example Docker Compose config with generic data-plane database names. +# Copy to config.docker.yml (or merge into your existing file) and enable sharding +# after mounting docker/postgres/init-sharding-dbs.sh on the db service. +# +# Align compose env DATABASE_URL with catalog for all services, e.g.: +# DATABASE_URL=postgresql://efficientai:password@db:5432/efficientai_catalog + +database: + url: "postgresql://efficientai:password@db:5432/efficientai_catalog" + catalog_url: "postgresql://efficientai:password@db:5432/efficientai_catalog" + pool_size: 3 + max_overflow: 5 + sharding: + enabled: true + row_chunk_size: 500 + shards: + - id: data-shard-01 + url: "postgresql://efficientai:password@db:5432/efficientai_data_01" + - id: data-shard-02 + url: "postgresql://efficientai:password@db:5432/efficientai_data_02" + - id: data-shard-03 + url: "postgresql://efficientai:password@db:5432/efficientai_data_03" diff --git a/config.yml.example b/config.yml.example index 882d050d..14feedfd 100644 --- a/config.yml.example +++ b/config.yml.example @@ -21,6 +21,7 @@ operational: # Database Configuration database: + # Legacy single-DB mode: use one database (e.g. efficientai). Sharding off (default). url: "postgresql://efficientai:password@localhost:5432/efficientai" # Alternative: specify individual components # user: "efficientai" @@ -29,6 +30,30 @@ database: # port: 5432 # db: "efficientai" + # Connection pool (per process; use smaller values when many data shards) + pool_size: 10 + max_overflow: 20 + + # Data-plane sharding (call-import rows today; same shard slots can host other + # high-volume row data later). Default off — single database.url only. + # + # When enabled, prefer generic physical names (not call-import-specific): + # catalog → efficientai_catalog (metadata, headers, registry, integrations) + # shards → efficientai_data_01, efficientai_data_02, … + # shard id → data-shard-01, data-shard-02, … (labels in config + registry) + # + # catalog_url: "postgresql://efficientai:password@localhost:5432/efficientai_catalog" + sharding: + enabled: false + row_chunk_size: 500 + # shards: + # - id: data-shard-01 + # url: "postgresql://efficientai:password@localhost:5432/efficientai_data_01" + # - id: data-shard-02 + # url: "postgresql://efficientai:password@localhost:5432/efficientai_data_02" + # - id: data-shard-03 + # url: "postgresql://efficientai:password@localhost:5432/efficientai_data_03" + # Redis Configuration redis: url: "redis://localhost:6379/0" diff --git a/docker-compose.yml b/docker-compose.yml index 0f873af3..770f03e5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -16,6 +16,9 @@ services: POSTGRES_DB: ${POSTGRES_DB:-efficientai} volumes: - postgres_data:/var/lib/postgresql/data + # Optional: generic catalog + data DBs for sharding tests (first init only): + # - ./docker/postgres/init-sharding-dbs.sh:/docker-entrypoint-initdb.d/02-init-sharding-dbs.sh:ro + # See config.docker.sharding.example.yml ports: - "5432:5432" healthcheck: @@ -114,10 +117,10 @@ services: # Dedicated worker for call-import + evaluation queues. Celery drains # ``imports`` (recording fetch) before ``diarization`` (manual diarise), - # then ``evaluations`` (LLM scoring). Fair eval dispatch + eval - # materialize run on the ``celery`` queue (``worker`` service) so a large - # import/eval fan-out in one workspace cannot head-of-line block dispatch - # for other workspaces. + # then ``eval-control`` (cancel/retry/materialize), then ``evaluations`` + # (fair dispatch + LLM scoring). + # The default ``worker`` service handles ``celery`` (legacy) and + # ``audio-metrics`` (Praat/UTMOS audio metric tasks). worker-imports: image: ghcr.io/efficientai-tech/efficientai-worker:${EFFICIENTAI_VERSION:-latest} build: @@ -145,7 +148,9 @@ services: - ./config.docker.yml:/app/config.yml:ro # Optional: mount GCP service account for GCS auth # - ./secrets/gcp-sa.json:/app/secrets/gcp-sa.json:ro - command: eai worker --config /app/config.yml --loglevel info --queues imports,diarization,evaluations --pool threads --concurrency 32 + # Thread pool: keep concurrency near Redis inflight caps + headroom when sharding + # (each task may hold catalog + shard connections for tens of seconds). + command: eai worker --config /app/config.yml --loglevel info --queues imports,diarization,eval-control,evaluations --pool threads --concurrency 12 volumes: postgres_data: diff --git a/docker/postgres/init-sharding-dbs.sh b/docker/postgres/init-sharding-dbs.sh new file mode 100644 index 00000000..417130fa --- /dev/null +++ b/docker/postgres/init-sharding-dbs.sh @@ -0,0 +1,10 @@ +#!/bin/bash +# Creates generic catalog + data-plane databases for local sharding tests. +# Runs once on first Postgres volume init (docker-entrypoint-initdb.d). +set -e +psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL + CREATE DATABASE efficientai_catalog; + CREATE DATABASE efficientai_data_01; + CREATE DATABASE efficientai_data_02; + CREATE DATABASE efficientai_data_03; +EOSQL diff --git a/docs-fumadocs/content/docs/advanced/call-import-sharding.mdx b/docs-fumadocs/content/docs/advanced/call-import-sharding.mdx new file mode 100644 index 00000000..b5a5d877 --- /dev/null +++ b/docs-fumadocs/content/docs/advanced/call-import-sharding.mdx @@ -0,0 +1,51 @@ +--- +id: call-import-sharding +title: Call Import Sharding +sidebar_position: 4 +description: Multi-database sharding for large call-import batches +--- + +# Call Import Sharding + +For large batches (10k+ rows), call-import row data can be spread across multiple PostgreSQL **data shards** with a **catalog** database for metadata, routing, and parent counters. + +## Architecture + +- **Catalog DB** — `CallImport`, `CallImportEvaluation`, shard slice registry, dispatch metadata +- **Data shards** — `CallImportRow`, `CallImportEvaluationRow` (partitioned by consistent hash on `row_index`) +- **Scatter/gather reads** — API and workers query each shard and merge results +- **Fair dispatch** — import/eval workers respect per-shard pending scans and Redis progress keys + +Enable sharding in `config.yml` under `database.sharding`. See `config.yml.example` and `config.docker.sharding.example.yml` for profiles. + +## Operations + +### Connection pools + +When `len(shards) > 1`, use smaller per-process pools (`pool_size` 3–5, `max_overflow` 5–10) so API + workers × shards stay under Postgres `max_connections`. + +### PgBouncer (optional) + +Point `database.url`, `catalog_url`, and each shard `url` at PgBouncer (`:6432`) in transaction pooling mode. Keep SQLAlchemy `pool_pre_ping` enabled. + +### Observability + +Row Celery tasks log `shard_id` on the import worker hot path. Watch Redis eval/import progress keys (`eval:progress`, `import:progress`) alongside catalog parent counters. + +### Rebalance + +Dry-run registry updates: + +```bash +python scripts/rebalance_call_import_shards.py --target-shard data-shard-02 +``` + +Use `--apply` only after pausing the import. The rebalance tool copies rows to the target shard, updates the catalog registry, then removes copies from the source shard. + +## Troubleshooting stalled evaluations + +If evaluation runs stall after recordings import (rows show `completed` import but no diarization/scoring): + +1. Restart API and `worker-imports` after deploying fixes. +2. **Retry evaluation** from the UI (use *Overwrite existing transcripts* if diarization previously failed). +3. Ensure the imports worker consumes **`imports,diarization,eval-control,evaluations`**. diff --git a/docs-fumadocs/content/docs/operations/call-import-sharding.mdx b/docs-fumadocs/content/docs/operations/call-import-sharding.mdx new file mode 100644 index 00000000..4e437196 --- /dev/null +++ b/docs-fumadocs/content/docs/operations/call-import-sharding.mdx @@ -0,0 +1,47 @@ +--- +title: Call-import sharding operations +description: Pools, PgBouncer, and observability for multi-node Postgres +--- + +# Call-import sharding operations + +## Configuration + +See `config.yml.example` under `database.sharding`. Default is **off** (`enabled: false`). + +**Naming:** use generic Postgres database names (`efficientai_catalog`, `efficientai_data_01`, …) +and config shard ids (`data-shard-01`, …). Routing is call-import-specific today; the same +data-plane slots can be reused for other row-heavy features later. + +| Profile | `enabled` | `shards` | `catalog_url` | +|---------|-----------|----------|---------------| +| Legacy | false | — | — | +| Single-node enterprise | true | 1 entry | optional | +| Standard | true | 3–5 | dedicated | +| High volume | true | 6+ | dedicated | + +## Connection pools + +When `len(shards) > 1`, use smaller per-process pools (`pool_size` 3–5, `max_overflow` 5–10) +so API + workers × shards stay under Postgres `max_connections`. + +## PgBouncer (optional) + +Point `database.url`, `catalog_url`, and each shard `url` at PgBouncer (`:6432`) in +transaction pooling mode. Keep SQLAlchemy `pool_pre_ping` enabled. + +## Observability + +Row Celery tasks log `shard_id` on the import worker hot path. Fair dispatch uses +per-shard pending scans when sharding is enabled—watch Redis eval/import progress keys +(`eval:progress`, `import:progress`) alongside catalog parent counters. + +## Rebalance + +Dry-run registry updates: + +```bash +python scripts/rebalance_call_import_shards.py --target-shard data-shard-02 +``` + +Use `--apply` only after pausing the import and copying row data between shards. diff --git a/docs-fumadocs/content/docs/products/call-imports.mdx b/docs-fumadocs/content/docs/products/call-imports.mdx index 54503aee..c8eac357 100644 --- a/docs-fumadocs/content/docs/products/call-imports.mdx +++ b/docs-fumadocs/content/docs/products/call-imports.mdx @@ -17,3 +17,30 @@ Call Imports lets you bulk-import production call recordings via CSV and run bat - Run metrics and insights across imported production calls Contact the EfficientAI team for a license. + +## Sharded evaluation pipeline (ops) + +If evaluation runs stall after recordings import (rows show `completed` import but no diarization/scoring): + +1. Restart API and `worker-imports` after deploying fixes. +2. **Retry evaluation** from the UI (use *Overwrite existing transcripts* if diarization previously failed). +3. Or clear stale dispatch locks on shard DBs: + +```sql +UPDATE call_import_evaluation_rows er +SET celery_task_id = NULL +FROM call_import_rows sr +WHERE er.call_import_row_id = sr.id + AND er.status = 'pending' + AND er.celery_task_id IS NOT NULL + AND sr.status = 'completed' + AND sr.recording_s3_key IS NOT NULL; +``` + +Abort and force-fail run on the `eval-control` queue (before `evaluations` on `worker-imports`) so they are not blocked behind large scoring backlogs. + +When running locally via `eai start` or `eai worker`, the imports worker must consume **`imports,diarization,eval-control,evaluations`**. If `eval-control` is missing, Run Evaluation will enqueue materialize tasks that never run and recording imports will not start. + +## Database sharding (enterprise scale) + +For large batches (10k+ rows), call-import row data can be spread across multiple PostgreSQL data shards with a catalog database for metadata and routing. See [Call Import Sharding](/docs/advanced/call-import-sharding/) for the full system design, configuration, and scaling guide. diff --git a/docs/operations/call-import-sharding-load-test.md b/docs/operations/call-import-sharding-load-test.md new file mode 100644 index 00000000..d279c638 --- /dev/null +++ b/docs/operations/call-import-sharding-load-test.md @@ -0,0 +1,21 @@ +# Call-import sharding load validation (Phase 11) + +Scenarios A–D from the internal Confluence runbook should be executed on AWS staging +(1 catalog + 5–6 row RDS instances) after enabling `database.sharding.enabled`. + +## Exit criteria + +- Row shard CPU ≤ 75% under scenario D (25k-row eval) +- Catalog CPU ≤ 50% +- No connection pool exhaustion (`pool_timeout` / PgBouncer queue depth stable) + +## Scenarios (summary) + +| ID | Description | +|----|-------------| +| A | Single import materialize + legacy import fetch | +| B | Unified eval pipeline, full metric set | +| C | Concurrent workspaces (fair dispatch) | +| D | 25k rows, max concurrency | + +Record results in the customer sign-off doc after GCP production sizing is confirmed. diff --git a/frontend/src/components/AIProviderModelPicker.tsx b/frontend/src/components/AIProviderModelPicker.tsx index 456c12f8..2b8b077d 100644 --- a/frontend/src/components/AIProviderModelPicker.tsx +++ b/frontend/src/components/AIProviderModelPicker.tsx @@ -30,6 +30,8 @@ type CredentialRow = { name: string | null source: 'aiprovider' | 'integration' gateway_model?: string | null + routing_mode?: string | null + effective_routing?: string | null } /** @@ -47,6 +49,7 @@ export default function AIProviderModelPicker({ onCredentialIdChange, llm_config, onLLMConfigChange, + onSelectionChange, disabled = false, size = 'md', showAdvancedOptions = true, @@ -59,6 +62,12 @@ export default function AIProviderModelPicker({ onCredentialIdChange?: (next: string) => void llm_config?: LLMGenerationConfig | null onLLMConfigChange?: (next: LLMGenerationConfig | null) => void + /** Fires once with the full selection when the credential dropdown changes. */ + onSelectionChange?: (next: { + provider: string + model: string + credentialId: string + }) => void disabled?: boolean size?: 'sm' | 'md' showAdvancedOptions?: boolean @@ -93,6 +102,8 @@ export default function AIProviderModelPicker({ name: p.name ?? null, source: 'aiprovider' as const, gateway_model: p.gateway_model ?? null, + routing_mode: p.routing_mode ?? null, + effective_routing: p.effective_routing ?? null, })) return [...aiRows, ...integrationRows] }, [aiProviders, integrations]) @@ -125,18 +136,38 @@ export default function AIProviderModelPicker({ ? aiProviders.find((p) => p.id === selectedCredential.id) : undefined + const gatewayCredential: AIProvider | undefined = useMemo(() => { + if (selectedAiProvider) return selectedAiProvider + if ( + selectedCredential?.source === 'aiprovider' && + selectedCredential.gateway_model?.trim() + ) { + return { + id: selectedCredential.id, + provider: selectedCredential.provider as AIProvider['provider'], + gateway_model: selectedCredential.gateway_model, + routing_mode: (selectedCredential.routing_mode as AIProvider['routing_mode']) ?? 'gateway', + effective_routing: selectedCredential.effective_routing as AIProvider['effective_routing'], + is_active: true, + created_at: '', + updated_at: '', + } + } + return undefined + }, [selectedAiProvider, selectedCredential]) + const resolvedProvider = selectedCredential?.provider || provider const { data: modelOptions } = useQuery({ queryKey: ['model-options', resolvedProvider], queryFn: () => apiClient.getModelOptions(resolvedProvider), - enabled: !!resolvedProvider, + enabled: !!resolvedProvider && !gatewayCredential?.gateway_model?.trim(), }) const rawCatalogModels: string[] = modelOptions?.llm ?? [] - const modelResolution = selectedAiProvider - ? resolveLLMModelsForCredential(selectedAiProvider, rawCatalogModels) + const modelResolution = gatewayCredential + ? resolveLLMModelsForCredential(gatewayCredential, rawCatalogModels) : { mode: 'catalog' as const, models: rawCatalogModels } const gatewayDirectModel = @@ -167,6 +198,10 @@ export default function AIProviderModelPicker({ const handleCredentialChange = (nextId: string) => { if (!nextId) { + if (onSelectionChange) { + onSelectionChange({ provider: '', model: '', credentialId: '' }) + return + } onCredentialIdChange?.('') onProviderChange('') onModelChange('') @@ -174,6 +209,14 @@ export default function AIProviderModelPicker({ } const row = activeCredentials.find((p) => p.id === nextId) if (!row) return + if (onSelectionChange) { + onSelectionChange({ + provider: row.provider, + model: '', + credentialId: row.id, + }) + return + } onCredentialIdChange?.(row.id) onProviderChange(row.provider) onModelChange('') diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index e4360b40..4f66da62 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -1803,9 +1803,17 @@ class ApiClient { async listCallImportEvaluationMetricClusterEligibleRows( callImportId: string, evaluationId: string, + options?: { + limit?: number + count_only?: boolean + }, ): Promise { + const params = new URLSearchParams() + if (options?.limit != null) params.set('limit', String(options.limit)) + if (options?.count_only) params.set('count_only', 'true') + const query = params.toString() const response = await this.client.get( - `/api/v1/call-imports/${callImportId}/evaluations/${evaluationId}/metric-clusters/eligible-rows`, + `/api/v1/call-imports/${callImportId}/evaluations/${evaluationId}/metric-clusters/eligible-rows${query ? `?${query}` : ''}`, ) return response.data } @@ -1843,6 +1851,7 @@ class ApiClient { credential_id?: string | null max_llm_calls?: number | null evaluation_row_ids?: string[] | null + row_limit?: number | null failure_policies?: Record< string, import('../types/api').MetricFailurePolicy @@ -1860,6 +1869,7 @@ class ApiClient { if (options?.evaluation_row_ids?.length) { body.evaluation_row_ids = options.evaluation_row_ids } + if (options?.row_limit != null) body.row_limit = options.row_limit if (options?.failure_policies) { body.failure_policies = options.failure_policies } diff --git a/frontend/src/lib/gatewayRouting.ts b/frontend/src/lib/gatewayRouting.ts index 0f85db7c..d753f650 100644 --- a/frontend/src/lib/gatewayRouting.ts +++ b/frontend/src/lib/gatewayRouting.ts @@ -33,8 +33,11 @@ export function resolveActiveAIProvider( providerKey: string, credentialId?: string | null, ): AIProvider | undefined { + const normalizedKey = providerKey.toLowerCase() const rows = aiProviders.filter( - (p) => p.is_active && p.provider.toLowerCase() === providerKey.toLowerCase(), + (p) => + p.is_active && + String(p.provider ?? '').toLowerCase() === normalizedKey, ) if (credentialId) { return rows.find((p) => p.id === credentialId) diff --git a/frontend/src/lib/llmModelOptions.ts b/frontend/src/lib/llmModelOptions.ts index 7e2ce6a2..d5a556e2 100644 --- a/frontend/src/lib/llmModelOptions.ts +++ b/frontend/src/lib/llmModelOptions.ts @@ -48,10 +48,8 @@ export function resolveLLMModelsForCredential( catalogModels: string[], ): LLMModelResolution { const gatewayModel = credential?.gateway_model?.trim() - if ( - credential?.provider?.toLowerCase() === 'custom' && - gatewayModel - ) { + const providerKey = String(credential?.provider ?? '').toLowerCase() + if (providerKey === 'custom' && gatewayModel) { return { mode: 'gateway_direct', model: gatewayModel } } if (credential && usesGatewayDirectModel(credential) && gatewayModel) { @@ -183,21 +181,62 @@ export interface LLMSelectionValue { credential_id?: string | null } -/** True when provider is set and a catalog or gateway model is resolved. */ -export function isLLMSelectionComplete( +const DEFAULT_LLM_MODELS: Record = { + openai: 'gpt-5-mini', + anthropic: 'claude-sonnet-4.6', + google: 'gemini-2.5-flash', + sarvam: 'sarvam-30b', +} + +function resolveCredentialForSelection( selection: LLMSelectionValue, aiProviders: AIProvider[], -): boolean { - if (!selection.provider) return false - if (selection.model?.trim()) return true - const credential = resolveActiveAIProvider( +): AIProvider | undefined { + if (selection.credential_id) { + const pinned = aiProviders.find( + (p) => p.is_active && p.id === selection.credential_id, + ) + if (pinned) return pinned + } + if (!selection.provider) return undefined + return resolveActiveAIProvider( aiProviders, selection.provider, selection.credential_id, ) - return ( - resolveLLMModelsForCredential(credential, []).mode === 'gateway_direct' +} + +/** True when provider is set and a catalog or gateway model is resolved. */ +export function isLLMSelectionComplete( + selection: LLMSelectionValue, + aiProviders: AIProvider[], +): boolean { + if (!selection.provider && !selection.credential_id) return false + if (selection.model?.trim()) return true + + const credential = resolveCredentialForSelection(selection, aiProviders) + if (!credential) return false + + const resolution = resolveLLMModelsForCredential(credential, []) + if (resolution.mode === 'gateway_direct') return true + + // Gateway-routed credentials without a pinned gateway_model still + // resolve at request time (same as partials / metrics surfaces). + if (routesViaGateway(credential)) return true + + return false +} + +/** True when the selection is half-filled (not empty default, not complete). */ +export function isLLMSelectionPartial( + selection: LLMSelectionValue, + aiProviders: AIProvider[], +): boolean { + const hasAnySelection = Boolean( + selection.provider || selection.model?.trim() || selection.credential_id, ) + if (!hasAnySelection) return false + return !isLLMSelectionComplete(selection, aiProviders) } /** Model string to send to APIs when gateway routing pins the model. */ @@ -206,12 +245,21 @@ export function resolveLLMModelForSubmit( aiProviders: AIProvider[], ): string | null { if (selection.model?.trim()) return selection.model.trim() - const credential = resolveActiveAIProvider( - aiProviders, - selection.provider ?? '', - selection.credential_id, - ) + + const credential = resolveCredentialForSelection(selection, aiProviders) + if (!credential) return null + const resolution = resolveLLMModelsForCredential(credential, []) if (resolution.mode === 'gateway_direct') return resolution.model + + if (routesViaGateway(credential)) { + const gatewayModel = credential.gateway_model?.trim() + if (gatewayModel) return gatewayModel + const providerKey = (credential.provider || selection.provider || '') + .toLowerCase() + .trim() + return DEFAULT_LLM_MODELS[providerKey] ?? 'gpt-5-mini' + } + return null } diff --git a/frontend/src/pages/callImports/CallImportDetail.tsx b/frontend/src/pages/callImports/CallImportDetail.tsx index 28f81ac9..41ae5a7c 100644 --- a/frontend/src/pages/callImports/CallImportDetail.tsx +++ b/frontend/src/pages/callImports/CallImportDetail.tsx @@ -60,8 +60,10 @@ import DiariseStatusPill from '../../components/callImports/DiariseStatusPill' import ProviderModelPicker, { type ProviderModelValue, } from '../../components/providers/ProviderModelPicker' +import AIProviderModelPicker from '../../components/AIProviderModelPicker' import { isLLMSelectionComplete, + isLLMSelectionPartial, resolveLLMModelForSubmit, } from '../../lib/llmModelOptions' import CallImportProgressBar from './components/CallImportProgressBar' @@ -578,6 +580,8 @@ export default function CallImportDetail() { const openRunEvaluationModal = useCallback(() => { setSelectedMetricIds([]) + setRunLLM({ provider: null, model: null, credential_id: null }) + void queryClient.invalidateQueries({ queryKey: ['ai-providers'] }) if (!evalDiariserLLM.provider) { setEvalDiariserLLM({ provider: 'openai', @@ -593,6 +597,7 @@ export default function CallImportDetail() { defaultDiarisationPrompt, evalDiariserLLM.provider, evalDiarisationPrompt, + queryClient, ]) const { data: aiProviders = [] } = useQuery({ @@ -1003,6 +1008,12 @@ export default function CallImportDetail() { onSuccess: (created) => { queryClient.invalidateQueries({ queryKey: ['call-import-evaluations', activeWorkspaceId, id] }) queryClient.invalidateQueries({ queryKey: ['call-import', activeWorkspaceId, id] }) + void queryClient.refetchQueries({ + queryKey: ['call-import-evaluations', activeWorkspaceId, id], + }) + void queryClient.refetchQueries({ + queryKey: ['call-import', activeWorkspaceId, id], + }) setShowRunEval(false) setSelectedMetricIds([]) setRunDraftName('') @@ -1028,6 +1039,9 @@ export default function CallImportDetail() { if (siblings.length > 0) { return } + void queryClient.refetchQueries({ + queryKey: ['call-import-evaluation', activeWorkspaceId, id, created.id], + }) // Land directly on the dedicated detail page for the new run. navigate(`/call-imports/${id}/evaluations/${created.id}`) }, @@ -4234,8 +4248,8 @@ export default function CallImportDetail() { {/* Run-level LLM config */} {(() => { const llmPartial = - Boolean(runLLM.provider) !== - Boolean(runLLM.model) + aiProviders.length > 0 && + isLLMSelectionPartial(runLLM, aiProviders) return (
- + setRunLLM((prev) => ({ + ...prev, + provider: next.provider || null, + model: next.model || null, + credential_id: next.credentialId || null, + })) + } + onProviderChange={(next) => + setRunLLM((prev) => ({ + ...prev, + provider: next || null, + })) + } + onCredentialIdChange={(next) => + setRunLLM((prev) => ({ + ...prev, + credential_id: next || null, + })) + } + onModelChange={(next) => + setRunLLM((prev) => ({ + ...prev, + model: next || null, + })) + } + llm_config={runLLM.llm_config ?? null} + onLLMConfigChange={(llm_config) => + setRunLLM((prev) => ({ ...prev, llm_config })) + } /> {llmPartial && (

- {runLLM.provider - ? 'Pick a model for this provider, or clear the provider to use the default.' + {runLLM.provider || runLLM.credential_id + ? 'Pick a complete LLM credential (including Custom gateway models), or clear the selection to use the default.' : 'Pick a provider for this model, or clear the model to use the default.'}

@@ -4291,6 +4333,34 @@ export default function CallImportDetail() { model: null, credential_id: null, } + const updateOverride = ( + patch: Partial, + ) => { + setMetricLLMOverrides((prev) => { + const copy = { ...prev } + const existing = + copy[target.id] || { + provider: null, + model: null, + credential_id: null, + } + const updated: ProviderModelValue = { + ...existing, + ...patch, + } + const cleared = + !updated.provider && + !updated.model?.trim() && + !updated.credential_id && + !updated.llm_config + if (cleared) { + delete copy[target.id] + } else { + copy[target.id] = updated + } + return copy + }) + } return (
)}

- { - setMetricLLMOverrides((prev) => { - const copy = { ...prev } - if (!next.provider && !next.model) { - delete copy[target.id] - } else { - copy[target.id] = next - } - return copy + + updateOverride({ + provider: next || null, + }) + } + onCredentialIdChange={(next) => + updateOverride({ + credential_id: next || null, }) - }} - defaultLabel="Use run default" + } + onModelChange={(next) => + updateOverride({ model: next || null }) + } + onLLMConfigChange={(llm_config) => + updateOverride({ llm_config }) + } + llm_config={override.llm_config ?? null} + size="sm" />
) @@ -4648,10 +4725,11 @@ export default function CallImportDetail() { ) } if ( - Boolean(runLLM.provider) !== Boolean(runLLM.model) + aiProviders.length > 0 && + isLLMSelectionPartial(runLLM, aiProviders) ) { disabledReasons.push( - 'Finish the Evaluation LLM selection (pick both a provider and a model, or clear both).', + 'Finish the Evaluation LLM selection (pick a provider and model, or a Custom gateway credential, or clear both to use the default).', ) } const isDisabled = @@ -4696,14 +4774,13 @@ export default function CallImportDetail() { } onClick={() => { // Build a clean overrides payload — drop any - // entries that didn't end up with both a - // provider and a model set so the API doesn't - // 400 on partial fills. We also discard - // entries for ids that are no longer in - // ``overrideTargets`` (e.g., the user set an - // override for a parent then deselected every - // one of its labels) so stale state doesn't - // trip backend validation. + // entries that didn't end up with a complete + // LLM selection so the API doesn't 400 on + // partial fills. We also discard entries for + // ids that are no longer in ``overrideTargets`` + // (e.g., the user set an override for a parent + // then deselected every one of its labels) so + // stale state doesn't trip backend validation. const overrides: Record< string, CallImportEvaluationLLMOverride @@ -4712,10 +4789,12 @@ export default function CallImportDetail() { metricLLMOverrides, )) { if (!overrideTargetIds.has(mid)) continue - if (val.provider && val.model) { + if (isLLMSelectionComplete(val, aiProviders)) { overrides[mid] = { - provider: val.provider, - model: val.model, + provider: val.provider!, + model: + resolveLLMModelForSubmit(val, aiProviders) ?? + val.model!, credential_id: val.credential_id || null, llm_config: val.llm_config || null, } @@ -4725,16 +4804,30 @@ export default function CallImportDetail() { } } } + const runLLMComplete = isLLMSelectionComplete( + runLLM, + aiProviders, + ) runEvaluationMutation.mutate({ metric_ids: selectedMetricIds, name: runDraftName.trim() || null, // Diarised is the only supported source // now; the backend rejects anything else. transcript_sources: ['diarised'], - llm_provider: runLLM.provider || null, - llm_model: runLLM.model || null, - llm_credential_id: runLLM.credential_id || null, - llm_config: runLLM.llm_config || null, + llm_provider: runLLMComplete + ? runLLM.provider || null + : null, + llm_model: runLLMComplete + ? resolveLLMModelForSubmit(runLLM, aiProviders) ?? + runLLM.model ?? + null + : null, + llm_credential_id: runLLMComplete + ? runLLM.credential_id || null + : null, + llm_config: runLLMComplete + ? runLLM.llm_config || null + : null, metric_llm_overrides: Object.keys(overrides).length ? overrides : null, diff --git a/frontend/src/pages/callImports/CallImportEvaluationDetail.tsx b/frontend/src/pages/callImports/CallImportEvaluationDetail.tsx index 74f0509d..a9f57c43 100644 --- a/frontend/src/pages/callImports/CallImportEvaluationDetail.tsx +++ b/frontend/src/pages/callImports/CallImportEvaluationDetail.tsx @@ -78,7 +78,6 @@ import type { MetricClustersRcaSummary, MetricFailurePolicy, MetricFailurePolicyMetricPreview, - MetricClusterEligibleRow, EvaluationUserInsightsState, EvaluationUserInsightItem, MetricPeriodDelta, @@ -259,6 +258,10 @@ function isUserInsightMetricName(name: string): boolean { } const ROWS_PAGE_SIZE = 50 +const EVAL_PROGRESS_POLL_MS = 3000 +const EVAL_PROGRESS_POLL_LARGE_MS = 10000 +const EVAL_LARGE_ROW_THRESHOLD = 5000 +const ROWS_REFETCH_WHILE_RUNNING_MS = 20000 const USER_INSIGHTS_SAMPLE_SIZE_OPTIONS = [50, 100, 150, 200, 300, 500] as const const DEFAULT_USER_INSIGHTS_SAMPLE_SIZE = 200 @@ -407,6 +410,9 @@ export default function CallImportEvaluationDetail() { queryClient.invalidateQueries({ queryKey: ['call-import-evaluations', activeWorkspaceId, id], }) + queryClient.invalidateQueries({ + queryKey: ['call-import', activeWorkspaceId, id], + }) } const bulkOperationConflictMessage = @@ -778,20 +784,29 @@ export default function CallImportEvaluationDetail() { queryKey: ['call-import', activeWorkspaceId, id], queryFn: () => apiClient.getCallImport(id!, { row_limit: 0, row_offset: 0 }), enabled: !!id, - // Poll while diarisation is in flight so the per-run "Diarising - // audio…" progress bar updates as the upstream transcribe / diarise - // worker churns through this batch's rows. Stops polling once - // everything settles to terminal states. + // Poll while import, diarisation, or the linked evaluation is in flight + // so all three summary bars update without a hard refresh. Mirrors the + // broader conditions on CallImportDetail plus active eval-run status. refetchInterval: (q) => { const ci = q.state.data as | { + status?: string diarised_pending_rows?: number diarised_running_rows?: number } | undefined + if (ci?.status === 'deleting') return 3000 + if (ci?.status === 'pending' || ci?.status === 'processing') return 5000 const inFlight = (ci?.diarised_pending_rows ?? 0) + (ci?.diarised_running_rows ?? 0) - return inFlight > 0 ? 4000 : false + if (inFlight > 0) return 4000 + const evaluation = queryClient.getQueryData(evaluationQueryKey) as + | { status?: string; bulk_operation?: unknown } + | undefined + if (evaluation?.bulk_operation) return EVALUATION_BULK_OPERATION_POLL_MS + const evalStatus = evaluation?.status + if (evalStatus === 'pending' || evalStatus === 'running') return 4000 + return false }, }) @@ -799,11 +814,16 @@ export default function CallImportEvaluationDetail() { queryKey: evaluationQueryKey, queryFn: () => apiClient.getCallImportEvaluation(id!, evalId!), enabled: !!id && !!evalId, + staleTime: 5000, refetchInterval: (q) => { const data = q.state.data if (data?.bulk_operation) return EVALUATION_BULK_OPERATION_POLL_MS const status = data?.status - return status === 'pending' || status === 'running' ? 3000 : false + if (status !== 'pending' && status !== 'running') return false + const totalRows = data?.total_rows ?? 0 + return totalRows > EVAL_LARGE_ROW_THRESHOLD + ? EVAL_PROGRESS_POLL_LARGE_MS + : EVAL_PROGRESS_POLL_MS }, }) @@ -849,13 +869,22 @@ export default function CallImportEvaluationDetail() { sort_by: sortBy || undefined, sort_dir: sortBy ? sortDir : undefined, }), - enabled: !!id && !!evalId, + enabled: + !!id && + !!evalId && + resultsTab === 'table' && + (evaluationQuery.data?.status === 'running' || + evaluationQuery.data?.status === 'completed' || + evaluationQuery.data?.status === 'partial' || + evaluationQuery.data?.status === 'failed'), refetchInterval: () => { if (evaluationQuery.data?.bulk_operation) { return EVALUATION_BULK_OPERATION_POLL_MS } - const status = evaluationQuery.data?.status - return status === 'pending' || status === 'running' ? 3000 : false + if (evaluationQuery.data?.status === 'running') { + return ROWS_REFETCH_WHILE_RUNNING_MS + } + return false }, }) @@ -882,24 +911,6 @@ export default function CallImportEvaluationDetail() { deepLinkRowId, ]) - const pendingRowsQuery = useQuery({ - queryKey: ['call-import-evaluation-pending-rows-count', activeWorkspaceId, id, evalId], - queryFn: () => - apiClient.listCallImportEvaluationRows(id!, evalId!, { - page: 1, - page_size: 1, - status: 'pending', - }), - enabled: !!id && !!evalId, - refetchInterval: () => { - if (evaluationQuery.data?.bulk_operation) { - return EVALUATION_BULK_OPERATION_POLL_MS - } - const status = evaluationQuery.data?.status - return status === 'pending' || status === 'running' ? 3000 : false - }, - }) - // Lazy: only fetch aggregates when the user lands on the // visualizations tab. Refetches while the run is still in flight so // the chart fills in as workers complete rows. @@ -917,10 +928,7 @@ export default function CallImportEvaluationDetail() { evalId!, vizBaselineEvaluationId, ), - enabled: - !!id && - !!evalId && - (resultsTab === 'visualizations' || resultsTab === 'table'), + enabled: !!id && !!evalId && resultsTab === 'visualizations', refetchInterval: () => { const status = evaluationQuery.data?.status return status === 'pending' || status === 'running' ? 5000 : false @@ -2199,7 +2207,12 @@ export default function CallImportEvaluationDetail() { : `Evaluation ${evaluation.id.slice(0, 8)}` const bulkOperation = evaluation.bulk_operation ?? null const bulkOperationActive = bulkOperation !== null - const pendingRowCount = pendingRowsQuery.data?.total ?? 0 + const pendingRowCount = Math.max( + 0, + (evaluation.total_rows ?? 0) - + (evaluation.completed_rows ?? 0) - + (evaluation.failed_rows ?? 0), + ) const getMetricLlmLabel = (metricId: string): string => { const override = evaluation.metric_llm_overrides?.[metricId] const overrideProvider = override?.provider?.trim() @@ -8432,47 +8445,36 @@ function UserInsightsStatusBanner({ ) } -const METRIC_CLUSTER_ROW_PRESETS = [25, 50, 100, 200] as const +const METRIC_CLUSTER_ROW_PRESETS = [25, 50, 500] as const +type MetricClusterRowPreset = + (typeof METRIC_CLUSTER_ROW_PRESETS)[number] | 'all' + +function metricClusterSelectedCount( + totalEligible: number, + preset: MetricClusterRowPreset, +): number { + if (totalEligible <= 0) return 0 + if (preset === 'all') return totalEligible + return Math.min(preset, totalEligible) +} function MetricClusterRowPicker({ - rows, - selectedIds, - onChangeSelectedIds, + totalEligible, + preset, + onChangePreset, disabled, }: { - rows: MetricClusterEligibleRow[] - selectedIds: Set - onChangeSelectedIds: (next: Set) => void + totalEligible: number + preset: MetricClusterRowPreset + onChangePreset: (next: MetricClusterRowPreset) => void disabled?: boolean }) { - const selectFirstN = (n: number) => { - onChangeSelectedIds( - new Set(rows.slice(0, n).map((r) => r.evaluation_row_id)), - ) - } + const selectedCount = metricClusterSelectedCount(totalEligible, preset) - const selectAll = () => { - onChangeSelectedIds(new Set(rows.map((r) => r.evaluation_row_id))) - } - - const presetActive = (n: number) => { - const limit = Math.min(n, rows.length) - if (limit === 0 || selectedIds.size !== limit) return false - const firstIds = rows.slice(0, limit).map((r) => r.evaluation_row_id) - return firstIds.every((id) => selectedIds.has(id)) - } + const presetActive = (n: number) => + preset !== 'all' && preset === n && selectedCount === n - const allActive = - rows.length > 0 && - selectedIds.size === rows.length && - rows.every((r) => selectedIds.has(r.evaluation_row_id)) - - const toggleRow = (id: string) => { - const next = new Set(selectedIds) - if (next.has(id)) next.delete(id) - else next.add(id) - onChangeSelectedIds(next) - } + const allActive = preset === 'all' && totalEligible > 0 const presetButtonClass = (active: boolean) => 'rounded-full px-2 py-0.5 border text-[10px] font-medium transition-colors disabled:opacity-40 ' + @@ -8482,98 +8484,40 @@ function MetricClusterRowPicker({ return (
-
+

- Calls to include ({selectedIds.size} / {rows.length}) + Calls to include ({selectedCount} / {totalEligible} eligible)

-
- - | - -
- {rows.length > 0 ? ( + {totalEligible > 0 ? (
- Quick: - {METRIC_CLUSTER_ROW_PRESETS.filter((n) => n <= rows.length).map( - (n) => ( - - ), - )} - {rows.length > METRIC_CLUSTER_ROW_PRESETS[METRIC_CLUSTER_ROW_PRESETS.length - 1] ? ( + {METRIC_CLUSTER_ROW_PRESETS.map((n) => ( - ) : null} + ))} +
- ) : null} + ) : ( +

+ No completed calls with a flagged quality metric yet. +

+ )}
- {rows.length === 0 ? ( -

- No completed calls with a flagged quality metric yet. -

- ) : ( -
    - {rows.map((row) => { - const id = row.evaluation_row_id - const label = - row.conversation_id?.trim() || - (row.row_index != null ? `Row ${row.row_index}` : id.slice(0, 8)) - const metrics = row.flagged_metric_names.join(', ') - return ( -
  • - -
    • - ) - })} -
- )}
) } @@ -8811,8 +8755,7 @@ function MetricClusterGenerationModal({ const [error, setError] = useState(null) const [pickerProvider, setPickerProvider] = useState('') const [pickerModel, setPickerModel] = useState('') - const [selectedRowIds, setSelectedRowIds] = useState>(new Set()) - const [selectionTouched, setSelectionTouched] = useState(false) + const [rowPreset, setRowPreset] = useState(25) const [llmPickerTouched, setLlmPickerTouched] = useState(false) const [policies, setPolicies] = useState>( {}, @@ -8844,25 +8787,27 @@ function MetricClusterGenerationModal({ getActiveWorkspaceId(), callImportId, evaluationId, - policiesSource, - JSON.stringify(policies), ], queryFn: () => apiClient.listCallImportEvaluationMetricClusterEligibleRows( callImportId, evaluationId, + { count_only: true }, ), enabled: open && !!callImportId && !!evaluationId, staleTime: 30_000, }) - const eligibleRows = eligibleRowsQuery.data?.items ?? [] + const totalEligible = eligibleRowsQuery.data?.total ?? 0 + const selectedRowCount = metricClusterSelectedCount(totalEligible, rowPreset) + const hasExistingClusters = !!state?.groups?.length useEffect(() => { if (!open) return setError(null) onError?.(null) + setRowPreset(25) }, [open]) useEffect(() => { @@ -8916,33 +8861,13 @@ function MetricClusterGenerationModal({ llmPickerTouched, ]) - useEffect(() => { - if (!open || selectionTouched || eligibleRows.length === 0) return - const fromState = state?.selected_evaluation_row_ids - if (fromState?.length) { - const valid = fromState.filter((id) => - eligibleRows.some((r) => r.evaluation_row_id === id), - ) - if (valid.length) { - setSelectedRowIds(new Set(valid)) - return - } - } - setSelectedRowIds(new Set(eligibleRows.map((r) => r.evaluation_row_id))) - }, [ - open, - eligibleRows, - selectionTouched, - state?.selected_evaluation_row_ids, - ]) - const reportError = (message: string | null) => { setError(message) onError?.(message) } const handleGenerate = async () => { - if (selectedRowIds.size === 0) { + if (selectedRowCount === 0) { reportError('Select at least one call to cluster.') return } @@ -8964,9 +8889,6 @@ function MetricClusterGenerationModal({ reportError(null) try { const force = hasExistingClusters - const allSelected = - eligibleRows.length > 0 && - selectedRowIds.size === eligibleRows.length await apiClient.generateCallImportEvaluationMetricClusters( callImportId, evaluationId, @@ -8975,9 +8897,7 @@ function MetricClusterGenerationModal({ regenerate: force, provider: pickerProvider || undefined, model: pickerModel || undefined, - evaluation_row_ids: allSelected - ? undefined - : Array.from(selectedRowIds), + row_limit: rowPreset === 'all' ? undefined : rowPreset, failure_policies: policies, }, ) @@ -9048,12 +8968,9 @@ function MetricClusterGenerationModal({

Loading eligible calls…

) : ( { - setSelectionTouched(true) - setSelectedRowIds(next) - }} + totalEligible={totalEligible} + preset={rowPreset} + onChangePreset={setRowPreset} disabled={generating} /> )} @@ -9092,7 +9009,7 @@ function MetricClusterGenerationModal({ variant="primary" onClick={handleGenerate} isLoading={generating} - disabled={generating || selectedRowIds.size === 0} + disabled={generating || selectedRowCount === 0} > Generate clusters diff --git a/frontend/src/pages/callImports/components/CallImportProgressBar.tsx b/frontend/src/pages/callImports/components/CallImportProgressBar.tsx index f018d9d7..3d476a65 100644 --- a/frontend/src/pages/callImports/components/CallImportProgressBar.tsx +++ b/frontend/src/pages/callImports/components/CallImportProgressBar.tsx @@ -32,15 +32,17 @@ export default function CallImportProgressBar({ } const safeTotal = Math.max(total, 0) - const completedPct = safeTotal > 0 ? (completed / safeTotal) * 100 : 0 - const failedPct = safeTotal > 0 ? (failed / safeTotal) * 100 : 0 + const safeCompleted = safeTotal > 0 ? Math.min(completed, safeTotal) : Math.max(completed, 0) + const safeFailed = safeTotal > 0 ? Math.min(failed, Math.max(safeTotal - safeCompleted, 0)) : Math.max(failed, 0) + const completedPct = safeTotal > 0 ? (safeCompleted / safeTotal) * 100 : 0 + const failedPct = safeTotal > 0 ? (safeFailed / safeTotal) * 100 : 0 return (
@@ -55,11 +57,11 @@ export default function CallImportProgressBar({
{showLabel && (
- {completed} + {safeCompleted} / {safeTotal} - {failed > 0 && ( - {failed} failed + {safeFailed > 0 && ( + {safeFailed} failed )}
)} diff --git a/frontend/src/pages/configurations/Integrations.tsx b/frontend/src/pages/configurations/Integrations.tsx index d2e4e0b8..e462ed0c 100644 --- a/frontend/src/pages/configurations/Integrations.tsx +++ b/frontend/src/pages/configurations/Integrations.tsx @@ -170,7 +170,7 @@ export default function Integrations() { }) const { data: aiproviders = [] } = useQuery({ - queryKey: ['aiproviders'], + queryKey: ['ai-providers'], queryFn: () => apiClient.listAIProviders(), }) @@ -268,19 +268,19 @@ export default function Integrations() { const createAIProviderMutation = useMutation({ mutationFn: (data: AIProviderCreate) => apiClient.createAIProvider(data), - onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['aiproviders'] }); showToast('AI Provider configured successfully!', 'success'); resetForm() }, + onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['ai-providers'] }); showToast('AI Provider configured successfully!', 'success'); resetForm() }, onError: (error: any) => { showToast(`Failed to configure provider: ${error.response?.data?.detail || error.message}`, 'error') }, }) const updateAIProviderMutation = useMutation({ mutationFn: ({ id, data }: { id: string; data: Partial }) => apiClient.updateAIProvider(id, data), - onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['aiproviders'] }); showToast('AI Provider updated successfully!', 'success'); resetForm() }, + onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['ai-providers'] }); showToast('AI Provider updated successfully!', 'success'); resetForm() }, onError: (error: any) => { showToast(`Failed to update provider: ${error.response?.data?.detail || error.message}`, 'error') }, }) const deleteAIProviderMutation = useMutation({ mutationFn: (id: string) => apiClient.deleteAIProvider(id), - onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['aiproviders'] }); showToast('AI Provider deleted successfully!', 'success'); setShowDeleteAIProviderModal(false); setAIProviderToDelete(null) }, + onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['ai-providers'] }); showToast('AI Provider deleted successfully!', 'success'); setShowDeleteAIProviderModal(false); setAIProviderToDelete(null) }, onError: (error: any) => { showToast(`Failed to delete provider: ${error.response?.data?.detail || error.message}`, 'error') }, }) @@ -325,7 +325,7 @@ export default function Integrations() { const setDefaultAIProviderMutation = useMutation({ mutationFn: (id: string) => apiClient.setDefaultAIProvider(id), - onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['aiproviders'] }); showToast('Default AI provider updated', 'success') }, + onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['ai-providers'] }); showToast('Default AI provider updated', 'success') }, onError: (error: any) => { showToast(error?.response?.data?.detail || error?.message || 'Failed to set default', 'error') }, }) diff --git a/schema_er_diagram.png b/schema_er_diagram.png index 8d0c2e14..3d8ae608 100644 Binary files a/schema_er_diagram.png and b/schema_er_diagram.png differ diff --git a/scripts/backfill_catalog_rows_to_shards.py b/scripts/backfill_catalog_rows_to_shards.py new file mode 100644 index 00000000..20743baa --- /dev/null +++ b/scripts/backfill_catalog_rows_to_shards.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +""" +Copy call_import_rows (and related eval rows) from catalog DB into data shards. + +Use when sharding was enabled but historical rows still live on the catalog +database (e.g. efficientai used as catalog_url after monolith dev). + +Dry-run by default; pass --apply to insert on shards and register slices. +Does not delete catalog copies until you verify (delete manually if desired). +""" + +from __future__ import annotations + +import argparse +import uuid +from typing import List + +from sqlalchemy.orm import Session + +from app.config import load_config_from_file +from app.database import SessionLocal +from app.db_sharding.pool_manager import db_pool_manager +from app.db_sharding.registry import load_slice_registry_for_import +from app.db_sharding.row_ops import ( + partition_mappings_by_shard, + register_shard_slices, + _reset_shard_write_role, + _shard_write_without_catalog_fks, +) +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import ( + CallImport, + CallImportEvaluationRow, + CallImportRow, +) + + +def _row_to_mapping(row: CallImportRow) -> dict: + cols = {c.name for c in CallImportRow.__table__.columns} + return {name: getattr(row, name) for name in cols} + + +def _eval_row_to_mapping(row: CallImportEvaluationRow) -> dict: + cols = {c.name for c in CallImportEvaluationRow.__table__.columns} + return {name: getattr(row, name) for name in cols} + + +def backfill_import( + catalog_db: Session, + call_import_id: uuid.UUID, + *, + apply: bool, +) -> dict: + rows: List[CallImportRow] = ( + catalog_db.query(CallImportRow) + .filter(CallImportRow.call_import_id == call_import_id) + .order_by(CallImportRow.row_index.asc()) + .all() + ) + if not rows: + return {"import_id": str(call_import_id), "catalog_rows": 0, "copied": 0} + + mappings = [_row_to_mapping(r) for r in rows] + buckets = partition_mappings_by_shard(catalog_db, call_import_id, mappings) + + eval_rows = ( + catalog_db.query(CallImportEvaluationRow) + .join(CallImportRow, CallImportRow.id == CallImportEvaluationRow.call_import_row_id) + .filter(CallImportRow.call_import_id == call_import_id) + .all() + ) + eval_by_shard: dict[str, list] = {} + for er in eval_rows: + source = next((r for r in rows if r.id == er.call_import_row_id), None) + if source is None: + continue + for sid in partition_mappings_by_shard( + catalog_db, call_import_id, [_row_to_mapping(source)] + ): + if sid == "legacy": + continue + eval_by_shard.setdefault(sid, []).append(_eval_row_to_mapping(er)) + break + + plan = {sid: len(ms) for sid, ms in buckets.items()} + if not apply: + return { + "import_id": str(call_import_id), + "catalog_rows": len(rows), + "plan_by_shard": plan, + "eval_rows": len(eval_rows), + "dry_run": True, + } + + router = db_pool_manager.router + assert router is not None + for shard_id, shard_mappings in buckets.items(): + if shard_id == "legacy": + continue + factory = db_pool_manager.shard_session_factory(shard_id) + shard_db = factory() + try: + _shard_write_without_catalog_fks(shard_db) + for m in shard_mappings: + shard_db.merge(CallImportRow(**m)) + for em in eval_by_shard.get(shard_id, []): + shard_db.merge(CallImportEvaluationRow(**em)) + shard_db.commit() + _reset_shard_write_role(shard_db) + except Exception: + shard_db.rollback() + try: + _reset_shard_write_role(shard_db) + except Exception: + pass + raise + finally: + shard_db.close() + + register_shard_slices(catalog_db, call_import_id, len(rows)) + catalog_db.commit() + return { + "import_id": str(call_import_id), + "catalog_rows": len(rows), + "plan_by_shard": plan, + "copied": len(rows), + "dry_run": False, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description="Backfill catalog rows into data shards") + parser.add_argument("--config", default="config.yml") + parser.add_argument("--call-import-id", type=uuid.UUID, default=None) + parser.add_argument("--all", action="store_true", help="Every import with catalog rows") + parser.add_argument("--apply", action="store_true") + args = parser.parse_args() + + load_config_from_file(args.config) + if not is_sharding_enabled(): + print("Enable database.sharding in config first.") + return 1 + + db = SessionLocal() + try: + if args.call_import_id: + ids = [args.call_import_id] + elif args.all: + ids = [ + row[0] + for row in db.query(CallImportRow.call_import_id).distinct().all() + if row[0] is not None + ] + else: + print("Pass --call-import-id UUID or --all") + return 1 + + for cid in ids: + if load_slice_registry_for_import(db, cid): + print(f"Skip {cid}: shard_slices already registered") + continue + result = backfill_import(db, cid, apply=args.apply) + print(result) + return 0 + finally: + db.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/rebalance_call_import_shards.py b/scripts/rebalance_call_import_shards.py new file mode 100644 index 00000000..85e03a06 --- /dev/null +++ b/scripts/rebalance_call_import_shards.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Dry-run / execute rebalance of call-import row slices between shards.""" + +from __future__ import annotations + +import argparse +import uuid + +from app.database import SessionLocal +from app.db_sharding.registry import load_slice_registry_for_import +from app.db_sharding.sessions import is_sharding_enabled +from app.models.database import CallImport, CallImportShardSlice + + +def main() -> int: + parser = argparse.ArgumentParser(description="Rebalance call-import shard slices") + parser.add_argument("call_import_id", type=uuid.UUID) + parser.add_argument("--target-shard", required=True, help="Destination shard id") + parser.add_argument("--dry-run", action="store_true", default=True) + parser.add_argument("--apply", action="store_true", help="Persist registry updates") + args = parser.parse_args() + + if not is_sharding_enabled(): + print("Sharding is disabled; nothing to rebalance.") + return 1 + + db = SessionLocal() + try: + call_import = db.query(CallImport).filter(CallImport.id == args.call_import_id).first() + if call_import is None: + print("Call import not found.") + return 1 + registry = load_slice_registry_for_import(db, args.call_import_id) + slices = ( + db.query(CallImportShardSlice) + .filter(CallImportShardSlice.call_import_id == args.call_import_id) + .order_by(CallImportShardSlice.slice_id.asc()) + .all() + ) + print(f"Import {args.call_import_id}: {len(slices)} slice(s), registry keys={len(registry)}") + for sl in slices: + print( + f" slice {sl.slice_id}: rows {sl.row_index_min}-{sl.row_index_max} " + f"shard {sl.shard_id} -> {args.target_shard if args.apply else '(dry-run)'}" + ) + if args.apply and not args.dry_run: + sl.shard_id = args.target_shard + if args.apply and not args.dry_run: + db.commit() + print("Registry updated. Run row data copy separately before resuming import.") + else: + print("Dry run only. Pass --apply without --dry-run to update registry.") + return 0 + finally: + db.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/conftest.py b/tests/conftest.py index 28a0e135..f33deb1e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,6 +20,47 @@ # Ensure storage service singletons can initialize in test environments. os.environ["UPLOAD_DIR"] = "/tmp/efficientai-test-uploads" +_TASKS_PACKAGE_DIR = str( + Path(__file__).resolve().parents[1] / "app" / "workers" / "tasks" +) + + +@pytest.fixture(autouse=True) +def disable_db_sharding_for_tests(monkeypatch, request): + """Tests use one SQLAlchemy session; ignore production shard routing.""" + if request.node.get_closest_marker("integration"): + yield + return + + from app.config import settings + from app.db_sharding.pool_manager import db_pool_manager + + monkeypatch.setattr(settings, "DB_SHARDING_ENABLED", False) + db_pool_manager.reset() + yield + + +@pytest.fixture(autouse=True) +def ensure_workers_tasks_package(): + """Keep ``app.workers.tasks`` importable without eager Celery imports.""" + import importlib + + workers_pkg = importlib.import_module("app.workers") + tasks_pkg = sys.modules.get("app.workers.tasks") + if tasks_pkg is None: + tasks_pkg = types.ModuleType("app.workers.tasks") + sys.modules["app.workers.tasks"] = tasks_pkg + tasks_pkg.__path__ = [_TASKS_PACKAGE_DIR] + workers_pkg.tasks = tasks_pkg + + helpers_pkg = sys.modules.get("app.workers.tasks.helpers") + if helpers_pkg is None: + helpers_pkg = types.ModuleType("app.workers.tasks.helpers") + sys.modules["app.workers.tasks.helpers"] = helpers_pkg + helpers_pkg.__path__ = [ + str(Path(__file__).resolve().parents[1] / "app" / "workers" / "tasks" / "helpers") + ] + @pytest.fixture def org_id(): @@ -295,8 +336,8 @@ def update_threshold_defaults(self, *_args, **_kwargs): fake_workers_tasks_pkg = sys.modules.get("app.workers.tasks") if fake_workers_tasks_pkg is None: fake_workers_tasks_pkg = types.ModuleType("app.workers.tasks") - fake_workers_tasks_pkg.__path__ = [] sys.modules["app.workers.tasks"] = fake_workers_tasks_pkg + fake_workers_tasks_pkg.__path__ = [_TASKS_PACKAGE_DIR] # ``app.workers.tasks`` is stubbed with an empty ``__path__`` so Celery # task modules are not eagerly imported, but several API routes and tests diff --git a/tests/test_api/test_call_import_diarization_and_eval_llm.py b/tests/test_api/test_call_import_diarization_and_eval_llm.py index 6899be44..5bad31ef 100644 --- a/tests/test_api/test_call_import_diarization_and_eval_llm.py +++ b/tests/test_api/test_call_import_diarization_and_eval_llm.py @@ -57,16 +57,18 @@ def _make_fake_row(): def _patch_session(monkeypatch, task_module, fake_row): - fake_query = MagicMock() - fake_query.filter.return_value = fake_query - fake_query.first.return_value = fake_row - fake_db = SimpleNamespace( - query=lambda *_a, **_kw: fake_query, commit=lambda: None, close=lambda: None, + flush=lambda: None, ) - monkeypatch.setattr(task_module, "SessionLocal", lambda: fake_db) + + def _locate(_row_id): + return fake_db, fake_db, fake_row, "legacy" + + monkeypatch.setattr("app.db_sharding.row_ops.locate_call_import_row", _locate) + monkeypatch.setattr("app.db_sharding.row_ops.close_row_sessions", lambda *_a: None) + monkeypatch.setattr("app.database.SessionLocal", lambda: fake_db) return fake_db @@ -554,14 +556,7 @@ def test_select_rows_for_transcription_skips_rows_with_existing_transcripts(): ), ] - fake_query = MagicMock() - fake_query.options.return_value = fake_query - fake_query.filter.return_value = fake_query - fake_query.order_by.return_value = SimpleNamespace( - all=lambda: fake_rows - ) - - fake_db = SimpleNamespace(query=lambda *_: fake_query) + fake_db = SimpleNamespace() payload = CallImportTranscribeRequest( stt_provider="deepgram", @@ -571,9 +566,13 @@ def test_select_rows_for_transcription_skips_rows_with_existing_transcripts(): only_missing=True, overwrite_existing=False, ) - selected, skip_counts = _select_rows_for_transcription( - fake_db, call_import, payload - ) + with patch( + "app.db_sharding.scatter_gather.load_call_import_rows_for_transcription", + return_value=fake_rows, + ): + selected, skip_counts = _select_rows_for_transcription( + fake_db, call_import, payload + ) # Two rows selected: the bare row, and the one with only a production # transcript. @@ -599,13 +598,7 @@ def test_select_rows_for_transcription_overwrite_replaces_existing(): row_index=0, ), ] - fake_query = MagicMock() - fake_query.options.return_value = fake_query - fake_query.filter.return_value = fake_query - fake_query.order_by.return_value = SimpleNamespace( - all=lambda: fake_rows - ) - fake_db = SimpleNamespace(query=lambda *_: fake_query) + fake_db = SimpleNamespace() payload = CallImportTranscribeRequest( stt_provider="deepgram", @@ -615,9 +608,13 @@ def test_select_rows_for_transcription_overwrite_replaces_existing(): only_missing=True, overwrite_existing=True, ) - selected, skip_counts = _select_rows_for_transcription( - fake_db, call_import, payload - ) + with patch( + "app.db_sharding.scatter_gather.load_call_import_rows_for_transcription", + return_value=fake_rows, + ): + selected, skip_counts = _select_rows_for_transcription( + fake_db, call_import, payload + ) assert len(selected) == 1 assert skip_counts == {} @@ -630,11 +627,7 @@ def test_select_rows_for_transcription_raises_on_unknown_row_id(): call_import = SimpleNamespace(id=uuid4()) requested_id = uuid4() - fake_query = MagicMock() - fake_query.options.return_value = fake_query - fake_query.filter.return_value = fake_query - fake_query.order_by.return_value = SimpleNamespace(all=lambda: []) - fake_db = SimpleNamespace(query=lambda *_: fake_query) + fake_db = SimpleNamespace() payload = CallImportTranscribeRequest( stt_provider="deepgram", @@ -643,10 +636,14 @@ def test_select_rows_for_transcription_raises_on_unknown_row_id(): diarization_llm_model="gpt-4o-mini", ) - with pytest.raises(HTTPException) as exc: - _select_rows_for_transcription( - fake_db, call_import, payload, requested_row_ids=[requested_id] - ) + with patch( + "app.db_sharding.scatter_gather.load_call_import_rows_for_transcription", + return_value=[], + ): + with pytest.raises(HTTPException) as exc: + _select_rows_for_transcription( + fake_db, call_import, payload, requested_row_ids=[requested_id] + ) assert exc.value.status_code == 400 assert str(requested_id) in exc.value.detail diff --git a/tests/test_api/test_call_import_evaluation_serialize_progress.py b/tests/test_api/test_call_import_evaluation_serialize_progress.py new file mode 100644 index 00000000..f6c42aad --- /dev/null +++ b/tests/test_api/test_call_import_evaluation_serialize_progress.py @@ -0,0 +1,60 @@ +"""Tests for evaluation serialize progress counter merge.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import patch +from uuid import uuid4 + +from app.api.v1.routes.call_import_evaluations import _serialize_eval + + +def test_serialize_eval_merges_redis_deltas_without_clearing(db_session, org_id, seed_org): + now = datetime.now(timezone.utc) + evaluation = type( + "EvalStub", + (), + { + "id": uuid4(), + "call_import_id": uuid4(), + "organization_id": org_id, + "name": "run", + "selected_metric_ids": [], + "selected_metric_groups": None, + "status": "running", + "total_rows": 10, + "completed_rows": 3, + "failed_rows": 1, + "error_message": None, + "llm_provider": None, + "llm_model": None, + "llm_credential_id": None, + "llm_config": None, + "metric_llm_overrides": None, + "stt_provider": None, + "stt_model": None, + "stt_credential_id": None, + "transcript_source": "diarised", + "tldr_summary": None, + "started_at": None, + "finished_at": None, + "created_at": now, + "updated_at": now, + }, + )() + + with patch( + "app.services.call_imports.progress_counters.merge_eval_counters_for_ui", + return_value=(5, 2), + ) as mock_merge, patch( + "app.services.call_imports.progress_counters.clear_eval_progress_redis", + ) as mock_clear, patch( + "app.api.v1.routes.call_import_evaluations._metrics_for_ids", + return_value=[], + ): + response = _serialize_eval(db_session, evaluation) + + mock_merge.assert_called_once_with(evaluation) + mock_clear.assert_not_called() + assert response.completed_rows == 5 + assert response.failed_rows == 2 diff --git a/tests/test_api/test_call_import_metric_clusters_rows.py b/tests/test_api/test_call_import_metric_clusters_rows.py index b0404769..a1d1bd31 100644 --- a/tests/test_api/test_call_import_metric_clusters_rows.py +++ b/tests/test_api/test_call_import_metric_clusters_rows.py @@ -2,10 +2,10 @@ from __future__ import annotations -from datetime import datetime, timezone +import types -from app.models.database import CallImportEvaluation from tests.test_api.test_call_import_evaluation_insights import _seed_eval_with_data +from tests.test_api.test_call_import_evaluation_rows_sorting import _seed_eval_with_rows from tests.test_api.test_call_import_evaluations import _stub_celery_revoke @@ -24,6 +24,52 @@ def test_list_eligible_metric_cluster_rows_empty_scores( assert body["items"] == [] +def test_list_eligible_metric_cluster_rows_count_only( + authenticated_client, db_session, org_id, seed_org, make_ai_provider +): + make_ai_provider(provider="openai", is_active=True) + call_import, evaluation, _ = _seed_eval_with_rows( + db_session, + org_id, + rows=[ + {"conversation_id": f"c{i}", "status": "completed", "score_value": 0.2} + for i in range(5) + ], + ) + + response = authenticated_client.get( + f"/api/v1/call-imports/{call_import.id}/evaluations/{evaluation.id}/metric-clusters/eligible-rows", + params={"count_only": True}, + ) + assert response.status_code == 200 + body = response.json() + assert body["total"] == 5 + assert body["items"] == [] + + +def test_list_eligible_metric_cluster_rows_limit( + authenticated_client, db_session, org_id, seed_org, make_ai_provider +): + make_ai_provider(provider="openai", is_active=True) + call_import, evaluation, _ = _seed_eval_with_rows( + db_session, + org_id, + rows=[ + {"conversation_id": f"c{i}", "status": "completed", "score_value": 0.2} + for i in range(5) + ], + ) + + response = authenticated_client.get( + f"/api/v1/call-imports/{call_import.id}/evaluations/{evaluation.id}/metric-clusters/eligible-rows", + params={"limit": 2}, + ) + assert response.status_code == 200 + body = response.json() + assert body["total"] == 5 + assert len(body["items"]) == 2 + + def test_generate_metric_clusters_rejects_unknown_row_id( authenticated_client, db_session, org_id, seed_org, make_ai_provider ): @@ -38,6 +84,69 @@ def test_generate_metric_clusters_rejects_unknown_row_id( assert "evaluation_row_ids" in response.json()["detail"].lower() or "missing" in response.json()["detail"].lower() +def test_generate_metric_clusters_rejects_row_limit_with_row_ids( + authenticated_client, db_session, org_id, seed_org, make_ai_provider +): + make_ai_provider(provider="openai", is_active=True) + call_import, evaluation, _ = _seed_eval_with_rows( + db_session, + org_id, + rows=[ + {"conversation_id": "c0", "status": "completed", "score_value": 0.2}, + ], + ) + + response = authenticated_client.post( + f"/api/v1/call-imports/{call_import.id}/evaluations/{evaluation.id}/metric-clusters", + json={ + "row_limit": 1, + "evaluation_row_ids": ["00000000-0000-0000-0000-000000000001"], + }, + ) + assert response.status_code == 400 + assert "row_limit" in response.json()["detail"].lower() + + +def test_generate_metric_clusters_row_limit( + authenticated_client, + db_session, + org_id, + seed_org, + make_ai_provider, + monkeypatch, +): + make_ai_provider(provider="openai", is_active=True) + call_import, evaluation, _ = _seed_eval_with_rows( + db_session, + org_id, + rows=[ + {"conversation_id": f"c{i}", "status": "completed", "score_value": 0.2} + for i in range(5) + ], + ) + + captured: dict = {} + + def fake_apply_async(*, kwargs=None, **_kw): + captured.update(kwargs or {}) + return types.SimpleNamespace(id="cluster-task-1") + + monkeypatch.setattr( + "app.workers.tasks.generate_evaluation_metric_clusters.generate_evaluation_metric_clusters_task.apply_async", + fake_apply_async, + ) + + response = authenticated_client.post( + f"/api/v1/call-imports/{call_import.id}/evaluations/{evaluation.id}/metric-clusters", + json={"row_limit": 2}, + ) + assert response.status_code == 200 + body = response.json() + assert body["status"] == "running" + assert len(body["selected_evaluation_row_ids"]) == 2 + assert len(captured["evaluation_row_ids"]) == 2 + + def test_cancel_preserves_selected_row_ids_in_state( authenticated_client, db_session, org_id, seed_org, make_ai_provider, monkeypatch ): diff --git a/tests/test_api/test_call_imports_routes.py b/tests/test_api/test_call_imports_routes.py index ab9a70ed..9ae12f78 100644 --- a/tests/test_api/test_call_imports_routes.py +++ b/tests/test_api/test_call_imports_routes.py @@ -797,14 +797,19 @@ def test_audio_upload_single_file_creates_completed_import( ) db_session.add(tag) db_session.commit() + workspace_id = _default_workspace_id(db_session, org_id) fake_s3 = _fake_enabled_s3() with _patched_s3(fake_s3): - response = authenticated_client.post( - "/api/v1/call-imports/audio-upload", - files={"files": ("sales call.wav", b"RIFFfake-wav", "audio/wav")}, - data={"dataset": "Manual recordings", "tag_ids": str(tag.id)}, - ) + with patch( + "app.api.v1.routes.call_imports.is_sharding_enabled", + return_value=False, + ): + response = authenticated_client.post( + "/api/v1/call-imports/audio-upload", + files={"files": ("sales call.wav", b"RIFFfake-wav", "audio/wav")}, + data={"dataset": "Manual recordings", "tag_ids": str(tag.id)}, + ) assert response.status_code == 201, response.text body = response.json() @@ -824,6 +829,7 @@ def test_audio_upload_single_file_creates_completed_import( .filter(CallImportRow.call_import_id == call_import_id) .all() ) + assert row.workspace_id == workspace_id assert row.conversation_id == "sales_call" assert row.status == CallImportRowStatus.COMPLETED assert row.transcript is None @@ -833,21 +839,67 @@ def test_audio_upload_single_file_creates_completed_import( fake_s3.upload_file_by_key.assert_called_once() -def test_audio_upload_multiple_files_dedupes_filename_call_ids( - authenticated_client, db_session, org_id, seed_org +def test_audio_upload_uses_shard_insert_when_sharding_enabled( + authenticated_client, db_session, org_id, seed_org, monkeypatch ): fake_s3 = _fake_enabled_s3() + inserted: list = [] + registered: list = [] + + monkeypatch.setattr( + "app.api.v1.routes.call_imports.is_sharding_enabled", + lambda: True, + ) + + def _fake_bulk_insert(_db, call_import_id, mappings): + inserted.extend(mappings) + return len(mappings) + + def _fake_register_slices(_db, call_import_id, total_rows): + registered.append((call_import_id, total_rows)) + + monkeypatch.setattr( + "app.db_sharding.row_ops.bulk_insert_mappings_on_shards", + _fake_bulk_insert, + ) + monkeypatch.setattr( + "app.db_sharding.row_ops.register_shard_slices", + _fake_register_slices, + ) + with _patched_s3(fake_s3): response = authenticated_client.post( "/api/v1/call-imports/audio-upload", - files=[ - ("files", ("call.mp3", b"mp3-a", "audio/mpeg")), - ("files", ("call.mp3", b"mp3-b", "audio/mpeg")), - ("files", ("support flac.flac", b"flac", "audio/flac")), - ], + files={"files": ("call.wav", b"RIFFfake-wav", "audio/wav")}, data={"dataset": "Manual recordings"}, ) + assert response.status_code == 201, response.text + assert len(inserted) == 1 + assert inserted[0]["workspace_id"] is not None + assert inserted[0]["recording_s3_key"] + assert registered == [(UUID(response.json()["id"]), 1)] + + +def test_audio_upload_multiple_files_dedupes_filename_call_ids( + authenticated_client, db_session, org_id, seed_org +): + fake_s3 = _fake_enabled_s3() + with _patched_s3(fake_s3): + with patch( + "app.api.v1.routes.call_imports.is_sharding_enabled", + return_value=False, + ): + response = authenticated_client.post( + "/api/v1/call-imports/audio-upload", + files=[ + ("files", ("call.mp3", b"mp3-a", "audio/mpeg")), + ("files", ("call.mp3", b"mp3-b", "audio/mpeg")), + ("files", ("support flac.flac", b"flac", "audio/flac")), + ], + data={"dataset": "Manual recordings"}, + ) + assert response.status_code == 201, response.text rows = ( db_session.query(CallImportRow) diff --git a/tests/test_db_sharding/__init__.py b/tests/test_db_sharding/__init__.py new file mode 100644 index 00000000..62c1e5d2 --- /dev/null +++ b/tests/test_db_sharding/__init__.py @@ -0,0 +1 @@ +# db_sharding tests diff --git a/tests/test_db_sharding/test_diarised_aggregate.py b/tests/test_db_sharding/test_diarised_aggregate.py new file mode 100644 index 00000000..71bc2358 --- /dev/null +++ b/tests/test_db_sharding/test_diarised_aggregate.py @@ -0,0 +1,22 @@ +"""Tests for sharded diarisation status aggregation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.scatter_gather import aggregate_diarised_transcript_counts + + +def test_aggregate_diarised_mono_db(): + catalog_db = MagicMock() + call_import_id = uuid4() + catalog_db.query.return_value.filter.return_value.group_by.return_value.all.return_value = [ + ("pending", 2), + ("completed", 5), + ] + + with patch("app.db_sharding.scatter_gather.is_sharding_enabled", return_value=False): + counts = aggregate_diarised_transcript_counts(catalog_db, call_import_id) + + assert counts == {"pending": 2, "completed": 5} diff --git a/tests/test_db_sharding/test_eval_rows_helpers.py b/tests/test_db_sharding/test_eval_rows_helpers.py new file mode 100644 index 00000000..39fda32e --- /dev/null +++ b/tests/test_db_sharding/test_eval_rows_helpers.py @@ -0,0 +1,44 @@ +"""Tests for sharded evaluation row read helpers.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.eval_rows import ( + count_evaluation_rows_for_run, + load_evaluation_rows_for_run, +) + + +def test_load_evaluation_rows_delegates_to_pairs(): + catalog_db = MagicMock() + evaluation_id = uuid4() + eval_row = MagicMock() + source_row = MagicMock() + + with patch( + "app.db_sharding.eval_rows.load_evaluation_row_pairs", + return_value=[(eval_row, source_row)], + ): + rows = load_evaluation_rows_for_run(catalog_db, evaluation_id) + + assert rows == [eval_row] + + +def test_count_evaluation_rows_by_status(): + catalog_db = MagicMock() + evaluation_id = uuid4() + r1 = MagicMock(status="completed") + r2 = MagicMock(status="failed") + + with patch( + "app.db_sharding.eval_rows.load_evaluation_row_pairs", + return_value=[(r1, MagicMock()), (r2, MagicMock())], + ): + assert ( + count_evaluation_rows_for_run( + catalog_db, evaluation_id, statuses=["completed"] + ) + == 1 + ) diff --git a/tests/test_db_sharding/test_eval_rows_pagination.py b/tests/test_db_sharding/test_eval_rows_pagination.py new file mode 100644 index 00000000..d7ebb9d1 --- /dev/null +++ b/tests/test_db_sharding/test_eval_rows_pagination.py @@ -0,0 +1,92 @@ +"""Tests for sharded eval-row pagination helpers.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from app.db_sharding.eval_rows import fetch_evaluation_row_pairs_page + + +def _pair(row_index: int, status: str = "completed"): + eval_row = SimpleNamespace(status=status, metric_scores={}) + source_row = SimpleNamespace(row_index=row_index, conversation_id=f"c-{row_index}") + return eval_row, source_row + + +@patch("app.db_sharding.eval_rows.is_sharding_enabled", return_value=True) +@patch("app.db_sharding.eval_rows.scatter_gather_eval_query_count", return_value=4) +@patch("app.db_sharding.eval_rows.db_pool_manager") +def test_fetch_page_unbounded_sort_loads_all_shard_rows( + mock_pool_manager, + _mock_count, + _mock_sharding, +): + """Non-row_index sorts must not truncate per-shard results.""" + shard_a = MagicMock() + shard_b = MagicMock() + shard_a.all.return_value = [_pair(0, "failed"), _pair(1, "failed")] + shard_b.all.return_value = [_pair(2, "completed"), _pair(3, "completed")] + shard_a.limit.return_value = shard_a + shard_b.limit.return_value = shard_b + + mock_pool_manager.router.shard_ids = ["shard-a", "shard-b"] + session_a = MagicMock() + session_b = MagicMock() + + def _factory(shard_id): + if shard_id == "shard-a": + return lambda: session_a + return lambda: session_b + + mock_pool_manager.shard_session_factory.side_effect = _factory + + def _build_query(session): + return shard_a if session is session_a else shard_b + + total, rows = fetch_evaluation_row_pairs_page( + MagicMock(), + _build_query, + page=1, + page_size=2, + sort_key=lambda pair: pair[0].status or "", + bounded_shard_fetch=False, + ) + + assert total == 4 + assert len(rows) == 2 + shard_a.limit.assert_not_called() + shard_b.limit.assert_not_called() + assert [pair[0].status for pair in rows] == ["completed", "completed"] + + +@patch("app.db_sharding.eval_rows.is_sharding_enabled", return_value=True) +@patch("app.db_sharding.eval_rows.scatter_gather_eval_query_count", return_value=4) +@patch("app.db_sharding.eval_rows.db_pool_manager") +def test_fetch_page_bounded_sort_limits_each_shard( + mock_pool_manager, + _mock_count, + _mock_sharding, +): + shard_a = MagicMock() + shard_a.all.return_value = [_pair(0), _pair(1)] + shard_a.limit.return_value = shard_a + + mock_pool_manager.router.shard_ids = ["shard-a"] + session = MagicMock() + mock_pool_manager.shard_session_factory.side_effect = lambda _shard_id: ( + lambda: session + ) + + total, rows = fetch_evaluation_row_pairs_page( + MagicMock(), + lambda _session: shard_a, + page=1, + page_size=2, + sort_key=lambda pair: int(pair[1].row_index or 0), + bounded_shard_fetch=True, + ) + + assert total == 4 + assert len(rows) == 2 + shard_a.limit.assert_called_once_with(2) diff --git a/tests/test_db_sharding/test_eval_sharding_reads.py b/tests/test_db_sharding/test_eval_sharding_reads.py new file mode 100644 index 00000000..588cb322 --- /dev/null +++ b/tests/test_db_sharding/test_eval_sharding_reads.py @@ -0,0 +1,76 @@ +"""Evaluation read helpers with multi-shard row data (integration-style).""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.eval_rows import ( + find_evaluation_row_in_run, + gather_retry_targets_sharded, + load_evaluation_rows_for_run, +) + + +def test_load_evaluation_rows_from_scatter_pairs(): + catalog_db = MagicMock() + evaluation_id = uuid4() + er_a = MagicMock(id=uuid4()) + er_b = MagicMock(id=uuid4()) + + with patch( + "app.db_sharding.eval_rows.load_evaluation_row_pairs", + return_value=[(er_a, MagicMock()), (er_b, MagicMock())], + ): + rows = load_evaluation_rows_for_run(catalog_db, evaluation_id) + + assert rows == [er_a, er_b] + + +def test_find_evaluation_row_in_run_two_shard_pairs(): + catalog_db = MagicMock() + evaluation_id = uuid4() + target_id = uuid4() + other_id = uuid4() + target_row = MagicMock(id=target_id, evaluation_id=evaluation_id) + other_row = MagicMock(id=other_id, evaluation_id=evaluation_id) + + with patch( + "app.db_sharding.eval_rows.load_evaluation_row_pairs", + return_value=[ + (other_row, MagicMock()), + (target_row, MagicMock()), + ], + ): + found, _source = find_evaluation_row_in_run( + catalog_db, evaluation_id, target_id + ) + + assert found is target_row + + +def test_gather_retry_targets_failed_rows_from_two_shards(): + catalog_db = MagicMock() + evaluation = MagicMock(id=uuid4()) + er1 = MagicMock(id=uuid4(), status="failed") + er2 = MagicMock(id=uuid4(), status="failed") + er3 = MagicMock(id=uuid4(), status="running") + + with patch( + "app.db_sharding.eval_rows.load_evaluation_row_pairs", + return_value=[ + (er1, MagicMock()), + (er2, MagicMock()), + (er3, MagicMock()), + ], + ): + targets, skipped = gather_retry_targets_sharded( + catalog_db, + evaluation, + requested_ids=None, + include_completed=False, + ) + + assert len(targets) == 2 + assert {t[0].id for t in targets} == {er1.id, er2.id} + assert skipped == [] diff --git a/tests/test_db_sharding/test_import_row_filters.py b/tests/test_db_sharding/test_import_row_filters.py new file mode 100644 index 00000000..a6cc5d3b --- /dev/null +++ b/tests/test_db_sharding/test_import_row_filters.py @@ -0,0 +1,48 @@ +"""Sharded import row filter scatter-gather tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.scatter_gather import ( + fetch_call_import_rows_filtered_page, + merge_rows_by_index, +) + + +def test_merge_rows_by_index_sorts(): + rows = [ + SimpleNamespace(row_index=2), + SimpleNamespace(row_index=0), + SimpleNamespace(row_index=1), + ] + ordered = merge_rows_by_index(rows) + assert [int(r.row_index) for r in ordered] == [0, 1, 2] + + +def test_filtered_page_slices_merged_rows(): + catalog_db = MagicMock() + call_import_id = uuid4() + merged = [ + SimpleNamespace(row_index=0, id=uuid4()), + SimpleNamespace(row_index=1, id=uuid4()), + SimpleNamespace(row_index=2, id=uuid4()), + ] + + with patch( + "app.db_sharding.scatter_gather._merged_call_import_rows_for_import", + return_value=merged, + ): + page = fetch_call_import_rows_filtered_page( + catalog_db, + call_import_id, + search_term="foo", + diarised_status_filter="completed", + offset=1, + limit=1, + ) + + assert len(page) == 1 + assert page[0].row_index == 1 diff --git a/tests/test_db_sharding/test_pool_manager.py b/tests/test_db_sharding/test_pool_manager.py new file mode 100644 index 00000000..f6bc9b1e --- /dev/null +++ b/tests/test_db_sharding/test_pool_manager.py @@ -0,0 +1,78 @@ +import pytest +from sqlalchemy import create_engine, text + +from app.db_sharding.pool_manager import DatabasePoolManager + + +@pytest.fixture +def manager(): + m = DatabasePoolManager() + yield m + m.reset() + + +def test_legacy_single_engine(manager, monkeypatch): + url = "sqlite:///:memory:" + monkeypatch.setattr( + "app.config.settings", + type( + "S", + (), + { + "DATABASE_URL": url, + "DB_SHARDING_ENABLED": False, + "DB_POOL_SIZE": 5, + "DB_MAX_OVERFLOW": 5, + "DB_CATALOG_URL": None, + "DB_SHARD_ROW_CHUNK_SIZE": 500, + "DB_SHARD_ENTRIES": [], + }, + )(), + ) + eng = manager.catalog_engine + assert eng.url.database == ":memory:" + assert not manager.sharding_enabled + assert manager.router is None + assert len(manager.all_engines_for_migrations()) == 1 + + +def test_sharding_two_shards_dedupe_migrations(manager, monkeypatch): + url = "sqlite:///:memory:" + monkeypatch.setattr( + "app.config.settings", + type( + "S", + (), + { + "DATABASE_URL": url, + "DB_SHARDING_ENABLED": True, + "DB_CATALOG_URL": url, + "DB_POOL_SIZE": 2, + "DB_MAX_OVERFLOW": 2, + "DB_SHARD_ROW_CHUNK_SIZE": 500, + "DB_SHARD_ENTRIES": [ + {"id": "data-shard-01", "url": url}, + {"id": "data-shard-02", "url": url}, + ], + }, + )(), + ) + assert manager.sharding_enabled + assert manager.router is not None + assert manager.router.shard_count == 2 + engines = manager.all_engines_for_migrations() + assert len(engines) == 1 + + session = manager.shard_session_factory( + manager.router.shard_id_for_row( + "00000000-0000-0000-0000-000000000001", 0 + ) + )() + shard_id = manager.router.shard_id_for_row( + "00000000-0000-0000-0000-000000000001", 0 + ) + try: + session.execute(text("SELECT 1")) + assert shard_id in ("data-shard-01", "data-shard-02") + finally: + session.close() diff --git a/tests/test_db_sharding/test_rebalance.py b/tests/test_db_sharding/test_rebalance.py new file mode 100644 index 00000000..d0c334d9 --- /dev/null +++ b/tests/test_db_sharding/test_rebalance.py @@ -0,0 +1,134 @@ +"""Unit tests for shard slice rebalance / backfill tooling.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from app.db_sharding.rebalance import ( + RebalanceError, + RebalancePlan, + SliceInfo, + assert_import_rebalance_ready, + build_rebalance_plan, + execute_rebalance_slices, + filter_unlocked_call_import_ids, + list_shard_slices, +) +from app.models.enums import CallImportStatus + + +def test_slice_info_row_count(): + info = SliceInfo(slice_id=0, shard_id="s1", row_index_min=0, row_index_max=499) + assert info.row_count == 500 + + +def test_list_shard_slices_maps_registry_rows(): + catalog_db = MagicMock() + call_import_id = uuid4() + catalog_db.execute.return_value.all.return_value = [ + (0, "data-shard-01", 0, 499), + (1, "data-shard-02", 500, 999), + ] + slices = list_shard_slices(catalog_db, call_import_id) + assert len(slices) == 2 + assert slices[0].shard_id == "data-shard-01" + assert slices[1].row_index_min == 500 + + +def test_build_rebalance_plan_filters_by_shard_and_slice(): + catalog_db = MagicMock() + call_import_id = uuid4() + slices = [ + SliceInfo(0, "data-shard-01", 0, 499), + SliceInfo(1, "data-shard-01", 500, 999), + SliceInfo(2, "data-shard-02", 1000, 1499), + ] + + with patch("app.db_sharding.rebalance.is_sharding_enabled", return_value=True): + with patch("app.db_sharding.rebalance.list_shard_slices", return_value=slices): + with patch( + "app.db_sharding.rebalance._evaluation_ids_for_import", + return_value=[], + ): + with patch( + "app.db_sharding.rebalance._count_rows_on_shard", + return_value=(500, 0), + ): + with patch( + "app.db_sharding.rebalance._configured_shard_ids", + return_value={"data-shard-01", "data-shard-02"}, + ): + plan = build_rebalance_plan( + catalog_db, + call_import_id, + from_shard_id="data-shard-01", + to_shard_id="data-shard-02", + slice_ids=[1], + ) + + assert plan.from_shard_id == "data-shard-01" + assert plan.to_shard_id == "data-shard-02" + assert len(plan.slices) == 1 + assert plan.slices[0].slice_id == 1 + assert plan.import_row_count == 500 + + +def test_build_rebalance_plan_rejects_same_shard(): + catalog_db = MagicMock() + with patch("app.db_sharding.rebalance.is_sharding_enabled", return_value=True): + with patch( + "app.db_sharding.rebalance._configured_shard_ids", + return_value={"data-shard-01"}, + ): + with pytest.raises(RebalanceError, match="must differ"): + build_rebalance_plan( + catalog_db, + uuid4(), + from_shard_id="data-shard-01", + to_shard_id="data-shard-01", + ) + + +def test_assert_import_rebalance_ready_blocks_processing(): + catalog_db = MagicMock() + call_import = MagicMock(status=CallImportStatus.PROCESSING) + catalog_db.query.return_value.filter.return_value.first.return_value = call_import + + with pytest.raises(RebalanceError, match="terminal status"): + assert_import_rebalance_ready(catalog_db, uuid4(), force=False) + + +def test_execute_rebalance_slices_dry_run_skips_copy(): + catalog_db = MagicMock() + plan = RebalancePlan( + call_import_id=uuid4(), + from_shard_id="data-shard-01", + to_shard_id="data-shard-02", + slices=(SliceInfo(0, "data-shard-01", 0, 10),), + import_row_count=11, + eval_row_count=3, + ) + + with patch("app.db_sharding.rebalance.is_sharding_enabled", return_value=True): + with patch("app.db_sharding.rebalance.assert_import_rebalance_ready"): + with patch("app.db_sharding.rebalance._copy_rows_between_shards") as copy_mock: + result = execute_rebalance_slices(catalog_db, plan, dry_run=True) + + assert result.dry_run is True + assert result.import_rows_moved == 11 + copy_mock.assert_not_called() + catalog_db.commit.assert_not_called() + + +def test_filter_unlocked_call_import_ids(): + locked_id = uuid4() + open_id = uuid4() + with patch( + "app.db_sharding.rebalance.is_import_rebalance_locked", + side_effect=lambda cid: cid == locked_id, + ): + out = filter_unlocked_call_import_ids([locked_id, open_id]) + assert out == [open_id] diff --git a/tests/test_db_sharding/test_router.py b/tests/test_db_sharding/test_router.py new file mode 100644 index 00000000..5a16b56e --- /dev/null +++ b/tests/test_db_sharding/test_router.py @@ -0,0 +1,49 @@ +import uuid + +import pytest + +from app.db_sharding.router import ShardRouter + + +def test_slice_id_for_row_index(): + router = ShardRouter(["a", "b"], row_chunk_size=500) + assert router.slice_id_for_row_index(0) == 0 + assert router.slice_id_for_row_index(499) == 0 + assert router.slice_id_for_row_index(500) == 1 + + +def test_single_shard_always_same(): + router = ShardRouter(["only"], row_chunk_size=100) + cid = uuid.uuid4() + for idx in (0, 50, 9999): + assert router.shard_id_for_row(cid, idx) == "only" + + +def test_deterministic_routing(): + router = ShardRouter(["s1", "s2", "s3"], row_chunk_size=500) + cid = uuid.uuid4() + first = router.shard_id_for_row(cid, 1200) + assert first in ("s1", "s2", "s3") + assert router.shard_id_for_row(cid, 1200) == first + assert router.shard_id_for_row(str(cid), 1200) == first + + +def test_registry_override(): + router = ShardRouter(["s1", "s2"], row_chunk_size=500) + cid = uuid.uuid4() + slice_id = router.slice_id_for_row_index(750) + assert slice_id == 1 + hashed = router.shard_id_for_slice(cid, slice_id) + registry = {(str(cid), slice_id): "s2"} + assert router.shard_id_for_row(cid, 750, slice_registry=registry) == "s2" + assert router.shard_id_for_row(cid, 750) == hashed + + +def test_invalid_chunk_size(): + with pytest.raises(ValueError): + ShardRouter(["a"], row_chunk_size=0) + + +def test_empty_shard_list(): + with pytest.raises(ValueError): + ShardRouter([]) diff --git a/tests/test_db_sharding/test_row_ops_deferred_commit.py b/tests/test_db_sharding/test_row_ops_deferred_commit.py new file mode 100644 index 00000000..e1c9b246 --- /dev/null +++ b/tests/test_db_sharding/test_row_ops_deferred_commit.py @@ -0,0 +1,57 @@ +"""Tests for deferred shard bulk-insert commit ordering.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.row_ops import ( + bulk_insert_mappings_on_shards, + commit_pending_shard_sessions, + rollback_pending_shard_sessions, +) + + +@patch("app.db_sharding.row_ops.is_sharding_enabled", return_value=True) +@patch("app.db_sharding.row_ops.partition_mappings_by_shard") +@patch("app.db_sharding.row_ops.db_pool_manager") +def test_bulk_insert_defer_commit_stages_without_closing_session( + mock_pool_manager, + mock_partition, + _mock_sharding, +): + shard_db = MagicMock() + mock_pool_manager.shard_session_factory.return_value = lambda: shard_db + mock_partition.return_value = {"shard-a": [{"row_index": 0}]} + + inserted, pending = bulk_insert_mappings_on_shards( + MagicMock(), + uuid4(), + [{"row_index": 0}], + defer_commit=True, + ) + + assert inserted == 1 + assert pending == [shard_db] + shard_db.commit.assert_not_called() + shard_db.close.assert_not_called() + + +@patch("app.db_sharding.row_ops.commit_shard_row_session") +@patch("app.db_sharding.row_ops._reset_shard_write_role") +def test_commit_pending_shard_sessions_commits_and_closes( + _mock_reset, + mock_commit, +): + shard_db = MagicMock() + commit_pending_shard_sessions([shard_db]) + mock_commit.assert_called_once_with(shard_db) + shard_db.close.assert_called_once() + + +@patch("app.db_sharding.row_ops._reset_shard_write_role") +def test_rollback_pending_shard_sessions_rolls_back_and_closes(_mock_reset): + shard_db = MagicMock() + rollback_pending_shard_sessions([shard_db]) + shard_db.rollback.assert_called_once() + shard_db.close.assert_called_once() diff --git a/tests/test_db_sharding/test_row_ops_eval.py b/tests/test_db_sharding/test_row_ops_eval.py new file mode 100644 index 00000000..ef987642 --- /dev/null +++ b/tests/test_db_sharding/test_row_ops_eval.py @@ -0,0 +1,44 @@ +"""Sharded eval-row placement helpers.""" + +import uuid +from unittest.mock import MagicMock, patch + +from app.db_sharding.row_ops import partition_eval_mappings_by_shard + + +def test_partition_eval_mappings_by_shard_routes_by_source_row_index(): + call_import_id = uuid.uuid4() + source_a = uuid.uuid4() + source_b = uuid.uuid4() + index_by_id = {source_a: 0, source_b: 1200} + mappings = [ + { + "evaluation_id": uuid.uuid4(), + "call_import_row_id": source_a, + "status": "pending", + }, + { + "evaluation_id": uuid.uuid4(), + "call_import_row_id": source_b, + "status": "pending", + }, + ] + catalog_db = MagicMock() + router = MagicMock() + router.shard_id_for_row.side_effect = lambda cid, idx, **_: ( + "s1" if idx < 500 else "s2" + ) + with patch("app.db_sharding.row_ops.is_sharding_enabled", return_value=True): + with patch( + "app.db_sharding.row_ops.router_for_import", + return_value=(router, None), + ): + buckets = partition_eval_mappings_by_shard( + catalog_db, + call_import_id, + mappings, + index_by_source_id=index_by_id, + ) + assert len(buckets["s1"]) == 1 + assert len(buckets["s2"]) == 1 + assert buckets["s1"][0]["call_import_row_id"] == source_a diff --git a/tests/test_db_sharding/test_scatter_gather_api.py b/tests/test_db_sharding/test_scatter_gather_api.py new file mode 100644 index 00000000..3bf4a806 --- /dev/null +++ b/tests/test_db_sharding/test_scatter_gather_api.py @@ -0,0 +1,18 @@ +"""Scatter-gather helpers.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.db_sharding.scatter_gather import fetch_call_import_rows_page + + +def test_fetch_rows_page_when_sharding_off(): + call_import_id = uuid4() + db = MagicMock() + mock_rows = [MagicMock(row_index=0)] + db.query.return_value.filter.return_value.order_by.return_value.offset.return_value.limit.return_value.all.return_value = ( + mock_rows + ) + with patch("app.db_sharding.scatter_gather.is_sharding_enabled", return_value=False): + out = fetch_call_import_rows_page(db, call_import_id, offset=0, limit=10) + assert out == mock_rows diff --git a/tests/test_db_sharding/test_sharding_postgres_integration.py b/tests/test_db_sharding/test_sharding_postgres_integration.py new file mode 100644 index 00000000..f13f2486 --- /dev/null +++ b/tests/test_db_sharding/test_sharding_postgres_integration.py @@ -0,0 +1,176 @@ +"""Two-shard Postgres integration for scatter-gather import row reads. + +Runs in CI when SHARDING_INTEGRATION_TEST=1 and catalog + shard URLs are set. +""" + +from __future__ import annotations + +import os +from uuid import uuid4 + +import pytest +from sqlalchemy import create_engine + +pytestmark = pytest.mark.integration + + +def _require_env(name: str) -> str: + value = (os.getenv(name) or "").strip() + if not value: + pytest.skip(f"{name} is not set") + return value + + +@pytest.fixture(scope="module") +def sharding_postgres_env(): + if os.getenv("SHARDING_INTEGRATION_TEST") != "1": + pytest.skip("SHARDING_INTEGRATION_TEST is not enabled") + + catalog_url = _require_env("CATALOG_DATABASE_URL") + shard_01 = _require_env("SHARD_DATABASE_URL_01") + shard_02 = _require_env("SHARD_DATABASE_URL_02") + + from app.database import Base + from app.db_sharding.pool_manager import db_pool_manager + + # Register ORM models on Base.metadata before create_all (see init_db()). + import app.models.database # noqa: F401 + + for url in (catalog_url, shard_01, shard_02): + engine = create_engine(url, pool_pre_ping=True) + Base.metadata.create_all(bind=engine) + engine.dispose() + + db_pool_manager.reset() + + class _Settings: + DATABASE_URL = catalog_url + DB_CATALOG_URL = catalog_url + DB_SHARDING_ENABLED = True + DB_SHARD_ROW_CHUNK_SIZE = 500 + DB_POOL_SIZE = 2 + DB_MAX_OVERFLOW = 2 + DB_SHARD_ENTRIES = [ + {"id": "data-shard-01", "url": shard_01}, + {"id": "data-shard-02", "url": shard_02}, + ] + + import app.config as config_module + + prior_settings = config_module.settings + config_module.settings = _Settings() + db_pool_manager.reset() + + from app.db_sharding.sessions import is_sharding_enabled + + if not is_sharding_enabled(): + config_module.settings = prior_settings + db_pool_manager.reset() + pytest.skip("pool manager did not enable sharding") + + yield { + "catalog_url": catalog_url, + "shard_01": shard_01, + "shard_02": shard_02, + } + + config_module.settings = prior_settings + db_pool_manager.reset() + + +def test_import_row_filter_scatter_gather_two_shards(sharding_postgres_env): + from app.db_sharding.pool_manager import open_catalog_session + from app.db_sharding.row_ops import bulk_insert_mappings_on_shards, register_shard_slices + from app.db_sharding.scatter_gather import ( + count_call_import_rows_filtered, + list_call_import_row_ids_filtered, + ) + from app.models.database import CallImport, CallImportRow, Organization, Workspace + from app.models.enums import CallImportRowStatus, CallImportStatus + + org_id = uuid4() + workspace_id = uuid4() + call_import_id = uuid4() + row_a_id = uuid4() + row_b_id = uuid4() + + catalog = open_catalog_session() + try: + org = Organization(id=org_id, name="Sharding CI Org") + workspace = Workspace( + id=workspace_id, + organization=org, + name="Default", + slug="default", + is_default=True, + ) + catalog.add_all([org, workspace]) + catalog.flush() + + catalog.add( + CallImport( + id=call_import_id, + organization_id=org_id, + workspace_id=workspace_id, + total_rows=2, + status=CallImportStatus.PROCESSING, + ) + ) + catalog.commit() + + mappings = [ + { + "id": row_a_id, + "call_import_id": call_import_id, + "organization_id": org_id, + "workspace_id": workspace_id, + "row_index": 0, + "conversation_id": "match-alpha", + "status": CallImportRowStatus.COMPLETED, + "transcript_status": "idle", + "diarised_transcript_status": "completed", + }, + { + "id": row_b_id, + "call_import_id": call_import_id, + "organization_id": org_id, + "workspace_id": workspace_id, + "row_index": 1, + "conversation_id": "other-beta", + "status": CallImportRowStatus.COMPLETED, + "transcript_status": "idle", + "diarised_transcript_status": "pending", + }, + ] + inserted, pending_sessions = bulk_insert_mappings_on_shards( + catalog, call_import_id, mappings, orm_class=CallImportRow + ) + assert inserted == 2 + assert pending_sessions == [] + register_shard_slices(catalog, call_import_id, 2) + catalog.commit() + + assert ( + count_call_import_rows_filtered( + catalog, + call_import_id, + search_term="alpha", + ) + == 1 + ) + assert ( + count_call_import_rows_filtered( + catalog, + call_import_id, + diarised_status_filter="pending", + ) + == 1 + ) + ids = list_call_import_row_ids_filtered( + catalog, + call_import_id, + search_term="alpha", + ) + assert ids == [row_a_id] + finally: + catalog.close() diff --git a/tests/test_services/test_call_import_bulk_ops_unified.py b/tests/test_services/test_call_import_bulk_ops_unified.py index 578a45f5..b357fa9e 100644 --- a/tests/test_services/test_call_import_bulk_ops_unified.py +++ b/tests/test_services/test_call_import_bulk_ops_unified.py @@ -139,6 +139,10 @@ def test_bulk_diarization_stores_redis_params_before_pending_commit( db_session, ): """Rows must not be visible to the fair dispatcher until params exist.""" + monkeypatch.setattr( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + lambda: False, + ) org = Organization(id=uuid4(), name="Diar Org") ws = Workspace( id=uuid4(), diff --git a/tests/test_services/test_call_imports/__init__.py b/tests/test_services/test_call_imports/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_services/test_call_imports/test_bulk_ops_sharding.py b/tests/test_services/test_call_imports/test_bulk_ops_sharding.py new file mode 100644 index 00000000..2b2c13b7 --- /dev/null +++ b/tests/test_services/test_call_imports/test_bulk_ops_sharding.py @@ -0,0 +1,123 @@ +"""Sharding-aware bulk call-import operation tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from app.models.schemas import CallImportTranscribeRequest +from app.services.call_imports.bulk_ops import ( + execute_bulk_diarization, + select_rows_for_transcription, +) + + +def test_select_rows_for_transcription_uses_scatter_gather_loader(): + call_import = SimpleNamespace(id=uuid4()) + row_id = uuid4() + fake_rows = [ + SimpleNamespace( + id=row_id, + row_index=0, + recording_s3_key="s3-key", + diarised_transcript=None, + ) + ] + payload = CallImportTranscribeRequest( + stt_provider="deepgram", + stt_model="nova-2", + diarization_llm_provider="openai", + diarization_llm_model="gpt-4o-mini", + ) + fake_db = MagicMock() + + with patch( + "app.db_sharding.scatter_gather.load_call_import_rows_for_transcription", + return_value=fake_rows, + ) as mock_loader: + selected, skip_counts = select_rows_for_transcription( + fake_db, + call_import, + payload, + requested_row_ids=[row_id], + ) + + mock_loader.assert_called_once_with( + fake_db, + call_import.id, + requested_row_ids=[row_id], + ) + assert len(selected) == 1 + assert skip_counts == {} + + +def test_execute_bulk_diarization_updates_rows_on_shards(): + call_import = SimpleNamespace(id=uuid4()) + row_id = uuid4() + fake_row = SimpleNamespace( + id=row_id, + row_index=0, + recording_s3_key="s3-key", + diarised_transcript=None, + ) + payload = CallImportTranscribeRequest( + stt_provider="deepgram", + stt_model="nova-2", + diarization_llm_provider="openai", + diarization_llm_model="gpt-4o-mini", + ) + fake_db = MagicMock() + + with ( + patch( + "app.services.call_imports.bulk_ops.select_rows_for_transcription", + return_value=([fake_row], {}), + ), + patch( + "app.workers.concurrency.diarization_dispatch.build_diarization_params_from_request", + return_value={"mode": "stt_llm"}, + ), + patch( + "app.workers.concurrency.diarization_dispatch.store_row_diarization_params_batch", + return_value=([row_id], []), + ), + patch( + "app.db_sharding.row_ops.update_call_import_rows_on_shards", + ) as mock_update, + patch( + "app.workers.concurrency.fair_diarization_dispatch.schedule_fair_diarization_dispatch", + ), + ): + result = execute_bulk_diarization(fake_db, call_import, payload) + + assert result.queued == 1 + mock_update.assert_called_once() + updates = mock_update.call_args.args[2] + assert updates[0]["id"] == row_id + assert updates[0]["diarised_transcript_status"] == "pending" + fake_db.commit.assert_not_called() + + +def test_count_completed_source_rows_uses_scatter_gather_when_sharding_enabled(): + from app.services.call_imports.bulk_ops import count_completed_source_rows + + call_import_id = uuid4() + fake_db = MagicMock() + + with ( + patch( + "app.db_sharding.scatter_gather.is_sharding_enabled", + return_value=True, + ), + patch( + "app.db_sharding.scatter_gather.count_completed_call_import_rows", + return_value=42, + ) as mock_count, + ): + total = count_completed_source_rows(fake_db, call_import_id) + + assert total == 42 + mock_count.assert_called_once_with(fake_db, call_import_id) diff --git a/tests/test_services/test_call_imports/test_progress_counters.py b/tests/test_services/test_call_imports/test_progress_counters.py new file mode 100644 index 00000000..d11e4c79 --- /dev/null +++ b/tests/test_services/test_call_imports/test_progress_counters.py @@ -0,0 +1,68 @@ +"""Tests for progress counter flush ordering (PR #104 P1 fix).""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import redis + +from app.services.call_imports import progress_counters + + +def test_flush_eval_skips_catalog_when_redis_decrement_fails(): + db = MagicMock() + evaluation_id = uuid4() + evaluation = MagicMock(completed_rows=0, failed_rows=0) + + with patch.object( + progress_counters, + "read_eval_progress", + return_value=(3, 1), + ): + with patch.object(progress_counters, "_client") as client_factory: + client = MagicMock() + client.hincrby.side_effect = redis.RedisError("down") + client_factory.return_value = client + db.query.return_value.filter.return_value.first.return_value = evaluation + + progress_counters.flush_eval_progress_to_catalog(db, evaluation_id) + + db.flush.assert_not_called() + + +def test_flush_eval_restores_redis_when_catalog_flush_fails(): + db = MagicMock() + evaluation_id = uuid4() + db.query.return_value.filter.return_value.first.return_value = None + + with patch.object( + progress_counters, + "read_eval_progress", + return_value=(2, 0), + ): + with patch.object(progress_counters, "_client") as client_factory: + client = MagicMock() + client_factory.return_value = client + with patch.object(progress_counters, "record_eval_row_terminal") as restore: + progress_counters.flush_eval_progress_to_catalog(db, evaluation_id) + + restore.assert_called_once_with(evaluation_id, completed_delta=2, failed_delta=0) + + +def test_engine_role_uses_parsed_url_comparison(): + from app.core.migrations import _engine_role_for_url + + catalog = "postgresql://user:pass@localhost:5432/efficientai_catalog" + normalized = "postgresql+psycopg2://user:pass@localhost:5432/efficientai_catalog" + + class _Settings: + DB_SHARDING_ENABLED = True + DB_CATALOG_URL = catalog + DATABASE_URL = catalog + + with patch("app.config.settings", _Settings()): + assert _engine_role_for_url(normalized) == "catalog" + assert _engine_role_for_url( + "postgresql://user:pass@localhost:5432/efficientai_data_01" + ) == "shard" diff --git a/tests/test_services/test_call_imports_bulk_ops.py b/tests/test_services/test_call_imports_bulk_ops.py index 5236e8ca..37ddb9a3 100644 --- a/tests/test_services/test_call_imports_bulk_ops.py +++ b/tests/test_services/test_call_imports_bulk_ops.py @@ -2,6 +2,7 @@ from __future__ import annotations +from unittest.mock import patch from uuid import uuid4 import pytest @@ -62,11 +63,15 @@ def test_select_rows_for_transcription_skips_without_recording(): diarization_llm_model="gpt-4o-mini", ) - selected, skip_counts = select_rows_for_transcription( - db, # type: ignore[arg-type] - _FakeCallImport(import_id), # type: ignore[arg-type] - payload, - ) + with patch( + "app.db_sharding.scatter_gather.load_call_import_rows_for_transcription", + return_value=[row_with, row_without], + ): + selected, skip_counts = select_rows_for_transcription( + db, # type: ignore[arg-type] + _FakeCallImport(import_id), # type: ignore[arg-type] + payload, + ) assert len(selected) == 1 assert selected[0].id == row_with.id diff --git a/tests/test_services/test_evaluation_retry_sharding.py b/tests/test_services/test_evaluation_retry_sharding.py new file mode 100644 index 00000000..ab0bd917 --- /dev/null +++ b/tests/test_services/test_evaluation_retry_sharding.py @@ -0,0 +1,93 @@ +"""Evaluation retry persistence under row sharding.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.models.database import CallImportEvaluationRow, CallImportRow +from app.services.call_imports.bulk_ops import _persist_evaluation_retry_targets + + +def test_persist_evaluation_retry_targets_commits_per_shard(): + catalog_db = MagicMock() + evaluation = SimpleNamespace(call_import_id=uuid4()) + eval_row_id = uuid4() + source_row_id = uuid4() + eval_row = SimpleNamespace( + id=eval_row_id, + celery_task_id=None, + status="failed", + ) + source_row = SimpleNamespace(id=source_row_id, row_index=5) + + shard_db = MagicMock() + + def _query(model): + query = MagicMock() + if model is CallImportEvaluationRow: + query.filter.return_value.all.return_value = [eval_row] + elif model is CallImportRow: + query.filter.return_value.all.return_value = [source_row] + else: + query.filter.return_value.all.return_value = [] + return query + + shard_db.query.side_effect = _query + session_factory = MagicMock(return_value=shard_db) + + with patch( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + return_value=True, + ), patch( + "app.db_sharding.row_ops.shard_id_for_row", + return_value="shard-a", + ), patch( + "app.db_sharding.pool_manager.db_pool_manager" + ) as manager, patch( + "app.services.call_imports.bulk_ops._batch_revoke_celery_task_ids", + ), patch( + "app.api.v1.routes.call_import_evaluations._prepare_source_row_for_retry", + ) as prepare, patch( + "app.api.v1.routes.call_import_evaluations._reset_eval_row_for_retry", + ) as reset: + manager.router = SimpleNamespace(shard_ids=["shard-a"]) + manager.shard_session_factory = MagicMock(return_value=session_factory) + + _persist_evaluation_retry_targets( + catalog_db, + evaluation, + [(eval_row, source_row)], + ) + + prepare.assert_called_once_with(source_row, transcribe_overwrite=False) + reset.assert_called_once() + assert shard_db.commit.call_count == 1 + shard_db.close.assert_called_once() + catalog_db.commit.assert_not_called() + + +def test_persist_evaluation_retry_targets_uses_catalog_when_not_sharded(): + catalog_db = MagicMock() + evaluation = SimpleNamespace(call_import_id=uuid4()) + eval_row = SimpleNamespace(id=uuid4(), celery_task_id=None, status="failed") + source_row = SimpleNamespace(id=uuid4(), row_index=1) + + with patch( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + return_value=False, + ), patch( + "app.services.call_imports.bulk_ops._batch_revoke_celery_task_ids", + ), patch( + "app.api.v1.routes.call_import_evaluations._prepare_source_row_for_retry", + ), patch( + "app.api.v1.routes.call_import_evaluations._reset_eval_row_for_retry", + ): + _persist_evaluation_retry_targets( + catalog_db, + evaluation, + [(eval_row, source_row)], + ) + + catalog_db.commit.assert_called_once() diff --git a/tests/test_services/test_telephony/test_recording_download.py b/tests/test_services/test_telephony/test_recording_download.py index 8ef45ad8..30218415 100644 --- a/tests/test_services/test_telephony/test_recording_download.py +++ b/tests/test_services/test_telephony/test_recording_download.py @@ -9,6 +9,7 @@ CredentialedRecordingThrottledError, ExotelAuthError, ExotelInvalidContentError, + ExotelTransientError, ) from app.services.telephony import recording_download as module @@ -115,16 +116,19 @@ def test_download_recording_url_authenticated_401_is_throttled(monkeypatch): ) -def test_download_recording_url_authenticated_400_is_throttled(monkeypatch): +def test_download_recording_url_authenticated_400_is_transient_without_penalty( + monkeypatch, +): monkeypatch.setattr( module.settings, "RECORDING_URL_ALLOWED_HOST_SUFFIXES", ["exotel.com"], raising=False, ) + penalize = MagicMock(return_value=20) monkeypatch.setattr( "app.workers.concurrency.telephony_credential_rate_limit.penalize_telephony_credential", - lambda *_a, **_kw: 20, + penalize, ) mock_client = MagicMock() @@ -135,13 +139,35 @@ def test_download_recording_url_authenticated_400_is_throttled(monkeypatch): with patch.object(module.socket, "getaddrinfo") as mock_getaddrinfo: mock_getaddrinfo.return_value = [(None, None, None, None, ("52.0.0.1", 0))] with patch.object(module.httpx, "Client", return_value=mock_client): - with pytest.raises(CredentialedRecordingThrottledError, match="400"): + with pytest.raises(ExotelTransientError, match="400"): module.download_recording_url( "https://api.exotel.com/recording.mp3", auth=("user", "pass"), credential_fingerprint="fp-test", ) + penalize.assert_not_called() + + +def test_download_recording_url_public_429_is_transient(monkeypatch): + monkeypatch.setattr( + module.settings, + "RECORDING_URL_ALLOWED_HOST_SUFFIXES", + ["exotel.com"], + raising=False, + ) + + mock_client = MagicMock() + mock_client.get.return_value = _mock_authenticated_response(429, text="Too Many Requests") + mock_client.__enter__.return_value = mock_client + mock_client.__exit__.return_value = False + + with patch.object(module.socket, "getaddrinfo") as mock_getaddrinfo: + mock_getaddrinfo.return_value = [(None, None, None, None, ("52.0.0.1", 0))] + with patch.object(module.httpx, "Client", return_value=mock_client): + with pytest.raises(ExotelTransientError, match="429"): + module.download_public_recording("https://api.exotel.com/recording.mp3") + def test_download_recording_url_public_401_stays_non_retryable(monkeypatch): monkeypatch.setattr( diff --git a/tests/test_workers/test_call_import_unified_dispatch.py b/tests/test_workers/test_call_import_unified_dispatch.py index 49afbdc3..17e21054 100644 --- a/tests/test_workers/test_call_import_unified_dispatch.py +++ b/tests/test_workers/test_call_import_unified_dispatch.py @@ -3,6 +3,8 @@ from types import SimpleNamespace from uuid import uuid4 +import pytest + from app.models.enums import CallImportRowStatus from app.workers.concurrency.eval_dispatch import ( EvalDispatchOutcome, @@ -15,6 +17,7 @@ def _evaluation(**kwargs): defaults = dict( id=uuid4(), + call_import_id=uuid4(), status="running", workspace_id=uuid4(), organization_id=uuid4(), @@ -32,8 +35,15 @@ def _evaluation(**kwargs): def _source_row(**kwargs): + call_import_id = kwargs.pop("call_import_id", None) + if call_import_id is None: + call_import = kwargs.pop("call_import", None) + if call_import is not None: + call_import_id = call_import.id defaults = dict( id=uuid4(), + call_import_id=call_import_id or uuid4(), + row_index=0, recording_s3_key=None, recording_url="https://example.com/rec.mp3", status=CallImportRowStatus.PENDING, @@ -82,16 +92,21 @@ def test_needs_transcribe_after_recording_ready(): def test_try_dispatch_enqueues_import_for_pending_row(monkeypatch): + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: False, + ) evaluation = _evaluation() eval_row = SimpleNamespace(id=uuid4(), celery_task_id=None, status="pending") + call_import = SimpleNamespace( + id=evaluation.call_import_id, + organization_id=evaluation.organization_id, + workspace_id=evaluation.workspace_id, + provider=None, + telephony_integration_id=None, + ) source_row = _source_row( - call_import=SimpleNamespace( - id=uuid4(), - organization_id=evaluation.organization_id, - workspace_id=evaluation.workspace_id, - provider=None, - telephony_integration_id=None, - ), + call_import_id=evaluation.call_import_id, ) captured = {} @@ -99,12 +114,8 @@ class _AsyncResult: id = "import-task-123" monkeypatch.setattr( - "app.workers.tasks.process_call_import_row.process_call_import_row_task", - SimpleNamespace( - apply_async=lambda *a, **kw: ( - captured.update({"apply_async_kwargs": kw}) or _AsyncResult() - ), - ), + "app.workers.tasks.process_call_import_row.process_call_import_row_task.apply_async", + lambda *a, **kw: captured.update({"apply_async_kwargs": kw}) or _AsyncResult(), ) def fake_reserve(**kwargs): @@ -121,6 +132,7 @@ def fake_reserve(**kwargs): evaluation=evaluation, eval_row=eval_row, source_row=source_row, + call_import=call_import, ) assert result == EvalDispatchOutcome("dispatched") diff --git a/tests/test_workers/test_eval_chain_transcribe_cleanup.py b/tests/test_workers/test_eval_chain_transcribe_cleanup.py new file mode 100644 index 00000000..4b4926d3 --- /dev/null +++ b/tests/test_workers/test_eval_chain_transcribe_cleanup.py @@ -0,0 +1,111 @@ +"""Regression tests for eval-chain diarisation cleanup parent rollup.""" + +from __future__ import annotations + +from app.workers.tasks.evaluate_call_import_row_core import ( + _apply_parent_status_from_counters, + reconcile_evaluation_counters, + rollup_parent, +) +from app.workers.tasks.transcribe_call_import_row import ( + _apply_eval_chain_transcribe_cleanup, +) +from tests.test_workers.test_evaluate_call_import_row import ( + _patch_row_location, + _seed, +) + + +def test_eval_chain_transcribe_cleanup_rollup_partial_after_retry( + db_session, + monkeypatch, +): + """Diarisation failures during retry must roll up parent counters to partial.""" + _patch_row_location(monkeypatch, db_session) + + _, _, _, source_rows, evaluation, eval_rows = _seed(db_session, row_count=4) + + eval_rows[0].status = "completed" + eval_rows[1].status = "completed" + eval_rows[2].status = "failed" + eval_rows[3].status = "failed" + evaluation.status = "partial" + evaluation.completed_rows = 2 + evaluation.failed_rows = 2 + db_session.commit() + + for row in (eval_rows[2], eval_rows[3]): + row.status = "pending" + row.error_message = None + row.finished_at = None + db_session.flush() + reconcile_evaluation_counters(db_session, evaluation) + _apply_parent_status_from_counters(evaluation) + db_session.commit() + + assert evaluation.status == "running" + assert evaluation.completed_rows == 2 + assert evaluation.failed_rows == 0 + + eval_rows[2].status = "completed" + rollup_parent( + db_session, + evaluation, + previous_row_status="pending", + new_row_status="completed", + ) + db_session.commit() + db_session.refresh(evaluation) + + assert evaluation.completed_rows == 3 + assert evaluation.failed_rows == 0 + assert evaluation.status == "running" + + source_rows[3].diarised_transcript_status = "failed" + source_rows[3].diarised_transcript_error = "STT timeout" + eval_rows[3].status = "pending" + db_session.commit() + + _apply_eval_chain_transcribe_cleanup(str(eval_rows[3].id)) + + db_session.refresh(evaluation) + db_session.refresh(eval_rows[3]) + + assert eval_rows[3].status == "failed" + assert evaluation.completed_rows == 3 + assert evaluation.failed_rows == 1 + assert evaluation.status == "partial" + + +def test_fail_eval_row_for_import_rollup(db_session, monkeypatch): + """Import fetch failures during eval dispatch must roll up parent counters.""" + from app.models.enums import CallImportRowStatus + from app.workers.concurrency.eval_dispatch import _fail_eval_row_for_import + + _patch_row_location(monkeypatch, db_session) + + _, _, _, source_rows, evaluation, eval_rows = _seed(db_session, row_count=2) + eval_rows[0].status = "completed" + evaluation.status = "running" + evaluation.completed_rows = 1 + evaluation.failed_rows = 0 + source_rows[1].status = CallImportRowStatus.FAILED + source_rows[1].recording_s3_key = None + source_rows[1].error_message = "Recording fetch failed" + db_session.commit() + + _fail_eval_row_for_import( + db_session, + eval_rows[1], + source_rows[1], + catalog_db=db_session, + evaluation=evaluation, + ) + + db_session.refresh(evaluation) + db_session.refresh(eval_rows[1]) + + assert eval_rows[1].status == "failed" + assert evaluation.completed_rows == 1 + assert evaluation.failed_rows == 1 + assert evaluation.status == "partial" diff --git a/tests/test_workers/test_eval_dispatch_import.py b/tests/test_workers/test_eval_dispatch_import.py new file mode 100644 index 00000000..084b6af1 --- /dev/null +++ b/tests/test_workers/test_eval_dispatch_import.py @@ -0,0 +1,167 @@ +"""Eval-chain import failure handling.""" + +from __future__ import annotations + +from uuid import uuid4 + +from app.models.database import ( + CallImport, + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, + Metric, + Organization, + TelephonyIntegration, + Workspace, +) +from app.models.enums import CallImportRowStatus, CallImportStatus, TelephonyProvider +from app.workers.concurrency.eval_dispatch import ( + _try_dispatch_single_row, + recover_eval_row_for_eval_chain, + source_row_import_blocks_eval, +) + + +def _seed_eval_row(db_session): + org = Organization(id=uuid4(), name="Eval Import Org") + db_session.add(org) + workspace = Workspace( + id=uuid4(), + organization_id=org.id, + name="Default", + slug="default", + is_default=True, + ) + db_session.add(workspace) + integration = TelephonyIntegration( + organization_id=org.id, + provider=TelephonyProvider.EXOTEL.value, + auth_id="enc", + auth_token="enc", + voice_app_id="sid", + is_active=True, + ) + db_session.add(integration) + db_session.flush() + + metric = Metric( + id=uuid4(), + organization_id=org.id, + workspace_id=workspace.id, + name="Quality", + metric_type="rating", + trigger="always", + enabled=True, + supported_surfaces=["agent"], + enabled_surfaces=["agent"], + ) + db_session.add(metric) + + call_import = CallImport( + organization_id=org.id, + workspace_id=workspace.id, + provider=TelephonyProvider.EXOTEL.value, + telephony_integration_id=integration.id, + original_filename="batch.csv", + total_rows=1, + completed_rows=0, + failed_rows=0, + status=CallImportStatus.PROCESSING, + ) + db_session.add(call_import) + db_session.flush() + + source_row = CallImportRow( + call_import_id=call_import.id, + organization_id=org.id, + workspace_id=workspace.id, + row_index=0, + conversation_id="call-0", + recording_url="https://api.exotel.com/recordings/0.mp3", + transcript="hello", + status=CallImportRowStatus.FAILED, + error_message="recording URL: old failure", + recording_s3_key="audio/org/test.mp3", + diarised_transcript="Agent: hi\nUser: hello", + diarised_transcript_status="completed", + ) + db_session.add(source_row) + db_session.flush() + + evaluation = CallImportEvaluation( + call_import_id=call_import.id, + organization_id=org.id, + workspace_id=workspace.id, + name="Run", + selected_metric_ids=[str(metric.id)], + status="running", + total_rows=1, + completed_rows=0, + failed_rows=0, + llm_provider="openai", + llm_model="gpt-4o", + ) + db_session.add(evaluation) + db_session.flush() + + eval_row = CallImportEvaluationRow( + evaluation_id=evaluation.id, + call_import_row_id=source_row.id, + status="pending", + ) + db_session.add(eval_row) + db_session.commit() + return call_import, evaluation, eval_row, source_row + + +def test_source_row_import_blocks_eval_false_when_audio_present(db_session): + _, _, _, source_row = _seed_eval_row(db_session) + assert source_row_import_blocks_eval(source_row) is False + + +def test_source_row_import_blocks_eval_true_when_failed_without_audio(db_session): + _, _, _, source_row = _seed_eval_row(db_session) + source_row.recording_s3_key = None + source_row.status = CallImportRowStatus.FAILED + assert source_row_import_blocks_eval(source_row) is True + + +def test_dispatch_does_not_fail_eval_when_source_failed_but_audio_present( + db_session, monkeypatch +): + call_import, evaluation, eval_row, source_row = _seed_eval_row(db_session) + + monkeypatch.setattr( + "app.services.call_imports.evaluation_bulk_op.get_evaluation_bulk_operation", + lambda _evaluation_id: None, + ) + monkeypatch.setattr( + "app.workers.concurrency.eval_dispatch.acquire_eval_slot", + lambda **_kwargs: False, + ) + + outcome = _try_dispatch_single_row( + db=db_session, + evaluation=evaluation, + eval_row=eval_row, + source_row=source_row, + call_import=call_import, + ) + + db_session.refresh(eval_row) + assert eval_row.status == "pending" + assert eval_row.error_message is None + assert outcome.result in {"at_capacity", "dispatched", "skip"} + + +def test_recover_eval_row_for_eval_chain_clears_stale_import_failure(db_session): + _, _, eval_row, _ = _seed_eval_row(db_session) + eval_row.status = "failed" + eval_row.error_message = "Transient: Telephony credential throttled for 15s" + db_session.commit() + + recover_eval_row_for_eval_chain(eval_row) + + assert eval_row.status == "pending" + assert eval_row.error_message is None + assert eval_row.finished_at is None diff --git a/tests/test_workers/test_eval_dispatch_sharding.py b/tests/test_workers/test_eval_dispatch_sharding.py new file mode 100644 index 00000000..0b08c0d3 --- /dev/null +++ b/tests/test_workers/test_eval_dispatch_sharding.py @@ -0,0 +1,102 @@ +"""Sharding commit paths for eval dispatch must bypass catalog-only FK parents.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from app.models.database import ( + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, +) + + +def _make_eval_bundle(): + evaluation = CallImportEvaluation( + id=uuid4(), + call_import_id=uuid4(), + organization_id=uuid4(), + workspace_id=uuid4(), + name="Eval", + selected_metric_ids=[], + status="running", + stt_provider="google", + stt_model="chirp", + ) + source_row = CallImportRow( + id=uuid4(), + call_import_id=evaluation.call_import_id, + organization_id=evaluation.organization_id, + workspace_id=evaluation.workspace_id, + row_index=0, + conversation_id="c1", + recording_url="https://example.com/a.mp3", + recording_s3_key="audio/a.mp3", + status="completed", + ) + eval_row = CallImportEvaluationRow( + id=uuid4(), + evaluation_id=evaluation.id, + call_import_row_id=source_row.id, + workspace_id=evaluation.workspace_id, + status="pending", + celery_task_id="old-task", + ) + return evaluation, eval_row, source_row + + +@patch("app.workers.concurrency.eval_dispatch.build_eval_chain_transcribe_apply_async") +@patch("app.db_sharding.row_ops.shard_row_write_context") +def test_enqueue_eval_chain_uses_shard_write_context(mock_context, mock_build_async): + from app.workers.concurrency.eval_dispatch import ( + enqueue_eval_chain_transcribe_after_import, + ) + + evaluation, eval_row, source_row = _make_eval_bundle() + db = MagicMock() + mock_build_async.return_value = MagicMock(id="transcribe-task-id") + mock_context.return_value.__enter__ = MagicMock(return_value=None) + mock_context.return_value.__exit__ = MagicMock(return_value=False) + + result = enqueue_eval_chain_transcribe_after_import( + db, + evaluation=evaluation, + eval_row=eval_row, + source_row=source_row, + slot_task_id="slot-task-id", + ) + + assert result is True + mock_context.assert_called_once_with(db) + db.flush.assert_called_once() + db.commit.assert_called_once() + assert eval_row.celery_task_id == "transcribe-task-id" + assert source_row.celery_task_id == "slot-task-id" + + +@patch("app.workers.concurrency.eval_dispatch.acquire_eval_slot", return_value=True) +@patch("app.db_sharding.row_ops.shard_row_write_context") +def test_reserve_slot_and_enqueue_uses_shard_write_context( + mock_context, _mock_acquire +): + from app.workers.concurrency.eval_dispatch import _reserve_slot_and_enqueue + + evaluation, eval_row, _source_row = _make_eval_bundle() + db = MagicMock() + mock_context.return_value.__enter__ = MagicMock(return_value=None) + mock_context.return_value.__exit__ = MagicMock(return_value=False) + + def enqueue_fn(_task_id: str): + return MagicMock(id="queued-task-id") + + assert ( + _reserve_slot_and_enqueue( + evaluation=evaluation, + eval_row=eval_row, + db=db, + enqueue_fn=enqueue_fn, + ) + is True + ) + mock_context.assert_called_once_with(db) + db.commit.assert_called_once() + assert eval_row.celery_task_id == "queued-task-id" diff --git a/tests/test_workers/test_eval_queue_routing.py b/tests/test_workers/test_eval_queue_routing.py index abc24533..e3bdd6ec 100644 --- a/tests/test_workers/test_eval_queue_routing.py +++ b/tests/test_workers/test_eval_queue_routing.py @@ -24,7 +24,7 @@ def test_dispatch_evaluation_rows_routes_to_evaluations_queue(): def test_dispatch_fair_eval_rows_routes_to_evaluations_queue(): routes = celery_app.conf.task_routes - assert routes["dispatch_fair_eval_rows"]["queue"] == "celery" + assert routes["dispatch_fair_eval_rows"]["queue"] == "evaluations" def test_manual_transcribe_default_route_routes_to_diarization_queue(): diff --git a/tests/test_workers/test_eval_transcribe_queue_routing.py b/tests/test_workers/test_eval_transcribe_queue_routing.py index 3320c988..b1a6a991 100644 --- a/tests/test_workers/test_eval_transcribe_queue_routing.py +++ b/tests/test_workers/test_eval_transcribe_queue_routing.py @@ -77,7 +77,12 @@ def _fake_async_result(task_id: str): def test_eval_chain_transcribe_uses_diarization_queue( _mock_acquire, stub_worker_task_modules, + monkeypatch, ): + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: False, + ) stub_worker_task_modules.apply_async.side_effect = ( lambda **kwargs: _fake_async_result(kwargs["task_id"]) ) @@ -86,6 +91,7 @@ def test_eval_chain_transcribe_uses_diarization_queue( source_row_id = uuid4() evaluation = SimpleNamespace( id=uuid4(), + call_import_id=uuid4(), workspace_id=uuid4(), organization_id=uuid4(), status="pending", @@ -101,6 +107,8 @@ def test_eval_chain_transcribe_uses_diarization_queue( eval_row = SimpleNamespace(id=eval_row_id, celery_task_id=None) source_row = SimpleNamespace( id=source_row_id, + call_import_id=evaluation.call_import_id, + row_index=0, status=CallImportRowStatus.COMPLETED, recording_s3_key="audio/test.wav", diarised_transcript="", @@ -133,6 +141,7 @@ def test_eval_dispatch_skips_failed_diarization_without_overwrite( ): evaluation = SimpleNamespace( id=uuid4(), + call_import_id=uuid4(), workspace_id=uuid4(), organization_id=uuid4(), status="pending", @@ -148,6 +157,8 @@ def test_eval_dispatch_skips_failed_diarization_without_overwrite( eval_row = SimpleNamespace(id=uuid4(), celery_task_id=None) source_row = SimpleNamespace( id=uuid4(), + call_import_id=evaluation.call_import_id, + row_index=0, status=CallImportRowStatus.COMPLETED, recording_s3_key="audio/test.wav", diarised_transcript="", diff --git a/tests/test_workers/test_evaluate_call_import_row.py b/tests/test_workers/test_evaluate_call_import_row.py index e1f6f20d..f3dd7242 100644 --- a/tests/test_workers/test_evaluate_call_import_row.py +++ b/tests/test_workers/test_evaluate_call_import_row.py @@ -2,7 +2,8 @@ from __future__ import annotations -from uuid import uuid4 +import sys +from uuid import UUID, uuid4 import pytest @@ -123,14 +124,74 @@ def _seed(db_session, *, row_count: int = 1, metric_count: int = 1): return org, call_import, metrics, source_rows, evaluation, eval_rows -def _patch_dependencies(monkeypatch, db_session, *, evaluate_with_llm=None): - """Stub SessionLocal and the LLM helper inside the eval task module.""" - from app.workers.tasks import evaluate_call_import_row as task_module +def _bound_celery_self(): + class _Request: + id = "test-celery-id" + + class _Self: + request = _Request() + max_retries = 2 + + return _Self() + + +def _ensure_bound_task_run(task): + """API tests stub Celery with ``fn.run = fn``; restore bind=True calling.""" + underlying = task.run + + def _run(*args, **kwargs): + try: + return underlying(*args, **kwargs) + except TypeError as exc: + if "required positional argument" not in str(exc): + raise + return underlying(_bound_celery_self(), *args, **kwargs) + + task.run = _run + return task + + +def _patch_row_location(monkeypatch, db_session): + """Route shard-aware row lookup to the pytest session.""" + + def _locate(eval_row_id): + eid = eval_row_id if isinstance(eval_row_id, UUID) else UUID(str(eval_row_id)) + row = ( + db_session.query(CallImportEvaluationRow, CallImportRow) + .join( + CallImportRow, + CallImportRow.id == CallImportEvaluationRow.call_import_row_id, + ) + .filter(CallImportEvaluationRow.id == eid) + .first() + ) + if row is None: + raise LookupError(f"call_import_evaluation_row {eid} not found") + eval_row, source_row = row + wrapped = _NonClosingSession(db_session) + return wrapped, wrapped, eval_row, source_row, "legacy" monkeypatch.setattr( - task_module, "SessionLocal", lambda: _NonClosingSession(db_session) + "app.db_sharding.row_ops.locate_call_import_evaluation_row", + _locate, + ) + monkeypatch.setattr( + "app.db_sharding.eval_rows.locate_call_import_evaluation_row", + _locate, ) + +def _patch_dependencies(monkeypatch, db_session, *, evaluate_with_llm=None): + """Stub SessionLocal and the LLM helper inside the eval task module.""" + import importlib + + session_factory = lambda: _NonClosingSession(db_session) + monkeypatch.setattr("app.database.SessionLocal", session_factory) + _patch_row_location(monkeypatch, db_session) + + task_module = importlib.import_module("app.workers.tasks.evaluate_call_import_row") + monkeypatch.setattr(task_module, "SessionLocal", session_factory) + def _default_eval(*_args, **_kwargs): metrics = _kwargs.get("llm_metrics") or (_args[1] if len(_args) > 1 else []) scores = { @@ -148,10 +209,27 @@ def _default_eval(*_args, **_kwargs): "evaluate_with_llm", evaluate_with_llm or _default_eval, ) + _ensure_bound_task_run(task_module.evaluate_call_import_row_task) return task_module +def _patch_audio_task(monkeypatch, db_session): + import importlib + + _patch_row_location(monkeypatch, db_session) + for mod_name in ( + "app.workers.tasks.evaluate_call_import_row_audio", + "app.workers.tasks.evaluate_call_import_row", + ): + sys.modules.pop(mod_name, None) + audio_module = importlib.import_module( + "app.workers.tasks.evaluate_call_import_row_audio" + ) + _ensure_bound_task_run(audio_module.evaluate_call_import_row_audio_task) + return audio_module + + def test_evaluate_call_import_row_happy_path(db_session, monkeypatch): _, _ci, metrics, _source_rows, evaluation, eval_rows = _seed(db_session) eval_row = eval_rows[0] @@ -1121,11 +1199,7 @@ def test_evaluate_call_import_row_audio_only_completes(db_session, monkeypatch): metrics[0].name = "MOS Score" db_session.commit() - from app.workers.tasks import evaluate_call_import_row_audio as audio_module - - monkeypatch.setattr( - audio_module, "SessionLocal", lambda: _NonClosingSession(db_session) - ) + audio_module = _patch_audio_task(monkeypatch, db_session) monkeypatch.setattr( "app.workers.concurrency.fair_dispatch.finish_eval_work_and_redispatch", lambda *args, **kwargs: None, @@ -1134,6 +1208,12 @@ def test_evaluate_call_import_row_audio_only_completes(db_session, monkeypatch): source_rows[0].recording_s3_key = "audio/key.mp3" db_session.commit() + import importlib + + audio_eval = importlib.import_module( + "app.workers.tasks.helpers.audio_evaluation" + ) + def _fake_audio(*_args, **_kwargs): mid = str(metrics[0].id) return { @@ -1144,9 +1224,12 @@ def _fake_audio(*_args, **_kwargs): } } - monkeypatch.setattr( - "app.workers.tasks.helpers.audio_evaluation.evaluate_audio_metrics", - _fake_audio, + monkeypatch.setattr(audio_eval, "evaluate_audio_metrics", _fake_audio) + + import importlib + + llm_task_module = importlib.import_module( + "app.workers.tasks.evaluate_call_import_row" ) chained = {"called": False} @@ -1155,7 +1238,8 @@ def _no_chain(*_args, **_kwargs): chained["called"] = True monkeypatch.setattr( - "app.workers.tasks.evaluate_call_import_row.evaluate_call_import_row_task.apply_async", + llm_task_module.evaluate_call_import_row_task, + "apply_async", _no_chain, ) @@ -1183,11 +1267,7 @@ def test_evaluate_call_import_row_audio_chain_enqueue_failure_marks_failed( evaluation.status = "running" db_session.commit() - from app.workers.tasks import evaluate_call_import_row_audio as audio_module - - monkeypatch.setattr( - audio_module, "SessionLocal", lambda: _NonClosingSession(db_session) - ) + audio_module = _patch_audio_task(monkeypatch, db_session) redispatched = {"called": False} def _track_redispatch(*_args, **_kwargs): @@ -1198,6 +1278,12 @@ def _track_redispatch(*_args, **_kwargs): _track_redispatch, ) + import importlib + + audio_eval = importlib.import_module( + "app.workers.tasks.helpers.audio_evaluation" + ) + def _fake_audio(*_args, **_kwargs): return { str(audio_metric.id): { @@ -1207,16 +1293,20 @@ def _fake_audio(*_args, **_kwargs): } } - monkeypatch.setattr( - "app.workers.tasks.helpers.audio_evaluation.evaluate_audio_metrics", - _fake_audio, + monkeypatch.setattr(audio_eval, "evaluate_audio_metrics", _fake_audio) + + import importlib + + llm_task_module = importlib.import_module( + "app.workers.tasks.evaluate_call_import_row" ) def _raise_enqueue(*_args, **_kwargs): raise ConnectionError("broker unavailable") monkeypatch.setattr( - "app.workers.tasks.evaluate_call_import_row.evaluate_call_import_row_task.apply_async", + llm_task_module.evaluate_call_import_row_task, + "apply_async", _raise_enqueue, ) diff --git a/tests/test_workers/test_fair_dispatch.py b/tests/test_workers/test_fair_dispatch.py index 5dcac4ff..bcc16a0b 100644 --- a/tests/test_workers/test_fair_dispatch.py +++ b/tests/test_workers/test_fair_dispatch.py @@ -38,6 +38,7 @@ def test_dispatch_batch_interleaves_evaluations_in_same_workspace(): workspace_id = uuid4() eval_a = uuid4() eval_b = uuid4() + call_import_id = uuid4() row_a1 = SimpleNamespace(id=uuid4(), status="pending", celery_task_id=None) row_b1 = SimpleNamespace(id=uuid4(), status="pending", celery_task_id=None) @@ -49,8 +50,12 @@ def test_dispatch_batch_interleaves_evaluations_in_same_workspace(): source_a2 = SimpleNamespace(id=uuid4()) source_b2 = SimpleNamespace(id=uuid4()) - evaluation_a = SimpleNamespace(id=eval_a, status="running") - evaluation_b = SimpleNamespace(id=eval_b, status="running") + evaluation_a = SimpleNamespace( + id=eval_a, status="running", call_import_id=call_import_id + ) + evaluation_b = SimpleNamespace( + id=eval_b, status="running", call_import_id=call_import_id + ) pending_by_eval = { eval_a: [ @@ -64,7 +69,9 @@ def test_dispatch_batch_interleaves_evaluations_in_same_workspace(): } dispatch_order: list[uuid4] = [] - def _pending_rows_for_evaluation(_db, evaluation_id, *, limit): + def _pending_rows_for_evaluation( + _db, evaluation_id, *, limit, shard_cache=None + ): rows = pending_by_eval.get(evaluation_id, []) return rows[:limit] @@ -80,6 +87,13 @@ def _try_dispatch_single_row(**kwargs): return EvalDispatchOutcome("dispatched") db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = SimpleNamespace( + id=call_import_id, + organization_id=uuid4(), + workspace_id=workspace_id, + provider=None, + telephony_integration_id=None, + ) with patch.object( fair_dispatch_module, "_evaluations_with_pending_rows", @@ -121,13 +135,18 @@ def test_dispatch_batch_job2_gets_rows_while_job1_has_backlog(): workspace_id = uuid4() eval_job1 = uuid4() eval_job2 = uuid4() + call_import_id = uuid4() row_job1 = SimpleNamespace(id=uuid4(), status="pending", celery_task_id=None) row_job2 = SimpleNamespace(id=uuid4(), status="pending", celery_task_id=None) source_job1 = SimpleNamespace(id=uuid4()) source_job2 = SimpleNamespace(id=uuid4()) - evaluation_job1 = SimpleNamespace(id=eval_job1, status="running") - evaluation_job2 = SimpleNamespace(id=eval_job2, status="pending") + evaluation_job1 = SimpleNamespace( + id=eval_job1, status="running", call_import_id=call_import_id + ) + evaluation_job2 = SimpleNamespace( + id=eval_job2, status="pending", call_import_id=call_import_id + ) pending_by_eval = { eval_job1: [(row_job1, source_job1, evaluation_job1)] * 5, @@ -135,7 +154,9 @@ def test_dispatch_batch_job2_gets_rows_while_job1_has_backlog(): } dispatched_evaluations: list[uuid4] = [] - def _pending_rows_for_evaluation(_db, evaluation_id, *, limit): + def _pending_rows_for_evaluation( + _db, evaluation_id, *, limit, shard_cache=None + ): rows = pending_by_eval.get(evaluation_id, []) return rows[:limit] @@ -146,6 +167,13 @@ def _try_dispatch_single_row(**kwargs): return EvalDispatchOutcome("dispatched") db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = SimpleNamespace( + id=call_import_id, + organization_id=uuid4(), + workspace_id=workspace_id, + provider=None, + telephony_integration_id=None, + ) with patch.object( fair_dispatch_module, "_evaluations_with_pending_rows", diff --git a/tests/test_workers/test_process_call_import_row.py b/tests/test_workers/test_process_call_import_row.py index e01f9ec8..2a7d3418 100644 --- a/tests/test_workers/test_process_call_import_row.py +++ b/tests/test_workers/test_process_call_import_row.py @@ -69,6 +69,7 @@ def _seed(db_session, *, row_count: int = 1): row = CallImportRow( call_import_id=call_import.id, organization_id=org.id, + workspace_id=workspace.id, row_index=idx, conversation_id=f"call-{idx}", recording_url=f"https://api.exotel.com/recordings/{idx}.mp3", @@ -251,7 +252,7 @@ def _load_task_module(): """ module_name = "app.workers.tasks.process_call_import_row" existing = sys.modules.get(module_name) - if existing is not None and hasattr(existing, "SessionLocal"): + if existing is not None and hasattr(existing, "process_call_import_row_task"): task = getattr(existing, "process_call_import_row_task", None) if isinstance(task, types.FunctionType) or getattr(task, "_bind_task_wrapped", False): return existing @@ -279,11 +280,32 @@ def _load_task_module(): def _patch_dependencies(monkeypatch, db_session, fake_client, fake_s3): - """Wire up SessionLocal + the lazily-imported services the task uses.""" + """Wire up row lookup + the lazily-imported services the task uses.""" + from uuid import UUID + + from app.models.database import CallImportRow + task_module = _load_task_module() + def fake_locate_call_import_row(row_id): + rid = row_id if isinstance(row_id, UUID) else UUID(str(row_id)) + row = db_session.query(CallImportRow).filter(CallImportRow.id == rid).first() + if row is None: + raise LookupError(f"call_import_row {rid} not found") + session = _NonClosingSession(db_session) + return session, session, row, "legacy" + + monkeypatch.setattr( + "app.db_sharding.row_ops.locate_call_import_row", + fake_locate_call_import_row, + ) + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: False, + ) monkeypatch.setattr( - task_module, "SessionLocal", lambda: _NonClosingSession(db_session) + "app.services.call_imports.bulk_ops.is_sharding_enabled", + lambda: False, ) # Telephony service: return our fake client regardless of provider. @@ -427,6 +449,77 @@ def _capture_retry(exc, countdown): assert captured["countdown"] == 18 +def test_eval_chain_import_throttle_does_not_fail_eval_row( + db_session, monkeypatch +): + from app.models.database import CallImportEvaluation, CallImportEvaluationRow + + org, call_import, rows = _seed(db_session, row_count=1) + row = rows[0] + + evaluation = CallImportEvaluation( + call_import_id=call_import.id, + organization_id=org.id, + workspace_id=call_import.workspace_id, + name="Eval", + selected_metric_ids=[], + status="running", + total_rows=1, + completed_rows=0, + failed_rows=0, + ) + db_session.add(evaluation) + db_session.flush() + eval_row = CallImportEvaluationRow( + evaluation_id=evaluation.id, + call_import_row_id=row.id, + status="pending", + ) + db_session.add(eval_row) + db_session.commit() + + class _FingerprintedClient: + _credential_fingerprint = "fp-denied" + + def download_recording(self, _url): + raise AssertionError("download should not run when credit is denied") + + fake_s3 = _FakeS3(enabled=True) + task_module = _patch_dependencies( + monkeypatch, db_session, _FingerprintedClient(), fake_s3 + ) + from app.workers.concurrency.telephony_credential_rate_limit import CreditStatus + + monkeypatch.setattr( + "app.workers.concurrency.telephony_credential_rate_limit.consume_telephony_import_credit", + lambda _fp: CreditStatus(allowed=False, wait_seconds=15, remaining=0), + ) + monkeypatch.setattr( + "app.workers.concurrency.limits.slot_registered_for_task", + lambda _task_id: True, + ) + monkeypatch.setattr( + "app.workers.concurrency.fair_dispatch.finish_eval_work_and_redispatch", + lambda _task_id: None, + ) + + def _capture_retry(exc, countdown): + raise RetryCalled((exc, countdown)) + + monkeypatch.setattr(task_module.process_call_import_row_task, "retry", _capture_retry) + + with pytest.raises(RetryCalled): + task_module.process_call_import_row_task.run( + str(row.id), + _eval_slot_task_id="slot-eval-chain", + run_eval_row_id=str(eval_row.id), + ) + + db_session.refresh(eval_row) + assert eval_row.status == "pending" + assert eval_row.error_message is None + + def test_process_call_import_row_marks_failed_on_auth_error_without_retry(db_session, monkeypatch): _, call_import, rows = _seed(db_session, row_count=1) row = rows[0] @@ -909,9 +1002,14 @@ def test_process_call_import_row_releases_slot_and_redispatches(db_session, monk finish_mock.assert_called_once_with("slot-task-abc") -def test_rollup_parent_status_preserves_deleting(db_session): +def test_rollup_parent_status_preserves_deleting(db_session, monkeypatch): from app.workers.tasks.process_call_import_row import _rollup_parent_status + monkeypatch.setattr( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + lambda: False, + ) + _, call_import, rows = _seed(db_session, row_count=2) call_import.status = CallImportStatus.DELETING rows[0].status = CallImportRowStatus.COMPLETED diff --git a/tests/test_workers/test_process_call_import_row_sharding.py b/tests/test_workers/test_process_call_import_row_sharding.py new file mode 100644 index 00000000..12037ea4 --- /dev/null +++ b/tests/test_workers/test_process_call_import_row_sharding.py @@ -0,0 +1,264 @@ +"""Sharding-specific tests for eval-chain import → diarize handoff.""" + +from unittest.mock import MagicMock +from uuid import uuid4 + +from app.models.database import ( + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, +) +from app.models.enums import CallImportRowStatus +from tests.test_workers.test_process_call_import_row import ( + _FakeExotelClient, + _FakeS3, + _NonClosingSession, + _patch_dependencies, + _seed, +) + + +def _seed_eval_chain(db_session, *, org, call_import, row): + evaluation = CallImportEvaluation( + id=uuid4(), + call_import_id=call_import.id, + organization_id=org.id, + workspace_id=call_import.workspace_id, + name="Eval run", + selected_metric_ids=[], + status="running", + total_rows=1, + ) + db_session.add(evaluation) + eval_row = CallImportEvaluationRow( + id=uuid4(), + evaluation_id=evaluation.id, + call_import_row_id=row.id, + workspace_id=call_import.workspace_id, + status="pending", + celery_task_id="stale-import-task-id", + ) + db_session.add(eval_row) + db_session.commit() + return evaluation, eval_row + + +def test_eval_chain_loads_evaluation_from_catalog_session( + db_session, monkeypatch +): + """With sharding, evaluation headers must be read from catalog, not shard.""" + org, call_import, rows = _seed(db_session, row_count=1) + row = rows[0] + evaluation, eval_row = _seed_eval_chain( + db_session, org=org, call_import=call_import, row=row + ) + + fake_client = _FakeExotelClient(audio=b"hello-audio", content_type="audio/mpeg") + fake_s3 = _FakeS3(enabled=True) + task_module = _patch_dependencies(monkeypatch, db_session, fake_client, fake_s3) + + catalog_session = _NonClosingSession(db_session) + chain_calls = [] + + def fake_locate_call_import_row(row_id): + return db_session, catalog_session, row, "data-shard-01" + + def fake_enqueue( + shard_db, + *, + evaluation, + eval_row, + source_row, + slot_task_id, + restricted_metric_ids=None, + transcribe_overwrite=False, + ): + chain_calls.append( + { + "evaluation_id": evaluation.id, + "eval_row_id": eval_row.id, + "source_row_id": source_row.id, + "slot_task_id": slot_task_id, + "shard_db": shard_db, + } + ) + return True + + monkeypatch.setattr( + "app.db_sharding.row_ops.locate_call_import_row", + fake_locate_call_import_row, + ) + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: True, + ) + monkeypatch.setattr( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + lambda: False, + ) + monkeypatch.setattr( + "app.workers.tasks.process_call_import_row._rollup_parent_status", + lambda _db, _call_import: None, + ) + monkeypatch.setattr( + "app.workers.concurrency.eval_dispatch.enqueue_eval_chain_transcribe_after_import", + fake_enqueue, + ) + monkeypatch.setattr( + "app.workers.concurrency.limits.slot_registered_for_task", + lambda _task_id: False, + ) + + result = task_module.process_call_import_row_task.run( + str(row.id), + _eval_slot_task_id="slot-task-abc", + run_eval_row_id=str(eval_row.id), + ) + + assert result["status"] == "completed" + assert len(chain_calls) == 1 + assert chain_calls[0]["evaluation_id"] == evaluation.id + assert chain_calls[0]["eval_row_id"] == eval_row.id + assert chain_calls[0]["slot_task_id"] == "slot-task-abc" + assert chain_calls[0]["shard_db"] is db_session + + +def test_eval_chain_cleanup_clears_stale_task_ids_and_redispatches( + db_session, monkeypatch +): + """When direct chain misses, clear celery_task_id so fair dispatch can resume.""" + org, call_import, rows = _seed(db_session, row_count=1) + row = rows[0] + _evaluation, eval_row = _seed_eval_chain( + db_session, org=org, call_import=call_import, row=row + ) + row_id = row.id + eval_row_id = eval_row.id + row.celery_task_id = "import-task-id" + db_session.commit() + + fake_client = _FakeExotelClient(audio=b"hello-audio", content_type="audio/mpeg") + fake_s3 = _FakeS3(enabled=True) + task_module = _patch_dependencies(monkeypatch, db_session, fake_client, fake_s3) + + finish_mock = MagicMock() + locate_calls = {"count": 0} + + class _CatalogWithoutEvalHeader(_NonClosingSession): + def query(self, *entities): + if entities and entities[0] is CallImportEvaluation: + empty = MagicMock() + empty.filter.return_value.first.return_value = None + return empty + return self._session.query(*entities) + + catalog_session = _CatalogWithoutEvalHeader(db_session) + + def fake_locate_call_import_row(_row_id): + located_row = ( + db_session.query(CallImportRow).filter(CallImportRow.id == row_id).one() + ) + return db_session, catalog_session, located_row, "data-shard-01" + + def fake_locate_call_import_evaluation_row(_eval_row_id): + locate_calls["count"] += 1 + refreshed_eval_row = ( + db_session.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.id == eval_row_id) + .one() + ) + refreshed_source = ( + db_session.query(CallImportRow) + .filter(CallImportRow.id == row_id) + .one() + ) + return db_session, db_session, refreshed_eval_row, refreshed_source, "legacy" + + chain_called = {"value": False} + + def fake_enqueue(*args, **kwargs): + chain_called["value"] = True + return True + + monkeypatch.setattr( + "app.db_sharding.row_ops.locate_call_import_row", + fake_locate_call_import_row, + ) + monkeypatch.setattr( + "app.db_sharding.row_ops.locate_call_import_evaluation_row", + fake_locate_call_import_evaluation_row, + ) + monkeypatch.setattr( + "app.workers.concurrency.eval_dispatch.enqueue_eval_chain_transcribe_after_import", + fake_enqueue, + ) + monkeypatch.setattr( + "app.workers.concurrency.limits.slot_registered_for_task", + lambda _task_id: True, + ) + monkeypatch.setattr( + "app.workers.concurrency.fair_dispatch.finish_eval_work_and_redispatch", + finish_mock, + ) + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: True, + ) + monkeypatch.setattr( + "app.services.call_imports.bulk_ops.is_sharding_enabled", + lambda: False, + ) + monkeypatch.setattr( + "app.workers.tasks.process_call_import_row._rollup_parent_status", + lambda _db, _call_import: None, + ) + + result = task_module.process_call_import_row_task.run( + str(row_id), + _eval_slot_task_id="slot-task-abc", + run_eval_row_id=str(eval_row_id), + ) + + assert result["status"] == "completed" + assert chain_called["value"] is False + assert locate_calls["count"] == 1 + eval_row_fresh = ( + db_session.query(CallImportEvaluationRow) + .filter(CallImportEvaluationRow.id == eval_row_id) + .one() + ) + row_fresh = db_session.query(CallImportRow).filter(CallImportRow.id == row_id).one() + assert eval_row_fresh.celery_task_id is None + assert row_fresh.celery_task_id is None + assert row_fresh.status == CallImportRowStatus.COMPLETED + finish_mock.assert_called_once_with("slot-task-abc") + + +def test_try_dispatch_single_row_skips_when_bulk_operation_active( + db_session, monkeypatch +): + from app.workers.concurrency.eval_dispatch import _try_dispatch_single_row + + org, call_import, rows = _seed(db_session, row_count=1) + row = rows[0] + evaluation, eval_row = _seed_eval_chain( + db_session, org=org, call_import=call_import, row=row + ) + + monkeypatch.setattr( + "app.services.call_imports.evaluation_bulk_op.get_evaluation_bulk_operation", + lambda _evaluation_id: "abort", + ) + monkeypatch.setattr( + "app.db_sharding.sessions.is_sharding_enabled", + lambda: False, + ) + + outcome = _try_dispatch_single_row( + db=db_session, + evaluation=evaluation, + eval_row=eval_row, + source_row=row, + ) + + assert outcome.result == "skip"