add worker tasks discover_workflows_from_modules and discover_all_workflows

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add ewoksjob options under `EWOKSJOB_OPTIONS`: ``log_memory_usage`` and ``detect_memory_leaks``.
 - Added ability to handle remote ewoksutils exception types on the client side.
 - Add `client.task_utils.TaskSubmitter` helper to execute a single task.
+- Added worker tasks `discover_workflows_from_modules` and `discover_all_workflows`.
 
 ### Fixed
 

pyproject.toml

@@ -30,7 +30,7 @@ Changelog = "https://github.com/ewoks-kit/ewoksjob/-/blob/main/CHANGELOG.md"
 
 [project.optional-dependencies]
 worker = [
-    "ewoks >=4.0.0",
+    "ewoks >=6.0.0",
     "objgraph",
 ]
 redis = [
@@ -95,6 +95,9 @@ package-dir = { "" = "src" }
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[tool.setuptools.package-data]
+"*" = ["*.json"]
+
 [tool.coverage.run]
 omit = ['*/tests/*']
 

src/ewoksjob/apps/_decorators/ewoks_tasks.py

@@ -5,6 +5,7 @@
 
 import ewoks
 from ewokscore import task_discovery
+from ewokscore import workflow_discovery
 
 from ...worker.executor import get_execute_method
 from ..errors import replace_exception_for_client
@@ -49,6 +50,8 @@ def new_celery_task(*args, **kwargs) -> Any:
     "convert_graph": ewoks.convert_graph,
     "discover_tasks_from_modules": task_discovery.discover_tasks_from_modules,
     "discover_all_tasks": task_discovery.discover_all_tasks,
+    "discover_workflows_from_modules": workflow_discovery.discover_workflows_from_modules,
+    "discover_all_workflows": workflow_discovery.discover_all_workflows,
 }
 
 _BOUND_TASKS = {"execute_graph"}

src/ewoksjob/apps/ewoks.py

@@ -44,3 +44,17 @@ def discover_tasks_from_modules(*args, **kwargs) -> List[dict]:
 @ewoks_tasks.ewoks_task
 def discover_all_tasks(*args, **kwargs) -> List[dict]:
     return ewoks_tasks.get_native_ewoks_task("discover_all_tasks")(*args, **kwargs)
+
+
+@app.task(bind=False)
+@ewoks_tasks.ewoks_task
+def discover_workflows_from_modules(*args, **kwargs) -> List[dict]:
+    return ewoks_tasks.get_native_ewoks_task("discover_workflows_from_modules")(
+        *args, **kwargs
+    )
+
+
+@app.task(bind=False)
+@ewoks_tasks.ewoks_task
+def discover_all_workflows(*args, **kwargs) -> List[dict]:
+    return ewoks_tasks.get_native_ewoks_task("discover_all_workflows")(*args, **kwargs)

src/ewoksjob/client/celery/tasks.py

@@ -12,6 +12,8 @@
     "convert_workflow",
     "discover_tasks_from_modules",
     "discover_all_tasks",
+    "discover_workflows_from_modules",
+    "discover_all_workflows",
 ]
 
 
@@ -55,3 +57,15 @@ def discover_tasks_from_modules(**kw) -> CeleryFuture:
 def discover_all_tasks(**kw) -> CeleryFuture:
     async_result = send_task("ewoksjob.apps.ewoks.discover_all_tasks", **kw)
     return CeleryFuture(async_result.id, async_result)
+
+
+def discover_workflows_from_modules(**kw) -> CeleryFuture:
+    async_result = send_task(
+        "ewoksjob.apps.ewoks.discover_workflows_from_modules", **kw
+    )
+    return CeleryFuture(async_result.id, async_result)
+
+
+def discover_all_workflows(**kw) -> CeleryFuture:
+    async_result = send_task("ewoksjob.apps.ewoks.discover_all_workflows", **kw)
+    return CeleryFuture(async_result.id, async_result)

src/ewoksjob/client/local/tasks.py

@@ -6,6 +6,7 @@
 
 import ewoks
 from ewokscore import task_discovery
+from ewokscore import workflow_discovery
 
 from ..dummy_workflow import dummy_workflow
 from .futures import LocalFuture
@@ -18,6 +19,8 @@
     "convert_workflow",
     "discover_tasks_from_modules",
     "discover_all_tasks",
+    "discover_workflows_from_modules",
+    "discover_all_workflows",
 ]
 
 
@@ -74,3 +77,21 @@ def discover_all_tasks(
 ) -> LocalFuture:
     pool = get_active_pool()
     return pool.submit(task_discovery.discover_all_tasks, args=args, kwargs=kwargs)
+
+
+def discover_workflows_from_modules(
+    args: Optional[Tuple] = tuple(), kwargs: Optional[Mapping] = None
+) -> LocalFuture:
+    pool = get_active_pool()
+    return pool.submit(
+        workflow_discovery.discover_workflows_from_modules, args=args, kwargs=kwargs
+    )
+
+
+def discover_all_workflows(
+    args: Optional[Tuple] = tuple(), kwargs: Optional[Mapping] = None
+) -> LocalFuture:
+    pool = get_active_pool()
+    return pool.submit(
+        workflow_discovery.discover_all_workflows, args=args, kwargs=kwargs
+    )

src/ewoksjob/tests/_loadtest/graph.json

@@ -0,0 +1,41 @@
+{
+  "graph": {
+    "id": "graph",
+    "schema_version": "1.1"
+  },
+  "nodes": [
+    {
+      "id": "node1",
+      "task_type": "method",
+      "task_identifier": "dummy",
+      "default_inputs": [
+        {
+          "name": "name",
+          "value": "node1"
+        },
+        {
+          "name": "value",
+          "value": 0
+        }
+      ]
+    },
+    {
+      "id": "node2",
+      "task_type": "graph",
+      "task_identifier": "subgraph"
+    }
+  ],
+  "links": [
+    {
+      "source": "node1",
+      "target": "node2",
+      "sub_target": "in",
+      "data_mapping": [
+        {
+          "target_input": "value",
+          "source_output": "return_value"
+        }
+      ]
+    }
+  ]
+}

src/ewoksjob/tests/_loadtest/subgraph.json

@@ -0,0 +1,29 @@
+{
+  "graph": {
+    "id": "subgraph",
+    "schema_version": "1.1",
+    "input_nodes": [
+      {
+        "id": "in",
+        "node": "subnode1"
+      }
+    ]
+  },
+  "nodes": [
+    {
+      "id": "subnode1",
+      "task_type": "method",
+      "task_identifier": "dummy",
+      "default_inputs": [
+        {
+          "name": "name",
+          "value": "subnode1"
+        },
+        {
+          "name": "value",
+          "value": 0
+        }
+      ]
+    }
+  ]
+}

src/ewoksjob/tests/conftest.py

@@ -3,6 +3,7 @@
 
 import pytest
 from ewokscore import events
+from ewokscore import workflow_discovery
 
 from ewoksjob.events.readers import read_ewoks_events
 from ewoksjob.worker import options as worker_options
@@ -130,6 +131,24 @@ def local_ewoks_worker(slurm_client_kwargs):
             assert len(pool._tasks) == 0, str(list(pool._tasks.values()))
 
 
+@pytest.fixture
+def local_patched_ewoks_worker(monkeypatch):
+    monkeypatch.setattr(workflow_discovery, "entry_points", _mock_entry_points)
+
+    with local.pool_context(pool_type="thread"):
+        yield
+
+
+class _MockEntryPoint:
+    def __init__(self, name):
+        self.name = name
+
+
+def _mock_entry_points(group):
+    assert group == "ewoks.workflows"
+    return [_MockEntryPoint("ewoksjob.tests._loadtest.*")]
+
+
 @pytest.fixture()
 def sqlite3_ewoks_events(tmp_path):
     uri = f"file:{tmp_path / 'ewoks_events.db'}"

src/ewoksjob/tests/test_workflow_discovery.py

@@ -0,0 +1,44 @@
+from ..client import celery
+from ..client import local
+
+
+def test_submit(ewoks_worker):
+    expected = [
+        "ewoksjob.tests._loadtest.graph",
+        "ewoksjob.tests._loadtest.subgraph",
+    ]
+    assert_submit(celery, [])
+    assert_submit(celery, expected, "ewoksjob.tests._loadt*.*")
+
+
+def test_submit_local(local_ewoks_worker):
+    expected = [
+        "ewoksjob.tests._loadtest.graph",
+        "ewoksjob.tests._loadtest.subgraph",
+    ]
+    assert_submit(local, [])
+    assert_submit(local, expected, "ewoksjob.tests._loadt*.*")
+
+
+def test_submit_local_patched(local_patched_ewoks_worker):
+    expected = [
+        "ewoksjob.tests._loadtest.graph",
+        "ewoksjob.tests._loadtest.subgraph",
+    ]
+
+    assert_submit(local, expected)
+    assert_submit(local, expected, "ewoksjob.tests._loadt*.*")
+
+
+def assert_submit(mod, expected, *modules):
+    if modules:
+        future = mod.discover_workflows_from_modules(args=modules)
+    else:
+        future = mod.discover_all_workflows()
+
+    results = future.result(timeout=60)
+
+    if modules:
+        assert set(results) == set(expected)
+    else:
+        assert set(expected).issubset(results), results