zenml-nightly 0.72.0.dev20250115__py3-none-any.whl → 0.72.0.dev20250117__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zenml/VERSION +1 -1
- zenml/cli/login.py +126 -50
- zenml/cli/server.py +24 -5
- zenml/config/server_config.py +142 -17
- zenml/constants.py +2 -11
- zenml/login/credentials.py +38 -14
- zenml/login/credentials_store.py +53 -18
- zenml/login/pro/client.py +3 -7
- zenml/login/pro/constants.py +0 -6
- zenml/login/pro/tenant/models.py +4 -2
- zenml/login/pro/utils.py +11 -25
- zenml/login/server_info.py +52 -0
- zenml/login/web_login.py +11 -6
- zenml/models/v2/misc/auth_models.py +1 -1
- zenml/models/v2/misc/server_models.py +44 -0
- zenml/zen_server/auth.py +97 -8
- zenml/zen_server/cloud_utils.py +79 -87
- zenml/zen_server/csrf.py +91 -0
- zenml/zen_server/deploy/helm/templates/NOTES.txt +22 -0
- zenml/zen_server/deploy/helm/templates/_environment.tpl +50 -24
- zenml/zen_server/deploy/helm/templates/server-secret.yaml +11 -0
- zenml/zen_server/deploy/helm/values.yaml +76 -7
- zenml/zen_server/feature_gate/feature_gate_interface.py +1 -1
- zenml/zen_server/jwt.py +16 -1
- zenml/zen_server/rbac/endpoint_utils.py +3 -3
- zenml/zen_server/routers/auth_endpoints.py +44 -21
- zenml/zen_server/routers/models_endpoints.py +1 -2
- zenml/zen_server/routers/pipelines_endpoints.py +2 -2
- zenml/zen_server/routers/stack_deployment_endpoints.py +5 -5
- zenml/zen_server/routers/workspaces_endpoints.py +2 -2
- zenml/zen_server/utils.py +64 -0
- zenml/zen_server/zen_server_api.py +5 -0
- zenml/zen_stores/base_zen_store.py +19 -1
- zenml/zen_stores/rest_zen_store.py +30 -20
- {zenml_nightly-0.72.0.dev20250115.dist-info → zenml_nightly-0.72.0.dev20250117.dist-info}/METADATA +3 -1
- {zenml_nightly-0.72.0.dev20250115.dist-info → zenml_nightly-0.72.0.dev20250117.dist-info}/RECORD +39 -37
- {zenml_nightly-0.72.0.dev20250115.dist-info → zenml_nightly-0.72.0.dev20250117.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.72.0.dev20250115.dist-info → zenml_nightly-0.72.0.dev20250117.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.72.0.dev20250115.dist-info → zenml_nightly-0.72.0.dev20250117.dist-info}/entry_points.txt +0 -0
zenml/VERSION
CHANGED
@@ -1 +1 @@
|
|
1
|
-
0.72.0.
|
1
|
+
0.72.0.dev20250117
|
zenml/cli/login.py
CHANGED
@@ -18,7 +18,7 @@ import os
|
|
18
18
|
import re
|
19
19
|
import sys
|
20
20
|
import time
|
21
|
-
from typing import Any, Dict, Optional, Union
|
21
|
+
from typing import Any, Dict, Optional, Tuple, Union
|
22
22
|
from uuid import UUID
|
23
23
|
|
24
24
|
import click
|
@@ -34,7 +34,8 @@ from zenml.exceptions import (
|
|
34
34
|
IllegalOperationError,
|
35
35
|
)
|
36
36
|
from zenml.logger import get_logger
|
37
|
-
from zenml.login.
|
37
|
+
from zenml.login.credentials import ServerType
|
38
|
+
from zenml.login.pro.constants import ZENML_PRO_API_URL
|
38
39
|
from zenml.login.web_login import web_login
|
39
40
|
from zenml.zen_server.utils import (
|
40
41
|
connected_to_local_server,
|
@@ -145,6 +146,7 @@ def connect_to_server(
|
|
145
146
|
api_key: Optional[str] = None,
|
146
147
|
verify_ssl: Union[str, bool] = True,
|
147
148
|
refresh: bool = False,
|
149
|
+
pro_server: bool = False,
|
148
150
|
) -> None:
|
149
151
|
"""Connect the client to a ZenML server or a SQL database.
|
150
152
|
|
@@ -154,6 +156,7 @@ def connect_to_server(
|
|
154
156
|
verify_ssl: Whether to verify the server's TLS certificate. If a string
|
155
157
|
is passed, it is interpreted as the path to a CA bundle file.
|
156
158
|
refresh: Whether to force a new login flow with the ZenML server.
|
159
|
+
pro_server: Whether the server is a ZenML Pro server.
|
157
160
|
"""
|
158
161
|
from zenml.login.credentials_store import get_credentials_store
|
159
162
|
from zenml.zen_stores.base_zen_store import BaseZenStore
|
@@ -170,7 +173,12 @@ def connect_to_server(
|
|
170
173
|
f"Authenticating to ZenML server '{url}' using an API key..."
|
171
174
|
)
|
172
175
|
credentials_store.set_api_key(url, api_key)
|
173
|
-
elif
|
176
|
+
elif pro_server:
|
177
|
+
# We don't have to do anything here assuming the user has already
|
178
|
+
# logged in to the ZenML Pro server using the ZenML Pro web login
|
179
|
+
# flow.
|
180
|
+
cli_utils.declare(f"Authenticating to ZenML server '{url}'...")
|
181
|
+
else:
|
174
182
|
if refresh or not credentials_store.has_valid_authentication(url):
|
175
183
|
cli_utils.declare(
|
176
184
|
f"Authenticating to ZenML server '{url}' using the web "
|
@@ -179,11 +187,6 @@ def connect_to_server(
|
|
179
187
|
web_login(url=url, verify_ssl=verify_ssl)
|
180
188
|
else:
|
181
189
|
cli_utils.declare(f"Connecting to ZenML server '{url}'...")
|
182
|
-
else:
|
183
|
-
# We don't have to do anything here assuming the user has already
|
184
|
-
# logged in to the ZenML Pro server using the ZenML Pro web login
|
185
|
-
# flow.
|
186
|
-
cli_utils.declare(f"Authenticating to ZenML server '{url}'...")
|
187
190
|
|
188
191
|
rest_store_config = RestZenStoreConfiguration(
|
189
192
|
url=url,
|
@@ -225,6 +228,7 @@ def connect_to_pro_server(
|
|
225
228
|
pro_server: Optional[str] = None,
|
226
229
|
api_key: Optional[str] = None,
|
227
230
|
refresh: bool = False,
|
231
|
+
pro_api_url: Optional[str] = None,
|
228
232
|
) -> None:
|
229
233
|
"""Connect the client to a ZenML Pro server.
|
230
234
|
|
@@ -233,6 +237,7 @@ def connect_to_pro_server(
|
|
233
237
|
If not provided, the web login flow will be initiated.
|
234
238
|
api_key: The API key to use to authenticate with the ZenML Pro server.
|
235
239
|
refresh: Whether to force a new login flow with the ZenML Pro server.
|
240
|
+
pro_api_url: The URL for the ZenML Pro API.
|
236
241
|
|
237
242
|
Raises:
|
238
243
|
ValueError: If incorrect parameters are provided.
|
@@ -243,6 +248,8 @@ def connect_to_pro_server(
|
|
243
248
|
from zenml.login.pro.client import ZenMLProClient
|
244
249
|
from zenml.login.pro.tenant.models import TenantStatus
|
245
250
|
|
251
|
+
pro_api_url = pro_api_url or ZENML_PRO_API_URL
|
252
|
+
|
246
253
|
server_id, server_url, server_name = None, None, None
|
247
254
|
login = False
|
248
255
|
if not pro_server:
|
@@ -264,20 +271,15 @@ def connect_to_pro_server(
|
|
264
271
|
server_name = pro_server
|
265
272
|
else:
|
266
273
|
server_url = pro_server
|
267
|
-
if not is_zenml_pro_server_url(server_url):
|
268
|
-
raise ValueError(
|
269
|
-
f"The URL '{server_url}' does not seem to belong to a ZenML Pro "
|
270
|
-
"server. Please check the server URL and try again."
|
271
|
-
)
|
272
274
|
|
273
275
|
credentials_store = get_credentials_store()
|
274
|
-
if not credentials_store.has_valid_pro_authentication():
|
276
|
+
if not credentials_store.has_valid_pro_authentication(pro_api_url):
|
275
277
|
# Without valid ZenML Pro credentials, we can only connect to a ZenML
|
276
278
|
# Pro server with an API key and we also need to know the URL of the
|
277
279
|
# server to connect to.
|
278
280
|
if api_key:
|
279
281
|
if server_url:
|
280
|
-
connect_to_server(server_url, api_key=api_key)
|
282
|
+
connect_to_server(server_url, api_key=api_key, pro_server=True)
|
281
283
|
return
|
282
284
|
else:
|
283
285
|
raise ValueError(
|
@@ -289,7 +291,9 @@ def connect_to_pro_server(
|
|
289
291
|
|
290
292
|
if login or refresh:
|
291
293
|
try:
|
292
|
-
token = web_login(
|
294
|
+
token = web_login(
|
295
|
+
pro_api_url=pro_api_url,
|
296
|
+
)
|
293
297
|
except AuthorizationException as e:
|
294
298
|
cli_utils.error(f"Authorization error: {e}")
|
295
299
|
|
@@ -320,7 +324,7 @@ def connect_to_pro_server(
|
|
320
324
|
# server argument passed to the command.
|
321
325
|
server_id = UUID(tenant_id)
|
322
326
|
|
323
|
-
client = ZenMLProClient()
|
327
|
+
client = ZenMLProClient(pro_api_url)
|
324
328
|
|
325
329
|
if server_id:
|
326
330
|
server = client.tenant.get(server_id)
|
@@ -405,7 +409,7 @@ def connect_to_pro_server(
|
|
405
409
|
f"Connecting to ZenML Pro server: {server.name} [{str(server.id)}] "
|
406
410
|
)
|
407
411
|
|
408
|
-
connect_to_server(server.url, api_key=api_key)
|
412
|
+
connect_to_server(server.url, api_key=api_key, pro_server=True)
|
409
413
|
|
410
414
|
# Update the stored server info with more accurate data taken from the
|
411
415
|
# ZenML Pro tenant object.
|
@@ -414,6 +418,42 @@ def connect_to_pro_server(
|
|
414
418
|
cli_utils.declare(f"Connected to ZenML Pro server: {server.name}.")
|
415
419
|
|
416
420
|
|
421
|
+
def is_pro_server(
|
422
|
+
url: str,
|
423
|
+
) -> Tuple[Optional[bool], Optional[str]]:
|
424
|
+
"""Check if the server at the given URL is a ZenML Pro server.
|
425
|
+
|
426
|
+
Args:
|
427
|
+
url: The URL of the server to check.
|
428
|
+
|
429
|
+
Returns:
|
430
|
+
True if the server is a ZenML Pro server, False otherwise, and the
|
431
|
+
extracted pro API URL if the server is a ZenML Pro server, or None if
|
432
|
+
no information could be extracted.
|
433
|
+
"""
|
434
|
+
from zenml.login.credentials_store import get_credentials_store
|
435
|
+
from zenml.login.server_info import get_server_info
|
436
|
+
|
437
|
+
# First, check the credentials store
|
438
|
+
credentials_store = get_credentials_store()
|
439
|
+
credentials = credentials_store.get_credentials(url)
|
440
|
+
if credentials:
|
441
|
+
if credentials.type == ServerType.PRO:
|
442
|
+
return True, credentials.pro_api_url
|
443
|
+
else:
|
444
|
+
return False, None
|
445
|
+
|
446
|
+
# Next, make a request to the server itself
|
447
|
+
server_info = get_server_info(url)
|
448
|
+
if not server_info:
|
449
|
+
return None, None
|
450
|
+
|
451
|
+
if server_info.is_pro_server():
|
452
|
+
return True, server_info.pro_api_url
|
453
|
+
|
454
|
+
return False, None
|
455
|
+
|
456
|
+
|
417
457
|
def _fail_if_authentication_environment_variables_set() -> None:
|
418
458
|
"""Fail if any of the authentication environment variables are set."""
|
419
459
|
environment_variables = [
|
@@ -650,6 +690,13 @@ def _fail_if_authentication_environment_variables_set() -> None:
|
|
650
690
|
"dashboard on a public domain. Primarily used for accessing the "
|
651
691
|
"dashboard in Colab. Only used when running `zenml login --local`.",
|
652
692
|
)
|
693
|
+
@click.option(
|
694
|
+
"--pro-api-url",
|
695
|
+
type=str,
|
696
|
+
default=None,
|
697
|
+
help="Custom URL for the ZenML Pro API. Useful when connecting "
|
698
|
+
"to a self-hosted ZenML Pro deployment.",
|
699
|
+
)
|
653
700
|
def login(
|
654
701
|
server: Optional[str] = None,
|
655
702
|
pro: bool = False,
|
@@ -667,6 +714,7 @@ def login(
|
|
667
714
|
blocking: bool = False,
|
668
715
|
image: Optional[str] = None,
|
669
716
|
ngrok_token: Optional[str] = None,
|
717
|
+
pro_api_url: Optional[str] = None,
|
670
718
|
) -> None:
|
671
719
|
"""Connect to a remote ZenML server.
|
672
720
|
|
@@ -691,6 +739,7 @@ def login(
|
|
691
739
|
ngrok_token: An ngrok auth token to use for exposing the local ZenML
|
692
740
|
dashboard on a public domain. Primarily used for accessing the
|
693
741
|
dashboard in Colab.
|
742
|
+
pro_api_url: Custom URL for the ZenML Pro API.
|
694
743
|
"""
|
695
744
|
_fail_if_authentication_environment_variables_set()
|
696
745
|
|
@@ -715,6 +764,7 @@ def login(
|
|
715
764
|
connect_to_pro_server(
|
716
765
|
pro_server=server,
|
717
766
|
refresh=True,
|
767
|
+
pro_api_url=pro_api_url,
|
718
768
|
)
|
719
769
|
return
|
720
770
|
|
@@ -740,29 +790,45 @@ def login(
|
|
740
790
|
)
|
741
791
|
|
742
792
|
if server is not None:
|
743
|
-
if
|
744
|
-
|
745
|
-
):
|
746
|
-
# The server argument is a ZenML Pro server URL, server name or UUID
|
793
|
+
if not re.match(r"^(https?|mysql)://", server):
|
794
|
+
# The server argument is a ZenML Pro server name or UUID
|
747
795
|
connect_to_pro_server(
|
748
796
|
pro_server=server,
|
749
797
|
api_key=api_key_value,
|
750
798
|
refresh=refresh,
|
799
|
+
pro_api_url=pro_api_url,
|
751
800
|
)
|
752
801
|
else:
|
753
|
-
|
754
|
-
|
755
|
-
|
756
|
-
|
757
|
-
|
758
|
-
|
802
|
+
# The server argument is a server URL
|
803
|
+
|
804
|
+
# First, try to discover if the server is a ZenML Pro server or not
|
805
|
+
server_is_pro, server_pro_api_url = is_pro_server(server)
|
806
|
+
if server_is_pro:
|
807
|
+
connect_to_pro_server(
|
808
|
+
pro_server=server,
|
809
|
+
api_key=api_key_value,
|
810
|
+
refresh=refresh,
|
811
|
+
# Prefer the pro API URL extracted from the server info if
|
812
|
+
# available
|
813
|
+
pro_api_url=server_pro_api_url or pro_api_url,
|
814
|
+
)
|
815
|
+
else:
|
816
|
+
connect_to_server(
|
817
|
+
url=server,
|
818
|
+
api_key=api_key_value,
|
819
|
+
verify_ssl=verify_ssl,
|
820
|
+
refresh=refresh,
|
821
|
+
)
|
759
822
|
|
760
823
|
elif current_non_local_server:
|
761
824
|
# The server argument is not provided, so we default to
|
762
825
|
# re-authenticating to the current non-local server that the client is
|
763
826
|
# connected to.
|
764
827
|
server = current_non_local_server
|
765
|
-
if
|
828
|
+
# First, try to discover if the server is a ZenML Pro server or not
|
829
|
+
server_is_pro, server_pro_api_url = is_pro_server(server)
|
830
|
+
|
831
|
+
if server_is_pro:
|
766
832
|
cli_utils.declare(
|
767
833
|
"No server argument was provided. Re-authenticating to "
|
768
834
|
"ZenML Pro...\n"
|
@@ -773,6 +839,9 @@ def login(
|
|
773
839
|
pro_server=server,
|
774
840
|
api_key=api_key_value,
|
775
841
|
refresh=True,
|
842
|
+
# Prefer the pro API URL extracted from the server info if
|
843
|
+
# available
|
844
|
+
pro_api_url=server_pro_api_url or pro_api_url,
|
776
845
|
)
|
777
846
|
else:
|
778
847
|
cli_utils.declare(
|
@@ -801,6 +870,7 @@ def login(
|
|
801
870
|
)
|
802
871
|
connect_to_pro_server(
|
803
872
|
api_key=api_key_value,
|
873
|
+
pro_api_url=pro_api_url,
|
804
874
|
)
|
805
875
|
|
806
876
|
|
@@ -857,11 +927,19 @@ def login(
|
|
857
927
|
default=False,
|
858
928
|
type=click.BOOL,
|
859
929
|
)
|
930
|
+
@click.option(
|
931
|
+
"--pro-api-url",
|
932
|
+
type=str,
|
933
|
+
default=None,
|
934
|
+
help="Custom URL for the ZenML Pro API. Useful when disconnecting "
|
935
|
+
"from a self-hosted ZenML Pro deployment.",
|
936
|
+
)
|
860
937
|
def logout(
|
861
938
|
server: Optional[str] = None,
|
862
939
|
local: bool = False,
|
863
940
|
clear: bool = False,
|
864
941
|
pro: bool = False,
|
942
|
+
pro_api_url: Optional[str] = None,
|
865
943
|
) -> None:
|
866
944
|
"""Disconnect from a ZenML server.
|
867
945
|
|
@@ -870,6 +948,7 @@ def logout(
|
|
870
948
|
clear: Clear all stored credentials and tokens.
|
871
949
|
local: Disconnect from the local ZenML server.
|
872
950
|
pro: Log out from ZenML Pro.
|
951
|
+
pro_api_url: Custom URL for the ZenML Pro API.
|
873
952
|
"""
|
874
953
|
from zenml.login.credentials_store import get_credentials_store
|
875
954
|
|
@@ -885,8 +964,9 @@ def logout(
|
|
885
964
|
"The `--pro` flag cannot be used with a specific server URL."
|
886
965
|
)
|
887
966
|
|
888
|
-
|
889
|
-
|
967
|
+
pro_api_url = pro_api_url or ZENML_PRO_API_URL
|
968
|
+
if credentials_store.has_valid_pro_authentication(pro_api_url):
|
969
|
+
credentials_store.clear_pro_credentials(pro_api_url)
|
890
970
|
cli_utils.declare("Logged out from ZenML Pro.")
|
891
971
|
else:
|
892
972
|
cli_utils.declare(
|
@@ -894,10 +974,13 @@ def logout(
|
|
894
974
|
)
|
895
975
|
|
896
976
|
if clear:
|
897
|
-
if
|
977
|
+
# Try to determine if the client is currently connected to a ZenML
|
978
|
+
# Pro server with the given pro API URL
|
979
|
+
credentials = credentials_store.get_credentials(store_cfg.url)
|
980
|
+
if credentials and credentials.pro_api_url == pro_api_url:
|
898
981
|
gc.set_default_store()
|
899
982
|
|
900
|
-
credentials_store.clear_all_pro_tokens()
|
983
|
+
credentials_store.clear_all_pro_tokens(pro_api_url)
|
901
984
|
cli_utils.declare("Logged out from all ZenML Pro servers.")
|
902
985
|
return
|
903
986
|
|
@@ -936,10 +1019,11 @@ def logout(
|
|
936
1019
|
|
937
1020
|
assert server is not None
|
938
1021
|
|
939
|
-
|
940
|
-
|
941
|
-
|
942
|
-
|
1022
|
+
gc.set_default_store()
|
1023
|
+
credentials = credentials_store.get_credentials(server)
|
1024
|
+
|
1025
|
+
if credentials and (clear or store_cfg.url == server):
|
1026
|
+
if credentials.type == ServerType.PRO:
|
943
1027
|
cli_utils.declare(
|
944
1028
|
f"Logging out from ZenML Pro server '{credentials.server_name}'."
|
945
1029
|
)
|
@@ -953,14 +1037,6 @@ def logout(
|
|
953
1037
|
"with 'zenml login <server-id-name-or-url>'."
|
954
1038
|
)
|
955
1039
|
else:
|
956
|
-
cli_utils.declare(
|
957
|
-
f"The client is not currently connected to the ZenML Pro server "
|
958
|
-
f"at '{server}'."
|
959
|
-
)
|
960
|
-
else:
|
961
|
-
gc.set_default_store()
|
962
|
-
credentials = credentials_store.get_credentials(server)
|
963
|
-
if credentials and (clear or store_cfg.url == server):
|
964
1040
|
cli_utils.declare(f"Logging out from {server}.")
|
965
1041
|
if clear:
|
966
1042
|
credentials_store.clear_credentials(server_url=server)
|
@@ -970,8 +1046,8 @@ def logout(
|
|
970
1046
|
"to the same server or 'zenml server list' to view other available "
|
971
1047
|
"servers that you can connect to with 'zenml login <server-url>'."
|
972
1048
|
)
|
973
|
-
|
974
|
-
|
975
|
-
|
976
|
-
|
977
|
-
|
1049
|
+
else:
|
1050
|
+
cli_utils.declare(
|
1051
|
+
f"The client is not currently connected to the ZenML server at "
|
1052
|
+
f"'{server}'."
|
1053
|
+
)
|
zenml/cli/server.py
CHANGED
@@ -187,6 +187,7 @@ def status() -> None:
|
|
187
187
|
"""Show details about the current configuration."""
|
188
188
|
from zenml.login.credentials_store import get_credentials_store
|
189
189
|
from zenml.login.pro.client import ZenMLProClient
|
190
|
+
from zenml.login.pro.constants import ZENML_PRO_API_URL
|
190
191
|
|
191
192
|
gc = GlobalConfiguration()
|
192
193
|
client = Client()
|
@@ -214,10 +215,11 @@ def status() -> None:
|
|
214
215
|
if server.type == ServerType.PRO:
|
215
216
|
# If connected to a ZenML Pro server, refresh the server info
|
216
217
|
pro_credentials = credentials_store.get_pro_credentials(
|
217
|
-
|
218
|
+
pro_api_url=server.pro_api_url or ZENML_PRO_API_URL,
|
219
|
+
allow_expired=False,
|
218
220
|
)
|
219
221
|
if pro_credentials:
|
220
|
-
pro_client = ZenMLProClient()
|
222
|
+
pro_client = ZenMLProClient(pro_credentials.url)
|
221
223
|
pro_servers = pro_client.tenant.list(
|
222
224
|
url=store_cfg.url, member_only=True
|
223
225
|
)
|
@@ -551,19 +553,36 @@ def server() -> None:
|
|
551
553
|
help="Show all ZenML servers, including those that are not running "
|
552
554
|
"and those with an expired authentication.",
|
553
555
|
)
|
554
|
-
|
556
|
+
@click.option(
|
557
|
+
"--pro-api-url",
|
558
|
+
type=str,
|
559
|
+
default=None,
|
560
|
+
help="Custom URL for the ZenML Pro API. Useful when disconnecting "
|
561
|
+
"from a self-hosted ZenML Pro deployment.",
|
562
|
+
)
|
563
|
+
def server_list(
|
564
|
+
verbose: bool = False,
|
565
|
+
all: bool = False,
|
566
|
+
pro_api_url: Optional[str] = None,
|
567
|
+
) -> None:
|
555
568
|
"""List all ZenML servers that this client is authorized to access.
|
556
569
|
|
557
570
|
Args:
|
558
571
|
verbose: Whether to show verbose output.
|
559
572
|
all: Whether to show all ZenML servers.
|
573
|
+
pro_api_url: Custom URL for the ZenML Pro API.
|
560
574
|
"""
|
561
575
|
from zenml.login.credentials_store import get_credentials_store
|
562
576
|
from zenml.login.pro.client import ZenMLProClient
|
577
|
+
from zenml.login.pro.constants import ZENML_PRO_API_URL
|
563
578
|
from zenml.login.pro.tenant.models import TenantRead, TenantStatus
|
564
579
|
|
580
|
+
pro_api_url = pro_api_url or ZENML_PRO_API_URL
|
581
|
+
|
565
582
|
credentials_store = get_credentials_store()
|
566
|
-
pro_token = credentials_store.get_pro_token(
|
583
|
+
pro_token = credentials_store.get_pro_token(
|
584
|
+
allow_expired=True, pro_api_url=pro_api_url
|
585
|
+
)
|
567
586
|
current_store_config = GlobalConfiguration().store_configuration
|
568
587
|
|
569
588
|
# The list of ZenML Pro servers kept in the credentials store
|
@@ -583,7 +602,7 @@ def server_list(verbose: bool = False, all: bool = False) -> None:
|
|
583
602
|
|
584
603
|
accessible_pro_servers: List[TenantRead] = []
|
585
604
|
try:
|
586
|
-
client = ZenMLProClient()
|
605
|
+
client = ZenMLProClient(pro_api_url)
|
587
606
|
accessible_pro_servers = client.tenant.list(member_only=not all)
|
588
607
|
except AuthorizationException as e:
|
589
608
|
cli_utils.warning(f"ZenML Pro authorization error: {e}")
|
zenml/config/server_config.py
CHANGED
@@ -11,7 +11,7 @@
|
|
11
11
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
12
12
|
# or implied. See the License for the specific language governing
|
13
13
|
# permissions and limitations under the License.
|
14
|
-
"""Functionality to support ZenML
|
14
|
+
"""Functionality to support ZenML Server Configuration."""
|
15
15
|
|
16
16
|
import json
|
17
17
|
import os
|
@@ -19,9 +19,17 @@ from secrets import token_hex
|
|
19
19
|
from typing import Any, Dict, List, Optional, Union
|
20
20
|
from uuid import UUID
|
21
21
|
|
22
|
-
from pydantic import
|
22
|
+
from pydantic import (
|
23
|
+
BaseModel,
|
24
|
+
ConfigDict,
|
25
|
+
Field,
|
26
|
+
PositiveInt,
|
27
|
+
field_validator,
|
28
|
+
model_validator,
|
29
|
+
)
|
23
30
|
|
24
31
|
from zenml.constants import (
|
32
|
+
DEFAULT_REPORTABLE_RESOURCES,
|
25
33
|
DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
|
26
34
|
DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
|
27
35
|
DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
|
@@ -44,6 +52,7 @@ from zenml.constants import (
|
|
44
52
|
DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
|
45
53
|
DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE,
|
46
54
|
ENV_ZENML_SERVER_PREFIX,
|
55
|
+
ENV_ZENML_SERVER_PRO_PREFIX,
|
47
56
|
)
|
48
57
|
from zenml.enums import AuthScheme
|
49
58
|
from zenml.logger import get_logger
|
@@ -127,10 +136,6 @@ class ServerConfiguration(BaseModel):
|
|
127
136
|
to use with the `EXTERNAL` authentication scheme.
|
128
137
|
external_user_info_url: The user info URL of an external authenticator
|
129
138
|
service to use with the `EXTERNAL` authentication scheme.
|
130
|
-
external_cookie_name: The name of the http-only cookie used to store the
|
131
|
-
bearer token used to authenticate with the external authenticator
|
132
|
-
service. Must be specified if the `EXTERNAL` authentication scheme
|
133
|
-
is used.
|
134
139
|
external_server_id: The ID of the ZenML server to use with the
|
135
140
|
`EXTERNAL` authentication scheme. If not specified, the regular
|
136
141
|
ZenML server ID is used.
|
@@ -246,7 +251,7 @@ class ServerConfiguration(BaseModel):
|
|
246
251
|
server_url: Optional[str] = None
|
247
252
|
dashboard_url: Optional[str] = None
|
248
253
|
root_url_path: str = ""
|
249
|
-
metadata: Dict[str,
|
254
|
+
metadata: Dict[str, str] = {}
|
250
255
|
auth_scheme: AuthScheme = AuthScheme.OAUTH2_PASSWORD_BEARER
|
251
256
|
jwt_token_algorithm: str = DEFAULT_ZENML_JWT_TOKEN_ALGORITHM
|
252
257
|
jwt_token_issuer: Optional[str] = None
|
@@ -276,11 +281,11 @@ class ServerConfiguration(BaseModel):
|
|
276
281
|
|
277
282
|
external_login_url: Optional[str] = None
|
278
283
|
external_user_info_url: Optional[str] = None
|
279
|
-
external_cookie_name: Optional[str] = None
|
280
284
|
external_server_id: Optional[UUID] = None
|
281
285
|
|
282
286
|
rbac_implementation_source: Optional[str] = None
|
283
287
|
feature_gate_implementation_source: Optional[str] = None
|
288
|
+
reportable_resources: List[str] = []
|
284
289
|
workload_manager_implementation_source: Optional[str] = None
|
285
290
|
pipeline_run_auth_window: int = (
|
286
291
|
DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW
|
@@ -370,14 +375,6 @@ class ServerConfiguration(BaseModel):
|
|
370
375
|
"authentication scheme."
|
371
376
|
)
|
372
377
|
|
373
|
-
# If the authentication scheme is set to `EXTERNAL`, the
|
374
|
-
# external cookie name must be specified.
|
375
|
-
if not data.get("external_cookie_name"):
|
376
|
-
raise ValueError(
|
377
|
-
"The external cookie name must be specified when "
|
378
|
-
"using the EXTERNAL authentication scheme."
|
379
|
-
)
|
380
|
-
|
381
378
|
if cors_allow_origins := data.get("cors_allow_origins"):
|
382
379
|
origins = cors_allow_origins.split(",")
|
383
380
|
data["cors_allow_origins"] = origins
|
@@ -552,12 +549,140 @@ class ServerConfiguration(BaseModel):
|
|
552
549
|
for k, v in os.environ.items():
|
553
550
|
if v == "":
|
554
551
|
continue
|
552
|
+
if k.startswith(ENV_ZENML_SERVER_PRO_PREFIX):
|
553
|
+
# Skip Pro configuration
|
554
|
+
continue
|
555
555
|
if k.startswith(ENV_ZENML_SERVER_PREFIX):
|
556
556
|
env_server_config[
|
557
557
|
k[len(ENV_ZENML_SERVER_PREFIX) :].lower()
|
558
558
|
] = v
|
559
559
|
|
560
|
-
|
560
|
+
server_config = ServerConfiguration(**env_server_config)
|
561
|
+
|
562
|
+
if server_config.deployment_type == ServerDeploymentType.CLOUD:
|
563
|
+
# If the zenml server is a Pro server, we will apply the Pro
|
564
|
+
# configuration overrides to the server config automatically.
|
565
|
+
# TODO: these should be retrieved dynamically from the ZenML Pro
|
566
|
+
# API.
|
567
|
+
server_pro_config = ServerProConfiguration.get_server_config()
|
568
|
+
server_config.auth_scheme = AuthScheme.EXTERNAL
|
569
|
+
server_config.external_login_url = (
|
570
|
+
f"{server_pro_config.dashboard_url}/api/auth/login"
|
571
|
+
)
|
572
|
+
server_config.external_user_info_url = (
|
573
|
+
f"{server_pro_config.api_url}/users/authorize_server"
|
574
|
+
)
|
575
|
+
server_config.external_server_id = server_pro_config.tenant_id
|
576
|
+
server_config.rbac_implementation_source = (
|
577
|
+
"zenml.zen_server.rbac.zenml_cloud_rbac.ZenMLCloudRBAC"
|
578
|
+
)
|
579
|
+
server_config.feature_gate_implementation_source = "zenml.zen_server.feature_gate.zenml_cloud_feature_gate.ZenMLCloudFeatureGateInterface"
|
580
|
+
server_config.reportable_resources = DEFAULT_REPORTABLE_RESOURCES
|
581
|
+
server_config.dashboard_url = f"{server_pro_config.dashboard_url}/organizations/{server_pro_config.organization_id}/tenants/{server_pro_config.tenant_id}"
|
582
|
+
server_config.metadata.update(
|
583
|
+
dict(
|
584
|
+
account_id=str(server_pro_config.organization_id),
|
585
|
+
organization_id=str(server_pro_config.organization_id),
|
586
|
+
tenant_id=str(server_pro_config.tenant_id),
|
587
|
+
)
|
588
|
+
)
|
589
|
+
if server_pro_config.tenant_name:
|
590
|
+
server_config.metadata.update(
|
591
|
+
dict(tenant_name=server_pro_config.tenant_name)
|
592
|
+
)
|
593
|
+
|
594
|
+
extra_cors_allow_origins = [
|
595
|
+
server_pro_config.dashboard_url,
|
596
|
+
server_pro_config.api_url,
|
597
|
+
]
|
598
|
+
if server_config.server_url:
|
599
|
+
extra_cors_allow_origins.append(server_config.server_url)
|
600
|
+
if (
|
601
|
+
not server_config.cors_allow_origins
|
602
|
+
or server_config.cors_allow_origins == ["*"]
|
603
|
+
):
|
604
|
+
server_config.cors_allow_origins = extra_cors_allow_origins
|
605
|
+
else:
|
606
|
+
server_config.cors_allow_origins += extra_cors_allow_origins
|
607
|
+
if "*" in server_config.cors_allow_origins:
|
608
|
+
# Remove the wildcard from the list
|
609
|
+
server_config.cors_allow_origins.remove("*")
|
610
|
+
|
611
|
+
# Remove duplicates
|
612
|
+
server_config.cors_allow_origins = list(
|
613
|
+
set(server_config.cors_allow_origins)
|
614
|
+
)
|
615
|
+
|
616
|
+
if server_config.jwt_token_expire_minutes is None:
|
617
|
+
server_config.jwt_token_expire_minutes = 60
|
618
|
+
|
619
|
+
return server_config
|
620
|
+
|
621
|
+
model_config = ConfigDict(
|
622
|
+
# Allow extra attributes from configs of previous ZenML versions to
|
623
|
+
# permit downgrading
|
624
|
+
extra="allow",
|
625
|
+
)
|
626
|
+
|
627
|
+
|
628
|
+
class ServerProConfiguration(BaseModel):
|
629
|
+
"""ZenML Server Pro configuration attributes.
|
630
|
+
|
631
|
+
All these attributes can be set through the environment with the
|
632
|
+
`ZENML_SERVER_PRO_`-Prefix. E.g. the value of the `ZENML_SERVER_PRO_API_URL`
|
633
|
+
environment variable will be extracted to api_url.
|
634
|
+
|
635
|
+
Attributes:
|
636
|
+
api_url: The ZenML Pro API URL.
|
637
|
+
dashboard_url: The ZenML Pro dashboard URL.
|
638
|
+
oauth2_client_secret: The ZenML Pro OAuth2 client secret used to
|
639
|
+
authenticate the ZenML server with the ZenML Pro API.
|
640
|
+
oauth2_audience: The OAuth2 audience.
|
641
|
+
organization_id: The ZenML Pro organization ID.
|
642
|
+
organization_name: The ZenML Pro organization name.
|
643
|
+
tenant_id: The ZenML Pro tenant ID.
|
644
|
+
tenant_name: The ZenML Pro tenant name.
|
645
|
+
"""
|
646
|
+
|
647
|
+
api_url: str
|
648
|
+
dashboard_url: str
|
649
|
+
oauth2_client_secret: str
|
650
|
+
oauth2_audience: str
|
651
|
+
organization_id: UUID
|
652
|
+
organization_name: Optional[str] = None
|
653
|
+
tenant_id: UUID
|
654
|
+
tenant_name: Optional[str] = None
|
655
|
+
|
656
|
+
@field_validator("api_url", "dashboard_url")
|
657
|
+
@classmethod
|
658
|
+
def _strip_trailing_slashes_url(cls, url: str) -> str:
|
659
|
+
"""Strip any trailing slashes on the API URL.
|
660
|
+
|
661
|
+
Args:
|
662
|
+
url: The API URL.
|
663
|
+
|
664
|
+
Returns:
|
665
|
+
The API URL with potential trailing slashes removed.
|
666
|
+
"""
|
667
|
+
return url.rstrip("/")
|
668
|
+
|
669
|
+
@classmethod
|
670
|
+
def get_server_config(cls) -> "ServerProConfiguration":
|
671
|
+
"""Get the server Pro configuration.
|
672
|
+
|
673
|
+
Returns:
|
674
|
+
The server Pro configuration.
|
675
|
+
"""
|
676
|
+
env_server_config: Dict[str, Any] = {}
|
677
|
+
for k, v in os.environ.items():
|
678
|
+
if v == "":
|
679
|
+
continue
|
680
|
+
if k.startswith(ENV_ZENML_SERVER_PRO_PREFIX):
|
681
|
+
env_server_config[
|
682
|
+
k[len(ENV_ZENML_SERVER_PRO_PREFIX) :].lower()
|
683
|
+
] = v
|
684
|
+
|
685
|
+
return ServerProConfiguration(**env_server_config)
|
561
686
|
|
562
687
|
model_config = ConfigDict(
|
563
688
|
# Allow extra attributes from configs of previous ZenML versions to
|