|
|
@@ -1,8 +1,10 @@
|
|
|
from pathlib import Path
|
|
|
+import subprocess
|
|
|
import sys
|
|
|
-from typing import Dict,Any
|
|
|
+from typing import Dict,Any, List, Optional
|
|
|
|
|
|
import httpx
|
|
|
+from pydantic import BaseModel,field_validator,Field
|
|
|
import redis
|
|
|
from src.services.subscription_manager import SubscriptionManager
|
|
|
from utils.config import config,APP_PATH
|
|
|
@@ -11,12 +13,39 @@ from utils.process_mgr import process_manager
|
|
|
import asyncio
|
|
|
from utils.logu import get_logger,logger
|
|
|
import os
|
|
|
+WORKER_DIR_BASE = APP_PATH.parent.parent
|
|
|
+py_client: Optional[Dict[str,Any]] = {
|
|
|
+ 'search': WORKER_DIR_BASE / 'worker\celery\client.py' ,
|
|
|
+ 'crawl': WORKER_DIR_BASE / 'worker\celery\crawl_client.py',
|
|
|
+ 'convert': WORKER_DIR_BASE / 'worker\celery\html_convert_tasks.py'
|
|
|
+}
|
|
|
|
|
|
|
|
|
+class WorkerModel(BaseModel):
|
|
|
+ name: str
|
|
|
+ queue_name: Optional[str] = Field(default=None,validate_default=True)
|
|
|
+ cmd: Optional[List[str]] = None
|
|
|
+ pid: Optional[int] = None
|
|
|
+
|
|
|
+ @field_validator("queue_name", mode="after")
|
|
|
+ @classmethod
|
|
|
+ def set_queue_name(cls, v: Optional[str], values) -> str:
|
|
|
+ if v is None:
|
|
|
+ if "name" not in values.data:
|
|
|
+ raise ValueError("name 字段缺失,无法生成 queue_name")
|
|
|
+ return f"{values.data['name']}_queue"
|
|
|
+ return v
|
|
|
+
|
|
|
+
|
|
|
class CeleryWorker:
|
|
|
+ def __init__(self, python_exe: str=sys.executable):
|
|
|
+ self.workers_model: Dict[str, WorkerModel] = {}
|
|
|
+ for worker_name in py_client.keys():
|
|
|
+ model = WorkerModel(name=worker_name)
|
|
|
+ model.cmd = [python_exe, '-m', 'celery', '-A', 'worker.celery.app', 'worker', '-Q',model.queue_name, f'--hostname={worker_name}@%h']
|
|
|
+ self.workers_model[worker_name] = model
|
|
|
async def run(self):
|
|
|
python_exe = sys.executable
|
|
|
- WORKER_DIR_BASE = APP_PATH.parent.parent
|
|
|
logger.info(f"{WORKER_DIR_BASE}")
|
|
|
# return
|
|
|
redis_cmd = [config.redis_exe]
|
|
|
@@ -33,12 +62,40 @@ class CeleryWorker:
|
|
|
await process_manager.start_process("flower", flower_cmd, cwd=WORKER_DIR_BASE)
|
|
|
proces = process_manager.processes.get("flower").get('process')
|
|
|
|
|
|
- search_worker_name = 'search'
|
|
|
- crawl_worker_name = 'crawl'
|
|
|
- convert_worker_name = 'convert'
|
|
|
- worker_list = [search_worker_name, crawl_worker_name, convert_worker_name]
|
|
|
- for worker_name in worker_list:
|
|
|
- await process_manager.start_process(f"{worker_name}_worker", [python_exe, '-m', 'celery', '-A', 'worker.celery.app', 'worker', '-Q',f'{worker_name}_queue', f'--hostname={worker_name}@%h'], cwd=WORKER_DIR_BASE)
|
|
|
+ async def start_all_workers(self):
|
|
|
+ for worker_name,worker_model in self.workers_model.items():
|
|
|
+ pid = await process_manager.start_process(worker_model.name, worker_model.cmd, cwd=WORKER_DIR_BASE)
|
|
|
+ worker_model.pid = pid
|
|
|
+ async def start_worker(self, name: str, in_cmd_windows: bool = False) -> WorkerModel:
|
|
|
+ if name in process_manager.processes:
|
|
|
+ return True
|
|
|
+ worker_model = self.workers_model.get(name)
|
|
|
+ if not worker_model:
|
|
|
+ raise ValueError(f"Invalid worker name: {name}")
|
|
|
+ if in_cmd_windows:
|
|
|
+ cmd = ['start','cmd', '/k' ]
|
|
|
+ cmd.extend(worker_model.cmd)
|
|
|
+ logger.info(f"run {' '.join(cmd)}")
|
|
|
+ process = subprocess.Popen(cmd, shell=True)
|
|
|
+ worker_model.pid = process.pid
|
|
|
+ else:
|
|
|
+ worker_model.pid = await process_manager.start_process(name, worker_model.cmd, cwd=WORKER_DIR_BASE)
|
|
|
+ return worker_model
|
|
|
+
|
|
|
+ async def stop_worker(self, name: str):
|
|
|
+ worker_model = self.workers_model.get(name)
|
|
|
+ if not worker_model:
|
|
|
+ raise ValueError(f"Invalid worker name: {name}")
|
|
|
+ await process_manager.stop_process(worker_model.name)
|
|
|
+ worker_model.pid = None
|
|
|
+ return worker_model
|
|
|
+
|
|
|
+ async def clean_worker_queue(self, name: str):
|
|
|
+ worker_model = self.workers_model.get(name)
|
|
|
+ if not worker_model:
|
|
|
+ raise ValueError(f"Invalid worker name: {name}")
|
|
|
+ queue_name = worker_model.queue_name
|
|
|
+ return subprocess.run([sys.executable, "-m", "celery", "-A", "worker.celery.app", "purge", "-Q", queue_name])
|
|
|
|
|
|
async def check_worker_status(self) -> Dict[str, Any]:
|
|
|
flower_url = "http://127.0.0.1:5555/workers?json=1"
|
|
|
@@ -69,3 +126,4 @@ class CeleryWorker:
|
|
|
except Exception as e:
|
|
|
return {"err": 1, "msg": f"Failed to connect to Redis: {e}"}
|
|
|
|
|
|
+celery_worker = CeleryWorker()
|