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()