From 8f0bde7f52ff7fda4f51d7513e48f882f9aa0b0c Mon Sep 17 00:00:00 2001 From: Alputer Date: Thu, 21 Nov 2024 23:03:55 +0100 Subject: [PATCH] feat(utils): add central function for Dask K8s component names (#478) Closes reanahub/reana#841 --- reana_commons/utils.py | 26 ++++++++++++++++++++++++++ tests/test_utils.py | 20 ++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/reana_commons/utils.py b/reana_commons/utils.py index 16051681..01064e67 100644 --- a/reana_commons/utils.py +++ b/reana_commons/utils.py @@ -414,6 +414,32 @@ def build_unique_component_name(component_type, id=None): ) +def get_trimmed_workflow_id(workflow_id: str, trim_level: int) -> str: + """Trim the given workflow id and return it. For a workflow with id '9eef9a08-5629-420d-8e97-29d498d88e20' with trim level 4, it returns '9eef9a08'.""" + return str(workflow_id).rsplit("-", trim_level)[0] + + +def get_dask_component_name( + workflow_id: str, name_type: str, dashboard_service_namespace: str = "default" +) -> str: + """Generate the name of a Dask-related component based on the workflow ID and name type. Note that we trim the end of the uuid so uniqueness is not 100% guaranteed.""" + trimmed_workflow_id = get_trimmed_workflow_id(str(workflow_id), 4) + name_map = { + "cluster": f"reana-run-dask-{trimmed_workflow_id}", + "autoscaler": f"dask-autoscaler-{trimmed_workflow_id}", + "dashboard_ingress": f"dask-dashboard-ingress-{trimmed_workflow_id}", + "dashboard_service": f"reana-run-dask-{trimmed_workflow_id}-scheduler", + "dashboard_service_uri": f"reana-run-dask-{trimmed_workflow_id}-scheduler.{dashboard_service_namespace}.svc.cluster.local:8786", + "dashboard_ingress_middleware": f"dask-dashboard-ingress-{trimmed_workflow_id}-replacepath", + "database_model_service": f"dask-service-{trimmed_workflow_id}", + } + if name_type not in name_map: + raise ValueError( + f"Invalid name type: '{name_type}'. Valid types are: {', '.join(name_map.keys())}." + ) + return name_map[name_type] + + def get_usage_percentage(usage: int, limit: int) -> str: """Usage percentage.""" if limit == 0: diff --git a/tests/test_utils.py b/tests/test_utils.py index a36e4e5b..aa150866 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -24,6 +24,7 @@ click_table_printer, format_cmd, get_workflow_status_change_verb, + get_trimmed_workflow_id, ) @@ -150,3 +151,22 @@ def test_get_workflow_status_change_verb_invalid(): """Test get_workflow_status_change_verb with an invalid status.""" with pytest.raises(ValueError, match="invalid"): get_workflow_status_change_verb("invalid") + + +@pytest.mark.parametrize( + "workflow_id, trim_level, expected", + [ + ("9eef9a08-5629-420d-8e97-29d498d88e20", 4, "9eef9a08"), + ("9eef9a08-5629-420d-8e97-29d498d88e20", 3, "9eef9a08-5629"), + ("9eef9a08-5629-420d-8e97-29d498d88e20", 2, "9eef9a08-5629-420d"), + ("9eef9a08-5629-420d-8e97-29d498d88e20", 1, "9eef9a08-5629-420d-8e97"), + ( + "9eef9a08-5629-420d-8e97-29d498d88e20", + 0, + "9eef9a08-5629-420d-8e97-29d498d88e20", + ), + ], +) +def test_get_trimmed_workflow_id(workflow_id, trim_level, expected): + """Test get_trimmed_workflow_id function with several different inputs.""" + assert get_trimmed_workflow_id(workflow_id, trim_level) == expected