282 lines
12 KiB
Python
282 lines
12 KiB
Python
import os
|
||
import yaml
|
||
import asyncio
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core.logging import get_logger
|
||
from app.db.models import Dataset
|
||
# ✅ 确保引入 TaskExecution
|
||
from app.db.models.data_collection import CollectionTask, CollectionTemplate, TaskExecution
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.module.collection.client.datax_client import DataxClient
|
||
from app.module.collection.schema.collection import SyncMode, create_execute_record
|
||
from app.module.collection.schedule import validate_schedule_expression
|
||
from app.module.dataset.service.service import Service
|
||
from app.module.shared.schema import TaskStatus, NodeType, EdgeType
|
||
from app.module.shared.common.lineage import LineageService
|
||
from app.db.models.base_entity import LineageNode, LineageEdge
|
||
# ✅ 确保引入 RuntimeClient
|
||
from app.module.shared.runtime_client import RuntimeClient
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
# 执行模式开关：'local' (默认) | 'runtime'
|
||
COLLECTION_EXECUTION_MODE = os.getenv("COLLECTION_EXECUTION_MODE", "local")
|
||
|
||
class _TempExecution:
|
||
"""用于生成配置时的临时占位符"""
|
||
def __init__(self, task_id):
|
||
self.id = task_id
|
||
|
||
class CollectionTaskService:
|
||
def __init__(self, db: AsyncSession):
|
||
self.db = db
|
||
|
||
async def create_task(self, task: CollectionTask, dataset: Dataset) -> CollectionTask:
|
||
self.db.add(task)
|
||
if task.sync_mode == SyncMode.ONCE:
|
||
task.status = TaskStatus.RUNNING.name
|
||
await self.db.commit()
|
||
asyncio.create_task(CollectionTaskService.run_async(task.id, dataset.id if dataset else None))
|
||
elif task.sync_mode == SyncMode.SCHEDULED:
|
||
if not task.schedule_expression:
|
||
raise ValueError("schedule_expression is required for scheduled tasks")
|
||
validate_schedule_expression(task.schedule_expression)
|
||
return task
|
||
|
||
@staticmethod
|
||
async def run_async(task_id: str, dataset_id: str = None):
|
||
"""任务执行入口"""
|
||
logger.info(f"Start to execute task {task_id}")
|
||
|
||
# 1. 获取任务基础信息
|
||
async with AsyncSessionLocal() as session:
|
||
task_result = await session.execute(select(CollectionTask).where(CollectionTask.id == task_id))
|
||
task = task_result.scalar_one_or_none()
|
||
if not task:
|
||
logger.error(f"task {task_id} not exist")
|
||
return
|
||
|
||
template_result = await session.execute(select(CollectionTemplate).where(CollectionTemplate.id == task.template_id))
|
||
template = template_result.scalar_one_or_none()
|
||
if not template:
|
||
logger.error(f"template {task.template_name} not exist")
|
||
return
|
||
|
||
# 2. 创建执行记录 (此时 task_execution.log_path 已经根据 schema 规则被设置好)
|
||
task_execution = create_execute_record(task)
|
||
session.add(task_execution)
|
||
await session.commit()
|
||
await session.refresh(task_execution)
|
||
exec_id = task_execution.id
|
||
|
||
# 3. 根据模式分发
|
||
try:
|
||
if COLLECTION_EXECUTION_MODE == "runtime":
|
||
# ✅ 关键修改：将 task_execution 对象传入，以便提取 execution_id
|
||
await CollectionTaskService._execute_via_runtime(
|
||
task_id, dataset_id, task, template, task_execution
|
||
)
|
||
|
||
# 提交成功后更新状态为 SUBMITTED
|
||
async with AsyncSessionLocal() as update_session:
|
||
exec_record = await update_session.get(TaskExecution, exec_id)
|
||
if exec_record:
|
||
exec_record.status = "SUBMITTED"
|
||
await update_session.commit()
|
||
|
||
logger.info(f"Task {task_id} submitted to Runtime. Starting background poller.")
|
||
# 启动后台轮询
|
||
asyncio.create_task(
|
||
CollectionTaskService._poll_runtime_status(task_id, dataset_id, exec_id)
|
||
)
|
||
else:
|
||
# Local 模式
|
||
await CollectionTaskService._execute_locally(task_id, dataset_id, task, template, exec_id)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Task execution failed for {task_id}: {e}")
|
||
async with AsyncSessionLocal() as fail_session:
|
||
exec_record = await fail_session.get(TaskExecution, exec_id)
|
||
if exec_record:
|
||
exec_record.status = TaskStatus.FAILED.name
|
||
exec_record.error_message = str(e)
|
||
await fail_session.commit()
|
||
|
||
# ==========================================
|
||
# 🟢 模式 A：本地执行逻辑
|
||
# ==========================================
|
||
@staticmethod
|
||
async def _execute_locally(task_id, dataset_id, task, template, exec_id):
|
||
"""本地直接运行 DataX"""
|
||
logger.info(f"[Local Mode] Executing DataX for task {task_id}")
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
exec_record = await session.get(TaskExecution, exec_id)
|
||
# DataX 运行是阻塞的，放在线程池中运行
|
||
datax_client = DataxClient(task=task, execution=exec_record, template=template)
|
||
await asyncio.to_thread(datax_client.run_datax_job)
|
||
|
||
# 同步 Task 状态
|
||
task.status = exec_record.status
|
||
session.add(task)
|
||
|
||
if exec_record.status == TaskStatus.COMPLETED.name:
|
||
await CollectionTaskService._finalize_task(session, task_id, dataset_id, task)
|
||
|
||
await session.commit()
|
||
|
||
# ==========================================
|
||
# 🔵 模式 B：Runtime -> Ray 执行逻辑
|
||
# ==========================================
|
||
@staticmethod
|
||
async def _execute_via_runtime(task_id, dataset_id, task, template, task_execution):
|
||
"""
|
||
生成配置并提交给 Runtime。
|
||
不需要 session，因为 task_execution 对象已包含所有必要信息。
|
||
"""
|
||
logger.info(f"[Runtime Mode] Submitting task {task_id} to deepdata-runtime")
|
||
|
||
# ✅ 关键步骤：提取 execution_id，这是日志文件名的关键
|
||
execution_id = task_execution.id
|
||
log_path = task_execution.log_path
|
||
|
||
# 1. 生成 DataX JSON 配置
|
||
datax_client = DataxClient(task=task, execution=_TempExecution(task_id), template=template)
|
||
datax_json_path = datax_client.generate_config()
|
||
|
||
# 2. 构造 process.yaml 指令文件
|
||
flow_dir = Path(f"/flow/{task_id}")
|
||
flow_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
process_cfg = {
|
||
"executor_type": "deepdata",
|
||
"instance_id": task_id,
|
||
"execution_id": execution_id,
|
||
"log_path": log_path,
|
||
"dataset_id": dataset_id,
|
||
"cli_config": {
|
||
"command": "python",
|
||
"args": ["/opt/datax/bin/datax.py", datax_json_path],
|
||
"timeout": 7200
|
||
}
|
||
}
|
||
|
||
# 写入指令文件
|
||
with open(flow_dir / "process.yaml", "w", encoding="utf-8") as f:
|
||
yaml.safe_dump(process_cfg, f)
|
||
|
||
# 3. 提交给 Runtime
|
||
runtime_client = RuntimeClient()
|
||
success = await runtime_client.submit_task(task_id)
|
||
if not success:
|
||
raise RuntimeError("Runtime submission returned False")
|
||
|
||
@staticmethod
|
||
async def _poll_runtime_status(task_id, dataset_id, exec_id):
|
||
"""后台轮询 Runtime 状态"""
|
||
runtime_client = RuntimeClient()
|
||
|
||
for _ in range(2000):
|
||
try:
|
||
status_info = await runtime_client.get_task_status(task_id)
|
||
|
||
if status_info:
|
||
current_status = (status_info.get("status") or "").upper()
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
exec_record = await session.get(TaskExecution, exec_id)
|
||
if not exec_record: return
|
||
|
||
task_result = await session.execute(select(CollectionTask).where(CollectionTask.id == task_id))
|
||
task = task_result.scalar_one_or_none()
|
||
|
||
if current_status in ["COMPLETED", "SUCCEEDED", "SUCCESS"]:
|
||
exec_record.status = TaskStatus.COMPLETED.name
|
||
await session.commit()
|
||
logger.info(f"Task {task_id} completed in Runtime. Finalizing.")
|
||
if task:
|
||
await CollectionTaskService._finalize_task(session, task_id, dataset_id, task)
|
||
return
|
||
|
||
elif current_status in ["FAILED", "ERROR"]:
|
||
exec_record.status = TaskStatus.FAILED.name
|
||
exec_record.error_message = status_info.get("error", "Runtime task failed")
|
||
await session.commit()
|
||
logger.error(f"Task {task_id} failed in Runtime: {exec_record.error_message}")
|
||
return
|
||
|
||
await asyncio.sleep(5)
|
||
except Exception as e:
|
||
logger.warning(f"Polling error for {task_id}: {e}")
|
||
await asyncio.sleep(5)
|
||
|
||
# ==========================================
|
||
# 🛠️ 公共收尾逻辑
|
||
# ==========================================
|
||
@staticmethod
|
||
async def _finalize_task(session, task_id, dataset_id, task):
|
||
"""任务执行成功后的后续处理"""
|
||
if not dataset_id: return
|
||
|
||
# 🐛 修复：对于 S3/远端写入器，本地 target_path 可能不存在
|
||
# 应当检查 task.target_path 是否为本地路径，或者跳过本地扫描逻辑
|
||
target_path = Path(task.target_path)
|
||
if target_path.exists() and target_path.is_dir():
|
||
service = Service(db=session)
|
||
source_paths = []
|
||
for file_path in target_path.rglob('*'):
|
||
if file_path.is_file():
|
||
source_paths.append(str(file_path.absolute()))
|
||
|
||
if source_paths:
|
||
await service.add_files_to_dataset(dataset_id=dataset_id, source_paths=source_paths)
|
||
await CollectionTaskService._add_dataset_to_graph(session, dataset_id, task)
|
||
else:
|
||
logger.warning(f"[Finalize] Target path {target_path} does not exist locally. "
|
||
f"Skipping dataset import. (This is expected for s3writer/nfswriter)")
|
||
|
||
@staticmethod
|
||
async def _add_dataset_to_graph(session, dataset_id, task):
|
||
"""将数据集加入血缘图"""
|
||
try:
|
||
dataset = await session.get(Dataset, dataset_id)
|
||
if not dataset:
|
||
logger.warning(f"dataset {dataset_id} not found when building lineage graph")
|
||
return
|
||
|
||
dataset_node = LineageNode(
|
||
id=dataset.id,
|
||
node_type=NodeType.DATASET.value,
|
||
name=dataset.name,
|
||
description=dataset.description
|
||
)
|
||
|
||
collection_node = LineageNode(
|
||
id=task.id,
|
||
node_type=NodeType.DATASOURCE.value,
|
||
name=task.name,
|
||
description=task.description
|
||
)
|
||
|
||
collection_edge = LineageEdge(
|
||
process_id=task.id,
|
||
name=task.name,
|
||
edge_type=EdgeType.DATA_COLLECTION.value,
|
||
description=dataset.description,
|
||
from_node_id=task.id,
|
||
to_node_id=dataset.id
|
||
)
|
||
|
||
lineage_service = LineageService(db=session)
|
||
await lineage_service.generate_graph(collection_node, collection_edge, dataset_node)
|
||
await session.commit()
|
||
|
||
except Exception as exc:
|
||
logger.error(f"Failed to add dataset lineage graph: {exc}")
|
||
await session.rollback()
|