From a6b76234857c0e14e8cbe73482b6f56344c030ea Mon Sep 17 00:00:00 2001 From: Serhii Koropets <33310880+koropets@users.noreply.github.com> Date: Mon, 5 Jun 2023 18:18:16 +0300 Subject: [PATCH] Additional labels for builder (#1325) * Add additional labels for model-builder workflow step * Additional tests --- gordo/cli/workflow_generator.py | 30 +++++++++++--- .../resources/argo-workflow.yml.template | 12 ++++-- .../test_workflow_generator.py | 41 +++++++++++++++++++ 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/gordo/cli/workflow_generator.py b/gordo/cli/workflow_generator.py index 3f7d10083..cd77b4fb6 100644 --- a/gordo/cli/workflow_generator.py +++ b/gordo/cli/workflow_generator.py @@ -88,23 +88,24 @@ def prepare_keda_prometheus_query(context): return template.render(**kwargs) -def prepare_resources_labels(value: str) -> List[Tuple[str, Any]]: +def prepare_resources_labels( + value: str, argument: str = "--resources-labels" +) -> List[Tuple[str, Any]]: resources_labels: List[Tuple[str, Any]] = [] if value: try: json_value = json.loads(value) except json.JSONDecodeError as e: raise click.ClickException( - '"--resources-labels=%s" contains invalid JSON value: %s' - % (value, str(e)) + '"%s=%s" contains invalid JSON value: %s' % (argument, value, str(e)) ) if isinstance(json_value, dict): resources_labels = cast(List[Tuple[str, Any]], list(json_value.items())) else: type_name = type(json_value).__name__ raise click.ClickException( - '"--resources-labels=%s" contains value with type "%s" instead "dict"' - % (value, type_name) + '"%s=%s" contains value with type %s instead of dict' + % (argument, value, type_name) ) return resources_labels @@ -312,6 +313,18 @@ def workflow_cli(gordo_ctx): envvar=f"{PREFIX}_RESOURCE_LABELS", default="", ) +@click.option( + "--model-builder-labels", + help="Additional labels for model-builder workflow step. Have to be empty string or a dictionary in JSON format", + envvar=f"{PREFIX}_MODEL_BUILDER_LABELS", + default="", +) +@click.option( + "--server-labels", + help="Additional labels for gordo-server. Have to be empty string or a dictionary in JSON format", + envvar=f"{PREFIX}_SERVER_LABELS", + default="", +) @click.option( "--server-termination-grace-period", help="terminationGracePeriodSeconds for the gordo server", @@ -409,6 +422,13 @@ def workflow_generator_cli(gordo_ctx, **ctx): context["resources_labels"] = prepare_resources_labels(context["resources_labels"]) + context["model_builder_labels"] = prepare_resources_labels( + context["model_builder_labels"], "--model-builder-labels" + ) + context["server_labels"] = prepare_resources_labels( + context["server_labels"], "--server-labels" + ) + if context["pod_security_context"]: pod_security_context = cast(PodSecurityContext, context["pod_security_context"]) context["pod_security_context"] = pod_security_context.dict(exclude_none=True) diff --git a/gordo/workflow/workflow_generator/resources/argo-workflow.yml.template b/gordo/workflow/workflow_generator/resources/argo-workflow.yml.template index aaeba8640..ffb067647 100644 --- a/gordo/workflow/workflow_generator/resources/argo-workflow.yml.template +++ b/gordo/workflow/workflow_generator/resources/argo-workflow.yml.template @@ -7,6 +7,9 @@ # model_builder_resources_limits_memory: Memory limit of the model builder in unit `M` # model_builder_resources_limits_cpu: CPU limit of the model builder in unit `m` +{% macro print_labels(labels) -%}{% if labels %}{% for label, value in labels %}"{{label}}": "{{value}}" +{% endfor %}{% endif %}{% endmacro %} + apiVersion: argoproj.io/v1alpha1 kind: Workflow metadata: @@ -718,8 +721,8 @@ spec:{% if argo_version.major == 2 %} applications.gordo.equinor.com/project-workflow: "{{project_workflow}}"{% endif %} applications.gordo.equinor.com/model-name: "{{'{{inputs.parameters.machine-name}}'}}"{% if "metadata" in builder_runtime and "labels" in builder_runtime["metadata"] %} {{ builder_runtime["metadata"]["labels"] | yaml | indent(8, True) }}{% endif %} - {% if resources_labels is defined %}{% for label, value in resources_labels %}"{{label}}": "{{value}}" - {% endfor %}{% endif %} +{{ print_labels(model_builder_labels) | indent(8, True) }} +{{ print_labels(resources_labels) | indent(8, True) }} tolerations: - effect: NoSchedule key: kubernetes.azure.com/scalesetpriority @@ -1065,8 +1068,8 @@ spec:{% if argo_version.major == 2 %} app.kubernetes.io/managed-by: gordo applications.gordo.equinor.com/project-name: "{{project_name}}" applications.gordo.equinor.com/project-revision: "{{project_revision}}" - {% if resources_labels is defined %}{% for label, value in resources_labels %}"{{label}}": "{{value}}" - {% endfor %}{% endif %} +{{ print_labels(resources_labels) | indent(22, True) }} +{{ print_labels(server_labels) | indent(22, True) }} {% if owner_references is defined %} ownerReferences: {{owner_references}} {% endif %} @@ -1084,6 +1087,7 @@ spec:{% if argo_version.major == 2 %} project: "{{project_name}}" {% if resources_labels is defined %}{% for label, value in resources_labels %}"{{label}}": "{{value}}" {% endfor %}{% endif %} +{{ print_labels(server_labels) | indent(22, True) }} spec: priorityClassName: server-priority terminationGracePeriodSeconds: {{ server_termination_grace_period|default(60, true) }} diff --git a/tests/gordo/workflow/test_workflow_generator/test_workflow_generator.py b/tests/gordo/workflow/test_workflow_generator/test_workflow_generator.py index 5ce4d71b2..991eb1bab 100644 --- a/tests/gordo/workflow/test_workflow_generator/test_workflow_generator.py +++ b/tests/gordo/workflow/test_workflow_generator/test_workflow_generator.py @@ -789,3 +789,44 @@ def test_argo3_binary(path_to_config_files: str): argo_binary="argo3", argo_version="3.1.0", ) + + +def test_model_builder_labels(path_to_config_files: str): + workflow_str = _generate_test_workflow_str( + path_to_config_files, + "config-test-simple.yml", + args=[ + "--model-builder-labels", + '{"model-builder-label1": "value1", "model-builder-label2": "value2"}', + ], + ) + workflow = yaml.safe_load(workflow_str) + templates = workflow["spec"]["templates"] + model_builder_step = [ + step for step in templates if step["name"] == "model-builder" + ][0] + labels = model_builder_step["metadata"]["labels"] + assert labels["model-builder-label1"] == "value1" + assert labels["model-builder-label2"] == "value2" + + +def test_model_server_(path_to_config_files: str): + workflow_str = _generate_test_workflow_str( + path_to_config_files, + "config-test-simple.yml", + args=[ + "--server-labels", + '{"server-label1": "value1", "server-label2": "value2"}', + ], + ) + workflow = yaml.safe_load(workflow_str) + templates = workflow["spec"]["templates"] + gordo_server_step = [ + step for step in templates if step["name"] == "gordo-server-deployment" + ][0] + parameters = gordo_server_step["steps"][0][0]["arguments"]["parameters"] + resource_step = [v for v in parameters if v["name"] == "resource"][0] + resource = yaml.safe_load(resource_step["value"]) + labels = resource["metadata"]["labels"] + assert labels["server-label1"] == "value1" + assert labels["server-label2"] == "value2"