def handler(*args, **kwargs):
"""
Wraps rpc errors as Flyte exceptions and handles authentication the client.
:param args:
:param kwargs:
:return:
"""
max_retries = 3
max_wait_time = 1000
try:
for i in range(max_retries):
try:
return fn(*args, **kwargs)
except _RpcError as e:
if e.code() == _GrpcStatusCode.UNAUTHENTICATED:
# Always retry auth errors.
if i == (max_retries - 1):
# Exit the loop and wrap the authentication error.
raise _user_exceptions.FlyteAuthenticationException(_six.text_type(e))
refresh_handler_fn = _get_refresh_handler(_creds_config.AUTH_MODE.get())
refresh_handler_fn(args[0])
else:
# No more retries if retry=False or max_retries reached.
if (retry is False) or i == (max_retries - 1):
raise
else:
# Retry: Start with 200ms wait-time and exponentially back-off upto 1 second.
wait_time = min(200 * (2 ** i), max_wait_time)
_logging.error(f"Non-auth RPC error {e}, sleeping {wait_time}ms and retrying")
time.sleep(wait_time / 1000)
except _RpcError as e:
if e.code() == _GrpcStatusCode.ALREADY_EXISTS:
raise _user_exceptions.FlyteEntityAlreadyExistsException(_six.text_type(e))
else:
raise
def handler(*args, **kwargs):
"""
Wraps rpc errors as Flyte exceptions and handles authentication the client.
:param args:
:param kwargs:
:return:
"""
retries = 2
try:
for i in range(retries):
try:
return fn(*args, **kwargs)
except _RpcError as e:
if e.code() == _GrpcStatusCode.UNAUTHENTICATED:
if i == (retries - 1):
# Exit the loop and wrap the authentication error.
raise _user_exceptions.FlyteAuthenticationException(
_six.text_type(e))
refresh_handler_fn = _get_refresh_handler(
_creds_config.AUTH_MODE.get())
refresh_handler_fn(args[0])
else:
raise
except _RpcError as e:
if e.code() == _GrpcStatusCode.ALREADY_EXISTS:
raise _user_exceptions.FlyteEntityAlreadyExistsException(
_six.text_type(e))
else:
raise
def _refresh_credentials_from_command(flyte_client):
"""
This function is used when the configuration value for AUTH_MODE is set to 'external_process'.
It reads an id token generated by an external process started by running the 'command'.
:param flyte_client: RawSynchronousFlyteClient
:return:
"""
command = _COMMAND.get()
cli_logger.debug(
"Starting external process to generate id token. Command {}".format(
command))
try:
output = subprocess.run(command,
capture_output=True,
text=True,
check=True)
except subprocess.CalledProcessError as e:
cli_logger.error(
"Failed to generate token from command {}".format(command))
raise _user_exceptions.FlyteAuthenticationException(
"Problems refreshing token with command: " + _six.text_type(e))
flyte_client.set_access_token(output.stdout.strip())
