Repository: mlflow/mlflow
Pull request: [BUG] add support for file removal for local artifact repo
<details><summary>&#x1F6E0 DevTools &#x1F6E0</summary>
<p>

[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/rzalawad/mlflow/pull/13005?quickstart=1)

#### Install mlflow from this PR

```
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/13005/merge
```

#### Checkout with GitHub CLI

```
gh pr checkout 13005
```

</p>
</details>

### Related Issues/PRs
Resolve https://github.com/mlflow/mlflow/issues/4185
<!-- Uncomment 'Resolve' if this PR can close the linked items. -->

### What changes are proposed in this pull request?

<!-- Please fill in changes proposed in this PR. -->

### How is this PR tested?

- [ ] Existing unit/integration tests
- [x] New unit/integration tests
- [ ] Manual tests

<!-- Attach code, screenshot, video used for manual testing here. -->
<img width="1710" alt="image" src="https://github.com/user-attachments/assets/5313083a-4aa8-4491-9335-f5b48150fc13">

### Does this PR require documentation update?

- [x] No. You can skip the rest of this section.
- [ ] Yes. I've updated:
  - [ ] Examples
  - [ ] API references
  - [ ] Instructions

### Release Notes
* Update local artifact repo to handle file deletions.

#### Is this a user-facing change?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Give a description of this change to be included in the release notes for MLflow users.

<!-- Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change. -->

#### What component(s), interfaces, languages, and integrations does this PR affect?
This PR affects the delete_artifacts feature of Local artifact repository. The user can now delete artifacts rather than just folders.

Components

- [x] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->

<a name="release-note-category"></a>

#### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [x] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes

#### Should this PR be included in the next patch release?

`Yes` should be selected for bug fixes, documentation updates, and other small changes. `No` should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

<details>
<summary>What is a minor/patch release?</summary>

- Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
  Bug fixes, doc updates and new features usually go into minor releases.
- Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
  Bug fixes and doc updates usually go into patch releases.

</details>

<!-- patch -->

- [x] Yes (this PR will be cherry-picked and included in the next patch release)
- [ ] No (this PR will be included in the next minor release)


Problem statement:
[FR] Add delete_artifact method to mlflow API

## Willingness to contribute
The MLflow Community encourages new feature contributions. Would you or another member of your organization be willing to contribute an implementation of this feature (either as an MLflow Plugin or an enhancement to the MLflow code base)?

- [ ] Yes. I can contribute this feature independently.
- [ x] Yes. I would be willing to contribute this feature with guidance from the MLflow community.
- [ ] No. I cannot contribute this feature at this time.

## Proposal Summary

We need to add a mlflow.delete_artifact method to the API as an analog to mlflow.log_artifact. 

## Motivation

If clients can add an artifact they should be able to delete it.

### What component(s), interfaces, languages, and integrations does this feature affect?
Components 
- [ x] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Hints from issue discussion:
@amesar Thanks for filing this. Yeah, this is an interesting enhancement as it poses a question about how long should we retain experiment runs and associated artifacts for auditing purposes. Currently, we just move the experiment into a _.trash_ folder under _mlruns_ and have `mlflow gc...` and mark the internal status of runs as DELETED and have it clean up in the OSS version, at least. 

This would be a good proposal to address this retention policy and document it as policy, as I understand. 

cc: area owners @arjundc-db @dbczumar 
Hey @dmatrix I'd like to make my first contribution if the issue is free
I've searched for the module but have not found it. Could you redirect me to the relevant code section as I am not familiar (yet) with MLflow

@dmatrix bump on this - any openness to this (I'm happy to implement if so)? Alternative as a user I'd also be fine with going directly to the aritfact repository - but as per for example https://github.com/mlflow/mlflow/blob/312f5f9f46f4ca4da6e0bf594bb7c23bbec4dbd2/mlflow/store/artifact/local_artifact_repo.py#LL114 it seems `delete_artifacts` is designed specifically to delete directories, not individual files. Would be open to changing this to support individual files too
Hi there, I think it would be great if the `delete_artifacts` method could just be changed so that it works with either files or directories. It seems like a simple enough fix.

Commit message: add file removal for local artifact repo and add test cases

Signed-off-by: Ridham Zalawadia <zalavadiaridham@gmail.com>

Patch to review:
diff --git a/mlflow/store/artifact/local_artifact_repo.py b/mlflow/store/artifact/local_artifact_repo.py
index 28112451bbfbe..b1ae5fb0e6439 100644
--- a/mlflow/store/artifact/local_artifact_repo.py
+++ b/mlflow/store/artifact/local_artifact_repo.py
@@ -122,7 +122,10 @@ def delete_artifacts(self, artifact_path=None):
         )
 
         if os.path.exists(artifact_path):
-            shutil.rmtree(artifact_path)
+            if os.path.isfile(artifact_path):
+                os.remove(artifact_path)
+            else:
+                shutil.rmtree(artifact_path)
 
     def download_trace_data(self) -> Dict[str, Any]:
         """
diff --git a/tests/store/artifact/test_local_artifact_repo.py b/tests/store/artifact/test_local_artifact_repo.py
index 4fa45a959ad39..491013d3661e6 100644
--- a/tests/store/artifact/test_local_artifact_repo.py
+++ b/tests/store/artifact/test_local_artifact_repo.py
@@ -181,7 +181,7 @@ def test_hidden_files_are_logged_correctly(local_artifact_repo):
             assert f.read() == "42"
 
 
-def test_delete_artifacts(local_artifact_repo):
+def test_delete_artifacts_folder(local_artifact_repo):
     with TempDir() as local_dir:
         os.mkdir(local_dir.path("subdir"))
         os.mkdir(local_dir.path("subdir", "nested"))
@@ -199,6 +199,31 @@ def test_delete_artifacts(local_artifact_repo):
         assert not os.path.exists(os.path.join(local_artifact_repo._artifact_dir))
 
 
+def test_delete_artifacts_files(local_artifact_repo):
+    with TempDir() as local_dir:
+        os.mkdir(local_dir.path("subdir"))
+        os.mkdir(local_dir.path("subdir", "nested"))
+        with open(local_dir.path("subdir", "a.txt"), "w") as f:
+            f.write("A")
+        with open(local_dir.path("subdir", "b.txt"), "w") as f:
+            f.write("B")
+        with open(local_dir.path("subdir", "nested", "c.txt"), "w") as f:
+            f.write("C")
+        local_artifact_repo.log_artifacts(local_dir.path("subdir"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "nested"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "a.txt"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "b.txt"))
+        local_artifact_repo.delete_artifacts(artifact_path="nested/c.txt")
+        local_artifact_repo.delete_artifacts(artifact_path="b.txt")
+        assert not os.path.exists(
+            os.path.join(local_artifact_repo._artifact_dir, "nested", "c.txt")
+        )
+        assert not os.path.exists(
+            os.path.join(local_artifact_repo._artifact_dir, "b.txt")
+        )
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "a.txt"))
+
+
 def test_delete_artifacts_with_nonexistent_path_succeeds(local_artifact_repo):
     local_artifact_repo.delete_artifacts("nonexistent")
 


Code review comments:
diff_hunk: @@ -199,6 +199,31 @@ def test_delete_artifacts(local_artifact_repo):
         assert not os.path.exists(os.path.join(local_artifact_repo._artifact_dir))
 
 
+def test_delete_artifacts_files(local_artifact_repo):
+    with TempDir() as local_dir:
+        os.mkdir(local_dir.path("subdir"))
+        os.mkdir(local_dir.path("subdir", "nested"))
+        with open(local_dir.path("subdir", "a.txt"), "w") as f:
+            f.write("A")
+        with open(local_dir.path("subdir", "b.txt"), "w") as f:
+            f.write("B")
+        with open(local_dir.path("subdir", "nested", "c.txt"), "w") as f:
+            f.write("C")
+        local_artifact_repo.log_artifacts(local_dir.path("subdir"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "nested"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "a.txt"))
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "b.txt"))
+        local_artifact_repo.delete_artifacts(artifact_path="nested/c.txt")
+        local_artifact_repo.delete_artifacts(artifact_path="b.txt")
+        assert not os.path.exists(
+            os.path.join(local_artifact_repo._artifact_dir, "nested", "c.txt")
+        )
+        assert not os.path.exists(
+            os.path.join(local_artifact_repo._artifact_dir, "b.txt")
+        )
+        assert os.path.exists(os.path.join(local_artifact_repo._artifact_dir, "a.txt"))
path: tests/store/artifact/test_local_artifact_repo.py
text: @user1:
/ai simplify the code using the `tmp_path` fixture.

@user2:
@user1

<!---->
Here is the simplified version of the above script using `tmp_path` fixture which is a part of the built-in pytest fixtures:

```
def test_delete_artifacts_files(local_artifact_repo, tmp_path):
    subdir = tmp_path / "subdir"
    nested = subdir / "nested"
    subdir.mkdir()
    nested.mkdir()

    (subdir / "a.txt").write_text("A")
    (subdir / "b.txt").write_text("B")
    (nested / "c.txt").write_text("C")

    local_artifact_repo.log_artifacts(str(subdir))
    
    assert (local_artifact_repo._artifact_dir / "nested").exists()
    assert (local_artifact_repo._artifact_dir / "a.txt").exists()
    assert (local_artifact_repo._artifact_dir / "b.txt").exists()

    local_artifact_repo.delete_artifacts(artifact_path="nested/c.txt")
    local_artifact_repo.delete_artifacts(artifact_path="b.txt")

    assert not (local_artifact_repo._artifact_dir / "nested" / "c.txt").exists()
    assert not (local_artifact_repo._artifact_dir / "b.txt").exists()
    assert (local_artifact_repo._artifact_dir / "a.txt").exists()
```
In this version of the code, we’re using `tmp_path` fixture which is a simple yet powerful way of dealing with temporary files and directories for tests and `pathlib` module for path manipulations which brings a more intuitive way to handle filesystem paths.
<!---->

<details><summary>Click here to see the usage data</summary>

`````json
{
  "prompt_tokens": 351,
  "completion_tokens": 314,
  "total_tokens": 665,
  "estimated_cost_in_usd": 0.02937
}
`````

</details>