Auto-prune log files older than LOG_MAX_DAYS (default 7) after each upload

Without pruning, parquet files accumulate indefinitely and risk exceeding
the HF private storage free tier. After each successful upload, old files
are deleted from the dataset repo using a date-prefix comparison on the
filename. Refactors log_inference into single-responsibility helpers
(_img_to_jpeg, _build_table, _upload_parquet, _prune_old_files) to make
the prune logic unit-testable without real network calls or waiting.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

logging_utils.py

@@ -1,10 +1,118 @@
 import os
 import uuid
 from io import BytesIO
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
 
 HF_TOKEN     = os.environ.get("HF_TOKEN")
 DATASET_REPO = os.environ.get("LOG_DATASET_REPO")
+MAX_LOG_DAYS = int(os.environ.get("LOG_MAX_DAYS", "7"))
+
+
+def _img_to_jpeg(img, quality=85):
+    if img is None:
+        return None
+    buf = BytesIO()
+    img.convert("RGB").save(buf, format="JPEG", quality=quality)
+    return buf.getvalue()
+
+
+def _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
+                 input_width, input_height, duration_seconds, success, error_message, now):
+    import json as _json
+    import pyarrow as pa
+
+    img_struct = pa.struct([("bytes", pa.binary()), ("path", pa.string())])
+    hf_meta = _json.dumps({"info": {"features": {
+        "timestamp":        {"dtype": "float64", "_type": "Value"},
+        "prompt":           {"dtype": "string",  "_type": "Value"},
+        "seed":             {"dtype": "int32",   "_type": "Value"},
+        "steps":            {"dtype": "int32",   "_type": "Value"},
+        "guidance_scale":   {"dtype": "float32", "_type": "Value"},
+        "input_images":     {"feature": {"_type": "Image"}, "_type": "Sequence"},
+        "output_image":     {"_type": "Image"},
+        "duration_seconds": {"dtype": "float32", "_type": "Value"},
+        "input_width":      {"dtype": "int32",   "_type": "Value"},
+        "input_height":     {"dtype": "int32",   "_type": "Value"},
+        "success":          {"dtype": "bool",    "_type": "Value"},
+        "error_message":    {"dtype": "string",  "_type": "Value"},
+    }}}).encode()
+    schema = pa.schema([
+        ("timestamp",        pa.float64()),
+        ("prompt",           pa.string()),
+        ("seed",             pa.int32()),
+        ("steps",            pa.int32()),
+        ("guidance_scale",   pa.float32()),
+        ("input_images",     pa.list_(img_struct)),
+        ("output_image",     img_struct),
+        ("duration_seconds", pa.float32()),
+        ("input_width",      pa.int32()),
+        ("input_height",     pa.int32()),
+        ("success",          pa.bool_()),
+        ("error_message",    pa.string()),
+    ], metadata={b"huggingface": hf_meta})
+
+    def _img(b):
+        return {"bytes": b, "path": None}
+
+    input_jpegs = [_img_to_jpeg(img) for img in pil_inputs]
+    output_jpeg = _img_to_jpeg(output_pil)
+
+    return pa.table({
+        "timestamp":        pa.array([now.timestamp()],                        type=pa.float64()),
+        "prompt":           pa.array([prompt],                                 type=pa.string()),
+        "seed":             pa.array([int(seed)],                              type=pa.int32()),
+        "steps":            pa.array([int(steps)],                             type=pa.int32()),
+        "guidance_scale":   pa.array([float(guidance_scale)],                  type=pa.float32()),
+        "input_images":     pa.array([[_img(b) for b in input_jpegs]],         type=pa.list_(img_struct)),
+        "output_image":     pa.array([_img(output_jpeg) if output_jpeg else None], type=img_struct),
+        "duration_seconds": pa.array([float(duration_seconds)],                type=pa.float32()),
+        "input_width":      pa.array([int(input_width)],                       type=pa.int32()),
+        "input_height":     pa.array([int(input_height)],                      type=pa.int32()),
+        "success":          pa.array([bool(success)],                          type=pa.bool_()),
+        "error_message":    pa.array([str(error_message)],                     type=pa.string()),
+    }, schema=schema)
+
+
+def _upload_parquet(api, repo_id, table, path_in_repo):
+    import tempfile
+    import pyarrow.parquet as pq
+
+    tmp_path = None
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+            tmp_path = tmp.name
+        pq.write_table(table, tmp_path)
+        print(f"[log] uploading {path_in_repo} ({os.path.getsize(tmp_path)//1024}KB)")
+        api.upload_file(
+            path_or_fileobj=tmp_path, path_in_repo=path_in_repo,
+            repo_id=repo_id, repo_type="dataset",
+        )
+        print(f"[log] upload done — {repo_id}/{path_in_repo}")
+    finally:
+        if tmp_path:
+            try:
+                os.unlink(tmp_path)
+            except Exception:
+                pass
+
+
+def _prune_old_files(api, repo_id, keep_days):
+    if keep_days <= 0:
+        return
+    cutoff = (datetime.now(timezone.utc) - timedelta(days=keep_days)).strftime("%Y-%m-%d")
+    try:
+        to_delete = [
+            f.path
+            for f in api.list_repo_tree(repo_id, repo_type="dataset", path_in_repo="data")
+            if f.path.endswith(".parquet") and os.path.basename(f.path)[:10] < cutoff
+        ]
+        for path in to_delete:
+            api.delete_file(path_in_repo=path, repo_id=repo_id, repo_type="dataset")
+            print(f"[log] pruned: {path}")
+        if to_delete:
+            print(f"[log] pruned {len(to_delete)} old file(s)")
+    except Exception as e:
+        print(f"[log] prune warning: {e}")
 
 
 def log_inference(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
@@ -12,94 +120,20 @@ def log_inference(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
     if not HF_TOKEN or not DATASET_REPO:
         print(f"[log] skipped — HF_TOKEN={'set' if HF_TOKEN else 'missing'}, DATASET_REPO={'set' if DATASET_REPO else 'missing'}")
         return
-    tmp_path = None
     try:
-        import tempfile
-        import json as _json
-        import pyarrow as pa
-        import pyarrow.parquet as pq
         from huggingface_hub import HfApi
 
-        img_struct = pa.struct([("bytes", pa.binary()), ("path", pa.string())])
-        hf_meta = _json.dumps({"info": {"features": {
-            "timestamp":        {"dtype": "float64", "_type": "Value"},
-            "prompt":           {"dtype": "string",  "_type": "Value"},
-            "seed":             {"dtype": "int32",   "_type": "Value"},
-            "steps":            {"dtype": "int32",   "_type": "Value"},
-            "guidance_scale":   {"dtype": "float32", "_type": "Value"},
-            "input_images":     {"feature": {"_type": "Image"}, "_type": "Sequence"},
-            "output_image":     {"_type": "Image"},
-            "duration_seconds": {"dtype": "float32", "_type": "Value"},
-            "input_width":      {"dtype": "int32",   "_type": "Value"},
-            "input_height":     {"dtype": "int32",   "_type": "Value"},
-            "success":          {"dtype": "bool",    "_type": "Value"},
-            "error_message":    {"dtype": "string",  "_type": "Value"},
-        }}}).encode()
-        schema = pa.schema([
-            ("timestamp",        pa.float64()),
-            ("prompt",           pa.string()),
-            ("seed",             pa.int32()),
-            ("steps",            pa.int32()),
-            ("guidance_scale",   pa.float32()),
-            ("input_images",     pa.list_(img_struct)),
-            ("output_image",     img_struct),
-            ("duration_seconds", pa.float32()),
-            ("input_width",      pa.int32()),
-            ("input_height",     pa.int32()),
-            ("success",          pa.bool_()),
-            ("error_message",    pa.string()),
-        ], metadata={b"huggingface": hf_meta})
-
-        def _to_jpeg(img, quality=85):
-            if img is None:
-                return None
-            buf = BytesIO()
-            img.convert("RGB").save(buf, format="JPEG", quality=quality)
-            return buf.getvalue()
-
-        def _img(b):
-            return {"bytes": b, "path": None}
-
         now = datetime.now(timezone.utc)
-        input_jpegs = [_to_jpeg(img) for img in pil_inputs]
-        output_jpeg = _to_jpeg(output_pil)
-
-        table = pa.table({
-            "timestamp":        pa.array([now.timestamp()],                    type=pa.float64()),
-            "prompt":           pa.array([prompt],                             type=pa.string()),
-            "seed":             pa.array([int(seed)],                          type=pa.int32()),
-            "steps":            pa.array([int(steps)],                         type=pa.int32()),
-            "guidance_scale":   pa.array([float(guidance_scale)],              type=pa.float32()),
-            "input_images":     pa.array([[_img(b) for b in input_jpegs]],     type=pa.list_(img_struct)),
-            "output_image":     pa.array([_img(output_jpeg) if output_jpeg else None], type=img_struct),
-            "duration_seconds": pa.array([float(duration_seconds)],            type=pa.float32()),
-            "input_width":      pa.array([int(input_width)],                   type=pa.int32()),
-            "input_height":     pa.array([int(input_height)],                  type=pa.int32()),
-            "success":          pa.array([bool(success)],                      type=pa.bool_()),
-            "error_message":    pa.array([str(error_message)],                 type=pa.string()),
-        }, schema=schema)
+        table = _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
+                             input_width, input_height, duration_seconds, success, error_message, now)
 
         uid = uuid.uuid4().hex[:8]
         path_in_repo = f"data/{now.strftime('%Y-%m-%d-%H%M%S')}-{uid}.parquet"
 
-        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
-            tmp_path = tmp.name
-        pq.write_table(table, tmp_path)
-
         api = HfApi(token=HF_TOKEN)
         api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", private=True, exist_ok=True)
-        print(f"[log] uploading {path_in_repo} ({os.path.getsize(tmp_path)//1024}KB)")
-        api.upload_file(
-            path_or_fileobj=tmp_path, path_in_repo=path_in_repo,
-            repo_id=DATASET_REPO, repo_type="dataset",
-        )
-        print(f"[log] upload done — {DATASET_REPO}/{path_in_repo}")
+        _upload_parquet(api, DATASET_REPO, table, path_in_repo)
+        _prune_old_files(api, DATASET_REPO, MAX_LOG_DAYS)
     except Exception as log_err:
         import traceback as _tb
         print(f"[log] WARNING: {log_err}\n{_tb.format_exc()}")
-    finally:
-        if tmp_path:
-            try:
-                os.unlink(tmp_path)
-            except Exception:
-                pass

tests/test_logging_utils.py

@@ -0,0 +1,68 @@
+from datetime import datetime, timezone, timedelta
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+from logging_utils import _prune_old_files
+
+
+def _make_api(paths):
+    """Return a mock HfApi whose list_repo_tree yields RepoFile-like objects for `paths`."""
+    api = MagicMock()
+    api.list_repo_tree.return_value = [SimpleNamespace(path=p) for p in paths]
+    return api
+
+
+def _old(name):
+    return f"data/2020-01-01-000000-{name}.parquet"
+
+
+def _recent(name):
+    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
+    return f"data/{today}-000000-{name}.parquet"
+
+
+def test_old_files_are_deleted():
+    api = _make_api([_old("aaa"), _old("bbb")])
+    _prune_old_files(api, "org/repo", keep_days=7)
+    assert api.delete_file.call_count == 2
+    deleted = {c.kwargs["path_in_repo"] for c in api.delete_file.call_args_list}
+    assert deleted == {_old("aaa"), _old("bbb")}
+
+
+def test_recent_files_are_kept():
+    api = _make_api([_recent("aaa"), _recent("bbb")])
+    _prune_old_files(api, "org/repo", keep_days=7)
+    api.delete_file.assert_not_called()
+
+
+def test_only_old_files_deleted_in_mixed_set():
+    api = _make_api([_old("aaa"), _recent("bbb"), _old("ccc")])
+    _prune_old_files(api, "org/repo", keep_days=7)
+    deleted = {c.kwargs["path_in_repo"] for c in api.delete_file.call_args_list}
+    assert deleted == {_old("aaa"), _old("ccc")}
+
+
+def test_empty_directory_does_nothing():
+    api = _make_api([])
+    _prune_old_files(api, "org/repo", keep_days=7)
+    api.delete_file.assert_not_called()
+
+
+def test_keep_days_zero_skips_pruning():
+    api = _make_api([_old("aaa")])
+    _prune_old_files(api, "org/repo", keep_days=0)
+    api.list_repo_tree.assert_not_called()
+    api.delete_file.assert_not_called()
+
+
+def test_non_parquet_files_are_ignored():
+    api = _make_api(["data/2020-01-01-README.md", _old("aaa")])
+    _prune_old_files(api, "org/repo", keep_days=7)
+    deleted = {c.kwargs["path_in_repo"] for c in api.delete_file.call_args_list}
+    assert deleted == {_old("aaa")}
+
+
+def test_api_error_does_not_raise():
+    api = MagicMock()
+    api.list_repo_tree.side_effect = RuntimeError("network error")
+    _prune_old_files(api, "org/repo", keep_days=7)  # must not raise