Merge pull request #308 from OpenPipe/fix-logs

Remove failing step checks

Author: JonesAndrew

src/art/local/backend.py

@@ -1,5 +1,6 @@
 import json
 import math
+from datetime import datetime
 
 from art.utils.deploy_model import (
     LoRADeploymentJob,
@@ -241,7 +242,7 @@ async def _delete_checkpoints(
             )
             steps_to_keep.append(best_step)
         except FileNotFoundError:
-            pass
+            print(f'"{output_dir}/history.jsonl" not found')
         except pl.exceptions.ColumnNotFoundError:
             print(f'No "{benchmark}" metric found in history')
         delete_checkpoints(output_dir, steps_to_keep)
@@ -273,7 +274,7 @@ async def _log(
         os.makedirs(parent_dir, exist_ok=True)
 
         # Get the file name for the current iteration, or default to 0 for non-trainable models
-        iteration = self.__get_step(model) if isinstance(model, TrainableModel) else 0
+        iteration = self.__get_step(model)
         file_name = f"{iteration:04d}.jsonl"
 
         # Write the logs to the file
@@ -443,11 +444,20 @@ def _log_metrics(
         step: int | None = None,
     ) -> None:
         metrics = {f"{split}/{metric}": value for metric, value in metrics.items()}
-        step = (
-            step
-            if step is not None
-            else (self.__get_step(model) if isinstance(model, TrainableModel) else 0)
-        )
+        step = step if step is not None else self.__get_step(model)
+
+        with open(
+            f"{get_model_dir(model=model, art_path=self._path)}/history.jsonl", "a"
+        ) as f:
+            f.write(
+                json.dumps(
+                    {
+                        k: v for k, v in metrics.items() if v == v
+                    }  # Filter out NaN values
+                    | {"step": step, "recorded_at": datetime.now().isoformat()}
+                )
+                + "\n"
+            )
 
         # If we have a W&B run, log the data there
         if run := self._get_wandb_run(model):