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
@@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any, List, Optional
18
18
  from uuid import UUID
19
19
 
20
20
  from pydantic import ValidationError
21
- from sqlalchemy import TEXT, Column
21
+ from sqlalchemy import TEXT, Column, UniqueConstraint
22
22
  from sqlmodel import Field, Relationship
23
23
 
24
24
  from zenml.config.source import Source
@@ -65,6 +65,12 @@ class ArtifactSchema(NamedSchema, table=True):
65
65
  """SQL Model for artifacts."""
66
66
 
67
67
  __tablename__ = "artifact"
68
+ __table_args__ = (
69
+ UniqueConstraint(
70
+ "name",
71
+ name="unique_artifact_name",
72
+ ),
73
+ )
68
74
 
69
75
  # Fields
70
76
  has_custom_name: bool
@@ -167,6 +173,13 @@ class ArtifactVersionSchema(BaseSchema, table=True):
167
173
  """SQL Model for artifact versions."""
168
174
 
169
175
  __tablename__ = "artifact_version"
176
+ __table_args__ = (
177
+ UniqueConstraint(
178
+ "version",
179
+ "artifact_id",
180
+ name="unique_version_for_artifact_id",
181
+ ),
182
+ )
170
183
 
171
184
  # Fields
172
185
  version: str
@@ -263,9 +276,15 @@ class ArtifactVersionSchema(BaseSchema, table=True):
263
276
  Args:
264
277
  artifact_version_request: The request model to convert.
265
278
 
279
+ Raises:
280
+ ValueError: If the request does not specify a version number.
281
+
266
282
  Returns:
267
283
  The converted schema.
268
284
  """
285
+ if not artifact_version_request.version:
286
+ raise ValueError("Missing version for artifact version request.")
287
+
269
288
  try:
270
289
  version_number = int(artifact_version_request.version)
271
290
  except ValueError:
@@ -341,6 +360,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
341
360
  updated=self.updated,
342
361
  tags=[t.tag.to_model() for t in self.tags],
343
362
  producer_pipeline_run_id=producer_pipeline_run_id,
363
+ artifact_store_id=self.artifact_store_id,
344
364
  )
345
365
 
346
366
  # Create the metadata of the model
@@ -348,7 +368,6 @@ class ArtifactVersionSchema(BaseSchema, table=True):
348
368
  if include_metadata:
349
369
  metadata = ArtifactVersionResponseMetadata(
350
370
  workspace=self.workspace.to_model(),
351
- artifact_store_id=self.artifact_store_id,
352
371
  producer_step_run_id=producer_step_run_id,
353
372
  visualizations=[v.to_model() for v in self.visualizations],
354
373
  run_metadata={m.key: m.to_model() for m in self.run_metadata},
@@ -43,7 +43,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True):
43
43
  # Foreign Keys
44
44
  artifact_version_id: UUID = build_foreign_key_field(
45
45
  source=__tablename__,
46
- target=ArtifactVersionSchema.__tablename__,
46
+ target=ArtifactVersionSchema.__tablename__, # type: ignore[has-type]
47
47
  source_column="artifact_version_id",
48
48
  target_column="id",
49
49
  ondelete="CASCADE",
@@ -23,9 +23,11 @@ from sqlmodel import Relationship
23
23
 
24
24
  from zenml.enums import StackComponentType
25
25
  from zenml.models import (
26
+ ComponentRequest,
26
27
  ComponentResponse,
27
28
  ComponentResponseBody,
28
29
  ComponentResponseMetadata,
30
+ ComponentResponseResources,
29
31
  ComponentUpdate,
30
32
  )
31
33
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
@@ -113,6 +115,39 @@ class StackComponentSchema(NamedSchema, table=True):
113
115
 
114
116
  connector_resource_id: Optional[str]
115
117
 
118
+ @classmethod
119
+ def from_request(
120
+ cls,
121
+ request: "ComponentRequest",
122
+ service_connector: Optional[ServiceConnectorSchema] = None,
123
+ ) -> "StackComponentSchema":
124
+ """Create a component schema from a request.
125
+
126
+ Args:
127
+ request: The request from which to create the component.
128
+ service_connector: Optional service connector to link to the
129
+ component.
130
+
131
+ Returns:
132
+ The component schema.
133
+ """
134
+ return cls(
135
+ name=request.name,
136
+ workspace_id=request.workspace,
137
+ user_id=request.user,
138
+ component_spec_path=request.component_spec_path,
139
+ type=request.type,
140
+ flavor=request.flavor,
141
+ configuration=base64.b64encode(
142
+ json.dumps(request.configuration).encode("utf-8")
143
+ ),
144
+ labels=base64.b64encode(
145
+ json.dumps(request.labels).encode("utf-8")
146
+ ),
147
+ connector=service_connector,
148
+ connector_resource_id=request.connector_resource_id,
149
+ )
150
+
116
151
  def update(
117
152
  self, component_update: "ComponentUpdate"
118
153
  ) -> "StackComponentSchema":
@@ -135,11 +170,6 @@ class StackComponentSchema(NamedSchema, table=True):
135
170
  self.labels = base64.b64encode(
136
171
  json.dumps(component_update.labels).encode("utf-8")
137
172
  )
138
- elif field == "type":
139
- component_type = component_update.type
140
-
141
- if component_type is not None:
142
- self.type = component_type
143
173
  else:
144
174
  setattr(self, field, value)
145
175
 
@@ -159,13 +189,15 @@ class StackComponentSchema(NamedSchema, table=True):
159
189
  include_resources: Whether the resources will be filled.
160
190
  **kwargs: Keyword arguments to allow schema specific logic
161
191
 
192
+ Raises:
193
+ RuntimeError: If the flavor for the component is missing in the DB.
162
194
 
163
195
  Returns:
164
196
  A `ComponentModel`
165
197
  """
166
198
  body = ComponentResponseBody(
167
199
  type=StackComponentType(self.type),
168
- flavor=self.flavor,
200
+ flavor_name=self.flavor,
169
201
  user=self.user.to_model() if self.user else None,
170
202
  created=self.created,
171
203
  updated=self.updated,
@@ -192,9 +224,20 @@ class StackComponentSchema(NamedSchema, table=True):
192
224
  if self.connector
193
225
  else None,
194
226
  )
227
+ resources = None
228
+ if include_resources:
229
+ if not self.flavor_schema:
230
+ raise RuntimeError(
231
+ f"Missing flavor {self.flavor} for component {self.name}."
232
+ )
233
+
234
+ resources = ComponentResponseResources(
235
+ flavor=self.flavor_schema.to_model()
236
+ )
195
237
  return ComponentResponse(
196
238
  id=self.id,
197
239
  name=self.name,
198
240
  body=body,
199
241
  metadata=metadata,
242
+ resources=resources,
200
243
  )
@@ -247,18 +247,23 @@ class OAuthDeviceSchema(BaseSchema, table=True):
247
247
  )
248
248
 
249
249
  def to_internal_model(
250
- self, hydrate: bool = False
250
+ self,
251
+ include_metadata: bool = False,
252
+ include_resources: bool = False,
251
253
  ) -> OAuthDeviceInternalResponse:
252
254
  """Convert a device schema to an internal device response model.
253
255
 
254
256
  Args:
255
- hydrate: bool to decide whether to return a hydrated version of the
256
- model.
257
+ include_metadata: Whether the metadata will be filled.
258
+ include_resources: Whether the resources will be filled.
257
259
 
258
260
  Returns:
259
261
  The converted internal device response model.
260
262
  """
261
- device_model = self.to_model(include_metadata=hydrate)
263
+ device_model = self.to_model(
264
+ include_metadata=include_metadata,
265
+ include_resources=include_resources,
266
+ )
262
267
  return OAuthDeviceInternalResponse(
263
268
  id=device_model.id,
264
269
  body=device_model.body,
@@ -127,6 +127,7 @@ class FlavorSchema(NamedSchema, table=True):
127
127
  user=self.user.to_model() if self.user else None,
128
128
  type=StackComponentType(self.type),
129
129
  integration=self.integration,
130
+ source=self.source,
130
131
  logo_url=self.logo_url,
131
132
  created=self.created,
132
133
  updated=self.updated,
@@ -141,7 +142,6 @@ class FlavorSchema(NamedSchema, table=True):
141
142
  connector_type=self.connector_type,
142
143
  connector_resource_type=self.connector_resource_type,
143
144
  connector_resource_id_attr=self.connector_resource_id_attr,
144
- source=self.source,
145
145
  docs_url=self.docs_url,
146
146
  sdk_docs_url=self.sdk_docs_url,
147
147
  is_custom=self.is_custom,
@@ -498,7 +498,7 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
498
498
  )
499
499
  artifact_version_id: UUID = build_foreign_key_field(
500
500
  source=__tablename__,
501
- target=ArtifactVersionSchema.__tablename__,
501
+ target=ArtifactVersionSchema.__tablename__, # type: ignore[has-type]
502
502
  source_column="artifact_version_id",
503
503
  target_column="id",
504
504
  ondelete="CASCADE",
@@ -88,7 +88,7 @@ class ServiceSchema(NamedSchema, table=True):
88
88
  )
89
89
  model_version_id: Optional[UUID] = build_foreign_key_field(
90
90
  source=__tablename__,
91
- target=ModelVersionSchema.__tablename__,
91
+ target=ModelVersionSchema.__tablename__, # type: ignore[has-type]
92
92
  source_column="model_version_id",
93
93
  target_column="id",
94
94
  ondelete="SET NULL",
@@ -97,7 +97,7 @@ class StepRunSchema(NamedSchema, table=True):
97
97
  )
98
98
  pipeline_run_id: UUID = build_foreign_key_field(
99
99
  source=__tablename__,
100
- target=PipelineRunSchema.__tablename__,
100
+ target=PipelineRunSchema.__tablename__, # type: ignore[has-type]
101
101
  source_column="pipeline_run_id",
102
102
  target_column="id",
103
103
  ondelete="CASCADE",
@@ -77,7 +77,7 @@ class TriggerSchema(NamedSchema, table=True):
77
77
 
78
78
  event_source_id: Optional[UUID] = build_foreign_key_field(
79
79
  source=__tablename__,
80
- target=EventSourceSchema.__tablename__,
80
+ target=EventSourceSchema.__tablename__, # type: ignore[has-type]
81
81
  source_column="event_source_id",
82
82
  target_column="id",
83
83
  # This won't happen because the SQL zen store prevents an event source