From 7e66017e89de9748913cca4fa24b15b20748569d Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Wed, 16 Oct 2024 13:49:17 -0700 Subject: [PATCH 1/6] Add an optional timeout parameter to polling while loops --- src/codeocean/computation.py | 23 ++++++++++++++++--- src/codeocean/data_asset.py | 43 +++++++++++++++++++++++++++--------- 2 files changed, 53 insertions(+), 13 deletions(-) diff --git a/src/codeocean/computation.py b/src/codeocean/computation.py index bd72451..cb58e81 100644 --- a/src/codeocean/computation.py +++ b/src/codeocean/computation.py @@ -4,7 +4,7 @@ from dataclasses_json import dataclass_json from requests_toolbelt.sessions import BaseUrlSession from typing import Optional -from time import sleep +from time import sleep, time from codeocean.enum import StrEnum @@ -141,20 +141,37 @@ def run_capsule(self, run_params: RunParams) -> Computation: return Computation.from_dict(res.json()) - def wait_until_completed(self, computation: Computation, polling_interval: int = 5) -> Computation: + def wait_until_completed( + self, computation: Computation, polling_interval: int = 5, timeout: int | None = None, + ) -> Computation: """ - Polls the given computation until it reaches the 'Completed' or 'Failed' state. + Polls the given computation until it reaches the 'Completed' or 'Failed' + state. + + - `polling_interval` and `timeout` are in seconds """ if polling_interval < 5: raise ValueError( f"Polling interval {polling_interval} should be greater than or equal to 5" ) + if timeout is not None and timeout < polling_interval: + raise ValueError( + f"Timeout {timeout} should be greater than or equal to polling interval {polling_interval}" + ) + if timeout is not None and timeout < 0: + raise ValueError( + f"Timeout {timeout} should be greater than or equal to 0 (seconds), or None" + ) + t0 = time() while True: comp = self.get_computation(computation.id) if comp.state in [ComputationState.Completed, ComputationState.Failed]: return comp + if timeout is not None and (time() - t0) > timeout: + raise TimeoutError(f"Computation {computation.id} did not complete within {timeout} seconds") + sleep(polling_interval) def list_computation_results(self, computation_id: str, path: str = "") -> Folder: diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index e7783af..aea860c 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -1,13 +1,14 @@ from __future__ import annotations -from dataclasses_json import dataclass_json from dataclasses import dataclass -from requests_toolbelt.sessions import BaseUrlSession -from time import sleep +from time import sleep, time from typing import Optional -from codeocean.components import SortOrder, SearchFilter, Permissions -from codeocean.computation import PipelineProcess, Param +from dataclasses_json import dataclass_json +from requests_toolbelt.sessions import BaseUrlSession + +from codeocean.components import Permissions, SearchFilter, SortOrder +from codeocean.computation import Param, PipelineProcess from codeocean.enum import StrEnum @@ -227,8 +228,12 @@ def get_data_asset(self, data_asset_id: str) -> DataAsset: return DataAsset.from_dict(res.json()) - def update_metadata(self, data_asset_id: str, update_params: DataAssetUpdateParams) -> DataAsset: - res = self.client.put(f"data_assets/{data_asset_id}", json=update_params.to_dict()) + def update_metadata( + self, data_asset_id: str, update_params: DataAssetUpdateParams + ) -> DataAsset: + res = self.client.put( + f"data_assets/{data_asset_id}", json=update_params.to_dict() + ) return DataAsset.from_dict(res.json()) @@ -237,20 +242,36 @@ def create_data_asset(self, data_asset_params: DataAssetParams) -> DataAsset: return DataAsset.from_dict(res.json()) - def wait_until_ready(self, data_asset: DataAsset, polling_interval: int = 5) -> DataAsset: + def wait_until_ready( + self, data_asset: DataAsset, polling_interval: int = 5, timeout: int | None = None + ) -> DataAsset: """ Polls the given data asset until it reaches the 'Ready' or 'Failed' state. + + - `polling_interval` and `timeout` are in seconds """ if polling_interval < 5: raise ValueError( f"Polling interval {polling_interval} should be greater than or equal to 5" ) + if timeout is not None and timeout < polling_interval: + raise ValueError( + f"Timeout {timeout} should be greater than or equal to polling interval {polling_interval}" + ) + if timeout is not None and timeout < 0: + raise ValueError( + f"Timeout {timeout} should be greater than or equal to 0 (seconds), or None" + ) + t0 = time() while True: da = self.get_data_asset(data_asset.id) if da.state in [DataAssetState.Ready, DataAssetState.Failed]: return da - + + if timeout is not None and (time() - t0) > timeout: + raise TimeoutError(f"Data asset {data_asset.id} was not ready within {timeout} seconds") + sleep(polling_interval) def delete_data_asset(self, data_asset_id: str): @@ -268,7 +289,9 @@ def archive_data_asset(self, data_asset_id: str, archive: bool): params={"archive": archive}, ) - def search_data_assets(self, search_params: DataAssetSearchParams) -> DataAssetSearchResults: + def search_data_assets( + self, search_params: DataAssetSearchParams + ) -> DataAssetSearchResults: res = self.client.post("data_assets/search", json=search_params.to_dict()) return DataAssetSearchResults.from_dict(res.json()) From ead041c12d961f9dd5f7da75d43ea20088c88d3a Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Wed, 16 Oct 2024 13:55:59 -0700 Subject: [PATCH 2/6] Relax type annotations for arguments to `time` functions --- src/codeocean/computation.py | 2 +- src/codeocean/data_asset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/codeocean/computation.py b/src/codeocean/computation.py index cb58e81..343ab09 100644 --- a/src/codeocean/computation.py +++ b/src/codeocean/computation.py @@ -142,7 +142,7 @@ def run_capsule(self, run_params: RunParams) -> Computation: return Computation.from_dict(res.json()) def wait_until_completed( - self, computation: Computation, polling_interval: int = 5, timeout: int | None = None, + self, computation: Computation, polling_interval: float = 5, timeout: float | None = None, ) -> Computation: """ Polls the given computation until it reaches the 'Completed' or 'Failed' diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index aea860c..4bce0b2 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -243,7 +243,7 @@ def create_data_asset(self, data_asset_params: DataAssetParams) -> DataAsset: return DataAsset.from_dict(res.json()) def wait_until_ready( - self, data_asset: DataAsset, polling_interval: int = 5, timeout: int | None = None + self, data_asset: DataAsset, polling_interval: float = 5, timeout: float | None = None ) -> DataAsset: """ Polls the given data asset until it reaches the 'Ready' or 'Failed' state. From a94d8a320159436fdfbbcee79d6f864a519d4e11 Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Wed, 16 Oct 2024 15:31:34 -0700 Subject: [PATCH 3/6] Lint whitespace --- src/codeocean/computation.py | 6 +++--- src/codeocean/data_asset.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/codeocean/computation.py b/src/codeocean/computation.py index 343ab09..fc9c028 100644 --- a/src/codeocean/computation.py +++ b/src/codeocean/computation.py @@ -146,8 +146,8 @@ def wait_until_completed( ) -> Computation: """ Polls the given computation until it reaches the 'Completed' or 'Failed' - state. - + state. + - `polling_interval` and `timeout` are in seconds """ if polling_interval < 5: @@ -171,7 +171,7 @@ def wait_until_completed( if timeout is not None and (time() - t0) > timeout: raise TimeoutError(f"Computation {computation.id} did not complete within {timeout} seconds") - + sleep(polling_interval) def list_computation_results(self, computation_id: str, path: str = "") -> Folder: diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 4bce0b2..2aac31e 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -247,7 +247,7 @@ def wait_until_ready( ) -> DataAsset: """ Polls the given data asset until it reaches the 'Ready' or 'Failed' state. - + - `polling_interval` and `timeout` are in seconds """ if polling_interval < 5: @@ -268,10 +268,10 @@ def wait_until_ready( if da.state in [DataAssetState.Ready, DataAssetState.Failed]: return da - + if timeout is not None and (time() - t0) > timeout: raise TimeoutError(f"Data asset {data_asset.id} was not ready within {timeout} seconds") - + sleep(polling_interval) def delete_data_asset(self, data_asset_id: str): From 6e23c2f6477e06bef2e7d4ea5c0ee481bb865512 Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Thu, 17 Oct 2024 07:09:14 -0700 Subject: [PATCH 4/6] Format import groups --- src/codeocean/data_asset.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 2aac31e..e56494f 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -1,14 +1,13 @@ from __future__ import annotations +from dataclasses_json import dataclass_json from dataclasses import dataclass +from requests_toolbelt.sessions import BaseUrlSession from time import sleep, time from typing import Optional -from dataclasses_json import dataclass_json -from requests_toolbelt.sessions import BaseUrlSession - -from codeocean.components import Permissions, SearchFilter, SortOrder -from codeocean.computation import Param, PipelineProcess +from codeocean.components import SortOrder, SearchFilter, Permissions +from codeocean.computation import PipelineProcess, Param from codeocean.enum import StrEnum From b0a5b9588e1714cf87b2c827a2d5ea442e130c73 Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Thu, 17 Oct 2024 07:16:44 -0700 Subject: [PATCH 5/6] Format line breaks --- src/codeocean/computation.py | 8 +++++--- src/codeocean/data_asset.py | 5 ++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/codeocean/computation.py b/src/codeocean/computation.py index fc9c028..723834b 100644 --- a/src/codeocean/computation.py +++ b/src/codeocean/computation.py @@ -142,11 +142,13 @@ def run_capsule(self, run_params: RunParams) -> Computation: return Computation.from_dict(res.json()) def wait_until_completed( - self, computation: Computation, polling_interval: float = 5, timeout: float | None = None, + self, + computation: Computation, + polling_interval: float = 5, + timeout: Optional[float] = None, ) -> Computation: """ - Polls the given computation until it reaches the 'Completed' or 'Failed' - state. + Polls the given computation until it reaches the 'Completed' or 'Failed' state. - `polling_interval` and `timeout` are in seconds """ diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index e56494f..c9126ca 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -242,7 +242,10 @@ def create_data_asset(self, data_asset_params: DataAssetParams) -> DataAsset: return DataAsset.from_dict(res.json()) def wait_until_ready( - self, data_asset: DataAsset, polling_interval: float = 5, timeout: float | None = None + self, + data_asset: DataAsset, + polling_interval: float = 5, + timeout: float | None = None, ) -> DataAsset: """ Polls the given data asset until it reaches the 'Ready' or 'Failed' state. From 56d99ec7484b32f0ead4e3f65dbc80ec39095f71 Mon Sep 17 00:00:00 2001 From: bjhardcastle Date: Thu, 17 Oct 2024 10:39:00 -0700 Subject: [PATCH 6/6] Revert auto-formatting --- src/codeocean/data_asset.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index c9126ca..69f8cc7 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -227,12 +227,8 @@ def get_data_asset(self, data_asset_id: str) -> DataAsset: return DataAsset.from_dict(res.json()) - def update_metadata( - self, data_asset_id: str, update_params: DataAssetUpdateParams - ) -> DataAsset: - res = self.client.put( - f"data_assets/{data_asset_id}", json=update_params.to_dict() - ) + def update_metadata(self, data_asset_id: str, update_params: DataAssetUpdateParams) -> DataAsset: + res = self.client.put(f"data_assets/{data_asset_id}", json=update_params.to_dict()) return DataAsset.from_dict(res.json()) @@ -291,9 +287,7 @@ def archive_data_asset(self, data_asset_id: str, archive: bool): params={"archive": archive}, ) - def search_data_assets( - self, search_params: DataAssetSearchParams - ) -> DataAssetSearchResults: + def search_data_assets(self, search_params: DataAssetSearchParams) -> DataAssetSearchResults: res = self.client.post("data_assets/search", json=search_params.to_dict()) return DataAssetSearchResults.from_dict(res.json())