115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
import os
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Dict, Any, Optional
|
|
|
|
|
|
class WorkerFacade:
|
|
def __init__(
|
|
self,
|
|
process_core,
|
|
output_folder: str,
|
|
output_prefix: str,
|
|
output_extension: str,
|
|
step_name: str,
|
|
logger,
|
|
) -> None:
|
|
self.process_core = process_core
|
|
self.output_folder = output_folder
|
|
self.output_prefix = output_prefix
|
|
self.output_extension = output_extension
|
|
self.step_name = step_name
|
|
self.logger = logger
|
|
self.tasks: Dict[str, Dict[str, Any]] = {}
|
|
|
|
def submit(self, parameters: Dict[str, Any], input_paths: Dict[str, str]) -> str:
|
|
task_id = str(uuid.uuid4())
|
|
now = datetime.now()
|
|
self.tasks[task_id] = {
|
|
"task_id": task_id,
|
|
"status": "pending",
|
|
"progress": 0,
|
|
"message": "Task created",
|
|
"result": None,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"parameters": parameters,
|
|
"input_paths": input_paths,
|
|
}
|
|
return task_id
|
|
|
|
def run_task(self, task_id: str) -> None:
|
|
try:
|
|
task = self.tasks[task_id]
|
|
task["status"] = "running"
|
|
task["message"] = f"{self.step_name} running"
|
|
task["progress"] = 10
|
|
task["updated_at"] = datetime.now()
|
|
|
|
output_filename = f"{self.output_prefix}{task_id}{self.output_extension}"
|
|
output_path = os.path.join(self.output_folder, output_filename)
|
|
result = self.process_core(task["input_paths"], output_path, task["parameters"])
|
|
|
|
task["status"] = "completed"
|
|
task["message"] = f"{self.step_name} completed"
|
|
task["progress"] = 100
|
|
task["result"] = result
|
|
task["updated_at"] = datetime.now()
|
|
|
|
self.logger.info(f"Task completed: {task_id}")
|
|
|
|
except Exception as exc:
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task["status"] = "failed"
|
|
task["message"] = f"{self.step_name} failed: {str(exc)}"
|
|
task["updated_at"] = datetime.now()
|
|
self.logger.error(f"Task failed {task_id}: {str(exc)}")
|
|
|
|
def handle_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
|
|
job_id = message["job_id"]
|
|
input_paths = message.get("input_paths", {})
|
|
parameters = message.get("parameters", {})
|
|
output_dir = message.get("output_dir", self.output_folder)
|
|
output_filename = f"{self.output_prefix}{job_id}{self.output_extension}"
|
|
output_path = os.path.join(output_dir, output_filename)
|
|
|
|
try:
|
|
result = self.process_core(input_paths, output_path, parameters)
|
|
return {
|
|
"job_id": job_id,
|
|
"step": self.step_name,
|
|
"result": "ok",
|
|
"output_path": output_path,
|
|
"details": result,
|
|
}
|
|
except Exception as exc:
|
|
return {
|
|
"job_id": job_id,
|
|
"step": self.step_name,
|
|
"result": "fail",
|
|
"reason": str(exc),
|
|
}
|
|
|
|
def get_task(self, task_id: str) -> Optional[Dict[str, Any]]:
|
|
return self.tasks.get(task_id)
|
|
|
|
def list_tasks(self):
|
|
return list(self.tasks.values())
|
|
|
|
def cancel_task(self, task_id: str) -> Dict[str, Any]:
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
raise KeyError("Task not found")
|
|
if task["status"] in ["completed", "failed", "cancelled"]:
|
|
raise ValueError("Task cannot be cancelled")
|
|
task["status"] = "cancelled"
|
|
task["message"] = "Task cancelled"
|
|
task["updated_at"] = datetime.now()
|
|
return task
|
|
|
|
def active_tasks_count(self) -> int:
|
|
return sum(
|
|
1 for task in self.tasks.values() if task["status"] in ["pending", "running"]
|
|
)
|