Model/backend/app/modelling/router.py
Khalim Conn-Kowlessar cfedb8e9cf A task that already has sub_tasks refuses re-distribution 🟩
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 12:04:34 +00:00

117 lines
4.2 KiB
Python

"""The Modelling Run Distributor: POST /v1/modelling/trigger-run (ADR-0055).
Accepts a portfolio-scoped modelling request expressed as filters, resolves
them to a concrete property set (ADR-0056), pre-creates one batch sub_task per
SQS message under the app-owned task, and fans the batches out to the
modelling_e2e workers. Never models synchronously; owns nothing after the
fan-out — progress and terminal state roll up from the workers.
"""
import json
from collections.abc import Callable, Iterator
from typing import Any, cast
import boto3
from fastapi import APIRouter, Depends, HTTPException
from sqlmodel import Session
from backend.app.config import get_settings
from backend.app.db.connection import db_engine
from backend.app.dependencies import validate_token
from backend.app.modelling.batching import pack_postcode_batches
from backend.app.modelling.property_filters import (
FilteredProperty,
resolve_filtered_property_ids,
)
from backend.app.modelling.schemas import TriggerRunRequest
from domain.tasks.subtasks import SubTask
from repositories.tasks.subtask_postgres_repository import SubTaskPostgresRepository
# Sends pre-serialised message bodies to the modelling_e2e queue. A seam so
# tests record bodies instead of calling AWS.
MessageSender = Callable[[list[str]], None]
def get_session() -> Iterator[Session]:
with Session(db_engine) as session:
yield session
def get_message_sender() -> MessageSender:
settings = get_settings()
client: Any = cast(Any, boto3.client("sqs", settings.AWS_DEFAULT_REGION)) # pyright: ignore[reportUnknownMemberType]
queue_url = settings.MODELLING_E2E_SQS_URL
def send(bodies: list[str]) -> None:
# send_message_batch caps at 10 entries per call — chunk accordingly.
for start in range(0, len(bodies), 10):
chunk = bodies[start : start + 10]
client.send_message_batch(
QueueUrl=queue_url,
Entries=[
{"Id": str(index), "MessageBody": body}
for index, body in enumerate(chunk)
],
)
return send
router = APIRouter(
prefix="/modelling",
tags=["modelling"],
dependencies=[Depends(validate_token)],
)
@router.post("/trigger-run", status_code=202)
async def trigger_run(
body: TriggerRunRequest,
session: Session = Depends(get_session),
send_messages: MessageSender = Depends(get_message_sender),
) -> dict[str, str]:
subtask_repo = SubTaskPostgresRepository(session)
if subtask_repo.list_by_task(body.task_id):
raise HTTPException(
status_code=409,
detail=(
f"Task {body.task_id} already has sub_tasks — it has been "
"distributed. Check its progress instead of re-triggering."
),
)
properties: list[FilteredProperty] = resolve_filtered_property_ids(
session, body.portfolio_id, body.filters
)
batches = pack_postcode_batches(properties)
# Pre-create one sub_task per (scenario, batch) message under the
# app-owned task, each holding its exact message payload — the fixed
# progress denominator and the batch's re-run recipe (ADR-0055).
subtasks: list[SubTask] = []
payloads: list[dict[str, Any]] = []
for scenario_id in body.scenario_ids:
for batch in batches:
payload: dict[str, Any] = {
"task_id": str(body.task_id),
"property_ids": [p.property_id for p in batch],
"portfolio_id": body.portfolio_id,
"scenario_id": scenario_id,
# ADR-0055 pinned flags: live EPC fetch, live prediction,
# solar fetched only where no stored row exists.
"refetch_epc": True,
"repredict_epc": True,
"refetch_solar": True,
"dry_run": False,
}
subtasks.append(SubTask.create(task_id=body.task_id, inputs=payload))
payloads.append(payload)
subtask_repo.create_many(subtasks)
send_messages(
[
json.dumps({**payload, "subtask_id": str(subtask.id)})
for payload, subtask in zip(payloads, subtasks)
]
)
return {"message": "Modelling Run distributed"}