diff --git a/README.md b/README.md index 4f32c54b..9eb46bdf 100644 --- a/README.md +++ b/README.md @@ -319,6 +319,29 @@ On timeout, an `APITimeoutError` is thrown. Note that requests that time out are [retried twice by default](#retries). +#### Tracing Configuration + +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: + +```python +from openlayer.lib import trace, configure + +# Configure timeout and retries for tracing API calls +configure( + api_key="your_api_key_here", + inference_pipeline_id="your_pipeline_id_here", + timeout=30.0, # 30 seconds timeout for tracing API calls (int or float) + max_retries=5 # Maximum number of retries for failed requests (default: 2) +) + +# Now use the decorators normally +@trace() +def my_function(): + return "result" +``` + +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. + ## Advanced ### Logging diff --git a/src/openlayer/lib/tracing/tracer.py b/src/openlayer/lib/tracing/tracer.py index 1471c758..222d8b55 100644 --- a/src/openlayer/lib/tracing/tracer.py +++ b/src/openlayer/lib/tracing/tracer.py @@ -8,7 +8,7 @@ import traceback from contextlib import contextmanager from functools import wraps -from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple +from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple, Union from ..._base_client import DefaultHttpxClient from ..._client import Openlayer @@ -34,17 +34,21 @@ _configured_api_key: Optional[str] = None _configured_pipeline_id: Optional[str] = None _configured_base_url: Optional[str] = None +_configured_timeout: Optional[Union[int, float]] = None +_configured_max_retries: Optional[int] = None def configure( api_key: Optional[str] = None, inference_pipeline_id: Optional[str] = None, base_url: Optional[str] = None, + timeout: Optional[Union[int, float]] = None, + max_retries: Optional[int] = None, ) -> None: """Configure the Openlayer tracer with custom settings. This function allows you to programmatically set the API key, inference pipeline ID, - and base URL for the Openlayer client, instead of relying on environment variables. + base URL, timeout, and retry settings for the Openlayer client, instead of relying on environment variables. Args: api_key: The Openlayer API key. If not provided, falls back to OPENLAYER_API_KEY environment variable. @@ -52,6 +56,8 @@ def configure( If not provided, falls back to OPENLAYER_INFERENCE_PIPELINE_ID environment variable. base_url: The base URL for the Openlayer API. If not provided, falls back to OPENLAYER_BASE_URL environment variable or the default. + timeout: The timeout for the Openlayer API in seconds (int or float). Defaults to 60 seconds. + max_retries: The maximum number of retries for failed API requests. Defaults to 2. Examples: >>> import openlayer.lib.tracing.tracer as tracer @@ -62,11 +68,13 @@ def configure( >>> def my_function(): ... return "result" """ - global _configured_api_key, _configured_pipeline_id, _configured_base_url, _client + global _configured_api_key, _configured_pipeline_id, _configured_base_url, _configured_timeout, _configured_max_retries, _client _configured_api_key = api_key _configured_pipeline_id = inference_pipeline_id _configured_base_url = base_url + _configured_timeout = timeout + _configured_max_retries = max_retries # Reset the client so it gets recreated with new configuration _client = None @@ -90,6 +98,12 @@ def _get_client() -> Optional[Openlayer]: if _configured_base_url is not None: client_kwargs["base_url"] = _configured_base_url + if _configured_timeout is not None: + client_kwargs["timeout"] = _configured_timeout + + if _configured_max_retries is not None: + client_kwargs["max_retries"] = _configured_max_retries + if _verify_ssl: _client = Openlayer(**client_kwargs) else: diff --git a/tests/test_tracer_configuration.py b/tests/test_tracer_configuration.py index 7303f139..a83745c4 100644 --- a/tests/test_tracer_configuration.py +++ b/tests/test_tracer_configuration.py @@ -15,6 +15,8 @@ def teardown_method(self): tracer._configured_api_key = None tracer._configured_pipeline_id = None tracer._configured_base_url = None + tracer._configured_timeout = None + tracer._configured_max_retries = None tracer._client = None def test_configure_sets_global_variables(self): @@ -22,12 +24,16 @@ def test_configure_sets_global_variables(self): api_key = "test_api_key" pipeline_id = "test_pipeline_id" base_url = "https://test.api.com" + timeout = 30.5 + max_retries = 5 - tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url) + tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url, timeout=timeout, max_retries=max_retries) assert tracer._configured_api_key == api_key assert tracer._configured_pipeline_id == pipeline_id assert tracer._configured_base_url == base_url + assert tracer._configured_timeout == timeout + assert tracer._configured_max_retries == max_retries def test_configure_resets_client(self): """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 mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url) + @patch("openlayer.lib.tracing.tracer.Openlayer") + def test_get_client_uses_configured_timeout(self, mock_openlayer: Any) -> None: + """Test that _get_client() uses the configured timeout.""" + with patch.object(tracer, "_publish", True): + timeout = 45.5 + tracer.configure(timeout=timeout) + + tracer._get_client() + + mock_openlayer.assert_called_once_with(timeout=timeout) + + @patch("openlayer.lib.tracing.tracer.Openlayer") + def test_get_client_uses_configured_max_retries(self, mock_openlayer: Any) -> None: + """Test that _get_client() uses the configured max_retries.""" + with patch.object(tracer, "_publish", True): + max_retries = 10 + tracer.configure(max_retries=max_retries) + + tracer._get_client() + + mock_openlayer.assert_called_once_with(max_retries=max_retries) + + @patch("openlayer.lib.tracing.tracer.Openlayer") + def test_get_client_uses_all_configured_values(self, mock_openlayer: Any) -> None: + """Test that _get_client() uses all configured values together.""" + with patch.object(tracer, "_publish", True): + api_key = "configured_api_key" + base_url = "https://configured.api.com" + timeout = 25 + max_retries = 3 + tracer.configure(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries) + + tracer._get_client() + + mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries) + @patch("openlayer.lib.tracing.tracer.DefaultHttpxClient") @patch("openlayer.lib.tracing.tracer.Openlayer") 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): """Test that configure() with None values doesn't overwrite existing config.""" # Set initial configuration tracer.configure( - api_key="initial_key", inference_pipeline_id="initial_pipeline", base_url="https://initial.com" + api_key="initial_key", + inference_pipeline_id="initial_pipeline", + base_url="https://initial.com", + timeout=60.0, + max_retries=5 ) # Configure with None values - tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None) + tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None, timeout=None, max_retries=None) # Values should be set to None (this is the expected behavior) assert tracer._configured_api_key is None assert tracer._configured_pipeline_id is None assert tracer._configured_base_url is None + assert tracer._configured_timeout is None + assert tracer._configured_max_retries is None