Spaces:
Running on Zero
Running on Zero
| """Batch orchestration for the GCMD classifier MVP.""" | |
| from __future__ import annotations | |
| import logging | |
| import time | |
| from collections.abc import Callable | |
| from datetime import datetime, timezone | |
| from typing import Any | |
| from pydantic import BaseModel, ConfigDict, Field | |
| from gcmd_classifier.config import ModelSettings | |
| from gcmd_classifier.llm.base import ModelClient | |
| from gcmd_classifier.logging_config import get_logger, log_event | |
| from gcmd_classifier.models import ( | |
| ArticleClassificationOutcome, | |
| ArticleLoadResult, | |
| ArticleProcessingStatus, | |
| ArticleRecord, | |
| ArticleResult, | |
| ArticleValidationIssue, | |
| OutputError, | |
| RunSummary, | |
| ) | |
| from gcmd_classifier.persistence.cache import ( | |
| ArticleResultCache, | |
| build_cache_identity, | |
| ) | |
| from gcmd_classifier.persistence.json_store import JsonResultStore | |
| from gcmd_classifier.pipeline.service import ( | |
| APPLICATION_VERSION, | |
| classify_article, | |
| failed_article_result, | |
| ) | |
| from gcmd_classifier.vocabulary.index import VocabularyIndex | |
| ArticleStartCallback = Callable[[int, int, ArticleRecord], None] | |
| ArticleFinishCallback = Callable[[int, int, ArticleResult, float], None] | |
| class BatchRunResult(BaseModel): | |
| """Structured batch result with article outputs and run summary.""" | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| results: tuple[ArticleResult, ...] = Field(default_factory=tuple) | |
| summary: RunSummary | |
| def run_batch( | |
| *, | |
| article_load_result: ArticleLoadResult, | |
| vocabulary: VocabularyIndex, | |
| model_client: ModelClient, | |
| settings: ModelSettings, | |
| store: JsonResultStore, | |
| cache: ArticleResultCache | None = None, | |
| force_reprocess: bool = False, | |
| run_id: str | None = None, | |
| application_version: str = APPLICATION_VERSION, | |
| relevant_config: dict[str, Any] | None = None, | |
| logger: logging.Logger | None = None, | |
| on_article_start: ArticleStartCallback | None = None, | |
| on_article_finish: ArticleFinishCallback | None = None, | |
| ) -> BatchRunResult: | |
| """Process valid articles in source order while preserving completed results.""" | |
| active_logger = logger or get_logger(__name__) | |
| started = time.perf_counter() | |
| started_at = _now_iso() | |
| actual_run_id = run_id or f"run-{started_at}" | |
| invalid_errors = tuple(_issue_to_error(issue) for issue in article_load_result.errors) | |
| results: list[ArticleResult] = [] | |
| cache_hits = 0 | |
| cache_misses = 0 | |
| log_event( | |
| active_logger, | |
| "batch_started", | |
| stage="batch", | |
| source_records=article_load_result.source_count, | |
| valid_records=len(article_load_result.articles), | |
| invalid_records=_invalid_record_count(article_load_result.errors), | |
| ) | |
| total_articles = len(article_load_result.articles) | |
| for index, article in enumerate(article_load_result.articles): | |
| article_started = time.perf_counter() | |
| if on_article_start is not None: | |
| on_article_start(index, total_articles, article) | |
| try: | |
| result: ArticleResult | None = None | |
| identity = build_cache_identity( | |
| article=article, | |
| vocabulary=vocabulary, | |
| settings=settings, | |
| application_version=application_version, | |
| relevant_config=relevant_config, | |
| ) | |
| if cache is not None and not force_reprocess: | |
| result = cache.get(identity) | |
| if result is not None: | |
| cache_hits += 1 | |
| log_event( | |
| active_logger, | |
| "cache_hit", | |
| DOI=article.DOI, | |
| index=index, | |
| stage="cache", | |
| cache_key=identity.cache_key, | |
| ) | |
| else: | |
| cache_misses += 1 | |
| log_event( | |
| active_logger, | |
| "cache_miss", | |
| DOI=article.DOI, | |
| index=index, | |
| stage="cache", | |
| cache_key=identity.cache_key, | |
| ) | |
| elif cache is not None: | |
| cache_misses += 1 | |
| if result is None: | |
| result = classify_article( | |
| article=article, | |
| vocabulary=vocabulary, | |
| model_client=model_client, | |
| settings=settings, | |
| run_id=actual_run_id, | |
| application_version=application_version, | |
| relevant_config=relevant_config, | |
| logger=active_logger, | |
| ) | |
| if cache is not None: | |
| cache.put(identity, result) | |
| else: | |
| result = result.model_copy( | |
| update={ | |
| "processing_metadata": result.processing_metadata.model_copy( | |
| update={"run_id": actual_run_id} | |
| ) | |
| } | |
| ) | |
| except Exception as exc: | |
| result = failed_article_result( | |
| article=article, | |
| error=OutputError( | |
| code=exc.__class__.__name__, | |
| message=str(exc), | |
| stage="batch", | |
| DOI=article.DOI, | |
| ), | |
| vocabulary=vocabulary, | |
| settings=settings, | |
| run_id=actual_run_id, | |
| application_version=application_version, | |
| relevant_config=relevant_config, | |
| ) | |
| log_event( | |
| active_logger, | |
| "article_failed", | |
| DOI=article.DOI, | |
| index=index, | |
| stage="batch", | |
| error_code=exc.__class__.__name__, | |
| ) | |
| store.save_article_result(result) | |
| results.append(result) | |
| if on_article_finish is not None: | |
| on_article_finish(index, total_articles, result, time.perf_counter() - article_started) | |
| summary = _run_summary( | |
| run_id=actual_run_id, | |
| started_at=started_at, | |
| completed_at=_now_iso(), | |
| duration_seconds=time.perf_counter() - started, | |
| article_load_result=article_load_result, | |
| results=tuple(results), | |
| invalid_errors=invalid_errors, | |
| cache_hits=cache_hits, | |
| cache_misses=cache_misses, | |
| ) | |
| store.write_consolidated(results=tuple(results), summary=summary) | |
| log_event( | |
| active_logger, | |
| "batch_finished", | |
| stage="batch", | |
| run_id=actual_run_id, | |
| processed_articles=len(results), | |
| cache_hits=cache_hits, | |
| cache_misses=cache_misses, | |
| errors=summary.total_errors, | |
| ) | |
| return BatchRunResult(results=tuple(results), summary=summary) | |
| def _run_summary( | |
| *, | |
| run_id: str, | |
| started_at: str, | |
| completed_at: str, | |
| duration_seconds: float, | |
| article_load_result: ArticleLoadResult, | |
| results: tuple[ArticleResult, ...], | |
| invalid_errors: tuple[OutputError, ...], | |
| cache_hits: int, | |
| cache_misses: int, | |
| ) -> RunSummary: | |
| warnings_count = sum(len(result.warnings) for result in results) + sum( | |
| len(record.warnings) for result in results for record in result.classifications | |
| ) | |
| result_errors = tuple(error for result in results for error in result.errors) | |
| errors = (*invalid_errors, *result_errors) | |
| completed = sum( | |
| result.processing_status is ArticleProcessingStatus.COMPLETED for result in results | |
| ) | |
| partial = sum(result.processing_status is ArticleProcessingStatus.PARTIAL for result in results) | |
| failed = sum(result.processing_status is ArticleProcessingStatus.FAILED for result in results) | |
| skipped = sum(result.processing_status is ArticleProcessingStatus.SKIPPED for result in results) | |
| classifications = sum(len(result.classifications) for result in results) | |
| processed = len(results) | |
| return RunSummary( | |
| run_id=run_id, | |
| started_at=started_at, | |
| completed_at=completed_at, | |
| articles_received=article_load_result.source_count, | |
| valid_article_records=len(article_load_result.articles), | |
| invalid_source_records=_invalid_record_count(article_load_result.errors), | |
| processed_articles=processed, | |
| cache_hits=cache_hits, | |
| cache_misses=cache_misses, | |
| total_warnings=warnings_count, | |
| total_errors=len(errors), | |
| duration_seconds=duration_seconds, | |
| articles_completed=completed, | |
| articles_partial=partial, | |
| articles_failed=failed, | |
| articles_skipped=skipped, | |
| articles_not_classified=sum( | |
| result.classification_outcome is ArticleClassificationOutcome.NOT_CLASSIFIED | |
| for result in results | |
| ), | |
| articles_requiring_review=sum( | |
| any(record.review_required for record in result.classifications) for result in results | |
| ), | |
| accepted_classifications=classifications, | |
| average_classifications_per_article=(classifications / processed if processed else None), | |
| average_processing_time_seconds=(duration_seconds / processed if processed else None), | |
| total_model_calls=sum( | |
| result.processing_metadata.model_calls or 0 | |
| for result in results | |
| if not result.processing_metadata.cache_used | |
| ), | |
| errors=errors, | |
| ) | |
| def _issue_to_error(issue: ArticleValidationIssue) -> OutputError: | |
| return OutputError( | |
| code=issue.code, | |
| message=issue.message, | |
| stage="article_loading", | |
| details={"field": issue.field}, | |
| index=issue.index, | |
| DOI=issue.DOI, | |
| ) | |
| def _invalid_record_count(issues: tuple[ArticleValidationIssue, ...]) -> int: | |
| indices = {issue.index for issue in issues if issue.index is not None} | |
| return len(indices) | |
| def _now_iso() -> str: | |
| return datetime.now(timezone.utc).isoformat() # noqa: UP017 | |