feat(gc): include STALE and REMOVING dataset versions in GC cleanup (#1621)

* feat(gc): include STALE dataset versions in GC cleanup

Studio's lifecycle system (datachain-ai/studio#12515) marks excess
dataset versions as STALE instead of deleting them directly. The
datachain GC must then perform the actual warehouse and metastore
deletion.

Add DatasetStatus.STALE to get_incomplete_dataset_versions() so
STALE versions with a finished or absent job are returned and
removed by cleanup_failed_dataset_versions().

* Update metastore.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* feat(dataset): add REMOVING status for versions marked for deletion

Introduce DatasetStatus.REMOVING so versions marked for deletion are distinct from STALE. Include REMOVING in get_incomplete_dataset_versions cleanup query and update tests accordingly.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: amritghimire

src/datachain/data_storage/metastore.py

@@ -329,14 +329,17 @@ def get_incomplete_dataset_versions(
         self, job_id: str | None = None
     ) -> list[tuple[DatasetRecord, str]]:
         """
-        Get incomplete dataset versions to clean up.
+        Get incomplete, stale, or removed dataset versions to clean up.
 
         When job_id is provided, returns versions belonging to that specific
         job (used during job failure cleanup).
 
-        When job_id is None, returns all incomplete dataset versions
-        whose associated job is finished, plus versions with no job_id
-        that are older than STALE_CREATED_THRESHOLD_HOURS (used by gc).
+        When job_id is None, returns all versions that are safe to delete:
+        - Status CREATED, FAILED, STALE where either:
+          - the associated job has finished, or
+          - there is no associated job (job_id is NULL) and the version is
+            older than STALE_CREATED_THRESHOLD_HOURS
+        - Status REMOVING: marked for deletion
 
         Returns:
             List of (DatasetRecord, version_string) tuples. Each DatasetRecord
@@ -1654,7 +1657,14 @@ def get_incomplete_dataset_versions(
                 )
             )
             .where(
-                dv.c.status.in_([DatasetStatus.CREATED, DatasetStatus.FAILED]),
+                dv.c.status.in_(
+                    [
+                        DatasetStatus.CREATED,
+                        DatasetStatus.FAILED,
+                        DatasetStatus.STALE,
+                        DatasetStatus.REMOVING,
+                    ]
+                ),
                 or_(
                     # job is finished
                     j.c.status.in_(

src/datachain/dataset.py

@@ -247,6 +247,7 @@ class DatasetStatus:
     FAILED = 3
     COMPLETE = 4
     STALE = 6
+    REMOVING = 7
 
 
 @dataclass
@@ -340,6 +341,7 @@ def is_final_status(self) -> bool:
             DatasetStatus.FAILED,
             DatasetStatus.COMPLETE,
             DatasetStatus.STALE,
+            DatasetStatus.REMOVING,
         ]
 
     def update(self, **kwargs):

tests/unit/test_dataset_status_management.py

@@ -300,3 +300,42 @@ def test_public_api_move_dataset_rejects_non_complete(
     # Should raise error for FAILED dataset
     with pytest.raises(DatasetNotFoundError):
         move_dataset(dataset_failed.name, "new_name_failed", session=test_session)
+
+
+@pytest.fixture
+def dataset_marked_for_removal(test_session, job) -> DatasetRecord:
+    ds = dc.read_values(value=["v1"], session=test_session).save(
+        "ds_marked_for_removal"
+    )
+    dataset = ds.dataset
+    assert dataset is not None
+    dv = test_session.catalog.metastore._datasets_versions
+    test_session.catalog.metastore.db.execute(
+        dv.update()
+        .where(dv.c.dataset_id == dataset.id)
+        .values(status=DatasetStatus.REMOVING)
+    )
+    return test_session.catalog.get_dataset(dataset.name, include_incomplete=True)
+
+
+def test_get_incomplete_dataset_versions_includes_marked_for_removal(
+    test_session, job, dataset_marked_for_removal
+):
+    to_clean = test_session.catalog.metastore.get_incomplete_dataset_versions()
+    assert dataset_marked_for_removal.name not in {ds.name for ds, _ in to_clean}
+
+    test_session.catalog.metastore.set_job_status(job.id, JobStatus.COMPLETE)
+    to_clean = test_session.catalog.metastore.get_incomplete_dataset_versions()
+    assert dataset_marked_for_removal.name in {ds.name for ds, _ in to_clean}
+
+
+def test_cleanup_failed_dataset_versions_removes_marked_for_removal(
+    test_session, job, dataset_marked_for_removal
+):
+    test_session.catalog.metastore.set_job_status(job.id, JobStatus.COMPLETE)
+
+    num_removed = test_session.catalog.cleanup_failed_dataset_versions()
+    assert num_removed == 1
+
+    with pytest.raises(DatasetNotFoundError):
+        test_session.catalog.get_dataset(dataset_marked_for_removal.name)