Login

zenml.login

ZenML login utilities.

Modules

credentials

ZenML login credentials models.

Classes

APIToken

Bases: BaseModel

Cached API Token.

Attributes

expired: bool property

Check if the token is expired.

Returns:

Name	Type	Description
bool	bool	True if the token is expired, False otherwise.

expires_at_with_leeway: Optional[datetime] property

Get the token expiration time with leeway.

Returns:

Type	Description
Optional[datetime]	The token expiration time with leeway.

ServerCredentials

Bases: BaseModel

Cached Server Credentials.

Attributes

api_hyperlink: str property

Get the hyperlink to the ZenML OpenAPI dashboard for this workspace.

Returns:

Type	Description
str	The hyperlink to the ZenML OpenAPI dashboard for this workspace.

auth_status: str property

Get the authentication status.

Returns:

Type	Description
str	The authentication status.

dashboard_hyperlink: str property

Get the hyperlink to the ZenML dashboard for this workspace.

Returns:

Type	Description
str	The hyperlink to the ZenML dashboard for this workspace.

dashboard_organization_url: str property

Get the URL to the ZenML Pro dashboard for this workspace's organization.

Returns:

Type	Description
str	The URL to the ZenML Pro dashboard for this workspace's organization.

dashboard_url: str property

Get the URL to the ZenML dashboard for this server.

Returns:

Type	Description
str	The URL to the ZenML dashboard for this server.

expires_at: str property

Get the expiration time of the token as a string.

Returns:

Type	Description
str	The expiration time of the token as a string.

expires_in: str property

Get the time remaining until the token expires.

Returns:

Type	Description
str	The time remaining until the token expires.

id: str property

Get the server identifier.

Returns:

Type	Description
str	The server identifier.

is_available: bool property

Check if the server is available (running and authenticated).

Returns:

Type	Description
bool	True if the server is available, False otherwise.

organization_hyperlink: str property

Get the hyperlink to the ZenML Pro dashboard for this server's organization.

Returns:

Type	Description
str	The hyperlink to the ZenML Pro dashboard for this server's
str	organization.

organization_id_hyperlink: str property

Get the hyperlink to the ZenML Pro dashboard for this workspace's organization using its ID.

Returns:

Type	Description
str	The hyperlink to the ZenML Pro dashboard for this workspace's
str	organization using its ID.

organization_name_hyperlink: str property

Get the hyperlink to the ZenML Pro dashboard for this server's organization using its name.

Returns:

Type	Description
str	The hyperlink to the ZenML Pro dashboard for this server's
str	organization using its name.

server_id_hyperlink: str property

Get the hyperlink to the ZenML dashboard for this server using its ID.

Returns:

Type	Description
str	The hyperlink to the ZenML dashboard for this server using its ID.

server_name_hyperlink: str property

Get the hyperlink to the ZenML dashboard for this server using its name.

Returns:

Type	Description
str	The hyperlink to the ZenML dashboard for this server using its name.

type: ServerType property

Get the server type.

Returns:

Type	Description
ServerType	The server type.

Functions

update_server_info(server_info: Union[ServerModel, WorkspaceRead]) -> None

Update with server information received from the server itself or from a ZenML Pro workspace descriptor.

Parameters:

Name	Type	Description	Default
server_info	Union[ServerModel, WorkspaceRead]	The server information to update with.	required

Source code in src/zenml/login/credentials.py

def update_server_info(
    self, server_info: Union[ServerModel, WorkspaceRead]
) -> None:
    """Update with server information received from the server itself or from a ZenML Pro workspace descriptor.

    Args:
        server_info: The server information to update with.
    """
    if isinstance(server_info, ServerModel):
        # The server ID doesn't change during the lifetime of the server
        self.server_id = self.server_id or server_info.id
        # All other attributes can change during the lifetime of the server
        self.deployment_type = server_info.deployment_type
        server_name = (
            server_info.pro_workspace_name
            or server_info.metadata.get("workspace_name")
            or server_info.name
        )
        if server_name:
            self.server_name = server_name
        if server_info.pro_organization_id:
            self.organization_id = server_info.pro_organization_id
        if server_info.pro_workspace_id:
            self.server_id = server_info.pro_workspace_id
        if server_info.pro_organization_name:
            self.organization_name = server_info.pro_organization_name
        if server_info.pro_workspace_name:
            self.workspace_name = server_info.pro_workspace_name
        if server_info.pro_api_url:
            self.pro_api_url = server_info.pro_api_url
        if server_info.pro_dashboard_url:
            self.pro_dashboard_url = server_info.pro_dashboard_url
        self.version = server_info.version or self.version
        # The server information was retrieved from the server itself, so we
        # can assume that the server is available
        self.status = "available"
    else:
        self.deployment_type = ServerDeploymentType.CLOUD
        self.server_id = server_info.id
        self.server_name = server_info.name
        self.organization_name = server_info.organization_name
        self.organization_id = server_info.organization_id
        self.workspace_name = server_info.name
        self.workspace_id = server_info.id
        self.status = server_info.status
        self.version = server_info.version

ServerType

Bases: StrEnum

The type of server.

Functions

credentials_store

ZenML login credentials store support.

Classes

CredentialsStore()

Login credentials store.

This is a singleton object that maintains a cache of all API tokens and API keys that are configured for the ZenML servers that the client connects to throughout its lifetime.

The cache is persistent and it is backed by a credentials.yaml YAML file kept in the global configuration location. The Credentials Store cache is populated mainly in the following ways:

1. when the user runs `zenml login` to authenticate to a ZenML Pro
server, the ZenML Pro API token fetched from the web login flow is
stored in the Credentials Store.
2. when the user runs `zenml login` to authenticate to a regular ZenML
server with the web login flow, the ZenML server API token fetched
through the web login flow is stored in the Credentials Store
3. when the user runs `zenml login` to authenticate to any ZenML server
using an API key, the API key is stored in the Credentials Store
4. when the REST zen store is initialized, it starts up not yet
authenticated. Then, if/when it needs to authenticate or re-authenticate
to the remote server, it will use whatever form of credentials it finds
in the Credentials Store, in order of priority:

    * if it finds an API token that is not expired (e.g. authorized
    device API tokens fetched through the web login flow or short-lived
    session API tokens fetched through some other means of
    authentication), it will use that to authenticate
    * for ZenML servers that use an API key to authenticate, it will use
    that to fetch a short-lived ZenML Pro server API token that it also
    stores in the Credentials Store
    * for ZenML Pro servers, it exchanges the longer-lived ZenML Pro API
    token into a short lived ZenML Pro server API token

Alongside credentials, the Credentials Store is also used to store additional server information: * ZenML Pro workspace information populated by the zenml login command * ZenML server information populated by the REST zen store by fetching the server's information endpoint after authenticating

Initializes the login credentials store with values loaded from the credentials YAML file.

CredentialsStore is a singleton class: only one instance can exist. Calling this constructor multiple times will always yield the same instance.

Source code in src/zenml/login/credentials_store.py

def __init__(self) -> None:
    """Initializes the login credentials store with values loaded from the credentials YAML file.

    CredentialsStore is a singleton class: only one instance can exist.
    Calling this constructor multiple times will always yield the same
    instance.
    """
    self.credentials = {}
    self._load_credentials()

Functions

can_login(server_url: str) -> bool

Check if credentials to login to the given server exist.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to check the authentication.	required

Returns:

Type	Description
bool	True if the credentials store contains credentials that can be used
bool	to login to the given server URL, False otherwise.

Source code in src/zenml/login/credentials_store.py

def can_login(self, server_url: str) -> bool:
    """Check if credentials to login to the given server exist.

    Args:
        server_url: The server URL for which to check the authentication.

    Returns:
        True if the credentials store contains credentials that can be used
        to login to the given server URL, False otherwise.
    """
    self.check_and_reload_from_file()
    credentials = self.get_credentials(server_url)
    if not credentials:
        return False

    if credentials.api_key is not None:
        return True
    elif (
        credentials.username is not None
        and credentials.password is not None
    ):
        return True
    elif credentials.type == ServerType.PRO:
        pro_api_url = credentials.pro_api_url or ZENML_PRO_API_URL
        pro_token = self.get_pro_token(pro_api_url, allow_expired=False)
        if pro_token:
            return True

    return False

check_and_reload_from_file() -> None

Check if the credentials file has been modified and reload it if necessary.

Source code in src/zenml/login/credentials_store.py

def check_and_reload_from_file(self) -> None:
    """Check if the credentials file has been modified and reload it if necessary."""
    if handle_bool_env_var(ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING):
        return
    if not self.last_modified_time:
        return
    credentials_file = self._credentials_file
    try:
        last_modified_time = os.path.getmtime(credentials_file)
    except FileNotFoundError:
        # The credentials file has been deleted
        self.last_modified_time = None
        return
    if last_modified_time != self.last_modified_time:
        self._load_credentials()

clear_all_pro_tokens(pro_api_url: Optional[str] = None) -> List[ServerCredentials]

Delete all tokens from the store for ZenML Pro servers connected to a given API server.

Parameters:

Name	Type	Description	Default
pro_api_url	Optional[str]	The URL of the ZenML Pro API server.	None

Returns:

Type	Description
List[ServerCredentials]	A list of the credentials that were cleared.

Source code in src/zenml/login/credentials_store.py

def clear_all_pro_tokens(
    self, pro_api_url: Optional[str] = None
) -> List[ServerCredentials]:
    """Delete all tokens from the store for ZenML Pro servers connected to a given API server.

    Args:
        pro_api_url: The URL of the ZenML Pro API server.

    Returns:
        A list of the credentials that were cleared.
    """
    credentials_to_clear = []
    for server_url, server in self.credentials.copy().items():
        if (
            server.type == ServerType.PRO
            and server.pro_api_url
            and (pro_api_url is None or server.pro_api_url == pro_api_url)
        ):
            if server.api_key:
                continue
            self.clear_token(server_url)
            credentials_to_clear.append(server)
    return credentials_to_clear

clear_credentials(server_url: str) -> None

Delete all credentials from the store for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to delete the credentials.	required

Source code in src/zenml/login/credentials_store.py

def clear_credentials(self, server_url: str) -> None:
    """Delete all credentials from the store for a specific server URL.

    Args:
        server_url: The server URL for which to delete the credentials.
    """
    self.check_and_reload_from_file()

    if server_url in self.credentials:
        del self.credentials[server_url]
        self._save_credentials()

clear_pro_credentials(pro_api_url: str) -> None

Delete the token from the store for a ZenML Pro API server.

Parameters:

Name	Type	Description	Default
pro_api_url	str	The URL of the ZenML Pro API server.	required

Source code in src/zenml/login/credentials_store.py

def clear_pro_credentials(self, pro_api_url: str) -> None:
    """Delete the token from the store for a ZenML Pro API server.

    Args:
        pro_api_url: The URL of the ZenML Pro API server.
    """
    self.clear_token(pro_api_url)

clear_token(server_url: str) -> None

Delete a token from the store for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to delete the token.	required

Source code in src/zenml/login/credentials_store.py

def clear_token(self, server_url: str) -> None:
    """Delete a token from the store for a specific server URL.

    Args:
        server_url: The server URL for which to delete the token.
    """
    self.check_and_reload_from_file()

    if server_url in self.credentials:
        credential = self.credentials[server_url]
        if (
            not credential.api_key
            and not credential.username
            and not credential.password is not None
        ):
            # Only delete the credential entry if there is no API key or
            # username/password to fall back on
            del self.credentials[server_url]
        else:
            credential.api_token = None
        self._save_credentials()

get_api_key(server_url: str) -> Optional[str]

Retrieve an API key from the credentials store for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to retrieve the API key.	required

Returns:

Type	Description
Optional[str]	The stored API key if it exists, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_api_key(self, server_url: str) -> Optional[str]:
    """Retrieve an API key from the credentials store for a specific server URL.

    Args:
        server_url: The server URL for which to retrieve the API key.

    Returns:
        The stored API key if it exists, None otherwise.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(server_url)
    if credential:
        return credential.api_key
    return None

get_credentials(server_url: str) -> Optional[ServerCredentials]

Retrieve the credentials for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to retrieve the credentials.	required

Returns:

Type	Description
Optional[ServerCredentials]	The stored credentials if they exist, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_credentials(self, server_url: str) -> Optional[ServerCredentials]:
    """Retrieve the credentials for a specific server URL.

    Args:
        server_url: The server URL for which to retrieve the credentials.

    Returns:
        The stored credentials if they exist, None otherwise.
    """
    self.check_and_reload_from_file()
    return self.credentials.get(server_url)

get_instance() -> Optional[CredentialsStore] classmethod

Get the singleton instance of the CredentialsStore.

Returns:

Type	Description
Optional[CredentialsStore]	The singleton instance of the CredentialsStore.

Source code in src/zenml/login/credentials_store.py

@classmethod
def get_instance(cls) -> Optional["CredentialsStore"]:
    """Get the singleton instance of the CredentialsStore.

    Returns:
        The singleton instance of the CredentialsStore.
    """
    return cast(CredentialsStore, cls._instance())

get_password(server_url: str) -> Tuple[Optional[str], Optional[str]]

Retrieve the username and password from the credentials store for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to retrieve the username and password.	required

Returns:

Type	Description
Tuple[Optional[str], Optional[str]]	The stored username and password if they exist, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_password(
    self, server_url: str
) -> Tuple[Optional[str], Optional[str]]:
    """Retrieve the username and password from the credentials store for a specific server URL.

    Args:
        server_url: The server URL for which to retrieve the username and
            password.

    Returns:
        The stored username and password if they exist, None otherwise.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(server_url)
    if credential:
        return credential.username, credential.password
    return None, None

get_pro_credentials(pro_api_url: str, allow_expired: bool = False) -> Optional[ServerCredentials]

Retrieve a valid token from the credentials store for a ZenML Pro API server.

Parameters:

Name	Type	Description	Default
pro_api_url	str	The URL of the ZenML Pro API server.	required
allow_expired	bool	Whether to allow expired tokens to be returned. The default behavior is to return None if a token does exist but is expired.	False

Returns:

Type	Description
Optional[ServerCredentials]	The stored credentials if they exist and are not expired, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_pro_credentials(
    self, pro_api_url: str, allow_expired: bool = False
) -> Optional[ServerCredentials]:
    """Retrieve a valid token from the credentials store for a ZenML Pro API server.

    Args:
        pro_api_url: The URL of the ZenML Pro API server.
        allow_expired: Whether to allow expired tokens to be returned. The
            default behavior is to return None if a token does exist but is
            expired.

    Returns:
        The stored credentials if they exist and are not expired, None otherwise.
    """
    credential = self.get_credentials(pro_api_url)
    if (
        credential
        and credential.type == ServerType.PRO_API
        and credential.api_token
        and (not credential.api_token.expired or allow_expired)
    ):
        return credential
    return None

get_pro_token(pro_api_url: str, allow_expired: bool = False) -> Optional[APIToken]

Retrieve a valid token from the credentials store for a ZenML Pro API server.

Parameters:

Name	Type	Description	Default
pro_api_url	str	The URL of the ZenML Pro API server.	required
allow_expired	bool	Whether to allow expired tokens to be returned. The default behavior is to return None if a token does exist but is expired.	False

Returns:

Type	Description
Optional[APIToken]	The stored token if it exists and is not expired, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_pro_token(
    self, pro_api_url: str, allow_expired: bool = False
) -> Optional[APIToken]:
    """Retrieve a valid token from the credentials store for a ZenML Pro API server.

    Args:
        pro_api_url: The URL of the ZenML Pro API server.
        allow_expired: Whether to allow expired tokens to be returned. The
            default behavior is to return None if a token does exist but is
            expired.

    Returns:
        The stored token if it exists and is not expired, None otherwise.
    """
    credential = self.get_pro_credentials(pro_api_url, allow_expired)
    if credential:
        return credential.api_token
    return None

get_token(server_url: str, allow_expired: bool = False) -> Optional[APIToken]

Retrieve a valid token from the credentials store for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which to retrieve the token.	required
allow_expired	bool	Whether to allow expired tokens to be returned. The default behavior is to return None if a token does exist but is expired.	False

Returns:

Type	Description
Optional[APIToken]	The stored token if it exists and is not expired, None otherwise.

Source code in src/zenml/login/credentials_store.py

def get_token(
    self, server_url: str, allow_expired: bool = False
) -> Optional[APIToken]:
    """Retrieve a valid token from the credentials store for a specific server URL.

    Args:
        server_url: The server URL for which to retrieve the token.
        allow_expired: Whether to allow expired tokens to be returned. The
            default behavior is to return None if a token does exist but is
            expired.

    Returns:
        The stored token if it exists and is not expired, None otherwise.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(server_url)
    if credential:
        token = credential.api_token
        if token and (not token.expired or allow_expired):
            return token
    return None

has_valid_authentication(url: str) -> bool

Check if a valid authentication credential for the given server URL is stored.

Parameters:

Name	Type	Description	Default
url	str	The server URL for which to check the authentication.	required

Returns:

Name	Type	Description
bool	bool	True if a valid token or API key is stored, False otherwise.

Source code in src/zenml/login/credentials_store.py

def has_valid_authentication(self, url: str) -> bool:
    """Check if a valid authentication credential for the given server URL is stored.

    Args:
        url: The server URL for which to check the authentication.

    Returns:
        bool: True if a valid token or API key is stored, False otherwise.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(url)

    if not credential:
        return False
    if credential.api_key or (
        credential.username and credential.password is not None
    ):
        return True
    token = credential.api_token
    return token is not None and not token.expired

has_valid_pro_authentication(pro_api_url: str) -> bool

Check if a valid token for a ZenML Pro API server is stored.

Parameters:

Name	Type	Description	Default
pro_api_url	str	The URL of the ZenML Pro API server.	required

Returns:

Name	Type	Description
bool	bool	True if a valid token is stored, False otherwise.

Source code in src/zenml/login/credentials_store.py

def has_valid_pro_authentication(self, pro_api_url: str) -> bool:
    """Check if a valid token for a ZenML Pro API server is stored.

    Args:
        pro_api_url: The URL of the ZenML Pro API server.

    Returns:
        bool: True if a valid token is stored, False otherwise.
    """
    return self.get_pro_token(pro_api_url) is not None

list_credentials(type: Optional[ServerType] = None) -> List[ServerCredentials]

Get all credentials stored in the credentials store.

Parameters:

Name	Type	Description	Default
type	Optional[ServerType]	Optional server type to filter the credentials by.	None

Returns:

Type	Description
List[ServerCredentials]	A list of all credentials stored in the credentials store.

Source code in src/zenml/login/credentials_store.py

def list_credentials(
    self, type: Optional[ServerType] = None
) -> List[ServerCredentials]:
    """Get all credentials stored in the credentials store.

    Args:
        type: Optional server type to filter the credentials by.

    Returns:
        A list of all credentials stored in the credentials store.
    """
    self.check_and_reload_from_file()

    credentials = list(self.credentials.values())

    if type is not None:
        credentials = [c for c in credentials if c and c.type == type]

    return credentials

reset_instance(store: Optional[CredentialsStore] = None) -> None classmethod

Reset the singleton instance of the CredentialsStore.

Parameters:

Name	Type	Description	Default
store	Optional[CredentialsStore]	Optional instance of the CredentialsStore to set as the singleton instance. If None, a new instance will be created.	None

Source code in src/zenml/login/credentials_store.py

@classmethod
def reset_instance(
    cls, store: Optional["CredentialsStore"] = None
) -> None:
    """Reset the singleton instance of the CredentialsStore.

    Args:
        store: Optional instance of the CredentialsStore to set as the
            singleton instance. If None, a new instance will be created.
    """
    current_store = cls.get_instance()
    if current_store is not None and current_store is not store:
        # Delete the credentials file from disk if it exists, otherwise
        # the credentials will be reloaded from the file when the new
        # instance is created and this call will have no effect
        current_store._delete_credentials_file()

    cls._clear(store)  # type: ignore[arg-type]
    if store:
        store._save_credentials()

set_api_key(server_url: str, api_key: str) -> None

Store an API key in the credentials store for a specific server URL.

If an API token or a password is already stored for the server URL, they will be replaced by the API key.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which the token is to be stored.	required
api_key	str	The API key to store.	required

Source code in src/zenml/login/credentials_store.py

def set_api_key(
    self,
    server_url: str,
    api_key: str,
) -> None:
    """Store an API key in the credentials store for a specific server URL.

    If an API token or a password is already stored for the server URL, they
    will be replaced by the API key.

    Args:
        server_url: The server URL for which the token is to be stored.
        api_key: The API key to store.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(server_url)
    if credential and credential.api_key != api_key:
        # Reset the API token if a new or updated API key is set, because
        # the current token might have been issued for a different account
        credential.api_token = None
        credential.api_key = api_key
        credential.username = None
        credential.password = None
    else:
        self.credentials[server_url] = ServerCredentials(
            url=server_url, api_key=api_key
        )

    self._save_credentials()

set_bare_token(server_url: str, token: str) -> APIToken

Store a bare API token.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which the token is to be stored.	required
token	str	The token to store.	required

Returns:

Name	Type	Description
APIToken	APIToken	The stored token.

Source code in src/zenml/login/credentials_store.py

def set_bare_token(
    self,
    server_url: str,
    token: str,
) -> APIToken:
    """Store a bare API token.

    Args:
        server_url: The server URL for which the token is to be stored.
        token: The token to store.

    Returns:
        APIToken: The stored token.
    """
    self.check_and_reload_from_file()

    api_token = APIToken(
        access_token=token,
    )

    credential = self.credentials.get(server_url)
    if credential:
        credential.api_token = api_token
    else:
        self.credentials[server_url] = ServerCredentials(
            url=server_url, api_token=api_token
        )

    self._save_credentials()

    return api_token

set_password(server_url: str, username: str, password: str) -> None

Store a username and password in the credentials store for a specific server URL.

If an API token is already stored for the server URL, it will be replaced by the username and password.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which the token is to be stored.	required
username	str	The username to store.	required
password	str	The password to store.	required

Source code in src/zenml/login/credentials_store.py

def set_password(
    self,
    server_url: str,
    username: str,
    password: str,
) -> None:
    """Store a username and password in the credentials store for a specific server URL.

    If an API token is already stored for the server URL, it will be
    replaced by the username and password.

    Args:
        server_url: The server URL for which the token is to be stored.
        username: The username to store.
        password: The password to store.
    """
    self.check_and_reload_from_file()
    credential = self.credentials.get(server_url)
    if credential and (
        credential.username != username or credential.password != password
    ):
        # Reset the API token if a new or updated password is set, because
        # the current token might have been issued for a different account
        credential.api_token = None
        credential.username = username
        credential.password = password
        credential.api_key = None
    else:
        self.credentials[server_url] = ServerCredentials(
            url=server_url, username=username, password=password
        )

    self._save_credentials()

set_token(server_url: str, token_response: OAuthTokenResponse, is_zenml_pro: bool = False) -> APIToken

Store an API token received from an OAuth2 server.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which the token is to be stored.	required
token_response	OAuthTokenResponse	Token response received from an OAuth2 server.	required
is_zenml_pro	bool	Whether the token is for a ZenML Pro server.	False

Returns:

Name	Type	Description
APIToken	APIToken	The stored token.

Source code in src/zenml/login/credentials_store.py

def set_token(
    self,
    server_url: str,
    token_response: OAuthTokenResponse,
    is_zenml_pro: bool = False,
) -> APIToken:
    """Store an API token received from an OAuth2 server.

    Args:
        server_url: The server URL for which the token is to be stored.
        token_response: Token response received from an OAuth2 server.
        is_zenml_pro: Whether the token is for a ZenML Pro server.

    Returns:
        APIToken: The stored token.
    """
    self.check_and_reload_from_file()
    if token_response.expires_in:
        expires_at = utc_now_tz_aware() + timedelta(
            seconds=token_response.expires_in
        )
        # Best practice to calculate the leeway depending on the token
        # expiration time:
        #
        # - for short-lived tokens (less than 1 hour), use a fixed leeway of
        # a few seconds (e.g., 30 seconds)
        # - for longer-lived tokens (e.g., 1 hour or more), use a
        # percentage-based leeway of 5-10%
        if token_response.expires_in < 3600:
            leeway = 30
        else:
            leeway = token_response.expires_in // 20
    else:
        expires_at = None
        leeway = None

    api_token = APIToken(
        access_token=token_response.access_token,
        expires_in=token_response.expires_in,
        expires_at=expires_at,
        leeway=leeway,
        device_id=token_response.device_id,
        device_metadata=token_response.device_metadata,
    )

    credential = self.credentials.get(server_url)
    if credential:
        credential.api_token = api_token
    else:
        credential = self.credentials[server_url] = ServerCredentials(
            url=server_url, api_token=api_token
        )

    if is_zenml_pro:
        # This is how we encode that the token is for a ZenML Pro server
        credential.pro_api_url = server_url

    self._save_credentials()

    return api_token

update_server_info(server_url: str, server_info: Union[ServerModel, WorkspaceRead]) -> None

Update the server information stored for a specific server URL.

Parameters:

Name	Type	Description	Default
server_url	str	The server URL for which the server information is to be updated.	required
server_info	Union[ServerModel, WorkspaceRead]	Updated server information.	required

Source code in src/zenml/login/credentials_store.py

def update_server_info(
    self,
    server_url: str,
    server_info: Union[ServerModel, WorkspaceRead],
) -> None:
    """Update the server information stored for a specific server URL.

    Args:
        server_url: The server URL for which the server information is to be
            updated.
        server_info: Updated server information.
    """
    self.check_and_reload_from_file()

    credential = self.credentials.get(server_url)
    if not credential:
        # No credentials stored for this server URL, nothing to update
        return

    credential.update_server_info(server_info)
    self._save_credentials()

Functions

get_credentials_store() -> CredentialsStore

Get the global credentials store instance.

Returns:

Type	Description
CredentialsStore	The global credentials store instance.

Source code in src/zenml/login/credentials_store.py

def get_credentials_store() -> CredentialsStore:
    """Get the global credentials store instance.

    Returns:
        The global credentials store instance.
    """
    return CredentialsStore()

Modules

pro

ZenML Pro client.

Modules

client

ZenML Pro client.

Classes

ZenMLProClient(url: str, api_token: Optional[APIToken] = None)

ZenML Pro client.

Initialize the ZenML Pro client.

Parameters:

Name	Type	Description	Default
url	str	The URL of the ZenML Pro API server.	required
api_token	Optional[APIToken]	The API token to use for authentication. If not provided, the token is fetched from the credentials store.	None

Raises:

Type	Description
AuthorizationException	If no API token is provided and no token is found in the credentials store.

Source code in src/zenml/login/pro/client.py

def __init__(self, url: str, api_token: Optional[APIToken] = None) -> None:
    """Initialize the ZenML Pro client.

    Args:
        url: The URL of the ZenML Pro API server.
        api_token: The API token to use for authentication. If not provided,
            the token is fetched from the credentials store.

    Raises:
        AuthorizationException: If no API token is provided and no token
            is found in the credentials store.
    """
    self._url = url
    if api_token is None:
        logger.debug(
            "No ZenML Pro API token provided. Fetching from credentials "
            "store."
        )
        api_token = get_credentials_store().get_token(
            server_url=self._url, allow_expired=True
        )
        if api_token is None:
            raise AuthorizationException(
                "No ZenML Pro API token found. Please run 'zenml login' to "
                "login to ZenML Pro."
            )

    self._api_token = api_token

Attributes

api_token: str property

Get the API token.

Returns:

Type	Description
str	The API token.

organization: OrganizationClient property

Get the organization client.

Returns:

Type	Description
OrganizationClient	The organization client.

session: requests.Session property

Authenticate to the ZenML Pro API server.

Returns:

Type	Description
Session	A requests session with the authentication token.

workspace: WorkspaceClient property

Get the workspace client.

Returns:

Type	Description
WorkspaceClient	The workspace client.

Functions

delete(path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Json

Make a DELETE request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/login/pro/client.py

def delete(
    self,
    path: str,
    params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Json:
    """Make a DELETE request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        params: The query parameters to pass to the endpoint.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    logger.debug(f"Sending DELETE request to {path}...")
    return self._request(
        "DELETE",
        self._url + path,
        params=params,
        **kwargs,
    )

get(path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Json

Make a GET request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/login/pro/client.py

def get(
    self,
    path: str,
    params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Json:
    """Make a GET request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        params: The query parameters to pass to the endpoint.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    logger.debug(f"Sending GET request to {path}...")
    return self._request(
        "GET",
        self._url + path,
        params=params,
        **kwargs,
    )

patch(path: str, body: Optional[BaseRestAPIModel] = None, params: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Json

Make a PATCH request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	Optional[BaseRestAPIModel]	The body to send.	None
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/login/pro/client.py

def patch(
    self,
    path: str,
    body: Optional[BaseRestAPIModel] = None,
    params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Json:
    """Make a PATCH request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    logger.debug(f"Sending PATCH request to {path}...")
    json = (
        body.model_dump(mode="json", exclude_unset=True) if body else None
    )
    return self._request(
        "PATCH",
        self._url + path,
        json=json,
        params=params,
        **kwargs,
    )

post(path: str, body: BaseRestAPIModel, params: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Json

Make a POST request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	BaseRestAPIModel	The body to send.	required
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/login/pro/client.py

def post(
    self,
    path: str,
    body: BaseRestAPIModel,
    params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Json:
    """Make a POST request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    logger.debug(f"Sending POST request to {path}...")
    return self._request(
        "POST",
        self._url + path,
        json=body.model_dump(mode="json"),
        params=params,
        **kwargs,
    )

put(path: str, body: Optional[BaseRestAPIModel] = None, params: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Json

Make a PUT request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	Optional[BaseRestAPIModel]	The body to send.	None
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/login/pro/client.py

def put(
    self,
    path: str,
    body: Optional[BaseRestAPIModel] = None,
    params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Json:
    """Make a PUT request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    logger.debug(f"Sending PUT request to {path}...")
    json = (
        body.model_dump(mode="json", exclude_unset=True) if body else None
    )
    return self._request(
        "PUT",
        self._url + path,
        json=json,
        params=params,
        **kwargs,
    )

raise_on_expired_api_token() -> None

Raise an exception if the API token has expired.

Raises:

Type	Description
AuthorizationException	If the API token has expired.

Source code in src/zenml/login/pro/client.py

def raise_on_expired_api_token(self) -> None:
    """Raise an exception if the API token has expired.

    Raises:
        AuthorizationException: If the API token has expired.
    """
    if self._api_token and self._api_token.expired:
        raise AuthorizationException(
            "Your ZenML Pro authentication has expired. Please run "
            "'zenml login' to login again."
        )

Functions

constants

ZenML Pro login constants.

models

ZenML Pro base models.

Classes

BaseRestAPIModel

Bases: BaseModel

Base class for all REST API models.

organization

ZenML Pro organization client.

Modules

client

ZenML Pro organization client.

Classes

OrganizationClient(client: ZenMLProClient)

Organization management client.

Initialize the organization client.

Parameters:

Name	Type	Description	Default
client	ZenMLProClient	ZenML Pro client.	required

Source code in src/zenml/login/pro/organization/client.py

def __init__(
    self,
    client: ZenMLProClient,
):
    """Initialize the organization client.

    Args:
        client: ZenML Pro client.
    """
    self.client = client

Functions

get(id_or_name: Union[UUID, str]) -> OrganizationRead

Get an organization by id or name.

Parameters:

Name	Type	Description	Default
id_or_name	Union[UUID, str]	Id or name of the organization to retrieve.	required

Returns:

Type	Description
OrganizationRead	An organization.

Source code in src/zenml/login/pro/organization/client.py

def get(
    self,
    id_or_name: Union[UUID, str],
) -> OrganizationRead:
    """Get an organization by id or name.

    Args:
        id_or_name: Id or name of the organization to retrieve.

    Returns:
        An organization.
    """
    return self.client._get_resource(
        resource_id=id_or_name,
        route=ORGANIZATIONS_ROUTE,
        response_model=OrganizationRead,
    )

list(offset: int = 0, limit: int = 20) -> List[OrganizationRead] async

List organizations.

Parameters:

Name	Type	Description	Default
offset	int	Query offset.	0
limit	int	Query limit.	20

Returns:

Type	Description
List[OrganizationRead]	List of organizations.

Source code in src/zenml/login/pro/organization/client.py

async def list(
    self,
    offset: int = 0,
    limit: int = 20,
) -> List[OrganizationRead]:
    """List organizations.

    Args:
        offset: Query offset.
        limit: Query limit.

    Returns:
        List of organizations.
    """
    return self.client._list_resources(
        route=ORGANIZATIONS_ROUTE,
        response_model=OrganizationRead,
        offset=offset,
        limit=limit,
    )

Functions

models

ZenML Pro organization models.

Classes

OrganizationRead

Bases: BaseRestAPIModel

Model for viewing organizations.

utils

ZenML Pro login utils.

Classes Functions

get_troubleshooting_instructions(url: str) -> str

Get troubleshooting instructions for a given ZenML Pro server URL.

Parameters:

Name	Type	Description	Default
url	str	ZenML Pro server URL	required

Returns:

Type	Description
str	Troubleshooting instructions

Source code in src/zenml/login/pro/utils.py

def get_troubleshooting_instructions(url: str) -> str:
    """Get troubleshooting instructions for a given ZenML Pro server URL.

    Args:
        url: ZenML Pro server URL

    Returns:
        Troubleshooting instructions
    """
    credentials_store = get_credentials_store()

    credentials = credentials_store.get_credentials(url)
    pro_api_url = None
    if credentials and credentials.type == ServerType.PRO:
        pro_api_url = credentials.pro_api_url or ZENML_PRO_API_URL

    if pro_api_url and credentials_store.has_valid_pro_authentication(
        pro_api_url
    ):
        client = ZenMLProClient(pro_api_url)

        try:
            servers = client.workspace.list(url=url, member_only=False)
        except Exception as e:
            logger.debug(f"Failed to list workspaces: {e}")
        else:
            if servers:
                server = servers[0]
                if server.status == WorkspaceStatus.AVAILABLE:
                    return (
                        f"The '{server.name}' ZenML Pro server that the client "
                        "is connected to is currently running but you may not "
                        "have the necessary permissions to access it. Please "
                        "contact your ZenML Pro administrator for more "
                        "information or try to manage the server members "
                        "yourself if you have the necessary permissions by "
                        f"visiting the ZenML Pro workspace page at {server.dashboard_url}."
                    )
                if server.status == WorkspaceStatus.DEACTIVATED:
                    return (
                        f"The '{server.name}' ZenML Pro server that the client "
                        "is connected to has been deactivated. "
                        "Please contact your ZenML Pro administrator for more "
                        "information or to reactivate the server yourself if "
                        "you have the necessary permissions by visiting the "
                        f"ZenML Pro Organization page at {server.dashboard_organization_url}."
                    )
                if server.status == WorkspaceStatus.PENDING:
                    return (
                        f"The '{server.name}' ZenML Pro server that the client "
                        "is connected to is currently undergoing maintenance "
                        "(e.g. being deployed, upgraded or re-activated). "
                        "Please try again later or contact your ZenML Pro "
                        "administrator for more information. You can also "
                        f"visit the ZenML Pro workspace page at {server.dashboard_url}."
                    )
                return (
                    f"The '{server.name}' ZenML Pro server that the client "
                    "is connected to is currently in a failed "
                    "state. Please contact your ZenML Pro administrator for "
                    "more information or try to re-deploy the server "
                    "yourself if you have the necessary permissions by "
                    "visiting the ZenML Pro Organization page at "
                    f"{server.dashboard_organization_url}."
                )

            return (
                f"The ZenML Pro server at URL '{url}' that the client is "
                "connected to does not exist or you may not have access to it. "
                "Please check the URL and your permissions and try again or "
                "connect your client to a different server by running `zenml "
                "login` or by using a service account API key."
            )

    return (
        f"The ZenML Pro server at URL '{url}' that the client is connected to "
        "does not exist, is not running, or you do not have permissions to "
        "connect to it. Please check the URL and your permissions "
        "and try again. The ZenML Pro server might have been deactivated or is "
        "currently pending maintenance. Please contact your ZenML Pro "
        "administrator for more information or try to manage the server "
        "state by visiting the ZenML Pro dashboard."
    )

workspace

ZenML Pro workspace client.

Modules

client

ZenML Pro workspace client.

Classes

WorkspaceClient(client: ZenMLProClient)

Workspace management client.

Initialize the workspace client.

Parameters:

Name	Type	Description	Default
client	ZenMLProClient	ZenML Pro client.	required

Source code in src/zenml/login/pro/workspace/client.py

def __init__(
    self,
    client: ZenMLProClient,
):
    """Initialize the workspace client.

    Args:
        client: ZenML Pro client.
    """
    self.client = client

Functions

get(id: UUID) -> WorkspaceRead

Get a workspace by id.

Parameters:

Name	Type	Description	Default
id	UUID	Id. of the workspace to retrieve.	required

Returns:

Type	Description
WorkspaceRead	A workspace.

Source code in src/zenml/login/pro/workspace/client.py

def get(self, id: UUID) -> WorkspaceRead:
    """Get a workspace by id.

    Args:
        id: Id. of the workspace to retrieve.

    Returns:
        A workspace.
    """
    return self.client._get_resource(
        resource_id=id,
        route=WORKSPACES_ROUTE,
        response_model=WorkspaceRead,
    )

list(offset: int = 0, limit: int = 20, workspace_name: Optional[str] = None, url: Optional[str] = None, organization_id: Optional[UUID] = None, status: Optional[WorkspaceStatus] = None, member_only: bool = False) -> List[WorkspaceRead]

List workspaces.

Parameters:

Name	Type	Description	Default
offset	int	Offset to use for filtering.	0
limit	int	Limit used for filtering.	20
workspace_name	Optional[str]	Workspace name to filter by.	None
url	Optional[str]	Workspace service URL to filter by.	None
organization_id	Optional[UUID]	Organization ID to filter by.	None
status	Optional[WorkspaceStatus]	Filter for only workspaces with this status.	None
member_only	bool	If True, only list workspaces where the user is a member (i.e. users that can connect to the workspace).	False

Returns:

Type	Description
List[WorkspaceRead]	List of workspaces.

Source code in src/zenml/login/pro/workspace/client.py

def list(
    self,
    offset: int = 0,
    limit: int = 20,
    workspace_name: Optional[str] = None,
    url: Optional[str] = None,
    organization_id: Optional[UUID] = None,
    status: Optional[WorkspaceStatus] = None,
    member_only: bool = False,
) -> List[WorkspaceRead]:
    """List workspaces.

    Args:
        offset: Offset to use for filtering.
        limit: Limit used for filtering.
        workspace_name: Workspace name to filter by.
        url: Workspace service URL to filter by.
        organization_id: Organization ID to filter by.
        status: Filter for only workspaces with this status.
        member_only: If True, only list workspaces where the user is a member
            (i.e. users that can connect to the workspace).

    Returns:
        List of workspaces.
    """
    return self.client._list_resources(
        route=WORKSPACES_ROUTE,
        response_model=WorkspaceRead,
        offset=offset,
        limit=limit,
        workspace_name=workspace_name,
        url=url,
        organization_id=organization_id,
        status=status,
        member_only=member_only,
    )

Functions

models

ZenML Pro workspace models.

Classes

WorkspaceRead

Bases: BaseRestAPIModel

Pydantic Model for viewing a Workspace.

Attributes

dashboard_organization_url: str property

Get the URL to the ZenML Pro dashboard for this workspace's organization.

Returns:

Type	Description
str	The URL to the ZenML Pro dashboard for this workspace's organization.

dashboard_url: str property

Get the URL to the ZenML Pro dashboard for this workspace.

Returns:

Type	Description
str	The URL to the ZenML Pro dashboard for this workspace.

organization_id: UUID property

Get the organization id.

Returns:

Type	Description
UUID	The organization id.

organization_name: str property

Get the organization name.

Returns:

Type	Description
str	The organization name.

url: Optional[str] property

Get the ZenML server URL.

Returns:

Type	Description
Optional[str]	The ZenML server URL, if available.

version: Optional[str] property

Get the ZenML service version.

Returns:

Type	Description
Optional[str]	The ZenML service version.

WorkspaceStatus

Bases: StrEnum

Enum that represents the desired state or status of a workspace.

These values can be used in two places:

• in the desired_state field of a workspace object, to indicate the desired state of the workspace (with the exception of PENDING and FAILED which are not valid values for desired_state)
• in the status field of a workspace object, to indicate the current state of the workspace

ZenMLServiceConfiguration

Bases: BaseRestAPIModel

ZenML service configuration.

ZenMLServiceRead

Bases: BaseRestAPIModel

Pydantic Model for viewing a ZenML service.

ZenMLServiceStatus

Bases: BaseRestAPIModel

ZenML service status.

server_info

ZenML server information retrieval.

Classes

Functions

get_server_info(url: str) -> Optional[ServerModel]

Retrieve server information from a remote ZenML server.

Parameters:

Name	Type	Description	Default
url	str	The URL of the ZenML server.	required

Returns:

Type	Description
Optional[ServerModel]	The server information or None if the server info could not be fetched.

Source code in src/zenml/login/server_info.py

def get_server_info(url: str) -> Optional[ServerModel]:
    """Retrieve server information from a remote ZenML server.

    Args:
        url: The URL of the ZenML server.

    Returns:
        The server information or None if the server info could not be fetched.
    """
    # Here we try to leverage the existing RestZenStore support to fetch the
    # server info and only the server info, which doesn't actually need
    # any authentication.
    try:
        store = RestZenStore(
            config=RestZenStoreConfiguration(
                url=url,
            )
        )
        return store.server_info
    except Exception as e:
        logger.warning(
            f"Failed to fetch server info from the server running at {url}: {e}"
        )

    return None

web_login

ZenML OAuth2 device authorization grant client support.

Classes

Functions

web_login(url: Optional[str] = None, verify_ssl: Optional[Union[str, bool]] = None, pro_api_url: Optional[str] = None) -> APIToken

Implements the OAuth2 Device Authorization Grant flow.

This function implements the client side of the OAuth2 Device Authorization Grant flow as defined in https://tools.ietf.org/html/rfc8628, with the following customizations:

• the unique ZenML client ID (user_id in the global config) is used as the OAuth2 client ID value
• additional information is added to the user agent header to be used by users to identify the ZenML client

Upon completion of the flow, the access token is saved in the credentials store.

Parameters:

Name	Type	Description	Default
url	Optional[str]	The URL of the OAuth2 server. If not provided, the ZenML Pro API server is used by default.	None
verify_ssl	Optional[Union[str, bool]]	Whether to verify the SSL certificate of the OAuth2 server. If a string is passed, it is interpreted as the path to a CA bundle file.	None
pro_api_url	Optional[str]	The URL of the ZenML Pro API server. If not provided, the default ZenML Pro API server URL is used.	None

Returns:

Type	Description
APIToken	The response returned by the OAuth2 server.

Raises:

Type	Description
AuthorizationException	If an error occurred during the authorization process.

Source code in src/zenml/login/web_login.py

def web_login(
    url: Optional[str] = None,
    verify_ssl: Optional[Union[str, bool]] = None,
    pro_api_url: Optional[str] = None,
) -> APIToken:
    """Implements the OAuth2 Device Authorization Grant flow.

    This function implements the client side of the OAuth2 Device Authorization
    Grant flow as defined in https://tools.ietf.org/html/rfc8628, with the
    following customizations:

    * the unique ZenML client ID (`user_id` in the global config) is used
    as the OAuth2 client ID value
    * additional information is added to the user agent header to be used by
    users to identify the ZenML client

    Upon completion of the flow, the access token is saved in the credentials store.

    Args:
        url: The URL of the OAuth2 server. If not provided, the ZenML Pro API
            server is used by default.
        verify_ssl: Whether to verify the SSL certificate of the OAuth2 server.
            If a string is passed, it is interpreted as the path to a CA bundle
            file.
        pro_api_url: The URL of the ZenML Pro API server. If not provided, the
            default ZenML Pro API server URL is used.

    Returns:
        The response returned by the OAuth2 server.

    Raises:
        AuthorizationException: If an error occurred during the authorization
            process.
    """
    from zenml.login.credentials_store import get_credentials_store
    from zenml.models import (
        OAuthDeviceAuthorizationRequest,
        OAuthDeviceAuthorizationResponse,
        OAuthDeviceTokenRequest,
        OAuthDeviceUserAgentHeader,
        OAuthTokenResponse,
    )

    credentials_store = get_credentials_store()

    # Make a request to the OAuth2 server to get the device code and user code.
    # The client ID used for the request is the unique ID of the ZenML client.
    response: Optional[requests.Response] = None

    # Add the following information in the user agent header to be used by users
    # to identify the ZenML client:
    #
    # * the ZenML version
    # * the python version
    # * the OS type
    # * the hostname
    #
    user_agent_header = OAuthDeviceUserAgentHeader(
        hostname=platform.node(),
        zenml_version=__version__,
        python_version=platform.python_version(),
        os=platform.system(),
    )

    zenml_pro = False
    if not url:
        # If no URL is provided, we use the ZenML Pro API server by default
        zenml_pro = True
        url = base_url = pro_api_url or ZENML_PRO_API_URL
    else:
        # Get rid of any trailing slashes to prevent issues when having double
        # slashes in the URL
        url = url.rstrip("/")
        if pro_api_url:
            # This is a ZenML Pro server. The device authentication is done
            # through the ZenML Pro API.
            zenml_pro = True
            base_url = pro_api_url
        else:
            base_url = url

    auth_request = OAuthDeviceAuthorizationRequest(
        client_id=GlobalConfiguration().user_id
    )

    # If an existing token is found in the credentials store, we reuse its
    # device ID to avoid creating a new device ID for the same device.
    existing_token = credentials_store.get_token(url)
    if existing_token and existing_token.device_id:
        auth_request.device_id = existing_token.device_id

    if zenml_pro:
        auth_url = base_url + AUTH + DEVICE_AUTHORIZATION
        login_url = base_url + AUTH + LOGIN
    else:
        auth_url = base_url + API + VERSION_1 + DEVICE_AUTHORIZATION
        login_url = base_url + API + VERSION_1 + LOGIN

    try:
        response = requests.post(
            auth_url,
            headers={
                "Content-Type": "application/x-www-form-urlencoded",
                "User-Agent": user_agent_header.encode(),
            },
            data=auth_request.model_dump(exclude_none=True),
            verify=verify_ssl,
            timeout=DEFAULT_HTTP_TIMEOUT,
        )
        if response.status_code == 200:
            auth_response = OAuthDeviceAuthorizationResponse(**response.json())
        else:
            logger.info(f"Error: {response.status_code} {response.text}")
            raise AuthorizationException(
                f"Could not connect to {base_url}. Please check the URL."
            )
    except (requests.exceptions.JSONDecodeError, ValueError, TypeError):
        logger.exception("Bad response received from API server.")
        raise AuthorizationException(
            "Bad response received from API server. Please check the URL."
        )
    except requests.exceptions.RequestException:
        logger.exception("Could not connect to API server.")
        raise AuthorizationException(
            f"Could not connect to {base_url}. Please check the URL."
        )

    # Open the verification URL in the user's browser
    verification_uri = (
        auth_response.verification_uri_complete
        or auth_response.verification_uri
    )
    if verification_uri.startswith("/"):
        # If the verification URI is a relative path, we need to add the base
        # URL to it
        verification_uri = base_url + verification_uri
    webbrowser.open(verification_uri)
    logger.info(
        f"If your browser did not open automatically, please open the "
        f"following URL into your browser to proceed with the authentication:"
        f"\n\n{verification_uri}\n"
    )

    # Poll the OAuth2 server until the user has authorized the device
    token_request = OAuthDeviceTokenRequest(
        device_code=auth_response.device_code,
        client_id=auth_request.client_id,
    )
    expires_in = auth_response.expires_in
    interval = auth_response.interval
    token_response: OAuthTokenResponse
    while True:
        response = requests.post(
            login_url,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            data=token_request.model_dump(),
            verify=verify_ssl,
            timeout=DEFAULT_HTTP_TIMEOUT,
        )
        if response.status_code == 200:
            # The user has authorized the device, so we can extract the access token
            token_response = OAuthTokenResponse(**response.json())
            if zenml_pro:
                logger.info("Successfully logged in to ZenML Pro.")
            else:
                logger.info(f"Successfully logged in to {url}.")
            break
        elif response.status_code == 400:
            try:
                error_response = OAuthError(**response.json())
            except (
                requests.exceptions.JSONDecodeError,
                ValueError,
                TypeError,
            ):
                raise AuthorizationException(
                    f"Error received from {base_url}: {response.text}"
                )

            if error_response.error == "authorization_pending":
                # The user hasn't authorized the device yet, so we wait for the
                # interval and try again
                pass
            elif error_response.error == "slow_down":
                # The OAuth2 server is asking us to slow down our polling
                interval += 5
            else:
                # There was another error with the request
                raise AuthorizationException(
                    f"Error: {error_response.error} {error_response.error_description}"
                )

            expires_in -= interval
            if expires_in <= 0:
                raise AuthorizationException(
                    "User did not authorize the device in time."
                )
            time.sleep(interval)
        else:
            # There was another error with the request
            raise AuthorizationException(
                f"Error: {response.status_code} {response.text}"
            )

    # Save the token in the credentials store
    return credentials_store.set_token(
        url, token_response, is_zenml_pro=zenml_pro
    )