[Run] Fix logging missing artifacts [1.6.x] (mlrun#5764)

server/api/crud/runs.py

@@ -491,9 +491,10 @@ def _get_artifacts_from_uris(
             key_tag_iteration_pairs,
         )
 
-        if len(artifacts) != len(artifact_uris):
+        # DB artifacts result may contain more artifacts if the job is still running
+        if len(artifacts) < len(artifact_uris):
             missing_artifacts = set(artifact_uris.keys()) - {
-                artifact.key for artifact in artifacts
+                artifact["metadata"]["key"] for artifact in artifacts
             }
             logger.warning(
                 "Some artifacts are missing from final run response, they may have been deleted",

server/api/db/sqldb/db.py

@@ -705,7 +705,7 @@ def list_artifacts_for_producer_id(
         producer_id: str,
         project: str = None,
         key_tag_iteration_pairs: list[tuple] = "",
-    ):
+    ) -> ArtifactList:
         project = project or mlrun.mlconf.default_project
         artifact_records = self._find_artifacts_for_producer_id(
             session,

tests/api/crud/test_runs.py

@@ -463,7 +463,7 @@ def test_get_workflow_run_restore_artifacts_metadata(
         self._validate_run_artifacts(artifacts, db, project, run_uid)
 
     @pytest.mark.parametrize("workflow_id", [None, str(uuid.uuid4())])
-    def test_get_workflow_run_iteration_restore_artifacts_metadata(
+    def test_get_run_iteration_restore_artifacts_metadata(
         self, db: sqlalchemy.orm.Session, workflow_id
     ):
         project = "project-name"
@@ -595,6 +595,50 @@ def test_get_workflow_run_best_iteration_restore_artifacts_metadata(
             best_iteration_artifacts + parent_run_arts, db, project, run_uid
         )
 
+    @pytest.mark.parametrize("workflow_uid", [None, str(uuid.uuid4())])
+    def test_get_run_restore_artifacts_metadata_with_missing_artifact(
+        self, db: sqlalchemy.orm.Session, workflow_uid
+    ):
+        project = "project-name"
+        run_uid = str(uuid.uuid4())
+        artifacts = self._generate_artifacts(
+            project, run_uid, workflow_uid, artifacts_len=3
+        )
+
+        # Create only 2 artifacts
+        for artifact in artifacts[:2]:
+            server.api.crud.Artifacts().store_artifact(
+                db,
+                artifact["spec"]["db_key"],
+                artifact,
+                iter=artifact["metadata"]["iter"],
+                project=project,
+                producer_id=workflow_uid or run_uid,
+            )
+
+        labels = {"kind": "job"}
+        if workflow_uid:
+            labels["workflow"] = workflow_uid
+
+        server.api.crud.Runs().store_run(
+            db,
+            {
+                "metadata": {
+                    "name": "run-name",
+                    "uid": run_uid,
+                    "labels": labels,
+                },
+                "status": {
+                    "artifacts": artifacts,
+                },
+            },
+            run_uid,
+            project=project,
+        )
+
+        # Expect only the 2 artifacts to be restored
+        self._validate_run_artifacts(artifacts[:2], db, project, run_uid)
+
     @staticmethod
     def _generate_artifacts(
         project,