Add stdlib json fallback for free-threading compatibility (#365)

Summary:
Pull Request resolved: #365

D97564498 introduced orjson as a hard dependency, which broke 40+ tests on
CPython 3.14 free-threading builds because orjson does not yet support
free-threading (PEP 703).

This diff adds a `_json_compat` wrapper module that:
- Uses orjson when available (for performance on standard CPython)
- Falls back to stdlib json when orjson is unavailable (e.g., free-threading builds)
- Provides a unified API: `loads()`, `dumps()`, and `JSONDecodeError`
- Returns `str` from `dumps()` (no `.decode()` needed at call sites)

All 18 files that previously imported orjson directly now import from
`tritonparse._json_compat` instead. Two files (`sourcemap_utils.py` and
`reproducer/utils.py`) intentionally keep stdlib json and are unchanged.

Also moves orjson from a required dependency to an optional extra in
`pyproject.toml` for OSS compatibility.

Reviewed By: xuzhao9

Differential Revision: D98194947

fbshipit-source-id: bd9ff863a3156514b346a417403da18d390a5e22

Author: FindHao

tritonparse/_json_compat.py

@@ -0,0 +1,87 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+JSON compatibility layer for tritonparse.
+
+Provides a unified interface that uses orjson when available (for performance)
+and falls back to stdlib json (for environments where orjson is unavailable,
+e.g. CPython 3.14 free-threading builds).
+
+``loads()`` accepts ``str | bytes | bytearray | memoryview`` inputs.
+``dumps()`` returns ``str``.
+"""
+
+try:
+    import orjson as _orjson
+
+    _HAS_ORJSON = True
+except ImportError:
+    import json as _json
+
+    _HAS_ORJSON = False
+
+
+def _coerce_keys(obj):
+    """Recursively convert non-string dict keys to strings.
+
+    stdlib ``json.dumps`` raises ``TypeError`` on non-string keys, whereas
+    orjson's ``OPT_NON_STR_KEYS`` converts them automatically. This helper
+    replicates that behavior for the fallback path.
+    """
+    if isinstance(obj, dict):
+        return {str(k): _coerce_keys(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_coerce_keys(v) for v in obj]
+    return obj
+
+
+if _HAS_ORJSON:
+    JSONDecodeError = _orjson.JSONDecodeError
+
+    def loads(data):
+        """Deserialize JSON string/bytes to a Python object."""
+        return _orjson.loads(data)
+
+    def dumps(obj, *, indent=False, sort_keys=False):
+        """Serialize a Python object to a JSON ``str``.
+
+        Args:
+            obj: The object to serialize.
+            indent: If True, pretty-print with 2-space indent.
+            sort_keys: If True, sort dictionary keys.
+        """
+        option = _orjson.OPT_NON_STR_KEYS
+        if indent:
+            option |= _orjson.OPT_INDENT_2
+        if sort_keys:
+            option |= _orjson.OPT_SORT_KEYS
+        return _orjson.dumps(obj, option=option).decode()
+
+else:
+    from json import JSONDecodeError  # noqa: F401
+
+    def loads(data):
+        """Deserialize JSON string/bytes to a Python object."""
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            data = (
+                bytes(data).decode() if isinstance(data, memoryview) else data.decode()
+            )
+        return _json.loads(data)
+
+    def dumps(obj, *, indent=False, sort_keys=False):
+        """Serialize a Python object to a JSON ``str``.
+
+        Args:
+            obj: The object to serialize.
+            indent: If True, pretty-print with 2-space indent.
+            sort_keys: If True, sort dictionary keys.
+        """
+        obj = _coerce_keys(obj)
+        kwargs = {"ensure_ascii": False}
+        if indent:
+            kwargs["indent"] = 2
+        else:
+            kwargs["separators"] = (",", ":")
+        if sort_keys:
+            kwargs["sort_keys"] = True
+        return _json.dumps(obj, **kwargs)

tritonparse/ai/client.py

@@ -16,7 +16,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Iterator, List, Optional, Tuple
 
-import orjson
+from tritonparse._json_compat import JSONDecodeError, loads
 
 
 @dataclass
@@ -292,9 +292,9 @@ def chat(
             stdout_snippet = ""
             if result and result.stdout:
                 try:
-                    data = orjson.loads(result.stdout)
+                    data = loads(result.stdout)
                     stdout_snippet = data.get("result", result.stdout[:500])
-                except (orjson.JSONDecodeError, AttributeError):
+                except (JSONDecodeError, AttributeError):
                     stdout_snippet = result.stdout[:500]
             stderr_snippet = (
                 result.stderr[:500] if result and result.stderr else "(empty)"
@@ -384,7 +384,7 @@ def chat_stream(
                         continue
 
                     try:
-                        event = orjson.loads(line)
+                        event = loads(line)
                         event_type = event.get("type")
 
                         if event_type == "assistant":
@@ -398,7 +398,7 @@ def chat_stream(
                             if "session_id" in event:
                                 self.session_id = event["session_id"]
 
-                    except orjson.JSONDecodeError:
+                    except JSONDecodeError:
                         continue
 
                 process.wait(timeout=self.timeout)
@@ -441,14 +441,14 @@ def _parse_response(self, stdout: str) -> Response:
             Parsed Response object
         """
         try:
-            data = orjson.loads(stdout)
+            data = loads(stdout)
             self.session_id = data.get("session_id")
             return Response(
                 content=data.get("result", stdout),
                 session_id=self.session_id,
                 cost_usd=data.get("total_cost_usd"),
                 raw=data,
             )
-        except orjson.JSONDecodeError:
+        except JSONDecodeError:
             # Non-JSON output, return as-is
             return Response(content=stdout.strip())

tritonparse/ai/parsers.py

@@ -12,7 +12,7 @@
 import re
 from typing import Optional
 
-import orjson
+from tritonparse._json_compat import JSONDecodeError, loads
 
 
 def extract_json(text: str) -> Optional[dict]:
@@ -34,17 +34,17 @@ def extract_json(text: str) -> Optional[dict]:
 
     # Try direct JSON parsing
     try:
-        return orjson.loads(text.strip())
-    except orjson.JSONDecodeError:
+        return loads(text.strip())
+    except JSONDecodeError:
         pass
 
     # Try extracting from markdown code block
     pattern = r"```(?:json)?\s*\n?(.*?)\n?```"
     match = re.search(pattern, text, re.DOTALL)
     if match:
         try:
-            return orjson.loads(match.group(1).strip())
-        except orjson.JSONDecodeError:
+            return loads(match.group(1).strip())
+        except JSONDecodeError:
             pass
 
     return None

tritonparse/bisect/state.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
-import orjson
+from tritonparse._json_compat import dumps, loads
 
 
 class BisectPhase(Enum):
@@ -322,12 +322,7 @@ def save(
 
         # Write JSON
         with open(save_path, "w") as f:
-            f.write(
-                orjson.dumps(
-                    state.to_dict(),
-                    option=orjson.OPT_INDENT_2 | orjson.OPT_NON_STR_KEYS,
-                ).decode()
-            )
+            f.write(dumps(state.to_dict(), indent=True))
 
         return save_path
 
@@ -344,11 +339,11 @@ def load(path: str) -> BisectState:
 
         Raises:
             FileNotFoundError: If state file doesn't exist.
-            orjson.JSONDecodeError: If file is not valid JSON.
+            JSONDecodeError: If file is not valid JSON.
             ValueError: If state data is invalid.
         """
         with open(path, "r") as f:
-            data = orjson.loads(f.read())
+            data = loads(f.read())
         return BisectState.from_dict(data)
 
     @staticmethod

tritonparse/diff/core/event_matcher.py

@@ -10,7 +10,7 @@
 from collections import defaultdict
 from typing import Any
 
-import orjson
+from tritonparse._json_compat import dumps, loads
 from tritonparse.tools.prettify_ndjson import load_ndjson
 
 
@@ -89,9 +89,9 @@ def ensure_source_mappings(event: dict[str, Any]) -> dict[str, Any]:
     # Reuse parse module function to generate source_mappings
     from tritonparse.parse.trace_processor import parse_single_trace_content
 
-    event_str = orjson.dumps(event, option=orjson.OPT_NON_STR_KEYS).decode()
+    event_str = dumps(event)
     parsed_str = parse_single_trace_content(event_str)
-    return orjson.loads(parsed_str.strip())
+    return loads(parsed_str.strip())
 
 
 def match_events_by_index(

tritonparse/diff/output/event_writer.py

@@ -14,7 +14,7 @@
 from dataclasses import asdict
 from typing import Any
 
-import orjson
+from tritonparse._json_compat import dumps
 from tritonparse.diff.core.diff_types import (
     CompilationDiffResult,
     PythonLineDiff,
@@ -173,12 +173,7 @@ def append_diff_to_file(file_path: str, diff_event: dict[str, Any]) -> None:
         diff_event: The compilation_diff event to append.
     """
     with open(file_path, "a") as f:
-        f.write(
-            orjson.dumps(
-                _sanitize_non_finite_floats(diff_event), option=orjson.OPT_NON_STR_KEYS
-            ).decode()
-            + "\n"
-        )
+        f.write(dumps(_sanitize_non_finite_floats(diff_event)) + "\n")
 
 
 class ConsolidatedDiffWriter:
@@ -297,22 +292,10 @@ def write(self, output_path: str | None = None) -> None:
         with open(path, "w") as f:
             # Write all unique compilation events first
             for event in self._events.values():
-                f.write(
-                    orjson.dumps(
-                        _sanitize_non_finite_floats(event),
-                        option=orjson.OPT_NON_STR_KEYS,
-                    ).decode()
-                    + "\n"
-                )
+                f.write(dumps(_sanitize_non_finite_floats(event)) + "\n")
             # Then write all diff events
             for diff in self._diffs:
-                f.write(
-                    orjson.dumps(
-                        _sanitize_non_finite_floats(diff),
-                        option=orjson.OPT_NON_STR_KEYS,
-                    ).decode()
-                    + "\n"
-                )
+                f.write(dumps(_sanitize_non_finite_floats(diff)) + "\n")
 
     @property
     def event_count(self) -> int:

tritonparse/info/cli.py

@@ -169,11 +169,11 @@ def _parse_value(value: str) -> Any:
 
     # Check for list (e.g., "[3024, 10752]")
     if value.startswith("[") and value.endswith("]"):
-        import orjson
+        from tritonparse._json_compat import JSONDecodeError, loads
 
         try:
-            return orjson.loads(value)
-        except orjson.JSONDecodeError:
+            return loads(value)
+        except JSONDecodeError:
             return value
 
     # Try to convert to int

tritonparse/parse/common.py

@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
-import orjson
+from tritonparse._json_compat import dumps
 from tritonparse.shared_vars import (
     DEFAULT_TRACE_FILE_PREFIX_WITHOUT_USER as LOG_PREFIX,
     is_fbcode,
@@ -518,11 +518,7 @@ def parse_logs(
     # Save file mapping to parsed_log_dir
     log_file_list_path = os.path.join(parsed_log_dir, "log_file_list.json")
     with open(log_file_list_path, "w") as f:
-        f.write(
-            orjson.dumps(
-                file_mapping, option=orjson.OPT_INDENT_2 | orjson.OPT_NON_STR_KEYS
-            ).decode()
-        )
+        f.write(dumps(file_mapping, indent=True))
 
     # NOTICE: this print is required for tlparser-tritonparse integration
     # DON'T REMOVE THIS PRINT