253 lines
8.2 KiB
Python
253 lines
8.2 KiB
Python
"""Persistent runtime observations used by cost-aware environment placement."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
import sqlite3
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Mapping, Sequence
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
EXECUTION_HISTORY_FILENAME = "execution-history.sqlite3"
|
|
_RECENT_SAMPLE_LIMIT = 20
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExecutionObservation:
|
|
"""Describe one completed remote component invocation."""
|
|
|
|
component_signature: str
|
|
environment_id: str
|
|
provider: str
|
|
elapsed_seconds: float
|
|
recorded_at_epoch: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExecutionEstimate:
|
|
"""Summarize recent successful runtimes for one component/environment pair."""
|
|
|
|
execution_seconds: float
|
|
sample_count: int
|
|
|
|
|
|
@dataclass
|
|
class ExecutionHistory:
|
|
"""Store bounded component timing samples in a local SQLite database."""
|
|
|
|
database_path: Path
|
|
|
|
@classmethod
|
|
def for_user_directory(cls, user_directory: Path) -> "ExecutionHistory":
|
|
"""Create history storage beneath the node pack's ComfyUI user directory."""
|
|
return cls(
|
|
database_path=(
|
|
user_directory.expanduser().resolve()
|
|
/ "comfyui-modal"
|
|
/ EXECUTION_HISTORY_FILENAME
|
|
)
|
|
)
|
|
|
|
def record(self, observation: ExecutionObservation) -> None:
|
|
"""Persist one successful, finite runtime observation."""
|
|
_validate_observation(observation)
|
|
with self._connect() as connection:
|
|
_insert_observation(connection, observation)
|
|
_prune_observations(connection, observation)
|
|
|
|
def estimates(
|
|
self,
|
|
component_signature: str,
|
|
environment_ids: Sequence[str],
|
|
) -> Mapping[str, ExecutionEstimate]:
|
|
"""Return robust recent runtime estimates for the requested environments."""
|
|
normalized_ids = tuple(
|
|
dict.fromkeys(
|
|
environment_id.strip()
|
|
for environment_id in environment_ids
|
|
if environment_id.strip()
|
|
)
|
|
)
|
|
if not normalized_ids or not self.database_path.exists():
|
|
return {}
|
|
rows = self._load_estimate_rows(component_signature, normalized_ids)
|
|
samples_by_environment: dict[str, list[float]] = {}
|
|
for environment_id, elapsed_seconds in rows:
|
|
samples = samples_by_environment.setdefault(str(environment_id), [])
|
|
if len(samples) < _RECENT_SAMPLE_LIMIT:
|
|
samples.append(float(elapsed_seconds))
|
|
return {
|
|
environment_id: ExecutionEstimate(
|
|
execution_seconds=_median(samples),
|
|
sample_count=len(samples),
|
|
)
|
|
for environment_id, samples in samples_by_environment.items()
|
|
if samples
|
|
}
|
|
|
|
def _load_estimate_rows(
|
|
self,
|
|
component_signature: str,
|
|
environment_ids: tuple[str, ...],
|
|
) -> list[tuple[str, float]]:
|
|
"""Load recent matching observation rows, returning none on index failure."""
|
|
placeholders = ",".join("?" for _ in environment_ids)
|
|
try:
|
|
with self._connect() as connection:
|
|
return connection.execute(
|
|
f"""
|
|
SELECT environment_id, elapsed_seconds
|
|
FROM execution_observations
|
|
WHERE component_signature = ?
|
|
AND environment_id IN ({placeholders})
|
|
ORDER BY recorded_at_epoch DESC, id DESC
|
|
""",
|
|
(component_signature, *environment_ids),
|
|
).fetchall()
|
|
except (OSError, sqlite3.Error) as exc:
|
|
logger.warning("Unable to read remote execution history: %s", exc)
|
|
return []
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
"""Open an initialized short-lived SQLite connection."""
|
|
self.database_path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(self.database_path, timeout=5.0)
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA synchronous=NORMAL")
|
|
_initialize_schema(connection)
|
|
return connection
|
|
|
|
|
|
def _validate_observation(observation: ExecutionObservation) -> None:
|
|
"""Reject invalid runtime observations before touching persistent state."""
|
|
if not observation.component_signature.strip():
|
|
raise ValueError("component_signature must not be empty.")
|
|
if not observation.environment_id.strip():
|
|
raise ValueError("environment_id must not be empty.")
|
|
if not observation.provider.strip():
|
|
raise ValueError("provider must not be empty.")
|
|
if not math.isfinite(observation.elapsed_seconds) or observation.elapsed_seconds < 0:
|
|
raise ValueError("elapsed_seconds must be finite and non-negative.")
|
|
if (
|
|
not math.isfinite(observation.recorded_at_epoch)
|
|
or observation.recorded_at_epoch < 0
|
|
):
|
|
raise ValueError("recorded_at_epoch must be finite and non-negative.")
|
|
|
|
|
|
def _insert_observation(
|
|
connection: sqlite3.Connection,
|
|
observation: ExecutionObservation,
|
|
) -> None:
|
|
"""Insert one observation into an initialized history database."""
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO execution_observations (
|
|
component_signature,
|
|
environment_id,
|
|
provider,
|
|
elapsed_seconds,
|
|
recorded_at_epoch
|
|
) VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
observation.component_signature,
|
|
observation.environment_id,
|
|
observation.provider,
|
|
observation.elapsed_seconds,
|
|
observation.recorded_at_epoch,
|
|
),
|
|
)
|
|
|
|
|
|
def _prune_observations(
|
|
connection: sqlite3.Connection,
|
|
observation: ExecutionObservation,
|
|
) -> None:
|
|
"""Retain only the bounded recent window for one signature and environment."""
|
|
connection.execute(
|
|
"""
|
|
DELETE FROM execution_observations
|
|
WHERE id IN (
|
|
SELECT id
|
|
FROM execution_observations
|
|
WHERE component_signature = ? AND environment_id = ?
|
|
ORDER BY recorded_at_epoch DESC, id DESC
|
|
LIMIT -1 OFFSET ?
|
|
)
|
|
""",
|
|
(
|
|
observation.component_signature,
|
|
observation.environment_id,
|
|
_RECENT_SAMPLE_LIMIT,
|
|
),
|
|
)
|
|
|
|
|
|
def _initialize_schema(connection: sqlite3.Connection) -> None:
|
|
"""Create the timing table and lookup index when missing."""
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS execution_observations (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
component_signature TEXT NOT NULL,
|
|
environment_id TEXT NOT NULL,
|
|
provider TEXT NOT NULL,
|
|
elapsed_seconds REAL NOT NULL,
|
|
recorded_at_epoch REAL NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE INDEX IF NOT EXISTS execution_observations_lookup
|
|
ON execution_observations (
|
|
component_signature,
|
|
environment_id,
|
|
recorded_at_epoch DESC
|
|
)
|
|
"""
|
|
)
|
|
|
|
|
|
def record_completed_execution(
|
|
*,
|
|
history: ExecutionHistory | None,
|
|
component_signature: str | None,
|
|
environment_id: str,
|
|
provider: str,
|
|
elapsed_seconds: float,
|
|
) -> None:
|
|
"""Best-effort record one completed invocation without failing its result."""
|
|
if history is None or not component_signature:
|
|
return
|
|
try:
|
|
history.record(
|
|
ExecutionObservation(
|
|
component_signature=component_signature,
|
|
environment_id=environment_id,
|
|
provider=provider,
|
|
elapsed_seconds=elapsed_seconds,
|
|
recorded_at_epoch=time.time(),
|
|
)
|
|
)
|
|
except (OSError, sqlite3.Error, ValueError) as exc:
|
|
logger.warning(
|
|
"Unable to persist remote execution timing environment=%s: %s",
|
|
environment_id,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _median(values: Sequence[float]) -> float:
|
|
"""Return the median of one non-empty numeric sequence."""
|
|
ordered = sorted(values)
|
|
midpoint = len(ordered) // 2
|
|
if len(ordered) % 2:
|
|
return ordered[midpoint]
|
|
return (ordered[midpoint - 1] + ordered[midpoint]) / 2
|