From 9545d058823b38421f943edd4ff1407e199fecb3 Mon Sep 17 00:00:00 2001 From: Mason Daugherty Date: Wed, 20 May 2026 20:13:05 -0500 Subject: [PATCH] test(fireworks): stabilize integration tests with rate limiting and retries (#37590) Fireworks integration tests have been flaky against the live API with 429s. Adds a shared, xdist-aware rate limiter and a global retry policy so transient rate-limit errors no longer fail the suite. Mirrors the same fix recently applied to `langchain-mistralai`. --- libs/partners/fireworks/Makefile | 2 +- libs/partners/fireworks/pyproject.toml | 1 + .../tests/integration_tests/_rate_limiter.py | 18 ++++++++++++++++++ .../integration_tests/test_chat_models.py | 19 +++++++++++++------ .../tests/integration_tests/test_standard.py | 2 ++ libs/partners/fireworks/uv.lock | 14 ++++++++++++++ 6 files changed, 49 insertions(+), 7 deletions(-) create mode 100644 libs/partners/fireworks/tests/integration_tests/_rate_limiter.py diff --git a/libs/partners/fireworks/Makefile b/libs/partners/fireworks/Makefile index bb09a646233..6d8c3b19290 100644 --- a/libs/partners/fireworks/Makefile +++ b/libs/partners/fireworks/Makefile @@ -15,7 +15,7 @@ test tests: uv run --group test pytest $(PYTEST_EXTRA) --disable-socket --allow-unix-socket $(TEST_FILE) integration_test integration_tests: - uv run --group test --group test_integration pytest -v --tb=short -n auto $(TEST_FILE) + uv run --group test --group test_integration pytest -v --tb=short -n auto --retries 3 --retry-delay 2 $(TEST_FILE) test_watch: uv run --group test ptw --snapshot-update --now . -- -vv $(TEST_FILE) diff --git a/libs/partners/fireworks/pyproject.toml b/libs/partners/fireworks/pyproject.toml index 7de55b93c84..ee6aaeeb458 100644 --- a/libs/partners/fireworks/pyproject.toml +++ b/libs/partners/fireworks/pyproject.toml @@ -48,6 +48,7 @@ test = [ "syrupy>=5.0.0,<6.0.0", "pytest-watcher>=0.3.4,<1.0.0", "pytest-asyncio>=1.3.0,<2.0.0", + "pytest-retry>=1.7.0,<1.8.0", "pytest-socket>=0.7.0,<1.0.0", "pytest-xdist>=3.8.0,<4.0.0", "langchain-core>=1.4.0,<2.0.0", diff --git a/libs/partners/fireworks/tests/integration_tests/_rate_limiter.py b/libs/partners/fireworks/tests/integration_tests/_rate_limiter.py new file mode 100644 index 00000000000..fdcde8aa32c --- /dev/null +++ b/libs/partners/fireworks/tests/integration_tests/_rate_limiter.py @@ -0,0 +1,18 @@ +"""Shared rate limiter for Fireworks integration tests. + +Scaled by ``PYTEST_XDIST_WORKER_COUNT`` so aggregate QPS across all xdist +workers stays bounded near the target rate. +""" + +from __future__ import annotations + +import os + +from langchain_core.rate_limiters import InMemoryRateLimiter + +_TARGET_REQUESTS_PER_SECOND = 0.5 +_WORKER_COUNT = max(1, int(os.environ.get("PYTEST_XDIST_WORKER_COUNT", "1"))) + +rate_limiter = InMemoryRateLimiter( + requests_per_second=_TARGET_REQUESTS_PER_SECOND / _WORKER_COUNT, +) diff --git a/libs/partners/fireworks/tests/integration_tests/test_chat_models.py b/libs/partners/fireworks/tests/integration_tests/test_chat_models.py index a36fa4f6f6f..31ee824f737 100644 --- a/libs/partners/fireworks/tests/integration_tests/test_chat_models.py +++ b/libs/partners/fireworks/tests/integration_tests/test_chat_models.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, Field from typing_extensions import TypedDict from langchain_fireworks import ChatFireworks +from tests.integration_tests._rate_limiter import rate_limiter _MODEL = "accounts/fireworks/models/gpt-oss-120b" @@ -21,7 +22,9 @@ _MODEL = "accounts/fireworks/models/gpt-oss-120b" @pytest.mark.parametrize("strict", [None, True, False]) def test_tool_choice_bool(strict: bool | None) -> None: # noqa: FBT001 """Test that tool choice is respected with different strict values.""" - llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6") + llm = ChatFireworks( + model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter + ) class MyTool(BaseModel): name: str @@ -59,7 +62,9 @@ def test_tool_choice_bool(strict: bool | None) -> None: # noqa: FBT001 async def test_astream() -> None: """Test streaming tokens from ChatFireworks.""" - llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6") + llm = ChatFireworks( + model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter + ) full: BaseMessageChunk | None = None chunks_with_token_counts = 0 @@ -96,7 +101,7 @@ async def test_astream() -> None: async def test_abatch_tags() -> None: """Test batch tokens from ChatFireworks.""" - llm = ChatFireworks(model=_MODEL) + llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter) result = await llm.abatch( ["I'm Pickle Rick", "I'm not Pickle Rick"], config={"tags": ["foo"]} @@ -107,7 +112,7 @@ async def test_abatch_tags() -> None: async def test_ainvoke() -> None: """Test invoke tokens from ChatFireworks.""" - llm = ChatFireworks(model=_MODEL) + llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter) result = await llm.ainvoke("I'm Pickle Rick", config={"tags": ["foo"]}) assert isinstance(result.content, str) @@ -115,7 +120,7 @@ async def test_ainvoke() -> None: def test_invoke() -> None: """Test invoke tokens from ChatFireworks.""" - llm = ChatFireworks(model=_MODEL) + llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter) result = llm.invoke("I'm Pickle Rick", config={"tags": ["foo"]}) assert isinstance(result.content, str) @@ -157,7 +162,9 @@ def _get_joke_class( @pytest.mark.parametrize("schema_type", ["pydantic", "typeddict", "json_schema"]) def test_structured_output_json_schema(schema_type: str) -> None: - llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6") + llm = ChatFireworks( + model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter + ) schema, validation_function = _get_joke_class(schema_type) # type: ignore[arg-type] chat = llm.with_structured_output(schema, method="json_schema") diff --git a/libs/partners/fireworks/tests/integration_tests/test_standard.py b/libs/partners/fireworks/tests/integration_tests/test_standard.py index 7f1b1de5299..f6ee0aaf8d4 100644 --- a/libs/partners/fireworks/tests/integration_tests/test_standard.py +++ b/libs/partners/fireworks/tests/integration_tests/test_standard.py @@ -8,6 +8,7 @@ from langchain_tests.integration_tests import ( # type: ignore[import-not-found ) from langchain_fireworks import ChatFireworks +from tests.integration_tests._rate_limiter import rate_limiter class TestFireworksStandard(ChatModelIntegrationTests): @@ -20,6 +21,7 @@ class TestFireworksStandard(ChatModelIntegrationTests): return { "model": "accounts/fireworks/models/kimi-k2p6", "temperature": 0, + "rate_limiter": rate_limiter, } @pytest.mark.xfail(reason="Not yet implemented.") diff --git a/libs/partners/fireworks/uv.lock b/libs/partners/fireworks/uv.lock index 74f2181091f..816c54a4b2d 100644 --- a/libs/partners/fireworks/uv.lock +++ b/libs/partners/fireworks/uv.lock @@ -778,6 +778,7 @@ test = [ { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, + { name = "pytest-retry" }, { name = "pytest-socket" }, { name = "pytest-watcher" }, { name = "pytest-xdist" }, @@ -808,6 +809,7 @@ test = [ { name = "pytest", specifier = ">=9.0.3,<10.0.0" }, { name = "pytest-asyncio", specifier = ">=1.3.0,<2.0.0" }, { name = "pytest-mock", specifier = ">=3.10.0,<4.0.0" }, + { name = "pytest-retry", specifier = ">=1.7.0,<1.8.0" }, { name = "pytest-socket", specifier = ">=0.7.0,<1.0.0" }, { name = "pytest-watcher", specifier = ">=0.3.4,<1.0.0" }, { name = "pytest-xdist", specifier = ">=3.8.0,<4.0.0" }, @@ -1691,6 +1693,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/c2/ce34735972cc42d912173e79f200fe66530225190c06655c5632a9d88f1e/pytest_recording-0.13.4-py3-none-any.whl", hash = "sha256:ad49a434b51b1c4f78e85b1e6b74fdcc2a0a581ca16e52c798c6ace971f7f439", size = 13723, upload-time = "2025-05-08T10:41:09.684Z" }, ] +[[package]] +name = "pytest-retry" +version = "1.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/5b/607b017994cca28de3a1ad22a3eee8418e5d428dcd8ec25b26b18e995a73/pytest_retry-1.7.0.tar.gz", hash = "sha256:f8d52339f01e949df47c11ba9ee8d5b362f5824dff580d3870ec9ae0057df80f", size = 19977, upload-time = "2025-01-19T01:56:13.115Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/ff/3266c8a73b9b93c4b14160a7e2b31d1e1088e28ed29f4c2d93ae34093bfd/pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4", size = 13775, upload-time = "2025-01-19T01:56:11.199Z" }, +] + [[package]] name = "pytest-socket" version = "0.7.0"