def WaitForOperation(operation_service, operation,
max_retries=None,
retry_interval=None):
"""Wait until the operation is complete or times out.
Args:
operation_service: The apitools service type for operations
operation: The operation resource to wait on
max_retries: Maximum number of times to poll the operation
retry_interval: Frequency of polling
Returns:
The operation resource when it has completed
Raises:
OperationTimeoutError: when the operation polling times out
OperationError: when the operation completed with an error
"""
if max_retries is None:
max_retries = constants.DEFAULT_OPERATION_MAX_RETRIES
if retry_interval is None:
retry_interval = constants.DEFAULT_OPERATION_RETRY_INTERVAL
completed_operation = _PollUntilDone(operation_service, operation,
max_retries, retry_interval)
if not completed_operation:
raise OperationTimeoutError(('Operation [{0}] timed out. This operation '
'may still be underway.').format(
operation.name))
if completed_operation.error:
raise OperationError(requests.ExtractErrorMessage(
encoding.MessageToPyValue(completed_operation.error)))
return completed_operation
def WaitForOperation(operation_service, operation, registry=None):
"""Wait until the operation is complete or times out.
Args:
operation_service: The apitools service type for operations
operation: The operation resource to wait on
registry: A resource registry to use for operation get requests.
Returns:
The operation resource when it has completed
Raises:
OperationTimeoutError: when the operation polling times out
OperationError: when the operation completed with an error
"""
if operation.done:
return operation
if not registry:
registry = resources.REGISTRY
request = registry.Parse(operation.name.split('/')[-1],
collection='ml.projects.operations').Request()
try:
operation = retry.Retryer(max_wait_ms=60 * 60 * 1000).RetryOnResult(
operation_service.Get,
args=(request, ),
should_retry_if=lambda op, _: not op.done,
sleep_ms=5000)
if operation.error:
raise OperationError(
requests.ExtractErrorMessage(
encoding.MessageToPyValue(operation.error)))
return operation
except retry.WaitException:
raise OperationTimeoutError(
'Operation [{0}] timed out. This operation may still be underway.'.
format(operation.name))
def IsDone(self, operation):
"""Overrides."""
if operation.done:
log.debug('Operation [{0}] complete. Result: {1}'.format(
operation.name,
json.dumps(encoding.MessageToDict(operation), indent=4)))
if operation.error:
raise OperationError(
requests.ExtractErrorMessage(
encoding.MessageToPyValue(operation.error)))
return True
log.debug('Operation [{0}] not complete. Waiting to retry.'.format(
operation.name))
return False