From e0ee82c8f9d56fcdb77620b0b56d2570ed239119 Mon Sep 17 00:00:00 2001
From: Dodam Ih <dodam@zetta.ai>
Date: Tue, 21 Jan 2025 22:15:07 -0800
Subject: [PATCH] fix: change id_generation to cloudpickle due to dill
 recursion bug

---
 pyproject.toml                      |  3 ++-
 zetta_utils/mazepa/id_generation.py | 16 ++++------------
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 33f2c56cc..d3e2c6a13 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,6 +19,7 @@ dependencies = [
   "rich >= 12.6.0",
   "python-logging-loki >= 0.3.1",
   "neuroglancer >= 2.32",
+  "cloudpickle >= 3.1.1",
   "dill >= 0.3.6",
   "pyyaml ~= 6.0.1",
   "requests==2.31.0", # version conflicts otherwise
@@ -34,7 +35,7 @@ keywords = ["neuroscience connectomics EM"]
 license = {text = "MIT"}
 name = "zetta_utils"
 readme = "README.md"
-requires-python = ">3.9,<3.13"
+requires-python = ">3.10,<3.13"
 urls = {Homepage = "https://github.com/zettaai/zetta_utils"}
 version = "0.0.2"
 
diff --git a/zetta_utils/mazepa/id_generation.py b/zetta_utils/mazepa/id_generation.py
index 62884393a..bb87597a2 100644
--- a/zetta_utils/mazepa/id_generation.py
+++ b/zetta_utils/mazepa/id_generation.py
@@ -4,7 +4,7 @@
 import uuid
 from typing import Callable, Optional
 
-import dill
+import cloudpickle
 import xxhash
 from coolname import generate_slug
 
@@ -42,7 +42,7 @@ def generate_invocation_id(
     prefix: Optional[str] = None,
 ) -> str:
     """Generate a unique and deterministic ID for a function invocation.
-    The ID is generated using xxhash and dill to hash the function and its arguments.
+    The ID is generated using xxhash and cloudpickle to hash the function and its arguments.
 
     :param fn: the function, or really any Callable, defaults to None
     :param args: the function arguments, or any list, defaults to None
@@ -53,16 +53,8 @@ def generate_invocation_id(
     """
     x = xxhash.xxh128()
     try:
-        x.update(
-            dill.dumps(
-                (fn, args, kwargs),
-                protocol=dill.DEFAULT_PROTOCOL,
-                byref=False,
-                recurse=True,
-                fmode=dill.FILE_FMODE,
-            )
-        )
-    except dill.PicklingError as e:
+        x.update(cloudpickle.dumps((fn, args, kwargs)))
+    except Exception as e:  # pylint: disable=broad-exception-caught
         logger.warning(f"Failed to pickle {fn} with args {args} and kwargs {kwargs}: {e}")
         x.update(str(uuid.uuid4()))