updated map such that context vars are propagated through to the inne…

…r threads created.
unifyai · Oct 31, 2024 · 9c8ddaa · 9c8ddaa
1 parent 74f1ad1
commit 9c8ddaa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/notebook.py b/notebook.py
diff --git a/unify/utils/map.py b/unify/utils/map.py
@@ -1,5 +1,6 @@
 import asyncio
 import threading
+import contextvars
 from typing import Any, List
 
 
@@ -27,6 +28,9 @@ def map(fn: callable, *args, mode="threading", **kwargs) -> Any:
     if mode == "threading":
 
         def fn_w_indexing(rets: List[None], thread_idx: int, *a, **kw):
+            for var, value in kw["context"].items():
+                var.set(value)
+            del kw["context"]
             ret = fn(*a, **kw)
             rets[thread_idx] = ret
 
@@ -38,6 +42,7 @@ def fn_w_indexing(rets: List[None], thread_idx: int, *a, **kw):
                 k: v[i] if (isinstance(v, list) or isinstance(v, tuple)) else v
                 for k, v in kwargs.items()
             }
+            kw["context"] = contextvars.copy_context()
             thread = threading.Thread(
                 target=fn_w_indexing,
                 args=(returns, i, *a),