Skip to content

Commit cf2cb47

Browse files
w7-mgfcodew7-mgfcode
committed
feat(api,db): forecast champion selector slice B — async comparison & results (#360)
1 parent e7f4db7 commit cf2cb47

33 files changed

Lines changed: 3206 additions & 45 deletions

‎.env.example‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,5 +126,15 @@ BATCH_GLOBAL_MAX_PARALLEL=4
126126
# mid-call, so a long fit can stall the drain.
127127
BATCH_CANCEL_DRAIN_TIMEOUT_SECONDS=30
128128

129+
# Model selection (champion selector) async runner (Slice B)
130+
# Hard upper bound on concurrent candidate backtests across all active selection
131+
# runs on this host. Effective parallelism per run is min(this, candidates).
132+
# Set to 1 for sequential execution. Requires uvicorn restart to apply.
133+
MODEL_SELECTION_GLOBAL_MAX_PARALLEL=4
134+
# Max seconds DELETE /model-selection/{id} waits for in-flight candidates to
135+
# drain before returning RFC 7807 504. sklearn / LightGBM fits are uncancellable
136+
# mid-call, so a long fit can stall the drain.
137+
MODEL_SELECTION_CANCEL_DRAIN_TIMEOUT_SECONDS=30
138+
129139
# Frontend (Vite)
130140
VITE_API_BASE_URL=http://localhost:8123

‎alembic/versions/d3e4f5a6b7c8_add_model_selection_candidate_and_progress.py‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
"""add model_selection_candidate and async progress columns
2+
3+
Revision ID: d3e4f5a6b7c8
4+
Revises: b667d321603c
5+
Create Date: 2026-06-01 09:30:00.000000
6+
7+
Slice B of the Forecast Champion Selector (issue #360). Converts the selection
8+
run into a DB-backed async LRO:
9+
10+
- creates ``model_selection_candidate`` (one row per candidate, FK CASCADE to
11+
``model_selection_run.selection_id``) carrying per-candidate status, result
12+
JSONB, error, and timing — the live-progress + audit surface;
13+
- adds ``started_at`` + the four final count columns to ``model_selection_run``;
14+
- widens the run status CheckConstraint to include ``'cancelled'`` (forward-only
15+
drop + recreate of the named constraint).
16+
17+
Mirrors ``c1d2e3f40512_create_batch_tables`` for JSONB / index / FK style.
18+
"""
19+
20+
from collections.abc import Sequence
21+
22+
import sqlalchemy as sa
23+
from sqlalchemy.dialects import postgresql
24+
25+
from alembic import op
26+
27+
# revision identifiers, used by Alembic.
28+
revision: str = "d3e4f5a6b7c8"
29+
down_revision: str | None = "b667d321603c"
30+
branch_labels: str | Sequence[str] | None = None
31+
depends_on: str | Sequence[str] | None = None
32+
33+
_OLD_RUN_STATUS = "status IN ('pending', 'running', 'completed', 'partial', 'failed')"
34+
_NEW_RUN_STATUS = (
35+
"status IN ('pending', 'running', 'completed', 'partial', 'failed', 'cancelled')"
36+
)
37+
38+
39+
def upgrade() -> None:
40+
"""Apply migration."""
41+
# ------------------------------------------------------------------
42+
# 1. Widen the run status CheckConstraint to include 'cancelled'.
43+
# ------------------------------------------------------------------
44+
op.drop_constraint(
45+
"ck_model_selection_run_valid_status",
46+
"model_selection_run",
47+
type_="check",
48+
)
49+
op.create_check_constraint(
50+
"ck_model_selection_run_valid_status",
51+
"model_selection_run",
52+
_NEW_RUN_STATUS,
53+
)
54+
55+
# ------------------------------------------------------------------
56+
# 2. Additive progress columns on the parent run.
57+
# ------------------------------------------------------------------
58+
op.add_column(
59+
"model_selection_run",
60+
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
61+
)
62+
op.add_column(
63+
"model_selection_run",
64+
sa.Column("total_candidates", sa.Integer(), nullable=False, server_default="0"),
65+
)
66+
op.add_column(
67+
"model_selection_run",
68+
sa.Column(
69+
"completed_candidates", sa.Integer(), nullable=False, server_default="0"
70+
),
71+
)
72+
op.add_column(
73+
"model_selection_run",
74+
sa.Column("failed_candidates", sa.Integer(), nullable=False, server_default="0"),
75+
)
76+
op.add_column(
77+
"model_selection_run",
78+
sa.Column(
79+
"cancelled_candidates", sa.Integer(), nullable=False, server_default="0"
80+
),
81+
)
82+
83+
# ------------------------------------------------------------------
84+
# 3. Per-candidate execution child table (FK CASCADE on selection_id).
85+
# ------------------------------------------------------------------
86+
op.create_table(
87+
"model_selection_candidate",
88+
sa.Column("id", sa.Integer(), nullable=False),
89+
sa.Column("candidate_id", sa.String(length=32), nullable=False),
90+
sa.Column("selection_id", sa.String(length=32), nullable=False),
91+
sa.Column("ordinal", sa.Integer(), nullable=False),
92+
sa.Column("model_type", sa.String(length=40), nullable=False),
93+
sa.Column("params", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
94+
sa.Column("status", sa.String(length=20), nullable=False),
95+
sa.Column("result", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
96+
sa.Column("error_message", sa.String(length=2000), nullable=True),
97+
sa.Column("error_type", sa.String(length=100), nullable=True),
98+
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
99+
sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
100+
sa.Column("duration_ms", sa.Integer(), nullable=True),
101+
sa.Column(
102+
"created_at",
103+
sa.DateTime(timezone=True),
104+
server_default=sa.text("now()"),
105+
nullable=False,
106+
),
107+
sa.Column(
108+
"updated_at",
109+
sa.DateTime(timezone=True),
110+
server_default=sa.text("now()"),
111+
nullable=False,
112+
),
113+
sa.CheckConstraint(
114+
"status IN ('pending', 'running', 'completed', 'failed', 'cancelled')",
115+
name="ck_model_selection_candidate_valid_status",
116+
),
117+
sa.ForeignKeyConstraint(
118+
["selection_id"],
119+
["model_selection_run.selection_id"],
120+
ondelete="CASCADE",
121+
),
122+
sa.PrimaryKeyConstraint("id"),
123+
)
124+
op.create_index(
125+
op.f("ix_model_selection_candidate_candidate_id"),
126+
"model_selection_candidate",
127+
["candidate_id"],
128+
unique=True,
129+
)
130+
op.create_index(
131+
op.f("ix_model_selection_candidate_selection_id"),
132+
"model_selection_candidate",
133+
["selection_id"],
134+
unique=False,
135+
)
136+
op.create_index(
137+
op.f("ix_model_selection_candidate_status"),
138+
"model_selection_candidate",
139+
["status"],
140+
unique=False,
141+
)
142+
op.create_index(
143+
"ix_model_selection_candidate_selection_status",
144+
"model_selection_candidate",
145+
["selection_id", "status"],
146+
unique=False,
147+
)
148+
149+
150+
def downgrade() -> None:
151+
"""Revert migration."""
152+
op.drop_index(
153+
"ix_model_selection_candidate_selection_status",
154+
table_name="model_selection_candidate",
155+
)
156+
op.drop_index(
157+
op.f("ix_model_selection_candidate_status"),
158+
table_name="model_selection_candidate",
159+
)
160+
op.drop_index(
161+
op.f("ix_model_selection_candidate_selection_id"),
162+
table_name="model_selection_candidate",
163+
)
164+
op.drop_index(
165+
op.f("ix_model_selection_candidate_candidate_id"),
166+
table_name="model_selection_candidate",
167+
)
168+
op.drop_table("model_selection_candidate")
169+
170+
op.drop_column("model_selection_run", "cancelled_candidates")
171+
op.drop_column("model_selection_run", "failed_candidates")
172+
op.drop_column("model_selection_run", "completed_candidates")
173+
op.drop_column("model_selection_run", "total_candidates")
174+
op.drop_column("model_selection_run", "started_at")
175+
176+
op.drop_constraint(
177+
"ck_model_selection_run_valid_status",
178+
"model_selection_run",
179+
type_="check",
180+
)
181+
op.create_check_constraint(
182+
"ck_model_selection_run_valid_status",
183+
"model_selection_run",
184+
_OLD_RUN_STATUS,
185+
)

‎app/core/config.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,17 @@ class Settings(BaseSettings):
134134
# are uncancellable mid-call, so a long fit can stall the drain.
135135
batch_cancel_drain_timeout_seconds: int = 30
136136

137+
# Model selection (champion selector) async runner (Slice B) — mirrors the
138+
# batch runner. Hard upper bound on concurrent candidate backtests across
139+
# all active selection runs on this host; sized for the same Postgres pool
140+
# (pool_size=5, max_overflow=10). Setting this to 1 makes the runner
141+
# sequential. Env override: MODEL_SELECTION_GLOBAL_MAX_PARALLEL=8 (restart).
142+
model_selection_global_max_parallel: int = 4
143+
# Max seconds DELETE /model-selection/{id} waits for in-flight candidates to
144+
# settle before returning RFC 7807 504. In-flight sklearn/LightGBM fits are
145+
# uncancellable mid-call, so a long fit can stall the drain.
146+
model_selection_cancel_drain_timeout_seconds: int = 30
147+
137148
# RAG Embedding Configuration
138149
rag_embedding_provider: Literal["openai", "ollama"] = "openai"
139150
openai_api_key: str = ""

‎app/core/tests/test_config.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,15 @@ def test_settings_has_defaults(monkeypatch):
2323
assert settings.api_port == 8123
2424

2525

26+
def test_model_selection_runner_defaults(monkeypatch):
27+
"""Slice B async-runner settings default to the batch-mirrored values."""
28+
monkeypatch.delenv("MODEL_SELECTION_GLOBAL_MAX_PARALLEL", raising=False)
29+
monkeypatch.delenv("MODEL_SELECTION_CANCEL_DRAIN_TIMEOUT_SECONDS", raising=False)
30+
settings = Settings(_env_file=None)
31+
assert settings.model_selection_global_max_parallel == 4
32+
assert settings.model_selection_cancel_drain_timeout_seconds == 30
33+
34+
2635
def test_settings_is_development_property():
2736
"""is_development should return True for development env."""
2837
settings = Settings(app_env="development")

‎app/features/model_selection/models.py‎

Lines changed: 83 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from enum import Enum
1616
from typing import Any
1717

18-
from sqlalchemy import CheckConstraint, Date, DateTime, Index, Integer, String
18+
from sqlalchemy import CheckConstraint, Date, DateTime, ForeignKey, Index, Integer, String
1919
from sqlalchemy.dialects.postgresql import JSONB
2020
from sqlalchemy.orm import Mapped, mapped_column
2121

@@ -27,17 +27,42 @@ class ModelSelectionStatus(str, Enum):
2727
"""Lifecycle states of a selection run.
2828
2929
Transitions:
30-
- PENDING -> RUNNING -> {COMPLETED, PARTIAL, FAILED}
31-
- PARTIAL fires when >=1 candidate succeeded AND >=1 candidate failed.
30+
- PENDING -> RUNNING -> {COMPLETED, PARTIAL, FAILED, CANCELLED}
31+
- PARTIAL fires when >=1 candidate succeeded AND >=1 candidate failed/cancelled.
3232
- FAILED fires when availability is unusable (fail-fast) OR every
3333
candidate's backtest errored (no valid winner).
34+
- CANCELLED (Slice B) fires when a cancel drained before any candidate
35+
reached a non-cancelled terminal state.
3436
"""
3537

3638
PENDING = "pending"
3739
RUNNING = "running"
3840
COMPLETED = "completed"
3941
PARTIAL = "partial"
4042
FAILED = "failed"
43+
CANCELLED = "cancelled"
44+
45+
46+
# Statuses a selection run cannot transition out of — the DELETE-route 409 set
47+
# (Slice B). Mirrors ``batch.models.TERMINAL_BATCH_STATES``.
48+
TERMINAL_SELECTION_STATES: frozenset[str] = frozenset(
49+
{
50+
ModelSelectionStatus.COMPLETED.value,
51+
ModelSelectionStatus.PARTIAL.value,
52+
ModelSelectionStatus.FAILED.value,
53+
ModelSelectionStatus.CANCELLED.value,
54+
}
55+
)
56+
57+
58+
class CandidateStatus(str, Enum):
59+
"""Per-candidate execution states inside an async selection run (Slice B)."""
60+
61+
PENDING = "pending"
62+
RUNNING = "running"
63+
COMPLETED = "completed"
64+
FAILED = "failed"
65+
CANCELLED = "cancelled"
4166

4267

4368
class ModelSelectionRun(TimestampMixin, Base):
@@ -74,13 +99,21 @@ class ModelSelectionRun(TimestampMixin, Base):
7499
forecast_result: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True)
75100
business_summary: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True)
76101
error_message: Mapped[str | None] = mapped_column(String(2000), nullable=True)
102+
# Slice B (async) — set when the run starts executing; the four count
103+
# columns cache the FINAL per-status candidate tally written once at settle
104+
# (live progress is derived from a GROUP BY over the child rows).
105+
started_at: Mapped[_dt.datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
106+
total_candidates: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
107+
completed_candidates: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
108+
failed_candidates: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
109+
cancelled_candidates: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
77110
completed_at: Mapped[_dt.datetime | None] = mapped_column(
78111
DateTime(timezone=True), nullable=True
79112
)
80113

81114
__table_args__ = (
82115
CheckConstraint(
83-
"status IN ('pending', 'running', 'completed', 'partial', 'failed')",
116+
"status IN ('pending', 'running', 'completed', 'partial', 'failed', 'cancelled')",
84117
name="ck_model_selection_run_valid_status",
85118
),
86119
Index(
@@ -91,3 +124,49 @@ class ModelSelectionRun(TimestampMixin, Base):
91124
),
92125
Index("ix_model_selection_run_status_created", "status", "created_at"),
93126
)
127+
128+
129+
class ModelSelectionCandidate(TimestampMixin, Base):
130+
"""One candidate's async execution record inside a selection run (Slice B).
131+
132+
Concurrent candidate tasks each write their OWN row in their OWN session —
133+
no shared-row write race. ``result`` carries the full ``CandidateResult``
134+
JSONB (incl. folds) on success; failed/cancelled candidates keep their row
135+
so they stay visible in the results UI. Mirrors ``batch.BatchJobItem``.
136+
"""
137+
138+
__tablename__ = "model_selection_candidate"
139+
140+
id: Mapped[int] = mapped_column(Integer, primary_key=True)
141+
candidate_id: Mapped[str] = mapped_column(String(32), unique=True, index=True)
142+
selection_id: Mapped[str] = mapped_column(
143+
String(32),
144+
ForeignKey("model_selection_run.selection_id", ondelete="CASCADE"),
145+
index=True,
146+
)
147+
ordinal: Mapped[int] = mapped_column(Integer) # submit order — stable display
148+
model_type: Mapped[str] = mapped_column(String(40))
149+
params: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
150+
status: Mapped[str] = mapped_column(
151+
String(20), default=CandidateStatus.PENDING.value, index=True
152+
)
153+
result: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True)
154+
error_message: Mapped[str | None] = mapped_column(String(2000), nullable=True)
155+
error_type: Mapped[str | None] = mapped_column(String(100), nullable=True)
156+
started_at: Mapped[_dt.datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
157+
completed_at: Mapped[_dt.datetime | None] = mapped_column(
158+
DateTime(timezone=True), nullable=True
159+
)
160+
duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
161+
162+
__table_args__ = (
163+
CheckConstraint(
164+
"status IN ('pending', 'running', 'completed', 'failed', 'cancelled')",
165+
name="ck_model_selection_candidate_valid_status",
166+
),
167+
Index(
168+
"ix_model_selection_candidate_selection_status",
169+
"selection_id",
170+
"status",
171+
),
172+
)

0 commit comments

Comments
 (0)