From a2e631c4401b6a74e0c1c4eb561060ef2a72f562 Mon Sep 17 00:00:00 2001 From: MoeexT Date: Tue, 3 Nov 2026 16:15:02 +0800 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20=E8=AF=84=E4=BC=B0=E4=BB=BB=E5=8A=A1?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E6=97=B6=E4=B8=8D=E5=86=8D=E9=9D=99=E9=BB=98?= =?UTF-8?q?=E5=AE=8C=E6=88=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 数据项缺少 input 等字段时不再抛 TypeError，prompt 渲染统一转换为文本 - 单条评估失败记为 FAILED，全部失败时任务置为 FAILED 并返回 evaluation.0005 - 没有可评估数据项时报 evaluation.0004，避免进度 100% 但结果为空 - COT 校验只要求 instruction/output，JSON 单对象和 JSONL 空行可正常解析 - evaluate_item 不再自行 commit，避免并发协程共用会话时互相干扰 --- .../app/core/exception/codes.py | 6 + .../module/evaluation/service/evaluation.py | 133 ++++++-- .../app/module/shared/util/structured_file.py | 7 +- .../tests/test_module_evaluation.py | 317 ++++++++++++++++++ .../tests/test_module_shared.py | 52 ++- 5 files changed, 491 insertions(+), 24 deletions(-) diff --git a/runtime/datamate-python/app/core/exception/codes.py b/runtime/datamate-python/app/core/exception/codes.py index 5ef9d56eb..cf30d5e80 100644 --- a/runtime/datamate-python/app/core/exception/codes.py +++ b/runtime/datamate-python/app/core/exception/codes.py @@ -92,6 +92,12 @@ def __init__(self): EVALUATION_MODEL_NOT_FOUND: Final = ErrorCode( "evaluation.0003", "Evaluation model not found", 404 ) + EVALUATION_NO_ITEMS: Final = ErrorCode( + "evaluation.0004", "No evaluable items found in the selected data source", 400 + ) + EVALUATION_ALL_ITEMS_FAILED: Final = ErrorCode( + "evaluation.0005", "All evaluation items failed to be evaluated", 500 + ) # ========== 生成模块 ========== GENERATION_TASK_NOT_FOUND: Final = ErrorCode( diff --git a/runtime/datamate-python/app/module/evaluation/service/evaluation.py b/runtime/datamate-python/app/module/evaluation/service/evaluation.py index 9271f9c9f..68eaa7102 100644 --- a/runtime/datamate-python/app/module/evaluation/service/evaluation.py +++ b/runtime/datamate-python/app/module/evaluation/service/evaluation.py @@ -20,6 +20,9 @@ logger = get_logger(__name__) +_MAX_EVAL_ATTEMPTS = 3 + + class EvaluationExecutor: def __init__(self, db: AsyncSession, task: EvaluationTask): self.db = db @@ -28,17 +31,32 @@ def __init__(self, db: AsyncSession, task: EvaluationTask): async def save_eval_items(self): pass + @staticmethod + def _as_prompt_text(value) -> str: + """Render a dataset field as prompt text.""" + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + def get_eval_prompt(self, item: EvaluationItem) -> str: prompt_text = get_prompt(self.task.task_type, json.loads(self.task.eval_config).get("dimensions")) - eval_content = json.loads(item.eval_content) + eval_content = json.loads(item.eval_content or "{}") + if not isinstance(eval_content, dict): + logger.warning( + f"Eval content of item {item.id} is {type(eval_content).__name__}, not an object; " + "using it as the instruction text" + ) + eval_content = {"instruction": self._as_prompt_text(eval_content)} if self.task.task_type == "QA": - prompt_text = ((prompt_text.replace("{content}", eval_content.get("input")) - .replace("{question}", eval_content.get("instruction"))) - .replace("{answer}", eval_content.get("output"))) + prompt_text = ((prompt_text.replace("{content}", self._as_prompt_text(eval_content.get("input"))) + .replace("{question}", self._as_prompt_text(eval_content.get("instruction")))) + .replace("{answer}", self._as_prompt_text(eval_content.get("output")))) if self.task.task_type == "COT": - prompt_text = ((prompt_text.replace("{question}", eval_content.get("instruction")) - .replace("{conclusion}", eval_content.get("output"))) - .replace("{chain_of_thought}", eval_content.get("chain_of_thought"))) + prompt_text = ((prompt_text.replace("{question}", self._as_prompt_text(eval_content.get("instruction"))) + .replace("{conclusion}", self._as_prompt_text(eval_content.get("output")))) + .replace("{chain_of_thought}", self._as_prompt_text(eval_content.get("chain_of_thought")))) return prompt_text async def execute(self): @@ -51,6 +69,9 @@ async def execute(self): query = select(EvaluationItem).where(EvaluationItem.task_id == self.task.id) count_query = select(func.count()).select_from(query.subquery()) total = (await self.db.execute(count_query)).scalar_one() + if total <= 0: + logger.error(f"No evaluation items found for task {self.task.id}; aborting evaluation") + raise BusinessError(ErrorCodes.EVALUATION_NO_ITEMS, data={"task_id": self.task.id}) evaluated_count = 0 for file in files: items = (await self.db.execute(query.where(EvaluationItem.file_id == file.file_id))).scalars().all() @@ -58,21 +79,55 @@ async def execute(self): self.evaluate_item(models, item, semaphore) for item in items ] - await asyncio.gather(*tasks, return_exceptions=True) - file.evaluated_count = len(items) - evaluated_count += file.evaluated_count + results = await asyncio.gather(*tasks, return_exceptions=True) + for item, result in zip(items, results): + if isinstance(result, BaseException): + logger.error( + f"Unexpected error while evaluating task={self.task.id}, item={item.id}: {result!r}" + ) + item.status = TaskStatus.FAILED.value + completed_count = sum(1 for item in items if item.status == TaskStatus.COMPLETED.value) + file.evaluated_count = completed_count + evaluated_count += completed_count self.task.eval_process = evaluated_count / total await self.db.commit() + logger.info( + f"Evaluation finished for task {self.task.id}: {evaluated_count}/{total} items evaluated" + ) + if evaluated_count == 0: + logger.error(f"No item could be evaluated for task {self.task.id}; marking it as failed") + raise BusinessError( + ErrorCodes.EVALUATION_ALL_ITEMS_FAILED, data={"task_id": self.task.id} + ) async def evaluate_item(self, models, item: EvaluationItem, semaphore: asyncio.Semaphore): + """Evaluate one item and store the outcome on the ORM object. + + The caller owns the transaction: this method never commits, so concurrent + coroutines cannot race on the single AsyncSession shared by the executor. + """ async with semaphore: - max_try = 3 - while max_try > 0: + try: prompt_text = self.get_eval_prompt(item) - resp_text = await asyncio.to_thread( - call_openai_style_model, models.base_url, models.api_key, models.model_name, - prompt_text, + except Exception as e: + logger.error( + f"Failed to build evaluation prompt for task={self.task.id}, item={item.id}: {e}" ) + item.status = TaskStatus.FAILED.value + return + + for attempt in range(1, _MAX_EVAL_ATTEMPTS + 1): + try: + resp_text = await asyncio.to_thread( + call_openai_style_model, models.base_url, models.api_key, models.model_name, + prompt_text, + ) + except Exception as e: + logger.error( + f"Model call failed (attempt {attempt}/{_MAX_EVAL_ATTEMPTS}) for task={self.task.id}, " + f"file={item.file_id}, item={item.id}: {e}" + ) + continue resp_text = extract_json_substring(resp_text) try: json.loads(resp_text) @@ -80,13 +135,16 @@ async def evaluate_item(self, models, item: EvaluationItem, semaphore: asyncio.S logger.error( f"Failed to parse LLM answer as JSON for task={self.task.id}, file={item.file_id}: {e}. Raw answer: {resp_text!r}" ) - max_try -= 1 continue item.eval_result = resp_text item.status = TaskStatus.COMPLETED.value - await self.db.commit() return + logger.warning( + f"Evaluation gave up after {_MAX_EVAL_ATTEMPTS} attempts for task={self.task.id}, " + f"file={item.file_id}, item={item.id}" + ) + item.status = TaskStatus.FAILED.value def get_source_type(self) -> SourceType: pass @@ -101,11 +159,18 @@ async def save_eval_items(self): .where(DatasetFiles.dataset_id == self.task.source_id))) .scalars().all()) handler = StructuredFileHandlerFactory().get_handler(self.task.task_type) + parsed_count = 0 for dataset_file in dataset_files: if dataset_file.file_type.upper() != "JSON" and dataset_file.file_type.upper() != "JSONL": continue items = handler.get_items_from_file(dataset_file.file_path) logger.info(f"parse {len(items)} items from file {dataset_file.file_name}") + if not items: + logger.warning( + f"No valid {self.task.task_type} items parsed from file {dataset_file.file_name} " + f"({dataset_file.file_path}); file skipped" + ) + continue for item in items: self.db.add(EvaluationItem( id=str(uuid.uuid4()), @@ -128,6 +193,9 @@ async def save_eval_items(self): created_by=self.task.created_by, updated_by=self.task.updated_by, )) + parsed_count += len(items) + if parsed_count == 0: + raise BusinessError(ErrorCodes.EVALUATION_NO_ITEMS, data={"task_id": self.task.id}) def get_source_type(self) -> SourceType: return SourceType.DATASET @@ -141,11 +209,17 @@ async def save_eval_items(self): synthesis_files = ((await self.db.execute(select(DataSynthesisFileInstance) .where(DataSynthesisFileInstance.synthesis_instance_id == self.task.source_id))) .scalars().all()) + parsed_count = 0 for synthesis_file in synthesis_files: synthesis_datas = ((await self.db.execute(select(SynthesisData) .where(SynthesisData.synthesis_file_instance_id == synthesis_file.id))) .scalars().all()) logger.info(f"get {len(synthesis_datas)} items from file {synthesis_file.file_name}") + if not synthesis_datas: + logger.warning( + f"No synthesized data found for file {synthesis_file.file_name}; file skipped" + ) + continue for synthesis_data in synthesis_datas: self.db.add(EvaluationItem( id=str(uuid.uuid4()), @@ -168,7 +242,9 @@ async def save_eval_items(self): created_by=self.task.created_by, updated_by=self.task.updated_by, )) - pass + parsed_count += len(synthesis_datas) + if parsed_count == 0: + raise BusinessError(ErrorCodes.EVALUATION_NO_ITEMS, data={"task_id": self.task.id}) def get_source_type(self) -> SourceType: return SourceType.SYNTHESIS @@ -197,25 +273,41 @@ async def run_evaluation_task(task_id: str): - task_id: id of EvaluationTaskModel """ logger.info(f"Background evaluation worker started add items for task {task_id}") + items_ready = False async with AsyncSessionLocal() as session: + task = None try: task = await session.execute(select(EvaluationTask).where(EvaluationTask.id == task_id)) task = task.scalar_one_or_none() + if task is None: + logger.error(f"Evaluation task {task_id} not found, skip background worker") + return factory = EvaluationExecutorFactory(session, task) executor = factory.get_executor(task.source_type) await executor.save_eval_items() + task.eval_process = 0.0 task.status = TaskStatus.RUNNING.value + items_ready = True except Exception as e: logger.error(f"Background worker encountered error for task {task_id}: {e}") - task.status = TaskStatus.FAILED.value + if task is not None: + task.status = TaskStatus.FAILED.value finally: await session.commit() + if not items_ready: + logger.error(f"Skip evaluation for task {task_id} because item preparation failed") + return + logger.info(f"Background evaluation worker started for task {task_id}") async with AsyncSessionLocal() as session: + task = None try: task = await session.execute(select(EvaluationTask).where(EvaluationTask.id == task_id)) task = task.scalar_one_or_none() + if task is None: + logger.error(f"Evaluation task {task_id} not found, skip background worker") + return factory = EvaluationExecutorFactory(session, task) executor = factory.get_executor(task.source_type) await executor.execute() @@ -223,6 +315,7 @@ async def run_evaluation_task(task_id: str): task.status = TaskStatus.COMPLETED.value except Exception as e: logger.error(f"Background worker encountered error for task {task_id}: {e}") - task.status = TaskStatus.FAILED.value + if task is not None: + task.status = TaskStatus.FAILED.value finally: await session.commit() diff --git a/runtime/datamate-python/app/module/shared/util/structured_file.py b/runtime/datamate-python/app/module/shared/util/structured_file.py index ef3232c63..b6bef670e 100644 --- a/runtime/datamate-python/app/module/shared/util/structured_file.py +++ b/runtime/datamate-python/app/module/shared/util/structured_file.py @@ -26,10 +26,13 @@ def get_items_from_file(self, file_path: str) -> list[dict]: data = json.load(f) if not self.validate_json(data): return items - items = data + # A JSON file may hold a single item object instead of a list. + items = [data] if isinstance(data, dict) else data elif file_type == "JSONL": with open(file_path, "r", encoding="utf-8") as f: for line in f: + if not line.strip(): + continue data = json.loads(line) if not self.validate_json(data): continue @@ -77,7 +80,7 @@ def __init__(self): "output": {"type": "string"}, "chain_of_thought": {"type": "string"} }, - "required": ["question", "instruction", "output"], + "required": ["instruction", "output"], } self.schema_list = { "type": "array", diff --git a/runtime/datamate-python/tests/test_module_evaluation.py b/runtime/datamate-python/tests/test_module_evaluation.py index a95159c03..80534e545 100644 --- a/runtime/datamate-python/tests/test_module_evaluation.py +++ b/runtime/datamate-python/tests/test_module_evaluation.py @@ -1,5 +1,20 @@ +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from app.core.exception import BusinessError, ErrorCodes +from app.db.models import EvaluationItem, EvaluationTask +from app.db.models.data_evaluation import EvaluationFile from app.module.evaluation.schema.prompt import EVALUATION_PROMPT_TEMPLATE +from app.module.evaluation.service import evaluation as evaluation_module +from app.module.evaluation.service.evaluation import ( + DatasetEvaluationExecutor, + EvaluationExecutor, +) from app.module.evaluation.service.prompt_template_service import PromptTemplateService +from app.module.shared.schema import TaskStatus def test_get_prompt_templates_size_matches_source() -> None: @@ -46,3 +61,305 @@ def test_prompt_template_dimension_fields_are_non_none() -> None: for dim in item.defaultDimensions: assert dim.dimension is not None assert dim.description is not None + + +class _EmptyResult: + def scalars(self): + return self + + def all(self): + return [] + + def scalar_one(self): + return 0 + + +class _EmptySession: + async def execute(self, _statement): + return _EmptyResult() + + +def _build_task(task_type: str) -> EvaluationTask: + return EvaluationTask( + id="task-1", + name="evaluation-task", + task_type=task_type, + source_type="DATASET", + source_id="dataset-1", + source_name="dataset", + eval_config=json.dumps({"modelId": "model-1", "dimensions": []}), + status="PENDING", + eval_process=0.0, + ) + + +def test_execute_raises_business_error_when_task_has_no_items(monkeypatch) -> None: + async def fake_get_model_by_id(_db, _model_id): + return object() + + monkeypatch.setattr(evaluation_module, "get_model_by_id", fake_get_model_by_id) + executor = DatasetEvaluationExecutor(_EmptySession(), _build_task("COT")) + + with pytest.raises(BusinessError) as exc_info: + asyncio.run(executor.execute()) + + assert exc_info.value.code == ErrorCodes.EVALUATION_NO_ITEMS.code + + +def test_cot_eval_prompt_tolerates_missing_chain_of_thought() -> None: + executor = EvaluationExecutor(None, _build_task("COT")) + item = EvaluationItem( + id="item-1", + task_id="task-1", + file_id="file-1", + item_id="1", + eval_content=json.dumps({"instruction": "1+1=?", "output": "2"}), + ) + + prompt = executor.get_eval_prompt(item) + + assert "1+1=?" in prompt + assert "## 思维链:\n\n" in prompt + + +def test_qa_eval_prompt_tolerates_missing_input_field() -> None: + executor = EvaluationExecutor(None, _build_task("QA")) + item = EvaluationItem( + id="item-1", + task_id="task-1", + file_id="file-1", + item_id="1", + eval_content=json.dumps({"instruction": "什么是向量检索？", "output": "按相似度检索向量。"}), + ) + + prompt = executor.get_eval_prompt(item) + + assert "{content}" not in prompt + assert "{question}" not in prompt + assert "什么是向量检索？" in prompt + assert "按相似度检索向量。" in prompt + + +def test_qa_eval_prompt_handles_non_object_eval_content() -> None: + executor = EvaluationExecutor(None, _build_task("QA")) + item = EvaluationItem( + id="item-1", + task_id="task-1", + file_id="file-1", + item_id="1", + eval_content=json.dumps(["a", "b"]), + ) + + prompt = executor.get_eval_prompt(item) + + assert '["a", "b"]' in prompt + + +class _CountingSession: + def __init__(self) -> None: + self.commits = 0 + + async def commit(self) -> None: + self.commits += 1 + + +def test_evaluate_item_defers_commit_to_the_executor(monkeypatch) -> None: + monkeypatch.setattr( + evaluation_module, + "call_openai_style_model", + lambda *args, **kwargs: '{"result": {"问题是否独立": "Y"}, "evaluation": "ok"}', + ) + session = _CountingSession() + executor = EvaluationExecutor(session, _build_task("QA")) + item = EvaluationItem( + id="item-1", + task_id="task-1", + file_id="file-1", + item_id="1", + eval_content=json.dumps({"instruction": "什么是向量检索？", "output": "按相似度检索向量。"}), + eval_result="{}", + status=TaskStatus.PENDING.value, + ) + models = SimpleNamespace(base_url="http://example.test/v1", api_key="k", model_name="m") + + asyncio.run(executor.evaluate_item(models, item, asyncio.Semaphore(1))) + + assert item.status == TaskStatus.COMPLETED.value + assert json.loads(item.eval_result)["result"]["问题是否独立"] == "Y" + assert session.commits == 0 + + +def test_evaluate_item_marks_item_failed_when_model_keeps_failing(monkeypatch) -> None: + def broken_model(*args, **kwargs): + raise RuntimeError("connection refused") + + monkeypatch.setattr(evaluation_module, "call_openai_style_model", broken_model) + executor = EvaluationExecutor(_CountingSession(), _build_task("QA")) + item = EvaluationItem( + id="item-1", + task_id="task-1", + file_id="file-1", + item_id="1", + eval_content=json.dumps({"instruction": "q", "output": "a"}), + eval_result="{}", + status=TaskStatus.PENDING.value, + ) + models = SimpleNamespace(base_url="http://example.test/v1", api_key="k", model_name="m") + + asyncio.run(executor.evaluate_item(models, item, asyncio.Semaphore(1))) + + assert item.status == TaskStatus.FAILED.value + assert item.eval_result == "{}" + + +class _ScriptedResult: + def __init__(self, rows=(), scalar=None) -> None: + self._rows = list(rows) + self._scalar = scalar + + def scalars(self): + return self + + def all(self): + return self._rows + + def scalar_one(self): + return self._scalar + + +class _ScriptedSession: + def __init__(self, results) -> None: + self._results = list(results) + self.commits = 0 + + async def execute(self, _statement): + return self._results.pop(0) + + async def commit(self) -> None: + self.commits += 1 + + +def test_execute_reports_failure_when_every_item_fails(monkeypatch) -> None: + async def fake_get_model_by_id(_db, _model_id): + return SimpleNamespace(base_url="http://example.test/v1", api_key="k", model_name="m") + + async def fake_evaluate_item(self, _models, item, _semaphore): + item.status = TaskStatus.FAILED.value + + monkeypatch.setattr(evaluation_module, "get_model_by_id", fake_get_model_by_id) + monkeypatch.setattr(EvaluationExecutor, "evaluate_item", fake_evaluate_item) + + task = _build_task("QA") + eval_file = EvaluationFile( + id="file-row-1", + task_id=task.id, + file_id="file-1", + file_name="qa.jsonl", + total_count=2, + evaluated_count=2, + ) + items = [ + EvaluationItem( + id=f"item-{index}", + task_id=task.id, + file_id="file-1", + item_id=str(index), + eval_content=json.dumps({"instruction": "q", "output": "a"}), + eval_result="{}", + status=TaskStatus.COMPLETED.value, + ) + for index in range(2) + ] + session = _ScriptedSession( + [ + _ScriptedResult(rows=[eval_file]), + _ScriptedResult(scalar=2), + _ScriptedResult(rows=items), + ] + ) + executor = EvaluationExecutor(session, task) + + with pytest.raises(BusinessError) as exc_info: + asyncio.run(executor.execute()) + + assert exc_info.value.code == ErrorCodes.EVALUATION_ALL_ITEMS_FAILED.code + assert eval_file.evaluated_count == 0 + assert task.eval_process == 0 + assert session.commits == 1 + + +class _StubResult: + def __init__(self, rows): + self._rows = rows + + def scalars(self): + return self + + def all(self): + return self._rows + + +class _CapturingSession: + def __init__(self, rows): + self._rows = rows + self.added = [] + + async def execute(self, _statement): + return _StubResult(self._rows) + + def add(self, obj): + self.added.append(obj) + + +def _write_cot_jsonl(path, rows) -> None: + path.write_text( + "\n".join(json.dumps(row, ensure_ascii=False) for row in rows) + "\n", + encoding="utf-8", + ) + + +def test_dataset_executor_parses_cot_files_without_question_field(tmp_path) -> None: + file_path = tmp_path / "cot.jsonl" + _write_cot_jsonl( + file_path, + [ + {"instruction": "1+1=?", "output": "2", "chain_of_thought": "1 plus 1 equals 2"}, + {"instruction": "2+2=?", "output": "4", "chain_of_thought": "2 plus 2 equals 4"}, + ], + ) + dataset_file = SimpleNamespace( + id="file-1", + file_name="cot.jsonl", + file_type="JSONL", + file_path=str(file_path), + ) + session = _CapturingSession([dataset_file]) + executor = DatasetEvaluationExecutor(session, _build_task("COT")) + + asyncio.run(executor.save_eval_items()) + + items = [obj for obj in session.added if isinstance(obj, EvaluationItem)] + files = [obj for obj in session.added if isinstance(obj, EvaluationFile)] + assert len(items) == 2 + assert json.loads(items[0].eval_content)["output"] == "2" + assert len(files) == 1 + assert files[0].total_count == 2 + + +def test_dataset_executor_fails_clearly_when_no_items_can_be_parsed(tmp_path) -> None: + file_path = tmp_path / "cot.jsonl" + _write_cot_jsonl(file_path, [{"foo": "bar"}]) + dataset_file = SimpleNamespace( + id="file-1", + file_name="cot.jsonl", + file_type="JSONL", + file_path=str(file_path), + ) + session = _CapturingSession([dataset_file]) + executor = DatasetEvaluationExecutor(session, _build_task("COT")) + + with pytest.raises(BusinessError) as exc_info: + asyncio.run(executor.save_eval_items()) + + assert exc_info.value.code == ErrorCodes.EVALUATION_NO_ITEMS.code + assert session.added == [] diff --git a/runtime/datamate-python/tests/test_module_shared.py b/runtime/datamate-python/tests/test_module_shared.py index e81b7e499..81101bd7c 100644 --- a/runtime/datamate-python/tests/test_module_shared.py +++ b/runtime/datamate-python/tests/test_module_shared.py @@ -66,9 +66,57 @@ def test_get_items_from_json_file_for_qa(tmp_path) -> None: assert items[0]["instruction"] == "q1" -def test_cot_handler_validate_json_requires_question_field() -> None: +def test_cot_handler_validate_json_accepts_instruction_and_output() -> None: handler = COTItemHandler() - assert handler.validate_json({"instruction": "x", "output": "y"}) is False + assert handler.validate_json({"instruction": "x", "output": "y"}) is True + assert handler.validate_json( + {"instruction": "x", "output": "y", "chain_of_thought": "z"} + ) is True + + +def test_cot_handler_validate_json_rejects_missing_output() -> None: + handler = COTItemHandler() + assert handler.validate_json({"instruction": "x", "chain_of_thought": "z"}) is False + + +def test_get_items_from_jsonl_parses_cot_rows(tmp_path) -> None: + file_path = tmp_path / "cot.jsonl" + rows = [ + {"instruction": "1+1=?", "output": "2", "chain_of_thought": "1 plus 1 equals 2"}, + {"instruction": "missing output"}, + { + "instruction": "2+2=?", + "output": "4", + "chain_of_thought": "2 plus 2 equals 4", + }, + ] + file_path.write_text( + "\n".join(json.dumps(r, ensure_ascii=False) for r in rows) + "\n", + encoding="utf-8", + ) + + handler = COTItemHandler() + items = handler.get_items_from_file(str(file_path)) + + assert len(items) == 2 + assert items[1]["output"] == "4" + + +def test_get_items_from_json_file_wraps_single_object(tmp_path) -> None: + file_path = tmp_path / "cot.json" + file_path.write_text( + json.dumps( + {"instruction": "1+1=?", "output": "2", "chain_of_thought": "1 plus 1 equals 2"}, + ensure_ascii=False, + ), + encoding="utf-8", + ) + + handler = COTItemHandler() + items = handler.get_items_from_file(str(file_path)) + + assert len(items) == 1 + assert items[0]["instruction"] == "1+1=?" def test_factory_get_handler_returns_cot_handler() -> None: From 87430aec7daea2ddcf378ed0829f87c32c9d7751 Mon Sep 17 00:00:00 2001 From: MoeexT Date: Tue, 3 Nov 2026 17:27:56 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20=E6=B8=85=E6=B4=97=E4=BB=BB=E5=8A=A1?= =?UTF-8?q?=E8=BD=AE=E8=AF=A2=E4=B8=8D=E5=86=8D=E6=8A=8A=E8=BF=90=E8=A1=8C?= =?UTF-8?q?=E4=B8=AD=E7=9A=84=E4=BB=BB=E5=8A=A1=E8=AF=AF=E5=88=A4=E4=B8=BA?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 轮询拿不到结论（runtime 不可达、返回 NOT_FOUND 或异常响应）时保持 RUNNING，只有持续超时才兜底置 FAILED：默认 7200s，可用 CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS 覆盖。此前一次 60s 超时就判失败，而数据面稍后才写入成功终态，于是出现先失败、过一会儿又变成功的现象。 每轮轮询先回读数据库终态，数据面已写入终态时不再覆盖；写入终态改为带状态条件的更新，保证幂等。状态查询单独设置 10s 超时，不再沿用提交任务的 60s 超时。 另外禁止重复提交运行中的任务，避免起第二个 Ray Job 并产生重复清洗结果；需要重跑时先停止任务。 新增 test_module_cleaning_scheduler.py 与 test_module_cleaning_service.py 覆盖上述行为。 --- .../repository/cleaning_task_repository.py | 48 ++- .../app/module/cleaning/runtime_client.py | 6 +- .../service/cleaning_task_scheduler.py | 278 ++++++++---- .../cleaning/service/cleaning_task_service.py | 14 +- .../tests/test_module_cleaning_scheduler.py | 395 ++++++++++++++++++ .../tests/test_module_cleaning_service.py | 66 +++ 6 files changed, 727 insertions(+), 80 deletions(-) create mode 100644 runtime/datamate-python/tests/test_module_cleaning_scheduler.py create mode 100644 runtime/datamate-python/tests/test_module_cleaning_service.py diff --git a/runtime/datamate-python/app/module/cleaning/repository/cleaning_task_repository.py b/runtime/datamate-python/app/module/cleaning/repository/cleaning_task_repository.py index f4c6b50ed..dfc8219e9 100644 --- a/runtime/datamate-python/app/module/cleaning/repository/cleaning_task_repository.py +++ b/runtime/datamate-python/app/module/cleaning/repository/cleaning_task_repository.py @@ -1,10 +1,20 @@ +from datetime import datetime from typing import List, Optional from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy import select, delete, func +from sqlalchemy import select, delete, func, update from app.db.models.cleaning import CleaningTask from app.module.cleaning.schema import CleaningTaskDto +# 终态集合：写入终态后不再接受控制面的状态覆盖 +TERMINAL_TASK_STATUSES = ( + "COMPLETED", + "PARTIAL_SUCCESS", + "FAILED", + "STOPPED", +) + + class CleaningTaskRepository: """Repository for cleaning task operations""" @@ -127,6 +137,42 @@ async def update_task(self, db: AsyncSession, task: CleaningTaskDto) -> None: await db.flush() + async def find_status(self, db: AsyncSession, task_id: str) -> Optional[str]: + """Query task status only (cheap read for polling)""" + query = select(self.model.status).where(self.model.id == task_id) + result = await db.execute(query) + return result.scalar_one_or_none() + + async def finish_if_unfinished( + self, + db: AsyncSession, + task_id: str, + status: str, + finished_at: Optional[datetime] = None + ) -> bool: + """Write a terminal status only when the task is not terminal yet. + + Returns True when the row was updated, False when the task was missing + or already in a terminal state. This keeps the control plane from + overwriting a status the executor already wrote directly to the table. + """ + values = {"status": status} + if finished_at is not None: + values["finished_at"] = finished_at + + query = ( + update(self.model) + .where( + self.model.id == task_id, + self.model.status.notin_(TERMINAL_TASK_STATUSES), + ) + .values(**values) + .execution_options(synchronize_session=False) + ) + result = await db.execute(query) + await db.flush() + return bool(result.rowcount and result.rowcount > 0) + async def delete_task_by_id(self, db: AsyncSession, task_id: str) -> None: """Delete task by ID""" query = delete(self.model).where(self.model.id == task_id) diff --git a/runtime/datamate-python/app/module/cleaning/runtime_client.py b/runtime/datamate-python/app/module/cleaning/runtime_client.py index 97c3523c3..95fa31225 100644 --- a/runtime/datamate-python/app/module/cleaning/runtime_client.py +++ b/runtime/datamate-python/app/module/cleaning/runtime_client.py @@ -4,6 +4,10 @@ logger = get_logger(__name__) +# 状态查询是轻量的本地接口，不需要沿用提交任务的 60s 超时； +# 超时只会让这一轮轮询“无结论”，不会再据此判失败。 +STATUS_REQUEST_TIMEOUT_SECONDS = 10.0 + class RuntimeClient: """HTTP client for communicating with runtime service""" @@ -46,7 +50,7 @@ async def get_task_status(self, task_id: str) -> Optional[dict]: """Get task status from runtime""" try: url = f"{self.base_url}/api/task/{task_id}/status" - response = await self.client.get(url) + response = await self.client.get(url, timeout=STATUS_REQUEST_TIMEOUT_SECONDS) response.raise_for_status() return response.json() except httpx.HTTPError as e: diff --git a/runtime/datamate-python/app/module/cleaning/service/cleaning_task_scheduler.py b/runtime/datamate-python/app/module/cleaning/service/cleaning_task_scheduler.py index 94ca825bb..b8015115b 100644 --- a/runtime/datamate-python/app/module/cleaning/service/cleaning_task_scheduler.py +++ b/runtime/datamate-python/app/module/cleaning/service/cleaning_task_scheduler.py @@ -1,5 +1,11 @@ import asyncio +import os +import time +from datetime import datetime +from typing import Callable, Optional + from sqlalchemy.ext.asyncio import AsyncSession + from app.core.logging import get_logger from app.module.cleaning.repository import CleaningTaskRepository from app.module.cleaning.runtime_client import RuntimeClient @@ -9,6 +15,23 @@ # 模块级单例：进程内所有请求共享同一个调度器 _scheduler: "CleaningTaskScheduler | None" = None +POLL_INTERVAL_SECONDS = 2 +# runtime 报 RUNNING（或返回无法判定的结果）时，多久回读一次数据库。 +# 数据面（Ray Job 执行器）会直接把终态写进 t_clean_task，控制面只是跟随者。 +DB_RECHECK_INTERVAL_SECONDS = 10 +# runtime 长时间给不出任何结论（不可达 / NOT_FOUND）时，控制面兜底标记 FAILED 的超时。 +# 必须明显大于“sidecar 重启后仍在跑”的任务最长耗时，否则又会把运行中的任务判成失败。 +DEFAULT_ORPHAN_TIMEOUT_SECONDS = 7200 +# 同一个未决任务重复告警的最小间隔 +UNRESOLVED_LOG_INTERVAL_SECONDS = 60 + +# runtime 明确给出终态时才写库，避免控制面覆盖数据面已写入的结果 +RUNTIME_COMPLETED_STATUSES = {"completed", "succeeded", "success"} +RUNTIME_FAILED_STATUSES = {"failed", "error"} +RUNTIME_STOPPED_STATUSES = {"cancelled", "canceled", "stopped"} +RUNTIME_ACTIVE_STATUSES = {"pending", "running", "submitted", "submitting", "queued"} +RUNTIME_NOT_FOUND_STATUS = "not_found" + def get_scheduler() -> "CleaningTaskScheduler": """获取全局调度器单例，不存在则创建""" @@ -21,21 +44,66 @@ def get_scheduler() -> "CleaningTaskScheduler": return _scheduler +def _orphan_timeout_seconds() -> float: + """runtime 无法给出结论时，控制面兜底失败的等待时间（秒）""" + raw = os.getenv("CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS") + if raw is None or not raw.strip(): + return DEFAULT_ORPHAN_TIMEOUT_SECONDS + try: + value = float(raw) + except (TypeError, ValueError): + logger.warning( + f"[Polling] Invalid CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS={raw!r}, " + f"using default {DEFAULT_ORPHAN_TIMEOUT_SECONDS}s" + ) + return DEFAULT_ORPHAN_TIMEOUT_SECONDS + if value <= 0: + logger.warning( + f"[Polling] CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS must be positive, " + f"using default {DEFAULT_ORPHAN_TIMEOUT_SECONDS}s" + ) + return DEFAULT_ORPHAN_TIMEOUT_SECONDS + return value + + class CleaningTaskScheduler: """Scheduler for executing cleaning tasks""" - def __init__(self, task_repo: CleaningTaskRepository, runtime_client: RuntimeClient): + def __init__( + self, + task_repo: CleaningTaskRepository, + runtime_client: RuntimeClient, + session_factory: Optional[Callable[[], AsyncSession]] = None, + ): self.task_repo = task_repo self.runtime_client = runtime_client + self._session_factory = session_factory self._polling_task_ids: set[str] = set() # 待轮询的任务 ID 集合 self._polling_started: bool = False - self._poll_failure_count: dict[str, int] = {} # runtime 不可达连续失败计数 - self._MAX_POLL_FAILURES = 1 # 一次超时(60s) → 标记 FAILED + # runtime 无法给出结论的起始时刻（monotonic），用于兜底超时 + self._unresolved_since: dict[str, float] = {} + # 兜底告警 / 数据库回读的节流时间戳 + self._last_log_at: dict[str, float] = {} + self._last_db_check_at: dict[str, float] = {} + + def _get_session_factory(self) -> Callable[[], AsyncSession]: + if self._session_factory is not None: + return self._session_factory + + from app.db.session import AsyncSessionLocal + + return AsyncSessionLocal + + def _forget(self, task_id: str) -> None: + """任务结束轮询，清理全部内存状态""" + self._polling_task_ids.discard(task_id) + self._unresolved_since.pop(task_id, None) + self._last_log_at.pop(task_id, None) + self._last_db_check_at.pop(task_id, None) async def execute_task(self, db: AsyncSession, task_id: str, retry_count: int) -> bool: """Execute cleaning task""" from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus - from datetime import datetime task = CleaningTaskDto() task.id = task_id @@ -47,6 +115,7 @@ async def execute_task(self, db: AsyncSession, task_id: str, retry_count: int) - if submitted: await self.task_repo.update_task(db, task) + self._forget(task_id) self._polling_task_ids.add(task_id) return submitted @@ -58,28 +127,24 @@ async def execute_task(self, db: AsyncSession, task_id: str, retry_count: int) - async def stop_task(self, db: AsyncSession, task_id: str) -> bool: """Stop cleaning task""" - from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus + from app.module.cleaning.schema import CleaningTaskStatus - self._polling_task_ids.discard(task_id) - self._poll_failure_count.pop(task_id, None) + self._forget(task_id) await self.runtime_client.stop_task(task_id) - task = CleaningTaskDto() - task.id = task_id - task.status = CleaningTaskStatus.STOPPED - - await self.task_repo.update_task(db, task) + await self.task_repo.finish_if_unfinished( + db, task_id, CleaningTaskStatus.STOPPED, datetime.now() + ) return True async def startup(self): """进程启动时调用：恢复未完成任务的轮询 + 启动全局轮询协程""" from app.module.cleaning.schema import CleaningTaskStatus - from app.db.session import AsyncSessionLocal # 从数据库恢复所有 RUNNING 状态的任务 try: - async with AsyncSessionLocal() as db: + async with self._get_session_factory()() as db: tasks = await self.task_repo.find_tasks(db, status=CleaningTaskStatus.RUNNING) for task in tasks: logger.info(f"[Polling] Recovered RUNNING task from DB: {task.id}") @@ -94,71 +159,134 @@ async def startup(self): logger.info("[Polling] Global polling loop started") async def _poll_all_tasks(self): - """全局轮询协程：每 2 秒轮询所有 RUNNING 任务的 runtime 状态""" - from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus - from app.db.session import AsyncSessionLocal - from datetime import datetime - + """全局轮询协程：轮询所有 RUNNING 任务的 runtime 状态""" logger.info("[Polling] Global status polling loop started") - terminal_statuses = {"completed", "failed", "cancelled", "stopped", "not_found"} while True: - task_ids = list(self._polling_task_ids) - if not task_ids: - # 没有待轮询任务，休眠后继续等 - await asyncio.sleep(2) - continue - - for task_id in task_ids: + for task_id in list(self._polling_task_ids): try: - status_data = await self.runtime_client.get_task_status(task_id) - if status_data is None: - # runtime 不可达，累计失败计数 - count = self._poll_failure_count.get(task_id, 0) + 1 - self._poll_failure_count[task_id] = count - logger.warning( - f"[Polling] Task {task_id} unreachable ({count}/{self._MAX_POLL_FAILURES})" - ) - if count >= self._MAX_POLL_FAILURES: - logger.error( - f"[Polling] Task {task_id} marked FAILED: " - f"runtime unreachable after {count} attempts" - ) - async with AsyncSessionLocal() as db: - task_dto = CleaningTaskDto() - task_dto.id = task_id - task_dto.status = CleaningTaskStatus.FAILED - task_dto.finished_at = datetime.now() - await self.task_repo.update_task(db, task_dto) - await db.commit() - self._polling_task_ids.discard(task_id) - self._poll_failure_count.pop(task_id, None) - continue - - # runtime 可达，重置失败计数 - self._poll_failure_count.pop(task_id, None) - - current_status = (status_data.get("status", "") or "").lower() - logger.debug(f"[Polling] Task {task_id} status: {current_status}") - - if current_status in terminal_statuses: - async with AsyncSessionLocal() as db: - task = CleaningTaskDto() - task.id = task_id - if current_status == "completed": - task.status = CleaningTaskStatus.COMPLETED - else: - task.status = CleaningTaskStatus.FAILED - task.finished_at = datetime.now() - await self.task_repo.update_task(db, task) - await db.commit() - logger.info( - f"[Polling] Task {task_id} finished: {current_status}" - ) - self._polling_task_ids.discard(task_id) - self._poll_failure_count.pop(task_id, None) - + await self._poll_task(task_id) except Exception as e: + # 单次轮询异常（数据库抖动等）不能把任务判成失败，下个周期重试 logger.error(f"[Polling] Error polling task {task_id}: {e}") - await asyncio.sleep(2) + await asyncio.sleep(POLL_INTERVAL_SECONDS) + + async def _poll_task(self, task_id: str): + """轮询单个任务：先看数据库，再问 runtime""" + from app.module.cleaning.schema import CleaningTaskStatus + + if await self._task_already_finished(task_id): + return + + status_data = await self.runtime_client.get_task_status(task_id) + if status_data is None: + # runtime 不可达（重启中 / 网络抖动）：保持 RUNNING，等它回来 + await self._handle_unresolved(task_id, "runtime unreachable") + return + + runtime_status = str(status_data.get("status") or "").strip().lower() + logger.debug(f"[Polling] Task {task_id} status: {runtime_status!r}") + + if runtime_status in RUNTIME_COMPLETED_STATUSES: + await self._finish_task(task_id, CleaningTaskStatus.COMPLETED, runtime_status) + self._forget(task_id) + return + + if runtime_status in RUNTIME_FAILED_STATUSES: + error = status_data.get("error") + if error: + logger.error(f"[Polling] Task {task_id} failed in runtime: {error}") + await self._finish_task(task_id, CleaningTaskStatus.FAILED, runtime_status) + self._forget(task_id) + return + + if runtime_status in RUNTIME_STOPPED_STATUSES: + await self._finish_task(task_id, CleaningTaskStatus.STOPPED, runtime_status) + self._forget(task_id) + return + + if runtime_status in RUNTIME_ACTIVE_STATUSES: + # runtime 说还在跑：清掉未决计时，继续等数据面写终态 + self._unresolved_since.pop(task_id, None) + self._last_log_at.pop(task_id, None) + return + + if runtime_status == RUNTIME_NOT_FOUND_STATUS: + # sidecar 重启后内存任务表为空，但 Ray Job 仍在跑并会自己写终态 + await self._handle_unresolved(task_id, "runtime reports NOT_FOUND") + return + + # 空响应 / 非预期内容（例如配置文件缺失导致的错误体），同样按未决处理 + await self._handle_unresolved(task_id, f"unexpected runtime status {runtime_status!r}") + + async def _task_already_finished(self, task_id: str) -> bool: + """数据面可能已经直接写了终态，此时不再轮询""" + from app.module.cleaning.schema import CleaningTaskStatus + + now = time.monotonic() + if now - self._last_db_check_at.get(task_id, 0.0) < DB_RECHECK_INTERVAL_SECONDS: + return False + self._last_db_check_at[task_id] = now + + db_status = await self._read_db_status(task_id) + if db_status is None: + logger.info(f"[Polling] Task {task_id} no longer exists, stop polling") + self._forget(task_id) + return True + + if db_status in { + CleaningTaskStatus.COMPLETED, + CleaningTaskStatus.PARTIAL_SUCCESS, + CleaningTaskStatus.FAILED, + CleaningTaskStatus.STOPPED, + }: + logger.info(f"[Polling] Task {task_id} already finished in DB: {db_status}") + self._forget(task_id) + return True + + return False + + async def _read_db_status(self, task_id: str) -> Optional[str]: + async with self._get_session_factory()() as db: + return await self.task_repo.find_status(db, task_id) + + async def _handle_unresolved(self, task_id: str, reason: str): + """runtime 给不出结论：保持 RUNNING，超时后才兜底置 FAILED""" + from app.module.cleaning.schema import CleaningTaskStatus + + now = time.monotonic() + since = self._unresolved_since.setdefault(task_id, now) + elapsed = now - since + timeout = _orphan_timeout_seconds() + + if elapsed >= timeout: + logger.error( + f"[Polling] Task {task_id} marked FAILED: unresolved for {int(elapsed)}s " + f"({reason}, timeout={int(timeout)}s)" + ) + await self._finish_task(task_id, CleaningTaskStatus.FAILED, reason) + self._forget(task_id) + return + + if now - self._last_log_at.get(task_id, 0.0) >= UNRESOLVED_LOG_INTERVAL_SECONDS: + self._last_log_at[task_id] = now + logger.warning( + f"[Polling] Task {task_id} still unresolved after {int(elapsed)}s ({reason}); " + f"keeping status unchanged, will fall back to FAILED after {int(timeout)}s" + ) + + async def _finish_task(self, task_id: str, status: str, reason: str = "") -> None: + """写入终态；已有终态时不做覆盖（幂等）""" + async with self._get_session_factory()() as db: + updated = await self.task_repo.finish_if_unfinished( + db, task_id, status, datetime.now() + ) + await db.commit() + + if updated: + logger.info(f"[Polling] Task {task_id} finished: {status} ({reason})") + else: + logger.info( + f"[Polling] Task {task_id} already has a terminal status, skip writing {status}" + ) diff --git a/runtime/datamate-python/app/module/cleaning/service/cleaning_task_service.py b/runtime/datamate-python/app/module/cleaning/service/cleaning_task_service.py index 76f9fc774..ae578628c 100644 --- a/runtime/datamate-python/app/module/cleaning/service/cleaning_task_service.py +++ b/runtime/datamate-python/app/module/cleaning/service/cleaning_task_service.py @@ -503,13 +503,21 @@ async def delete_task(self, db: AsyncSession, task_id: str) -> None: async def execute_task(self, db: AsyncSession, task_id: str) -> bool: """Execute task""" - succeeded = await self.result_repo.find_by_instance_id(db, task_id, "COMPLETED") - succeed_set = {res.src_file_id for res in succeeded} - task = await self.task_repo.find_task_by_id(db, task_id) if not task: raise BusinessError(ErrorCodes.CLEANING_TASK_NOT_FOUND, task_id) + # 运行中的任务不能重复提交：否则会起第二个 Ray Job，并产生重复的清洗结果。 + # 需要重跑时先停止任务（状态置 STOPPED），再执行。 + if task.status == CleaningTaskStatus.RUNNING: + raise BusinessError( + ErrorCodes.CLEANING_TASK_STATUS_INVALID, + "Task is running, cannot be executed again. Please stop the task first." + ) + + succeeded = await self.result_repo.find_by_instance_id(db, task_id, "COMPLETED") + succeed_set = {res.src_file_id for res in succeeded} + src_dataset = await self.dataset_service.get_dataset(task.src_dataset_id) if src_dataset: task.before_size = src_dataset.totalSize diff --git a/runtime/datamate-python/tests/test_module_cleaning_scheduler.py b/runtime/datamate-python/tests/test_module_cleaning_scheduler.py new file mode 100644 index 000000000..cfad27c46 --- /dev/null +++ b/runtime/datamate-python/tests/test_module_cleaning_scheduler.py @@ -0,0 +1,395 @@ +"""Polling behaviour of the cleaning task scheduler. + +The control plane must follow the runtime and the data plane instead of +declaring a task FAILED on the first hiccup: a single unreachable poll or a +NOT_FOUND after the runtime sidecar restarts used to mark long running tasks +as failed while the Ray Job kept running and later wrote COMPLETED. +""" + +import asyncio +import time +from datetime import datetime + +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.models.cleaning import CleaningTask +from app.module.cleaning.repository.cleaning_task_repository import ( + CleaningTaskRepository, + TERMINAL_TASK_STATUSES, +) +from app.module.cleaning.schema.cleaning import CleaningTaskStatus +from app.module.cleaning.service import cleaning_task_scheduler as scheduler_module +from app.module.cleaning.service.cleaning_task_scheduler import ( + DB_RECHECK_INTERVAL_SECONDS, + DEFAULT_ORPHAN_TIMEOUT_SECONDS, + CleaningTaskScheduler, + _orphan_timeout_seconds, +) + +TASK_ID = "task-1" + + +class FakeSession: + def __init__(self) -> None: + self.committed = False + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + return False + + async def commit(self) -> None: + self.committed = True + + +class FakeRepository: + def __init__(self, statuses: dict[str, str] | None = None, update_result: bool = True): + self.statuses = dict(statuses or {}) + self.update_result = update_result + self.finish_calls: list[tuple[str, str]] = [] + self.status_reads = 0 + + async def find_status(self, db, task_id: str): + self.status_reads += 1 + return self.statuses.get(task_id) + + async def finish_if_unfinished(self, db, task_id: str, status: str, finished_at=None) -> bool: + self.finish_calls.append((task_id, status)) + if not self.update_result: + return False + self.statuses[task_id] = status + return True + + +class FakeRuntimeClient: + def __init__(self, responses: list[dict | None]): + self.responses = list(responses) + self.calls = 0 + + async def get_task_status(self, task_id: str): + self.calls += 1 + if not self.responses: + return None + if len(self.responses) == 1: + return self.responses[0] + return self.responses.pop(0) + + +def _build(responses, statuses=None, update_result=True): + repo = FakeRepository(statuses, update_result) + client = FakeRuntimeClient(responses) + sessions: list[FakeSession] = [] + + def factory() -> FakeSession: + session = FakeSession() + sessions.append(session) + return session + + scheduler = CleaningTaskScheduler(repo, client, session_factory=factory) + scheduler._polling_task_ids.add(TASK_ID) + return scheduler, repo, client, sessions + + +def test_not_found_keeps_task_running() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "NOT_FOUND"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [] + assert TASK_ID in scheduler._polling_task_ids + assert TASK_ID in scheduler._unresolved_since + + +def test_unreachable_runtime_keeps_task_running() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [None], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [] + assert TASK_ID in scheduler._polling_task_ids + + +def test_unresolved_task_falls_back_to_failed_after_timeout() -> None: + async def run(): + scheduler, repo, _client, sessions = _build( + [{"status": "NOT_FOUND"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + scheduler._unresolved_since[TASK_ID] = time.monotonic() - ( + DEFAULT_ORPHAN_TIMEOUT_SECONDS + 1 + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo, sessions + + scheduler, repo, sessions = asyncio.run(run()) + + assert repo.finish_calls == [(TASK_ID, CleaningTaskStatus.FAILED)] + assert TASK_ID not in scheduler._polling_task_ids + assert sessions and sessions[-1].committed is True + + +def test_not_found_then_data_plane_completes_is_followed() -> None: + """Reported symptom: NOT_FOUND after a runtime restart, COMPLETED later. + + The control plane must keep the task RUNNING in between and must not write + FAILED, otherwise the UI flips back to COMPLETED once the data plane writes. + """ + + async def run(): + scheduler, repo, client, _sessions = _build( + [{"status": "NOT_FOUND"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + + for _ in range(3): + # 模拟每 10 秒回读一次数据库 + scheduler._last_db_check_at.pop(TASK_ID, None) + await scheduler._poll_task(TASK_ID) + + assert repo.finish_calls == [] + assert TASK_ID in scheduler._polling_task_ids + + # 数据面（Ray Job 执行器）直接写入终态 + repo.statuses[TASK_ID] = CleaningTaskStatus.COMPLETED + scheduler._last_db_check_at.pop(TASK_ID, None) + await scheduler._poll_task(TASK_ID) + return scheduler, repo, client + + scheduler, repo, client = asyncio.run(run()) + + assert repo.finish_calls == [] + assert client.calls == 3 + assert TASK_ID not in scheduler._polling_task_ids + + +def test_running_status_clears_unresolved_timer() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "running"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + scheduler._unresolved_since[TASK_ID] = time.monotonic() - 600 + scheduler._last_log_at[TASK_ID] = time.monotonic() + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [] + assert TASK_ID not in scheduler._unresolved_since + assert TASK_ID in scheduler._polling_task_ids + + +def test_terminal_db_status_stops_polling_without_calling_runtime() -> None: + async def run(): + scheduler, repo, client, _sessions = _build( + [{"status": "running"}], {TASK_ID: CleaningTaskStatus.COMPLETED} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo, client + + scheduler, repo, client = asyncio.run(run()) + + assert client.calls == 0 + assert repo.finish_calls == [] + assert TASK_ID not in scheduler._polling_task_ids + + +def test_data_plane_partial_success_stops_polling() -> None: + async def run(): + scheduler, _repo, client, _sessions = _build( + [{"status": "running"}], {TASK_ID: CleaningTaskStatus.PARTIAL_SUCCESS} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, client + + scheduler, client = asyncio.run(run()) + + assert client.calls == 0 + assert TASK_ID not in scheduler._polling_task_ids + + +def test_deleted_task_stops_polling() -> None: + async def run(): + scheduler, _repo, client, _sessions = _build([{"status": "running"}], {}) + await scheduler._poll_task(TASK_ID) + return scheduler, client + + scheduler, client = asyncio.run(run()) + + assert client.calls == 0 + assert TASK_ID not in scheduler._polling_task_ids + + +def test_completed_status_is_written() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "COMPLETED"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [(TASK_ID, CleaningTaskStatus.COMPLETED)] + assert TASK_ID not in scheduler._polling_task_ids + + +def test_failed_status_is_written() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "failed", "error": "boom"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [(TASK_ID, CleaningTaskStatus.FAILED)] + assert TASK_ID not in scheduler._polling_task_ids + + +def test_cancelled_status_maps_to_stopped() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "cancelled"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [(TASK_ID, CleaningTaskStatus.STOPPED)] + + +def test_existing_terminal_status_is_not_overwritten() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "failed"}], {TASK_ID: CleaningTaskStatus.RUNNING}, update_result=False + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [(TASK_ID, CleaningTaskStatus.FAILED)] + assert TASK_ID not in scheduler._polling_task_ids + + +def test_unexpected_payload_keeps_task_running() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"code": "runtime.0009", "message": "config missing"}], + {TASK_ID: CleaningTaskStatus.RUNNING}, + ) + await scheduler._poll_task(TASK_ID) + return scheduler, repo + + scheduler, repo = asyncio.run(run()) + + assert repo.finish_calls == [] + assert TASK_ID in scheduler._polling_task_ids + + +def test_db_recheck_is_throttled() -> None: + async def run(): + scheduler, repo, _client, _sessions = _build( + [{"status": "running"}], {TASK_ID: CleaningTaskStatus.RUNNING} + ) + await scheduler._poll_task(TASK_ID) + await scheduler._poll_task(TASK_ID) + return repo + + repo = asyncio.run(run()) + + assert DB_RECHECK_INTERVAL_SECONDS > 0 + assert repo.status_reads == 1 + + +def test_orphan_timeout_defaults(monkeypatch) -> None: + monkeypatch.delenv("CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS", raising=False) + assert _orphan_timeout_seconds() == DEFAULT_ORPHAN_TIMEOUT_SECONDS + + monkeypatch.setenv("CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS", "600") + assert _orphan_timeout_seconds() == 600.0 + + +def test_orphan_timeout_ignores_invalid_values(monkeypatch) -> None: + monkeypatch.setenv("CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS", "not-a-number") + assert _orphan_timeout_seconds() == DEFAULT_ORPHAN_TIMEOUT_SECONDS + + monkeypatch.setenv("CLEANING_POLL_ORPHAN_TIMEOUT_SECONDS", "-5") + assert _orphan_timeout_seconds() == DEFAULT_ORPHAN_TIMEOUT_SECONDS + + +def test_terminal_statuses_cover_data_plane_results() -> None: + assert set(TERMINAL_TASK_STATUSES) == { + CleaningTaskStatus.COMPLETED, + CleaningTaskStatus.PARTIAL_SUCCESS, + CleaningTaskStatus.FAILED, + CleaningTaskStatus.STOPPED, + } + + +def test_scheduler_module_exposes_singleton() -> None: + scheduler_module._scheduler = None + first = scheduler_module.get_scheduler() + second = scheduler_module.get_scheduler() + + assert first is second + + scheduler_module._scheduler = None + + +def test_finish_if_unfinished_keeps_first_terminal_status() -> None: + """The guarded update must not overwrite a status the executor already wrote.""" + + async def run(): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(lambda sync_conn: CleaningTask.__table__.create(sync_conn)) + + factory = async_sessionmaker(engine, expire_on_commit=False) + repo = CleaningTaskRepository() + + async with factory() as db: + db.add( + CleaningTask( + id=TASK_ID, + name="demo", + status=CleaningTaskStatus.RUNNING, + src_dataset_id="dataset-1", + src_dataset_name="dataset", + retry_count=0, + ) + ) + await db.commit() + + assert await repo.find_status(db, TASK_ID) == CleaningTaskStatus.RUNNING + assert await repo.find_status(db, "missing") is None + + written = await repo.finish_if_unfinished( + db, TASK_ID, CleaningTaskStatus.COMPLETED, datetime.now() + ) + await db.commit() + assert written is True + assert await repo.find_status(db, TASK_ID) == CleaningTaskStatus.COMPLETED + + written_again = await repo.finish_if_unfinished( + db, TASK_ID, CleaningTaskStatus.FAILED, datetime.now() + ) + await db.commit() + assert written_again is False + assert await repo.find_status(db, TASK_ID) == CleaningTaskStatus.COMPLETED + + await engine.dispose() + + asyncio.run(run()) diff --git a/runtime/datamate-python/tests/test_module_cleaning_service.py b/runtime/datamate-python/tests/test_module_cleaning_service.py new file mode 100644 index 000000000..d4fedd91f --- /dev/null +++ b/runtime/datamate-python/tests/test_module_cleaning_service.py @@ -0,0 +1,66 @@ +"""Guards around re-executing a cleaning task.""" + +import asyncio + +import pytest + +from app.core.exception import BusinessError, ErrorCodes +from app.module.cleaning.schema.cleaning import CleaningTaskDto, CleaningTaskStatus +from app.module.cleaning.service.cleaning_task_service import CleaningTaskService + +TASK_ID = "11111111-1111-1111-1111-111111111111" + + +class FakeTaskRepository: + def __init__(self, status: str) -> None: + self.status = status + + async def find_task_by_id(self, db, task_id: str): + return CleaningTaskDto( + id=task_id, + name="demo", + status=self.status, + src_dataset_id="dataset-1", + retry_count=0, + ) + + +class UnexpectedCallRepository: + async def find_by_instance_id(self, db, task_id: str, status: str): + raise AssertionError("result repository must not be touched for a running task") + + +class UnexpectedScheduler: + async def execute_task(self, db, task_id: str, retry_count: int) -> bool: + raise AssertionError("scheduler must not resubmit a running task") + + +def _service(status: str) -> CleaningTaskService: + return CleaningTaskService( + task_repo=FakeTaskRepository(status), + result_repo=UnexpectedCallRepository(), + operator_instance_repo=None, + operator_service=None, + scheduler=UnexpectedScheduler(), + validator=None, + dataset_service=None, + lineage_service=None, + ) + + +def test_execute_task_rejects_running_task() -> None: + with pytest.raises(BusinessError) as exc_info: + asyncio.run(_service(CleaningTaskStatus.RUNNING).execute_task(None, TASK_ID)) + + assert exc_info.value.error_code is ErrorCodes.CLEANING_TASK_STATUS_INVALID + assert "running" in str(exc_info.value.data).lower() + + +def test_execute_task_rejects_pending_resubmission_of_running_task_only() -> None: + """A finished task may be re-executed; only RUNNING is blocked.""" + service = _service(CleaningTaskStatus.COMPLETED) + + with pytest.raises(AssertionError): + # The guard must not fire for COMPLETED: the flow continues into the + # result repository, which is faked to fail loudly. + asyncio.run(service.execute_task(None, TASK_ID))