197 lines
7.7 KiB
Python
197 lines
7.7 KiB
Python
"""Batch orchestrator: discover jobs, run them across a bounded thread pool,
|
|
and write one run report per day.
|
|
|
|
One `ThreadPoolExecutor` runs every discovered (field, disease) job so a slow
|
|
crop cannot serialise the others; `pipeline.job.run_job` isolates each job's
|
|
failures from the rest. This module owns everything above that: worklist
|
|
discovery + CLI filters, the `--skip-existing-advice` resume check, the
|
|
09:00-style deadline, and writing the run report.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import date, datetime
|
|
from datetime import time as dtime
|
|
|
|
from pipeline.config import Settings
|
|
from pipeline.db import fetch_all
|
|
from pipeline.job import JobResult, run_job
|
|
from pipeline.output import write_job_output, write_run_report
|
|
from pipeline.prompts import PromptRegistry
|
|
from pipeline.resources import Resources
|
|
from pipeline.vocab import Vocabulary
|
|
from pipeline.worklist import Job, build_jobs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_EXISTING_ADVICE_SQL = "SELECT field_id, disease FROM advice WHERE [date] = ?"
|
|
|
|
|
|
def _parse_deadline(deadline_str: str) -> datetime:
|
|
"""
|
|
Build today's deadline datetime from an "HH:MM" string.
|
|
|
|
Deliberately anchored to the wall-clock date the batch is actually
|
|
running on, not `as_of`: the SLA ("advice ready by 09:00") is about when
|
|
the run happens, and in production `as_of` is always today anyway.
|
|
Anchoring to `as_of` instead would make every historical/backfill run
|
|
(`--as-of` in the past) look like its deadline had already passed.
|
|
"""
|
|
try:
|
|
hour_str, minute_str = deadline_str.split(":", 1)
|
|
return datetime.combine(date.today(), dtime(hour=int(hour_str), minute=int(minute_str)))
|
|
except (ValueError, TypeError) as exc:
|
|
raise ValueError(f"Invalid deadline '{deadline_str}'. Expected HH:MM.") from exc
|
|
|
|
|
|
def filter_jobs(
|
|
jobs: list[Job],
|
|
*,
|
|
crop: str | None = None,
|
|
disease: str | None = None,
|
|
field_ids: list[int] | None = None,
|
|
) -> list[Job]:
|
|
"""Narrow the discovered worklist to CLI-selected crop / disease / field IDs."""
|
|
filtered = jobs
|
|
if crop:
|
|
filtered = [j for j in filtered if j.crop_english.lower() == crop.strip().lower()]
|
|
if disease:
|
|
filtered = [j for j in filtered if j.disease_english.lower() == disease.strip().lower()]
|
|
if field_ids:
|
|
wanted = set(field_ids)
|
|
filtered = [j for j in filtered if j.field_id in wanted]
|
|
return filtered
|
|
|
|
|
|
def _existing_advice_keys(resources: Resources, as_of: date) -> set[tuple[int, str]]:
|
|
"""One bulk SELECT instead of one per job, for --skip-existing-advice."""
|
|
conn = resources.sql_connection()
|
|
rows = fetch_all(conn, _EXISTING_ADVICE_SQL, (as_of,))
|
|
return {(int(row.field_id), str(row.disease)) for row in rows}
|
|
|
|
|
|
def run_batch(
|
|
settings: Settings,
|
|
crop_vocab: Vocabulary,
|
|
as_of: date,
|
|
*,
|
|
crop: str | None = None,
|
|
disease: str | None = None,
|
|
field_ids: list[int] | None = None,
|
|
deadline_override: str | None = None,
|
|
skip_existing_advice: bool = False,
|
|
dry_run: bool = False,
|
|
workers_override: int | None = None,
|
|
) -> int:
|
|
"""Run every discovered job and write the run report. Returns the process exit code."""
|
|
prompts = PromptRegistry(settings.prompts_dir)
|
|
configured_pairs = [(cp.crop, dp.disease) for cp in settings.crops for dp in cp.diseases]
|
|
fallbacks = prompts.validate(configured_pairs)
|
|
if fallbacks:
|
|
logger.warning(
|
|
"%d crop/disease pair(s) have no dedicated advice prompt and will fall back "
|
|
"to prompts/advice/_default/: %s",
|
|
len(fallbacks),
|
|
fallbacks,
|
|
)
|
|
|
|
resources = Resources(settings, prompts)
|
|
try:
|
|
jobs = build_jobs(resources.sql_connection(), settings.crops, settings.worklist, crop_vocab)
|
|
jobs = filter_jobs(jobs, crop=crop, disease=disease, field_ids=field_ids)
|
|
logger.info("Worklist discovered %d job(s) for as_of=%s", len(jobs), as_of.isoformat())
|
|
|
|
if not jobs:
|
|
logger.warning(
|
|
"No jobs to run for as_of=%s; nothing matched the configured "
|
|
"crops/diseases (or the CLI filters).",
|
|
as_of.isoformat(),
|
|
)
|
|
write_run_report(settings.output_dir, as_of, [])
|
|
return 0
|
|
|
|
resources.load_product_index()
|
|
|
|
existing: set[tuple[int, str]] = set()
|
|
if skip_existing_advice:
|
|
existing = _existing_advice_keys(resources, as_of)
|
|
logger.info(
|
|
"%d field/disease pair(s) already have advice for %s and will be skipped",
|
|
len(existing),
|
|
as_of.isoformat(),
|
|
)
|
|
|
|
deadline_str = deadline_override or settings.schedule.deadline
|
|
deadline = _parse_deadline(deadline_str)
|
|
batch_start = datetime.now()
|
|
total_budget_s = (deadline - batch_start).total_seconds()
|
|
warned_80pct = False
|
|
warn_lock = threading.Lock()
|
|
|
|
def _run_one(job: Job) -> JobResult:
|
|
nonlocal warned_80pct
|
|
|
|
now = datetime.now()
|
|
if now >= deadline:
|
|
logger.error(
|
|
"Deadline %s reached; skipping field=%s crop=%s disease=%s",
|
|
deadline.isoformat(timespec="minutes"),
|
|
job.field_id,
|
|
job.crop_english,
|
|
job.disease_english,
|
|
)
|
|
return JobResult(job=job, status="skipped_deadline")
|
|
|
|
if total_budget_s > 0:
|
|
elapsed_s = (now - batch_start).total_seconds()
|
|
if elapsed_s >= 0.8 * total_budget_s:
|
|
with warn_lock:
|
|
already_warned, warned_80pct = warned_80pct, True
|
|
if not already_warned:
|
|
logger.warning(
|
|
"80%% of the time budget before the %s deadline has been "
|
|
"used; some of the remaining %d job(s) may be skipped.",
|
|
deadline_str,
|
|
len(jobs),
|
|
)
|
|
|
|
if (job.field_id, job.disease_english) in existing:
|
|
return JobResult(job=job, status="skipped_existing")
|
|
|
|
return run_job(resources, job, as_of, dry_run=dry_run)
|
|
|
|
workers = workers_override or settings.concurrency.workers
|
|
results: list[JobResult] = []
|
|
with ThreadPoolExecutor(max_workers=max(1, workers)) as pool:
|
|
futures = {pool.submit(_run_one, job): job for job in jobs}
|
|
for future in as_completed(futures):
|
|
job = futures[future]
|
|
try:
|
|
result = future.result()
|
|
except BaseException as exc: # noqa: BLE001 - defense in depth; run_job already isolates failures
|
|
logger.exception(
|
|
"Unexpected error scheduling job field=%s crop=%s disease=%s",
|
|
job.field_id,
|
|
job.crop_english,
|
|
job.disease_english,
|
|
)
|
|
result = JobResult(
|
|
job=job,
|
|
status="failed",
|
|
error_type=type(exc).__name__,
|
|
error_message=str(exc),
|
|
)
|
|
|
|
results.append(result)
|
|
if result.payload is not None:
|
|
write_job_output(settings.output_dir, as_of, result.payload)
|
|
|
|
report = write_run_report(settings.output_dir, as_of, results)
|
|
logger.info("Batch complete: %s", {k: v for k, v in report["totals"].items() if v})
|
|
return 0 if report["totals"].get("failed", 0) == 0 else 1
|
|
finally:
|
|
resources.close()
|