Implement cleanup of run files upon model session and run deletion

DashAI/back/api/api_v1/endpoints/model_sessions.py

@@ -10,7 +10,8 @@
     ColumnsValidationParams,
     ModelSessionParams,
 )
-from DashAI.back.dependencies.database.models import Dataset, ModelSession
+from DashAI.back.api.utils import remove_path
+from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import sessionmaker
@@ -271,6 +272,8 @@ async def delete_model_session(
     -------
     Response with code 204 NO_CONTENT
     """
+    import os
+
     with session_factory() as db:
         try:
             model_session = db.get(ModelSession, model_session_id)
@@ -279,16 +282,42 @@ async def delete_model_session(
                     status_code=status.HTTP_404_NOT_FOUND,
                     detail="Model session not found",
                 )
+
+            # Snapshot the on-disk paths of the runs that the FK cascade is
+            # about to delete, so they can be cleaned up after the commit.
+            runs = db.query(Run).filter(Run.model_session_id == model_session_id).all()
+            paths_to_clean = [
+                path
+                for run in runs
+                for path in (
+                    run.run_path,
+                    run.plot_history_path,
+                    run.plot_slice_path,
+                    run.plot_contour_path,
+                    run.plot_importance_path,
+                )
+            ]
+
             db.delete(model_session)
             db.commit()
-            return Response(status_code=status.HTTP_204_NO_CONTENT)
         except exc.SQLAlchemyError as e:
             log.exception(e)
             raise HTTPException(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail="Internal database error",
             ) from e
 
+    # Best-effort disk cleanup, done only once the DB delete has committed:
+    # a cleanup failure here must not undo an otherwise-successful deletion.
+    for path in paths_to_clean:
+        if path and os.path.exists(path):
+            try:
+                remove_path(path)
+            except (OSError, ValueError) as e:
+                log.warning(f"Failed to delete path {path}: {e}")
+
+    return Response(status_code=status.HTTP_204_NO_CONTENT)
+
 
 @router.patch("/{model_session_id}")
 @inject

DashAI/back/api/api_v1/endpoints/runs.py

@@ -8,6 +8,7 @@
 from sqlalchemy import exc, select
 
 from DashAI.back.api.api_v1.schemas.runs_params import RunParams, UpdateRunParams
+from DashAI.back.api.utils import remove_path
 from DashAI.back.core.enums.metrics import LevelEnum
 from DashAI.back.dependencies.database.models import (
     GlobalExplainer,
@@ -381,8 +382,12 @@ async def delete_run(
                     status_code=status.HTTP_404_NOT_FOUND, detail="Run not found"
                 )
             db.delete(run)
-            if run.status == RunStatus.FINISHED:
-                os.remove(run.run_path)
+            if (
+                run.status == RunStatus.FINISHED
+                and run.run_path
+                and os.path.exists(run.run_path)
+            ):
+                remove_path(run.run_path)
             db.commit()
             return Response(status_code=status.HTTP_204_NO_CONTENT)
         except exc.SQLAlchemyError as e:
@@ -391,7 +396,7 @@ async def delete_run(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail="Internal database error",
             ) from e
-        except OSError as e:
+        except (OSError, ValueError) as e:
             log.exception(e)
             raise HTTPException(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -735,30 +740,6 @@ def reset_run(run):
         setattr(run, "plot_importance_path", None)
 
 
-def remove_path(path):
-    """Removes a file or directory
-
-    Parameters
-    ----------
-    path : str
-        The path to the file or directory to remove.
-
-    Raises
-    ------
-    ValueError
-        Raised if the path is not a file, directory, or symbolic link.
-    """
-    import os
-    import shutil
-
-    if os.path.isfile(path) or os.path.islink(path):
-        os.remove(path)
-    elif os.path.isdir(path):
-        shutil.rmtree(path)
-    else:
-        raise ValueError("file {} is not a file or dir.".format(path))
-
-
 # ─── Fold metrics ─────────────────────────────────────────────
 
 

DashAI/back/api/utils.py

@@ -84,3 +84,27 @@ def create_random_state():
     import numpy as np
 
     return np.random.RandomState()
+
+
+def remove_path(path):
+    """Removes a file or directory
+
+    Parameters
+    ----------
+    path : str
+        The path to the file or directory to remove.
+
+    Raises
+    ------
+    ValueError
+        Raised if the path is not a file, directory, or symbolic link.
+    """
+    import os
+    import shutil
+
+    if os.path.isfile(path) or os.path.islink(path):
+        os.remove(path)
+    elif os.path.isdir(path):
+        shutil.rmtree(path)
+    else:
+        raise ValueError("file {} is not a file or dir.".format(path))

DashAI/back/job/model_job.py

@@ -4,6 +4,7 @@
 from kink import inject
 from sqlalchemy import exc
 
+from DashAI.back.api.utils import remove_path
 from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run
 from DashAI.back.dependencies.downloads.nested import missing_downloads
 from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy
@@ -187,6 +188,8 @@ def run(
                 self.report_progress(0.95, "Saving model")
                 try:
                     run_path = os.path.join(config["RUNS_PATH"], str(run.id))
+                    if os.path.exists(run_path):
+                        remove_path(run_path)
                     model.save(run_path)
                 except Exception as e:
                     log.exception(e)

tests/back/api/test_jobs.py

@@ -317,3 +317,116 @@ def test_job_with_wrong_run(client: TestClient):
     )
     assert response.status_code == 500, response.text
     assert response.status_code == 500, response.text
+
+
+def test_execute_job_with_stale_run_path(client: TestClient, model_session_id: int):
+    """A leftover file/dir at run_path (e.g. from a run that reused this id)
+    must not prevent training from finishing."""
+    container = client.app.container
+    session_factory = container["session_factory"]
+    config = container["config"]
+
+    with session_factory() as db:
+        run = Run(
+            model_session_id=model_session_id,
+            model_name="DummyModel",
+            parameters={},
+            optimizer_name="",
+            optimizer_parameters={
+                "n_trials": 10,
+                "sampler": "TPESampler",
+                "pruner": "None",
+            },
+            goal_metric="",
+            name="DummyRunStalePath",
+        )
+        db.add(run)
+        db.commit()
+        db.refresh(run)
+        run_id = run.id
+
+    # Simulate a directory left behind by a previously-deleted run that
+    # happened to reuse this same id.
+    stale_run_path = os.path.join(config["RUNS_PATH"], str(run_id))
+    os.makedirs(stale_run_path, exist_ok=True)
+    with open(os.path.join(stale_run_path, "leftover.txt"), "w") as f:
+        f.write("stale")
+
+    response = client.post(
+        "/api/v1/job/",
+        data={"job_type": "ModelJob", "kwargs": json.dumps({"run_id": run_id})},
+    )
+    assert response.status_code == 201, response.text
+    job_id = response.json()["id"]
+
+    response = client.get(f"/api/v1/job/status/{job_id}")
+    assert response.status_code == 200, response.text
+    assert response.json()["status"] == "finished", response.json()
+
+    response = client.get(f"/api/v1/run/{run_id}")
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert data["status"] == 3
+    assert os.path.isfile(data["run_path"])
+
+    client.delete(f"/api/v1/run/{run_id}")
+
+
+def test_delete_model_session_cleans_up_run_files(
+    client: TestClient, dataset_id: int, test_registry, tmp_path
+):
+    """Deleting a ModelSession cascades its Runs in the DB, but the
+    on-disk run_path/plot files must also be cleaned up, not orphaned."""
+    container = client.app.container
+    session_factory = container["session_factory"]
+
+    run_dir = tmp_path / "cleanup_run"
+    run_dir.mkdir()
+    plot_file = tmp_path / "plot.png"
+    plot_file.write_text("x")
+
+    with session_factory() as db:
+        model_session = ModelSession(
+            dataset_id=dataset_id,
+            name="CleanupSession",
+            task_name="DummyTask",
+            input_columns=["SepalLengthCm"],
+            output_columns=["Species"],
+            train_metrics=[],
+            validation_metrics=[],
+            test_metrics=[],
+            evaluation_strategy="HoldoutEvaluationStrategy",
+            splits=json.dumps(
+                {
+                    "train": 0.5,
+                    "test": 0.2,
+                    "validation": 0.3,
+                    "seed": 42,
+                    "shuffle": True,
+                    "stratify": False,
+                }
+            ),
+        )
+        db.add(model_session)
+        db.commit()
+        db.refresh(model_session)
+        ms_id = model_session.id
+
+        run = Run(
+            model_session_id=ms_id,
+            model_name="DummyModel",
+            parameters={},
+            optimizer_name="",
+            optimizer_parameters={},
+            goal_metric="",
+            name="CleanupRun",
+            run_path=str(run_dir),
+            plot_history_path=str(plot_file),
+        )
+        db.add(run)
+        db.commit()
+
+    response = client.delete(f"/api/v1/model-session/{ms_id}")
+    assert response.status_code == 204, response.text
+    assert not run_dir.exists()
+    assert not plot_file.exists()

tests/back/api/test_runs_api.py

@@ -3,7 +3,7 @@
 import pytest
 from fastapi.testclient import TestClient
 
-from DashAI.back.dependencies.database.models import Dataset
+from DashAI.back.dependencies.database.models import Dataset, Run, RunStatus
 
 
 @pytest.fixture(scope="module", name="dataset_id")
@@ -276,6 +276,39 @@ def test_modify_run_model(client: TestClient):
     assert response.status_code == 304
 
 
+def test_delete_run_with_directory_run_path(
+    client: TestClient, model_session_id: int, tmp_path
+):
+    """A FINISHED run whose run_path is a directory (Hugging Face-style
+    models) must be deletable, not raise IsADirectoryError/PermissionError."""
+    run_dir = tmp_path / "hf_style_run_dir"
+    run_dir.mkdir()
+    (run_dir / "config.json").write_text("{}")
+
+    container = client.app.container
+    session_factory = container["session_factory"]
+    with session_factory() as db:
+        run = Run(
+            model_session_id=model_session_id,
+            model_name="SomeHFModel",
+            parameters={},
+            optimizer_name="",
+            optimizer_parameters={},
+            goal_metric="",
+            name="DirRun",
+            status=RunStatus.FINISHED,
+            run_path=str(run_dir),
+        )
+        db.add(run)
+        db.commit()
+        db.refresh(run)
+        run_id = run.id
+
+    response = client.delete(f"/api/v1/run/{run_id}")
+    assert response.status_code == 204, response.text
+    assert not run_dir.exists()
+
+
 @pytest.mark.order(-1)
 def test_delete_run(client: TestClient):
     # Delete all the runs in the db