Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support multiple tokens locally #2549

Merged
merged 20 commits into from
Oct 1, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 158 additions & 19 deletions src/huggingface_hub/_login.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import os
import subprocess
import warnings
from functools import partial
from getpass import getpass
from pathlib import Path
Expand All @@ -33,7 +34,13 @@
set_git_credential,
unset_git_credential,
)
from .utils._token import _get_token_from_environment, _get_token_from_google_colab
from .utils._auth_profiles import _read_profiles, _save_profiles
from .utils._token import (
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
_get_token_from_environment,
_get_token_from_google_colab,
_get_token_from_profile,
_save_token_to_profile,
)


logger = logging.get_logger(__name__)
Expand All @@ -49,6 +56,7 @@

def login(
token: Optional[str] = None,
profile_name: Optional[str] = None,
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
add_to_git_credential: bool = False,
new_session: bool = True,
write_permission: bool = False,
Expand Down Expand Up @@ -82,6 +90,8 @@ def login(
Args:
token (`str`, *optional*):
User access token to generate from https://huggingface.co/settings/token.
profile_name (`str`, *optional*):
Name of the profile to add or update. If `None`, will add or update the "default" profile.
add_to_git_credential (`bool`, defaults to `False`):
If `True`, token will be set as git credential. If no git credential helper
is configured, a warning will be displayed to the user. If `token` is `None`,
Expand All @@ -108,31 +118,58 @@ def login(
"`--add-to-git-credential` if using via `huggingface-cli` if "
"you want to set the git credential as well."
)
_login(token, add_to_git_credential=add_to_git_credential, write_permission=write_permission)
_login(
token,
add_to_git_credential=add_to_git_credential,
write_permission=write_permission,
profile_name=profile_name,
)
elif is_notebook():
notebook_login(new_session=new_session, write_permission=write_permission)
else:
interpreter_login(new_session=new_session, write_permission=write_permission)


def logout() -> None:
def logout(profile_name: Optional[str] = None, all: bool = False) -> None:
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
"""Logout the machine from the Hub.

Token is deleted from the machine and removed from git credential.

Args:
profile_name (`str`, *optional*):
Name of the profile to logout from. If `None`, will logout from the active profile.
all (`bool`, defaults to `False`):
If `True`, all profiles are deleted.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the profile name is not found.
"""
if get_token() is None:
print("Not logged in!")
return

if all:
# Delete all profiles and token
for file_path in (constants.HF_TOKEN_PATH, constants.HF_PROFILES_PATH):
try:
Path(file_path).unlink()
except FileNotFoundError:
pass

print("Successfully logged out from all profiles.")

if profile_name is None:
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
# Logout from the active profile, i.e. delete the token file
try:
Path(constants.HF_TOKEN_PATH).unlink()
except FileNotFoundError:
pass
else:
_logout_from_profile(profile_name)
print(f"Successfully removed profile: `{profile_name}`")
return
# Delete token from git credentials
unset_git_credential()

# Delete token file
try:
Path(constants.HF_TOKEN_PATH).unlink()
except FileNotFoundError:
pass

# Check if still logged in
if _get_token_from_google_colab() is not None:
raise EnvironmentError(
Expand All @@ -144,10 +181,60 @@ def logout() -> None:
"Token has been deleted from your machine but you are still logged in.\n"
"To log out, you must clear out both `HF_TOKEN` and `HUGGING_FACE_HUB_TOKEN` environment variables."
)

print("Successfully logged out.")


def auth_switch(profile_name: str) -> None:
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
"""Switch to a different profile.

Args:
profile_name (`str`):
Name of the profile to switch to.

Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the profile name is not found.
"""
token = _get_token_from_profile(profile_name)
if token:
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
# Write token to HF_TOKEN_PATH
_set_active_profile(profile_name)
print(f"Switched to profile: {profile_name}")
if _get_token_from_environment():
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
warnings.warn(
"The environment variable `HF_TOKEN` is set and will override " "the token from the profile."
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
)
else:
raise ValueError(f"Profile {profile_name} not found in {constants.HF_PROFILES_PATH}")


def auth_list() -> None:
"""List all available profiles."""
profiles = _read_profiles()
current_profile = None

if not profiles.sections():
print("No profiles found.")
return
for profile in profiles.sections():
if profiles.get(profile, "hf_token") == get_token():
current_profile = profile
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
# Print header
print(f"{'Profile Name':^20} {'Token':^20}")
print(f"{'-'*20} {'-'*20}")

# Print profiles
for profile in profiles.sections():
token = profiles.get(profile, "hf_token", fallback="<not set>")
masked_token = f"{token[:4]}{'*' * 8}" if token != "<not set>" else token
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
is_current = "*" if profile == current_profile else " "

print(f"{is_current} {profile:^19} {masked_token:^20}")

if _get_token_from_environment():
print("\nNote: Environment variable `HF_TOKEN` is set and is the current active token.")
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved


###
# Interpreter-based login (text)
###
Expand Down Expand Up @@ -189,9 +276,15 @@ def interpreter_login(new_session: bool = True, write_permission: bool = False)
if os.name == "nt":
print("Token can be pasted using 'Right-Click'.")
token = getpass("Enter your token (input will not be visible): ")
profile_name = input("Enter profile name (default: 'default'): ") or "default"
add_to_git_credential = _ask_for_confirmation_no_tui("Add token as git credential?")

_login(token=token, add_to_git_credential=add_to_git_credential, write_permission=write_permission)
_login(
token=token,
profile_name=profile_name,
add_to_git_credential=add_to_git_credential,
write_permission=write_permission,
)


###
Expand Down Expand Up @@ -296,7 +389,12 @@ def login_token_event(t, write_permission: bool = False):
###


def _login(token: str, add_to_git_credential: bool, write_permission: bool = False) -> None:
def _login(
token: str,
add_to_git_credential: bool,
write_permission: bool = False,
profile_name: Optional[str] = None,
) -> None:
from .hf_api import get_token_permission # avoid circular import

if token.startswith("api_org"):
Expand All @@ -321,13 +419,54 @@ def _login(token: str, add_to_git_credential: bool, write_permission: bool = Fal
)
else:
print("Token has not been saved to git credential helper.")
profile_name = profile_name or "default"
# Save token to profiles file
_save_token_to_profile(token, profile_name)
# Set active profile
_set_active_profile(profile_name)
print("Login successful.")
if _get_token_from_environment():
print("Note: Environment variable`HF_TOKEN` is set and is the current active token.")
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved
else:
print(f"The current active profile is: `{profile_name}`")


# Save token
path = Path(constants.HF_TOKEN_PATH)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(token)
print(f"Your token has been saved to {constants.HF_TOKEN_PATH}")
print("Login successful")
def _logout_from_profile(profile_name: str) -> None:
"""Logout from a profile.

Args:
profile_name (`str`):
The name of the profile to logout from.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the profile name is not found.
"""
config = _read_profiles()
if config.get(profile_name, "hf_token") == get_token():
warnings.warn(f"Active profile {profile_name} will been deleted.")
if profile_name in config:
config.remove_section(profile_name)
_save_profiles(config)
else:
raise ValueError(f"Profile {profile_name} not found in {constants.HF_PROFILES_PATH}")
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved


def _set_active_profile(profile_name: str) -> None:
"""Set the active profile.

Args:
profile_name (`str`):
The name of the profile to set as active.
"""
token = _get_token_from_profile(profile_name)
if token:
# Write token to HF_TOKEN_PATH
path = Path(constants.HF_TOKEN_PATH)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(token)
print(f"Your token has been saved to {constants.HF_TOKEN_PATH}")
else:
raise ValueError(f"Profile {profile_name} not found in {constants.HF_PROFILES_PATH}")
hanouticelina marked this conversation as resolved.
Show resolved Hide resolved


def _current_token_okay(write_permission: bool = False):
Expand Down
54 changes: 46 additions & 8 deletions src/huggingface_hub/commands/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,15 @@
from requests.exceptions import HTTPError

from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.constants import (
ENDPOINT,
REPO_TYPES,
REPO_TYPES_URL_PREFIXES,
SPACES_SDK_TYPES,
)
from huggingface_hub.constants import ENDPOINT, REPO_TYPES, REPO_TYPES_URL_PREFIXES, SPACES_SDK_TYPES
from huggingface_hub.hf_api import HfApi

from .._login import ( # noqa: F401 # for backward compatibility # noqa: F401 # for backward compatibility
NOTEBOOK_LOGIN_PASSWORD_HTML,
NOTEBOOK_LOGIN_TOKEN_HTML_END,
NOTEBOOK_LOGIN_TOKEN_HTML_START,
auth_list,
auth_switch,
login,
logout,
notebook_login,
Expand All @@ -46,6 +43,11 @@ def register_subcommand(parser: _SubParsersAction):
type=str,
help="Token generated from https://huggingface.co/settings/tokens",
)
login_parser.add_argument(
"--profile-name",
type=str,
help="Optional: Name of the profile to log in to.",
)
login_parser.add_argument(
"--add-to-git-credential",
action="store_true",
Expand All @@ -54,9 +56,31 @@ def register_subcommand(parser: _SubParsersAction):
login_parser.set_defaults(func=lambda args: LoginCommand(args))
whoami_parser = parser.add_parser("whoami", help="Find out which huggingface.co account you are logged in as.")
whoami_parser.set_defaults(func=lambda args: WhoamiCommand(args))

logout_parser = parser.add_parser("logout", help="Log out")
logout_parser.add_argument(
"--profile-name",
type=str,
help="Optional: Name of the profile to log out from.",
)
logout_parser.add_argument(
"--all",
action="store_true",
help="Optional: Log out from all profiles.",
)
logout_parser.set_defaults(func=lambda args: LogoutCommand(args))

auth_parser = parser.add_parser("auth", help="Other authentication related commands")
auth_subparsers = auth_parser.add_subparsers(help="Authentication subcommands")
auth_switch_parser = auth_subparsers.add_parser("switch", help="Switch between profiles")
auth_switch_parser.add_argument(
"--profile-name",
type=str,
help="Optional: Name of the profile to switch to.",
)
auth_switch_parser.set_defaults(func=lambda args: AuthSwitchCommand(args))
auth_list_parser = auth_subparsers.add_parser("list", help="List all profiles")
auth_list_parser.set_defaults(func=lambda args: AuthListCommand(args))
# new system: git-based repo system
repo_parser = parser.add_parser("repo", help="{create} Commands to interact with your huggingface.co repos.")
repo_subparsers = repo_parser.add_subparsers(help="huggingface.co repos related commands")
Expand Down Expand Up @@ -95,12 +119,26 @@ def __init__(self, args):

class LoginCommand(BaseUserCommand):
def run(self):
login(token=self.args.token, add_to_git_credential=self.args.add_to_git_credential)
login(
token=self.args.token,
profile_name=self.args.profile_name,
add_to_git_credential=self.args.add_to_git_credential,
)


class LogoutCommand(BaseUserCommand):
def run(self):
logout()
logout(profile_name=self.args.profile_name, all=self.args.all)


class AuthSwitchCommand(BaseUserCommand):
def run(self):
auth_switch(profile_name=self.args.profile_name)


class AuthListCommand(BaseUserCommand):
def run(self):
auth_list()


class WhoamiCommand(BaseUserCommand):
Expand Down
2 changes: 1 addition & 1 deletion src/huggingface_hub/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def _as_int(value: Optional[str]) -> Optional[int]:
# See https://github.com/huggingface/huggingface_hub/issues/1232
_OLD_HF_TOKEN_PATH = os.path.expanduser("~/.huggingface/token")
HF_TOKEN_PATH = os.environ.get("HF_TOKEN_PATH", os.path.join(HF_HOME, "token"))

HF_PROFILES_PATH = os.environ.get("HF_PROFILES_PATH", os.path.join(HF_HOME, "profiles"))

if _staging_mode:
# In staging mode, we use a different cache to ensure we don't mix up production and staging data or tokens
Expand Down
Loading
Loading