fix: improve metadata serialization (#244)

improve metadata serialization with orjson instead of torch serialization. | Method | Time Taken (1000 calls) | |--- |--- | | torch | 70ms | | json | 2 ms | | orjson | 0.3 ms | | ujson. | 0.8 ms| --------- Signed-off-by: kbatra <[email protected]>
numaproj · Aug 16, 2023 · 52a65c0 · 52a65c0
1 parent 8482c01
commit 52a65c0
Show file tree

Hide file tree

Showing 3 changed files with 398 additions and 131 deletions.
diff --git a/numalogic/registry/redis_registry.py b/numalogic/registry/redis_registry.py
@@ -13,6 +13,7 @@
 import time
 from datetime import datetime, timedelta
 from typing import Optional
+import orjson
 
 from redis.exceptions import RedisError
 
@@ -128,7 +129,7 @@ def __get_artifact_data(
  deserialized_artifact = loads(serialized_artifact)
  deserialized_metadata = None
  if serialized_metadata:
- deserialized_metadata = loads(serialized_metadata)
+ deserialized_metadata = orjson.loads(serialized_metadata)
  return ArtifactData(
  artifact=deserialized_artifact,
  metadata=deserialized_metadata,
@@ -185,7 +186,7 @@ def __save_artifact(
  _LOGGER.debug("Setting latest key : %s ,to this new key = %s", latest_key, new_version_key)
  serialized_metadata = ""
  if metadata:
- serialized_metadata = dumps(deserialized_object=metadata)
+ serialized_metadata = orjson.dumps(metadata)
  serialized_artifact = dumps(deserialized_object=artifact)
  pipe.hset(
  name=new_version_key,