(Feature) dspy.asyncify

dsp/utils/settings.py

@@ -24,6 +24,7 @@
     experimental=False,
     backoff_time=10,
     callbacks=[],
+    async_max_workers=8,
 )
 
 

dspy/__init__.py

@@ -7,11 +7,13 @@
 from .signatures import *
 
 # Functional must be imported after primitives, predict and signatures
-from .functional import *  # isort: skip
+from .functional import * # isort: skip
 from dspy.evaluate import Evaluate # isort: skip
 from dspy.clients import * # isort: skip
 from dspy.adapters import * # isort: skip
 from dspy.utils.logging_utils import configure_dspy_loggers, disable_logging, enable_logging
+from dspy.utils.asyncify import asyncify
+
 settings = dsp.settings
 
 configure_dspy_loggers(__name__)
@@ -71,4 +73,4 @@
 BetterTogether = dspy.teleprompt.BetterTogether
 COPRO = dspy.teleprompt.COPRO
 MIPROv2 = dspy.teleprompt.MIPROv2
-Ensemble = dspy.teleprompt.Ensemble
+Ensemble = dspy.teleprompt.Ensemble

dspy/utils/asyncify.py

@@ -0,0 +1,27 @@
+from anyio import CapacityLimiter
+import asyncer
+
+
+_limiter = None
+
+
+def get_async_max_workers():
+    import dspy
+
+    return dspy.settings.async_max_workers
+
+
+def get_limiter():
+    async_max_workers = get_async_max_workers()
+
+    global _limiter
+    if _limiter is None:
+        _limiter = CapacityLimiter(async_max_workers)
+    elif _limiter.total_tokens != async_max_workers:
+        _limiter.total_tokens = async_max_workers
+
+    return _limiter
+
+
+def asyncify(program):
+    return asyncer.asyncify(program, abandon_on_cancel=True, limiter=get_limiter())

pyproject.toml

@@ -20,12 +20,10 @@ classifiers = [
     "Programming Language :: Python :: 3",    # removed 3.8
     "Programming Language :: Python :: 3.9",
 ]
-
 # We have both project and tool.poetry.dependencies. Should we remove one?
 # tool.poetry.dependencies is a convenience thing for poetry users.
 # project dependencies function similarly to requirements.txt,
 # `pip install .` will pull from pyproject.toml dependencies
-
 dependencies = [
     "backoff~=2.2",
     "joblib~=1.3",
@@ -44,6 +42,8 @@ dependencies = [
     "diskcache",
     "json-repair",
     "tenacity>=8.2.3",
+    "anyio",
+    "asyncer==0.0.8",
 ]
 
 [project.optional-dependencies]
@@ -86,9 +86,8 @@ readme = "README.md"
 homepage = "https://github.com/stanfordnlp/dspy"
 repository = "https://github.com/stanfordnlp/dspy"
 keywords = ["dspy", "ai", "language models", "llm", "openai"]
-# may be a bit much
-
 
+# may be a bit much
 [tool.poetry.dependencies]
 python = ">=3.9,<3.13"
 pydantic = "^2.0"
@@ -134,17 +133,14 @@ jinja2 = "^3.1.3"
 magicattr = "^0.1.6"
 litellm = "1.51.0"
 diskcache = "^5.6.0"
-
 redis = "^5.1.1"
 falkordb = "^1.0.9"
-
-
 json-repair = "^0.30.0"
 tenacity = ">=8.2.3"
-
+asyncer = "0.0.8"
 
 [tool.poetry.group.dev.dependencies]
-pytest = "^6.2.5"
+pytest = "^8.3.3"
 transformers = "^4.38.2"
 torch = "^2.2.1"
 pytest-mock = "^3.12.0"
@@ -233,11 +229,9 @@ select = [
     "F",  # Pyflakes
     "E",  # Pycodestyle
 ]
-
 ignore = [
     "E501",  # Line too long
 ]
-
 # Allow fix for all enabled rules (when `--fix`) is provided.
 fixable = ["ALL"]
 unfixable = []

requirements-dev.txt

@@ -1,6 +1,6 @@
 black==24.2.0
 pre-commit==3.7.0
-pytest==8.2.1
+pytest==8.3.3
 pytest-env==1.1.3
 pytest-mock==3.12.0
 ruff==0.3.0

requirements.txt

@@ -15,3 +15,5 @@ requests
 tenacity>=8.2.3
 tqdm
 ujson
+anyio
+asyncer==0.0.8

tests/conftest.py

@@ -11,3 +11,8 @@ def clear_settings():
     yield
 
     dspy.settings.configure(**DEFAULT_CONFIG, inherit_config=False)
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"

tests/predict/test_chain_of_thought.py

@@ -1,5 +1,3 @@
-import textwrap
-
 import dspy
 from dspy import ChainOfThought
 from dspy.utils import DummyLM

tests/predict/test_multi_chain_comparison.py

@@ -2,29 +2,31 @@
 from dspy.utils.dummies import DummyLM
 
 
-def test_basic_example():
-    class BasicQA(dspy.Signature):
-        """Answer questions with short factoid answers."""
-
-        question = dspy.InputField()
-        answer = dspy.OutputField(desc="often between 1 and 5 words")
-
-    # Example completions generated by a model for reference
-    completions = [
-        dspy.Prediction(
-            rationale="I recall that during clear days, the sky often appears this color.",
-            answer="blue",
-        ),
-        dspy.Prediction(
-            rationale="Based on common knowledge, I believe the sky is typically seen as this color.",
-            answer="green",
-        ),
-        dspy.Prediction(
-            rationale="From images and depictions in media, the sky is frequently represented with this hue.",
-            answer="blue",
-        ),
-    ]
+class BasicQA(dspy.Signature):
+    """Answer questions with short factoid answers."""
+
+    question = dspy.InputField()
+    answer = dspy.OutputField(desc="often between 1 and 5 words")
+
 
+# Example completions generated by a model for reference
+completions = [
+    dspy.Prediction(
+        rationale="I recall that during clear days, the sky often appears this color.",
+        answer="blue",
+    ),
+    dspy.Prediction(
+        rationale="Based on common knowledge, I believe the sky is typically seen as this color.",
+        answer="green",
+    ),
+    dspy.Prediction(
+        rationale="From images and depictions in media, the sky is frequently represented with this hue.",
+        answer="blue",
+    ),
+]
+
+
+def test_basic_example():
     # Pass signature to MultiChainComparison module
     compare_answers = dspy.MultiChainComparison(BasicQA)
 

tests/predict/test_predict.py

@@ -1,5 +1,4 @@
 import copy
-import textwrap
 
 import pydantic
 import ujson
@@ -126,9 +125,11 @@ class Output(pydantic.BaseModel):
     # Demos don't need to keep the same types after saving and loading the state.
     assert new_instance.demos[0]["input"] == original_instance.demos[0].input.model_dump_json()
 
+
 def test_signature_fields_after_dump_and_load_state(tmp_path):
     class CustomSignature(dspy.Signature):
         """I am just an instruction."""
+
         sentence = dspy.InputField(desc="I am an innocent input!")
         sentiment = dspy.OutputField()
 
@@ -138,6 +139,7 @@ class CustomSignature(dspy.Signature):
 
     class CustomSignature2(dspy.Signature):
         """I am not a pure instruction."""
+
         sentence = dspy.InputField(desc="I am a malicious input!")
         sentiment = dspy.OutputField(desc="I am a malicious output!", prefix="I am a prefix!")
 
@@ -220,46 +222,39 @@ class OutputOnlySignature(dspy.Signature):
     assert predictor().output == "short answer"
 
 
-
 def test_chainable_load(tmp_path):
     """Test both traditional and chainable load methods."""
-    
+
     file_path = tmp_path / "test_chainable.json"
-
-
+
     original = Predict("question -> answer")
     original.demos = [{"question": "test", "answer": "response"}]
     original.save(file_path)
-
-
+
     traditional = Predict("question -> answer")
     traditional.load(file_path)
     assert traditional.demos == original.demos
-
-
+
     chainable = Predict("question -> answer").load(file_path, return_self=True)
-    assert chainable is not None  
+    assert chainable is not None
     assert chainable.demos == original.demos
-
-
+
     assert chainable.signature.dump_state() == original.signature.dump_state()
-
-
+
     result = Predict("question -> answer").load(file_path)
     assert result is None
 
+
 def test_load_state_chaining():
     """Test that load_state returns self for chaining."""
     original = Predict("question -> answer")
     original.demos = [{"question": "test", "answer": "response"}]
     state = original.dump_state()
-
-
+
     new_instance = Predict("question -> answer").load_state(state)
     assert new_instance is not None
     assert new_instance.demos == original.demos
-
-
+
     legacy_instance = Predict("question -> answer").load_state(state, use_legacy_loading=True)
     assert legacy_instance is not None
-    assert legacy_instance.demos == original.demos
+    assert legacy_instance.demos == original.demos

tests/predict/test_program_of_thought.py

@@ -1,5 +1,3 @@
-import textwrap
-
 import dspy
 from dspy import ProgramOfThought, Signature
 from dspy.utils import DummyLM

tests/utils/test_asyncify.py

@@ -0,0 +1,55 @@
+from time import time, sleep
+import asyncio
+
+import math
+import pytest
+
+import dspy
+from dspy.utils.asyncify import get_limiter
+
+
+@pytest.mark.anyio
+async def test_async_limiter():
+    limiter = get_limiter()
+    assert limiter.total_tokens == 8, "Default async capacity should be 8"
+    assert get_limiter() == limiter, "AsyncLimiter should be a singleton"
+
+    dspy.settings.configure(async_max_workers=16)
+    assert get_limiter() == limiter, "AsyncLimiter should be a singleton"
+    assert get_limiter().total_tokens == 16, "Async capacity should be 16"
+    assert get_limiter() == get_limiter(), "AsyncLimiter should be a singleton"
+
+
+@pytest.mark.anyio
+async def test_asyncify():
+    def the_answer_to_life_the_universe_and_everything(wait: float):
+        sleep(wait)
+        return 42
+
+    ask_the_question = dspy.asyncify(the_answer_to_life_the_universe_and_everything)
+
+    async def run_n_tasks(n: int, wait: float):
+        await asyncio.gather(*[ask_the_question(wait) for _ in range(n)])
+
+    async def verify_asyncify(capacity: int, number_of_tasks: int, wait: float = 0.5):
+        dspy.settings.configure(async_max_workers=capacity)
+
+        start = time()
+        await run_n_tasks(number_of_tasks, wait)
+        end = time()
+        total_time = end - start
+
+        # If asyncify is working correctly, the total time should be less than the total number of loops
+        # `(number_of_tasks / capacity)` times wait time, plus the computational overhead. The lower bound should
+        # be `math.floor(number_of_tasks * 1.0 / capacity) * wait` because there are more than
+        # `math.floor(number_of_tasks * 1.0 / capacity)` loops.
+        lower_bound = math.floor(number_of_tasks * 1.0 / capacity) * wait
+        upper_bound = math.ceil(number_of_tasks * 1.0 / capacity) * wait + 2 * wait  # 2*wait for buffer
+
+        assert lower_bound < total_time < upper_bound
+
+    await verify_asyncify(4, 10)
+    await verify_asyncify(8, 15)
+    await verify_asyncify(8, 30)
+
+