From 2d8bf69b4e00e3cfa29fade50b70219674619384 Mon Sep 17 00:00:00 2001 From: Kexin Chen <22360275@zju.edu.cn> Date: Fri, 23 Oct 2026 22:00:53 +0800 Subject: [PATCH] fix(workflow): include isolated tasks in DAG results --- ms_agent/workflow/dag_workflow.py | 15 ++-- tests/workflow/test_dag_workflow.py | 107 ++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 6 deletions(-) create mode 100644 tests/workflow/test_dag_workflow.py diff --git a/ms_agent/workflow/dag_workflow.py b/ms_agent/workflow/dag_workflow.py index 14ea1d671..ad0101e14 100644 --- a/ms_agent/workflow/dag_workflow.py +++ b/ms_agent/workflow/dag_workflow.py @@ -23,10 +23,15 @@ def build_workflow(self): # Build adjacency list graph and in-degree map self.graph: Dict[str, List[str]] = defaultdict(list) indegree: Dict[str, int] = defaultdict(int) - tasks: Set[str] = set(self.config.keys()) + tasks: Set[str] = { + name + for name, task_config in self.config.items() + if isinstance(task_config, DictConfig) + and 'agent_config' in task_config + } for task_name, task_config in self.config.items(): - if 'next' in task_config: + if task_name in tasks and 'next' in task_config: next_tasks = task_config['next'] if isinstance(next_tasks, str): next_tasks = [next_tasks] @@ -34,12 +39,10 @@ def build_workflow(self): self.graph[task_name].append(nxt) indegree[nxt] += 1 - self.nodes = set(list(self.graph.keys()) + list(indegree.keys())) + self.nodes = tasks # Find root tasks (indegree==0) - self.roots = [ - t for t in tasks if 'next' in self.config[t] and indegree[t] == 0 - ] + self.roots = [t for t in tasks if indegree[t] == 0] if not self.roots: raise ValueError('No root task found for DagWorkflow') diff --git a/tests/workflow/test_dag_workflow.py b/tests/workflow/test_dag_workflow.py new file mode 100644 index 000000000..603e2bd01 --- /dev/null +++ b/tests/workflow/test_dag_workflow.py @@ -0,0 +1,107 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise DAG scheduling through local workflow and built-in agent loading.""" +import tempfile +import unittest +from pathlib import Path + +from omegaconf import OmegaConf + +from ms_agent.workflow import DagWorkflow, WorkflowLoader + + +class TestDagWorkflow(unittest.IsolatedAsyncioTestCase): + + def setUp(self): + self.directory = tempfile.TemporaryDirectory(prefix='ms-agent-dag-') + self.addCleanup(self.directory.cleanup) + self.folder = Path(self.directory.name) + self.value = {'value': 23} + OmegaConf.save( + OmegaConf.create({ + 'type': 'CodeAgent', + 'save_history': False, + 'output_dir': str(self.folder / 'output'), + }), self.folder / 'identity.yaml') + + def build(self, tasks): + config = { + 'type': 'DagWorkflow', + **{ + name: dict(entry, agent_config='identity.yaml') + for name, entry in tasks.items() + }, + } + path = self.folder / 'workflow.yaml' + OmegaConf.save(OmegaConf.create(config), path) + return WorkflowLoader.build( + config_dir_or_id=str(path), + trust_remote_code=False, + load_cache=False) + + async def check_workflow(self, tasks, expected): + workflow = self.build(tasks) + self.assertIsInstance(workflow, DagWorkflow) + actual = await workflow.run(dict(self.value)) + self.assertEqual(actual, expected) + self.assertEqual(set(workflow.topo_order), set(tasks)) + # Config.from_task also adds local_dir, name, tools and callbacks. + # Those entries describe the workflow, not additional agent tasks. + self.assertEqual(workflow.nodes, set(tasks)) + + async def test_single_task(self): + await self.check_workflow({'only': {}}, {'only': self.value}) + + async def test_single_task_with_empty_successors(self): + await self.check_workflow({'only': {'next': []}}, + {'only': self.value}) + + async def test_independent_tasks(self): + await self.check_workflow({'a': {}, 'b': {}}, + {'a': self.value, 'b': self.value}) + + async def test_independent_tasks_with_empty_successors(self): + await self.check_workflow({'a': {'next': []}, 'b': {'next': []}}, + {'a': self.value, 'b': self.value}) + + async def test_connected_and_isolated_tasks(self): + await self.check_workflow( + {'solo': {}, 'a': {'next': 'b'}, 'b': {}}, + {'solo': self.value, 'b': self.value}) + + async def test_connected_and_isolated_task_with_empty_successors(self): + await self.check_workflow( + {'a': {'next': 'b'}, 'b': {}, 'solo': {'next': []}}, + {'b': self.value, 'solo': self.value}) + + async def test_chain(self): + await self.check_workflow({'a': {'next': 'b'}, 'b': {}}, + {'b': self.value}) + + async def test_fork(self): + await self.check_workflow( + {'a': {'next': ['b', 'c']}, 'b': {}, 'c': {}}, + {'b': self.value, 'c': self.value}) + + async def test_join(self): + await self.check_workflow( + {'a': {'next': 'c'}, 'b': {'next': 'c'}, 'c': {}}, + {'c': [self.value, self.value]}) + + async def test_diamond(self): + await self.check_workflow({ + 'a': {'next': ['b', 'c']}, + 'b': {'next': 'd'}, + 'c': {'next': 'd'}, + 'd': {}, + }, {'d': [self.value, self.value]}) + + async def test_cycles_without_roots_are_rejected(self): + for tasks in ({'a': {'next': 'a'}}, + {'a': {'next': 'b'}, 'b': {'next': 'a'}}): + with self.subTest(tasks=tasks): + with self.assertRaisesRegex(ValueError, 'No root task found'): + self.build(tasks) + + +if __name__ == '__main__': + unittest.main()