Relax configure allowance in async tasks (#8444)

* init

* add tests

Author: chenmoneygithub

dspy/dsp/utils/settings.py

@@ -31,6 +31,7 @@
 # Global base configuration and owner tracking
 main_thread_config = copy.deepcopy(DEFAULT_CONFIG)
 config_owner_thread_id = None
+config_owner_async_task = None
 
 # Global lock for settings configuration
 global_lock = threading.Lock()
@@ -105,13 +106,12 @@ def config(self):
         return self.copy()
 
     def _ensure_configure_allowed(self):
-        global main_thread_config, config_owner_thread_id
+        global main_thread_config, config_owner_thread_id, config_owner_async_task
         current_thread_id = threading.get_ident()
 
         if config_owner_thread_id is None:
-            # First `configure` call is always allowed.
+            # First `configure` call assigns the owner thread id.
             config_owner_thread_id = current_thread_id
-            return
 
         if config_owner_thread_id != current_thread_id:
             # Disallow a second `configure` calls from other threads.
@@ -130,6 +130,11 @@ def _ensure_configure_allowed(self):
         if not is_async_task:
             return
 
+        if config_owner_async_task is None:
+            # First `configure` call assigns the owner async task.
+            config_owner_async_task = asyncio.current_task()
+            return
+
         # We are in an async task. Now check for IPython and allow calling `configure` from IPython.
         in_ipython = False
         try:
@@ -142,10 +147,10 @@ def _ensure_configure_allowed(self):
             # If `IPython` is not installed or `get_ipython` failed, we are not in an IPython environment.
             in_ipython = False
 
-        if not in_ipython:
+        if not in_ipython and config_owner_async_task != asyncio.current_task():
             raise RuntimeError(
-                "dspy.settings.configure(...) cannot be called a second time from an async task. Use "
-                "`dspy.context(...)` instead."
+                "dspy.settings.configure(...) can only be called from the same async task that called it first. Please "
+                "use `dspy.context(...)` in other async tasks instead."
             )
 
     def configure(self, **kwargs):

tests/utils/test_settings.py

@@ -20,27 +20,13 @@ def test_forbid_configure_call_in_child_thread():
     dspy.configure(lm=dspy.LM("openai/gpt-4o"), adapter=dspy.JSONAdapter(), callbacks=[lambda x: x])
 
     def worker():
-        with pytest.raises(RuntimeError, match="Cannot call dspy.configure in a child thread"):
+        with pytest.raises(RuntimeError, match="Cannot call dspy.configure"):
             dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"), callbacks=[])
 
     with ThreadPoolExecutor(max_workers=1) as executor:
         executor.submit(worker)
 
 
-@pytest.mark.asyncio
-async def test_forbid_configure_call_in_async_function():
-    with pytest.raises(
-        RuntimeError,
-        match=r"dspy.settings.configure\(\.\.\.\) cannot be called a second time from*",
-    ):
-        dspy.configure(lm=dspy.LM("openai/gpt-4o"), adapter=dspy.JSONAdapter(), callbacks=[lambda x: x])
-        dspy.configure(lm=dspy.LM("openai/gpt-4o"), adapter=dspy.JSONAdapter(), callbacks=[lambda x: x])
-
-    with dspy.context(lm=dspy.LM("openai/gpt-4o-mini"), callbacks=[]):
-        # context is allowed
-        pass
-
-
 def test_dspy_context():
     dspy.configure(lm=dspy.LM("openai/gpt-4o"), adapter=dspy.JSONAdapter(), callbacks=[lambda x: x])
     with dspy.context(lm=dspy.LM("openai/gpt-4o-mini"), callbacks=[]):
@@ -165,3 +151,43 @@ async def aforward(self, question: str) -> str:
         # The main thread is not affected by the context
         assert dspy.settings.lm.model == "openai/gpt-4.1"
         assert dspy.settings.trace == []
+
+
+@pytest.mark.asyncio
+async def test_dspy_configure_allowance_async():
+    def bar1():
+        # `dspy.configure` is disallowed in different async tasks from the initial one.
+        # In this case, foo1 (async) calls bar1 (sync), and bar1 uses the async task from foo1.
+        with pytest.raises(RuntimeError) as e:
+            dspy.configure(lm=dspy.LM("openai/gpt-4o"))
+        assert "dspy.settings.configure(...) can only be called from the same async" in str(e.value)
+
+    async def foo1():
+        bar1()
+        await asyncio.sleep(0.1)
+
+    async def foo2():
+        # `dspy.configure` is disallowed in different async tasks from the initial one.
+        with pytest.raises(RuntimeError) as e:
+            dspy.configure(lm=dspy.LM("openai/gpt-4o"))
+        assert "dspy.settings.configure(...) can only be called from the same async" in str(e.value)
+        await asyncio.sleep(0.1)
+
+    async def foo3():
+        # `dspy.context` is allowed in different async tasks from the initial one.
+        with dspy.context(lm=dspy.LM("openai/gpt-4o")):
+            await asyncio.sleep(0.1)
+
+    async def foo4():
+        # foo4 is directly invoked by the entry task, so it has the same async task as the entry task.
+        dspy.configure(lm=dspy.LM("openai/gpt-4o"))
+        await asyncio.sleep(0.1)
+
+    # `dspy.configure` is allowed to be called multiple times in the same async task.
+    dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))
+    dspy.configure(lm=dspy.LM("openai/gpt-4o"))
+    dspy.configure(adapter=dspy.JSONAdapter())
+
+    await asyncio.gather(foo1(), foo2(), foo3())
+
+    foo4()