Skip to content

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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
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
88
89
90
91
92
93
94
95
96
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
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
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
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
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
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
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
631
632
633
634
635
636
637
638
639
640
641
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
336
337
338
339
340
341
342
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
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
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
244
245
246
247
248
249
250
251
252
253
254
255
256
257
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
281
282
283
284
285
286
287
288
289
290
291
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
119
120
121
122
123
124
125
126
@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
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
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
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
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
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
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
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
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
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
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
389
390
391
392
393
394
395
396
397
398
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
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
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
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
@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
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
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
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
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
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
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
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
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
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
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
664
665
666
667
668
669
670
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
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
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
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
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
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
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
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
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
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
126
127
128
129
130
131
132
133
134
135
136
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
31
32
33
34
35
36
37
38
39
40
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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
31
32
33
34
35
36
37
38
39
40
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
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
    )