Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

ZenML Pro web login implementation #3141

Merged
merged 31 commits into from
Oct 29, 2024
Merged

Conversation

stefannica
Copy link
Contributor

@stefannica stefannica commented Oct 24, 2024

Describe changes

Implements the zenml login CLI command and all associated ZenML Pro web login features:

The new zenml login CLI command can be used to manage the servers that your ZenML client is authenticated and connected to. Note that "connecting" and "authenticating" are no longer equivalent: you only need to authenticate (login) to a ZenML Server once, then connect/re-connect to it as many times as necessary while the authentication is valid, because all temporary credentials are now cached locally in the Credentials Store. This allows users to easily switch between the local server and the remote servers without having to go through the web login authentication flow.

zenml login can be used to:

  • authenticate to ZenML Pro and gain access to all your ZenML Pro servers
  • authenticate to any number of self-hosted remote ZenML servers (previously done with zenml connect)
  • start/restart/re-configure and connect to the local ZenML server (previously done with zenml up/zenml down)
  • refresh your ZenML Pro authentication session
  • refresh the authentication session for any self-hosted ZenML server (zenml connect always refreshed the session)

zenml logout does the reverse and can be used to:

  • log out from ZenML Pro and clear the cached authentication credentials
  • log out from any self-hosted ZenML server and clear the cached authentication credentials (previously done with zenml disconnect)
  • shut down the local ZenML server (previously done with zenml down)

zenml server list can be used to list all the servers that the client can be connected to, with options to include ZenML Pro servers that are inactive or not accessible and ZenML servers that have expired credentials.

The existing commands zenml connect, zenml up, zenml down, zenml disconnect and zenml show are kept in place but are deprecated and call zenml login, zenml logout and zenml server show under the hood.

Implementation details

At the core of the new features lies a new design component called the 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. 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 in the following ways:

  1. when the user runs zenml login to authenticate to a ZenML Pro server, it stores the ZenML Pro API token fetched from the web login flow in the Credentials Store
  2. when the user runs zenml login to authenticate to a regular ZenML server with the web login flow, it stores the ZenML API token fetched through the web login flow in the Credentials Store
  3. when the user runs zenml login to authenticate to any ZenML server using an API key, it stores the API key in the Credentials Store
  4. the REST zen store is initialized non-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:
    • for ZenML servers that use an API key to authenticate, it converts that into a short-lived ZenML Pro server API token that it also stores in the Credentials Store
    • for ZenML Pro servers, it exchanges the long-lived ZenML Pro API token into a short lived ZenML Pro server API token
    • for regular ZenML servers, it will use the cached API token that was last fetched through the web login flow
    • that's it ! no other forms of authentication are possible

Alongside credentials, the Credentials Store is also used to store additional server information:
* ZenML Pro tenant 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

This information is used in commands like zenml status and zenml server list to display as much information as possible about cached servers.

How to test this

export ZENML_PRO_URL=https://staging.cloud.zenml.io
export ZENML_PRO_API_URL=https://staging.cloudapi.zenml.io
zenml login

Examples

$ zenml login
No server argument was provided. Logging to ZenML Pro...
If your browser did not open automatically, please open the following URL into your browser to proceed with the authentication:

https://staging.cloud.zenml.io/devices/verify?user_code=wB_7qQSqMfXcRHmMLTwxSQ

Successfully logged in to ZenML Pro.
You can now run 'zenml server list' to view the available ZenML Pro servers and then 'zenml login <server-url-name-or-id>' to connect to a specific server without having to log in again until your session 
expires.
Connecting to ZenML Pro server: test-zenml-login [16f8a35d-5c2f-44aa-a564-b34186fbf6d6] 
Authenticating to ZenML server 'https://64f4e58f-zenml.staging.cloudinfra.zenml.io'...
Setting the global active workspace to 'default'.
Setting the global active stack to default.
Updated the global store configuration.
Connected to ZenML Pro server: test-zenml-login.
$ zenml login --local
Creating database tables
Creating default workspace 'default' ...
Creating default stack in workspace default...
Deploying a local daemon ZenML server.
Connecting to the local daemon ZenML server (http://127.0.0.1:8237).
The current global active stack is no longer available. Resetting the active stack to default.
Updated the global store configuration.
Connected to the local daemon ZenML server (http://127.0.0.1:8237).
The local ZenML dashboard is available at 'http://127.0.0.1:8237'.
Automatically opening the dashboard in your browser. To disable this, set the env variable AUTO_OPEN_DASHBOARD=false.
$ zenml server list
┏━━━━━━━━┯━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━┓
┃ ACTIVE │ TYPE  │ ID                                   │ NAME                │ ORGANIZATION            │ VERSION ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃   👉   │ PRO   │ 97c9a968-15f8-42c9-8594-348e1e45594e │ second-tenant       │ Silent Assassin's Creed │ 0.67.0  ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃        │ PRO   │ 16f8a35d-5c2f-44aa-a564-b34186fbf6d6 │ test-zenml-login    │ Silent Assassin's Creed │ 0.67.0  ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃        │ PRO   │ 0abfc085-5154-429c-87f1-f47526a422a9 │ new-tenant          │ ScandinavianAI          │ 0.67.0  ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃        │ PRO   │ 300b7ddb-eb2e-4e06-ab43-ff4c0065c4b5 │ filters             │ Arbyte                  │ 0.66.0  ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃        │ PRO   │ 3e299f77-170e-4aba-9338-5a444446cf66 │ arbyte-tenant       │ Arbyte                  │ 0.66.0  ┃
┠────────┼───────┼──────────────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┨
┃        │ LOCAL │ 3d34ce75-8c33-48b5-8bd4-262b1ed95eed │ local daemon server │ N/A                     │ 0.67.0  ┃
┗━━━━━━━━┷━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━┛
$ zenml server list -av
┏━━━━━━━━┯━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━┯━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ ACTIVE │ TYPE  │ ID                          │ NAME                │ ORGANIZATION            │ VERSION │ STATUS      │ DASHBOARD URL               │ API URL                      │ AUTH STATUS                 ┃
┠────────┼───────┼─────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┼─────────────┼─────────────────────────────┼──────────────────────────────┼─────────────────────────────┨
┃   👉   │ PRO   │ 97c9a968-15f8-42c9-8594-348 │ second-tenant       │ Silent Assassin's Creed │ 0.67.0  │ available   │ https://staging.cloud.zenml │ https://ea892d2b-zenml.stagi │ valid until 2024-10-25      ┃
┃        │       │ e1e45594e                   │                     │                         │         │             │ .io/organizations/bf873af9- │ ng.cloudinfra.zenml.io       │ 14:11:59 CEST (in           ┃
┃        │       │                             │                     │                         │         │             │ aaf9-4ad1-a08e-3dc6d910d590 │                              │ 22h42m20s)                  ┃
┃        │       │                             │                     │                         │         │             │ /tenants/97c9a968-15f8-42c9 │                              │                             ┃
┃        │       │                             │                     │                         │         │             │ -8594-348e1e45594e          │                              │                             ┃
┠────────┼───────┼─────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┼─────────────┼─────────────────────────────┼──────────────────────────────┼─────────────────────────────┨
┃        │ PRO   │ 16f8a35d-5c2f-44aa-a564-b34 │ test-zenml-login    │ Silent Assassin's Creed │ 0.67.0  │ available   │ https://staging.cloud.zenml │ https://64f4e58f-zenml.stagi │ valid until 2024-10-25      ┃
┃        │       │ 186fbf6d6                   │                     │                         │         │             │ .io/organizations/bf873af9- │ ng.cloudinfra.zenml.io       │ 14:11:59 CEST (in           ┃
┃        │       │                             │                     │                         │         │             │ aaf9-4ad1-a08e-3dc6d910d590 │                              │ 22h42m20s)                  ┃
┃        │       │                             │                     │                         │         │             │ /tenants/16f8a35d-5c2f-44aa │                              │                             ┃
┃        │       │                             │                     │                         │         │             │ -a564-b34186fbf6d6          │                              │                             ┃
┠────────┼───────┼─────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┼─────────────┼─────────────────────────────┼──────────────────────────────┼─────────────────────────────┨
┃        │ PRO   │ 6a11ed24-4fb0-479b-878e-0f9 │ TeamATeam           │ Hackathon               │ 0.58.2  │ deactivated │ https://staging.cloud.zenml │ https://d6f927b8-zenml.stagi │ valid until 2024-10-25      ┃
┃        │       │ 1857dba60                   │                     │                         │         │             │ .io/organizations/04962642- │ ng.cloudinfra.zenml.io       │ 14:11:59 CEST (in           ┃
┃        │       │                             │                     │                         │         │             │ c46c-42ae-a483-1b282952169b │                              │ 22h42m20s)                  ┃
┃        │       │                             │                     │                         │         │             │ /tenants/6a11ed24-4fb0-479b │                              │                             ┃
┃        │       │                             │                     │                         │         │             │ -878e-0f91857dba60          │                              │                             ┃
┠────────┼───────┼─────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┼─────────────┼─────────────────────────────┼──────────────────────────────┼─────────────────────────────┨
┃        │ PRO   │ 3e299f77-170e-4aba-9338-5a4 │ arbyte-tenant       │ Arbyte                  │ 0.66.0  │ available   │ https://staging.cloud.zenml │ https://5cb0b6d9-zenml.stagi │ valid until 2024-10-25      ┃
┃        │       │ 44446cf66                   │                     │                         │         │             │ .io/organizations/745b40f0- │ ng.cloudinfra.zenml.io       │ 14:11:59 CEST (in           ┃
┃        │       │                             │                     │                         │         │             │ 4e59-4a0e-9afd-cd55c655232c │                              │ 22h42m20s)                  ┃
┃        │       │                             │                     │                         │         │             │ /tenants/3e299f77-170e-4aba │                              │                             ┃
┃        │       │                             │                     │                         │         │             │ -9338-5a444446cf66          │                              │                             ┃
┠────────┼───────┼─────────────────────────────┼─────────────────────┼─────────────────────────┼─────────┼─────────────┼─────────────────────────────┼──────────────────────────────┼─────────────────────────────┨
┃        │ LOCAL │ 3d34ce75-8c33-48b5-8bd4-262 │ local daemon server │ N/A                     │ 0.67.0  │ active      │ http://127.0.0.1:8237       │ http://127.0.0.1:8237        │ no authentication required  ┃
┃        │       │ b1ed95eed                   │                     │                         │         │             │                             │                              │                             ┃
┗━━━━━━━━┷━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━┷━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
$ zenml status
-----ZenML Client Status-----
Connected to a ZenML Pro server: `second-tenant` [97c9a968-15f8-42c9-8594-348e1e45594e]
  ZenML Pro Organization: Silent Assassin's Creed
  ZenML Pro authentication: valid until 2024-10-25 14:11:59 CEST (in 22h42m48s)
  Dashboard: https://staging.cloud.zenml.io/organizations/bf873af9-aaf9-4ad1-a08e-3dc6d910d590/tenants/97c9a968-15f8-42c9-8594-348e1e45594e
  API: https://ea892d2b-zenml.staging.cloudinfra.zenml.io
  Server status: 'available'
  Server authentication: valid until 2024-10-24 16:23:44 CEST (in 54m33s)
  The active user is: '[email protected]'
  The active stack is: 'default' (global)
Using configuration from: '/home/stefan/.config/zenml'
Local store files are located at: '/home/stefan/.config/zenml/local_stores'

-----Local ZenML Server Status-----
The local daemon server is running at: http://127.0.0.1:8237
$ zenml login test-zenml-login
Connecting to ZenML Pro server: test-zenml-login [16f8a35d-5c2f-44aa-a564-b34186fbf6d6] 
Authenticating to ZenML server 'https://64f4e58f-zenml.staging.cloudinfra.zenml.io'...
The current global active stack is no longer available. Resetting the active stack to default.
Updated the global store configuration.
Connected to ZenML Pro server: test-zenml-login.

Pre-requisites

Please ensure you have done the following:

  • I have read the CONTRIBUTING.md document.
  • If my change requires a change to docs, I have updated the documentation accordingly.
  • I have added tests to cover my changes.
  • I have based my new branch on develop and the open PR is targeting develop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.
  • If my changes require changes to the dashboard, these changes are communicated/requested.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Other (add details above)

Copy link
Contributor

coderabbitai bot commented Oct 24, 2024

Important

Review skipped

Auto reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.


Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Generate unit testing code for this file.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai generate unit testing code for this file.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and generate unit testing code.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

@github-actions github-actions bot added internal To filter out internal PRs and issues enhancement New feature or request labels Oct 24, 2024
Copy link
Contributor

@schustmi schustmi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a pleasure to read, so much better than before!

src/zenml/analytics/context.py Outdated Show resolved Hide resolved
src/zenml/services/local/local_daemon_entrypoint.py Outdated Show resolved Hide resolved
src/zenml/zen_stores/base_zen_store.py Outdated Show resolved Hide resolved
src/zenml/zen_stores/base_zen_store.py Outdated Show resolved Hide resolved
src/zenml/zen_stores/rest_zen_store.py Outdated Show resolved Hide resolved
src/zenml/login/pro/tenant/models.py Outdated Show resolved Hide resolved
src/zenml/login/credentials.py Outdated Show resolved Hide resolved
src/zenml/login/credentials.py Outdated Show resolved Hide resolved
src/zenml/login/credentials.py Outdated Show resolved Hide resolved
src/zenml/login/credentials.py Outdated Show resolved Hide resolved
@schustmi
Copy link
Contributor

@stefannica I just tried this, and after trusting my device for 30 days, I get the following output in the CLI

2024-11-23      ┃
┃        │       │ 385a9795                     │                     │                  │         │             │ io/organizations/00c4d26e-d2 │ ng.cloudinfra.zenml.io       │ 00:18:16 CET (in            ┃
┃        │       │                              │                     │                  │         │             │ 65-4f65-80c5-767737e6b8b7/te │                              │ 28d11h59m0s)  

Seems like there is one day missing somehow?

@stefannica
Copy link
Contributor Author

I just tried this, and after trusting my device for 30 days, I get the following output in the CLI

This is a best practice I implemented regarding API tokens: you extract a leeway (percentage off token validity) from the original validity period in order to trigger or ask for a re-login before the tokens actually expire. 1 day out of 30 might be a bit to much though, but we can find a proper solution for this if you think it's not right. The current algorithm is implemented in the credentials store:

            # 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

@schustmi
Copy link
Contributor

I just tried this, and after trusting my device for 30 days, I get the following output in the CLI

This is a best practice I implemented regarding API tokens: you extract a leeway (percentage off token validity) from the original validity period in order to trigger or ask for a re-login before the tokens actually expire. 1 day out of 30 might be a bit to much though, but we can find a proper solution for this if you think it's not right. The current algorithm is implemented in the credentials store:

            # 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

No all good, just wanted to mention that this seemed a little odd. Wouldn't it make sense to ask for a re-login with the leeway included, but display the actual expiration date in the CLI?

In any case, none of that matters all too much, I'll leave it to you as the expert :)

@schustmi
Copy link
Contributor

@stefannica I have another question where I don't really understand what happens: When running a pipeline remotely, we set the API token in zenml.orchestrator.utils.get_config_environment_vars(...) using the ZENML_STORE_API_TOKEN. How is this being picked up currently?

@stefannica
Copy link
Contributor Author

@stefannica I have another question where I don't really understand what happens: When running a pipeline remotely, we set the API token in zenml.orchestrator.utils.get_config_environment_vars(...) using the ZENML_STORE_API_TOKEN. How is this being picked up currently?

Before I answer this, there's another change that is related to this that you should be aware of: the REST Zen Store configuration no longer stores an API key or an API token. It is still allowed to configure the ZENML_STORE_API_KEY and ZENML_STORE_API_TOKEN env vars and these attributes can still be loaded from the zenml global configuration file (for backwards compatibility), but the API key and API token are moved from the REST Zen Store configuration to the Credentials Store when it is initialized (check the new REST Zen Store configuration validator).

When a pipeline is run remotely, get_config_environment_vars calls get_api_token to generate a new API token scoped to the pipeline. This hasn't changed. What did change is that the API key, if configured, is no longer picked from the store configuration, it is taken from the Credentials Store.

@schustmi
Copy link
Contributor

Yeah I think when I was trying this, the latest commit which fixes this was not available yet. Now I see the API token being picked up correctly with set_bare_token()

src/zenml/analytics/context.py Outdated Show resolved Hide resolved
Copy link
Contributor

@schustmi schustmi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only two small questions left 🎉

src/zenml/zen_stores/rest_zen_store.py Show resolved Hide resolved
src/zenml/login/credentials.py Outdated Show resolved Hide resolved
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Copy link
Contributor

Classification template updates in examples/mlops_starter have been pushed.

Copy link
Contributor

E2E template updates in examples/e2e have been pushed.

Copy link
Contributor

NLP template updates in examples/e2e_nlp have been pushed.

@stefannica stefannica changed the title WIP: ZenML Pro web login implementation ZenML Pro web login implementation Oct 28, 2024
@schustmi schustmi merged commit 00d571e into develop Oct 29, 2024
70 of 71 checks passed
@schustmi schustmi deleted the feature/GROW-105-zenml-login branch October 29, 2024 13:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
breaking-change enhancement New feature or request internal To filter out internal PRs and issues run-slow-ci
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants