equinor · koropets · Jun 5, 2023 · Jun 4, 2023 · Jun 5, 2023
diff --git 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
@@ -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
@@ -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"