busya
4205618ee6
Warnings fixed: - datetime.utcnow() → datetime.now(UTC) across 48+ files (src/ + tests/) - datetime.utcnow (callback ref) → lambda: datetime.now(UTC) in model fields (18 files) - Pydantic class Config → model_config = ConfigDict(...) (16 files) - Pydantic .dict() → .model_dump() (8 files) - ConfigDict(allow_population_by_field_name=True) → validate_by_name=True - SQLAlchemy declarative_base() import path updated - FastAPI on_event → lifespan context manager (app.py) - Import sorting (ruff I001) auto-fixed across all files - Fixed broken re-export chains that ruff F401 cleanup broke: _validate_bcp47: service.py now imports from dictionary_validation directly job_to_response: _job_routes.py and test imports from service_utils directly fetch_datasource_metadata: restored re-export in service.py - Added missing TranslateJobService import in _job_routes.py (was deleted by F401) - Added ConfigDict(protected_namespaces=()) for DashboardDatasetItem schema field - pytest.ini: replaced deprecated importmode with asyncio_mode All 440 tests pass with zero deprecation warnings.
353 lines
17 KiB
Python
353 lines
17 KiB
Python
# #region JobLifecycleModule [C:5] [TYPE Module] [SEMANTICS task,lifecycle,state,machine,execution]
|
|
# @BRIEF Task creation, execution, pause/resume, and completion transitions for plugin-backed
|
|
# jobs — extracted from the monolithic TaskManager.
|
|
# @LAYER Core
|
|
# @RELATION DEPENDS_ON -> [TaskGraph]
|
|
# @RELATION DEPENDS_ON -> [EventBus]
|
|
# @RELATION DEPENDS_ON -> [TaskContext]
|
|
# @RELATION DEPENDS_ON -> [PluginLoader]
|
|
# @RELATION DEPENDS_ON -> [TaskPersistenceService]
|
|
# @RATIONALE Extracted from TaskManager to satisfy INV_7. JobLifecycle encodes the task state
|
|
# machine with plugin execution, context injection, backward-compatible plugin
|
|
# dispatch, and input/mapping pause/resume transitions.
|
|
# @REJECTED Keeping lifecycle/execution logic inside TaskManager was rejected — it mixes
|
|
# async execution orchestration with passive registry and log buffering concerns,
|
|
# violating single-responsibility and INV_7.
|
|
# @PRE Requested plugin ids resolve to executable plugins; task state transitions target known
|
|
# TaskStatus values.
|
|
# @POST Every scheduled task transitions through a valid lifecycle path ending in persisted
|
|
# terminal or waiting state.
|
|
# @SIDE_EFFECT Schedules async execution, pauses on input/mapping requests, mutates persisted
|
|
# task lifecycle markers.
|
|
# @INVARIANT A task cannot be resumed from a waiting state unless a matching future exists or
|
|
# a new wait future is created.
|
|
|
|
import asyncio
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import UTC, datetime
|
|
import inspect
|
|
from typing import Any
|
|
|
|
from ..cot_logger import seed_trace_id
|
|
from ..logger import belief_scope, logger
|
|
from .context import TaskContext
|
|
from .graph import TaskGraph
|
|
from .models import Task, TaskStatus
|
|
|
|
|
|
# #region JobLifecycle [C:5] [TYPE Class] [SEMANTICS task,lifecycle,execution,state,machine]
|
|
# @BRIEF Encodes task creation, execution, pause/resume, and completion transitions for
|
|
# plugin-backed jobs.
|
|
# @RELATION DEPENDS_ON -> [TaskGraph]
|
|
# @RELATION DEPENDS_ON -> [TaskContext]
|
|
# @RELATION DEPENDS_ON -> [PluginLoader]
|
|
# @RELATION DEPENDS_ON -> [TaskPersistenceService]
|
|
# @PRE plugin_loader resolves plugin ids; graph and event_bus are initialized.
|
|
# @POST Every scheduled task transitions through a valid lifecycle path ending in persisted
|
|
# terminal or waiting state.
|
|
# @RATIONALE Extracted as standalone lifecycle state machine to isolate async execution
|
|
# orchestration (plugin dispatch, context injection, pause/resume futures) from
|
|
# passive registry and log buffering concerns.
|
|
# @REJECTED Keeping execution logic inside TaskManager was rejected — it mixed async task
|
|
# orchestration with synchronous registry operations and background threading,
|
|
# violating single-responsibility and making the 708-line module unmaintainable.
|
|
class JobLifecycle:
|
|
"""Task state machine: creation, execution, pause/resume, completion."""
|
|
|
|
# #region __init__ [TYPE Function]
|
|
# @BRIEF Initialize with dependencies.
|
|
def __init__(
|
|
self,
|
|
plugin_loader: Any,
|
|
graph: TaskGraph,
|
|
event_bus: Any,
|
|
persistence_service: Any,
|
|
executor: ThreadPoolExecutor,
|
|
loop: asyncio.AbstractEventLoop,
|
|
):
|
|
self.plugin_loader = plugin_loader
|
|
self.graph = graph
|
|
self.event_bus = event_bus
|
|
self.persistence_service = persistence_service
|
|
self.executor = executor
|
|
self.loop = loop
|
|
self._dataset_subscribers: dict[str, list[asyncio.Queue]] = {}
|
|
# #endregion __init__
|
|
|
|
# #region create_task [C:4] [TYPE Function] [SEMANTICS task,create,persist,schedule]
|
|
# @BRIEF Creates and queues a new task for execution.
|
|
# @PRE Plugin with plugin_id exists. Params are valid dict.
|
|
# @POST Task is created, added to registry, and scheduled for execution.
|
|
# @RAISES ValueError if plugin not found or params invalid.
|
|
async def create_task(
|
|
self, plugin_id: str, params: dict[str, Any], user_id: str | None = None,
|
|
add_log_callback=None,
|
|
) -> Task:
|
|
with belief_scope("JobLifecycle.create_task", f"plugin_id={plugin_id}"):
|
|
if not self.plugin_loader.has_plugin(plugin_id):
|
|
logger.error(f"Plugin with ID '{plugin_id}' not found.")
|
|
raise ValueError(f"Plugin with ID '{plugin_id}' not found.")
|
|
self.plugin_loader.get_plugin(plugin_id)
|
|
if not isinstance(params, dict):
|
|
logger.error("Task parameters must be a dictionary.")
|
|
raise ValueError("Task parameters must be a dictionary.")
|
|
logger.reason("Creating task node and scheduling execution")
|
|
task = Task(plugin_id=plugin_id, params=params, user_id=user_id)
|
|
self.graph.add_task(task)
|
|
self.persistence_service.persist_task(task)
|
|
logger.info(f"Task {task.id} created and scheduled for execution")
|
|
self.loop.create_task(
|
|
self._run_task(task.id, add_log_callback=add_log_callback)
|
|
)
|
|
logger.reflect(
|
|
"Task creation persisted and execution scheduled",
|
|
extra={"task_id": task.id, "plugin_id": plugin_id},
|
|
)
|
|
return task
|
|
# #endregion create_task
|
|
|
|
# #region _run_task [C:4] [TYPE Function] [SEMANTICS task,execute,context,dispatch]
|
|
# @BRIEF Internal method to execute a task with TaskContext support.
|
|
# @PRE Task exists in registry.
|
|
# @POST Task is executed, status updated to SUCCESS or FAILED.
|
|
# @SIDE_EFFECT Executes plugin code, updates task status and persistence, emits logs.
|
|
async def _run_task(self, task_id: str, add_log_callback=None):
|
|
seed_trace_id()
|
|
with belief_scope("JobLifecycle._run_task", f"task_id={task_id}"):
|
|
task = self.graph.get_task(task_id)
|
|
if task is None:
|
|
logger.error(f"Task {task_id} not found in registry")
|
|
return
|
|
plugin = self.plugin_loader.get_plugin(task.plugin_id)
|
|
logger.reason(
|
|
"Transitioning task to running state",
|
|
extra={"task_id": task_id, "plugin_id": task.plugin_id},
|
|
)
|
|
logger.info(
|
|
f"Starting execution of task {task_id} for plugin '{plugin.name}'"
|
|
)
|
|
task.status = TaskStatus.RUNNING
|
|
task.started_at = datetime.now(UTC)
|
|
self.persistence_service.persist_task(task)
|
|
if add_log_callback:
|
|
add_log_callback(
|
|
task_id, "INFO",
|
|
f"Task started for plugin '{plugin.name}'",
|
|
source="system",
|
|
)
|
|
try:
|
|
params = {**task.params, "_task_id": task_id}
|
|
sig = inspect.signature(plugin.execute)
|
|
accepts_context = "context" in sig.parameters
|
|
if accepts_context:
|
|
# Create TaskContext for new-style plugins
|
|
context = TaskContext(
|
|
task_id=task_id,
|
|
add_log_fn=add_log_callback,
|
|
params=params,
|
|
default_source="plugin",
|
|
background_tasks=None,
|
|
)
|
|
if asyncio.iscoroutinefunction(plugin.execute):
|
|
task.result = await plugin.execute(params, context=context)
|
|
else:
|
|
task.result = await self.loop.run_in_executor(
|
|
self.executor,
|
|
lambda: plugin.execute(params, context=context),
|
|
)
|
|
else:
|
|
if asyncio.iscoroutinefunction(plugin.execute):
|
|
task.result = await plugin.execute(params)
|
|
else:
|
|
task.result = await self.loop.run_in_executor(
|
|
self.executor, plugin.execute, params
|
|
)
|
|
logger.info(f"Task {task_id} completed successfully")
|
|
task.status = TaskStatus.SUCCESS
|
|
if add_log_callback:
|
|
add_log_callback(
|
|
task_id, "INFO",
|
|
f"Task completed successfully for plugin '{plugin.name}'",
|
|
source="system",
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Task {task_id} failed: {e}")
|
|
task.status = TaskStatus.FAILED
|
|
if add_log_callback:
|
|
add_log_callback(
|
|
task_id, "ERROR",
|
|
f"Task failed: {e}",
|
|
source="system",
|
|
metadata={"error_type": type(e).__name__},
|
|
)
|
|
finally:
|
|
task.finished_at = datetime.now(UTC)
|
|
self.event_bus.flush_task_logs(task_id)
|
|
self.persistence_service.persist_task(task)
|
|
logger.info(
|
|
f"Task {task_id} execution finished with status: {task.status}"
|
|
)
|
|
logger.reflect(
|
|
"Task lifecycle reached persisted terminal state",
|
|
extra={"task_id": task_id, "status": str(task.status)},
|
|
)
|
|
# Broadcast dataset.updated for mapping and documentation tasks (FR-024, R4)
|
|
if task.plugin_id in ("dataset-mapper", "llm_documentation"):
|
|
dataset_ids = task.params.get("dataset_ids") or [task.params.get("dataset_id")]
|
|
env_id = task.params.get("env") or task.params.get("environment_id", "")
|
|
if dataset_ids and env_id:
|
|
self._broadcast_dataset_updated(env_id, dataset_ids)
|
|
# #endregion _run_task
|
|
|
|
# #region resolve_task [C:3] [TYPE Function] [SEMANTICS task,resolve,mapping,resume]
|
|
# @BRIEF Resumes a task that is awaiting mapping.
|
|
# @PRE Task exists and is in AWAITING_MAPPING state.
|
|
# @POST Task status updated to RUNNING, params updated, execution resumed.
|
|
# @RAISES ValueError if task not found or not awaiting mapping.
|
|
async def resolve_task(
|
|
self, task_id: str, resolution_params: dict[str, Any],
|
|
) -> None:
|
|
with belief_scope("JobLifecycle.resolve_task", f"task_id={task_id}"):
|
|
task = self.graph.get_task(task_id)
|
|
if not task or task.status != TaskStatus.AWAITING_MAPPING:
|
|
raise ValueError("Task is not awaiting mapping.")
|
|
task.params.update(resolution_params)
|
|
task.status = TaskStatus.RUNNING
|
|
self.persistence_service.persist_task(task)
|
|
self.graph.resolve_future(task_id, True)
|
|
# #endregion resolve_task
|
|
|
|
# #region wait_for_resolution [C:3] [TYPE Function] [SEMANTICS task,wait,mapping,future]
|
|
# @BRIEF Pauses execution and waits for a resolution signal.
|
|
# @PRE Task exists.
|
|
# @POST Execution pauses until future is set.
|
|
async def wait_for_resolution(self, task_id: str) -> None:
|
|
with belief_scope("JobLifecycle.wait_for_resolution", f"task_id={task_id}"):
|
|
task = self.graph.get_task(task_id)
|
|
if not task:
|
|
return
|
|
task.status = TaskStatus.AWAITING_MAPPING
|
|
self.persistence_service.persist_task(task)
|
|
future = self.loop.create_future()
|
|
self.graph.create_future(task_id, future)
|
|
try:
|
|
await future
|
|
finally:
|
|
self.graph.remove_future(task_id)
|
|
# #endregion wait_for_resolution
|
|
|
|
# #region wait_for_input [C:3] [TYPE Function] [SEMANTICS task,wait,input,future]
|
|
# @BRIEF Pauses execution and waits for user input.
|
|
# @PRE Task exists.
|
|
# @POST Execution pauses until future is set via resume_task_with_password.
|
|
async def wait_for_input(self, task_id: str) -> None:
|
|
with belief_scope("JobLifecycle.wait_for_input", f"task_id={task_id}"):
|
|
task = self.graph.get_task(task_id)
|
|
if not task:
|
|
return
|
|
future = self.loop.create_future()
|
|
self.graph.create_future(task_id, future)
|
|
try:
|
|
await future
|
|
finally:
|
|
self.graph.remove_future(task_id)
|
|
# #endregion wait_for_input
|
|
|
|
# #region await_input [C:3] [TYPE Function] [SEMANTICS task,input,pause,state]
|
|
# @BRIEF Transition a task to AWAITING_INPUT state with input request.
|
|
# @PRE Task exists and is in RUNNING state.
|
|
# @POST Task status changed to AWAITING_INPUT, input_request set, persisted.
|
|
# @RAISES ValueError if task not found or not RUNNING.
|
|
def await_input(
|
|
self, task_id: str, input_request: dict[str, Any],
|
|
add_log_callback=None,
|
|
) -> None:
|
|
with belief_scope("JobLifecycle.await_input", f"task_id={task_id}"):
|
|
task = self.graph.get_task(task_id)
|
|
if not task:
|
|
raise ValueError(f"Task {task_id} not found")
|
|
if task.status != TaskStatus.RUNNING:
|
|
raise ValueError(
|
|
f"Task {task_id} is not RUNNING (current: {task.status})"
|
|
)
|
|
task.status = TaskStatus.AWAITING_INPUT
|
|
task.input_required = True
|
|
task.input_request = input_request
|
|
self.persistence_service.persist_task(task)
|
|
if add_log_callback:
|
|
add_log_callback(
|
|
task_id, "INFO",
|
|
"Task paused for user input",
|
|
metadata={"input_request": input_request},
|
|
)
|
|
# #endregion await_input
|
|
|
|
# #region resume_task_with_password [C:3] [TYPE Function] [SEMANTICS task,resume,password,input]
|
|
# @BRIEF Resume a task that is awaiting input with provided passwords.
|
|
# @PRE Task exists and is in AWAITING_INPUT state.
|
|
# @POST Task status changed to RUNNING, passwords injected, task resumed.
|
|
# @RAISES ValueError if task not found, not awaiting input, or passwords invalid.
|
|
def resume_task_with_password(
|
|
self, task_id: str, passwords: dict[str, str],
|
|
add_log_callback=None,
|
|
) -> None:
|
|
with belief_scope(
|
|
"JobLifecycle.resume_task_with_password", f"task_id={task_id}"
|
|
):
|
|
task = self.graph.get_task(task_id)
|
|
if not task:
|
|
raise ValueError(f"Task {task_id} not found")
|
|
if task.status != TaskStatus.AWAITING_INPUT:
|
|
raise ValueError(
|
|
f"Task {task_id} is not AWAITING_INPUT (current: {task.status})"
|
|
)
|
|
if not isinstance(passwords, dict) or not passwords:
|
|
raise ValueError("Passwords must be a non-empty dictionary")
|
|
task.params["passwords"] = passwords
|
|
task.input_required = False
|
|
task.input_request = None
|
|
task.status = TaskStatus.RUNNING
|
|
self.persistence_service.persist_task(task)
|
|
if add_log_callback:
|
|
add_log_callback(
|
|
task_id, "INFO",
|
|
"Task resumed with passwords",
|
|
metadata={"databases": list(passwords.keys())},
|
|
)
|
|
self.graph.resolve_future(task_id, True)
|
|
# #endregion resume_task_with_password
|
|
|
|
# #region subscribe_dataset_events [C:2] [TYPE Function] [SEMANTICS dataset,subscribe,events,websocket]
|
|
# @BRIEF Subscribe to dataset.updated events for a specific environment.
|
|
async def subscribe_dataset_events(self, env_id: str) -> asyncio.Queue:
|
|
queue = asyncio.Queue()
|
|
if env_id not in self._dataset_subscribers:
|
|
self._dataset_subscribers[env_id] = []
|
|
self._dataset_subscribers[env_id].append(queue)
|
|
return queue
|
|
# #endregion subscribe_dataset_events
|
|
|
|
# #region unsubscribe_dataset_events [C:2] [TYPE Function] [SEMANTICS dataset,unsubscribe,events]
|
|
# @BRIEF Unsubscribe from dataset.updated events for a specific environment.
|
|
def unsubscribe_dataset_events(self, env_id: str, queue: asyncio.Queue):
|
|
if env_id in self._dataset_subscribers:
|
|
if queue in self._dataset_subscribers[env_id]:
|
|
self._dataset_subscribers[env_id].remove(queue)
|
|
if not self._dataset_subscribers[env_id]:
|
|
del self._dataset_subscribers[env_id]
|
|
# #endregion unsubscribe_dataset_events
|
|
|
|
# #region _broadcast_dataset_updated [C:3] [TYPE Function] [SEMANTICS dataset,broadcast,event,websocket]
|
|
# @BRIEF Broadcast dataset.updated event to all subscribers for a given environment.
|
|
def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]):
|
|
event = {"type": "dataset.updated", "payload": {"env_id": env_id, "dataset_ids": dataset_ids}}
|
|
if env_id in self._dataset_subscribers:
|
|
for queue in self._dataset_subscribers[env_id]:
|
|
try:
|
|
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
|
except Exception:
|
|
pass
|
|
# #endregion _broadcast_dataset_updated
|
|
# #endregion JobLifecycle
|
|
# #endregion JobLifecycleModule
|