feat(config): set maximum limit for number of Dask workers (#712)

reanahub · Nov 11, 2024 · d36b59b · d36b59b
1 parent 95ea2df
commit d36b59b
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 0 deletions.
diff --git a/docs/openapi.json b/docs/openapi.json
@@ -445,6 +445,10 @@
                   "title": "The maximum memory limit for Dask clusters created by users",
                   "value": "16Gi"
                 },
+                "dask_cluster_max_number_of_workers": {
+                  "title": "The maximum number of workers that users can ask for the single Dask cluster",
+                  "value": "20"
+                },
                 "dask_cluster_max_single_worker_memory": {
                   "title": "The maximum amount of memory that users can ask for the single Dask worker",
                   "value": "8Gi"
@@ -547,6 +551,17 @@
                   },
                   "type": "object"
                 },
+                "dask_cluster_max_number_of_workers": {
+                  "properties": {
+                    "title": {
+                      "type": "string"
+                    },
+                    "value": {
+                      "type": "string"
+                    }
+                  },
+                  "type": "object"
+                },
                 "dask_cluster_max_single_worker_memory": {
                   "properties": {
                     "title": {

diff --git a/reana_server/config.py b/reana_server/config.py
@@ -74,6 +74,11 @@
 )
 """Number of workers in Dask cluster by default """
 
+REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS = int(
+    os.getenv("REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS", 20)
+)
+"""Maximum number of workers in Dask cluster."""
+
 REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY = os.getenv(
     "REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY", "2Gi"
 )

diff --git a/reana_server/rest/info.py b/reana_server/rest/info.py
@@ -30,6 +30,7 @@
     REANA_DASK_CLUSTER_MAX_MEMORY_LIMIT,
     REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY,
     REANA_DASK_CLUSTER_MAX_SINGLE_WORKER_MEMORY,
+    REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS,
 )
 from reana_server.decorators import signin_required
 
@@ -173,6 +174,13 @@ def info(user, **kwargs):  # noqa
                   value:
                     type: string
                 type: object
+              dask_cluster_max_number_of_workers:
+                properties:
+                  title:
+                    type: string
+                  value:
+                    type: string
+                type: object
             type: object
           examples:
             application/json:
@@ -237,6 +245,10 @@ def info(user, **kwargs):  # noqa
                     "title": "The maximum amount of memory that users can ask for the single Dask worker",
                     "value": "8Gi"
                 },
+                "dask_cluster_max_number_of_workers": {
+                    "title": "The maximum number of workers that users can ask for the single Dask cluster",
+                    "value": "20"
+                },
               }
         500:
           description: >-
@@ -315,6 +327,10 @@ def info(user, **kwargs):  # noqa
                 title="The maximum amount of memory that users can ask for the single Dask worker",
                 value=REANA_DASK_CLUSTER_MAX_SINGLE_WORKER_MEMORY,
             )
+            cluster_information["dask_cluster_max_number_of_workers"] = dict(
+                title="The maximum number of workers that users can ask for the single Dask cluster",
+                value=REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS,
+            )
 
         return InfoSchema().dump(cluster_information)
 
@@ -366,3 +382,4 @@ class InfoSchema(Schema):
         dask_cluster_max_memory_limit = fields.Nested(StringInfoValue)
         dask_cluster_default_single_worker_memory = fields.Nested(StringInfoValue)
         dask_cluster_max_single_worker_memory = fields.Nested(StringInfoValue)
+        dask_cluster_max_number_of_workers = fields.Nested(StringInfoValue)
diff --git a/reana_server/validation.py b/reana_server/validation.py
@@ -26,6 +26,7 @@
     DASK_ENABLED,
     REANA_DASK_CLUSTER_MAX_MEMORY_LIMIT,
     REANA_DASK_CLUSTER_DEFAULT_NUMBER_OF_WORKERS,
+    REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS,
     REANA_DASK_CLUSTER_DEFAULT_SINGLE_WORKER_MEMORY,
     REANA_DASK_CLUSTER_MAX_SINGLE_WORKER_MEMORY,
 )
@@ -188,6 +189,12 @@ def validate_dask_memory_and_cores_limits(reana_yaml: Dict) -> None:
                 "number_of_workers", REANA_DASK_CLUSTER_DEFAULT_NUMBER_OF_WORKERS
             )
         )
+
+        if number_of_workers > REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS:
+            raise REANAValidationError(
+                f"The number of requested Dask workers ({number_of_workers}) exceeds the maximum limit ({REANA_DASK_CLUSTER_MAX_NUMBER_OF_WORKERS})."
+            )
+
         requested_dask_cluster_memory = (
             kubernetes_memory_to_bytes(single_worker_memory) * number_of_workers
         )