added remove_checkpoints instead of single checkpoint

Author: ilongin

src/datachain/catalog/catalog.py

@@ -2145,8 +2145,7 @@ def cleanup_checkpoints(self, ttl_seconds: int | None = None) -> int:
                 self._cleanup_udf_tables(f"udf_{group_id}_", "_input")
 
         checkpoints = list(self.metastore.list_checkpoints(job_id=inactive_job_ids))
-        for ch in checkpoints:
-            self.metastore.remove_checkpoint(ch.id)
+        self.metastore.remove_checkpoints([ch.id for ch in checkpoints])
 
         logger.info(
             "Checkpoint cleanup complete: removed %d checkpoints from %d jobs",

src/datachain/data_storage/metastore.py

@@ -591,8 +591,10 @@ def get_or_create_checkpoint(
         """Get or create checkpoint. Must be atomic and idempotent."""
 
     @abstractmethod
-    def remove_checkpoint(self, checkpoint_id: str, conn: Any | None = None) -> None:
-        """Removes a checkpoint by ID"""
+    def remove_checkpoints(
+        self, checkpoint_ids: list[str], conn: Any | None = None
+    ) -> None:
+        """Soft-delete checkpoints by IDs."""
 
     @abstractmethod
     def get_jobs_with_expired_checkpoints(
@@ -2743,10 +2745,14 @@ def get_dataset_version_for_job_ancestry(
 
         return self.dataset_version_class.parse(*results[0])
 
-    def remove_checkpoint(self, checkpoint_id: str, conn: Any | None = None) -> None:
+    def remove_checkpoints(
+        self, checkpoint_ids: list[str], conn: Any | None = None
+    ) -> None:
+        if not checkpoint_ids:
+            return
         self.db.execute(
             self._checkpoints.update()
-            .where(self._checkpoints.c.id == checkpoint_id)
+            .where(self._checkpoints.c.id.in_(checkpoint_ids))
             .values(status=CheckpointStatus.DELETED),
             conn=conn,
         )

src/datachain/query/dataset.py

@@ -1155,7 +1155,7 @@ def _run_from_scratch(
         )
 
         if partial_checkpoint:
-            self.metastore.remove_checkpoint(partial_checkpoint.id)
+            self.metastore.remove_checkpoints([partial_checkpoint.id])
             self.metastore.get_or_create_checkpoint(self.job.id, hash_output)
             logger.debug(
                 "UDF(%s) [job=%s run_group=%s]: Promoted partial to final, hash=%s",
@@ -1291,7 +1291,7 @@ def _continue_udf(
             partial_table, udf_output_table_name(self.job.id, hash_output)
         )
 
-        self.metastore.remove_checkpoint(partial_checkpoint.id)
+        self.metastore.remove_checkpoints([partial_checkpoint.id])
         self.metastore.get_or_create_checkpoint(self.job.id, hash_output)
         logger.debug(
             "UDF(%s) [job=%s run_group=%s]: Promoted partial to final, hash=%s",