feat(kubernetes): Allow adding custom labels to worker pods

Add a config property for the `KubernetesSenderConfig` to allow adding
custom labels for the created jobs.

Signed-off-by: Martin Nonnenmacher <martin.nonnenmacher@doubleopen.org>

Author: mnonnenmacher

transport/kubernetes/src/main/kotlin/KubernetesMessageSender.kt

@@ -99,7 +99,7 @@ internal class KubernetesMessageSender<T : Any>(
 
         val msgConfig = config.forMessage(message)
         val envVars = createEnvironment()
-        val labels = createTraceIdLabels(traceId) + mapOf(
+        val labels = config.labels + createTraceIdLabels(traceId) + mapOf(
             RUN_ID_LABEL to message.header.ortRunId.toString(),
             WORKER_LABEL to endpoint.configPrefix
         )

transport/kubernetes/src/main/kotlin/KubernetesSenderConfig.kt

@@ -100,6 +100,12 @@ data class KubernetesSenderConfig(
     /** A list with [PvcVolumeMount]s to be added to the new pod. */
     val pvcVolumes: List<PvcVolumeMount> = emptyList(),
 
+    /**
+     * A list of labels to add to the new pod. The labels `ort-worker`, `run-id` and any `trace-id-*` are inserted
+     * automatically by the Kubernetes sender and cannot be overridden; matching entries in this map are ignored.
+     */
+    val labels: Map<String, String> = emptyMap(),
+
     /**
      * A map with annotations to be added to the new pod. The map defines the keys and values of annotation.
      * Corresponding annotations are added to the _template_ section of newly created jobs, so that they are present
@@ -182,6 +188,14 @@ data class KubernetesSenderConfig(
          */
         private const val MOUNT_PVCS_PROPERTY = "mountPvcs"
 
+        /**
+         * The name of the configuration property defining the labels to add to new pods. The value of this property is
+         * interpreted as a comma-separated list of labels, ignoring whitespace around the labels. Each label must have
+         * the format _key=value_. The labels `ort-worker`, `run-id` and any `trace-id-*` are inserted automatically by
+         * the Kubernetes sender and cannot be overridden; matching entries in this list are ignored.
+         */
+        private const val LABELS_PROPERTY = "labels"
+
         /**
          * The name of the configuration property that allows defining annotations based on environment variables.
          * If available, the value of this property is interpreted as a comma-separated list of environment variable
@@ -249,6 +263,9 @@ data class KubernetesSenderConfig(
         /** A regular expression to parse a PVC-based volume mount declaration. */
         private val mountPvcDeclarationRegex = Regex("""(\S+)\s*->\s*([^,]+),([RrWw])""")
 
+        /** A regular expression to split the list of labels. */
+        private val splitLabelsRegex = splitRegex(LIST_SEPARATOR)
+
         /** A regular expression to split the list with environment variables defining annotations. */
         private val splitAnnotationVariablesRegex = splitRegex(LIST_SEPARATOR)
 
@@ -270,6 +287,7 @@ data class KubernetesSenderConfig(
                 args = config.getStringOrDefault(ARGS_PROPERTY, "").splitAtWhitespace(),
                 secretVolumes = config.parseSecretVolumeMounts(),
                 pvcVolumes = config.parsePvcVolumeMounts(),
+                labels = createLabels(config.getStringOrDefault(LABELS_PROPERTY, "")),
                 annotations = createAnnotations(config.getStringOrDefault(ANNOTATIONS_VARIABLES_PROPERTY, "")),
                 serviceAccountName = config.getStringOrNull(SERVICE_ACCOUNT_PROPERTY),
                 enableDebugLogging = config.getBooleanOrDefault(ENABLE_DEBUG_LOGGING_PROPERTY, false),
@@ -329,6 +347,36 @@ data class KubernetesSenderConfig(
                 PvcVolumeMount(claimName, mountPath, readOnly.lowercase() == "r")
             }
 
+        /**
+         * Create the map with labels based on the given [labels]. Extract the labels from the comma-delimited string,
+         * ignoring whitespace around the key value pairs. Each label must have the format _key=value_. Ignore invalid
+         * labels, but log a warning for them. Also ignore labels with reserved keys, but log a warning for them as
+         * well.
+         */
+        private fun createLabels(labels: String): Map<String, String> {
+            if (labels.isEmpty()) return emptyMap()
+
+            val reservedLabels = listOf("ort-worker", "run-id")
+
+            return labels.split(splitLabelsRegex).mapNotNull { label ->
+                val keyValue = label.split(splitKeyValueRegex, limit = 2)
+                if (keyValue.size != 2 || keyValue[0].isEmpty() || keyValue[1].isEmpty()) {
+                    logger.warn("Ignore invalid label declaration: '$label'. Labels must have the format 'key=value'.")
+                    return@mapNotNull null
+                }
+
+                if (keyValue[0] in reservedLabels || keyValue[0].startsWith("trace-id-")) {
+                    logger.warn(
+                        "Ignore label with reserved key '${keyValue[0]}'. The keys 'ort-worker', 'run-id' and " +
+                                "'trace-id-*' are reserved and cannot be used in custom labels."
+                    )
+                    return@mapNotNull null
+                }
+
+                keyValue[0] to keyValue[1]
+            }.toMap()
+        }
+
         /**
          * Create the map with annotations based on the given string with [variableNames]. Extract the names from the
          * comma-delimited string. Look up the values of the referenced variables and extract the keys and values of

transport/kubernetes/src/test/kotlin/KubernetesMessageSenderFactoryTest.kt

@@ -47,6 +47,7 @@ private const val ARGS = "run \"all tests\" fast"
 private const val SECRET_MOUNTS =
     "secret1->/mnt/sec1|sub1 \"secret2->/path/with/white space\" \"secret3 -> /mnt/other | sub2\""
 private const val PVC_MOUNTS = "pvc1->/mnt/pvc1,R \"pvc2->/path/with/white space,W\" \"pvc3 -> /mnt/other,r\""
+private const val LABELS = "label1=value1 , label2 = value2"
 private const val SERVICE_ACCOUNT = "test_service_account"
 
 private val annotationVariables = mapOf(
@@ -74,6 +75,7 @@ class KubernetesMessageSenderFactoryTest : StringSpec({
             "$keyPrefix.enableDebugLogging" to "true",
             "$keyPrefix.mountSecrets" to SECRET_MOUNTS,
             "$keyPrefix.mountPvcs" to PVC_MOUNTS,
+            "$keyPrefix.labels" to LABELS,
             "$keyPrefix.annotationVariables" to annotationVariables.keys.joinToString(),
             "$keyPrefix.serviceAccount" to SERVICE_ACCOUNT
         )
@@ -104,6 +106,10 @@ class KubernetesMessageSenderFactoryTest : StringSpec({
                 PvcVolumeMount("pvc2", "/path/with/white space", readOnly = false),
                 PvcVolumeMount("pvc3", "/mnt/other", readOnly = true)
             )
+            labels shouldContainExactly mapOf(
+                "label1" to "value1",
+                "label2" to "value2"
+            )
             annotations shouldContainExactly mapOf(
                 "ort-server.org/test" to "true",
                 "ort-server.org/performance" to "fast"
@@ -137,6 +143,7 @@ class KubernetesMessageSenderFactoryTest : StringSpec({
             restartPolicy shouldBe "OnFailure"
             imagePullSecret should beNull()
             secretVolumes should beEmpty()
+            labels.keys should beEmpty()
             annotations.keys should beEmpty()
             serviceAccountName should beNull()
             enableDebugLogging shouldBe false
@@ -186,6 +193,46 @@ class KubernetesMessageSenderFactoryTest : StringSpec({
         )
     }
 
+    "Invalid labels are ignored" {
+        val keyPrefix = "analyzer.sender"
+        val labels = "label1=value1 , invalid label, label2 = value2, =invalid, invalid="
+        val configMap = mapOf(
+            "$keyPrefix.type" to KubernetesSenderConfig.TRANSPORT_NAME,
+            "$keyPrefix.namespace" to NAMESPACE,
+            "$keyPrefix.imageName" to IMAGE_NAME,
+            "$keyPrefix.labels" to labels
+        )
+        val configManager = ConfigManager.create(ConfigFactory.parseMap(configMap))
+
+        val sender = MessageSenderFactory.createSender(AnalyzerEndpoint, configManager)
+
+        sender.shouldBeTypeOf<KubernetesMessageSender<AnalyzerEndpoint>>()
+        sender.config.labels shouldContainExactly mapOf(
+            "label1" to "value1",
+            "label2" to "value2"
+        )
+    }
+
+    "Reserved labels are ignored" {
+        val keyPrefix = "analyzer.sender"
+        val labels = "label1=value1,label2=value2,ort-worker=invalid,run-id=invalid,trace-id-0=invalid"
+        val configMap = mapOf(
+            "$keyPrefix.type" to KubernetesSenderConfig.TRANSPORT_NAME,
+            "$keyPrefix.namespace" to NAMESPACE,
+            "$keyPrefix.imageName" to IMAGE_NAME,
+            "$keyPrefix.labels" to labels
+        )
+        val configManager = ConfigManager.create(ConfigFactory.parseMap(configMap))
+
+        val sender = MessageSenderFactory.createSender(AnalyzerEndpoint, configManager)
+
+        sender.shouldBeTypeOf<KubernetesMessageSender<AnalyzerEndpoint>>()
+        sender.config.labels shouldContainExactly mapOf(
+            "label1" to "value1",
+            "label2" to "value2"
+        )
+    }
+
     "Invalid variables defining annotations are ignored" {
         val keyPrefix = "analyzer.sender"
         val validVariable = "validVariable"

transport/kubernetes/src/test/kotlin/KubernetesMessageSenderTest.kt

@@ -119,15 +119,17 @@ class KubernetesMessageSenderTest : StringSpec({
 
         verifyLabels(
             actualLabels = job.metadata?.labels.orEmpty(),
-            expectedRunId = message.header.ortRunId
+            expectedRunId = message.header.ortRunId,
+            customLabels = senderConfig.labels
         )
 
         val jobAnnotations = job.spec?.template?.metadata?.annotations.orEmpty()
         jobAnnotations shouldBe annotations
 
         verifyLabels(
             actualLabels = job.spec?.template?.metadata?.labels.orEmpty(),
-            expectedRunId = message.header.ortRunId
+            expectedRunId = message.header.ortRunId,
+            customLabels = senderConfig.labels
         )
     }
 
@@ -278,6 +280,7 @@ private fun createConfig(vararg overrides: Pair<String, String>): KubernetesSend
         "imagePullSecret" to "image_pull_secret",
         "mountSecrets" to "secretService->/mnt/secret topSecret->/mnt/top/secret|sub-secret",
         "mountPvcs" to "pvc1->/mnt/readOnly,R pvc2->/mnt/data,W",
+        "labels" to "label1=value1,label2=value2,ort-worker=invalid,run-id=invalid,trace-id-0=invalid",
         "annotationVariables" to "v1,v2",
         "serviceAccountName" to "test_service_account"
     )
@@ -296,12 +299,16 @@ private fun createConfig(vararg overrides: Pair<String, String>): KubernetesSend
 private fun messageWithProperties(vararg properties: Pair<String, String>): Message<AnalyzerRequest> =
     message.copy(header = header.copy(transportProperties = mapOf(*properties)))
 
-private fun verifyLabels(actualLabels: Map<String, String>, expectedRunId: Long) {
+private fun verifyLabels(actualLabels: Map<String, String>, expectedRunId: Long, customLabels: Map<String, String>) {
     val traceIdFromLabels = (0..3).fold("") { id, idx ->
         id + actualLabels.getValue("trace-id-$idx")
     }
     traceIdFromLabels shouldBe traceId
 
     actualLabels["run-id"] shouldBe expectedRunId.toString()
     actualLabels["ort-worker"] shouldBe "analyzer"
+
+    customLabels.forEach { (key, value) ->
+        actualLabels[key] shouldBe value
+    }
 }

website/docs/admin-guide/infrastructure/transport/kubernetes.md

@@ -26,6 +26,7 @@ endpoint {
         commands = "/bin/sh"
         args = "-c java my.pkg.MyClass"
         mountSecrets = "server-secrets->/mnt/secrets server-certificates->/mnt/certificates"
+        labels = "label1=value1,label2=value2"
         annotationVariables = "ANNOTATION_VAR1, ANNOTATION_VAR2"
         serviceAccount = "my_service_account"
         cpuRequest = 250m
@@ -108,6 +109,14 @@ On this path, for each key of the secret a file is created whose content is the
 To achieve this, the Kubernetes Transport implementation generates corresponding `volume` and `volumeMount` declarations in the pod configuration.
 This mechanism is useful not only for secrets but also for other kinds of external data that should be accessible from a pod, for instance, custom certificates.
 
+### `labels`
+
+**Default: `empty`**
+
+A comma-separated list of labels to be added to the job.
+Labels are key-value pairs and must have the form `key=value`.
+The labels `ort-worker`, `run-id` and any `trace-id-*` are inserted automatically by the Kubernetes sender and cannot be overridden; matching entries in this list are ignored.
+
 ### `annotationVariables`
 
 **Default: `empty`**