Store Kubernetes step job names in run metadata (#4180)

* Store Kubernetes step job names in run metadata

* Tests

* Mypy

Author: schustmi

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

@@ -28,7 +28,7 @@
 from zenml.entrypoints.step_entrypoint_configuration import (
     StepEntrypointConfiguration,
 )
-from zenml.enums import ExecutionMode, ExecutionStatus
+from zenml.enums import ExecutionMode, ExecutionStatus, MetadataResourceTypes
 from zenml.exceptions import AuthorizationException
 from zenml.integrations.kubernetes import kube_utils
 from zenml.integrations.kubernetes.constants import (
@@ -61,6 +61,7 @@
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineSnapshotResponse,
+    RunMetadataResource,
 )
 from zenml.orchestrators import publish_utils
 from zenml.orchestrators.step_run_utils import (
@@ -524,6 +525,16 @@ def start_step_job(node: Node) -> NodeStatus:
                 job_manifest=job_manifest,
             )
 
+            Client().create_run_metadata(
+                metadata={"step_jobs": {step_name: job_name}},
+                resources=[
+                    RunMetadataResource(
+                        id=pipeline_run.id,
+                        type=MetadataResourceTypes.PIPELINE_RUN,
+                    )
+                ],
+            )
+
             node.metadata["job_name"] = job_name
 
             return NodeStatus.RUNNING

src/zenml/zen_stores/schemas/utils.py

@@ -127,10 +127,23 @@ def fetch_metadata(self, **kwargs: Any) -> Dict[str, MetadataType]:
                 and the values represent the latest entry with this key.
         """
         metadata_collection = self.fetch_metadata_collection(**kwargs)
-        return {
-            k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
-            for k, v in metadata_collection.items()
-        }
+        metadata: Dict[str, MetadataType] = {}
+
+        for key, values in metadata_collection.items():
+            values = sorted(values, key=lambda x: x.created, reverse=False)
+
+            if all(isinstance(item.value, dict) for item in values):
+                # All metadata values for this key are dictionaries, so we can
+                # merge them into a single dictionary
+                metadata[key] = {
+                    k: v
+                    for item in values
+                    for k, v in item.value.items()  # type: ignore[union-attr]
+                }
+            else:
+                metadata[key] = values[-1].value
+
+        return metadata
 
 
 def get_resource_type_name(schema_class: Type[BaseSchema]) -> str:

tests/integration/functional/test_client.py

@@ -538,6 +538,56 @@ def test_create_run_metadata_for_artifact(clean_client_with_run: Client):
     assert rm["axel"] == "is awesome"
 
 
+def test_create_run_metadata_overwrite(clean_client_with_run: Client):
+    """Test creating run metadata with the same key but different value."""
+    pipeline_run = clean_client_with_run.list_pipeline_runs()[0]
+
+    clean_client_with_run.create_run_metadata(
+        metadata={"axel": "is not awesome"},
+        resources=[
+            RunMetadataResource(
+                id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
+            )
+        ],
+    )
+    clean_client_with_run.create_run_metadata(
+        metadata={"axel": "is awesome"},
+        resources=[
+            RunMetadataResource(
+                id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
+            )
+        ],
+    )
+
+    rm = clean_client_with_run.get_pipeline_run(pipeline_run.id).run_metadata
+    assert rm["axel"] == "is awesome"
+
+
+def test_create_run_metadata_dict_merge(clean_client_with_run: Client):
+    """Test creating run metadata with the same key merges dictionaries."""
+    pipeline_run = clean_client_with_run.list_pipeline_runs()[0]
+
+    clean_client_with_run.create_run_metadata(
+        metadata={"key": {"a": 1}},
+        resources=[
+            RunMetadataResource(
+                id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
+            )
+        ],
+    )
+    clean_client_with_run.create_run_metadata(
+        metadata={"key": {"a": 2, "b": 2}},
+        resources=[
+            RunMetadataResource(
+                id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
+            )
+        ],
+    )
+
+    rm = clean_client_with_run.get_pipeline_run(pipeline_run.id).run_metadata
+    assert rm["key"] == {"a": 2, "b": 2}
+
+
 # .---------.
 # | SECRETS |
 # '---------'