Skip to content

Commit aca3dab

Browse files
feat: add Openlayer's client custom settings for tracer (#536)
* feat: add timeout configuration for Openlayer tracer * feat: add optional timeout configuration to tracer * test: enhance tracer configuration tests to include timeout validation * docs: update README to include tracing timeout configuration details * feat: add max retries configuration to Openlayer tracer * docs: update README to clarify tracing configuration with timeouts and retries * test: enhance tracer configuration tests to include max retries validation Added tests to verify the correct configuration and usage of the `max_retries` parameter in the tracer. Updated existing tests to incorporate `max_retries` in the configuration and assertions, ensuring that the tracer behaves as expected with this new parameter.
1 parent 5030f6e commit aca3dab

File tree

3 files changed

+91
-6
lines changed

3 files changed

+91
-6
lines changed

README.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,29 @@ On timeout, an `APITimeoutError` is thrown.
319319

320320
Note that requests that time out are [retried twice by default](#retries).
321321

322+
#### Tracing Configuration
323+
324+
When using the tracing decorators (`@trace()` and `@trace_async()`), you can configure timeouts and retries for the underlying API calls to Openlayer using the `configure()` function:
325+
326+
```python
327+
from openlayer.lib import trace, configure
328+
329+
# Configure timeout and retries for tracing API calls
330+
configure(
331+
api_key="your_api_key_here",
332+
inference_pipeline_id="your_pipeline_id_here",
333+
timeout=30.0, # 30 seconds timeout for tracing API calls (int or float)
334+
max_retries=5 # Maximum number of retries for failed requests (default: 2)
335+
)
336+
337+
# Now use the decorators normally
338+
@trace()
339+
def my_function():
340+
return "result"
341+
```
342+
343+
These settings apply to all API calls made by the tracer when streaming trace data to Openlayer. If not specified, the defaults are 60 seconds timeout and 2 retries.
344+
322345
## Advanced
323346

324347
### Logging

src/openlayer/lib/tracing/tracer.py

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import traceback
99
from contextlib import contextmanager
1010
from functools import wraps
11-
from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple
11+
from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple, Union
1212

1313
from ..._base_client import DefaultHttpxClient
1414
from ..._client import Openlayer
@@ -34,24 +34,30 @@
3434
_configured_api_key: Optional[str] = None
3535
_configured_pipeline_id: Optional[str] = None
3636
_configured_base_url: Optional[str] = None
37+
_configured_timeout: Optional[Union[int, float]] = None
38+
_configured_max_retries: Optional[int] = None
3739

3840

3941
def configure(
4042
api_key: Optional[str] = None,
4143
inference_pipeline_id: Optional[str] = None,
4244
base_url: Optional[str] = None,
45+
timeout: Optional[Union[int, float]] = None,
46+
max_retries: Optional[int] = None,
4347
) -> None:
4448
"""Configure the Openlayer tracer with custom settings.
4549
4650
This function allows you to programmatically set the API key, inference pipeline ID,
47-
and base URL for the Openlayer client, instead of relying on environment variables.
51+
base URL, timeout, and retry settings for the Openlayer client, instead of relying on environment variables.
4852
4953
Args:
5054
api_key: The Openlayer API key. If not provided, falls back to OPENLAYER_API_KEY environment variable.
5155
inference_pipeline_id: The default inference pipeline ID to use for tracing.
5256
If not provided, falls back to OPENLAYER_INFERENCE_PIPELINE_ID environment variable.
5357
base_url: The base URL for the Openlayer API. If not provided, falls back to
5458
OPENLAYER_BASE_URL environment variable or the default.
59+
timeout: The timeout for the Openlayer API in seconds (int or float). Defaults to 60 seconds.
60+
max_retries: The maximum number of retries for failed API requests. Defaults to 2.
5561
5662
Examples:
5763
>>> import openlayer.lib.tracing.tracer as tracer
@@ -62,11 +68,13 @@ def configure(
6268
>>> def my_function():
6369
... return "result"
6470
"""
65-
global _configured_api_key, _configured_pipeline_id, _configured_base_url, _client
71+
global _configured_api_key, _configured_pipeline_id, _configured_base_url, _configured_timeout, _configured_max_retries, _client
6672

6773
_configured_api_key = api_key
6874
_configured_pipeline_id = inference_pipeline_id
6975
_configured_base_url = base_url
76+
_configured_timeout = timeout
77+
_configured_max_retries = max_retries
7078

7179
# Reset the client so it gets recreated with new configuration
7280
_client = None
@@ -90,6 +98,12 @@ def _get_client() -> Optional[Openlayer]:
9098
if _configured_base_url is not None:
9199
client_kwargs["base_url"] = _configured_base_url
92100

101+
if _configured_timeout is not None:
102+
client_kwargs["timeout"] = _configured_timeout
103+
104+
if _configured_max_retries is not None:
105+
client_kwargs["max_retries"] = _configured_max_retries
106+
93107
if _verify_ssl:
94108
_client = Openlayer(**client_kwargs)
95109
else:

tests/test_tracer_configuration.py

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,25 @@ def teardown_method(self):
1515
tracer._configured_api_key = None
1616
tracer._configured_pipeline_id = None
1717
tracer._configured_base_url = None
18+
tracer._configured_timeout = None
19+
tracer._configured_max_retries = None
1820
tracer._client = None
1921

2022
def test_configure_sets_global_variables(self):
2123
"""Test that configure() sets the global configuration variables."""
2224
api_key = "test_api_key"
2325
pipeline_id = "test_pipeline_id"
2426
base_url = "https://test.api.com"
27+
timeout = 30.5
28+
max_retries = 5
2529

26-
tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url)
30+
tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url, timeout=timeout, max_retries=max_retries)
2731

2832
assert tracer._configured_api_key == api_key
2933
assert tracer._configured_pipeline_id == pipeline_id
3034
assert tracer._configured_base_url == base_url
35+
assert tracer._configured_timeout == timeout
36+
assert tracer._configured_max_retries == max_retries
3137

3238
def test_configure_resets_client(self):
3339
"""Test that configure() resets the client to force recreation."""
@@ -77,6 +83,42 @@ def test_get_client_uses_both_configured_values(self, mock_openlayer: Any) -> No
7783

7884
mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url)
7985

86+
@patch("openlayer.lib.tracing.tracer.Openlayer")
87+
def test_get_client_uses_configured_timeout(self, mock_openlayer: Any) -> None:
88+
"""Test that _get_client() uses the configured timeout."""
89+
with patch.object(tracer, "_publish", True):
90+
timeout = 45.5
91+
tracer.configure(timeout=timeout)
92+
93+
tracer._get_client()
94+
95+
mock_openlayer.assert_called_once_with(timeout=timeout)
96+
97+
@patch("openlayer.lib.tracing.tracer.Openlayer")
98+
def test_get_client_uses_configured_max_retries(self, mock_openlayer: Any) -> None:
99+
"""Test that _get_client() uses the configured max_retries."""
100+
with patch.object(tracer, "_publish", True):
101+
max_retries = 10
102+
tracer.configure(max_retries=max_retries)
103+
104+
tracer._get_client()
105+
106+
mock_openlayer.assert_called_once_with(max_retries=max_retries)
107+
108+
@patch("openlayer.lib.tracing.tracer.Openlayer")
109+
def test_get_client_uses_all_configured_values(self, mock_openlayer: Any) -> None:
110+
"""Test that _get_client() uses all configured values together."""
111+
with patch.object(tracer, "_publish", True):
112+
api_key = "configured_api_key"
113+
base_url = "https://configured.api.com"
114+
timeout = 25
115+
max_retries = 3
116+
tracer.configure(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)
117+
118+
tracer._get_client()
119+
120+
mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)
121+
80122
@patch("openlayer.lib.tracing.tracer.DefaultHttpxClient")
81123
@patch("openlayer.lib.tracing.tracer.Openlayer")
82124
def test_get_client_with_ssl_disabled_and_config(self, mock_openlayer: Any, mock_http_client: Any) -> None:
@@ -150,13 +192,19 @@ def test_configure_with_none_values(self):
150192
"""Test that configure() with None values doesn't overwrite existing config."""
151193
# Set initial configuration
152194
tracer.configure(
153-
api_key="initial_key", inference_pipeline_id="initial_pipeline", base_url="https://initial.com"
195+
api_key="initial_key",
196+
inference_pipeline_id="initial_pipeline",
197+
base_url="https://initial.com",
198+
timeout=60.0,
199+
max_retries=5
154200
)
155201

156202
# Configure with None values
157-
tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None)
203+
tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None, timeout=None, max_retries=None)
158204

159205
# Values should be set to None (this is the expected behavior)
160206
assert tracer._configured_api_key is None
161207
assert tracer._configured_pipeline_id is None
162208
assert tracer._configured_base_url is None
209+
assert tracer._configured_timeout is None
210+
assert tracer._configured_max_retries is None

0 commit comments

Comments
 (0)