From d36b59bece781b2c4d11b20ea0d9a47630ef2601 Mon Sep 17 00:00:00 2001
From: Alputer <alp.tuna.453@gmail.com>
Date: Tue, 5 Nov 2024 16:37:12 +0100
Subject: [PATCH] feat(config): set maximum limit for number of Dask workers
 (#712)

---
 docs/openapi.json          | 15 +++++++++++++++
 reana_server/config.py     |  5 +++++
 reana_server/rest/info.py  | 17 +++++++++++++++++
 reana_server/validation.py |  7 +++++++
 4 files changed, 44 insertions(+)

diff --git a/docs/openapi.json b/docs/openapi.json
index d8c95fc7..fd83bc8f 100644
--- 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
index 79a01cff..7a74b9a0 100644
--- 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
index a3660a60..b7ff1fab 100644
--- 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
index 4e85830c..ed9ef2d1 100644
--- 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
         )