def wait_for_ready_queued_resource()

in xlml/utils/tpu.py [0:0]

• 17 lines of code
• 3 McCabe index (conditional complexity)

  def wait_for_ready_queued_resource(qualified_name: str):
    creds, _ = google.auth.default()
    client = tpu_api.TpuClient(credentials=creds)

    qr = client.get_queued_resource(name=qualified_name)
    state = qr.state.state
    logging.info(f'Queued resource state: {state.name}')
    if qr.state.state == tpu_api.QueuedResourceState.State.ACTIVE:
      return True
    elif qr.state.state in [
        tpu_api.QueuedResourceState.State.CREATING,
        tpu_api.QueuedResourceState.State.WAITING_FOR_RESOURCES,
        tpu_api.QueuedResourceState.State.ACCEPTED,
        tpu_api.QueuedResourceState.State.PROVISIONING,
    ]:
      return False
    else:
      raise RuntimeError(f'Bad queued resource state {state.name}')