Merge branch 'develop' into feature/populate-imagepullsecret-from-container-registry

Author: AlexejPenner

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -46,6 +46,9 @@
 from kubernetes import config as k8s_config
 
 from zenml.config.base_settings import BaseSettings
+from zenml.constants import (
+    METADATA_ORCHESTRATOR_RUN_ID,
+)
 from zenml.enums import StackComponentType
 from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import (
     KubernetesOrchestratorConfig,
@@ -62,6 +65,7 @@
 )
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
 from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
@@ -468,9 +472,7 @@ def submit_pipeline(
         # This will internally also build the command/args for all step pods.
         command = KubernetesOrchestratorEntrypointConfiguration.get_entrypoint_command()
         args = KubernetesOrchestratorEntrypointConfiguration.get_entrypoint_arguments(
-            run_name=orchestrator_run_name,
             deployment_id=deployment.id,
-            kubernetes_namespace=self.config.kubernetes_namespace,
             run_id=placeholder_run.id if placeholder_run else None,
         )
 
@@ -509,6 +511,18 @@ def submit_pipeline(
                 }
             )
 
+        orchestrator_pod_labels = {
+            "pipeline": kube_utils.sanitize_label(pipeline_name),
+        }
+
+        if placeholder_run:
+            orchestrator_pod_labels["run_id"] = kube_utils.sanitize_label(
+                str(placeholder_run.id)
+            )
+            orchestrator_pod_labels["run_name"] = kube_utils.sanitize_label(
+                str(placeholder_run.name)
+            )
+
         # Schedule as CRON job if CRON schedule is given.
         if deployment.schedule:
             if not deployment.schedule.cron_expression:
@@ -526,9 +540,7 @@ def submit_pipeline(
 
             cron_job_manifest, secret_manifests = build_cron_job_manifest(
                 cron_expression=cron_expression,
-                run_name=orchestrator_run_name,
                 pod_name=pod_name,
-                pipeline_name=pipeline_name,
                 image_name=image,
                 command=command,
                 args=args,
@@ -545,6 +557,7 @@ def submit_pipeline(
                 core_api=self._k8s_core_api,
                 registry_credentials=registry_credentials,
                 local_stores_path=local_stores_path,
+                labels=orchestrator_pod_labels,
             )
 
             # Create imagePullSecrets first
@@ -569,14 +582,14 @@ def submit_pipeline(
             pod_manifest, secret_manifests = build_pod_manifest(
                 run_name=orchestrator_run_name,
                 pod_name=pod_name,
-                pipeline_name=pipeline_name,
                 image_name=image,
                 command=command,
                 args=args,
                 privileged=False,
                 pod_settings=orchestrator_pod_settings,
                 service_account_name=service_account_name,
                 env=environment,
+                labels=orchestrator_pod_labels,
                 mount_local_stores=self.config.is_local,
                 namespace=self.config.kubernetes_namespace,
                 auto_generate_image_pull_secrets=settings.auto_generate_image_pull_secrets,
@@ -600,6 +613,11 @@ def submit_pipeline(
                 startup_timeout=settings.pod_startup_timeout,
             )
 
+            metadata: Dict[str, MetadataType] = {
+                METADATA_ORCHESTRATOR_RUN_ID: pod_name,
+            }
+
+            # Wait for the orchestrator pod to finish and stream logs.
             if settings.synchronous:
 
                 def _wait_for_run_to_finish() -> None:
@@ -616,7 +634,8 @@ def _wait_for_run_to_finish() -> None:
                     )
 
                 return SubmissionResult(
-                    wait_for_completion=_wait_for_run_to_finish
+                    metadata=metadata,
+                    wait_for_completion=_wait_for_run_to_finish,
                 )
             else:
                 logger.info(
@@ -625,7 +644,9 @@ def _wait_for_run_to_finish() -> None:
                     f"Run the following command to inspect the logs: "
                     f"`kubectl logs {pod_name} -n {self.config.kubernetes_namespace}`."
                 )
-                return None
+                return SubmissionResult(
+                    metadata=metadata,
+                )
 
     def _get_service_account_name(
         self, settings: KubernetesOrchestratorSettings
@@ -770,3 +791,18 @@ def get_kubernetes_image_pull_secret_data(
             f"Returning ImagePullSecret data: registry='{registry_uri}', username='{username[:3]}***'"
         )
         return registry_uri, username, password
+
+    def get_pipeline_run_metadata(
+        self, run_id: UUID
+    ) -> Dict[str, "MetadataType"]:
+        """Get general component-specific metadata for a pipeline run.
+
+        Args:
+            run_id: The ID of the pipeline run.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        return {
+            METADATA_ORCHESTRATOR_RUN_ID: self.get_orchestrator_run_id(),
+        }