synth-ai 0.2.14__py3-none-any.whl → 0.4.4__py3-none-any.whl

synth_ai/cli/local/experiment_queue/database.py

@@ -0,0 +1,175 @@
+"""SQLAlchemy database utilities for the experiment queue."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+from .config import load_config
+
+
+class Base(DeclarativeBase):
+    """Declarative base for experiment queue ORM models."""
+
+
+_engine: Engine | None = None
+_session_factory: sessionmaker[Session] | None = None
+
+
+def get_engine() -> Engine:
+    """Get (or lazily create) the SQLAlchemy engine."""
+    global _engine
+    if _engine is None:
+        config = load_config()
+        # Enable WAL mode for SQLite to allow concurrent reads/writes
+        # This is critical for Celery broker + our ORM to work together
+        connect_args = {
+            "check_same_thread": False,
+            "timeout": 5.0,  # 5 second timeout - fast enough but allows WAL mode to work
+        }
+        # Enable WAL mode via PRAGMA after connection
+        def _enable_wal(dbapi_conn, connection_record):
+            # dbapi_conn is the raw sqlite3 connection, not SQLAlchemy
+            dbapi_conn.execute("PRAGMA journal_mode=WAL")
+            dbapi_conn.execute("PRAGMA synchronous=NORMAL")
+        
+        _engine = create_engine(
+            config.sqlalchemy_url,
+            echo=False,
+            future=True,
+            connect_args=connect_args,
+            pool_pre_ping=True,
+            poolclass=None,  # Use default pool
+        )
+        # Register event listener to enable WAL mode on each connection
+        from sqlalchemy import event
+        event.listen(_engine, "connect", _enable_wal)
+    return _engine
+
+
+def get_session() -> Session:
+    """Return a new Session bound to the shared engine."""
+    global _session_factory
+    if _session_factory is None:
+        _session_factory = sessionmaker(
+            bind=get_engine(),
+            class_=Session,
+            expire_on_commit=False,
+            autoflush=False,
+            future=True,
+        )
+    return _session_factory()
+
+
+def init_db() -> None:
+    """Create tables if they do not exist and ensure WAL mode is enabled.
+    
+    This initializes the application database (SQLite) for experiment queue data.
+    Celery broker uses Redis, so no SQLite locking conflicts.
+    """
+    engine = get_engine()
+    
+    # Enable WAL mode on the database file before creating tables
+    # This must happen before Celery's broker connection tries to use it
+    # Use multiple connections to ensure WAL mode is persistent
+    with engine.connect() as conn:
+        # Enable WAL mode (allows concurrent reads/writes)
+        result = conn.execute(text("PRAGMA journal_mode=WAL"))
+        wal_result = result.scalar()
+        if wal_result != "wal":
+            # Database might be locked or in use - try again
+            import time
+            time.sleep(0.1)
+            result = conn.execute(text("PRAGMA journal_mode=WAL"))
+            wal_result = result.scalar()
+        conn.execute(text("PRAGMA synchronous=NORMAL"))
+        conn.execute(text("PRAGMA busy_timeout=5000"))  # 5 seconds - allows concurrent access
+        conn.commit()
+        
+        # Verify WAL mode is enabled
+        if wal_result != "wal":
+            import warnings
+            warnings.warn(
+                f"WAL mode not enabled! Got: {wal_result}. "
+                f"This may cause SQLite locking errors with Celery.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+    
+    # Create our application tables
+    Base.metadata.create_all(engine)
+    
+    # Migrate schema: add status_json column if it doesn't exist
+    with engine.begin() as conn:  # Use begin() to ensure transaction is committed
+        # Check if table exists first
+        table_exists = conn.execute(text(
+            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='experiment_jobs'"
+        )).scalar() > 0
+        
+        if table_exists:
+            # Check if status_json column exists
+            result = conn.execute(text(
+                "SELECT COUNT(*) FROM pragma_table_info('experiment_jobs') WHERE name='status_json'"
+            ))
+            column_exists = result.scalar() > 0
+            
+            if not column_exists:
+                # Add status_json column
+                conn.execute(text(
+                    "ALTER TABLE experiment_jobs ADD COLUMN status_json TEXT"
+                ))
+    
+    # Migrate schema: create job_execution_logs table if it doesn't exist
+    with engine.connect() as conn:
+        # Check if job_execution_logs table exists
+        result = conn.execute(text(
+            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='job_execution_logs'"
+        ))
+        table_exists = result.scalar() > 0
+        
+        if not table_exists:
+            # Create job_execution_logs table
+            conn.execute(text("""
+                CREATE TABLE job_execution_logs (
+                    log_id VARCHAR(64) PRIMARY KEY,
+                    job_id VARCHAR(64) NOT NULL,
+                    command TEXT NOT NULL,
+                    working_directory TEXT NOT NULL,
+                    returncode INTEGER NOT NULL,
+                    stdout TEXT NOT NULL DEFAULT '',
+                    stderr TEXT NOT NULL DEFAULT '',
+                    python_executable VARCHAR(255),
+                    environment_keys TEXT,
+                    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (job_id) REFERENCES experiment_jobs(job_id) ON DELETE CASCADE
+                )
+            """))
+            # Create indexes
+            conn.execute(text("CREATE INDEX idx_job_execution_logs_job ON job_execution_logs(job_id)"))
+            conn.execute(text("CREATE INDEX idx_job_execution_logs_returncode ON job_execution_logs(returncode)"))
+            conn.execute(text("CREATE INDEX idx_job_execution_logs_created ON job_execution_logs(created_at)"))
+            conn.commit()
+    
+    # Force WAL mode one more time after table creation
+    # This ensures it's set even if table creation changed something
+    with engine.connect() as conn:
+        conn.execute(text("PRAGMA journal_mode=WAL"))
+        conn.commit()
+
+
+@contextmanager
+def session_scope() -> Iterator[Session]:
+    """Provide a transactional scope for DB operations."""
+    session = get_session()
+    try:
+        yield session
+        session.commit()
+    except Exception:
+        session.rollback()
+        raise
+    finally:
+        session.close()

synth_ai/cli/local/experiment_queue/dispatcher.py

@@ -0,0 +1,119 @@
+"""Utility helpers to enqueue Celery tasks while respecting parallelism limits."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import List
+
+from sqlalchemy import Select, func, select
+from sqlalchemy.orm import Session
+
+# Clear config cache if env vars are set (must happen before other imports)
+if os.getenv("EXPERIMENT_QUEUE_DB_PATH") or os.getenv("EXPERIMENT_QUEUE_TRAIN_CMD"):
+    from . import config as queue_config
+
+    queue_config.reset_config_cache()
+
+from .celery_app import get_celery_app
+from .models import Experiment, ExperimentJob, ExperimentJobStatus
+
+logger = logging.getLogger(__name__)
+
+TASK_NAME = "synth_ai.cli.local.experiment_queue.run_experiment_job"
+
+
+def _active_jobs_query(experiment_id: str) -> Select[tuple[int]]:
+    """Build a SQLAlchemy query to count active jobs for an experiment.
+    
+    Counts jobs that are:
+    - QUEUED or RUNNING status
+    - Have a Celery task ID assigned (dispatched)
+    - Not yet completed (completed_at is None)
+    
+    Args:
+        experiment_id: Experiment identifier
+        
+    Returns:
+        SQLAlchemy Select query that returns a count
+    """
+    return (
+        select(func.count(ExperimentJob.job_id))
+        .where(
+            ExperimentJob.experiment_id == experiment_id,
+            ExperimentJob.status.in_(
+                [ExperimentJobStatus.RUNNING, ExperimentJobStatus.QUEUED]
+            ),
+            ExperimentJob.celery_task_id.is_not(None),
+            ExperimentJob.completed_at.is_(None),
+        )
+    )
+
+
+def dispatch_available_jobs(session: Session, experiment_id: str) -> list[str]:
+    """Send ready jobs to Celery while honoring experiment parallelism."""
+    experiment = session.get(Experiment, experiment_id)
+    if not experiment:
+        return []
+
+    limit = max(experiment.parallelism_limit or 1, 1)
+    active = session.execute(_active_jobs_query(experiment_id)).scalar_one()
+    remaining_slots = max(limit - active, 0)
+    if remaining_slots <= 0:
+        return []
+
+    ready_jobs: List[ExperimentJob] = (
+        session.query(ExperimentJob)
+        .filter(
+            ExperimentJob.experiment_id == experiment_id,
+            ExperimentJob.status == ExperimentJobStatus.QUEUED,
+            ExperimentJob.celery_task_id.is_(None),
+        )
+        .order_by(ExperimentJob.created_at.asc())
+        .limit(remaining_slots)
+        .all()
+    )
+    if not ready_jobs:
+        return []
+
+    app = get_celery_app()
+    # Get broker URL for logging (may not be available in test stubs)
+    try:
+        broker_url = app.conf.broker_url
+    except AttributeError:
+        broker_url = "unknown"
+    logger.info(
+        "Dispatching %d jobs using Celery broker: %s",
+        len(ready_jobs),
+        broker_url,
+    )
+    
+    # Extract job IDs before committing session
+    job_ids = [job.job_id for job in ready_jobs]
+    
+    # Commit session before calling send_task
+    # Redis broker handles concurrency natively, no locking issues
+    session.flush()
+    session.commit()
+    
+    dispatched: list[str] = []
+    for job_id in job_ids:
+        # send_task sends message to Redis broker
+        result = app.send_task(TASK_NAME, args=[job_id])
+        
+        # Update celery_task_id in a new transaction
+        job = session.query(ExperimentJob).filter(ExperimentJob.job_id == job_id).first()
+        if job:
+            job.celery_task_id = result.id
+            logger.debug(
+                "Dispatched job %s with task_id %s",
+                job_id,
+                result.id,
+            )
+            dispatched.append(job_id)
+    
+    # Flush to persist celery_task_id updates
+    session.flush()
+    # Note: Task message is sent to Redis broker
+    # The final commit happens when session_scope exits in create_experiment()
+    return dispatched

synth_ai/cli/local/experiment_queue/models.py

@@ -0,0 +1,231 @@
+"""SQLAlchemy ORM models for the experiment queue."""
+
+from __future__ import annotations
+
+import enum
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import (
+    JSON,
+    DateTime,
+    Float,
+    ForeignKey,
+    Index,
+    Integer,
+    String,
+    Text,
+    UniqueConstraint,
+    func,
+)
+from sqlalchemy import (
+    Enum as SAEnum,
+)
+from sqlalchemy.ext.mutable import MutableDict
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from .database import Base
+
+
+class ExperimentStatus(str, enum.Enum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+class ExperimentJobType(str, enum.Enum):
+    GEPA = "gepa"
+    MIPRO = "mipro"
+
+
+class ExperimentJobStatus(str, enum.Enum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+class TrialStatus(str, enum.Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class Experiment(Base):
+    __tablename__ = "experiments"
+
+    experiment_id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    name: Mapped[str] = mapped_column(String(255), nullable=False)
+    description: Mapped[str | None] = mapped_column(Text())
+    status: Mapped[ExperimentStatus] = mapped_column(
+        SAEnum(ExperimentStatus, name="experiment_status"),
+        default=ExperimentStatus.QUEUED,
+        nullable=False,
+    )
+    parallelism_limit: Mapped[int] = mapped_column(Integer, default=3, nullable=False)
+    config_toml: Mapped[str] = mapped_column(Text(), nullable=False)
+    metadata_json: Mapped[dict[str, Any]] = mapped_column(
+        "metadata",
+        MutableDict.as_mutable(JSON),
+        default=dict,
+        nullable=False,
+    )
+    error: Mapped[str | None] = mapped_column(Text())
+
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False
+    )
+    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+
+    jobs: Mapped[list[ExperimentJob]] = relationship(
+        back_populates="experiment", cascade="all, delete-orphan"
+    )
+    trials: Mapped[list[Trial]] = relationship(
+        back_populates="experiment", cascade="all, delete-orphan"
+    )
+
+    __table_args__ = (
+        Index("idx_experiments_status", "status"),
+        Index("idx_experiments_created", "created_at"),
+    )
+
+class ExperimentJob(Base):
+    __tablename__ = "experiment_jobs"
+
+    job_id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    experiment_id: Mapped[str] = mapped_column(
+        String(64), ForeignKey("experiments.experiment_id", ondelete="CASCADE"), nullable=False
+    )
+    job_type: Mapped[ExperimentJobType] = mapped_column(
+        SAEnum(ExperimentJobType, name="experiment_job_type"),
+        nullable=False,
+    )
+    config_path: Mapped[str] = mapped_column(Text(), nullable=False)
+    config_overrides: Mapped[dict[str, Any]] = mapped_column(
+        MutableDict.as_mutable(JSON), default=dict, nullable=False
+    )
+    status: Mapped[ExperimentJobStatus] = mapped_column(
+        SAEnum(ExperimentJobStatus, name="experiment_job_status"),
+        default=ExperimentJobStatus.QUEUED,
+        nullable=False,
+    )
+    celery_task_id: Mapped[str | None] = mapped_column(String(128))
+    backend_job_id: Mapped[str | None] = mapped_column(String(128))
+    result: Mapped[dict[str, Any] | None] = mapped_column(MutableDict.as_mutable(JSON))
+    status_json: Mapped[dict[str, Any] | None] = mapped_column(MutableDict.as_mutable(JSON))
+    error: Mapped[str | None] = mapped_column(Text())
+
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False
+    )
+    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+
+    experiment: Mapped[Experiment] = relationship(back_populates="jobs")
+    trials: Mapped[list[Trial]] = relationship(back_populates="job")
+    execution_logs: Mapped[list[JobExecutionLog]] = relationship(
+        back_populates="job", cascade="all, delete-orphan"
+    )
+
+    __table_args__ = (
+        Index("idx_experiment_jobs_status", "status"),
+        Index("idx_experiment_jobs_experiment", "experiment_id"),
+        Index("idx_experiment_jobs_celery_task", "celery_task_id"),
+    )
+
+
+class Trial(Base):
+    __tablename__ = "trials"
+
+    trial_id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    experiment_id: Mapped[str] = mapped_column(
+        String(64), ForeignKey("experiments.experiment_id", ondelete="CASCADE"), nullable=False
+    )
+    job_id: Mapped[str | None] = mapped_column(
+        String(64), ForeignKey("experiment_jobs.job_id", ondelete="SET NULL")
+    )
+    trial_number: Mapped[int | None] = mapped_column(Integer)
+    system_name: Mapped[str | None] = mapped_column(String(128))
+    prompt_id: Mapped[str | None] = mapped_column(String(128))
+    session_id: Mapped[str | None] = mapped_column(String(128))
+    trace_stored_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+
+    reward: Mapped[float | None] = mapped_column(Float)
+    aggregate_score: Mapped[float | None] = mapped_column(Float)
+    rewards: Mapped[dict[str, Any] | None] = mapped_column(MutableDict.as_mutable(JSON))
+    metadata_json: Mapped[dict[str, Any]] = mapped_column(
+        "metadata",
+        MutableDict.as_mutable(JSON),
+        default=dict,
+        nullable=False,
+    )
+    status: Mapped[TrialStatus] = mapped_column(
+        SAEnum(TrialStatus, name="experiment_trial_status"),
+        default=TrialStatus.PENDING,
+        nullable=False,
+    )
+
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False
+    )
+    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+
+    experiment: Mapped[Experiment] = relationship(back_populates="trials")
+    job: Mapped[ExperimentJob | None] = relationship(back_populates="trials")
+
+    __table_args__ = (
+        Index("idx_trials_experiment", "experiment_id"),
+        Index("idx_trials_job", "job_id"),
+        Index("idx_trials_status", "status"),
+        Index("idx_trials_reward", "reward"),
+        Index("idx_trials_aggregate_score", "aggregate_score"),
+        Index("idx_trials_experiment_trial_number", "experiment_id", "trial_number"),
+        UniqueConstraint("experiment_id", "trial_number", name="uq_trial_number_per_experiment"),
+    )
+
+
+class JobExecutionLog(Base):
+    """Detailed execution logs for job subprocess runs.
+    
+    Stores full stdout, stderr, command, and environment info for ALL jobs (successful and failed).
+    This allows querying failures directly from the database without guessing.
+    """
+    __tablename__ = "job_execution_logs"
+
+    log_id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    job_id: Mapped[str] = mapped_column(
+        String(64), ForeignKey("experiment_jobs.job_id", ondelete="CASCADE"), nullable=False
+    )
+    
+    # Execution details
+    command: Mapped[str] = mapped_column(Text(), nullable=False)  # Full command executed
+    working_directory: Mapped[str] = mapped_column(Text(), nullable=False)
+    returncode: Mapped[int] = mapped_column(Integer, nullable=False)
+    
+    # Output (stored as Text to handle large outputs)
+    stdout: Mapped[str] = mapped_column(Text(), nullable=False, default="")
+    stderr: Mapped[str] = mapped_column(Text(), nullable=False, default="")
+    
+    # Environment info (for debugging)
+    python_executable: Mapped[str | None] = mapped_column(String(255))
+    environment_keys: Mapped[list[str] | None] = mapped_column(JSON)  # List of env var keys present
+    
+    # Timestamps
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False
+    )
+    
+    # Relationship
+    job: Mapped[ExperimentJob] = relationship(back_populates="execution_logs")
+
+    __table_args__ = (
+        Index("idx_job_execution_logs_job", "job_id"),
+        Index("idx_job_execution_logs_returncode", "returncode"),
+        Index("idx_job_execution_logs_created", "created_at"),
+    )

synth_ai/cli/local/experiment_queue/progress_info.py

@@ -0,0 +1,160 @@
+"""Type-safe progress information dataclass."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True, slots=True)
+class ProgressInfo:
+    """Type-safe representation of progress information parsed from output.
+    
+    This replaces dict[str, Any] usage in status_tracker.py for better type safety.
+    """
+
+    rollouts_completed: int | None = None
+    total_rollouts: int | None = None
+    best_score: float | None = None
+    trials_completed: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate progress data after initialization."""
+        # Validate rollouts_completed
+        if self.rollouts_completed is not None:
+            assert isinstance(self.rollouts_completed, int), (
+                f"rollouts_completed must be int, got {type(self.rollouts_completed).__name__}"
+            )
+            assert self.rollouts_completed >= 0, (
+                f"rollouts_completed must be >= 0, got {self.rollouts_completed}"
+            )
+            if self.total_rollouts is not None:
+                assert self.rollouts_completed <= self.total_rollouts, (
+                    f"rollouts_completed ({self.rollouts_completed}) > total_rollouts ({self.total_rollouts})"
+                )
+
+        # Validate total_rollouts
+        if self.total_rollouts is not None:
+            assert isinstance(self.total_rollouts, int), (
+                f"total_rollouts must be int, got {type(self.total_rollouts).__name__}"
+            )
+            assert self.total_rollouts > 0, (
+                f"total_rollouts must be > 0, got {self.total_rollouts}"
+            )
+
+        # Validate best_score
+        if self.best_score is not None:
+            assert isinstance(self.best_score, int | float), (
+                f"best_score must be int | float, got {type(self.best_score).__name__}"
+            )
+            assert 0 <= self.best_score <= 1, (
+                f"best_score must be in [0, 1], got {self.best_score}"
+            )
+
+        # Validate trials_completed
+        if self.trials_completed is not None:
+            assert isinstance(self.trials_completed, int), (
+                f"trials_completed must be int, got {type(self.trials_completed).__name__}"
+            )
+            assert self.trials_completed >= 0, (
+                f"trials_completed must be >= 0, got {self.trials_completed}"
+            )
+
+    def to_dict(self) -> dict[str, int | float | None]:
+        """Convert to dictionary for compatibility with existing code."""
+        return {
+            "rollouts_completed": self.rollouts_completed,
+            "total_rollouts": self.total_rollouts,
+            "best_score": self.best_score,
+            "trials_completed": self.trials_completed,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ProgressInfo:
+        """Create from dictionary with type coercion and validation.
+        
+        Args:
+            data: Dictionary with progress information (may have wrong types)
+            
+        Returns:
+            ProgressInfo instance with validated and coerced types
+            
+        Raises:
+            AssertionError: If data cannot be coerced or is invalid
+        """
+        # Coerce and validate each field
+        rollouts_completed = None
+        if "rollouts_completed" in data and data["rollouts_completed"] is not None:
+            val = data["rollouts_completed"]
+            if isinstance(val, int | float | str):
+                try:
+                    coerced = int(float(val))
+                    assert coerced >= 0, f"rollouts_completed must be >= 0, got {coerced}"
+                    rollouts_completed = coerced
+                except (ValueError, TypeError, AssertionError) as e:
+                    raise AssertionError(
+                        f"Cannot coerce rollouts_completed to int: {val} (type: {type(val).__name__})"
+                    ) from e
+            else:
+                raise AssertionError(
+                    f"rollouts_completed must be int | float | str, got {type(val).__name__}: {val}"
+                )
+
+        total_rollouts = None
+        if "total_rollouts" in data and data["total_rollouts"] is not None:
+            val = data["total_rollouts"]
+            if isinstance(val, int | float | str):
+                try:
+                    total_rollouts = int(float(val))
+                except (ValueError, TypeError) as e:
+                    raise AssertionError(
+                        f"Cannot coerce total_rollouts to int: {val} (type: {type(val).__name__})"
+                    ) from e
+            else:
+                raise AssertionError(
+                    f"total_rollouts must be int | float | str, got {type(val).__name__}: {val}"
+                )
+            if total_rollouts is not None:
+                assert total_rollouts > 0, f"total_rollouts must be > 0, got {total_rollouts}"
+
+        best_score = None
+        if "best_score" in data and data["best_score"] is not None:
+            val = data["best_score"]
+            if isinstance(val, int | float | str):
+                try:
+                    coerced = float(val)
+                    assert 0 <= coerced <= 1, f"best_score must be in [0, 1], got {coerced}"
+                    best_score = coerced
+                except (ValueError, TypeError, AssertionError) as e:
+                    raise AssertionError(
+                        f"Cannot coerce best_score to float: {val} (type: {type(val).__name__})"
+                    ) from e
+            else:
+                raise AssertionError(
+                    f"best_score must be int | float | str, got {type(val).__name__}: {val}"
+                )
+
+        trials_completed = None
+        if "trials_completed" in data and data["trials_completed"] is not None:
+            val = data["trials_completed"]
+            if isinstance(val, int | float | str):
+                try:
+                    coerced = int(float(val))
+                    assert coerced >= 0, f"trials_completed must be >= 0, got {coerced}"
+                    trials_completed = coerced
+                except (ValueError, TypeError, AssertionError) as e:
+                    raise AssertionError(
+                        f"Cannot coerce trials_completed to int: {val} (type: {type(val).__name__})"
+                    ) from e
+            else:
+                raise AssertionError(
+                    f"trials_completed must be int | float | str, got {type(val).__name__}: {val}"
+                )
+
+        return cls(
+            rollouts_completed=rollouts_completed,
+            total_rollouts=total_rollouts,
+            best_score=best_score,
+            trials_completed=trials_completed,
+        )
+