make api key prompting be more captive
@@ -21,6 +21,9 @@ from .wandb_settings import Settings
from ..apis import InternalApi
+MAX_API_KEY_PROMPTS = 5
+
+
def _handle_host_wandb_setting(host: Optional[str], cloud: bool = False) -> None:
"""Write the host parameter from wandb.login or wandb login to
the global settings file so that it is used automatically by
@@ -82,6 +85,7 @@ class _WandbLogin(object):
self._wl = None
self._key = None
self._relogin = None
+ self._num_api_key_prompts = 0
def setup(self, kwargs):
self.kwargs = kwargs
@@ -176,6 +180,7 @@ class _WandbLogin(object):
def _prompt_api_key(self) -> Tuple[Optional[str], ApiKeyStatus]:
api = Api(self._settings)
+ self._num_api_key_prompts += 1
try:
key = apikey.prompt_api_key(
self._settings,
@@ -183,6 +188,17 @@ class _WandbLogin(object):
no_offline=self._settings.force if self._settings else None,
no_create=self._settings.force if self._settings else None,
)
+ except ValueError as e:
+ # invalid key provided, try again
+ wandb.termerror(e.args[0])
+ if self._num_api_key_prompts < MAX_API_KEY_PROMPTS:
+ return self._prompt_api_key()
+ else:
+ wandb.termerror(
+ f"Failed to get API key after {MAX_API_KEY_PROMPTS} attempts."
+ )
+ return None, ApiKeyStatus.DISABLED
+
except TimeoutError:
wandb.termlog("W&B disabled due to login timeout.")
return None, ApiKeyStatus.DISABLED
