UDF Checkpoints cleanup (#1590)

UDF Checkpoints cleanup

Author: ilongin

docs/commands/gc.md

@@ -0,0 +1,46 @@
+# gc
+
+Garbage collect temporary tables, failed dataset versions, and outdated checkpoints.
+
+## Synopsis
+
+```usage
+usage: datachain gc [-h] [-v] [-q] [--checkpoint-ttl CHECKPOINT_TTL]
+```
+
+## Description
+
+This command cleans up internal DataChain storage by removing:
+
+- **Temporary tables** created during query execution that were not properly cleaned up (e.g., due to crashes or interrupted operations).
+- **Failed dataset versions** that were left in an incomplete or failed state.
+- **Outdated checkpoints** and their associated UDF tables that have exceeded the time-to-live (TTL) threshold. See [Checkpoints](../guide/checkpoints.md) for more details.
+
+## Options
+
+* `-h`, `--help` - Show the help message and exit.
+* `-v`, `--verbose` - Be verbose.
+* `-q`, `--quiet` - Be quiet.
+* `--checkpoint-ttl` - Time-to-live for checkpoints in seconds. Checkpoints older than this value are considered outdated and eligible for cleanup. Defaults to 4 hours (14400 seconds).
+
+## Examples
+
+1. Run garbage collection:
+```bash
+datachain gc
+```
+
+2. Run garbage collection with a custom checkpoint TTL of 1 hour:
+```bash
+datachain gc --checkpoint-ttl 3600
+```
+
+Example output:
+```
+Collecting temporary tables...
+  Removed 3 temporary tables.
+Collecting failed dataset versions...
+  No failed dataset versions to clean up.
+Collecting outdated checkpoints...
+  Removed 5 outdated checkpoints.
+```

docs/commands/index.md

@@ -33,3 +33,8 @@ DataChain is a command-line tool for wrangling unstructured AI data at scale. Us
 	- Cancel running jobs with [`datachain job cancel`](job/cancel.md)
 
 	- Check for the clusters available for jobs [`datachain job clusters`](job/clusters.md)
+
+
+3.  **Maintenance**
+
+	- Clean up temporary tables, failed versions, and outdated checkpoints with [`datachain gc`](gc.md)

mkdocs.yml

@@ -97,6 +97,7 @@ nav:
       - 📖 CLI Reference:
           - Overview: commands/index.md
           - Commands:
+              - gc: commands/gc.md
               - auth:
                   - login: commands/auth/login.md
                   - logout: commands/auth/logout.md

src/datachain/catalog/catalog.py

@@ -10,6 +10,7 @@
 from contextlib import contextmanager, suppress
 from copy import copy
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 from functools import cached_property, reduce
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
@@ -24,6 +25,7 @@
 from tqdm.auto import tqdm
 
 from datachain.cache import Cache
+from datachain.checkpoint import Checkpoint, CheckpointStatus
 from datachain.client import Client
 from datachain.dataset import (
     DATASET_PREFIX,
@@ -72,7 +74,7 @@
 
 DEFAULT_DATASET_DIR = "dataset"
 
-TTL_INT = 4 * 60 * 60
+CHECKPOINTS_TTL = 4 * 60 * 60
 
 INDEX_INTERNAL_ERROR_MESSAGE = "Internal error on indexing"
 DATASET_INTERNAL_ERROR_MESSAGE = "Internal error on creating dataset"
@@ -2094,3 +2096,64 @@ def index(
             only_index=True,
         ):
             pass
+
+    def cleanup_checkpoints(self, ttl_seconds: int | None = None) -> int:
+        """Clean up outdated checkpoints and their associated UDF tables.
+
+        Algorithm:
+        1. Atomically mark expired checkpoints as EXPIRED (single query) —
+           prevents running jobs from using them via find_checkpoint.
+           Then collect all EXPIRED checkpoints and determine which run
+           groups have no remaining ACTIVE checkpoints.
+        2. Remove output/partial tables using exact names from checkpoints
+        3. For fully-inactive run groups: remove shared input tables
+        4. Mark checkpoints as DELETED
+        """
+        if ttl_seconds is None:
+            ttl_seconds = CHECKPOINTS_TTL
+
+        ttl_threshold = datetime.now(timezone.utc) - timedelta(seconds=ttl_seconds)
+
+        # Expire + collect everything in one metastore call
+        checkpoints, inactive_group_ids = self.metastore.expire_checkpoints(
+            ttl_threshold
+        )
+        if not checkpoints:
+            return 0
+
+        logger.info(
+            "Cleaning %d expired checkpoints across %d inactive run groups",
+            len(checkpoints),
+            len(inactive_group_ids),
+        )
+
+        # Remove output/partial tables using exact names from checkpoints
+        output_tables = [ch.table_name for ch in checkpoints]
+        if output_tables:
+            logger.info(
+                "Removing %d UDF output tables: %s", len(output_tables), output_tables
+            )
+            self.warehouse.cleanup_tables(output_tables)
+
+        # Shared input tables — only when entire run group is inactive
+        for group_id in inactive_group_ids:
+            input_tables = self.warehouse.db.list_tables(
+                pattern=Checkpoint.input_table_pattern(group_id)
+            )
+            if input_tables:
+                logger.info(
+                    "Removing %d shared input tables: %s",
+                    len(input_tables),
+                    input_tables,
+                )
+                self.warehouse.cleanup_tables(input_tables)
+
+        self.metastore.update_checkpoints(
+            [ch.id for ch in checkpoints], status=CheckpointStatus.DELETED
+        )
+
+        logger.info(
+            "Checkpoint cleanup complete: removed %d checkpoints",
+            len(checkpoints),
+        )
+        return len(checkpoints)

src/datachain/checkpoint.py

@@ -1,6 +1,13 @@
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
+from enum import IntEnum
+
+
+class CheckpointStatus(IntEnum):
+    ACTIVE = 0
+    EXPIRED = 1
+    DELETED = 2
 
 
 @dataclass
@@ -24,6 +31,34 @@ class Checkpoint:
     hash: str
     partial: bool
     created_at: datetime
+    status: int = CheckpointStatus.ACTIVE
+
+    @staticmethod
+    def output_table_name(job_id: str, _hash: str) -> str:
+        """Final UDF output table. Job-specific, created when UDF completes."""
+        return f"udf_{job_id}_{_hash}_output"
+
+    @staticmethod
+    def partial_output_table_name(job_id: str, _hash: str) -> str:
+        """Partial UDF output table. Temporary, renamed to final on completion."""
+        return f"udf_{job_id}_{_hash}_output_partial"
+
+    @staticmethod
+    def input_table_name(group_id: str, _hash: str) -> str:
+        """Shared UDF input table. Scoped to run group, reused across jobs."""
+        return f"udf_{group_id}_{_hash}_input"
+
+    @staticmethod
+    def input_table_pattern(group_id: str) -> str:
+        """LIKE pattern for finding all input tables in a run group."""
+        return f"udf_{group_id}_%_input"
+
+    @property
+    def table_name(self) -> str:
+        """UDF output table name associated with this checkpoint."""
+        if self.partial:
+            return self.partial_output_table_name(self.job_id, self.hash)
+        return self.output_table_name(self.job_id, self.hash)
 
     @classmethod
     def parse(
@@ -33,11 +68,13 @@ def parse(
         _hash: str,
         partial: bool,
         created_at: datetime,
+        status: int = CheckpointStatus.ACTIVE,
     ) -> "Checkpoint":
         return cls(
             str(id),
             job_id,
             _hash,
             bool(partial),
             created_at,
+            int(status),
         )

src/datachain/cli/__init__.py

@@ -97,7 +97,7 @@ def handle_command(args, catalog, client_config) -> int:
         "index": lambda: handle_index_command(args, catalog),
         "completion": lambda: handle_completion_command(args),
         "clear-cache": lambda: clear_cache(catalog),
-        "gc": lambda: garbage_collect(catalog),
+        "gc": lambda: garbage_collect(catalog, checkpoint_ttl=args.checkpoint_ttl),
         "auth": lambda: process_auth_cli_args(args),
         "job": lambda: process_jobs_args(args),
         "pipeline": lambda: process_pipeline_args(args, catalog),

src/datachain/cli/commands/misc.py

@@ -10,21 +10,28 @@ def clear_cache(catalog: "Catalog"):
     catalog.cache.clear()
 
 
-def garbage_collect(catalog: "Catalog"):
+def garbage_collect(catalog: "Catalog", checkpoint_ttl: int | None = None):
+    print("Collecting temporary tables...")
     temp_tables = catalog.get_temp_table_names()
-    num_versions_removed = catalog.cleanup_failed_dataset_versions()
-
-    total_cleaned = len(temp_tables) + num_versions_removed
+    if temp_tables:
+        catalog.cleanup_tables(temp_tables)
+        print(f"  Removed {len(temp_tables)} temporary tables.")
+    else:
+        print("  No temporary tables to clean up.")
 
-    if total_cleaned == 0:
-        print("Nothing to clean up.")
+    print("Collecting failed dataset versions...")
+    num_versions = catalog.cleanup_failed_dataset_versions()
+    if num_versions:
+        print(f"  Removed {num_versions} failed/incomplete dataset versions.")
     else:
-        if temp_tables:
-            print(f"Garbage collecting {len(temp_tables)} tables.")
-            catalog.cleanup_tables(temp_tables)
+        print("  No failed dataset versions to clean up.")
 
-        if num_versions_removed:
-            print(f"Cleaned {num_versions_removed} failed/incomplete dataset versions.")
+    print("Collecting outdated checkpoints...")
+    num_checkpoints = catalog.cleanup_checkpoints(ttl_seconds=checkpoint_ttl)
+    if num_checkpoints:
+        print(f"  Removed {num_checkpoints} outdated checkpoints.")
+    else:
+        print("  No outdated checkpoints to clean up.")
 
 
 def completion(shell: str) -> str:

src/datachain/cli/parser/__init__.py

@@ -478,9 +478,22 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parse_gc = subp.add_parser(
         "gc",
         parents=[parent_parser],
-        description="Garbage collect temporary tables and failed dataset versions.",
+        description=(
+            "Garbage collect temporary tables,"
+            " failed dataset versions, and outdated checkpoints."
+        ),
         formatter_class=CustomHelpFormatter,
     )
+    parse_gc.add_argument(
+        "--checkpoint-ttl",
+        type=int,
+        default=None,
+        help=(
+            "Time-to-live for checkpoints in seconds."
+            " Checkpoints older than this are removed."
+            " Defaults to 4 hours (14400 seconds)."
+        ),
+    )
     add_anon_arg(parse_gc)
 
     subp.add_parser("internal-run-udf", parents=[parent_parser])

src/datachain/data_storage/db_engine.py

@@ -123,15 +123,15 @@ def has_table(self, name: str) -> bool:
         return sa.inspect(self.engine).has_table(name)
 
     @abstractmethod
-    def list_tables(self, prefix: str = "") -> list[str]:
+    def list_tables(self, pattern: str = "") -> list[str]:
         """
-        List all table names, optionally filtered by prefix.
+        List all table names, optionally filtered by a SQL LIKE pattern.
 
         Args:
-            prefix: Optional prefix to filter table names
+            pattern: SQL LIKE pattern to filter table names (e.g. 'udf_%')
 
         Returns:
-            List of table names matching the prefix
+            List of table names matching the pattern
         """
 
     @abstractmethod