use orjson for faster JSON parsing (#253)

Co-authored-by: Abubakar Abid <abubakar@huggingface.co>

Author: AkiSakurai

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "gradio>=5.46.0,<6.0.0",
     "numpy<3.0.0",
     "pillow<12.0.0",
+    "orjson>=3.0,<4.0.0"
 ]
 classifiers = [
     "Programming Language :: Python :: 3",

tests/test_sqlite_storage.py

@@ -6,6 +6,7 @@
 import tempfile
 import time
 
+import orjson
 import pytest
 
 from trackio.sqlite_storage import SQLiteStorage
@@ -203,8 +204,6 @@ def test_config_storage_in_database(temp_dir):
 
 def test_old_database_without_configs_table(temp_dir):
     # To make sure that we can continue to work with projects created with older versions of Trackio.
-    import json
-
     db_path = SQLiteStorage.get_project_db_path("test")
     db_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -220,7 +219,7 @@ def test_old_database_without_configs_table(temp_dir):
         """)
         conn.execute(
             "INSERT INTO metrics (timestamp, run_name, step, metrics) VALUES (?, ?, ?, ?)",
-            ("2024-01-01", "test_run", 0, json.dumps({"loss": 0.5})),
+            ("2024-01-01", "test_run", 0, orjson.dumps({"loss": 0.5})),
         )
 
     config = SQLiteStorage.get_run_config("test", "test_run")

trackio/sqlite_storage.py

@@ -1,4 +1,3 @@
-import json
 import os
 import platform
 import sqlite3
@@ -13,6 +12,7 @@
     fcntl = None
 
 import huggingface_hub as hf
+import orjson
 import pandas as pd
 
 try:  # absolute imports when installed from PyPI
@@ -166,7 +166,7 @@ def export_to_parquet():
                 metrics = df["metrics"].copy()
                 metrics = pd.DataFrame(
                     metrics.apply(
-                        lambda x: deserialize_values(json.loads(x))
+                        lambda x: deserialize_values(orjson.loads(x))
                     ).values.tolist(),
                     index=df.index,
                 )
@@ -196,9 +196,9 @@ def import_from_parquet():
                     for col in other_cols:
                         del metrics[col]
                     # combine them all into a single metrics col
-                    metrics = json.loads(metrics.to_json(orient="records"))
+                    metrics = orjson.loads(metrics.to_json(orient="records"))
                     df["metrics"] = [
-                        json.dumps(serialize_values(row)) for row in metrics
+                        orjson.dumps(serialize_values(row)) for row in metrics
                     ]
                 df.to_sql("metrics", conn, if_exists="replace", index=False)
 
@@ -273,7 +273,7 @@ def log(project: str, run: str, metrics: dict, step: int | None = None):
                         current_timestamp,
                         run,
                         current_step,
-                        json.dumps(serialize_values(metrics)),
+                        orjson.dumps(serialize_values(metrics)),
                     ),
                 )
                 conn.commit()
@@ -335,7 +335,7 @@ def bulk_log(
                             timestamps[i],
                             run,
                             steps[i],
-                            json.dumps(serialize_values(metrics)),
+                            orjson.dumps(serialize_values(metrics)),
                         )
                     )
 
@@ -356,7 +356,11 @@ def bulk_log(
                         (run_name, config, created_at)
                         VALUES (?, ?, ?)
                         """,
-                        (run, json.dumps(serialize_values(config)), current_timestamp),
+                        (
+                            run,
+                            orjson.dumps(serialize_values(config)),
+                            current_timestamp,
+                        ),
                     )
 
                 conn.commit()
@@ -383,7 +387,7 @@ def get_logs(project: str, run: str) -> list[dict]:
             rows = cursor.fetchall()
             results = []
             for row in rows:
-                metrics = json.loads(row["metrics"])
+                metrics = orjson.loads(row["metrics"])
                 metrics = deserialize_values(metrics)
                 metrics["timestamp"] = row["timestamp"]
                 metrics["step"] = row["step"]
@@ -490,7 +494,7 @@ def store_config(project: str, run: str, config: dict) -> None:
                     (run_name, config, created_at)
                     VALUES (?, ?, ?)
                     """,
-                    (run, json.dumps(serialize_values(config)), current_timestamp),
+                    (run, orjson.dumps(serialize_values(config)), current_timestamp),
                 )
                 conn.commit()
 
@@ -513,7 +517,7 @@ def get_run_config(project: str, run: str) -> dict | None:
 
                 row = cursor.fetchone()
                 if row:
-                    config = json.loads(row["config"])
+                    config = orjson.loads(row["config"])
                     return deserialize_values(config)
                 return None
             except sqlite3.OperationalError as e:
@@ -557,7 +561,7 @@ def get_all_run_configs(project: str) -> dict[str, dict]:
 
                 results = {}
                 for row in cursor.fetchall():
-                    config = json.loads(row["config"])
+                    config = orjson.loads(row["config"])
                     results[row["run_name"]] = deserialize_values(config)
                 return results
             except sqlite3.OperationalError as e: