feat(dask): add DASK_SCHEDULER_URI env variable (#464)

reanahub · Oct 24, 2024 · d357643 · d357643
1 parent db9c258
commit d357643
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/AUTHORS.md b/AUTHORS.md
@@ -4,6 +4,7 @@ The list of contributors in alphabetical order:
 
 - [Adelina Lintuluoto](https://orcid.org/0000-0002-0726-1452)
 - [Agisilaos Kounelis](https://orcid.org/0000-0001-9312-3189)
+- [Alp Tuna](https://orcid.org/0009-0001-1915-3993)
 - [Anton Khodak](https://orcid.org/0000-0003-3263-4553)
 - [Audrius Mecionis](https://orcid.org/0000-0002-3759-1663)
 - [Chris Hollowell](https://orcid.org/0000-0003-1239-111X)

diff --git a/reana_job_controller/config.py b/reana_job_controller/config.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 # This file is part of REANA.
-# Copyright (C) 2017, 2018, 2019, 2020, 2021, 2022, 2023 CERN.
+# Copyright (C) 2017, 2018, 2019, 2020, 2021, 2022, 2023, 2024 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.
@@ -23,6 +23,9 @@
 CACHE_ENABLED = False
 """Determines if jobs caching is enabled."""
 
+DASK_SCHEDULER_URI = os.getenv("DASK_SCHEDULER_URI", "tcp://127.0.0.1:8080")
+"""Address of the Dask Scheduler."""
+
 SHARED_VOLUME_PATH_ROOT = os.getenv("SHARED_VOLUME_PATH_ROOT", "/var/reana")
 """Root path of the shared volume ."""
 

diff --git a/reana_job_controller/job_manager.py b/reana_job_controller/job_manager.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 # This file is part of REANA.
-# Copyright (C) 2019, 2020, 2021 CERN.
+# Copyright (C) 2019, 2020, 2021, 2024 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.
@@ -14,7 +14,7 @@
 from reana_db.database import Session
 from reana_db.models import Job as JobTable, JobCache, JobStatus, Workflow
 
-from reana_job_controller.config import CACHE_ENABLED
+from reana_job_controller.config import CACHE_ENABLED, DASK_SCHEDULER_URI
 
 
 class JobManager:
@@ -149,4 +149,5 @@ def _extend_env_vars(self, env_vars):
         prefix = "REANA"
         env_vars[prefix + "_WORKSPACE"] = self.workflow_workspace
         env_vars[prefix + "_WORKFLOW_UUID"] = str(self.workflow_uuid)
+        env_vars["DASK_SCHEDULER_URI"] = DASK_SCHEDULER_URI
         return env_vars
diff --git a/tests/test_job_manager.py b/tests/test_job_manager.py
@@ -81,15 +81,15 @@ def test_execute_kubernetes_job(
                 assert len(containers) == 2  # main job + sidecar
                 assert len(init_containers) == 1
                 assert init_containers[0]["name"] == KRB5_INIT_CONTAINER_NAME
-                assert len(env_vars) == 6  # KRB5CCNAME is added
+                assert len(env_vars) == 7  # KRB5CCNAME is added
                 assert "trap" in command[0] and expected_command in command[0]
                 assert "kinit -R" in containers[1]["args"][0]
                 assert containers[1]["name"] == KRB5_RENEW_CONTAINER_NAME
             else:
                 assert len(containers) == 1
                 assert len(init_containers) == 0
-                # custom env + REANA_WORKSPACE + REANA_WORKFLOW_UUID + two secrets
-                assert len(env_vars) == 5
+                # custom env + REANA_WORKSPACE + REANA_WORKFLOW_UUID + DASK_SCHEDULER_URI + two secrets
+                assert len(env_vars) == 6
                 assert command == [expected_command]