feat: add ability to rename runs (#428)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
Co-authored-by: Abubakar Abid <abubakar@huggingface.co>
Co-authored-by: Abubakar Abid <islamrealm@gmail.com>

Author: 3 people

.changeset/huge-meals-rescue.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:feat: add ability to rename runs

pyproject.toml

@@ -8,6 +8,7 @@ description = "A lightweight, local-first, and free experiment tracking library
 authors = [
     { name = "Abubakar Abid", email = "abubakar@huggingface.co" },
     { name = "Zach Nation", email = "zach@huggingface.co" },
+    { name = "Saba Noorassa", email = "saba@huggingface.co" },
 ]
 readme = "README.md"
 requires-python = ">=3.10"

tests/conftest.py

@@ -24,7 +24,12 @@ def temp_dir(monkeypatch):
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
         for name in ["trackio.sqlite_storage"]:
             monkeypatch.setattr(f"{name}.TRACKIO_DIR", Path(tmpdir))
-        for name in ["trackio.media.media", "trackio.media.utils", "trackio.utils"]:
+        for name in [
+            "trackio.media.media",
+            "trackio.media.utils",
+            "trackio.utils",
+            "trackio.sqlite_storage",
+        ]:
             monkeypatch.setattr(f"{name}.MEDIA_DIR", Path(tmpdir) / "media")
         context_vars.current_run.set(None)
         context_vars.current_project.set(None)

tests/e2e-local/test_api.py

@@ -101,3 +101,65 @@ def test_move_run(temp_dir, image_ndarray):
 
     target_config = SQLiteStorage.get_run_config(project=target_project, run=run_name)
     assert target_config is not None
+
+
+def test_rename_run(temp_dir, image_ndarray):
+    project = "test_rename_project"
+    old_name = "old_run_name"
+    new_name = "new_run_name"
+
+    trackio.init(project=project, name=old_name)
+
+    image1 = trackio.Image(image_ndarray, caption="test_image_1")
+    image2 = trackio.Image(image_ndarray, caption="test_image_2")
+
+    trackio.log(metrics={"loss": 0.1, "acc": 0.9, "img1": image1})
+    trackio.log(metrics={"loss": 0.2, "acc": 0.95, "img2": image2})
+    trackio.finish()
+
+    old_logs = SQLiteStorage.get_logs(project=project, run=old_name)
+    assert len(old_logs) == 2
+    assert old_logs[0]["loss"] == 0.1
+    assert old_logs[1]["loss"] == 0.2
+
+    image1_path = old_logs[0]["img1"].get("file_path")
+    assert image1_path is not None
+    normalized_path = str(image1_path).replace("\\", "/")
+    assert normalized_path.startswith(f"{project}/{old_name}/")
+
+    api = Api()
+    runs = api.runs(project)
+    run = runs[0]
+    assert run.name == old_name
+
+    result = run.rename(new_name)
+    assert result is run
+    assert run.name == new_name
+
+    new_logs = SQLiteStorage.get_logs(project=project, run=new_name)
+    assert len(new_logs) == 2
+    assert new_logs[0]["loss"] == 0.1
+    assert new_logs[1]["loss"] == 0.2
+
+    new_image1_path = new_logs[0]["img1"].get("file_path")
+    assert new_image1_path is not None
+    normalized_new_path1 = str(new_image1_path).replace("\\", "/")
+    assert normalized_new_path1.startswith(f"{project}/{new_name}/")
+
+    new_image2_path = new_logs[1]["img2"].get("file_path")
+    assert new_image2_path is not None
+    normalized_new_path2 = str(new_image2_path).replace("\\", "/")
+    assert normalized_new_path2.startswith(f"{project}/{new_name}/")
+
+    old_logs_after = SQLiteStorage.get_logs(project=project, run=old_name)
+    assert len(old_logs_after) == 0
+
+    runs_after = SQLiteStorage.get_runs(project=project)
+    assert old_name not in runs_after
+    assert new_name in runs_after
+
+    old_config_after = SQLiteStorage.get_run_config(project=project, run=old_name)
+    assert old_config_after is None
+
+    new_config = SQLiteStorage.get_run_config(project=project, run=new_name)
+    assert new_config is not None

tests/unit/test_sqlite_storage.py

@@ -274,3 +274,126 @@ def test_get_runs_returns_chronological_order(temp_dir):
 
     runs = SQLiteStorage.get_runs("proj")
     assert runs == ["run-z", "run-a", "run-m"]
+
+
+def test_rename_run(temp_dir):
+    project = "test_project"
+    old_name = "old_run"
+    new_name = "new_run"
+
+    config = {"param1": "value1", "_Created": "2023-01-01T00:00:00"}
+    metrics = [{"accuracy": 0.95, "loss": 0.1}]
+    SQLiteStorage.bulk_log(project, old_name, metrics, config=config)
+
+    assert SQLiteStorage.get_run_config(project, old_name) is not None
+    assert len(SQLiteStorage.get_logs(project, old_name)) > 0
+
+    SQLiteStorage.rename_run(project, old_name, new_name)
+
+    assert SQLiteStorage.get_run_config(project, old_name) is None
+    assert len(SQLiteStorage.get_logs(project, old_name)) == 0
+
+    assert SQLiteStorage.get_run_config(project, new_name) is not None
+    assert len(SQLiteStorage.get_logs(project, new_name)) > 0
+
+    new_logs = SQLiteStorage.get_logs(project, new_name)
+    assert new_logs[0]["accuracy"] == 0.95
+    assert new_logs[0]["loss"] == 0.1
+
+
+def test_rename_run_duplicate_name(temp_dir):
+    project = "test_project"
+    run1 = "run1"
+    run2 = "run2"
+
+    SQLiteStorage.bulk_log(project, run1, [{"a": 1}])
+    SQLiteStorage.bulk_log(project, run2, [{"b": 2}])
+
+    with pytest.raises(ValueError, match="already exists"):
+        SQLiteStorage.rename_run(project, run1, run2)
+
+    assert len(SQLiteStorage.get_logs(project, run1)) > 0
+    assert len(SQLiteStorage.get_logs(project, run2)) > 0
+
+
+def test_rename_run_with_media(temp_dir):
+    from trackio.utils import MEDIA_DIR
+
+    project = "test_project"
+    old_name = "old_run"
+    new_name = "new_run"
+
+    media_dir = MEDIA_DIR / project / old_name
+    media_dir.mkdir(parents=True, exist_ok=True)
+    test_file = media_dir / "test.txt"
+    test_file.write_text("test content")
+
+    metrics = [
+        {
+            "image": {
+                "_type": "trackio.image",
+                "file_path": f"{project}/{old_name}/test.txt",
+                "caption": "test",
+            }
+        }
+    ]
+    SQLiteStorage.bulk_log(project, old_name, metrics)
+
+    SQLiteStorage.rename_run(project, old_name, new_name)
+
+    new_media_dir = MEDIA_DIR / project / new_name
+    assert new_media_dir.exists()
+    assert (new_media_dir / "test.txt").exists()
+
+    old_media_dir = MEDIA_DIR / project / old_name
+    assert not old_media_dir.exists()
+
+    new_logs = SQLiteStorage.get_logs(project, new_name)
+    assert len(new_logs) > 0
+    assert "image" in new_logs[0]
+    assert new_logs[0]["image"]["file_path"].startswith(f"{project}/{new_name}/")
+
+
+def test_rename_run_nonexistent(temp_dir):
+    project = "test_project"
+    old_name = "nonexistent_run"
+    new_name = "new_run"
+
+    with pytest.raises(ValueError, match="does not exist"):
+        SQLiteStorage.rename_run(project, old_name, new_name)
+
+
+def test_rename_run_empty_name(temp_dir):
+    project = "test_project"
+    old_name = "old_run"
+
+    SQLiteStorage.bulk_log(project, old_name, [{"a": 1}])
+
+    with pytest.raises(ValueError, match="cannot be empty"):
+        SQLiteStorage.rename_run(project, old_name, "")
+
+    with pytest.raises(ValueError, match="cannot be empty"):
+        SQLiteStorage.rename_run(project, old_name, "   ")
+
+    assert len(SQLiteStorage.get_logs(project, old_name)) > 0
+
+
+def test_rename_run_with_system_metrics(temp_dir):
+    project = "test_project"
+    old_name = "old_run"
+    new_name = "new_run"
+
+    metrics = [{"accuracy": 0.95}]
+    SQLiteStorage.bulk_log(project, old_name, metrics)
+
+    system_metrics = [{"gpu_usage": 80.5}]
+    SQLiteStorage.bulk_log_system(project, old_name, system_metrics)
+
+    SQLiteStorage.rename_run(project, old_name, new_name)
+
+    assert len(SQLiteStorage.get_logs(project, new_name)) > 0
+    assert len(SQLiteStorage.get_system_logs(project, new_name)) > 0
+    assert len(SQLiteStorage.get_system_logs(project, old_name)) == 0
+
+    new_system_logs = SQLiteStorage.get_system_logs(project, new_name)
+    assert new_system_logs[0]["gpu_usage"] == 80.5

trackio/api.py

@@ -28,6 +28,11 @@ def move(self, new_project: str) -> bool:
             self.project = new_project
         return success
 
+    def rename(self, new_name: str) -> "Run":
+        SQLiteStorage.rename_run(self.project, self.name, new_name)
+        self.name = new_name
+        return self
+
     def __repr__(self) -> str:
         return f"<Run {self.name} in project {self.project}>"