zenml-nightly 0.70.0.dev20241122__py3-none-any.whl → 0.70.0.dev20241201__py3-none-any.whl

zenml/zen_server/deploy/helm/templates/_environment.tpl

@@ -3,6 +3,123 @@ Helpers for environment variables configured in ZenML deployments and secrets st
 */}}
 
 
+{{/*
+ZenML store configuration options (non-secret values).
+
+This template constructs a dictionary that is similar to the python values that
+can be configured in the zenml.zen_store.sql_zen_store.SqlZenStoreConfiguration
+class. Only non-secret values are included in this dictionary.
+
+The dictionary is then converted into deployment environment variables by other
+templates and inserted where it is needed.
+
+The input is taken from a .ZenML dict that is passed to the template and
+contains the values configured in the values.yaml file for the ZenML server.
+
+Args:
+  .ZenML: A dictionary with the ZenML configuration values configured for the
+  ZenML server.
+Returns:
+  A dictionary with the non-secret values configured for the ZenML store.
+*/}}
+{{- define "zenml.storeConfigurationAttrs" -}}
+{{- if .ZenML.database.url }}
+type: sql
+ssl_verify_server_cert: {{ .ZenML.database.sslVerifyServerCert | default "false" | quote }}
+{{- if .ZenML.database.backupStrategy }}
+backup_strategy: {{ .ZenML.database.backupStrategy | quote }}
+{{- if eq .ZenML.database.backupStrategy "database" }}
+backup_database: {{ .ZenML.database.backupDatabase | quote }}
+{{- else if eq .ZenML.database.backupStrategy "dump-file" }}
+backup_directory: "/backups"
+{{- end }}
+{{- end }}
+{{- if .ZenML.database.poolSize }}
+pool_size: {{ .ZenML.database.poolSize | quote }}
+{{- end }}
+{{- if .ZenML.database.maxOverflow }}
+max_overflow: {{ .ZenML.database.maxOverflow | quote }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+ZenML store configuration options (secret values).
+
+This template constructs a dictionary that is similar to the python values that
+can be configured in the zenml.zen_store.sql_zen_store.SqlZenStoreConfiguration
+class. Only secret values are included in this dictionary.
+
+The dictionary is then converted into deployment environment variables by other
+templates and inserted where it is needed.
+
+The input is taken from a .ZenML dict that is passed to the template and
+contains the values configured in the values.yaml file for the ZenML server.
+
+Args:
+  .ZenML: A dictionary with the ZenML configuration values configured for the
+  ZenML server.
+Returns:
+  A dictionary with the secret values configured for the ZenML store.
+*/}}
+{{- define "zenml.storeSecretConfigurationAttrs" -}}
+{{- if .ZenML.database.url }}
+url: {{ .ZenML.database.url | quote }}
+{{- if .ZenML.database.sslCa }}
+ssl_ca: {{ .Files.Get .ZenML.database.sslCa }}
+{{- end }}
+{{- if .ZenML.database.sslCert }}
+ssl_cert: {{ .Files.Get .ZenML.database.sslCert }}
+{{- end }}
+{{- if .ZenML.database.sslKey }}
+ssl_key: {{ .Files.Get .ZenML.database.sslKey }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+Store configuration environment variables (non-secret values).
+
+Passes the .Values.zenml dict as input to the `zenml.storeConfigurationAttrs`
+template and converts the output into a dictionary of environment variables that
+need to be configured for the store.
+
+Args:
+  .Values: The values.yaml file for the ZenML deployment.
+Returns:
+  A dictionary with the non-secret environment variables that are configured for
+  the store (i.e. keys starting with `ZENML_STORE_`).
+*/}}
+{{- define "zenml.storeEnvVariables" -}}
+{{ $zenml := dict "ZenML" .Values.zenml }}
+{{- range $k, $v := include "zenml.storeConfigurationAttrs" $zenml | fromYaml }}
+ZENML_STORE_{{ $k | upper }}: {{ $v | quote }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+Store configuration environment variables (secret values).
+
+Passes the .Values.zenml dict as input to the `zenml.storeSecretConfigurationAttrs`
+template and converts the output into a dictionary of environment variables that
+need to be configured for the store.
+
+Args:
+  .Values: The values.yaml file for the ZenML deployment.
+Returns:
+  A dictionary with the secret environment variables that are configured for
+  the store (i.e. keys starting with `ZENML_STORE_`).
+*/}}
+{{- define "zenml.storeSecretEnvVariables" -}}
+{{ $zenml := dict "ZenML" .Values.zenml }}
+{{- range $k, $v := include "zenml.storeSecretConfigurationAttrs" $zenml | fromYaml }}
+ZENML_STORE_{{ $k | upper }}: {{ $v | quote }}
+{{- end }}
+{{- end }}
+
 {{/*
 ZenML server configuration options (non-secret values).
 

zenml/zen_server/deploy/helm/templates/server-db-job.yaml

@@ -79,20 +79,9 @@ spec:
             {{- end }}
             - name: ZENML_DEFAULT_PROJECT_NAME
               value: {{ .Values.zenml.defaultProject | quote }}
-            - name: ZENML_STORE_TYPE
-              value: sql
-            - name: ZENML_STORE_SSL_VERIFY_SERVER_CERT
-              value: {{ .Values.zenml.database.sslVerifyServerCert | default "false" | quote }}
-            {{- if .Values.zenml.database.backupStrategy }}
-            - name: ZENML_STORE_BACKUP_STRATEGY
-              value: {{ .Values.zenml.database.backupStrategy | quote }}
-            {{- if eq .Values.zenml.database.backupStrategy "database" }}
-            - name: ZENML_STORE_BACKUP_DATABASE
-              value: {{ .Values.zenml.database.backupDatabase | quote }}
-            {{- else if eq .Values.zenml.database.backupStrategy "dump-file" }}
-            - name: ZENML_STORE_BACKUP_DIRECTORY
-              value: /backups
-            {{- end }}
+            {{- range $k, $v := include "zenml.storeEnvVariables" . | fromYaml }}
+            - name: {{ $k }}
+              value: {{ $v | quote }}
             {{- end }}
             {{- range $k, $v := include "zenml.serverEnvVariables" . | fromYaml }}
             - name: {{ $k }}

zenml/zen_server/deploy/helm/templates/server-deployment.yaml

@@ -63,14 +63,14 @@ spec:
               value: "True"
             {{- end }}
             {{- if .Values.zenml.database.url }}
-            - name: ZENML_STORE_TYPE
-              value: sql
             - name: DISABLE_DATABASE_MIGRATION
               value: "True"
-            - name: ZENML_STORE_SSL_VERIFY_SERVER_CERT
-              value: {{ .Values.zenml.database.sslVerifyServerCert | default "false" | quote }}
             {{- end }}
 
+            {{- range $k, $v := include "zenml.storeEnvVariables" . | fromYaml }}
+            - name: {{ $k }}
+              value: {{ $v | quote }}
+            {{- end }}
             {{- range $k, $v := include "zenml.serverEnvVariables" . | fromYaml }}
             - name: {{ $k }}
               value: {{ $v | quote }}
@@ -104,6 +104,18 @@ spec:
             httpGet:
               path: /health
               port: http
+          lifecycle:
+            preStop:
+              exec:
+                # Give the process 15 more seconds before the SIGTERM signal is
+                # sent. This allows the endpoint removal to reach the ingress
+                # controller in time and for traffic to be routed away from the
+                # pod before it is shut down. This eliminates the number of 502
+                # errors returned to the user.
+                #
+                # See https://learnk8s.io/graceful-shutdown for more information.
+                #
+                command: ["sleep", "15"]
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
       {{- with .Values.nodeSelector }}

zenml/zen_server/deploy/helm/templates/server-secret.yaml

@@ -10,23 +10,8 @@ data:
   {{- else }}
   ZENML_SERVER_JWT_SECRET_KEY: {{ $prevServerSecret.data.ZENML_SERVER_JWT_SECRET_KEY | default (randAlphaNum 32 | b64enc | quote) }}
   {{- end }}
-  {{- if .Values.zenml.database.url }}
-  ZENML_STORE_URL: {{ .Values.zenml.database.url | b64enc | quote }}
-  {{- if .Values.zenml.database.sslCa }}
-  ZENML_STORE_SSL_CA: {{ .Files.Get .Values.zenml.database.sslCa | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslCert }}
-  ZENML_STORE_SSL_CERT: {{ .Files.Get .Values.zenml.database.sslCert | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslKey }}
-  ZENML_STORE_SSL_KEY: {{ .Files.Get .Values.zenml.database.sslKey | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.poolSize }}
-  ZENML_STORE_POOL_SIZE: {{ .Values.zenml.database.poolSize | b64enc | quote }}
-  {{- end }}
-  {{- if .Values.zenml.database.maxOverflow }}
-  ZENML_STORE_MAX_OVERFLOW: {{ .Values.zenml.database.maxOverflow | b64enc | quote }}
-  {{- end }}
+  {{- range $k, $v := include "zenml.storeSecretEnvVariables" . | fromYaml}}
+  {{ $k }}: {{ $v | b64enc | quote }}
   {{- end }}
   {{- range $k, $v := include "zenml.secretsStoreSecretEnvVariables" . | fromYaml}}
   {{ $k }}: {{ $v | b64enc | quote }}

zenml/zen_server/rbac/endpoint_utils.py

@@ -189,7 +189,7 @@ def verify_permissions_and_list_entities(
 def verify_permissions_and_update_entity(
     id: UUIDOrStr,
     update_model: AnyUpdate,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
 ) -> AnyResponse:
     """Verify permissions and update an entity.
@@ -203,7 +203,8 @@ def verify_permissions_and_update_entity(
     Returns:
         A model of the updated entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.UPDATE)
     updated_model = update_method(model.id, update_model)
     return dehydrate_response_model(updated_model)
@@ -211,7 +212,7 @@ def verify_permissions_and_update_entity(
 
 def verify_permissions_and_delete_entity(
     id: UUIDOrStr,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     delete_method: Callable[[UUIDOrStr], None],
 ) -> AnyResponse:
     """Verify permissions and delete an entity.
@@ -224,7 +225,8 @@ def verify_permissions_and_delete_entity(
     Returns:
         The deleted entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.DELETE)
     delete_method(model.id)
 

zenml/zen_server/rbac/models.py

@@ -59,7 +59,6 @@ class ResourceType(StrEnum):
     PIPELINE_DEPLOYMENT = "pipeline_deployment"
     PIPELINE_BUILD = "pipeline_build"
     RUN_TEMPLATE = "run_template"
-    USER = "user"
     SERVICE = "service"
     RUN_METADATA = "run_metadata"
     SECRET = "secret"
@@ -70,7 +69,9 @@ class ResourceType(StrEnum):
     TAG = "tag"
     TRIGGER = "trigger"
     TRIGGER_EXECUTION = "trigger_execution"
-    WORKSPACE = "workspace"
+    # Deactivated for now
+    # USER = "user"
+    # WORKSPACE = "workspace"
 
 
 class Resource(BaseModel):

zenml/zen_server/rbac/utils.py

@@ -413,8 +413,6 @@ def get_resource_type_for_model(
         TagResponse,
         TriggerExecutionResponse,
         TriggerResponse,
-        UserResponse,
-        WorkspaceResponse,
     )
 
     mapping: Dict[
@@ -434,8 +432,8 @@ def get_resource_type_for_model(
         ModelVersionResponse: ResourceType.MODEL_VERSION,
         ArtifactResponse: ResourceType.ARTIFACT,
         ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
-        WorkspaceResponse: ResourceType.WORKSPACE,
-        UserResponse: ResourceType.USER,
+        # WorkspaceResponse: ResourceType.WORKSPACE,
+        # UserResponse: ResourceType.USER,
         PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,
@@ -570,7 +568,6 @@ def get_schema_for_resource_type(
         TriggerExecutionSchema,
         TriggerSchema,
         UserSchema,
-        WorkspaceSchema,
     )
 
     mapping: Dict[ResourceType, Type["BaseSchema"]] = {
@@ -588,13 +585,13 @@ def get_schema_for_resource_type(
         ResourceType.SERVICE: ServiceSchema,
         ResourceType.TAG: TagSchema,
         ResourceType.SERVICE_ACCOUNT: UserSchema,
-        ResourceType.WORKSPACE: WorkspaceSchema,
+        # ResourceType.WORKSPACE: WorkspaceSchema,
         ResourceType.PIPELINE_RUN: PipelineRunSchema,
         ResourceType.PIPELINE_DEPLOYMENT: PipelineDeploymentSchema,
         ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
         ResourceType.RUN_TEMPLATE: RunTemplateSchema,
         ResourceType.RUN_METADATA: RunMetadataSchema,
-        ResourceType.USER: UserSchema,
+        # ResourceType.USER: UserSchema,
         ResourceType.ACTION: ActionSchema,
         ResourceType.EVENT_SOURCE: EventSourceSchema,
         ResourceType.TRIGGER: TriggerSchema,

zenml/zen_server/routers/server_endpoints.py

@@ -22,6 +22,7 @@ from zenml.constants import (
     ACTIVATE,
     API,
     INFO,
+    LOAD_INFO,
     ONBOARDING_STATE,
     SERVER_SETTINGS,
     VERSION_1,
@@ -30,6 +31,7 @@ from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     ServerActivationRequest,
+    ServerLoadInfo,
     ServerModel,
     ServerSettingsResponse,
     ServerSettingsUpdate,
@@ -71,6 +73,51 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    LOAD_INFO,
+    response_model=ServerLoadInfo,
+)
+@handle_exceptions
+def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
+    """Get information about the server load.
+
+    Returns:
+        Information about the server load.
+    """
+    import threading
+
+    # Get the current number of threads
+    num_threads = len(threading.enumerate())
+
+    store = zen_store()
+
+    if store.config.driver == "sqlite":
+        # SQLite doesn't have a connection pool
+        return ServerLoadInfo(
+            threads=num_threads,
+            db_connections_total=0,
+            db_connections_active=0,
+            db_connections_overflow=0,
+        )
+
+    from sqlalchemy.pool import QueuePool
+
+    # Get the number of connections
+    pool = store.engine.pool
+    assert isinstance(pool, QueuePool)
+    idle_conn = pool.checkedin()
+    active_conn = pool.checkedout()
+    overflow_conn = max(0, pool.overflow())
+    total_conn = idle_conn + active_conn
+
+    return ServerLoadInfo(
+        threads=num_threads,
+        db_connections_total=total_conn,
+        db_connections_active=active_conn,
+        db_connections_overflow=overflow_conn,
+    )
+
+
 @router.get(
     ONBOARDING_STATE,
     responses={

zenml/zen_server/routers/users_endpoints.py

@@ -46,14 +46,10 @@ from zenml.zen_server.auth import (
 )
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rate_limit import RequestLimiter
-from zenml.zen_server.rbac.endpoint_utils import (
-    verify_permissions_and_create_entity,
-)
 from zenml.zen_server.rbac.models import Action, Resource, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
     dehydrate_response_model,
-    get_allowed_resource_ids,
     get_schema_for_resource_type,
     update_resource_membership,
     verify_permission_for_model,
@@ -112,17 +108,18 @@ def list_users(
     Returns:
         A list of all users.
     """
-    allowed_ids = get_allowed_resource_ids(resource_type=ResourceType.USER)
-    if allowed_ids is not None:
-        # Make sure users can see themselves
-        allowed_ids.add(auth_context.user.id)
-    else:
-        if not auth_context.user.is_admin and not server_config().rbac_enabled:
-            allowed_ids = {auth_context.user.id}
-
-    user_filter_model.configure_rbac(
-        authenticated_user_id=auth_context.user.id, id=allowed_ids
-    )
+    # allowed_ids = get_allowed_resource_ids(resource_type=ResourceType.USER)
+    # if allowed_ids is not None:
+    #     # Make sure users can see themselves
+    #     allowed_ids.add(auth_context.user.id)
+    # else:
+    #     if not auth_context.user.is_admin and not server_config().rbac_enabled:
+    #         allowed_ids = {auth_context.user.id}
+    if not auth_context.user.is_admin and not server_config().rbac_enabled:
+        user_filter_model.configure_rbac(
+            authenticated_user_id=auth_context.user.id,
+            id={auth_context.user.id},
+        )
 
     page = zen_store().list_users(
         user_filter_model=user_filter_model, hydrate=hydrate
@@ -175,11 +172,12 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             auth_context.user.is_admin, "create user"
         )
 
-        new_user = verify_permissions_and_create_entity(
-            request_model=user,
-            resource_type=ResourceType.USER,
-            create_method=zen_store().create_user,
-        )
+        # new_user = verify_permissions_and_create_entity(
+        #     request_model=user,
+        #     resource_type=ResourceType.USER,
+        #     create_method=zen_store().create_user,
+        # )
+        new_user = zen_store().create_user(user)
 
         # add back the original unhashed activation token, if generated, to
         # send it back to the client
@@ -217,10 +215,10 @@ def get_user(
         verify_admin_status_if_no_rbac(
             auth_context.user.is_admin, "get other user"
         )
-        verify_permission_for_model(
-            user,
-            action=Action.READ,
-        )
+        # verify_permission_for_model(
+        #     user,
+        #     action=Action.READ,
+        # )
 
     return dehydrate_response_model(user)
 
@@ -304,10 +302,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             verify_admin_status_if_no_rbac(
                 auth_context.user.is_admin, "update other user account"
             )
-            verify_permission_for_model(
-                user,
-                action=Action.UPDATE,
-            )
+            # verify_permission_for_model(
+            #     user,
+            #     action=Action.UPDATE,
+            # )
 
         # Validate a password change
         if user_update.password is not None:
@@ -497,10 +495,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
         verify_admin_status_if_no_rbac(
             auth_context.user.is_admin, "deactivate user"
         )
-        verify_permission_for_model(
-            user,
-            action=Action.UPDATE,
-        )
+        # verify_permission_for_model(
+        #     user,
+        #     action=Action.UPDATE,
+        # )
 
         user_update = UserUpdate(
             active=False,
@@ -548,10 +546,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             verify_admin_status_if_no_rbac(
                 auth_context.user.is_admin, "delete user"
             )
-            verify_permission_for_model(
-                user,
-                action=Action.DELETE,
-            )
+            # verify_permission_for_model(
+            #     user,
+            #     action=Action.DELETE,
+            # )
 
         zen_store().delete_user(user_name_or_id=user_name_or_id)
 
@@ -746,7 +744,7 @@ if server_config().rbac_enabled:
             KeyError: If no resource with the given type and ID exists.
         """
         user = zen_store().get_user(user_name_or_id)
-        verify_permission_for_model(user, action=Action.READ)
+        # verify_permission_for_model(user, action=Action.READ)
 
         if user.id == auth_context.user.id:
             raise ValueError(