feat(dask): add nthreads option to dask-worker command (#636)

Following the integration of Dask into REANA and discussions with
several REANA users, we identified the need to make the number of
threads configurable.

In this commit, we propagate the specified number of threads per
worker or the default value to the `dask-worker` command.

Closes reanahub/reana#874

Author: Alputer

reana_workflow_controller/config.py

@@ -346,6 +346,11 @@ def _parse_interactive_sessions_environments(env_var):
 )
 """Maximum memory for one Dask worker."""
 
+REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_THREADS = int(
+    os.getenv("REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_THREADS", 4)
+)
+"""Number of threads for one Dask worker by default."""
+
 VOMSPROXY_CONTAINER_IMAGE = os.getenv(
     "VOMSPROXY_CONTAINER_IMAGE", "docker.io/reanahub/reana-auth-vomsproxy:1.3.0"
 )

reana_workflow_controller/dask.py

@@ -58,6 +58,7 @@ def __init__(
         workflow_workspace,
         user_id,
         num_of_workers,
+        num_of_threads,
         single_worker_memory,
         kerberos=False,
         voms_proxy=False,
@@ -77,6 +78,7 @@ def __init__(
         self.cluster_name = get_dask_component_name(workflow_id, "cluster")
         self.num_of_workers = num_of_workers
         self.single_worker_memory = single_worker_memory
+        self.num_of_threads = num_of_threads
         self.workflow_spec = workflow_spec
         self.workflow_workspace = workflow_workspace
         self.workflow_id = str(workflow_id)
@@ -179,9 +181,9 @@ def _prepare_cluster(self):
         ] = self.cluster_image
 
         # Create the worker command
-        self.cluster_body["spec"]["worker"]["spec"]["containers"][0]["args"][
-            0
-        ] = f'cd {self.workflow_workspace} && {self.cluster_body["spec"]["worker"]["spec"]["containers"][0]["args"][0]}'
+        self.cluster_body["spec"]["worker"]["spec"]["containers"][0]["args"] = [
+            f"cd {self.workflow_workspace} && exec dask-worker --name $(DASK_WORKER_NAME) --dashboard --dashboard-address 8788 --nthreads {self.num_of_threads} --memory-limit {self.single_worker_memory}"
+        ]
 
         # Set resource limits for workers
         self.cluster_body["spec"]["worker"]["spec"]["containers"][0]["resources"] = {

reana_workflow_controller/templates/dask_cluster.yaml

@@ -8,8 +8,6 @@ spec:
         - name: worker
           imagePullPolicy: "IfNotPresent"
           command: ["/bin/sh", "-c"]
-          args:
-            - exec dask-worker --name $(DASK_WORKER_NAME) --dashboard --dashboard-address 8788
           ports:
             - name: http-dashboard
               containerPort: 8788

reana_workflow_controller/workflow_run_manager.py

@@ -91,6 +91,7 @@
     WORKFLOW_ENGINE_YADAGE_ENV_VARS,
     REANA_DASK_CLUSTER_DEFAULT_NUMBER_OF_WORKERS,
     REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY,
+    REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_THREADS,
 )
 
 
@@ -415,6 +416,13 @@ def start_batch_workflow_run(
                         "single_worker_memory",
                         REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY,
                     ),
+                    num_of_threads=self.workflow.reana_specification["workflow"]
+                    .get("resources", {})
+                    .get("dask", {})
+                    .get(
+                        "single_worker_threads",
+                        REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_THREADS,
+                    ),
                     kerberos=self.requires_kerberos(),
                     voms_proxy=self.requires_voms_proxy(),
                     rucio=self.requires_rucio(),

setup.py

@@ -51,7 +51,7 @@
     "marshmallow>2.13.0,<3.0.0",  # same upper pin as reana-server
     "opensearch-py>=2.7.0,<2.8.0",
     "packaging>=18.0",
-    "reana-commons[kubernetes]>=0.95.0a8,<0.96.0",
+    "reana-commons[kubernetes]>=0.95.0a9,<0.96.0",
     "reana-db>=0.95.0a5,<0.96.0",
     "requests>=2.25.0",
     "sqlalchemy-utils>=0.31.0",

tests/conftest.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 # This file is part of REANA.
-# Copyright (C) 2017, 2018, 2019, 2020, 2021, 2022, 2024 CERN.
+# Copyright (C) 2017, 2018, 2019, 2020, 2021, 2022, 2024, 2025 CERN.
 #
 # REANA is free software; you can redistribute it and/or modify it
 # under the terms of the MIT License; see LICENSE file for more details.
@@ -209,6 +209,7 @@ def dask_resource_manager(sample_serial_workflow_in_db_with_dask, mock_user_secr
         workflow_workspace="/path/to/workspace",
         user_id="user-123",
         num_of_workers=2,
+        num_of_threads=8,
         single_worker_memory="256Mi",
     )
     return manager

tests/test_dask.py

@@ -1,5 +1,5 @@
 # This file is part of REANA.
-# Copyright (C) 2024 CERN.
+# Copyright (C) 2024, 2025 CERN.
 #
 # REANA is free software; you can redistribute it and/or modify it
 # under the terms of the MIT License; see LICENSE file for more details.
@@ -425,9 +425,7 @@ def test_prepare_cluster(dask_resource_manager):
             "env"
         ]
 
-        expected_command = (
-            f"cd {dask_resource_manager.workflow_workspace} && worker-command"
-        )
+        expected_command = f"cd {dask_resource_manager.workflow_workspace} && exec dask-worker --name $(DASK_WORKER_NAME) --dashboard --dashboard-address 8788 --nthreads 8 --memory-limit 256Mi"
         assert (
             dask_resource_manager.cluster_body["spec"]["worker"]["spec"]["containers"][
                 0