Skip to content

Commit

Permalink
add wait_function
Browse files Browse the repository at this point in the history
Signed-off-by: SdgJlbl <sarah.diot-girard@owkin.com>
  • Loading branch information
SdgJlbl committed Feb 16, 2024
1 parent cc1dada commit e74e4cc
Showing 1 changed file with 37 additions and 0 deletions.
37 changes: 37 additions & 0 deletions substra/sdk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1076,6 +1076,43 @@ def wait_task(
timeout=timeout,
)

@logit
def wait_function(
self, key: str, *, timeout: Optional[float] = None, polling_period: float = 1.0, raise_on_failure: bool = True
) -> models.Task:
"""Wait for the build of the given function to finish.
It is considered finished when the status is ready, failed or cancelled.
Args:
key (str): the key of the task to wait for.
timeout (float, optional): maximum time to wait, in seconds. If set to None, will hang until completion.
polling_period (float): time to wait between two checks, in seconds. Defaults to 2.0.
raise_on_failure (bool): whether to raise an exception if the execution fails. Defaults to True.
Returns:
models.Task: the task after completion
Raises:
exceptions.FutureFailureError: The task failed or have been cancelled.
exceptions.FutureTimeoutError: The task took more than the duration set in the timeout to complete.
Not raised when `timeout == None`
"""
asset_getter = self.get_function
status_canceled = models.FunctionStatus.canceled.value
status_failed = models.FunctionStatus.failed.value
statuses_stopped = (models.FunctionStatus.ready.value, models.FunctionStatus.canceled.value)
return self._wait(
key=key,
asset_getter=asset_getter,
polling_period=polling_period,
raise_on_failure=raise_on_failure,
status_canceled=status_canceled,
status_failed=status_failed,
statuses_stopped=statuses_stopped,
timeout=timeout,
)

def _wait(
self,
*,
Expand Down

0 comments on commit e74e4cc

Please sign in to comment.