Model/tests/orchestration/test_task_orchestrator.py
Jun-te Kim b1ff711260 perf(modelling_e2e): batch SubTask bookkeeping to stop per-property writes
Even after batching the data writes, the handler still wrote to the DB per
property through the orchestrator's SubTask bookkeeping: create + start +
complete each self-committed, and _cascade re-listed every sibling and re-saved
the parent on every transition — ~5 writes per property plus an O(N^2) cascade.

- TaskOrchestrator.run_subtasks: create all children in one INSERT, run each
  (failures isolated per child), then persist all terminal states in one bulk
  save and cascade the parent once. Children go WAITING -> terminal; the
  transient IN_PROGRESS row is never written.
- SubTaskRepository.create_many / save_many (bulk INSERT / bulk fetch + update).
- _cascade short-circuits when the Task is already FAILED (terminal) — skips the
  sibling roll-up entirely.
- modelling_e2e handler fans out via run_subtasks instead of per-property
  create_child_subtask + run_subtask.

Per N-property batch the SubTask bookkeeping drops from ~5N writes + an O(N^2)
cascade to ~2 writes + 1 cascade.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-24 19:26:42 +00:00

281 lines
9.2 KiB
Python

from collections.abc import Iterator
from dataclasses import dataclass
import pytest
from sqlalchemy import Engine
from sqlmodel import Session
from domain.tasks.subtasks import SubTask, SubTaskStatus
from domain.tasks.tasks import Source, TaskStatus
from orchestration.task_orchestrator import TaskOrchestrator
from repositories.tasks.subtask_postgres_repository import SubTaskPostgresRepository
from repositories.tasks.task_postgres_repository import TaskPostgresRepository
@dataclass
class Harness:
orchestrator: TaskOrchestrator
tasks: TaskPostgresRepository
subtasks: SubTaskPostgresRepository
@pytest.fixture
def harness(db_engine: Engine) -> Iterator[Harness]:
with Session(db_engine) as session:
tasks = TaskPostgresRepository(session=session)
subtasks = SubTaskPostgresRepository(session=session)
yield Harness(
orchestrator=TaskOrchestrator(task_repo=tasks, subtask_repo=subtasks),
tasks=tasks,
subtasks=subtasks,
)
def test_create_task_with_subtask_creates_both_in_waiting(
harness: Harness,
) -> None:
# act
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test",
inputs={"foo": "bar"},
source=Source.PORTFOLIO,
source_id="abc",
)
# assert
assert task.status is TaskStatus.WAITING
assert subtask.status is SubTaskStatus.WAITING
assert subtask.task_id == task.id
assert subtask.inputs == {"foo": "bar"}
def test_start_subtask_cascades_to_in_progress(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
started = harness.orchestrator.start_subtask(
subtask.id, cloud_logs_url="https://example/log"
)
# assert
assert started.status is SubTaskStatus.IN_PROGRESS
assert started.cloud_logs_url == "https://example/log"
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_complete_subtask_cascades_to_complete(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.start_subtask(subtask.id)
# act
harness.orchestrator.complete_subtask(subtask.id, {"value": 42})
# assert
done_subtask = harness.subtasks.get(subtask.id)
done_task = harness.tasks.get(task.id)
assert done_subtask.outputs == {"result": {"value": 42}}
assert done_task.status is TaskStatus.COMPLETE
assert done_task.job_completed is not None
def test_fail_subtask_cascades_to_failed(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
harness.orchestrator.fail_subtask(subtask.id, RuntimeError("boom"))
# assert
failed_subtask = harness.subtasks.get(subtask.id)
failed_task = harness.tasks.get(task.id)
assert failed_subtask.outputs == {"error": "boom"}
assert failed_task.status is TaskStatus.FAILED
def test_failed_subtask_locks_task_failed_even_with_others_complete(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = SubTask.create(task_id=task.id)
harness.subtasks.create(second)
# act
harness.orchestrator.complete_subtask(first.id)
harness.orchestrator.fail_subtask(second.id, RuntimeError("nope"))
# assert
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_mixed_complete_and_in_progress_keeps_task_in_progress(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = SubTask.create(task_id=task.id)
harness.subtasks.create(second)
# act
harness.orchestrator.complete_subtask(first.id)
harness.orchestrator.start_subtask(second.id)
# assert
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_run_subtask_happy_path_returns_result_and_cascades_complete(
harness: Harness,
) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
result = harness.orchestrator.run_subtask(subtask.id, work=lambda: {"answer": 42})
# assert
assert result == {"answer": 42}
assert harness.subtasks.get(subtask.id).status is SubTaskStatus.COMPLETE
assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE
def test_create_child_subtask_adds_waiting_child_without_changing_parent_status(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.start_subtask(first.id)
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
# act
child = harness.orchestrator.create_child_subtask(
task.id, inputs={"split": "a"}
)
# assert
persisted_child = harness.subtasks.get(child.id)
assert persisted_child.task_id == task.id
assert persisted_child.status is SubTaskStatus.WAITING
assert persisted_child.inputs == {"split": "a"}
assert persisted_child.id != first.id
# Cascade is a no-op: parent stays IN_PROGRESS.
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_run_subtask_failing_work_marks_failed_and_reraises(
harness: Harness,
) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
def boom() -> None:
raise RuntimeError("boom")
# act / assert
with pytest.raises(RuntimeError, match="boom"):
harness.orchestrator.run_subtask(subtask.id, work=boom)
assert harness.subtasks.get(subtask.id).status is SubTaskStatus.FAILED
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_run_subtasks_creates_runs_and_completes_a_whole_batch(
harness: Harness,
) -> None:
"""run_subtasks fans the parent into one child per item, runs each, and
leaves every child COMPLETE with the parent cascaded to COMPLETE."""
# arrange — a parent task whose coordinator subtask is done
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.complete_subtask(coordinator.id)
# act
results = harness.orchestrator.run_subtasks(
task.id,
[{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
work=lambda st: (st.inputs or {})["property_id"] * 10,
)
# assert — results in order, all children COMPLETE, parent COMPLETE
assert results == [10, 20, 30]
children = [s for s in harness.subtasks.list_by_task(task.id) if s.id != coordinator.id]
assert len(children) == 3
assert all(c.status is SubTaskStatus.COMPLETE for c in children)
assert {(c.inputs or {})["property_id"] for c in children} == {1, 2, 3}
assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE
def test_run_subtasks_isolates_a_failing_item_and_continues(
harness: Harness,
) -> None:
"""A raising item is marked FAILED with its error recorded, its siblings still
complete (result None for the failure), and the parent cascades to FAILED."""
# arrange
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.complete_subtask(coordinator.id)
def _work(st: SubTask) -> int:
pid = (st.inputs or {})["property_id"]
if pid == 2:
raise RuntimeError("property 2 exploded")
return pid
# act — does NOT raise; failure is isolated
results = harness.orchestrator.run_subtasks(
task.id,
[{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
work=_work,
)
# assert
assert results == [1, None, 3]
children = {
(c.inputs or {})["property_id"]: c
for c in harness.subtasks.list_by_task(task.id)
if c.id != coordinator.id
}
assert children[1].status is SubTaskStatus.COMPLETE
assert children[3].status is SubTaskStatus.COMPLETE
assert children[2].status is SubTaskStatus.FAILED
assert children[2].outputs == {"error": "property 2 exploded"}
# any FAILED child → parent FAILED
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_cascade_short_circuits_once_task_already_failed(harness: Harness) -> None:
"""Once the Task is FAILED, completing another SubTask leaves it FAILED — the
terminal state is not recomputed away."""
# arrange — two children; fail the first so the task is FAILED
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = harness.orchestrator.create_child_subtask(task.id)
harness.orchestrator.fail_subtask(coordinator.id, RuntimeError("boom"))
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
# act — complete the other child
harness.orchestrator.complete_subtask(second.id)
# assert — task stays FAILED (cascade short-circuited on the terminal state)
assert harness.tasks.get(task.id).status is TaskStatus.FAILED