zenml-nightly 0.68.0.dev20241028__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.dev20241028.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/METADATA +18 -12
  121. {zenml_nightly-0.68.0.dev20241028.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.dev20241028.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/LICENSE +0 -0
  124. {zenml_nightly-0.68.0.dev20241028.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/WHEEL +0 -0
  125. {zenml_nightly-0.68.0.dev20241028.dist-info → zenml_nightly-0.68.1.dev20241101.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,496 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro client."""
15
+
16
+ from typing import (
17
+ TYPE_CHECKING,
18
+ Any,
19
+ Dict,
20
+ List,
21
+ Optional,
22
+ Type,
23
+ TypeVar,
24
+ Union,
25
+ )
26
+ from uuid import UUID
27
+
28
+ import requests
29
+ from requests.adapters import HTTPAdapter, Retry
30
+
31
+ from zenml.analytics import source_context
32
+ from zenml.exceptions import AuthorizationException
33
+ from zenml.logger import get_logger
34
+ from zenml.login.credentials import APIToken
35
+ from zenml.login.credentials_store import get_credentials_store
36
+ from zenml.login.pro.constants import ZENML_PRO_API_URL
37
+ from zenml.login.pro.models import BaseRestAPIModel
38
+ from zenml.utils.singleton import SingletonMetaClass
39
+ from zenml.zen_server.exceptions import exception_from_response
40
+
41
+ logger = get_logger(__name__)
42
+
43
+ if TYPE_CHECKING:
44
+ from zenml.login.pro.organization.client import OrganizationClient
45
+ from zenml.login.pro.tenant.client import TenantClient
46
+
47
+ # type alias for possible json payloads (the Anys are recursive Json instances)
48
+ Json = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
49
+
50
+
51
+ AnyResponse = TypeVar("AnyResponse", bound=BaseRestAPIModel)
52
+
53
+
54
+ class ZenMLProClient(metaclass=SingletonMetaClass):
55
+ """ZenML Pro client."""
56
+
57
+ _url: str
58
+ _api_token: APIToken
59
+ _session: Optional[requests.Session] = None
60
+ _tenant: Optional["TenantClient"] = None
61
+ _organization: Optional["OrganizationClient"] = None
62
+
63
+ def __init__(
64
+ self, url: Optional[str] = None, api_token: Optional[APIToken] = None
65
+ ) -> None:
66
+ """Initialize the ZenML Pro client.
67
+
68
+ Args:
69
+ url: The URL of the ZenML Pro API server. If not provided, the
70
+ default ZenML Pro API server URL is used.
71
+ api_token: The API token to use for authentication. If not provided,
72
+ the token is fetched from the credentials store.
73
+
74
+ Raises:
75
+ AuthorizationException: If no API token is provided and no token
76
+ is found in the credentials store.
77
+ """
78
+ self._url = url or ZENML_PRO_API_URL
79
+ if api_token is None:
80
+ logger.debug(
81
+ "No ZenML Pro API token provided. Fetching from credentials "
82
+ "store."
83
+ )
84
+ api_token = get_credentials_store().get_token(
85
+ server_url=self._url, allow_expired=True
86
+ )
87
+ if api_token is None:
88
+ raise AuthorizationException(
89
+ "No ZenML Pro API token found. Please run 'zenml login' to "
90
+ "login to ZenML Pro."
91
+ )
92
+
93
+ self._api_token = api_token
94
+
95
+ @property
96
+ def tenant(self) -> "TenantClient":
97
+ """Get the tenant client.
98
+
99
+ Returns:
100
+ The tenant client.
101
+ """
102
+ if self._tenant is None:
103
+ from zenml.login.pro.tenant.client import TenantClient
104
+
105
+ self._tenant = TenantClient(client=self)
106
+ return self._tenant
107
+
108
+ @property
109
+ def organization(self) -> "OrganizationClient":
110
+ """Get the organization client.
111
+
112
+ Returns:
113
+ The organization client.
114
+ """
115
+ if self._organization is None:
116
+ from zenml.login.pro.organization.client import OrganizationClient
117
+
118
+ self._organization = OrganizationClient(client=self)
119
+ return self._organization
120
+
121
+ @property
122
+ def api_token(self) -> str:
123
+ """Get the API token.
124
+
125
+ Returns:
126
+ The API token.
127
+ """
128
+ return self._api_token.access_token
129
+
130
+ def raise_on_expired_api_token(self) -> None:
131
+ """Raise an exception if the API token has expired.
132
+
133
+ Raises:
134
+ AuthorizationException: If the API token has expired.
135
+ """
136
+ if self._api_token and self._api_token.expired:
137
+ raise AuthorizationException(
138
+ "Your ZenML Pro authentication has expired. Please run "
139
+ "'zenml login' to login again."
140
+ )
141
+
142
+ @property
143
+ def session(self) -> requests.Session:
144
+ """Authenticate to the ZenML Pro API server.
145
+
146
+ Returns:
147
+ A requests session with the authentication token.
148
+ """
149
+ # Check if the API token has expired before every call to the server.
150
+ # This prevents unwanted authorization errors from being raised during
151
+ # the call itself.
152
+ self.raise_on_expired_api_token()
153
+ if self._session is None:
154
+ self._session = requests.Session()
155
+ retries = Retry(backoff_factor=0.1, connect=5)
156
+ self._session.mount("https://", HTTPAdapter(max_retries=retries))
157
+ self._session.mount("http://", HTTPAdapter(max_retries=retries))
158
+ self._session.headers.update(
159
+ {"Authorization": "Bearer " + self.api_token}
160
+ )
161
+ logger.debug("Authenticated to ZenML Pro server.")
162
+ return self._session
163
+
164
+ @staticmethod
165
+ def _handle_response(response: requests.Response) -> Json:
166
+ """Handle API response, translating http status codes to Exception.
167
+
168
+ Args:
169
+ response: The response to handle.
170
+
171
+ Returns:
172
+ The parsed response.
173
+
174
+ Raises:
175
+ ValueError: if the response is not in the right format.
176
+ RuntimeError: if an error response is received from the server
177
+ and a more specific exception cannot be determined.
178
+ exc: the exception converted from an error response, if one
179
+ is returned from the server.
180
+ """
181
+ if 200 <= response.status_code < 300:
182
+ try:
183
+ payload: Json = response.json()
184
+ return payload
185
+ except requests.exceptions.JSONDecodeError:
186
+ raise ValueError(
187
+ "Bad response from API. Expected json, got\n"
188
+ f"{response.text}"
189
+ )
190
+ elif response.status_code >= 400:
191
+ exc = exception_from_response(response)
192
+ if exc is not None:
193
+ raise exc
194
+ else:
195
+ raise RuntimeError(
196
+ f"{response.status_code} HTTP Error received from server: "
197
+ f"{response.text}"
198
+ )
199
+ else:
200
+ raise RuntimeError(
201
+ "Error retrieving from API. Got response "
202
+ f"{response.status_code} with body:\n{response.text}"
203
+ )
204
+
205
+ def _request(
206
+ self,
207
+ method: str,
208
+ url: str,
209
+ params: Optional[Dict[str, Any]] = None,
210
+ **kwargs: Any,
211
+ ) -> Json:
212
+ """Make a request to the REST API.
213
+
214
+ Args:
215
+ method: The HTTP method to use.
216
+ url: The URL to request.
217
+ params: The query parameters to pass to the endpoint.
218
+ kwargs: Additional keyword arguments to pass to the request.
219
+
220
+ Returns:
221
+ The parsed response.
222
+
223
+ Raises:
224
+ AuthorizationException: if the request fails due to an expired
225
+ authentication token.
226
+ """
227
+ params = {k: str(v) for k, v in params.items()} if params else {}
228
+
229
+ self.session.headers.update(
230
+ {source_context.name: source_context.get().value}
231
+ )
232
+
233
+ try:
234
+ return self._handle_response(
235
+ self.session.request(
236
+ method,
237
+ url,
238
+ params=params,
239
+ **kwargs,
240
+ )
241
+ )
242
+ except AuthorizationException:
243
+ # Check if this is caused by an expired API token.
244
+ self.raise_on_expired_api_token()
245
+
246
+ # If not, raise the exception.
247
+ raise
248
+
249
+ def get(
250
+ self,
251
+ path: str,
252
+ params: Optional[Dict[str, Any]] = None,
253
+ **kwargs: Any,
254
+ ) -> Json:
255
+ """Make a GET request to the given endpoint path.
256
+
257
+ Args:
258
+ path: The path to the endpoint.
259
+ params: The query parameters to pass to the endpoint.
260
+ kwargs: Additional keyword arguments to pass to the request.
261
+
262
+ Returns:
263
+ The response body.
264
+ """
265
+ logger.debug(f"Sending GET request to {path}...")
266
+ return self._request(
267
+ "GET",
268
+ self._url + path,
269
+ params=params,
270
+ **kwargs,
271
+ )
272
+
273
+ def delete(
274
+ self,
275
+ path: str,
276
+ params: Optional[Dict[str, Any]] = None,
277
+ **kwargs: Any,
278
+ ) -> Json:
279
+ """Make a DELETE request to the given endpoint path.
280
+
281
+ Args:
282
+ path: The path to the endpoint.
283
+ params: The query parameters to pass to the endpoint.
284
+ kwargs: Additional keyword arguments to pass to the request.
285
+
286
+ Returns:
287
+ The response body.
288
+ """
289
+ logger.debug(f"Sending DELETE request to {path}...")
290
+ return self._request(
291
+ "DELETE",
292
+ self._url + path,
293
+ params=params,
294
+ **kwargs,
295
+ )
296
+
297
+ def post(
298
+ self,
299
+ path: str,
300
+ body: BaseRestAPIModel,
301
+ params: Optional[Dict[str, Any]] = None,
302
+ **kwargs: Any,
303
+ ) -> Json:
304
+ """Make a POST request to the given endpoint path.
305
+
306
+ Args:
307
+ path: The path to the endpoint.
308
+ body: The body to send.
309
+ params: The query parameters to pass to the endpoint.
310
+ kwargs: Additional keyword arguments to pass to the request.
311
+
312
+ Returns:
313
+ The response body.
314
+ """
315
+ logger.debug(f"Sending POST request to {path}...")
316
+ return self._request(
317
+ "POST",
318
+ self._url + path,
319
+ json=body.model_dump(mode="json"),
320
+ params=params,
321
+ **kwargs,
322
+ )
323
+
324
+ def put(
325
+ self,
326
+ path: str,
327
+ body: Optional[BaseRestAPIModel] = None,
328
+ params: Optional[Dict[str, Any]] = None,
329
+ **kwargs: Any,
330
+ ) -> Json:
331
+ """Make a PUT request to the given endpoint path.
332
+
333
+ Args:
334
+ path: The path to the endpoint.
335
+ body: The body to send.
336
+ params: The query parameters to pass to the endpoint.
337
+ kwargs: Additional keyword arguments to pass to the request.
338
+
339
+ Returns:
340
+ The response body.
341
+ """
342
+ logger.debug(f"Sending PUT request to {path}...")
343
+ json = (
344
+ body.model_dump(mode="json", exclude_unset=True) if body else None
345
+ )
346
+ return self._request(
347
+ "PUT",
348
+ self._url + path,
349
+ json=json,
350
+ params=params,
351
+ **kwargs,
352
+ )
353
+
354
+ def patch(
355
+ self,
356
+ path: str,
357
+ body: Optional[BaseRestAPIModel] = None,
358
+ params: Optional[Dict[str, Any]] = None,
359
+ **kwargs: Any,
360
+ ) -> Json:
361
+ """Make a PATCH request to the given endpoint path.
362
+
363
+ Args:
364
+ path: The path to the endpoint.
365
+ body: The body to send.
366
+ params: The query parameters to pass to the endpoint.
367
+ kwargs: Additional keyword arguments to pass to the request.
368
+
369
+ Returns:
370
+ The response body.
371
+ """
372
+ logger.debug(f"Sending PATCH request to {path}...")
373
+ json = (
374
+ body.model_dump(mode="json", exclude_unset=True) if body else None
375
+ )
376
+ return self._request(
377
+ "PATCH",
378
+ self._url + path,
379
+ json=json,
380
+ params=params,
381
+ **kwargs,
382
+ )
383
+
384
+ def _create_resource(
385
+ self,
386
+ resource: BaseRestAPIModel,
387
+ response_model: Type[AnyResponse],
388
+ route: str,
389
+ params: Optional[Dict[str, Any]] = None,
390
+ ) -> AnyResponse:
391
+ """Create a new resource.
392
+
393
+ Args:
394
+ resource: The resource to create.
395
+ route: The resource REST API route to use.
396
+ response_model: Optional model to use to deserialize the response
397
+ body. If not provided, the resource class itself will be used.
398
+ params: Optional query parameters to pass to the endpoint.
399
+
400
+ Returns:
401
+ The created resource.
402
+ """
403
+ response_body = self.post(f"{route}", body=resource, params=params)
404
+
405
+ return response_model.model_validate(response_body)
406
+
407
+ def _get_resource(
408
+ self,
409
+ resource_id: Union[str, int, UUID],
410
+ route: str,
411
+ response_model: Type[AnyResponse],
412
+ **params: Any,
413
+ ) -> AnyResponse:
414
+ """Retrieve a single resource.
415
+
416
+ Args:
417
+ resource_id: The ID of the resource to retrieve.
418
+ route: The resource REST API route to use.
419
+ response_model: Model to use to serialize the response body.
420
+ params: Optional query parameters to pass to the endpoint.
421
+
422
+ Returns:
423
+ The retrieved resource.
424
+ """
425
+ # leave out filter params that are not supplied
426
+ params = dict(filter(lambda x: x[1] is not None, params.items()))
427
+ body = self.get(f"{route}/{str(resource_id)}", params=params)
428
+ return response_model.model_validate(body)
429
+
430
+ def _list_resources(
431
+ self,
432
+ route: str,
433
+ response_model: Type[AnyResponse],
434
+ **params: Any,
435
+ ) -> List[AnyResponse]:
436
+ """Retrieve a list of resources filtered by some criteria.
437
+
438
+ Args:
439
+ route: The resource REST API route to use.
440
+ response_model: Model to use to serialize the response body.
441
+ params: Filter parameters to use in the query.
442
+
443
+ Returns:
444
+ List of retrieved resources matching the filter criteria.
445
+
446
+ Raises:
447
+ ValueError: If the value returned by the server is not a list.
448
+ """
449
+ # leave out filter params that are not supplied
450
+ params = dict(filter(lambda x: x[1] is not None, params.items()))
451
+ body = self.get(f"{route}", params=params)
452
+ if not isinstance(body, list):
453
+ raise ValueError(
454
+ f"Bad API Response. Expected list, got {type(body)}"
455
+ )
456
+ return [response_model.model_validate(entry) for entry in body]
457
+
458
+ def _update_resource(
459
+ self,
460
+ resource_id: Union[str, int, UUID],
461
+ resource_update: BaseRestAPIModel,
462
+ response_model: Type[AnyResponse],
463
+ route: str,
464
+ **params: Any,
465
+ ) -> AnyResponse:
466
+ """Update an existing resource.
467
+
468
+ Args:
469
+ resource_id: The id of the resource to update.
470
+ resource_update: The resource update.
471
+ response_model: Optional model to use to deserialize the response
472
+ body. If not provided, the resource class itself will be used.
473
+ route: The resource REST API route to use.
474
+ params: Optional query parameters to pass to the endpoint.
475
+
476
+ Returns:
477
+ The updated resource.
478
+ """
479
+ # leave out filter params that are not supplied
480
+ params = dict(filter(lambda x: x[1] is not None, params.items()))
481
+ response_body = self.put(
482
+ f"{route}/{str(resource_id)}", body=resource_update, params=params
483
+ )
484
+
485
+ return response_model.model_validate(response_body)
486
+
487
+ def _delete_resource(
488
+ self, resource_id: Union[str, UUID], route: str
489
+ ) -> None:
490
+ """Delete a resource.
491
+
492
+ Args:
493
+ resource_id: The ID of the resource to delete.
494
+ route: The resource REST API route to use.
495
+ """
496
+ self.delete(f"{route}/{str(resource_id)}")
@@ -0,0 +1,34 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro login constants."""
15
+
16
+ import os
17
+
18
+ ENV_ZENML_PRO_API_URL = "ZENML_PRO_API_URL"
19
+ DEFAULT_ZENML_PRO_API_URL = "https://cloudapi.zenml.io"
20
+
21
+ ZENML_PRO_API_URL = os.getenv(
22
+ ENV_ZENML_PRO_API_URL, default=DEFAULT_ZENML_PRO_API_URL
23
+ )
24
+
25
+ ENV_ZENML_PRO_URL = "ZENML_PRO_URL"
26
+ DEFAULT_ZENML_PRO_URL = "https://cloud.zenml.io"
27
+
28
+ ZENML_PRO_URL = os.getenv(ENV_ZENML_PRO_URL, default=DEFAULT_ZENML_PRO_URL)
29
+
30
+ ENV_ZENML_PRO_SERVER_SUBDOMAIN = "ZENML_PRO_SERVER_SUBDOMAIN"
31
+ DEFAULT_ZENML_PRO_SERVER_SUBDOMAIN = "cloudinfra.zenml.io"
32
+ ZENML_PRO_SERVER_SUBDOMAIN = os.getenv(
33
+ ENV_ZENML_PRO_SERVER_SUBDOMAIN, default=DEFAULT_ZENML_PRO_SERVER_SUBDOMAIN
34
+ )
@@ -0,0 +1,25 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro base models."""
15
+
16
+ from pydantic import BaseModel, ConfigDict
17
+
18
+
19
+ class BaseRestAPIModel(BaseModel):
20
+ """Base class for all REST API models."""
21
+
22
+ model_config = ConfigDict(
23
+ # Allow extra attributes to allow compatibility with future versions
24
+ extra="allow",
25
+ )
@@ -0,0 +1,14 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro organization client."""
@@ -0,0 +1,79 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro organization client."""
15
+
16
+ from typing import List, Union
17
+ from uuid import UUID
18
+
19
+ from zenml.logger import get_logger
20
+ from zenml.login.pro.client import ZenMLProClient
21
+ from zenml.login.pro.organization.models import OrganizationRead
22
+
23
+ logger = get_logger(__name__)
24
+
25
+ ORGANIZATIONS_ROUTE = "/organizations"
26
+
27
+
28
+ class OrganizationClient:
29
+ """Organization management client."""
30
+
31
+ def __init__(
32
+ self,
33
+ client: ZenMLProClient,
34
+ ):
35
+ """Initialize the organization client.
36
+
37
+ Args:
38
+ client: ZenML Pro client.
39
+ """
40
+ self.client = client
41
+
42
+ def get(
43
+ self,
44
+ id_or_name: Union[UUID, str],
45
+ ) -> OrganizationRead:
46
+ """Get an organization by id or name.
47
+
48
+ Args:
49
+ id_or_name: Id or name of the organization to retrieve.
50
+
51
+ Returns:
52
+ An organization.
53
+ """
54
+ return self.client._get_resource(
55
+ resource_id=id_or_name,
56
+ route=ORGANIZATIONS_ROUTE,
57
+ response_model=OrganizationRead,
58
+ )
59
+
60
+ async def list(
61
+ self,
62
+ offset: int = 0,
63
+ limit: int = 20,
64
+ ) -> List[OrganizationRead]:
65
+ """List organizations.
66
+
67
+ Args:
68
+ offset: Query offset.
69
+ limit: Query limit.
70
+
71
+ Returns:
72
+ List of organizations.
73
+ """
74
+ return self.client._list_resources(
75
+ route=ORGANIZATIONS_ROUTE,
76
+ response_model=OrganizationRead,
77
+ offset=offset,
78
+ limit=limit,
79
+ )
@@ -0,0 +1,32 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro organization models."""
15
+
16
+ from datetime import datetime
17
+ from typing import Optional
18
+ from uuid import UUID
19
+
20
+ from zenml.login.pro.models import BaseRestAPIModel
21
+
22
+
23
+ class OrganizationRead(BaseRestAPIModel):
24
+ """Model for viewing organizations."""
25
+
26
+ id: UUID
27
+
28
+ name: str
29
+ description: Optional[str] = None
30
+
31
+ created: datetime
32
+ updated: datetime
@@ -0,0 +1,14 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """ZenML Pro tenant client."""