zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__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.
Files changed (125) hide show
  1. README.md +17 -11
  2. RELEASE_NOTES.md +9 -0
  3. zenml/VERSION +1 -1
  4. zenml/__init__.py +1 -1
  5. zenml/analytics/context.py +16 -1
  6. zenml/analytics/utils.py +18 -7
  7. zenml/artifacts/utils.py +40 -216
  8. zenml/cli/__init__.py +63 -90
  9. zenml/cli/base.py +3 -3
  10. zenml/cli/login.py +951 -0
  11. zenml/cli/server.py +462 -353
  12. zenml/cli/service_accounts.py +4 -4
  13. zenml/cli/stack.py +77 -2
  14. zenml/cli/stack_components.py +5 -16
  15. zenml/cli/user_management.py +0 -12
  16. zenml/cli/utils.py +24 -77
  17. zenml/client.py +46 -14
  18. zenml/config/compiler.py +1 -0
  19. zenml/config/global_config.py +9 -0
  20. zenml/config/pipeline_configurations.py +2 -1
  21. zenml/config/pipeline_run_configuration.py +2 -1
  22. zenml/constants.py +3 -9
  23. zenml/enums.py +1 -1
  24. zenml/exceptions.py +11 -0
  25. zenml/integrations/github/code_repositories/github_code_repository.py +1 -1
  26. zenml/login/__init__.py +16 -0
  27. zenml/login/credentials.py +346 -0
  28. zenml/login/credentials_store.py +603 -0
  29. zenml/login/pro/__init__.py +16 -0
  30. zenml/login/pro/client.py +496 -0
  31. zenml/login/pro/constants.py +34 -0
  32. zenml/login/pro/models.py +25 -0
  33. zenml/login/pro/organization/__init__.py +14 -0
  34. zenml/login/pro/organization/client.py +79 -0
  35. zenml/login/pro/organization/models.py +32 -0
  36. zenml/login/pro/tenant/__init__.py +14 -0
  37. zenml/login/pro/tenant/client.py +92 -0
  38. zenml/login/pro/tenant/models.py +174 -0
  39. zenml/login/pro/utils.py +121 -0
  40. zenml/{cli → login}/web_login.py +64 -28
  41. zenml/materializers/base_materializer.py +43 -9
  42. zenml/materializers/built_in_materializer.py +1 -1
  43. zenml/metadata/metadata_types.py +49 -0
  44. zenml/model/model.py +0 -38
  45. zenml/models/__init__.py +3 -0
  46. zenml/models/v2/base/base.py +12 -8
  47. zenml/models/v2/base/filter.py +9 -0
  48. zenml/models/v2/core/artifact_version.py +49 -10
  49. zenml/models/v2/core/component.py +54 -19
  50. zenml/models/v2/core/flavor.py +13 -13
  51. zenml/models/v2/core/model.py +3 -1
  52. zenml/models/v2/core/model_version.py +3 -5
  53. zenml/models/v2/core/model_version_artifact.py +3 -1
  54. zenml/models/v2/core/model_version_pipeline_run.py +3 -1
  55. zenml/models/v2/core/pipeline.py +3 -1
  56. zenml/models/v2/core/pipeline_run.py +23 -1
  57. zenml/models/v2/core/run_template.py +3 -1
  58. zenml/models/v2/core/stack.py +7 -3
  59. zenml/models/v2/core/step_run.py +43 -2
  60. zenml/models/v2/misc/auth_models.py +11 -2
  61. zenml/models/v2/misc/server_models.py +2 -0
  62. zenml/orchestrators/base_orchestrator.py +8 -4
  63. zenml/orchestrators/step_launcher.py +1 -0
  64. zenml/orchestrators/step_run_utils.py +10 -2
  65. zenml/orchestrators/step_runner.py +67 -55
  66. zenml/orchestrators/utils.py +45 -22
  67. zenml/pipelines/pipeline_decorator.py +5 -0
  68. zenml/pipelines/pipeline_definition.py +206 -160
  69. zenml/pipelines/run_utils.py +11 -10
  70. zenml/services/local/local_daemon_entrypoint.py +4 -4
  71. zenml/services/service.py +2 -2
  72. zenml/stack/stack.py +2 -6
  73. zenml/stack/stack_component.py +2 -7
  74. zenml/stack/utils.py +26 -14
  75. zenml/steps/base_step.py +8 -2
  76. zenml/steps/step_context.py +0 -3
  77. zenml/steps/step_invocation.py +14 -5
  78. zenml/steps/utils.py +1 -0
  79. zenml/utils/materializer_utils.py +1 -1
  80. zenml/utils/requirements_utils.py +71 -0
  81. zenml/utils/singleton.py +15 -3
  82. zenml/utils/source_utils.py +39 -2
  83. zenml/utils/visualization_utils.py +1 -1
  84. zenml/zen_server/auth.py +44 -39
  85. zenml/zen_server/deploy/__init__.py +7 -7
  86. zenml/zen_server/deploy/base_provider.py +46 -73
  87. zenml/zen_server/deploy/{local → daemon}/__init__.py +3 -3
  88. zenml/zen_server/deploy/{local/local_provider.py → daemon/daemon_provider.py} +44 -63
  89. zenml/zen_server/deploy/{local/local_zen_server.py → daemon/daemon_zen_server.py} +50 -22
  90. zenml/zen_server/deploy/deployer.py +90 -171
  91. zenml/zen_server/deploy/deployment.py +20 -12
  92. zenml/zen_server/deploy/docker/docker_provider.py +9 -28
  93. zenml/zen_server/deploy/docker/docker_zen_server.py +19 -3
  94. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  95. zenml/zen_server/deploy/helm/README.md +2 -2
  96. zenml/zen_server/exceptions.py +11 -0
  97. zenml/zen_server/jwt.py +9 -9
  98. zenml/zen_server/routers/auth_endpoints.py +30 -8
  99. zenml/zen_server/routers/stack_components_endpoints.py +1 -1
  100. zenml/zen_server/routers/workspaces_endpoints.py +1 -1
  101. zenml/zen_server/template_execution/runner_entrypoint_configuration.py +7 -4
  102. zenml/zen_server/template_execution/utils.py +6 -61
  103. zenml/zen_server/utils.py +64 -36
  104. zenml/zen_stores/base_zen_store.py +4 -49
  105. zenml/zen_stores/migrations/versions/0.68.1_release.py +23 -0
  106. zenml/zen_stores/migrations/versions/c22561cbb3a9_add_artifact_unique_constraints.py +86 -0
  107. zenml/zen_stores/rest_zen_store.py +325 -147
  108. zenml/zen_stores/schemas/api_key_schemas.py +9 -4
  109. zenml/zen_stores/schemas/artifact_schemas.py +21 -2
  110. zenml/zen_stores/schemas/artifact_visualization_schemas.py +1 -1
  111. zenml/zen_stores/schemas/component_schemas.py +49 -6
  112. zenml/zen_stores/schemas/device_schemas.py +9 -4
  113. zenml/zen_stores/schemas/flavor_schemas.py +1 -1
  114. zenml/zen_stores/schemas/model_schemas.py +1 -1
  115. zenml/zen_stores/schemas/service_schemas.py +1 -1
  116. zenml/zen_stores/schemas/step_run_schemas.py +1 -1
  117. zenml/zen_stores/schemas/trigger_schemas.py +1 -1
  118. zenml/zen_stores/sql_zen_store.py +393 -140
  119. zenml/zen_stores/template_utils.py +3 -1
  120. {zenml_nightly-0.68.0.dev20241027.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/METADATA +18 -12
  121. {zenml_nightly-0.68.0.dev20241027.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/RECORD +124 -107
  122. zenml/api.py +0 -60
  123. {zenml_nightly-0.68.0.dev20241027.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/LICENSE +0 -0
  124. {zenml_nightly-0.68.0.dev20241027.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/WHEEL +0 -0
  125. {zenml_nightly-0.68.0.dev20241027.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/entry_points.txt +0 -0
@@ -37,6 +37,7 @@ from pydantic import (
37
37
  BaseModel,
38
38
  ConfigDict,
39
39
  Field,
40
+ ValidationError,
40
41
  field_validator,
41
42
  model_validator,
42
43
  )
@@ -112,9 +113,19 @@ from zenml.enums import (
112
113
  StackDeploymentProvider,
113
114
  StoreType,
114
115
  )
115
- from zenml.exceptions import AuthorizationException, MethodNotAllowedError
116
+ from zenml.exceptions import (
117
+ AuthorizationException,
118
+ CredentialsNotValid,
119
+ MethodNotAllowedError,
120
+ )
116
121
  from zenml.io import fileio
117
122
  from zenml.logger import get_logger
123
+ from zenml.login.credentials import APIToken
124
+ from zenml.login.credentials_store import get_credentials_store
125
+ from zenml.login.pro.utils import (
126
+ get_troubleshooting_instructions,
127
+ is_zenml_pro_server_url,
128
+ )
118
129
  from zenml.models import (
119
130
  ActionFilter,
120
131
  ActionRequest,
@@ -173,6 +184,7 @@ from zenml.models import (
173
184
  OAuthDeviceFilter,
174
185
  OAuthDeviceResponse,
175
186
  OAuthDeviceUpdate,
187
+ OAuthTokenResponse,
176
188
  Page,
177
189
  PipelineBuildFilter,
178
190
  PipelineBuildRequest,
@@ -256,6 +268,7 @@ from zenml.service_connectors.service_connector_registry import (
256
268
  from zenml.utils.networking_utils import (
257
269
  replace_localhost_with_internal_hostname,
258
270
  )
271
+ from zenml.utils.pydantic_utils import before_validator_handler
259
272
  from zenml.zen_server.exceptions import exception_from_response
260
273
  from zenml.zen_stores.base_zen_store import BaseZenStore
261
274
 
@@ -281,10 +294,10 @@ class RestZenStoreConfiguration(StoreConfiguration):
281
294
  username: The username to use to connect to the Zen server.
282
295
  password: The password to use to connect to the Zen server.
283
296
  api_key: The service account API key to use to connect to the Zen
284
- server.
285
- api_token: The API token to use to connect to the Zen server. Generated
286
- by the client and stored in the configuration file on the first
287
- login and every time the API key is refreshed.
297
+ server. This is only set if the API key is configured explicitly via
298
+ environment variables or the ZenML global configuration file. API
299
+ keys configured via the CLI are stored in the credentials store
300
+ instead.
288
301
  verify_ssl: Either a boolean, in which case it controls whether we
289
302
  verify the server's TLS certificate, or a string, in which case it
290
303
  must be a path to a CA bundle to use or the CA bundle value itself.
@@ -294,32 +307,11 @@ class RestZenStoreConfiguration(StoreConfiguration):
294
307
 
295
308
  type: StoreType = StoreType.REST
296
309
 
297
- username: Optional[str] = None
298
- password: Optional[str] = None
299
- api_key: Optional[str] = None
300
- api_token: Optional[str] = None
301
- verify_ssl: Union[bool, str] = Field(True, union_mode="left_to_right")
310
+ verify_ssl: Union[bool, str] = Field(
311
+ default=True, union_mode="left_to_right"
312
+ )
302
313
  http_timeout: int = DEFAULT_HTTP_TIMEOUT
303
314
 
304
- @model_validator(mode="after")
305
- def validate_credentials(self) -> "RestZenStoreConfiguration":
306
- """Validates the credentials provided in the values dictionary.
307
-
308
- Raises:
309
- ValueError: If neither api_token nor username nor api_key is set.
310
-
311
- Returns:
312
- The values dictionary.
313
- """
314
- # Check if the values dictionary contains either an API token, an API
315
- # key or a username as non-empty strings.
316
- if self.api_token or self.username or self.api_key:
317
- return self
318
- raise ValueError(
319
- "Neither api_token nor username nor api_key is set in the "
320
- "store config."
321
- )
322
-
323
315
  @field_validator("url")
324
316
  @classmethod
325
317
  def validate_url(cls, url: str) -> str:
@@ -406,13 +398,47 @@ class RestZenStoreConfiguration(StoreConfiguration):
406
398
  with open(self.verify_ssl, "r") as f:
407
399
  self.verify_ssl = f.read()
408
400
 
401
+ @model_validator(mode="before")
402
+ @classmethod
403
+ @before_validator_handler
404
+ def _move_credentials(cls, data: Dict[str, Any]) -> Dict[str, Any]:
405
+ """Moves credentials (API keys, API tokens, passwords) from the config to the credentials store.
406
+
407
+ Args:
408
+ data: The values dict used to instantiate the model.
409
+
410
+ Returns:
411
+ The values dict without credentials.
412
+ """
413
+ url = data.get("url")
414
+ if not url:
415
+ return data
416
+
417
+ url = replace_localhost_with_internal_hostname(url)
418
+
419
+ if api_token := data.pop("api_token", None):
420
+ credentials_store = get_credentials_store()
421
+ credentials_store.set_bare_token(url, api_token)
422
+
423
+ username = data.pop("username", None)
424
+ password = data.pop("password", None)
425
+ if username is not None and password is not None:
426
+ credentials_store = get_credentials_store()
427
+ credentials_store.set_password(url, username, password)
428
+
429
+ if api_key := data.pop("api_key", None):
430
+ credentials_store = get_credentials_store()
431
+ credentials_store.set_api_key(url, api_key)
432
+
433
+ return data
434
+
409
435
  model_config = ConfigDict(
410
436
  # Don't validate attributes when assigning them. This is necessary
411
437
  # because the `verify_ssl` attribute can be expanded to the contents
412
438
  # of the certificate file.
413
439
  validate_assignment=False,
414
- # Forbid extra attributes set in the class.
415
- extra="forbid",
440
+ # Ignore extra attributes set in the class.
441
+ extra="ignore",
416
442
  )
417
443
 
418
444
 
@@ -422,7 +448,7 @@ class RestZenStore(BaseZenStore):
422
448
  config: RestZenStoreConfiguration
423
449
  TYPE: ClassVar[StoreType] = StoreType.REST
424
450
  CONFIG_TYPE: ClassVar[Type[StoreConfiguration]] = RestZenStoreConfiguration
425
- _api_token: Optional[str] = None
451
+ _api_token: Optional[APIToken] = None
426
452
  _session: Optional[requests.Session] = None
427
453
 
428
454
  # ====================================
@@ -434,9 +460,54 @@ class RestZenStore(BaseZenStore):
434
460
  # --------------------------------
435
461
 
436
462
  def _initialize(self) -> None:
437
- """Initialize the REST store."""
438
- client_version = zenml.__version__
439
- server_version = self.get_store_info().version
463
+ """Initialize the REST store.
464
+
465
+ Raises:
466
+ RuntimeError: If the store cannot be initialized.
467
+ AuthorizationException: If the store cannot be initialized due to
468
+ authentication errors.
469
+ """
470
+ try:
471
+ client_version = zenml.__version__
472
+ server_version = self.get_store_info().version
473
+
474
+ # Handle cases where the ZenML server is not available
475
+ except ConnectionError as e:
476
+ error_message = (
477
+ f"Cannot connect to the ZenML server at {self.url}."
478
+ )
479
+ if urlparse(self.url).hostname in [
480
+ "localhost",
481
+ "127.0.0.1",
482
+ "host.docker.internal",
483
+ ]:
484
+ recommendation = (
485
+ "Please run `zenml login --local --restart` to restart the "
486
+ "server."
487
+ )
488
+ else:
489
+ recommendation = (
490
+ f"Please run `zenml login {self.url}` to reconnect to the "
491
+ "server."
492
+ )
493
+ raise RuntimeError(f"{error_message}\n{recommendation}") from e
494
+
495
+ except AuthorizationException as e:
496
+ raise AuthorizationException(
497
+ f"Authorization failed for store at '{self.url}'. Please check "
498
+ f"your credentials: {str(e)}"
499
+ )
500
+
501
+ except Exception as e:
502
+ zenml_pro_extra = ""
503
+ if is_zenml_pro_server_url(self.url):
504
+ zenml_pro_extra = (
505
+ "\nHINT: " + get_troubleshooting_instructions(self.url)
506
+ )
507
+ raise RuntimeError(
508
+ f"Error connecting to URL "
509
+ f"'{self.url}': {str(e)}" + zenml_pro_extra
510
+ ) from e
440
511
 
441
512
  if not DISABLE_CLIENT_SERVER_MISMATCH_WARNING and (
442
513
  server_version != client_version
@@ -639,20 +710,6 @@ class RestZenStore(BaseZenStore):
639
710
  params={"hydrate": hydrate},
640
711
  )
641
712
 
642
- def set_api_key(self, api_key: str) -> None:
643
- """Set the API key to use for authentication.
644
-
645
- Args:
646
- api_key: The API key to use for authentication.
647
- """
648
- self.config.api_key = api_key
649
- self.clear_session()
650
- # TODO: find a way to persist the API key in the configuration file
651
- # without calling _write_config() here.
652
- # This is the only place where we need to explicitly call
653
- # _write_config() to persist the global configuration.
654
- GlobalConfiguration()._write_config()
655
-
656
713
  def list_api_keys(
657
714
  self,
658
715
  service_account_id: UUID,
@@ -3979,76 +4036,166 @@ class RestZenStore(BaseZenStore):
3979
4036
  # Internal helper methods
3980
4037
  # =======================
3981
4038
 
3982
- def _get_auth_token(self) -> str:
3983
- """Get the authentication token for the REST store.
4039
+ def get_or_generate_api_token(self) -> str:
4040
+ """Get or generate an API token.
3984
4041
 
3985
4042
  Returns:
3986
- The authentication token.
4043
+ The API token.
3987
4044
 
3988
4045
  Raises:
3989
- ValueError: if the response from the server isn't in the right
3990
- format.
3991
- """
3992
- if self._api_token is None:
3993
- # Check if the API token is already stored in the config
3994
- if self.config.api_token:
3995
- self._api_token = self.config.api_token
3996
- # Check if the username and password are provided in the config
3997
- elif (
3998
- self.config.username is not None
3999
- and self.config.password is not None
4000
- or self.config.api_key is not None
4001
- ):
4002
- data: Optional[Dict[str, str]] = None
4003
- if self.config.api_key is not None:
4004
- data = {
4005
- "grant_type": OAuthGrantTypes.ZENML_API_KEY.value,
4006
- "password": self.config.api_key,
4007
- }
4008
- elif (
4009
- self.config.username is not None
4010
- and self.config.password is not None
4011
- ):
4012
- data = {
4013
- "grant_type": OAuthGrantTypes.OAUTH_PASSWORD.value,
4014
- "username": self.config.username,
4015
- "password": self.config.password,
4016
- }
4017
-
4018
- response = self._handle_response(
4019
- requests.post(
4020
- self.url + API + VERSION_1 + LOGIN,
4021
- data=data,
4022
- verify=self.config.verify_ssl,
4023
- timeout=self.config.http_timeout,
4024
- )
4046
+ CredentialsNotValid: if an API token cannot be fetched or
4047
+ generated because the client credentials are not valid.
4048
+ """
4049
+ if self._api_token is None or self._api_token.expired:
4050
+ # Check if a valid API token is already in the cache
4051
+ credentials_store = get_credentials_store()
4052
+ credentials = credentials_store.get_credentials(self.url)
4053
+ token = credentials.api_token if credentials else None
4054
+ if credentials and token and not token.expired:
4055
+ self._api_token = token
4056
+
4057
+ # Populate the server info in the credentials store if it is
4058
+ # not already present
4059
+ if not credentials.server_id:
4060
+ try:
4061
+ server_info = self.get_store_info()
4062
+ except Exception as e:
4063
+ logger.warning(f"Failed to get server info: {e}.")
4064
+ else:
4065
+ credentials_store.update_server_info(
4066
+ self.url, server_info
4067
+ )
4068
+
4069
+ return self._api_token.access_token
4070
+
4071
+ # Token is expired or not found in the cache. Time to get a new one.
4072
+
4073
+ if not token:
4074
+ logger.debug(f"Authenticating to {self.url}")
4075
+ else:
4076
+ logger.debug(
4077
+ f"Authentication token for {self.url} expired; refreshing..."
4025
4078
  )
4026
- if (
4027
- not isinstance(response, dict)
4028
- or "access_token" not in response
4029
- ):
4030
- raise ValueError(
4031
- f"Bad API Response. Expected access token dict, got "
4032
- f"{type(response)}"
4079
+
4080
+ data: Optional[Dict[str, str]] = None
4081
+ headers: Dict[str, str] = {}
4082
+
4083
+ # Check if an API key is configured
4084
+ api_key = credentials_store.get_api_key(self.url)
4085
+
4086
+ # Check if username and password are configured
4087
+ username, password = credentials_store.get_password(self.url)
4088
+
4089
+ api_key_hint = (
4090
+ "\nHint: If you're getting this error in an automated, "
4091
+ "non-interactive workload like a pipeline run or a CI/CD job, "
4092
+ "you should use a service account API key to authenticate to "
4093
+ "the server instead of temporary CLI login credentials. For "
4094
+ "more information, see "
4095
+ "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
4096
+ )
4097
+
4098
+ if api_key is not None:
4099
+ # An API key is configured. Use it as a password to
4100
+ # authenticate.
4101
+ data = {
4102
+ "grant_type": OAuthGrantTypes.ZENML_API_KEY.value,
4103
+ "password": api_key,
4104
+ }
4105
+ elif username is not None and password is not None:
4106
+ # Username and password are configured. Use them to authenticate.
4107
+ data = {
4108
+ "grant_type": OAuthGrantTypes.OAUTH_PASSWORD.value,
4109
+ "username": username,
4110
+ "password": password,
4111
+ }
4112
+ elif is_zenml_pro_server_url(self.url):
4113
+ # ZenML Pro tenants use a proprietary authorization grant
4114
+ # where the ZenML Pro API session token is exchanged for a
4115
+ # regular ZenML server access token.
4116
+
4117
+ # Get the ZenML Pro API session token, if cached and valid
4118
+ pro_token = credentials_store.get_pro_token(allow_expired=True)
4119
+ if not pro_token:
4120
+ raise CredentialsNotValid(
4121
+ "You need to be logged in to ZenML Pro in order to "
4122
+ f"access the ZenML Pro server '{self.url}'. Please run "
4123
+ "'zenml login' to log in or choose a different server."
4124
+ + api_key_hint
4125
+ )
4126
+
4127
+ elif pro_token.expired:
4128
+ raise CredentialsNotValid(
4129
+ "Your ZenML Pro login session has expired. "
4130
+ "Please log in again using 'zenml login'."
4131
+ + api_key_hint
4033
4132
  )
4034
- self._api_token = response["access_token"]
4035
- self.config.api_token = self._api_token
4133
+
4134
+ data = {
4135
+ "grant_type": OAuthGrantTypes.ZENML_EXTERNAL.value,
4136
+ }
4137
+ headers.update(
4138
+ {"Authorization": "Bearer " + pro_token.access_token}
4139
+ )
4036
4140
  else:
4037
- raise ValueError(
4038
- "No API token, API key or username/password provided. "
4039
- "Please provide either an API token, an API key or a "
4040
- "username and password in the ZenML config."
4141
+ if not token:
4142
+ raise CredentialsNotValid(
4143
+ "No valid credentials found. Please run 'zenml login "
4144
+ f"--url {self.url}' to connect to the current server."
4145
+ + api_key_hint
4146
+ )
4147
+ elif token.expired:
4148
+ raise CredentialsNotValid(
4149
+ "Your authentication to the current server has expired. "
4150
+ "Please log in again using 'zenml login --url "
4151
+ f"{self.url}'." + api_key_hint
4152
+ )
4153
+
4154
+ response = self._handle_response(
4155
+ requests.post(
4156
+ self.url + API + VERSION_1 + LOGIN,
4157
+ data=data,
4158
+ verify=self.config.verify_ssl,
4159
+ timeout=self.config.http_timeout,
4160
+ headers=headers,
4041
4161
  )
4042
- return self._api_token
4162
+ )
4163
+ try:
4164
+ token_response = OAuthTokenResponse.model_validate(response)
4165
+ except ValidationError as e:
4166
+ raise CredentialsNotValid(
4167
+ "Unexpected response received while authenticating to "
4168
+ f"the server {e}"
4169
+ ) from e
4170
+
4171
+ # Cache the token
4172
+ self._api_token = credentials_store.set_token(
4173
+ self.url, token_response
4174
+ )
4175
+
4176
+ # Update the server info in the credentials store with the latest
4177
+ # information from the server.
4178
+ # NOTE: this is the best place to do this because we know that
4179
+ # the token is valid and the server is reachable.
4180
+ try:
4181
+ server_info = self.get_store_info()
4182
+ except Exception as e:
4183
+ logger.warning(f"Failed to get server info: {e}.")
4184
+ else:
4185
+ credentials_store.update_server_info(self.url, server_info)
4186
+
4187
+ return self._api_token.access_token
4043
4188
 
4044
4189
  @property
4045
4190
  def session(self) -> requests.Session:
4046
- """Authenticate to the ZenML server.
4191
+ """Initialize and return a requests session.
4047
4192
 
4048
4193
  Returns:
4049
- A requests session with the authentication token.
4194
+ A requests session.
4050
4195
  """
4051
4196
  if self._session is None:
4197
+ # We only need to initialize the session once over the lifetime
4198
+ # of the client. We can swap the token out when it expires.
4052
4199
  if self.config.verify_ssl is False:
4053
4200
  urllib3.disable_warnings(
4054
4201
  urllib3.exceptions.InsecureRequestWarning
@@ -4059,40 +4206,65 @@ class RestZenStore(BaseZenStore):
4059
4206
  self._session.mount("https://", HTTPAdapter(max_retries=retries))
4060
4207
  self._session.mount("http://", HTTPAdapter(max_retries=retries))
4061
4208
  self._session.verify = self.config.verify_ssl
4062
- token = self._get_auth_token()
4063
- self._session.headers.update({"Authorization": "Bearer " + token})
4064
- logger.debug("Authenticated to ZenML server.")
4209
+
4210
+ # Note that we return an unauthenticated session here. An API token
4211
+ # is only fetched and set in the authorization header when and if it is
4212
+ # needed.
4065
4213
  return self._session
4066
4214
 
4067
- def clear_session(self) -> None:
4068
- """Clear the authentication session and any cached API tokens.
4215
+ def authenticate(self, force: bool = False) -> None:
4216
+ """Authenticate or re-authenticate to the ZenML server.
4217
+
4218
+ Args:
4219
+ force: If True, force a re-authentication even if a valid API token
4220
+ is currently cached. This is useful when the current API token
4221
+ is known to be invalid or expired.
4222
+ """
4223
+ # This is called to trigger an authentication flow, either because
4224
+ # the current API token is expired or no longer valid, or because
4225
+ # a configuration change has happened or merely because an
4226
+ # authentication was never attempted before.
4227
+ #
4228
+ # 1. Drop the API token currently being used, if any.
4229
+ # 2. If force=True, clear the current API token from the credentials
4230
+ # store, if any, otherwise it will just be re-used on the next call.
4231
+ # 3. Get a new API token
4232
+
4233
+ # The authentication token could have expired or invalidated through
4234
+ # other means; refresh it and try again. This will clear any cached
4235
+ # token and trigger a new authentication flow.
4236
+ if self._api_token and not force:
4237
+ if self._api_token.expired:
4238
+ logger.info(
4239
+ "Authentication session expired; attempting to "
4240
+ "re-authenticate."
4241
+ )
4242
+ else:
4243
+ logger.info(
4244
+ "Authentication session was invalidated by the server; "
4245
+ "This can happen for example if the user's permissions "
4246
+ "have been revoked or if the server has been restarted "
4247
+ "and lost its session state. Attempting to "
4248
+ "re-authenticate."
4249
+ )
4250
+ else:
4251
+ if force:
4252
+ # Clear the current API token from the credentials store, if
4253
+ # any, to force a new authentication flow.
4254
+ get_credentials_store().clear_token(self.url)
4255
+ # Never authenticated since the client was created or the API token
4256
+ # was explicitly cleared.
4257
+ logger.debug(f"Authenticating to {self.url}...")
4069
4258
 
4070
- Raises:
4071
- AuthorizationException: If the API token can't be reset because
4072
- the store configuration does not contain username and password
4073
- or an API key to fetch a new token.
4074
- """
4075
- self._session = None
4076
4259
  self._api_token = None
4077
- # Clear the configured API token only if it's possible to fetch a new
4078
- # one from the server using other credentials (username/password or
4079
- # service account API key).
4080
- if (
4081
- self.config.username is not None
4082
- and self.config.password is not None
4083
- or self.config.api_key is not None
4084
- ):
4085
- self.config.api_token = None
4086
- elif self.config.api_token:
4087
- raise AuthorizationException(
4088
- "Unable to refresh invalid API token. This is probably "
4089
- "because you're connected to your ZenML server with device "
4090
- "authentication. Rerunning `zenml connect --url "
4091
- f"{self.config.url}` should solve this issue. "
4092
- "If you're seeing this error from an automated workload, "
4093
- "you should probably use a service account to start that "
4094
- "workload to prevent this error"
4095
- )
4260
+
4261
+ new_api_token = self.get_or_generate_api_token()
4262
+
4263
+ # Set or refresh the authentication token
4264
+ self.session.headers.update(
4265
+ {"Authorization": "Bearer " + new_api_token}
4266
+ )
4267
+ logger.debug(f"Authenticated to {self.url}")
4096
4268
 
4097
4269
  @staticmethod
4098
4270
  def _handle_response(response: requests.Response) -> Json:
@@ -4156,8 +4328,8 @@ class RestZenStore(BaseZenStore):
4156
4328
  The parsed response.
4157
4329
 
4158
4330
  Raises:
4159
- AuthorizationException: if the request fails due to an expired
4160
- authentication token.
4331
+ CredentialsNotValid: if the request fails due to invalid
4332
+ client credentials.
4161
4333
  """
4162
4334
  params = {k: str(v) for k, v in params.items()} if params else {}
4163
4335
 
@@ -4176,12 +4348,18 @@ class RestZenStore(BaseZenStore):
4176
4348
  **kwargs,
4177
4349
  )
4178
4350
  )
4179
- except AuthorizationException:
4180
- # The authentication token could have expired; refresh it and try
4181
- # again. This will clear any cached token and trigger a new
4182
- # authentication flow.
4183
- self.clear_session()
4184
- logger.info("Authentication token expired; refreshing...")
4351
+ except CredentialsNotValid:
4352
+ # NOTE: CredentialsNotValid is raised only when the server
4353
+ # explicitly indicates that the credentials are not valid and they
4354
+ # can be thrown away.
4355
+
4356
+ # We authenticate or re-authenticate here and then try the request
4357
+ # again, this time with a valid API token in the header.
4358
+ self.authenticate(
4359
+ # If the last request was authenticated with an API token,
4360
+ # we force a re-authentication to get a fresh token.
4361
+ force=self._api_token is not None
4362
+ )
4185
4363
 
4186
4364
  try:
4187
4365
  return self._handle_response(
@@ -4194,11 +4372,11 @@ class RestZenStore(BaseZenStore):
4194
4372
  **kwargs,
4195
4373
  )
4196
4374
  )
4197
- except AuthorizationException:
4198
- logger.info(
4199
- "Your authentication token has expired. Please re-authenticate."
4200
- )
4201
- raise
4375
+ except CredentialsNotValid as e:
4376
+ raise CredentialsNotValid(
4377
+ "The current credentials are no longer valid. Please log in "
4378
+ "again using 'zenml login'."
4379
+ ) from e
4202
4380
 
4203
4381
  def get(
4204
4382
  self,
@@ -155,20 +155,25 @@ class APIKeySchema(NamedSchema, table=True):
155
155
  )
156
156
 
157
157
  def to_internal_model(
158
- self, hydrate: bool = False
158
+ self,
159
+ include_metadata: bool = False,
160
+ include_resources: bool = False,
159
161
  ) -> APIKeyInternalResponse:
160
162
  """Convert a `APIKeySchema` to an `APIKeyInternalResponse`.
161
163
 
162
164
  The internal response model includes the hashed key values.
163
165
 
164
166
  Args:
165
- hydrate: bool to decide whether to return a hydrated version of the
166
- model.
167
+ include_metadata: Whether the metadata will be filled.
168
+ include_resources: Whether the resources will be filled.
167
169
 
168
170
  Returns:
169
171
  The created APIKeyInternalResponse.
170
172
  """
171
- model = self.to_model(include_metadata=hydrate)
173
+ model = self.to_model(
174
+ include_metadata=include_metadata,
175
+ include_resources=include_resources,
176
+ )
172
177
  model.get_body().key = self.key
173
178
 
174
179
  return APIKeyInternalResponse(