Skip to content

Commit

Permalink
Updates Burr tracking client to safely encode json
Browse files Browse the repository at this point in the history
There are often missing characters -- this enables the user to safely
encode them. They'll get replaced, but generally that's OK (and
certainly better than failing).
  • Loading branch information
elijahbenizzy committed May 22, 2024
1 parent e10db73 commit 63fb023
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions burr/tracking/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import os
import re
import traceback
from typing import Any, Dict, Optional
from typing import Any, Dict, Optional, io

from burr import system
from burr.common import types as burr_types
Expand Down Expand Up @@ -57,6 +57,12 @@ def _allowed_project_name(project_name: str, on_windows: bool) -> bool:
return bool(re.match(pattern, project_name))


def safe_json_write(line: str, f: io.TextIO):
# Every once in a while we'll hit a non-utf-8 character
# In this case we replace it and hope for the best
return json.loads(line.decode("utf-8", errors="replace"))


class LocalTrackingClient(
PostApplicationCreateHook,
PreRunStepHook,
Expand Down Expand Up @@ -209,15 +215,20 @@ def post_application_create(
**future_kwargs: Any,
):
self._ensure_dir_structure(app_id)
self.f = open(os.path.join(self.storage_dir, app_id, self.LOG_FILENAME), "a")
self.f = open(
os.path.join(self.storage_dir, app_id, self.LOG_FILENAME),
"a",
encoding="utf-8",
errors="replace",
)
graph_path = os.path.join(self.storage_dir, app_id, self.GRAPH_FILENAME)
if os.path.exists(graph_path):
logger.info(f"Graph already exists at {graph_path}. Not overwriting.")
return
graph = ApplicationModel.from_application_graph(
application_graph,
).model_dump()
with open(graph_path, "w") as f:
with open(graph_path, "w", encoding="utf-8", errors="replace") as f:
json.dump(graph, f)

metadata_path = os.path.join(self.storage_dir, app_id, self.METADATA_FILENAME)
Expand All @@ -230,7 +241,7 @@ def post_application_create(
if parent_pointer is not None
else None,
).model_dump()
with open(metadata_path, "w") as f:
with open(metadata_path, "w", errors="replace") as f:
json.dump(metadata, f)

def _append_write_line(self, model: pydantic.BaseModel):
Expand Down

0 comments on commit 63fb023

Please sign in to comment.