Add trackio.save() (#351)

Co-authored-by: Abubakar <abubakar@Abubakars-MacBook-Pro.local>
Co-authored-by: Abubakar Abid <aaabid93@gmail.com>

Author: 3 people

.changeset/sweet-suns-know.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:Add `trackio.save()`

docs/source/api.md

@@ -16,6 +16,10 @@
 
 [[autodoc]] finish
 
+## save
+
+[[autodoc]] save
+
 ## show
 
 [[autodoc]] show

examples/files/config1.yml

@@ -0,0 +1,17 @@
+model:
+  name: "simple_mlp"
+  layers: [128, 64, 32]
+  activation: "relu"
+  dropout: 0.2
+
+training:
+  epochs: 10
+  batch_size: 32
+  learning_rate: 0.001
+  optimizer: "adam"
+
+data:
+  dataset: "mnist"
+  train_split: 0.8
+  normalize: true
+

examples/files/config2.yml

@@ -0,0 +1,18 @@
+model:
+  name: "cnn"
+  layers: [32, 64, 128]
+  activation: "relu"
+  dropout: 0.3
+
+training:
+  epochs: 15
+  batch_size: 64
+  learning_rate: 0.0005
+  optimizer: "sgd"
+
+data:
+  dataset: "cifar10"
+  train_split: 0.75
+  normalize: true
+  augmentation: true
+

examples/files/models/model1.pth

@@ -0,0 +1,5 @@
+FAKE_MODEL_FILE_1
+This is a fake PyTorch model checkpoint file.
+In reality, this would contain binary model weights.
+For testing purposes, this is just a text placeholder.
+

examples/files/models/model2.pth

@@ -0,0 +1,5 @@
+FAKE_MODEL_FILE_2
+This is a fake PyTorch model checkpoint file.
+In reality, this would contain binary model weights.
+For testing purposes, this is just a text placeholder.
+

examples/files/training-and-saving-files.py

@@ -0,0 +1,46 @@
+import random
+import time
+from pathlib import Path
+
+import trackio as wandb
+
+PROJECT_ID = random.randint(100000, 999999)
+EPOCHS = 10
+
+current_dir = Path(__file__).parent
+
+
+def main():
+    wandb.init(
+        project=f"test-save-{PROJECT_ID}",
+        name="test-run",
+        config={
+            "epochs": EPOCHS,
+            "learning_rate": 0.001,
+            "batch_size": 32,
+        },
+    )
+
+    wandb.save(current_dir / "config1.yml")
+    wandb.save(current_dir / "config2.yml")
+    wandb.save(current_dir / "models/*.pth")
+
+    for epoch in range(EPOCHS):
+        loss = 2.0 * (1 - epoch / EPOCHS) + random.uniform(-0.1, 0.1)
+        accuracy = 0.5 + 0.4 * (epoch / EPOCHS) + random.uniform(-0.05, 0.05)
+
+        wandb.log(
+            {
+                "loss": round(loss, 4),
+                "accuracy": round(accuracy, 4),
+            }
+        )
+
+        time.sleep(0.1)
+
+    wandb.finish()
+    print(f"* Test completed. Check project 'test-save-{PROJECT_ID}' for saved files.")
+
+
+if __name__ == "__main__":
+    main()

examples/test-save.py

@@ -0,0 +1,43 @@
+import random
+import time
+
+import trackio as wandb
+
+PROJECT_ID = random.randint(100000, 999999)
+EPOCHS = 10
+
+
+def main():
+    wandb.init(
+        project=f"test-save-{PROJECT_ID}",
+        name="test-run",
+        config={
+            "epochs": EPOCHS,
+            "learning_rate": 0.001,
+            "batch_size": 32,
+        },
+    )
+
+    wandb.save("files/config1.yml")
+    wandb.save("files/config2.yml")
+    wandb.save("files/models/*.pth")
+
+    for epoch in range(EPOCHS):
+        loss = 2.0 * (1 - epoch / EPOCHS) + random.uniform(-0.1, 0.1)
+        accuracy = 0.5 + 0.4 * (epoch / EPOCHS) + random.uniform(-0.05, 0.05)
+
+        wandb.log(
+            {
+                "loss": round(loss, 4),
+                "accuracy": round(accuracy, 4),
+            }
+        )
+
+        time.sleep(0.1)
+
+    wandb.finish()
+    print(f"* Test completed. Check project 'test-save-{PROJECT_ID}' for saved files.")
+
+
+if __name__ == "__main__":
+    main()

tests/conftest.py

@@ -14,7 +14,7 @@ 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.file_storage"]:
+        for name in ["trackio.media.media", "trackio.media.utils"]:
             monkeypatch.setattr(f"{name}.MEDIA_DIR", Path(tmpdir) / "media")
         yield tmpdir