Skip to content

Commit

Permalink
[OPIK-555] Add retries for get dataset items operation (#839)
Browse files Browse the repository at this point in the history
* Add retry

* Update connection_retry decorator parameters

* Add dispatch option to cookbooks workflow
  • Loading branch information
alexkuzmik authored Dec 9, 2024
1 parent 4c1790f commit 222c25c
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 4 deletions.
14 changes: 12 additions & 2 deletions .github/workflows/documentation_cookbook_tests.yml
Original file line number Diff line number Diff line change
@@ -1,6 +1,15 @@
name: Documentation - Test cookbooks
on:
workflow_dispatch:
inputs:
install_opik:
description: 'Enable opik installation from source files'
required: false
default: 'false'
type: choice
options:
- 'false'
- 'true'
schedule:
- cron: '0 0 * * *' # Run once a day at midnight UTC
jobs:
Expand Down Expand Up @@ -43,8 +52,9 @@ jobs:
run: |
python -m pip install --upgrade pip
python -m pip install -U ipython nbconvert
# - name: Install opik # Uncomment if you want to run cookbooks with local version of opik
# run: pip install sdks/python
- name: Install opik from source files (optional)
if: ${{ github.event_name == 'workflow_dispatch' && inputs.install_opik == 'true' }}
run: pip install sdks/python
- name: Prepare env variables
run: |
directory=$(dirname -- "${NOTEBOOK_TO_TEST}")
Expand Down
3 changes: 2 additions & 1 deletion sdks/python/src/opik/api_objects/dataset/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from opik.rest_api.types import dataset_item_write as rest_dataset_item
from opik.message_processing.batching import sequence_splitter
from opik import exceptions, config

from opik.rest_client_configurator import retry_decorators
from .. import constants
from . import dataset_item, converters

Expand Down Expand Up @@ -206,6 +206,7 @@ def get_items(self, nb_samples: Optional[int] = None) -> List[Dict[str, Any]]:

return dataset_items_as_dicts

@retry_decorators.connection_retry
def __internal_api__get_items_as_dataclasses__(
self, nb_samples: Optional[int] = None
) -> List[dataset_item.DatasetItem]:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

connection_retry = tenacity.retry(
stop=tenacity.stop_after_attempt(3),
wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
wait=tenacity.wait_exponential(multiplier=2, min=3, max=10),
retry=tenacity.retry_if_exception_type(
(
httpx.RemoteProtocolError, # handle retries for expired connections
Expand Down

0 comments on commit 222c25c

Please # to comment.