Skip to content

Zen Server

zenml.zen_server

ZenML Server Implementation.

The ZenML Server is a centralized service meant for use in a collaborative setting in which stacks, stack components, flavors, pipeline and pipeline runs can be shared over the network with other users.

You can use the zenml server up command to spin up ZenML server instances that are either running locally as daemon processes or docker containers, or to deploy a ZenML server remotely on a managed cloud platform. The other CLI commands in the same zenml server group can be used to manage the server instances deployed from your local machine.

To connect the local ZenML client to one of the managed ZenML servers, call zenml server connect with the name of the server you want to connect to.

Modules

auth

Authentication module for ZenML server.

Classes
AuthContext

Bases: BaseModel

The authentication context.

CookieOAuth2TokenBearer

Bases: OAuth2PasswordBearer

OAuth2 token bearer authentication scheme that uses a cookie.

Functions
authenticate_api_key(api_key: str) -> AuthContext

Implement service account API key authentication.

Parameters:

Name Type Description Default
api_key str

The service account API key.

required

Returns:

Type Description
AuthContext

The authentication context reflecting the authenticated service account.

Raises:

Type Description
CredentialsNotValid

If the service account could not be authorized.

Source code in src/zenml/zen_server/auth.py
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
def authenticate_api_key(
    api_key: str,
) -> AuthContext:
    """Implement service account API key authentication.

    Args:
        api_key: The service account API key.


    Returns:
        The authentication context reflecting the authenticated service account.

    Raises:
        CredentialsNotValid: If the service account could not be authorized.
    """
    try:
        decoded_api_key = APIKey.decode_api_key(api_key)
    except ValueError:
        error = "Authentication error: error decoding API key"
        logger.exception(error)
        raise CredentialsNotValid(error)

    internal_api_key = _fetch_and_verify_api_key(
        api_key_id=decoded_api_key.id, key_to_verify=decoded_api_key.key
    )

    # For now, a lot of code still relies on the active user in the auth
    # context being a UserResponse object, which is a superset of the
    # ServiceAccountResponse object. So we need to convert the service
    # account to a user here.
    user_model = internal_api_key.service_account.to_user_model()
    return AuthContext(user=user_model, api_key=internal_api_key)
authenticate_credentials(user_name_or_id: Optional[Union[str, UUID]] = None, password: Optional[str] = None, access_token: Optional[str] = None, csrf_token: Optional[str] = None, activation_token: Optional[str] = None) -> AuthContext

Verify if user authentication credentials are valid.

This function can be used to validate all supplied user credentials to cover a range of possibilities:

  • username only - only when the no-auth scheme is used
  • username+password - for basic HTTP authentication or the OAuth2 password grant
  • access token (with embedded user id) - after successful authentication using one of the supported grants
  • username+activation token - for user activation

Parameters:

Name Type Description Default
user_name_or_id Optional[Union[str, UUID]]

The username or user ID.

None
password Optional[str]

The password.

None
access_token Optional[str]

The access token.

None
csrf_token Optional[str]

The CSRF token.

None
activation_token Optional[str]

The activation token.

None

Returns:

Type Description
AuthContext

The authenticated account details.

Raises:

Type Description
CredentialsNotValid

If the credentials are invalid.

Source code in src/zenml/zen_server/auth.py
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def authenticate_credentials(
    user_name_or_id: Optional[Union[str, UUID]] = None,
    password: Optional[str] = None,
    access_token: Optional[str] = None,
    csrf_token: Optional[str] = None,
    activation_token: Optional[str] = None,
) -> AuthContext:
    """Verify if user authentication credentials are valid.

    This function can be used to validate all supplied user credentials to
    cover a range of possibilities:

     * username only - only when the no-auth scheme is used
     * username+password - for basic HTTP authentication or the OAuth2 password
       grant
     * access token (with embedded user id) - after successful authentication
       using one of the supported grants
     * username+activation token - for user activation

    Args:
        user_name_or_id: The username or user ID.
        password: The password.
        access_token: The access token.
        csrf_token: The CSRF token.
        activation_token: The activation token.

    Returns:
        The authenticated account details.

    Raises:
        CredentialsNotValid: If the credentials are invalid.
    """
    user: Optional[UserAuthModel] = None
    auth_context: Optional[AuthContext] = None
    if user_name_or_id:
        try:
            # NOTE: this method will not return a user if the user name or ID
            # identifies a service account instead of a regular user. This
            # is intentional because service accounts are not allowed to
            # be used to authenticate to the API using a username and password,
            # or an activation token.
            user = zen_store().get_auth_user(user_name_or_id)
            user_model = zen_store().get_user(
                user_name_or_id=user_name_or_id, include_private=True
            )
            auth_context = AuthContext(user=user_model)
        except KeyError:
            # even when the user does not exist, we still want to execute the
            # password/token verification to protect against response discrepancy
            # attacks (https://cwe.mitre.org/data/definitions/204.html)
            logger.exception(
                f"Authentication error: error retrieving account "
                f"{user_name_or_id}"
            )
            pass

    if password is not None:
        if not UserAuthModel.verify_password(password, user):
            error = "Authentication error: invalid username or password"
            logger.error(error)
            raise CredentialsNotValid(error)
        if user and not user.active:
            error = f"Authentication error: user {user.name} is not active"
            logger.error(error)
            raise CredentialsNotValid(error)

    elif activation_token is not None:
        if not UserAuthModel.verify_activation_token(activation_token, user):
            error = (
                f"Authentication error: invalid activation token for user "
                f"{user_name_or_id}"
            )
            logger.error(error)
            raise CredentialsNotValid(error)

    elif access_token is not None:
        try:
            decoded_token = JWTToken.decode_token(
                token=access_token,
            )
        except CredentialsNotValid as e:
            error = f"Authentication error: error decoding access token: {e}."
            logger.exception(error)
            raise CredentialsNotValid(error)

        if decoded_token.session_id:
            if not csrf_token:
                error = "Authentication error: missing CSRF token"
                logger.error(error)
                raise CredentialsNotValid(error)

            decoded_csrf_token = CSRFToken.decode_token(csrf_token)

            if decoded_csrf_token.session_id != decoded_token.session_id:
                error = (
                    "Authentication error: CSRF token does not match the "
                    "access token"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

        try:
            user_model = zen_store().get_user(
                user_name_or_id=decoded_token.user_id, include_private=True
            )
        except KeyError:
            error = (
                f"Authentication error: error retrieving token account "
                f"{decoded_token.user_id}"
            )
            logger.error(error)
            raise CredentialsNotValid(error)

        if not user_model.active:
            error = (
                f"Authentication error: account {user_model.name} is not "
                f"active"
            )
            logger.error(error)
            raise CredentialsNotValid(error)

        api_key_model: Optional[APIKeyInternalResponse] = None
        if decoded_token.api_key_id:
            # The API token was generated from an API key. We still have to
            # verify if the API key hasn't been deactivated or deleted in the
            # meantime.
            api_key_model = _fetch_and_verify_api_key(decoded_token.api_key_id)

        device_model: Optional[OAuthDeviceInternalResponse] = None
        if decoded_token.device_id:
            if server_config().auth_scheme in [
                AuthScheme.NO_AUTH,
                AuthScheme.EXTERNAL,
            ]:
                error = "Authentication error: device authorization is not supported."
                logger.error(error)
                raise CredentialsNotValid(error)

            # Access tokens that have been issued for a device are only valid
            # for that device, so we need to check if the device ID matches any
            # of the valid devices in the database.
            try:
                device_model = zen_store().get_internal_authorized_device(
                    device_id=decoded_token.device_id
                )
            except KeyError:
                error = (
                    f"Authentication error: error retrieving token device "
                    f"{decoded_token.device_id}"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if (
                device_model.user is None
                or device_model.user.id != user_model.id
            ):
                error = (
                    f"Authentication error: device {decoded_token.device_id} "
                    f"does not belong to user {user_model.name}"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if device_model.status != OAuthDeviceStatus.ACTIVE:
                error = (
                    f"Authentication error: device {decoded_token.device_id} "
                    f"is not active"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if (
                device_model.expires
                and utc_now(tz_aware=device_model.expires)
                >= device_model.expires
            ):
                error = (
                    f"Authentication error: device {decoded_token.device_id} "
                    "has expired"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            zen_store().update_internal_authorized_device(
                device_id=device_model.id,
                update=OAuthDeviceInternalUpdate(
                    update_last_login=True,
                ),
            )

        if decoded_token.schedule_id:
            # If the token contains a schedule ID, we need to check if the
            # schedule still exists in the database. We use a cached version
            # of the schedule active status to avoid unnecessary database
            # queries.

            @cache_result(expiry=30)
            def get_schedule_active(schedule_id: UUID) -> Optional[bool]:
                """Get the active status of a schedule.

                Args:
                    schedule_id: The schedule ID.

                Returns:
                    The schedule active status or None if the schedule does not
                    exist.
                """
                try:
                    schedule = zen_store().get_schedule(
                        schedule_id, hydrate=False
                    )
                except KeyError:
                    return False

                return schedule.active

            schedule_active = get_schedule_active(decoded_token.schedule_id)
            if schedule_active is None:
                error = (
                    f"Authentication error: error retrieving token schedule "
                    f"{decoded_token.schedule_id}"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if not schedule_active:
                error = (
                    f"Authentication error: schedule {decoded_token.schedule_id} "
                    "is not active"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

        if decoded_token.pipeline_run_id:
            # If the token contains a pipeline run ID, we need to check if the
            # pipeline run exists in the database and the pipeline run has
            # not concluded. We use a cached version of the pipeline run status
            # to avoid unnecessary database queries.

            @cache_result(expiry=30)
            def get_pipeline_run_status(
                pipeline_run_id: UUID,
            ) -> Optional[ExecutionStatus]:
                """Get the status of a pipeline run.

                Args:
                    pipeline_run_id: The pipeline run ID.

                Returns:
                    The pipeline run status or None if the pipeline run does not
                    exist.
                """
                try:
                    pipeline_run = zen_store().get_run(
                        pipeline_run_id, hydrate=False
                    )
                except KeyError:
                    return None

                return pipeline_run.status

            pipeline_run_status = get_pipeline_run_status(
                decoded_token.pipeline_run_id
            )
            if pipeline_run_status is None:
                error = (
                    f"Authentication error: error retrieving token pipeline run "
                    f"{decoded_token.pipeline_run_id}"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if pipeline_run_status.is_finished:
                error = (
                    f"The execution of pipeline run "
                    f"{decoded_token.pipeline_run_id} has already concluded and "
                    "API tokens scoped to it are no longer valid."
                )
                logger.error(error)
                raise CredentialsNotValid(error)

        if decoded_token.step_run_id:
            # If the token contains a step run ID, we need to check if the
            # step run exists in the database and the step run has not concluded.
            # We use a cached version of the step run status to avoid unnecessary
            # database queries.

            @cache_result(expiry=30)
            def get_step_run_status(
                step_run_id: UUID,
            ) -> Optional[ExecutionStatus]:
                """Get the status of a step run.

                Args:
                    step_run_id: The step run ID.

                Returns:
                    The step run status or None if the step run does not exist.
                """
                try:
                    step_run = zen_store().get_run_step(
                        step_run_id, hydrate=False
                    )
                except KeyError:
                    return None

                return step_run.status

            step_run_status = get_step_run_status(decoded_token.step_run_id)
            if step_run_status is None:
                error = (
                    f"Authentication error: error retrieving token step run "
                    f"{decoded_token.step_run_id}"
                )
                logger.error(error)
                raise CredentialsNotValid(error)

            if step_run_status.is_finished:
                error = (
                    f"The execution of step run "
                    f"{decoded_token.step_run_id} has already concluded and "
                    "API tokens scoped to it are no longer valid."
                )
                logger.error(error)
                raise CredentialsNotValid(error)

        auth_context = AuthContext(
            user=user_model,
            access_token=decoded_token,
            encoded_access_token=access_token,
            device=device_model,
            api_key=api_key_model,
        )

    else:
        # IMPORTANT: the ONLY way we allow the authentication process to
        # continue without any credentials (i.e. no password, activation
        # token or access token) is if authentication is explicitly disabled
        # by setting the auth_scheme to NO_AUTH.
        if server_config().auth_scheme != AuthScheme.NO_AUTH:
            error = "Authentication error: no credentials provided"
            logger.error(error)
            raise CredentialsNotValid(error)

    if not auth_context:
        error = "Authentication error: invalid credentials"
        logger.error(error)
        raise CredentialsNotValid(error)

    return auth_context
authenticate_device(client_id: UUID, device_code: str) -> AuthContext

Verify if device authorization credentials are valid.

Parameters:

Name Type Description Default
client_id UUID

The OAuth2 client ID.

required
device_code str

The device code.

required

Returns:

Type Description
AuthContext

The authenticated account details.

Raises:

Type Description
OAuthError

If the device authorization credentials are invalid.

Source code in src/zenml/zen_server/auth.py
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
    """Verify if device authorization credentials are valid.

    Args:
        client_id: The OAuth2 client ID.
        device_code: The device code.

    Returns:
        The authenticated account details.

    Raises:
        OAuthError: If the device authorization credentials are invalid.
    """
    # This is the part of the OAuth2 device code grant flow where a client
    # device is continuously polling the server to check if the user has
    # authorized a device. The following needs to happen to successfully
    # authenticate the device and return a valid access token:
    #
    # 1. the device code and client ID must match a device in the DB
    # 2. the device must be in the VERIFIED state, meaning that the user
    # has successfully authorized the device via the user code but the
    # device client hasn't yet fetched the associated API access token yet.
    # 3. the device must not be expired

    config = server_config()
    store = zen_store()

    try:
        device_model = store.get_internal_authorized_device(
            client_id=client_id
        )
    except KeyError:
        error = (
            f"Authentication error: error retrieving device with client ID "
            f"{client_id}"
        )
        logger.error(error)
        raise OAuthError(
            error="invalid_client",
            error_description=error,
        )

    if device_model.status != OAuthDeviceStatus.VERIFIED:
        error = (
            f"Authentication error: device with client ID {client_id} is "
            f"{device_model.status.value}."
        )
        logger.error(error)
        if device_model.status == OAuthDeviceStatus.PENDING:
            oauth_error = "authorization_pending"
        elif device_model.status == OAuthDeviceStatus.LOCKED:
            oauth_error = "access_denied"
        else:
            oauth_error = "expired_token"
        raise OAuthError(
            error=oauth_error,
            error_description=error,
        )

    if (
        device_model.expires
        and utc_now(tz_aware=device_model.expires) >= device_model.expires
    ):
        error = (
            f"Authentication error: device for client ID {client_id} has "
            "expired"
        )
        logger.error(error)
        raise OAuthError(
            error="expired_token",
            error_description=error,
        )

    # Check the device code
    if not device_model.verify_device_code(device_code):
        # If the device code is invalid, increment the failed auth attempts
        # counter and lock the device if the maximum number of failed auth
        # attempts has been reached.
        failed_auth_attempts = device_model.failed_auth_attempts + 1
        update = OAuthDeviceInternalUpdate(
            failed_auth_attempts=failed_auth_attempts
        )
        if failed_auth_attempts >= config.max_failed_device_auth_attempts:
            update.locked = True

        store.update_internal_authorized_device(
            device_id=device_model.id,
            update=update,
        )

        if failed_auth_attempts >= config.max_failed_device_auth_attempts:
            error = (
                f"Authentication error: device for client ID {client_id} "
                "has been locked due to too many failed authentication "
                "attempts."
            )
        else:
            error = (
                f"Authentication error: device for client ID {client_id} "
                "has an invalid device code."
            )

        logger.error(error)
        raise OAuthError(
            error="access_denied",
            error_description=error,
        )

    # The device is valid, so we can return the user associated with it.
    # This is the one and only time we return an AuthContext authorized by
    # a device code in order to be exchanged for an access token. Subsequent
    # requests to the API will be authenticated using the access token.
    #
    # Update the device state to ACTIVE and set an expiration date for it
    # past which it can no longer be used for authentication. The expiration
    # date also determines the expiration date of the access token issued
    # for this device.
    expires_in: int = 0
    if config.jwt_token_expire_minutes:
        if device_model.trusted_device:
            expires_in = config.trusted_device_expiration_minutes or 0
        else:
            expires_in = config.device_expiration_minutes or 0

    update = OAuthDeviceInternalUpdate(
        status=OAuthDeviceStatus.ACTIVE,
        expires_in=expires_in * 60,
    )
    device_model = zen_store().update_internal_authorized_device(
        device_id=device_model.id,
        update=update,
    )

    # This can never happen because the VERIFIED state is only set if
    # a user verified and has been associated with the device.
    assert device_model.user is not None

    return AuthContext(user=device_model.user, device=device_model)
authenticate_external_user(external_access_token: str, request: Request) -> AuthContext

Implement external authentication.

Parameters:

Name Type Description Default
external_access_token str

The access token used to authenticate the user to the external authenticator.

required
request Request

The request object.

required

Returns:

Type Description
AuthContext

The authentication context reflecting the authenticated user.

Raises:

Type Description
AuthorizationException

If the external user could not be authorized.

Source code in src/zenml/zen_server/auth.py
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
def authenticate_external_user(
    external_access_token: str, request: Request
) -> AuthContext:
    """Implement external authentication.

    Args:
        external_access_token: The access token used to authenticate the user
            to the external authenticator.
        request: The request object.

    Returns:
        The authentication context reflecting the authenticated user.

    Raises:
        AuthorizationException: If the external user could not be authorized.
    """
    config = server_config()
    store = zen_store()

    assert config.external_user_info_url is not None

    # Use the external access token to extract the user information and
    # permissions

    # Get the user information from the external authenticator
    user_info_url = config.external_user_info_url
    headers = {"Authorization": "Bearer " + external_access_token}
    headers.update(get_zenml_headers())
    query_params = dict(server_id=str(config.get_external_server_id()))

    try:
        auth_response = requests.get(
            user_info_url,
            headers=headers,
            params=urlencode(query_params),
            timeout=EXTERNAL_AUTHENTICATOR_TIMEOUT,
        )
    except Exception as e:
        logger.exception(
            f"Error fetching user information from external authenticator: {e}"
        )
        raise AuthorizationException(
            "Error fetching user information from external authenticator."
        )

    external_user: Optional[ExternalUserModel] = None

    if 200 <= auth_response.status_code < 300:
        try:
            payload = auth_response.json()
        except requests.exceptions.JSONDecodeError:
            logger.exception(
                "Error decoding JSON response from external authenticator."
            )
            raise AuthorizationException(
                "Unknown external authenticator error"
            )

        if isinstance(payload, dict):
            try:
                external_user = ExternalUserModel.model_validate(payload)
            except Exception as e:
                logger.exception(
                    f"Error parsing user information from external "
                    f"authenticator: {e}"
                )
                pass

    elif auth_response.status_code in [401, 403]:
        raise AuthorizationException("Not authorized to access this server.")
    elif auth_response.status_code == 404:
        raise AuthorizationException(
            "External authenticator did not recognize this server."
        )
    else:
        logger.error(
            f"Error fetching user information from external authenticator. "
            f"Status code: {auth_response.status_code}, "
            f"Response: {auth_response.text}"
        )
        raise AuthorizationException(
            "Error fetching user information from external authenticator. "
        )

    if not external_user:
        raise AuthorizationException("Unknown external authenticator error")

    # With an external user object, we can now authenticate the user against
    # the ZenML server

    # Check if the external user already exists in the ZenML server database
    # If not, create a new user. If yes, update the existing user.
    try:
        user = store.get_external_user(user_id=external_user.id)

        # Update the user information
        user = store.update_user(
            user_id=user.id,
            user_update=UserUpdate(
                name=external_user.email,
                full_name=external_user.name or "",
                email_opted_in=True,
                active=True,
                email=external_user.email,
                is_admin=external_user.is_admin,
            ),
        )
    except KeyError:
        logger.info(
            f"External user with ID {external_user.id} not found in ZenML "
            f"server database. Creating a new user."
        )
        user = store.create_user(
            UserRequest(
                name=external_user.email,
                full_name=external_user.name or "",
                external_user_id=external_user.id,
                email_opted_in=True,
                active=True,
                email=external_user.email,
                is_admin=external_user.is_admin,
            )
        )

        with AnalyticsContext() as context:
            context.user_id = user.id
            context.identify(
                traits={
                    "email": external_user.email,
                    "source": "external_auth",
                }
            )
            context.alias(user_id=external_user.id, previous_id=user.id)

    # This is the best spot to update the onboarding state to mark the
    # "zenml login" step as completed for ZenML Pro servers, because the
    # user has just successfully logged in. However, we need to differentiate
    # between web clients (i.e. the dashboard) and CLI clients (i.e. the
    # zenml CLI).
    user_agent = request.headers.get("User-Agent", "").lower()
    if "zenml/" in user_agent:
        store.update_onboarding_state(
            completed_steps={OnboardingStep.DEVICE_VERIFIED}
        )

    return AuthContext(user=user)
authentication_provider() -> Callable[..., AuthContext]

Returns the authentication provider.

Returns:

Type Description
Callable[..., AuthContext]

The authentication provider.

Raises:

Type Description
ValueError

If the authentication scheme is not supported.

Source code in src/zenml/zen_server/auth.py
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
def authentication_provider() -> Callable[..., AuthContext]:
    """Returns the authentication provider.

    Returns:
        The authentication provider.

    Raises:
        ValueError: If the authentication scheme is not supported.
    """
    auth_scheme = server_config().auth_scheme
    if auth_scheme == AuthScheme.NO_AUTH:
        return no_authentication
    elif auth_scheme == AuthScheme.HTTP_BASIC:
        return http_authentication
    elif auth_scheme == AuthScheme.OAUTH2_PASSWORD_BEARER:
        return oauth2_authentication
    elif auth_scheme == AuthScheme.EXTERNAL:
        return oauth2_authentication
    else:
        raise ValueError(f"Unknown authentication scheme: {auth_scheme}")
generate_access_token(user_id: UUID, response: Optional[Response] = None, request: Optional[Request] = None, device: Optional[OAuthDeviceInternalResponse] = None, api_key: Optional[APIKeyInternalResponse] = None, expires_in: Optional[int] = None, schedule_id: Optional[UUID] = None, pipeline_run_id: Optional[UUID] = None, step_run_id: Optional[UUID] = None) -> OAuthTokenResponse

Generates an access token for the given user.

Parameters:

Name Type Description Default
user_id UUID

The ID of the user.

required
response Optional[Response]

The FastAPI response object. If passed, the access token will also be set as an HTTP only cookie in the response.

None
request Optional[Request]

The FastAPI request object. Used to determine the request origin and to decide whether to use cross-site security measures for the access token cookie.

None
device Optional[OAuthDeviceInternalResponse]

The device used for authentication.

None
api_key Optional[APIKeyInternalResponse]

The service account API key used for authentication.

None
expires_in Optional[int]

The number of seconds until the token expires. If not set, the default value is determined automatically based on the server configuration and type of token. If set to 0, the token will not expire.

None
schedule_id Optional[UUID]

The ID of the schedule to scope the token to.

None
pipeline_run_id Optional[UUID]

The ID of the pipeline run to scope the token to.

None
step_run_id Optional[UUID]

The ID of the step run to scope the token to.

None

Returns:

Type Description
OAuthTokenResponse

An authentication response with an access token.

Source code in src/zenml/zen_server/auth.py
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
def generate_access_token(
    user_id: UUID,
    response: Optional[Response] = None,
    request: Optional[Request] = None,
    device: Optional[OAuthDeviceInternalResponse] = None,
    api_key: Optional[APIKeyInternalResponse] = None,
    expires_in: Optional[int] = None,
    schedule_id: Optional[UUID] = None,
    pipeline_run_id: Optional[UUID] = None,
    step_run_id: Optional[UUID] = None,
) -> OAuthTokenResponse:
    """Generates an access token for the given user.

    Args:
        user_id: The ID of the user.
        response: The FastAPI response object. If passed, the access
            token will also be set as an HTTP only cookie in the response.
        request: The FastAPI request object. Used to determine the request
            origin and to decide whether to use cross-site security measures for
            the access token cookie.
        device: The device used for authentication.
        api_key: The service account API key used for authentication.
        expires_in: The number of seconds until the token expires. If not set,
            the default value is determined automatically based on the server
            configuration and type of token. If set to 0, the token will not
            expire.
        schedule_id: The ID of the schedule to scope the token to.
        pipeline_run_id: The ID of the pipeline run to scope the token to.
        step_run_id: The ID of the step run to scope the token to.

    Returns:
        An authentication response with an access token.
    """
    config = server_config()

    # If the expiration time is not supplied, the JWT tokens are set to expire
    # according to the values configured in the server config. Device tokens are
    # handled separately from regular user tokens.
    expires: Optional[datetime] = None
    if expires_in == 0:
        expires_in = None
    elif expires_in is not None:
        expires = utc_now() + timedelta(seconds=expires_in)
    elif device:
        # If a device was used for authentication, the token will expire
        # at the same time as the device.
        expires = device.expires
        if expires:
            expires_in = max(
                int(expires.timestamp() - utc_now().timestamp()),
                0,
            )
    elif config.jwt_token_expire_minutes:
        expires = utc_now() + timedelta(
            minutes=config.jwt_token_expire_minutes
        )
        expires_in = config.jwt_token_expire_minutes * 60

    # Figure out if this is a same-site request or a cross-site request
    same_site = True
    if response and request:
        # Extract the origin domain from the request; use the referer as a
        # fallback
        origin_domain: Optional[str] = None
        origin = request.headers.get("origin", request.headers.get("referer"))
        if origin:
            # If the request origin is known, we use it to determine whether
            # this is a cross-site request and enable additional security
            # measures.
            origin_domain = urlparse(origin).netloc

        server_domain: Optional[str] = config.auth_cookie_domain
        # If the server's cookie domain is not explicitly set in the
        # server's configuration, we use other sources to determine it:
        #
        # 1. the server's root URL, if set in the server's configuration
        # 2. the X-Forwarded-Host header, if set by the reverse proxy
        # 3. the request URL, if all else fails
        if not server_domain and config.server_url:
            server_domain = urlparse(config.server_url).netloc
        if not server_domain:
            server_domain = request.headers.get(
                "x-forwarded-host", request.url.netloc
            )

        # Same-site requests can come from the same domain or from a
        # subdomain of the domain used to issue cookies.
        if origin_domain and server_domain:
            same_site = is_same_or_subdomain(origin_domain, server_domain)

    csrf_token: Optional[str] = None
    session_id: Optional[UUID] = None
    if not same_site:
        # If responding to a cross-site login request, we need to generate and
        # sign a CSRF token associated with the authentication session.
        session_id = uuid4()
        csrf_token = CSRFToken(session_id=session_id).encode()

    access_token = JWTToken(
        user_id=user_id,
        device_id=device.id if device else None,
        api_key_id=api_key.id if api_key else None,
        schedule_id=schedule_id,
        pipeline_run_id=pipeline_run_id,
        step_run_id=step_run_id,
        # Set the session ID if this is a cross-site request
        session_id=session_id,
    ).encode(expires=expires)

    if response:
        # Also set the access token as an HTTP only cookie in the response
        response.set_cookie(
            key=config.get_auth_cookie_name(),
            value=access_token,
            httponly=True,
            secure=not same_site,
            samesite="lax" if same_site else "none",
            max_age=config.jwt_token_expire_minutes * 60
            if config.jwt_token_expire_minutes
            else None,
            domain=config.auth_cookie_domain,
        )

    return OAuthTokenResponse(
        access_token=access_token,
        expires_in=expires_in,
        token_type="bearer",
        csrf_token=csrf_token,
    )
get_auth_context() -> Optional[AuthContext]

Returns the current authentication context.

Returns:

Type Description
Optional[AuthContext]

The authentication context.

Source code in src/zenml/zen_server/auth.py
85
86
87
88
89
90
91
92
def get_auth_context() -> Optional["AuthContext"]:
    """Returns the current authentication context.

    Returns:
        The authentication context.
    """
    auth_context = _auth_context.get()
    return auth_context
http_authentication(credentials: HTTPBasicCredentials = Depends(HTTPBasic())) -> AuthContext

Authenticates any request to the ZenML Server with basic HTTP authentication.

Parameters:

Name Type Description Default
credentials HTTPBasicCredentials

HTTP basic auth credentials passed to the request.

Depends(HTTPBasic())

Returns:

Type Description
AuthContext

The authentication context reflecting the authenticated user.

noqa: DAR401
Source code in src/zenml/zen_server/auth.py
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
def http_authentication(
    credentials: HTTPBasicCredentials = Depends(HTTPBasic()),
) -> AuthContext:
    """Authenticates any request to the ZenML Server with basic HTTP authentication.

    Args:
        credentials: HTTP basic auth credentials passed to the request.

    Returns:
        The authentication context reflecting the authenticated user.

    # noqa: DAR401
    """
    try:
        return authenticate_credentials(
            user_name_or_id=credentials.username, password=credentials.password
        )
    except CredentialsNotValid as e:
        # We want to be very explicit here and return a CredentialsNotValid
        # exception encoded as a 401 Unauthorized error encoded, so that the
        # client can distinguish between a 401 error due to invalid credentials
        # and other 401 errors and handle them accordingly by throwing away the
        # current access token and re-authenticating.
        raise http_exception_from_error(e)
no_authentication() -> AuthContext

Doesn't authenticate requests to the ZenML server.

Returns:

Type Description
AuthContext

The authentication context reflecting the default user.

Source code in src/zenml/zen_server/auth.py
1072
1073
1074
1075
1076
1077
1078
def no_authentication() -> AuthContext:
    """Doesn't authenticate requests to the ZenML server.

    Returns:
        The authentication context reflecting the default user.
    """
    return authenticate_credentials(user_name_or_id=DEFAULT_USERNAME)
oauth2_authentication(request: Request, token: str = Depends(CookieOAuth2TokenBearer(tokenUrl=server_config().root_url_path + API + VERSION_1 + LOGIN))) -> AuthContext

Authenticates any request to the ZenML server with OAuth2 JWT tokens.

Parameters:

Name Type Description Default
token str

The JWT bearer token to be authenticated.

Depends(CookieOAuth2TokenBearer(tokenUrl=root_url_path + API + VERSION_1 + LOGIN))
request Request

The FastAPI request object.

required

Returns:

Type Description
AuthContext

The authentication context reflecting the authenticated user.

noqa: DAR401
Source code in src/zenml/zen_server/auth.py
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
def oauth2_authentication(
    request: Request,
    token: str = Depends(
        CookieOAuth2TokenBearer(
            tokenUrl=server_config().root_url_path + API + VERSION_1 + LOGIN,
        )
    ),
) -> AuthContext:
    """Authenticates any request to the ZenML server with OAuth2 JWT tokens.

    Args:
        token: The JWT bearer token to be authenticated.
        request: The FastAPI request object.

    Returns:
        The authentication context reflecting the authenticated user.

    # noqa: DAR401
    """
    csrf_token = request.headers.get("X-CSRF-Token")
    try:
        auth_context = authenticate_credentials(
            access_token=token, csrf_token=csrf_token
        )
    except CredentialsNotValid as e:
        # We want to be very explicit here and return a CredentialsNotValid
        # exception encoded as a 401 Unauthorized error encoded, so that the
        # client can distinguish between a 401 error due to invalid credentials
        # and other 401 errors and handle them accordingly by throwing away the
        # current access token and re-authenticating.
        raise http_exception_from_error(e)

    return auth_context
set_auth_context(auth_context: AuthContext) -> AuthContext

Sets the current authentication context.

Parameters:

Name Type Description Default
auth_context AuthContext

The authentication context.

required

Returns:

Type Description
AuthContext

The authentication context.

Source code in src/zenml/zen_server/auth.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
def set_auth_context(auth_context: "AuthContext") -> "AuthContext":
    """Sets the current authentication context.

    Args:
        auth_context: The authentication context.

    Returns:
        The authentication context.
    """
    _auth_context.set(auth_context)
    return auth_context

cache

Memory cache module for the ZenML server.

Classes
MemoryCache(max_capacity: int, default_expiry: int)

Simple in-memory cache with expiry and capacity management.

This cache is thread-safe and can be used in both synchronous and asynchronous contexts. It uses a simple LRU (Least Recently Used) eviction strategy to manage the cache size.

Each cache entry has a key, value, timestamp, and expiry. The cache automatically removes expired entries and evicts the oldest entry when the cache reaches its maximum capacity.

Usage Example:

cache = MemoryCache()
uuid_key = UUID("12345678123456781234567812345678")

if not cache.get(uuid_key):
    # Get the value from the database or other source
    value = get_value_from_database()
    cache.set(uuid_key, value, expiry=60)

Usage Example with decorator:

@cache_result(expiry=60)
def get_cached_value(key: UUID) -> Any:
    return get_value_from_database(key)

uuid_key = UUID("12345678123456781234567812345678")

value = get_cached_value(uuid_key)

Initialize the cache with a maximum capacity and default expiry time.

Parameters:

Name Type Description Default
max_capacity int

The maximum number of entries the cache can hold.

required
default_expiry int

The default expiry time in seconds.

required
Source code in src/zenml/zen_server/cache.py
86
87
88
89
90
91
92
93
94
95
96
def __init__(self, max_capacity: int, default_expiry: int) -> None:
    """Initialize the cache with a maximum capacity and default expiry time.

    Args:
        max_capacity: The maximum number of entries the cache can hold.
        default_expiry: The default expiry time in seconds.
    """
    self.cache: OrderedDictType[UUID, MemoryCacheEntry] = OrderedDict()
    self.max_capacity = max_capacity
    self.default_expiry = default_expiry
    self._lock = Lock()
Functions
get(key: UUID) -> Optional[Any]

Retrieve value if it's still valid; otherwise, return None.

Parameters:

Name Type Description Default
key UUID

The key to retrieve the value for.

required

Returns:

Type Description
Optional[Any]

The value if it's still valid; otherwise, None.

Source code in src/zenml/zen_server/cache.py
112
113
114
115
116
117
118
119
120
121
122
def get(self, key: UUID) -> Optional[Any]:
    """Retrieve value if it's still valid; otherwise, return None.

    Args:
        key: The key to retrieve the value for.

    Returns:
        The value if it's still valid; otherwise, None.
    """
    with self._lock:
        return self._get_internal(key)
set(key: UUID, value: Any, expiry: Optional[int] = None) -> None

Insert value into cache with optional custom expiry time in seconds.

Parameters:

Name Type Description Default
key UUID

The key to insert the value with.

required
value Any

The value to insert into the cache.

required
expiry Optional[int]

The expiry time in seconds. If None, uses the default expiry.

None
Source code in src/zenml/zen_server/cache.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def set(self, key: UUID, value: Any, expiry: Optional[int] = None) -> None:
    """Insert value into cache with optional custom expiry time in seconds.

    Args:
        key: The key to insert the value with.
        value: The value to insert into the cache.
        expiry: The expiry time in seconds. If None, uses the default expiry.
    """
    with self._lock:
        self.cache[key] = MemoryCacheEntry(
            value=value, expiry=expiry or self.default_expiry
        )
        self._cleanup()
MemoryCacheEntry(value: Any, expiry: int)

Simple class to hold cache entry data.

Initialize a cache entry with value and expiry time.

Parameters:

Name Type Description Default
value Any

The value to store in the cache.

required
expiry int

The expiry time in seconds.

required
Source code in src/zenml/zen_server/cache.py
32
33
34
35
36
37
38
39
40
41
def __init__(self, value: Any, expiry: int) -> None:
    """Initialize a cache entry with value and expiry time.

    Args:
        value: The value to store in the cache.
        expiry: The expiry time in seconds.
    """
    self.value: Any = value
    self.expiry: int = expiry
    self.timestamp: float = time.time()
Attributes
expired: bool property

Check if the cache entry has expired.

Returns:

Type Description
bool

True if the cache entry has expired; otherwise, False.

Functions
Functions
cache_result(expiry: Optional[int] = None) -> Callable[[F], F]

A decorator to cache the result of a function based on a UUID key argument.

Parameters:

Name Type Description Default
expiry Optional[int]

Custom time in seconds for the cache entry to expire. If None, uses the default expiry time.

None

Returns:

Type Description
Callable[[F], F]

A decorator that wraps a function, caching its results based on a UUID

Callable[[F], F]

key.

Source code in src/zenml/zen_server/cache.py
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def cache_result(
    expiry: Optional[int] = None,
) -> Callable[[F], F]:
    """A decorator to cache the result of a function based on a UUID key argument.

    Args:
        expiry: Custom time in seconds for the cache entry to expire. If None,
            uses the default expiry time.

    Returns:
        A decorator that wraps a function, caching its results based on a UUID
        key.
    """

    def decorator(func: F) -> F:
        """The actual decorator that wraps the function with caching logic.

        Args:
            func: The function to wrap.

        Returns:
            The wrapped function with caching logic.
        """

        def wrapper(key: UUID) -> Any:
            """The wrapped function with caching logic.

            Args:
                key: The key to use for caching.

            Returns:
                The result of the original function, either from cache or
                freshly computed.
            """
            from zenml.zen_server.utils import memcache

            cache = memcache()

            # Attempt to retrieve the result from cache
            cached_value = cache.get(key)
            if cached_value is not None:
                logger.debug(
                    f"Memory cache hit for key: {key} and func: {func.__name__}"
                )
                return cached_value

            # Call the original function and cache its result
            result = func(key)
            cache.set(key, result, expiry)
            return result

        return wrapper

    return decorator

cloud_utils

Utils concerning anything concerning the cloud control plane backend.

Classes
ZenMLCloudConnection()

Class to use for communication between server and control plane.

Initialize the RBAC component.

Source code in src/zenml/zen_server/cloud_utils.py
20
21
22
23
24
25
def __init__(self) -> None:
    """Initialize the RBAC component."""
    self._config = ServerProConfiguration.get_server_config()
    self._session: Optional[requests.Session] = None
    self._token: Optional[str] = None
    self._token_expires_at: Optional[datetime] = None
Attributes
session: requests.Session property

Authenticate to the ZenML Pro Management Plane.

Returns:

Type Description
Session

A requests session with the authentication token.

Functions
delete(endpoint: str, params: Optional[Dict[str, Any]] = None, data: Optional[Dict[str, Any]] = None) -> requests.Response

Send a DELETE request using the active session.

Parameters:

Name Type Description Default
endpoint str

The endpoint to send the request to. This will be appended to the base URL.

required
params Optional[Dict[str, Any]]

Parameters to include in the request.

None
data Optional[Dict[str, Any]]

Data to include in the request.

None

Returns:

Type Description
Response

The response.

Source code in src/zenml/zen_server/cloud_utils.py
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def delete(
    self,
    endpoint: str,
    params: Optional[Dict[str, Any]] = None,
    data: Optional[Dict[str, Any]] = None,
) -> requests.Response:
    """Send a DELETE request using the active session.

    Args:
        endpoint: The endpoint to send the request to. This will be appended
            to the base URL.
        params: Parameters to include in the request.
        data: Data to include in the request.

    Returns:
        The response.
    """
    return self.request(
        method="DELETE", endpoint=endpoint, params=params, data=data
    )
get(endpoint: str, params: Optional[Dict[str, Any]]) -> requests.Response

Send a GET request using the active session.

Parameters:

Name Type Description Default
endpoint str

The endpoint to send the request to. This will be appended to the base URL.

required
params Optional[Dict[str, Any]]

Parameters to include in the request.

required

Returns:

Type Description
Response

The response.

Source code in src/zenml/zen_server/cloud_utils.py
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def get(
    self, endpoint: str, params: Optional[Dict[str, Any]]
) -> requests.Response:
    """Send a GET request using the active session.

    Args:
        endpoint: The endpoint to send the request to. This will be appended
            to the base URL.
        params: Parameters to include in the request.

    Returns:
        The response.
    """
    return self.request(method="GET", endpoint=endpoint, params=params)
patch(endpoint: str, params: Optional[Dict[str, Any]] = None, data: Optional[Dict[str, Any]] = None) -> requests.Response

Send a PATCH request using the active session.

Parameters:

Name Type Description Default
endpoint str

The endpoint to send the request to. This will be appended to the base URL.

required
params Optional[Dict[str, Any]]

Parameters to include in the request.

None
data Optional[Dict[str, Any]]

Data to include in the request.

None

Returns:

Type Description
Response

The response.

Source code in src/zenml/zen_server/cloud_utils.py
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def patch(
    self,
    endpoint: str,
    params: Optional[Dict[str, Any]] = None,
    data: Optional[Dict[str, Any]] = None,
) -> requests.Response:
    """Send a PATCH request using the active session.

    Args:
        endpoint: The endpoint to send the request to. This will be appended
            to the base URL.
        params: Parameters to include in the request.
        data: Data to include in the request.

    Returns:
        The response.
    """
    return self.request(
        method="PATCH", endpoint=endpoint, params=params, data=data
    )
post(endpoint: str, params: Optional[Dict[str, Any]] = None, data: Optional[Dict[str, Any]] = None) -> requests.Response

Send a POST request using the active session.

Parameters:

Name Type Description Default
endpoint str

The endpoint to send the request to. This will be appended to the base URL.

required
params Optional[Dict[str, Any]]

Parameters to include in the request.

None
data Optional[Dict[str, Any]]

Data to include in the request.

None

Returns:

Type Description
Response

The response.

Source code in src/zenml/zen_server/cloud_utils.py
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def post(
    self,
    endpoint: str,
    params: Optional[Dict[str, Any]] = None,
    data: Optional[Dict[str, Any]] = None,
) -> requests.Response:
    """Send a POST request using the active session.

    Args:
        endpoint: The endpoint to send the request to. This will be appended
            to the base URL.
        params: Parameters to include in the request.
        data: Data to include in the request.

    Returns:
        The response.
    """
    return self.request(
        method="POST", endpoint=endpoint, params=params, data=data
    )
request(method: str, endpoint: str, params: Optional[Dict[str, Any]] = None, data: Optional[Dict[str, Any]] = None) -> requests.Response

Send a request using the active session.

Parameters:

Name Type Description Default
method str

The HTTP method to use.

required
endpoint str

The endpoint to send the request to. This will be appended to the base URL.

required
params Optional[Dict[str, Any]]

Parameters to include in the request.

None
data Optional[Dict[str, Any]]

Data to include in the request.

None

Raises:

Type Description
SubscriptionUpgradeRequiredError

If the current subscription tier is insufficient for the attempted operation.

RuntimeError

If the request failed.

Returns:

Type Description
Response

The response.

Source code in src/zenml/zen_server/cloud_utils.py
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def request(
    self,
    method: str,
    endpoint: str,
    params: Optional[Dict[str, Any]] = None,
    data: Optional[Dict[str, Any]] = None,
) -> requests.Response:
    """Send a request using the active session.

    Args:
        method: The HTTP method to use.
        endpoint: The endpoint to send the request to. This will be appended
            to the base URL.
        params: Parameters to include in the request.
        data: Data to include in the request.

    Raises:
        SubscriptionUpgradeRequiredError: If the current subscription tier
            is insufficient for the attempted operation.
        RuntimeError: If the request failed.

    Returns:
        The response.
    """
    url = self._config.api_url + endpoint

    response = self.session.request(
        method=method, url=url, params=params, json=data, timeout=7
    )
    if response.status_code == 401:
        # Refresh the auth token and try again
        self._clear_session()
        response = self.session.request(
            method=method, url=url, params=params, json=data, timeout=7
        )

    try:
        response.raise_for_status()
    except requests.HTTPError as e:
        if response.status_code == 402:
            raise SubscriptionUpgradeRequiredError(response.json())
        else:
            raise RuntimeError(
                f"Failed while trying to contact the central zenml pro "
                f"service: {e}"
            )

    return response
Functions
cloud_connection() -> ZenMLCloudConnection

Return the initialized cloud connection.

Returns:

Type Description
ZenMLCloudConnection

The cloud connection.

Source code in src/zenml/zen_server/cloud_utils.py
257
258
259
260
261
262
263
264
265
266
267
def cloud_connection() -> ZenMLCloudConnection:
    """Return the initialized cloud connection.

    Returns:
        The cloud connection.
    """
    global _cloud_connection
    if _cloud_connection is None:
        _cloud_connection = ZenMLCloudConnection()

    return _cloud_connection
send_pro_workspace_status_update() -> None

Send a workspace status update to the Cloud API.

Source code in src/zenml/zen_server/cloud_utils.py
270
271
272
def send_pro_workspace_status_update() -> None:
    """Send a workspace status update to the Cloud API."""
    cloud_connection().patch("/workspace_status")

csrf

CSRF token utilities module for ZenML server.

Classes
CSRFToken

Bases: BaseModel

Pydantic object representing a CSRF token.

Attributes:

Name Type Description
session_id UUID

The id of the authenticated session.

Functions
decode_token(token: str) -> CSRFToken classmethod

Decodes a CSRF token.

Decodes a CSRF access token and returns a CSRFToken object with the information retrieved from its contents.

Parameters:

Name Type Description Default
token str

The encoded CSRF token.

required

Returns:

Type Description
CSRFToken

The decoded CSRF token.

Raises:

Type Description
CredentialsNotValid

If the token is invalid.

Source code in src/zenml/zen_server/csrf.py
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@classmethod
def decode_token(
    cls,
    token: str,
) -> "CSRFToken":
    """Decodes a CSRF token.

    Decodes a CSRF access token and returns a `CSRFToken` object with the
    information retrieved from its contents.

    Args:
        token: The encoded CSRF token.

    Returns:
        The decoded CSRF token.

    Raises:
        CredentialsNotValid: If the token is invalid.
    """
    from itsdangerous import BadData, BadSignature, URLSafeSerializer

    config = server_config()

    serializer = URLSafeSerializer(config.jwt_secret_key)
    try:
        # Decode and verify the token
        data = serializer.loads(token)
    except BadSignature as e:
        raise CredentialsNotValid(
            "Invalid CSRF token: signature mismatch"
        ) from e
    except BadData as e:
        raise CredentialsNotValid("Invalid CSRF token") from e

    try:
        return CSRFToken(session_id=UUID(data))
    except ValueError as e:
        raise CredentialsNotValid(
            "Invalid CSRF token: the session ID is not a valid UUID"
        ) from e
encode() -> str

Creates a CSRF token.

Encodes, signs and returns a CSRF access token.

Returns:

Type Description
str

The generated CSRF token.

Source code in src/zenml/zen_server/csrf.py
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def encode(self) -> str:
    """Creates a CSRF token.

    Encodes, signs and returns a CSRF access token.

    Returns:
        The generated CSRF token.
    """
    from itsdangerous import URLSafeSerializer

    config = server_config()

    serializer = URLSafeSerializer(config.jwt_secret_key)
    token = serializer.dumps(str(self.session_id))
    return token
Functions

deploy

ZenML server deployments.

Classes
LocalServerDeployer

Local server deployer singleton.

This class is responsible for managing the various server provider implementations and for directing server deployment lifecycle requests to the responsible provider. It acts as a facade built on top of the various server providers.

Functions
connect_to_server() -> None

Connect to the local ZenML server instance.

Raises:

Type Description
ServerDeploymentError

If the local ZenML server is not running or is unreachable.

Source code in src/zenml/zen_server/deploy/deployer.py
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def connect_to_server(
    self,
) -> None:
    """Connect to the local ZenML server instance.

    Raises:
        ServerDeploymentError: If the local ZenML server is not running or
            is unreachable.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    server = self.get_server()
    provider_name = server.config.provider.value

    gc = GlobalConfiguration()
    if not server.status or not server.status.url:
        raise ServerDeploymentError(
            f"The local {provider_name} ZenML server is not currently "
            "running or is unreachable."
        )

    store_config = RestZenStoreConfiguration(
        url=server.status.url,
    )

    if gc.store_configuration == store_config:
        logger.info(
            "Your client is already connected to the local "
            f"{provider_name} ZenML server."
        )
        return

    logger.info(
        f"Connecting to the local {provider_name} ZenML server "
        f"({store_config.url})."
    )

    gc.set_store(store_config)

    logger.info(
        f"Connected to the local {provider_name} ZenML server "
        f"({store_config.url})."
    )
deploy_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None, restart: bool = False) -> LocalServerDeployment

Deploy the local ZenML server or update the existing deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the deployment is successful. If not supplied, the default timeout value specified by the provider is used.

None
restart bool

If True, the existing server deployment will be torn down and a new server will be deployed.

False

Returns:

Type Description
LocalServerDeployment

The local server deployment.

Source code in src/zenml/zen_server/deploy/deployer.py
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def deploy_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
    restart: bool = False,
) -> LocalServerDeployment:
    """Deploy the local ZenML server or update the existing deployment.

    Args:
        config: The server deployment configuration.
        timeout: The timeout in seconds to wait until the deployment is
            successful. If not supplied, the default timeout value specified
            by the provider is used.
        restart: If True, the existing server deployment will be torn down
            and a new server will be deployed.

    Returns:
        The local server deployment.
    """
    # Ensure that the local database is always initialized before any local
    # server is deployed or updated.
    self.initialize_local_database()

    try:
        self.get_server()
    except ServerDeploymentNotFoundError:
        pass
    else:
        return self.update_server(
            config=config, timeout=timeout, restart=restart
        )

    provider_name = config.provider.value
    provider = self.get_provider(config.provider)

    logger.info(f"Deploying a local {provider_name} ZenML server.")
    return provider.deploy_server(config, timeout=timeout)
disconnect_from_server() -> None

Disconnect from the ZenML server instance.

Source code in src/zenml/zen_server/deploy/deployer.py
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def disconnect_from_server(
    self,
) -> None:
    """Disconnect from the ZenML server instance."""
    gc = GlobalConfiguration()
    store_cfg = gc.store_configuration

    if store_cfg.type != StoreType.REST:
        logger.info(
            "Your client is not currently connected to a ZenML server."
        )
        return

    logger.info(
        f"Disconnecting from the local ({store_cfg.url}) ZenML server."
    )

    gc.set_default_store()

    logger.info("Disconnected from the local ZenML server.")
get_provider(provider_type: ServerProviderType) -> BaseServerProvider classmethod

Get the server provider associated with a provider type.

Parameters:

Name Type Description Default
provider_type ServerProviderType

The server provider type.

required

Returns:

Type Description
BaseServerProvider

The server provider associated with the provider type.

Raises:

Type Description
ServerProviderNotFoundError

If no provider is registered for the given provider type.

Source code in src/zenml/zen_server/deploy/deployer.py
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@classmethod
def get_provider(
    cls, provider_type: ServerProviderType
) -> BaseServerProvider:
    """Get the server provider associated with a provider type.

    Args:
        provider_type: The server provider type.

    Returns:
        The server provider associated with the provider type.

    Raises:
        ServerProviderNotFoundError: If no provider is registered for the
            given provider type.
    """
    if provider_type not in cls._providers:
        raise ServerProviderNotFoundError(
            f"Server provider '{provider_type}' is not registered."
        )
    return cls._providers[provider_type]
get_server() -> LocalServerDeployment

Get the local server deployment.

Returns:

Type Description
LocalServerDeployment

The local server deployment.

Raises:

Type Description
ServerDeploymentNotFoundError

If no local server deployment is found.

Source code in src/zenml/zen_server/deploy/deployer.py
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
def get_server(
    self,
) -> LocalServerDeployment:
    """Get the local server deployment.

    Returns:
        The local server deployment.

    Raises:
        ServerDeploymentNotFoundError: If no local server deployment is
            found.
    """
    for provider in self._providers.values():
        try:
            return provider.get_server(
                LocalServerDeploymentConfig(provider=provider.TYPE)
            )
        except ServerDeploymentNotFoundError:
            pass

    raise ServerDeploymentNotFoundError(
        "No local server deployment was found."
    )
get_server_logs(follow: bool = False, tail: Optional[int] = None) -> Generator[str, bool, None]

Retrieve the logs for the local ZenML server.

Parameters:

Name Type Description Default
follow bool

if True, the logs will be streamed as they are written

False
tail Optional[int]

only retrieve the last NUM lines of log output.

None

Returns:

Type Description
None

A generator that can be accessed to get the service logs.

Source code in src/zenml/zen_server/deploy/deployer.py
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def get_server_logs(
    self,
    follow: bool = False,
    tail: Optional[int] = None,
) -> Generator[str, bool, None]:
    """Retrieve the logs for the local ZenML server.

    Args:
        follow: if True, the logs will be streamed as they are written
        tail: only retrieve the last NUM lines of log output.

    Returns:
        A generator that can be accessed to get the service logs.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    server = self.get_server()

    provider_name = server.config.provider.value
    provider = self.get_provider(server.config.provider)

    logger.info(
        f"Fetching logs from the local {provider_name} ZenML server..."
    )
    return provider.get_server_logs(
        server.config, follow=follow, tail=tail
    )
initialize_local_database() -> None

Initialize the local ZenML database.

Source code in src/zenml/zen_server/deploy/deployer.py
88
89
90
91
def initialize_local_database(self) -> None:
    """Initialize the local ZenML database."""
    default_store_cfg = GlobalConfiguration().get_default_store()
    BaseZenStore.create_store(default_store_cfg)
is_connected_to_server() -> bool

Check if the ZenML client is currently connected to the local ZenML server.

Returns:

Type Description
bool

True if the ZenML client is connected to the local ZenML server, False

bool

otherwise.

Source code in src/zenml/zen_server/deploy/deployer.py
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def is_connected_to_server(self) -> bool:
    """Check if the ZenML client is currently connected to the local ZenML server.

    Returns:
        True if the ZenML client is connected to the local ZenML server, False
        otherwise.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    try:
        server = self.get_server()
    except ServerDeploymentNotFoundError:
        return False

    gc = GlobalConfiguration()
    return (
        server.status is not None
        and server.status.url is not None
        and gc.store_configuration.url == server.status.url
    )
register_provider(provider: Type[BaseServerProvider]) -> None classmethod

Register a server provider.

Parameters:

Name Type Description Default
provider Type[BaseServerProvider]

The server provider to register.

required

Raises:

Type Description
TypeError

If a provider with the same type is already registered.

Source code in src/zenml/zen_server/deploy/deployer.py
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@classmethod
def register_provider(cls, provider: Type[BaseServerProvider]) -> None:
    """Register a server provider.

    Args:
        provider: The server provider to register.

    Raises:
        TypeError: If a provider with the same type is already registered.
    """
    if provider.TYPE in cls._providers:
        raise TypeError(
            f"Server provider '{provider.TYPE}' is already registered."
        )
    logger.debug(f"Registering server provider '{provider.TYPE}'.")
    cls._providers[provider.TYPE] = provider()
remove_server(timeout: Optional[int] = None) -> None

Tears down and removes all resources and files associated with the local ZenML server deployment.

Parameters:

Name Type Description Default
timeout Optional[int]

The timeout in seconds to wait until the deployment is successfully torn down. If not supplied, a provider specific default timeout value is used.

None
Source code in src/zenml/zen_server/deploy/deployer.py
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def remove_server(
    self,
    timeout: Optional[int] = None,
) -> None:
    """Tears down and removes all resources and files associated with the local ZenML server deployment.

    Args:
        timeout: The timeout in seconds to wait until the deployment is
            successfully torn down. If not supplied, a provider specific
            default timeout value is used.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    try:
        server = self.get_server()
    except ServerDeploymentNotFoundError:
        return

    provider_name = server.config.provider.value
    provider = self.get_provider(server.config.provider)

    if self.is_connected_to_server():
        try:
            self.disconnect_from_server()
        except Exception as e:
            logger.warning(
                f"Failed to disconnect from the local server: {e}"
            )

    logger.info(f"Tearing down the local {provider_name} ZenML server.")
    provider.remove_server(server.config, timeout=timeout)
update_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None, restart: bool = False) -> LocalServerDeployment

Update an existing local ZenML server deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The new server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the deployment is successful. If not supplied, a default timeout value of 30 seconds is used.

None
restart bool

If True, the existing server deployment will be torn down and a new server will be deployed.

False

Returns:

Type Description
LocalServerDeployment

The updated server deployment.

Source code in src/zenml/zen_server/deploy/deployer.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def update_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
    restart: bool = False,
) -> LocalServerDeployment:
    """Update an existing local ZenML server deployment.

    Args:
        config: The new server deployment configuration.
        timeout: The timeout in seconds to wait until the deployment is
            successful. If not supplied, a default timeout value of 30
            seconds is used.
        restart: If True, the existing server deployment will be torn down
            and a new server will be deployed.

    Returns:
        The updated server deployment.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    existing_server = self.get_server()

    provider = self.get_provider(config.provider)
    if existing_server.config.provider != config.provider or restart:
        existing_provider = self.get_provider(
            existing_server.config.provider
        )

        # Tear down the existing server deployment
        existing_provider.remove_server(
            existing_server.config, timeout=timeout
        )

        # Deploy a new server with the new provider
        return provider.deploy_server(config, timeout=timeout)

    return provider.update_server(config, timeout=timeout)
LocalServerDeployment

Bases: BaseModel

Server deployment.

Attributes:

Name Type Description
config LocalServerDeploymentConfig

The server deployment configuration.

status Optional[LocalServerDeploymentStatus]

The server deployment status.

Attributes
is_running: bool property

Check if the server is running.

Returns:

Type Description
bool

Whether the server is running.

LocalServerDeploymentConfig

Bases: BaseModel

Generic local server deployment configuration.

All local server deployment configurations should inherit from this class and handle extra attributes as provider specific attributes.

Attributes:

Name Type Description
provider ServerProviderType

The server provider type.

Attributes
url: Optional[str] property

Get the configured server URL.

Returns:

Type Description
Optional[str]

The configured server URL.

Modules
base_provider

Base ZenML server provider class.

Classes
BaseServerProvider

Bases: ABC

Base ZenML server provider class.

All ZenML server providers must extend and implement this base class.

Functions
deploy_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None) -> LocalServerDeployment

Deploy a new ZenML server.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The generic server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the deployment is successful. If not supplied, the default timeout value specified by the provider is used.

None

Returns:

Type Description
LocalServerDeployment

The newly created server deployment.

Raises:

Type Description
ServerDeploymentExistsError

If a deployment already exists.

Source code in src/zenml/zen_server/deploy/base_provider.py
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def deploy_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
) -> LocalServerDeployment:
    """Deploy a new ZenML server.

    Args:
        config: The generic server deployment configuration.
        timeout: The timeout in seconds to wait until the deployment is
            successful. If not supplied, the default timeout value specified
            by the provider is used.

    Returns:
        The newly created server deployment.

    Raises:
        ServerDeploymentExistsError: If a deployment already exists.
    """
    try:
        self._get_service()
    except KeyError:
        pass
    else:
        raise ServerDeploymentExistsError(
            f"Local {self.TYPE.value} ZenML server deployment already exists"
        )

    # convert the generic deployment config to a provider specific
    # deployment config
    config = self._convert_config(config)
    service = self._create_service(config, timeout)
    return self._get_deployment(service)
get_server(config: LocalServerDeploymentConfig) -> LocalServerDeployment

Retrieve information about a ZenML server deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The generic server deployment configuration.

required

Returns:

Type Description
LocalServerDeployment

The server deployment.

Raises:

Type Description
ServerDeploymentNotFoundError

If a deployment doesn't exist.

Source code in src/zenml/zen_server/deploy/base_provider.py
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def get_server(
    self,
    config: LocalServerDeploymentConfig,
) -> LocalServerDeployment:
    """Retrieve information about a ZenML server deployment.

    Args:
        config: The generic server deployment configuration.

    Returns:
        The server deployment.

    Raises:
        ServerDeploymentNotFoundError: If a deployment doesn't exist.
    """
    try:
        service = self._get_service()
    except KeyError:
        raise ServerDeploymentNotFoundError(
            f"The local {self.TYPE.value} ZenML server deployment was not "
            f"found"
        )

    return self._get_deployment(service)
get_server_logs(config: LocalServerDeploymentConfig, follow: bool = False, tail: Optional[int] = None) -> Generator[str, bool, None]

Retrieve the logs of a ZenML server.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The generic server deployment configuration.

required
follow bool

if True, the logs will be streamed as they are written

False
tail Optional[int]

only retrieve the last NUM lines of log output.

None

Returns:

Type Description
None

A generator that can be accessed to get the service logs.

Raises:

Type Description
ServerDeploymentNotFoundError

If a deployment doesn't exist.

Source code in src/zenml/zen_server/deploy/base_provider.py
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def get_server_logs(
    self,
    config: LocalServerDeploymentConfig,
    follow: bool = False,
    tail: Optional[int] = None,
) -> Generator[str, bool, None]:
    """Retrieve the logs of a ZenML server.

    Args:
        config: The generic server deployment configuration.
        follow: if True, the logs will be streamed as they are written
        tail: only retrieve the last NUM lines of log output.

    Returns:
        A generator that can be accessed to get the service logs.

    Raises:
        ServerDeploymentNotFoundError: If a deployment doesn't exist.
    """
    try:
        service = self._get_service()
    except KeyError:
        raise ServerDeploymentNotFoundError(
            f"The local {self.TYPE.value} ZenML server deployment was not "
            f"found"
        )

    return service.get_logs(follow=follow, tail=tail)
register_as_provider() -> None classmethod

Register the class as a server provider.

Source code in src/zenml/zen_server/deploy/base_provider.py
55
56
57
58
59
60
@classmethod
def register_as_provider(cls) -> None:
    """Register the class as a server provider."""
    from zenml.zen_server.deploy.deployer import LocalServerDeployer

    LocalServerDeployer.register_provider(cls)
remove_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None) -> None

Tears down and removes all resources and files associated with a ZenML server deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The generic server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the server is removed. If not supplied, the default timeout value specified by the provider is used.

None

Raises:

Type Description
ServerDeploymentNotFoundError

If a deployment doesn't exist.

Source code in src/zenml/zen_server/deploy/base_provider.py
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def remove_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
) -> None:
    """Tears down and removes all resources and files associated with a ZenML server deployment.

    Args:
        config: The generic server deployment configuration.
        timeout: The timeout in seconds to wait until the server is
            removed. If not supplied, the default timeout value specified
            by the provider is used.

    Raises:
        ServerDeploymentNotFoundError: If a deployment doesn't exist.
    """
    try:
        service = self._get_service()
    except KeyError:
        raise ServerDeploymentNotFoundError(
            f"The local {self.TYPE.value} ZenML server deployment was not "
            f"found"
        )

    logger.info(f"Shutting down the local {self.TYPE.value} ZenML server.")
    self._delete_service(service, timeout)
update_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None) -> LocalServerDeployment

Update an existing ZenML server deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The new generic server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the update is successful. If not supplied, the default timeout value specified by the provider is used.

None

Returns:

Type Description
LocalServerDeployment

The updated server deployment.

Raises:

Type Description
ServerDeploymentNotFoundError

If a deployment doesn't exist.

Source code in src/zenml/zen_server/deploy/base_provider.py
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def update_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
) -> LocalServerDeployment:
    """Update an existing ZenML server deployment.

    Args:
        config: The new generic server deployment configuration.
        timeout: The timeout in seconds to wait until the update is
            successful. If not supplied, the default timeout value specified
            by the provider is used.

    Returns:
        The updated server deployment.

    Raises:
        ServerDeploymentNotFoundError: If a deployment doesn't exist.
    """
    try:
        service = self._get_service()
    except KeyError:
        raise ServerDeploymentNotFoundError(
            f"The local {self.TYPE.value} ZenML server deployment was not "
            f"found"
        )

    # convert the generic deployment config to a provider specific
    # deployment config
    config = self._convert_config(config)
    old_config = self._get_deployment_config(service)

    if old_config == config:
        logger.info(
            f"The local {self.TYPE.value} ZenML server is already "
            "configured with the same parameters."
        )
        service = self._start_service(service, timeout)
    else:
        logger.info(f"Updating the local {self.TYPE.value} ZenML server.")
        service = self._update_service(service, config, timeout)

    return self._get_deployment(service)
Functions
daemon

ZenML Server Local Daemon Deployment.

Classes
DaemonServerProvider

Bases: BaseServerProvider

Daemon ZenML server provider.

Functions
check_local_server_dependencies() -> None staticmethod

Check if local server dependencies are installed.

Raises:

Type Description
RuntimeError

If the dependencies are not installed.

Source code in src/zenml/zen_server/deploy/daemon/daemon_provider.py
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
@staticmethod
def check_local_server_dependencies() -> None:
    """Check if local server dependencies are installed.

    Raises:
        RuntimeError: If the dependencies are not installed.
    """
    try:
        # Make sure the ZenML Server dependencies are installed
        import fastapi  # noqa
        import jwt  # noqa
        import multipart  # noqa
        import uvicorn  # noqa
    except ImportError:
        # Unable to import the ZenML Server dependencies.
        raise RuntimeError(
            "The local daemon ZenML server provider is unavailable because the "
            "ZenML server requirements seems to be unavailable on your machine. "
            "This is probably because ZenML was installed without the optional "
            "ZenML Server dependencies. To install the missing dependencies "
            f'run `pip install "zenml[server]=={__version__}"`.'
        )
Modules
daemon_provider

Zen Server daemon provider implementation.

Classes
DaemonServerProvider

Bases: BaseServerProvider

Daemon ZenML server provider.

Functions
check_local_server_dependencies() -> None staticmethod

Check if local server dependencies are installed.

Raises:

Type Description
RuntimeError

If the dependencies are not installed.

Source code in src/zenml/zen_server/deploy/daemon/daemon_provider.py
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
@staticmethod
def check_local_server_dependencies() -> None:
    """Check if local server dependencies are installed.

    Raises:
        RuntimeError: If the dependencies are not installed.
    """
    try:
        # Make sure the ZenML Server dependencies are installed
        import fastapi  # noqa
        import jwt  # noqa
        import multipart  # noqa
        import uvicorn  # noqa
    except ImportError:
        # Unable to import the ZenML Server dependencies.
        raise RuntimeError(
            "The local daemon ZenML server provider is unavailable because the "
            "ZenML server requirements seems to be unavailable on your machine. "
            "This is probably because ZenML was installed without the optional "
            "ZenML Server dependencies. To install the missing dependencies "
            f'run `pip install "zenml[server]=={__version__}"`.'
        )
Functions
daemon_zen_server

Local daemon ZenML server deployment service implementation.

Classes
DaemonServerDeploymentConfig

Bases: LocalServerDeploymentConfig

Daemon server deployment configuration.

Attributes:

Name Type Description
port int

The TCP port number where the server is accepting connections.

address int

The IP address where the server is reachable.

blocking bool

Run the server in blocking mode instead of using a daemon process.

Attributes
url: Optional[str] property

Get the configured server URL.

Returns:

Type Description
Optional[str]

The configured server URL.

DaemonZenServer(**attrs: Any)

Bases: LocalDaemonService

Service daemon that can be used to start a local daemon ZenML server.

Attributes:

Name Type Description
config DaemonZenServerConfig

service configuration

endpoint LocalDaemonServiceEndpoint

optional service endpoint

Source code in src/zenml/services/service.py
188
189
190
191
192
193
194
195
196
197
198
def __init__(
    self,
    **attrs: Any,
) -> None:
    """Initialize the service instance.

    Args:
        **attrs: keyword arguments.
    """
    super().__init__(**attrs)
    self.config.name = self.config.name or self.__class__.__name__
Functions
config_path() -> str classmethod

Path to the directory where the local daemon ZenML server files are located.

Returns:

Type Description
str

Path to the local daemon ZenML server runtime directory.

Source code in src/zenml/zen_server/deploy/daemon/daemon_zen_server.py
114
115
116
117
118
119
120
121
122
123
124
125
@classmethod
def config_path(cls) -> str:
    """Path to the directory where the local daemon ZenML server files are located.

    Returns:
        Path to the local daemon ZenML server runtime directory.
    """
    return os.path.join(
        get_global_config_directory(),
        "zen_server",
        "daemon",
    )
get_service() -> Optional[DaemonZenServer] classmethod

Load and return the local daemon ZenML server service, if present.

Returns:

Type Description
Optional[DaemonZenServer]

The local daemon ZenML server service or None, if the local server

Optional[DaemonZenServer]

deployment is not found.

Source code in src/zenml/zen_server/deploy/daemon/daemon_zen_server.py
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
@classmethod
def get_service(cls) -> Optional["DaemonZenServer"]:
    """Load and return the local daemon ZenML server service, if present.

    Returns:
        The local daemon ZenML server service or None, if the local server
        deployment is not found.
    """
    config_filename = os.path.join(cls.config_path(), "service.json")
    try:
        with open(config_filename, "r") as f:
            return cast(
                "DaemonZenServer", DaemonZenServer.from_json(f.read())
            )
    except FileNotFoundError:
        return None
provision() -> None

Provision the service.

Source code in src/zenml/zen_server/deploy/daemon/daemon_zen_server.py
184
185
186
def provision(self) -> None:
    """Provision the service."""
    super().provision()
run() -> None

Run the ZenML Server.

Raises:

Type Description
ValueError

if started with a global configuration that connects to another ZenML server.

Source code in src/zenml/zen_server/deploy/daemon/daemon_zen_server.py
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def run(self) -> None:
    """Run the ZenML Server.

    Raises:
        ValueError: if started with a global configuration that connects to
            another ZenML server.
    """
    import uvicorn

    gc = GlobalConfiguration()
    if gc.store_configuration.type == StoreType.REST:
        raise ValueError(
            "The ZenML server cannot be started with REST store type."
        )
    logger.info(
        "Starting ZenML Server as blocking "
        "process... press CTRL+C once to stop it."
    )

    self.endpoint.prepare_for_start()

    try:
        uvicorn.run(
            ZEN_SERVER_ENTRYPOINT,
            host=self.endpoint.config.ip_address,
            port=self.endpoint.config.port or 8000,
            log_level="info",
            server_header=False,
        )
    except KeyboardInterrupt:
        logger.info("ZenML Server stopped. Resuming normal execution.")
start(timeout: int = 0) -> None

Start the service and optionally wait for it to become active.

Parameters:

Name Type Description Default
timeout int

amount of time to wait for the service to become active. If set to 0, the method will return immediately after checking the service status.

0
Source code in src/zenml/zen_server/deploy/daemon/daemon_zen_server.py
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
def start(self, timeout: int = 0) -> None:
    """Start the service and optionally wait for it to become active.

    Args:
        timeout: amount of time to wait for the service to become active.
            If set to 0, the method will return immediately after checking
            the service status.
    """
    if not self.config.blocking:
        super().start(timeout)
    else:
        gc = GlobalConfiguration()

        # In the blocking mode, we need to temporarily set the environment
        # variables for the running process to make it look like the server
        # is running in a separate environment (i.e. using a different
        # global configuration path). This is necessary to avoid polluting
        # the client environment with the server's configuration.
        local_stores_path = GlobalConfiguration().local_stores_path
        GlobalConfiguration._reset_instance()
        Client._reset_instance()
        original_config_path = os.environ.get(ENV_ZENML_CONFIG_PATH)
        os.environ[ENV_ZENML_SERVER] = "true"
        os.environ[ENV_ZENML_CONFIG_PATH] = self._global_config_path
        os.environ[ENV_ZENML_ANALYTICS_OPT_IN] = str(gc.analytics_opt_in)
        os.environ[ENV_ZENML_USER_ID] = str(gc.user_id)
        # Set the local stores path to the same path used by the client.
        # This ensures that the server's default store configuration is
        # initialized to point at the same local SQLite database as the
        # client.
        os.environ[ENV_ZENML_LOCAL_STORES_PATH] = local_stores_path
        os.environ[ENV_ZENML_SERVER_AUTH_SCHEME] = AuthScheme.NO_AUTH.value
        try:
            self.run()
        finally:
            # Restore the original client environment variables
            del os.environ[ENV_ZENML_SERVER]
            if original_config_path:
                os.environ[ENV_ZENML_CONFIG_PATH] = original_config_path
            else:
                del os.environ[ENV_ZENML_CONFIG_PATH]
            del os.environ[ENV_ZENML_LOCAL_STORES_PATH]
            del os.environ[ENV_ZENML_SERVER_AUTH_SCHEME]
            GlobalConfiguration._reset_instance()
            Client._reset_instance()
DaemonZenServerConfig(**data: Any)

Bases: LocalDaemonServiceConfig

Local daemon Zen server configuration.

Attributes:

Name Type Description
server DaemonServerDeploymentConfig

The deployment configuration.

Source code in src/zenml/services/service.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def __init__(self, **data: Any):
    """Initialize the service configuration.

    Args:
        **data: keyword arguments.

    Raises:
        ValueError: if neither 'name' nor 'model_name' is set.
    """
    super().__init__(**data)
    if self.name or self.model_name:
        self.service_name = data.get(
            "service_name",
            f"{ZENM_ENDPOINT_PREFIX}{self.name or self.model_name}",
        )
    else:
        raise ValueError("Either 'name' or 'model_name' must be set.")
Functions
deployer

ZenML server deployer singleton implementation.

Classes
LocalServerDeployer

Local server deployer singleton.

This class is responsible for managing the various server provider implementations and for directing server deployment lifecycle requests to the responsible provider. It acts as a facade built on top of the various server providers.

Functions
connect_to_server() -> None

Connect to the local ZenML server instance.

Raises:

Type Description
ServerDeploymentError

If the local ZenML server is not running or is unreachable.

Source code in src/zenml/zen_server/deploy/deployer.py
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def connect_to_server(
    self,
) -> None:
    """Connect to the local ZenML server instance.

    Raises:
        ServerDeploymentError: If the local ZenML server is not running or
            is unreachable.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    server = self.get_server()
    provider_name = server.config.provider.value

    gc = GlobalConfiguration()
    if not server.status or not server.status.url:
        raise ServerDeploymentError(
            f"The local {provider_name} ZenML server is not currently "
            "running or is unreachable."
        )

    store_config = RestZenStoreConfiguration(
        url=server.status.url,
    )

    if gc.store_configuration == store_config:
        logger.info(
            "Your client is already connected to the local "
            f"{provider_name} ZenML server."
        )
        return

    logger.info(
        f"Connecting to the local {provider_name} ZenML server "
        f"({store_config.url})."
    )

    gc.set_store(store_config)

    logger.info(
        f"Connected to the local {provider_name} ZenML server "
        f"({store_config.url})."
    )
deploy_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None, restart: bool = False) -> LocalServerDeployment

Deploy the local ZenML server or update the existing deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the deployment is successful. If not supplied, the default timeout value specified by the provider is used.

None
restart bool

If True, the existing server deployment will be torn down and a new server will be deployed.

False

Returns:

Type Description
LocalServerDeployment

The local server deployment.

Source code in src/zenml/zen_server/deploy/deployer.py
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def deploy_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
    restart: bool = False,
) -> LocalServerDeployment:
    """Deploy the local ZenML server or update the existing deployment.

    Args:
        config: The server deployment configuration.
        timeout: The timeout in seconds to wait until the deployment is
            successful. If not supplied, the default timeout value specified
            by the provider is used.
        restart: If True, the existing server deployment will be torn down
            and a new server will be deployed.

    Returns:
        The local server deployment.
    """
    # Ensure that the local database is always initialized before any local
    # server is deployed or updated.
    self.initialize_local_database()

    try:
        self.get_server()
    except ServerDeploymentNotFoundError:
        pass
    else:
        return self.update_server(
            config=config, timeout=timeout, restart=restart
        )

    provider_name = config.provider.value
    provider = self.get_provider(config.provider)

    logger.info(f"Deploying a local {provider_name} ZenML server.")
    return provider.deploy_server(config, timeout=timeout)
disconnect_from_server() -> None

Disconnect from the ZenML server instance.

Source code in src/zenml/zen_server/deploy/deployer.py
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def disconnect_from_server(
    self,
) -> None:
    """Disconnect from the ZenML server instance."""
    gc = GlobalConfiguration()
    store_cfg = gc.store_configuration

    if store_cfg.type != StoreType.REST:
        logger.info(
            "Your client is not currently connected to a ZenML server."
        )
        return

    logger.info(
        f"Disconnecting from the local ({store_cfg.url}) ZenML server."
    )

    gc.set_default_store()

    logger.info("Disconnected from the local ZenML server.")
get_provider(provider_type: ServerProviderType) -> BaseServerProvider classmethod

Get the server provider associated with a provider type.

Parameters:

Name Type Description Default
provider_type ServerProviderType

The server provider type.

required

Returns:

Type Description
BaseServerProvider

The server provider associated with the provider type.

Raises:

Type Description
ServerProviderNotFoundError

If no provider is registered for the given provider type.

Source code in src/zenml/zen_server/deploy/deployer.py
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@classmethod
def get_provider(
    cls, provider_type: ServerProviderType
) -> BaseServerProvider:
    """Get the server provider associated with a provider type.

    Args:
        provider_type: The server provider type.

    Returns:
        The server provider associated with the provider type.

    Raises:
        ServerProviderNotFoundError: If no provider is registered for the
            given provider type.
    """
    if provider_type not in cls._providers:
        raise ServerProviderNotFoundError(
            f"Server provider '{provider_type}' is not registered."
        )
    return cls._providers[provider_type]
get_server() -> LocalServerDeployment

Get the local server deployment.

Returns:

Type Description
LocalServerDeployment

The local server deployment.

Raises:

Type Description
ServerDeploymentNotFoundError

If no local server deployment is found.

Source code in src/zenml/zen_server/deploy/deployer.py
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
def get_server(
    self,
) -> LocalServerDeployment:
    """Get the local server deployment.

    Returns:
        The local server deployment.

    Raises:
        ServerDeploymentNotFoundError: If no local server deployment is
            found.
    """
    for provider in self._providers.values():
        try:
            return provider.get_server(
                LocalServerDeploymentConfig(provider=provider.TYPE)
            )
        except ServerDeploymentNotFoundError:
            pass

    raise ServerDeploymentNotFoundError(
        "No local server deployment was found."
    )
get_server_logs(follow: bool = False, tail: Optional[int] = None) -> Generator[str, bool, None]

Retrieve the logs for the local ZenML server.

Parameters:

Name Type Description Default
follow bool

if True, the logs will be streamed as they are written

False
tail Optional[int]

only retrieve the last NUM lines of log output.

None

Returns:

Type Description
None

A generator that can be accessed to get the service logs.

Source code in src/zenml/zen_server/deploy/deployer.py
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def get_server_logs(
    self,
    follow: bool = False,
    tail: Optional[int] = None,
) -> Generator[str, bool, None]:
    """Retrieve the logs for the local ZenML server.

    Args:
        follow: if True, the logs will be streamed as they are written
        tail: only retrieve the last NUM lines of log output.

    Returns:
        A generator that can be accessed to get the service logs.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    server = self.get_server()

    provider_name = server.config.provider.value
    provider = self.get_provider(server.config.provider)

    logger.info(
        f"Fetching logs from the local {provider_name} ZenML server..."
    )
    return provider.get_server_logs(
        server.config, follow=follow, tail=tail
    )
initialize_local_database() -> None

Initialize the local ZenML database.

Source code in src/zenml/zen_server/deploy/deployer.py
88
89
90
91
def initialize_local_database(self) -> None:
    """Initialize the local ZenML database."""
    default_store_cfg = GlobalConfiguration().get_default_store()
    BaseZenStore.create_store(default_store_cfg)
is_connected_to_server() -> bool

Check if the ZenML client is currently connected to the local ZenML server.

Returns:

Type Description
bool

True if the ZenML client is connected to the local ZenML server, False

bool

otherwise.

Source code in src/zenml/zen_server/deploy/deployer.py
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def is_connected_to_server(self) -> bool:
    """Check if the ZenML client is currently connected to the local ZenML server.

    Returns:
        True if the ZenML client is connected to the local ZenML server, False
        otherwise.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    try:
        server = self.get_server()
    except ServerDeploymentNotFoundError:
        return False

    gc = GlobalConfiguration()
    return (
        server.status is not None
        and server.status.url is not None
        and gc.store_configuration.url == server.status.url
    )
register_provider(provider: Type[BaseServerProvider]) -> None classmethod

Register a server provider.

Parameters:

Name Type Description Default
provider Type[BaseServerProvider]

The server provider to register.

required

Raises:

Type Description
TypeError

If a provider with the same type is already registered.

Source code in src/zenml/zen_server/deploy/deployer.py
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@classmethod
def register_provider(cls, provider: Type[BaseServerProvider]) -> None:
    """Register a server provider.

    Args:
        provider: The server provider to register.

    Raises:
        TypeError: If a provider with the same type is already registered.
    """
    if provider.TYPE in cls._providers:
        raise TypeError(
            f"Server provider '{provider.TYPE}' is already registered."
        )
    logger.debug(f"Registering server provider '{provider.TYPE}'.")
    cls._providers[provider.TYPE] = provider()
remove_server(timeout: Optional[int] = None) -> None

Tears down and removes all resources and files associated with the local ZenML server deployment.

Parameters:

Name Type Description Default
timeout Optional[int]

The timeout in seconds to wait until the deployment is successfully torn down. If not supplied, a provider specific default timeout value is used.

None
Source code in src/zenml/zen_server/deploy/deployer.py
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def remove_server(
    self,
    timeout: Optional[int] = None,
) -> None:
    """Tears down and removes all resources and files associated with the local ZenML server deployment.

    Args:
        timeout: The timeout in seconds to wait until the deployment is
            successfully torn down. If not supplied, a provider specific
            default timeout value is used.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    try:
        server = self.get_server()
    except ServerDeploymentNotFoundError:
        return

    provider_name = server.config.provider.value
    provider = self.get_provider(server.config.provider)

    if self.is_connected_to_server():
        try:
            self.disconnect_from_server()
        except Exception as e:
            logger.warning(
                f"Failed to disconnect from the local server: {e}"
            )

    logger.info(f"Tearing down the local {provider_name} ZenML server.")
    provider.remove_server(server.config, timeout=timeout)
update_server(config: LocalServerDeploymentConfig, timeout: Optional[int] = None, restart: bool = False) -> LocalServerDeployment

Update an existing local ZenML server deployment.

Parameters:

Name Type Description Default
config LocalServerDeploymentConfig

The new server deployment configuration.

required
timeout Optional[int]

The timeout in seconds to wait until the deployment is successful. If not supplied, a default timeout value of 30 seconds is used.

None
restart bool

If True, the existing server deployment will be torn down and a new server will be deployed.

False

Returns:

Type Description
LocalServerDeployment

The updated server deployment.

Source code in src/zenml/zen_server/deploy/deployer.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def update_server(
    self,
    config: LocalServerDeploymentConfig,
    timeout: Optional[int] = None,
    restart: bool = False,
) -> LocalServerDeployment:
    """Update an existing local ZenML server deployment.

    Args:
        config: The new server deployment configuration.
        timeout: The timeout in seconds to wait until the deployment is
            successful. If not supplied, a default timeout value of 30
            seconds is used.
        restart: If True, the existing server deployment will be torn down
            and a new server will be deployed.

    Returns:
        The updated server deployment.
    """
    # this will also raise ServerDeploymentNotFoundError if the server
    # does not exist
    existing_server = self.get_server()

    provider = self.get_provider(config.provider)
    if existing_server.config.provider != config.provider or restart:
        existing_provider = self.get_provider(
            existing_server.config.provider
        )

        # Tear down the existing server deployment
        existing_provider.remove_server(
            existing_server.config, timeout=timeout
        )

        # Deploy a new server with the new provider
        return provider.deploy_server(config, timeout=timeout)

    return provider.update_server(config, timeout=timeout)
Functions
deployment

Zen Server deployment definitions.

Classes
LocalServerDeployment

Bases: BaseModel

Server deployment.

Attributes:

Name Type Description
config LocalServerDeploymentConfig

The server deployment configuration.

status Optional[LocalServerDeploymentStatus]

The server deployment status.

Attributes
is_running: bool property

Check if the server is running.

Returns:

Type Description
bool

Whether the server is running.

LocalServerDeploymentConfig

Bases: BaseModel

Generic local server deployment configuration.

All local server deployment configurations should inherit from this class and handle extra attributes as provider specific attributes.

Attributes:

Name Type Description
provider ServerProviderType

The server provider type.

Attributes
url: Optional[str] property

Get the configured server URL.

Returns:

Type Description
Optional[str]

The configured server URL.

LocalServerDeploymentStatus

Bases: BaseModel

Local server deployment status.

Ideally this should convey the following information:

  • whether the server's deployment is managed by this client (i.e. if the server was deployed with zenml login --local)
  • for a managed deployment, the status of the deployment/tear-down, e.g. not deployed, deploying, running, deleting, deployment timeout/error, tear-down timeout/error etc.
  • for an unmanaged deployment, the operational status (i.e. whether the server is reachable)
  • the URL of the server

Attributes:

Name Type Description
status ServiceState

The status of the server deployment.

status_message Optional[str]

A message describing the last status.

connected bool

Whether the client is currently connected to this server.

url Optional[str]

The URL of the server.

docker

ZenML Server Docker Deployment.

Classes
DockerServerProvider

Bases: BaseServerProvider

Docker ZenML server provider.

Modules
docker_provider

Zen Server docker deployer implementation.

Classes
DockerServerProvider

Bases: BaseServerProvider

Docker ZenML server provider.

Functions
docker_zen_server

Service implementation for the ZenML docker server deployment.

Classes
DockerServerDeploymentConfig

Bases: LocalServerDeploymentConfig

Docker server deployment configuration.

Attributes:

Name Type Description
port int

The TCP port number where the server is accepting connections.

image str

The Docker image to use for the server.

Attributes
url: Optional[str] property

Get the configured server URL.

Returns:

Type Description
Optional[str]

The configured server URL.

DockerZenServer(**attrs: Any)

Bases: ContainerService

Service that can be used to start a docker ZenServer.

Attributes:

Name Type Description
config DockerZenServerConfig

service configuration

endpoint ContainerServiceEndpoint

service endpoint

Source code in src/zenml/services/service.py
188
189
190
191
192
193
194
195
196
197
198
def __init__(
    self,
    **attrs: Any,
) -> None:
    """Initialize the service instance.

    Args:
        **attrs: keyword arguments.
    """
    super().__init__(**attrs)
    self.config.name = self.config.name or self.__class__.__name__
Functions
config_path() -> str classmethod

Path to the directory where the docker ZenML server files are located.

Returns:

Type Description
str

Path to the docker ZenML server runtime directory.

Source code in src/zenml/zen_server/deploy/docker/docker_zen_server.py
116
117
118
119
120
121
122
123
124
125
126
127
@classmethod
def config_path(cls) -> str:
    """Path to the directory where the docker ZenML server files are located.

    Returns:
        Path to the docker ZenML server runtime directory.
    """
    return os.path.join(
        get_global_config_directory(),
        "zen_server",
        "docker",
    )
get_service() -> Optional[DockerZenServer] classmethod

Load and return the docker ZenML server service, if present.

Returns:

Type Description
Optional[DockerZenServer]

The docker ZenML server service or None, if the docker server

Optional[DockerZenServer]

deployment is not found.

Source code in src/zenml/zen_server/deploy/docker/docker_zen_server.py
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@classmethod
def get_service(cls) -> Optional["DockerZenServer"]:
    """Load and return the docker ZenML server service, if present.

    Returns:
        The docker ZenML server service or None, if the docker server
        deployment is not found.
    """
    config_filename = os.path.join(cls.config_path(), "service.json")
    try:
        with open(config_filename, "r") as f:
            return cast(
                "DockerZenServer", DockerZenServer.from_json(f.read())
            )
    except FileNotFoundError:
        return None
provision() -> None

Provision the service.

Source code in src/zenml/zen_server/deploy/docker/docker_zen_server.py
194
195
196
def provision(self) -> None:
    """Provision the service."""
    super().provision()
run() -> None

Run the ZenML Server.

Raises:

Type Description
ValueError

if started with a global configuration that connects to another ZenML server.

Source code in src/zenml/zen_server/deploy/docker/docker_zen_server.py
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def run(self) -> None:
    """Run the ZenML Server.

    Raises:
        ValueError: if started with a global configuration that connects to
            another ZenML server.
    """
    import uvicorn

    gc = GlobalConfiguration()
    if gc.store_configuration.type == StoreType.REST:
        raise ValueError(
            "The ZenML server cannot be started with REST store type."
        )
    logger.info(
        "Starting ZenML Server as blocking "
        "process... press CTRL+C once to stop it."
    )

    self.endpoint.prepare_for_start()

    try:
        uvicorn.run(
            ZEN_SERVER_ENTRYPOINT,
            host="0.0.0.0",  # nosec
            port=self.endpoint.config.port or 8000,
            log_level="info",
            server_header=False,
        )
    except KeyboardInterrupt:
        logger.info("ZenML Server stopped. Resuming normal execution.")
DockerZenServerConfig(**data: Any)

Bases: ContainerServiceConfig

Docker Zen server configuration.

Attributes:

Name Type Description
server DockerServerDeploymentConfig

The deployment configuration.

Source code in src/zenml/services/service.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def __init__(self, **data: Any):
    """Initialize the service configuration.

    Args:
        **data: keyword arguments.

    Raises:
        ValueError: if neither 'name' nor 'model_name' is set.
    """
    super().__init__(**data)
    if self.name or self.model_name:
        self.service_name = data.get(
            "service_name",
            f"{ZENM_ENDPOINT_PREFIX}{self.name or self.model_name}",
        )
    else:
        raise ValueError("Either 'name' or 'model_name' must be set.")
Functions
exceptions

ZenML server deployment exceptions.

Classes
ServerDeploymentConfigurationError(message: Optional[str] = None, url: Optional[str] = None)

Bases: ServerDeploymentError

Raised when there is a ZenML server deployment configuration error .

Source code in src/zenml/exceptions.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    message: Optional[str] = None,
    url: Optional[str] = None,
):
    """The BaseException used to format messages displayed to the user.

    Args:
        message: Message with details of exception. This message
                 will be appended with another message directing user to
                 `url` for more information. If `None`, then default
                 Exception behavior is used.
        url: URL to point to in exception message. If `None`, then no url
             is appended.
    """
    if message and url:
        message += f" For more information, visit {url}."
    super().__init__(message)
ServerDeploymentError(message: Optional[str] = None, url: Optional[str] = None)

Bases: ZenMLBaseException

Base exception class for all ZenML server deployment related errors.

Source code in src/zenml/exceptions.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    message: Optional[str] = None,
    url: Optional[str] = None,
):
    """The BaseException used to format messages displayed to the user.

    Args:
        message: Message with details of exception. This message
                 will be appended with another message directing user to
                 `url` for more information. If `None`, then default
                 Exception behavior is used.
        url: URL to point to in exception message. If `None`, then no url
             is appended.
    """
    if message and url:
        message += f" For more information, visit {url}."
    super().__init__(message)
ServerDeploymentExistsError(message: Optional[str] = None, url: Optional[str] = None)

Bases: ServerDeploymentError

Raised when trying to deploy a new ZenML server with the same name.

Source code in src/zenml/exceptions.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    message: Optional[str] = None,
    url: Optional[str] = None,
):
    """The BaseException used to format messages displayed to the user.

    Args:
        message: Message with details of exception. This message
                 will be appended with another message directing user to
                 `url` for more information. If `None`, then default
                 Exception behavior is used.
        url: URL to point to in exception message. If `None`, then no url
             is appended.
    """
    if message and url:
        message += f" For more information, visit {url}."
    super().__init__(message)
ServerDeploymentNotFoundError(message: Optional[str] = None, url: Optional[str] = None)

Bases: ServerDeploymentError

Raised when trying to fetch a ZenML server deployment that doesn't exist.

Source code in src/zenml/exceptions.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    message: Optional[str] = None,
    url: Optional[str] = None,
):
    """The BaseException used to format messages displayed to the user.

    Args:
        message: Message with details of exception. This message
                 will be appended with another message directing user to
                 `url` for more information. If `None`, then default
                 Exception behavior is used.
        url: URL to point to in exception message. If `None`, then no url
             is appended.
    """
    if message and url:
        message += f" For more information, visit {url}."
    super().__init__(message)
ServerProviderNotFoundError(message: Optional[str] = None, url: Optional[str] = None)

Bases: ServerDeploymentError

Raised when using a ZenML server provider that doesn't exist.

Source code in src/zenml/exceptions.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    message: Optional[str] = None,
    url: Optional[str] = None,
):
    """The BaseException used to format messages displayed to the user.

    Args:
        message: Message with details of exception. This message
                 will be appended with another message directing user to
                 `url` for more information. If `None`, then default
                 Exception behavior is used.
        url: URL to point to in exception message. If `None`, then no url
             is appended.
    """
    if message and url:
        message += f" For more information, visit {url}."
    super().__init__(message)

exceptions

REST API exception handling.

Classes
ErrorModel

Bases: BaseModel

Base class for error responses.

Functions
error_detail(error: Exception, exception_type: Optional[Type[Exception]] = None) -> List[str]

Convert an Exception to API representation.

Parameters:

Name Type Description Default
error Exception

Exception to convert.

required
exception_type Optional[Type[Exception]]

Exception type to use in the error response instead of the type of the supplied exception. This is useful when the raised exception is a subclass of an exception type that is properly handled by the REST API.

None

Returns:

Type Description
List[str]

List of strings representing the error.

Source code in src/zenml/zen_server/exceptions.py
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def error_detail(
    error: Exception, exception_type: Optional[Type[Exception]] = None
) -> List[str]:
    """Convert an Exception to API representation.

    Args:
        error: Exception to convert.
        exception_type: Exception type to use in the error response instead of
            the type of the supplied exception. This is useful when the raised
            exception is a subclass of an exception type that is properly
            handled by the REST API.

    Returns:
        List of strings representing the error.
    """
    class_name = (
        exception_type.__name__ if exception_type else type(error).__name__
    )
    return [class_name, str(error)]
exception_from_response(response: requests.Response) -> Optional[Exception]

Convert an error HTTP response to an exception.

Uses the REST_API_EXCEPTIONS list to determine the appropriate exception class to use based on the response status code and the exception class name embedded in the response body.

The last entry in the list of exceptions associated with a status code is used as a fallback if the exception class name in the response body is not found in the list.

Parameters:

Name Type Description Default
response Response

HTTP error response to convert.

required

Returns:

Type Description
Optional[Exception]

Exception with the appropriate type and arguments, or None if the

Optional[Exception]

response does not contain an error or the response cannot be unpacked

Optional[Exception]

into an exception.

Source code in src/zenml/zen_server/exceptions.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def exception_from_response(
    response: requests.Response,
) -> Optional[Exception]:
    """Convert an error HTTP response to an exception.

    Uses the REST_API_EXCEPTIONS list to determine the appropriate exception
    class to use based on the response status code and the exception class name
    embedded in the response body.

    The last entry in the list of exceptions associated with a status code is
    used as a fallback if the exception class name in the response body is not
    found in the list.

    Args:
        response: HTTP error response to convert.

    Returns:
        Exception with the appropriate type and arguments, or None if the
        response does not contain an error or the response cannot be unpacked
        into an exception.
    """

    def unpack_exc() -> Tuple[Optional[str], str]:
        """Unpack the response body into an exception name and message.

        Returns:
            Tuple of exception name and message.
        """
        try:
            response_json = response.json()
        except requests.exceptions.JSONDecodeError:
            return None, response.text

        if isinstance(response_json, dict):
            detail = response_json.get("detail", response.text)
        else:
            detail = response_json

        # The detail can also be a single string
        if isinstance(detail, str):
            return None, detail

        # The detail should be a list of strings encoding the exception
        # class name and the exception message
        if not isinstance(detail, list):
            return None, response.text

        # First detail item is the exception class name
        if len(detail) < 1 or not isinstance(detail[0], str):
            return None, response.text

        # Remaining detail items are the exception arguments
        message = ": ".join([str(arg) for arg in detail[1:]])
        return detail[0], message

    exc_name, exc_msg = unpack_exc()
    default_exc: Optional[Type[Exception]] = None

    for exception, status_code in REST_API_EXCEPTIONS:
        if response.status_code != status_code:
            continue
        default_exc = exception
        if exc_name == exception.__name__:
            # An entry was found that is an exact match for both the status
            # code and the exception class name.
            break
    else:
        # The exception class name extracted from the response body was not
        # found in the list of exceptions associated with the status code, so
        # use the last entry as a fallback.
        if default_exc is None:
            return None

        exception = default_exc

    # There is one special case where we want to return a specific exception:
    # 401 Unauthorized exceptions thrown directly by FastAPI in the course of
    # authentication are interpreted as AuthorizationException, but we want to
    # return CredentialsNotValid instead.
    if response.status_code == 401:
        if not isinstance(exception(), CredentialsNotValid):
            if response.headers.get("WWW-Authenticate"):
                return CredentialsNotValid(exc_msg)

    return exception(exc_msg)
http_exception_from_error(error: Exception) -> HTTPException

Convert an Exception to a HTTP error response.

Uses the REST_API_EXCEPTIONS list to determine the appropriate status code associated with the exception type. The exception class name and arguments are embedded in the HTTP error response body.

The lookup uses the first occurrence of the exception type in the list. If the exception type is not found in the list, the lookup uses isinstance to determine the most specific exception type corresponding to the supplied exception. This allows users to call this method with exception types that are not directly listed in the REST_API_EXCEPTIONS list.

Parameters:

Name Type Description Default
error Exception

Exception to convert.

required

Returns:

Type Description
HTTPException

HTTPException with the appropriate status code and error detail.

Source code in src/zenml/zen_server/exceptions.py
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def http_exception_from_error(error: Exception) -> "HTTPException":
    """Convert an Exception to a HTTP error response.

    Uses the REST_API_EXCEPTIONS list to determine the appropriate status code
    associated with the exception type. The exception class name and arguments
    are embedded in the HTTP error response body.

    The lookup uses the first occurrence of the exception type in the list. If
    the exception type is not found in the list, the lookup uses `isinstance`
    to determine the most specific exception type corresponding to the supplied
    exception. This allows users to call this method with exception types that
    are not directly listed in the REST_API_EXCEPTIONS list.

    Args:
        error: Exception to convert.

    Returns:
        HTTPException with the appropriate status code and error detail.
    """
    from fastapi import HTTPException

    status_code = 0
    matching_exception_type: Optional[Type[Exception]] = None

    for exception_type, exc_status_code in REST_API_EXCEPTIONS:
        if error.__class__ is exception_type:
            # Found an exact match
            matching_exception_type = exception_type
            status_code = exc_status_code
            break
        if isinstance(error, exception_type):
            # Found a matching exception
            if not matching_exception_type:
                # This is the first matching exception, so keep it
                matching_exception_type = exception_type
                status_code = exc_status_code
                continue

            # This is not the first matching exception, so check if it is more
            # specific than the previous matching exception
            if issubclass(
                exception_type,
                matching_exception_type,
            ):
                matching_exception_type = exception_type
                status_code = exc_status_code

    # When the matching exception is not found in the list, a 500 Internal
    # Server Error is returned
    status_code = status_code or 500
    matching_exception_type = matching_exception_type or RuntimeError

    return HTTPException(
        status_code=status_code,
        detail=error_detail(error, matching_exception_type),
    )

feature_gate

Modules
endpoint_utils

All endpoint utils for the feature gate implementations.

Classes Functions
check_entitlement(resource_type: ResourceType) -> None

Queries the feature gate to see if the operation falls within the Pro workspaces entitlements.

Raises an exception if the user is not entitled to create an instance of the resource. Otherwise, simply returns.

Parameters:

Name Type Description Default
resource_type ResourceType

The type of resource to check for.

required
Source code in src/zenml/zen_server/feature_gate/endpoint_utils.py
22
23
24
25
26
27
28
29
30
31
32
33
def check_entitlement(resource_type: ResourceType) -> None:
    """Queries the feature gate to see if the operation falls within the Pro workspaces entitlements.

    Raises an exception if the user is not entitled to create an instance of the
    resource. Otherwise, simply returns.

    Args:
        resource_type: The type of resource to check for.
    """
    if not server_config().feature_gate_enabled:
        return
    return feature_gate().check_entitlement(resource=resource_type)
report_decrement(resource_type: ResourceType, resource_id: UUID) -> None

Reports the deletion/deactivation of a feature/resource.

Parameters:

Name Type Description Default
resource_type ResourceType

The type of resource to report a decrement in count for.

required
resource_id UUID

ID of the resource that was deleted.

required
Source code in src/zenml/zen_server/feature_gate/endpoint_utils.py
50
51
52
53
54
55
56
57
58
59
60
61
def report_decrement(resource_type: ResourceType, resource_id: UUID) -> None:
    """Reports the deletion/deactivation of a feature/resource.

    Args:
        resource_type: The type of resource to report a decrement in count for.
        resource_id: ID of the resource that was deleted.
    """
    if not server_config().feature_gate_enabled:
        return
    feature_gate().report_event(
        resource=resource_type, resource_id=resource_id, is_decrement=True
    )
report_usage(resource_type: ResourceType, resource_id: UUID) -> None

Reports the creation/usage of a feature/resource.

Parameters:

Name Type Description Default
resource_type ResourceType

The type of resource to report a usage for

required
resource_id UUID

ID of the resource that was created.

required
Source code in src/zenml/zen_server/feature_gate/endpoint_utils.py
36
37
38
39
40
41
42
43
44
45
46
47
def report_usage(resource_type: ResourceType, resource_id: UUID) -> None:
    """Reports the creation/usage of a feature/resource.

    Args:
        resource_type: The type of resource to report a usage for
        resource_id: ID of the resource that was created.
    """
    if not server_config().feature_gate_enabled:
        return
    feature_gate().report_event(
        resource=resource_type, resource_id=resource_id
    )
feature_gate_interface

Definition of the feature gate interface.

Classes
FeatureGateInterface

Bases: ABC

Feature gate interface definition.

Functions
check_entitlement(resource: ResourceType) -> None abstractmethod

Checks if a user is entitled to create a resource.

Parameters:

Name Type Description Default
resource ResourceType

The resource the user wants to create

required
Source code in src/zenml/zen_server/feature_gate/feature_gate_interface.py
25
26
27
28
29
30
31
32
33
34
@abstractmethod
def check_entitlement(self, resource: ResourceType) -> None:
    """Checks if a user is entitled to create a resource.

    Args:
        resource: The resource the user wants to create

    Raises:
        UpgradeRequiredError in case a subscription limit is reached
    """
report_event(resource: ResourceType, resource_id: UUID, is_decrement: bool = False) -> None abstractmethod

Reports the usage of a feature to the aggregator backend.

Parameters:

Name Type Description Default
resource ResourceType

The resource the user created

required
resource_id UUID

ID of the resource that was created/deleted.

required
is_decrement bool

In case this event reports an actual decrement of usage

False
Source code in src/zenml/zen_server/feature_gate/feature_gate_interface.py
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@abstractmethod
def report_event(
    self,
    resource: ResourceType,
    resource_id: UUID,
    is_decrement: bool = False,
) -> None:
    """Reports the usage of a feature to the aggregator backend.

    Args:
        resource: The resource the user created
        resource_id: ID of the resource that was created/deleted.
        is_decrement: In case this event reports an actual decrement of usage
    """
zenml_cloud_feature_gate

ZenML Pro implementation of the feature gate.

Classes
RawUsageEvent

Bases: BaseModel

Model for reporting raw usage of a feature.

In case of consumables the UsageReport allows the Pricing Backend to increment the usage per time-frame by 1.

ZenMLCloudFeatureGateInterface()

Bases: FeatureGateInterface

ZenML Cloud Feature Gate implementation.

Initialize the object.

Source code in src/zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py
65
66
67
def __init__(self) -> None:
    """Initialize the object."""
    self._connection = cloud_connection()
Functions
check_entitlement(resource: ResourceType) -> None

Checks if a user is entitled to create a resource.

Parameters:

Name Type Description Default
resource ResourceType

The resource the user wants to create

required

Raises:

Type Description
SubscriptionUpgradeRequiredError

in case a subscription limit is reached

Source code in src/zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def check_entitlement(self, resource: ResourceType) -> None:
    """Checks if a user is entitled to create a resource.

    Args:
        resource: The resource the user wants to create

    Raises:
        SubscriptionUpgradeRequiredError: in case a subscription limit is reached
    """
    try:
        response = self._connection.get(
            endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
        )
    except SubscriptionUpgradeRequiredError:
        raise SubscriptionUpgradeRequiredError(
            f"Your subscription reached its `{resource}` limit. Please "
            f"upgrade your subscription or reach out to us."
        )

    if response.status_code != 200:
        logger.warning(
            "Unexpected response status code from entitlement "
            f"endpoint: {response.status_code}. Message: "
            f"{response.json()}"
        )
report_event(resource: ResourceType, resource_id: UUID, is_decrement: bool = False) -> None

Reports the usage of a feature to the aggregator backend.

Parameters:

Name Type Description Default
resource ResourceType

The resource the user created

required
resource_id UUID

ID of the resource that was created/deleted.

required
is_decrement bool

In case this event reports an actual decrement of usage

False
Source code in src/zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def report_event(
    self,
    resource: ResourceType,
    resource_id: UUID,
    is_decrement: bool = False,
) -> None:
    """Reports the usage of a feature to the aggregator backend.

    Args:
        resource: The resource the user created
        resource_id: ID of the resource that was created/deleted.
        is_decrement: In case this event reports an actual decrement of usage
    """
    data = RawUsageEvent(
        organization_id=ORGANIZATION_ID,
        feature=resource,
        total=1 if not is_decrement else -1,
        metadata={
            "workspace_id": str(server_config.get_external_server_id()),
            "resource_id": str(resource_id),
        },
    ).model_dump()
    response = self._connection.post(
        endpoint=USAGE_EVENT_ENDPOINT, data=data
    )
    if response.status_code != 200:
        logger.error(
            "Usage report not accepted by upstream backend. "
            f"Status Code: {response.status_code}, Message: "
            f"{response.json()}."
        )
Functions

jwt

JWT utilities module for ZenML server.

Classes
JWTToken

Bases: BaseModel

Pydantic object representing a JWT token.

Attributes:

Name Type Description
user_id UUID

The id of the authenticated User.

device_id Optional[UUID]

The id of the authenticated device.

api_key_id Optional[UUID]

The id of the authenticated API key for which this token was issued.

schedule_id Optional[UUID]

The id of the schedule for which the token was issued.

pipeline_run_id Optional[UUID]

The id of the pipeline run for which the token was issued.

step_run_id Optional[UUID]

The id of the step run for which the token was issued.

session_id Optional[UUID]

The id of the authenticated session (used for CSRF).

claims Dict[str, Any]

The original token claims.

Functions
decode_token(token: str, verify: bool = True) -> JWTToken classmethod

Decodes a JWT access token.

Decodes a JWT access token and returns a JWTToken object with the information retrieved from its subject claim.

Parameters:

Name Type Description Default
token str

The encoded JWT token.

required
verify bool

Whether to verify the signature of the token.

True

Returns:

Type Description
JWTToken

The decoded JWT access token.

Raises:

Type Description
CredentialsNotValid

If the token is invalid.

Source code in src/zenml/zen_server/jwt.py
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@classmethod
def decode_token(
    cls,
    token: str,
    verify: bool = True,
) -> "JWTToken":
    """Decodes a JWT access token.

    Decodes a JWT access token and returns a `JWTToken` object with the
    information retrieved from its subject claim.

    Args:
        token: The encoded JWT token.
        verify: Whether to verify the signature of the token.

    Returns:
        The decoded JWT access token.

    Raises:
        CredentialsNotValid: If the token is invalid.
    """
    config = server_config()

    try:
        claims_data = jwt.decode(
            token,
            config.jwt_secret_key,
            algorithms=[config.jwt_token_algorithm],
            audience=config.get_jwt_token_audience(),
            issuer=config.get_jwt_token_issuer(),
            verify=verify,
            leeway=timedelta(seconds=config.jwt_token_leeway_seconds),
        )
        claims = cast(Dict[str, Any], claims_data)
    except jwt.PyJWTError as e:
        raise CredentialsNotValid(f"Invalid JWT token: {e}") from e

    subject: str = claims.pop("sub", "")
    if not subject:
        raise CredentialsNotValid(
            "Invalid JWT token: the subject claim is missing"
        )

    try:
        user_id = UUID(subject)
    except ValueError:
        raise CredentialsNotValid(
            "Invalid JWT token: the subject claim is not a valid UUID"
        )

    device_id: Optional[UUID] = None
    if "device_id" in claims:
        try:
            device_id = UUID(claims.pop("device_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the device_id claim is not a valid "
                "UUID"
            )

    api_key_id: Optional[UUID] = None
    if "api_key_id" in claims:
        try:
            api_key_id = UUID(claims.pop("api_key_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the api_key_id claim is not a valid "
                "UUID"
            )

    schedule_id: Optional[UUID] = None
    if "schedule_id" in claims:
        try:
            schedule_id = UUID(claims.pop("schedule_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the schedule_id claim is not a valid "
                "UUID"
            )

    pipeline_run_id: Optional[UUID] = None
    if "pipeline_run_id" in claims:
        try:
            pipeline_run_id = UUID(claims.pop("pipeline_run_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the pipeline_run_id claim is not a valid "
                "UUID"
            )

    step_run_id: Optional[UUID] = None
    if "step_run_id" in claims:
        try:
            step_run_id = UUID(claims.pop("step_run_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the step_run_id claim is not a valid "
                "UUID"
            )

    session_id: Optional[UUID] = None
    if "session_id" in claims:
        try:
            session_id = UUID(claims.pop("session_id"))
        except ValueError:
            raise CredentialsNotValid(
                "Invalid JWT token: the session_id claim is not a valid "
                "UUID"
            )

    return JWTToken(
        user_id=user_id,
        device_id=device_id,
        api_key_id=api_key_id,
        schedule_id=schedule_id,
        pipeline_run_id=pipeline_run_id,
        step_run_id=step_run_id,
        session_id=session_id,
        claims=claims,
    )
encode(expires: Optional[datetime] = None) -> str

Creates a JWT access token.

Encodes, signs and returns a JWT access token.

Parameters:

Name Type Description Default
expires Optional[datetime]

Datetime after which the token will expire. If not provided, the JWT token will not be set to expire.

None

Returns:

Type Description
str

The generated access token.

Source code in src/zenml/zen_server/jwt.py
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def encode(self, expires: Optional[datetime] = None) -> str:
    """Creates a JWT access token.

    Encodes, signs and returns a JWT access token.

    Args:
        expires: Datetime after which the token will expire. If not
            provided, the JWT token will not be set to expire.

    Returns:
        The generated access token.
    """
    config = server_config()

    claims: Dict[str, Any] = self.claims.copy()

    claims["sub"] = str(self.user_id)
    claims["iss"] = config.get_jwt_token_issuer()
    claims["aud"] = config.get_jwt_token_audience()

    if expires:
        claims["exp"] = expires
    else:
        claims.pop("exp", None)

    if self.device_id:
        claims["device_id"] = str(self.device_id)
    if self.api_key_id:
        claims["api_key_id"] = str(self.api_key_id)
    if self.schedule_id:
        claims["schedule_id"] = str(self.schedule_id)
    if self.pipeline_run_id:
        claims["pipeline_run_id"] = str(self.pipeline_run_id)
    if self.step_run_id:
        claims["step_run_id"] = str(self.step_run_id)
    if self.session_id:
        claims["session_id"] = str(self.session_id)

    return jwt.encode(
        claims,
        config.jwt_secret_key,
        algorithm=config.jwt_token_algorithm,
    )
Functions

rate_limit

Rate limiting for the ZenML Server.

Classes
RequestLimiter(day_limit: Optional[int] = None, minute_limit: Optional[int] = None)

Simple in-memory rate limiter.

Initializes the limiter.

Parameters:

Name Type Description Default
day_limit Optional[int]

The number of requests allowed per day.

None
minute_limit Optional[int]

The number of requests allowed per minute.

None

Raises:

Type Description
ValueError

If both day_limit and minute_limit are None.

Source code in src/zenml/zen_server/rate_limit.py
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    day_limit: Optional[int] = None,
    minute_limit: Optional[int] = None,
):
    """Initializes the limiter.

    Args:
        day_limit: The number of requests allowed per day.
        minute_limit: The number of requests allowed per minute.

    Raises:
        ValueError: If both day_limit and minute_limit are None.
    """
    self.limiting_enabled = server_config().rate_limit_enabled
    if not self.limiting_enabled:
        return
    if day_limit is None and minute_limit is None:
        raise ValueError("Pass either day or minuter limits, or both.")
    self.day_limit = day_limit
    self.minute_limit = minute_limit
    self.limiter: Dict[str, List[float]] = defaultdict(list)
Functions
hit_limiter(request: Request) -> None

Increase the number of hits in the limiter.

Parameters:

Name Type Description Default
request Request

Request object.

required

Raises:

Type Description
HTTPException

If the request limit is exceeded.

Source code in src/zenml/zen_server/rate_limit.py
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def hit_limiter(self, request: Request) -> None:
    """Increase the number of hits in the limiter.

    Args:
        request: Request object.

    Raises:
        HTTPException: If the request limit is exceeded.
    """
    if not self.limiting_enabled:
        return
    from fastapi import HTTPException

    requester = self._get_ipaddr(request)
    now = time.time()
    minute_ago = now - 60
    day_ago = now - 60 * 60 * 24
    self.limiter[requester].append(now)

    from bisect import bisect_left

    # remove failures older than a day
    older_index = bisect_left(self.limiter[requester], day_ago)
    self.limiter[requester] = self.limiter[requester][older_index:]

    if self.day_limit and len(self.limiter[requester]) > self.day_limit:
        raise HTTPException(
            status_code=429, detail="Daily request limit exceeded."
        )
    minute_requests = len(
        [
            limiter_hit
            for limiter_hit in self.limiter[requester][::-1]
            if limiter_hit >= minute_ago
        ]
    )
    if self.minute_limit and minute_requests > self.minute_limit:
        raise HTTPException(
            status_code=429, detail="Minute request limit exceeded."
        )
limit_failed_requests(request: Request) -> Generator[None, Any, Any]

Limits the number of failed requests.

Parameters:

Name Type Description Default
request Request

Request object.

required

Yields:

Type Description
None

None

Source code in src/zenml/zen_server/rate_limit.py
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@contextmanager
def limit_failed_requests(
    self, request: Request
) -> Generator[None, Any, Any]:
    """Limits the number of failed requests.

    Args:
        request: Request object.

    Yields:
        None
    """
    self.hit_limiter(request)

    yield

    # if request was successful - reset limiter
    self.reset_limiter(request)
reset_limiter(request: Request) -> None

Resets the limiter on successful request.

Parameters:

Name Type Description Default
request Request

Request object.

required
Source code in src/zenml/zen_server/rate_limit.py
108
109
110
111
112
113
114
115
116
117
def reset_limiter(self, request: Request) -> None:
    """Resets the limiter on successful request.

    Args:
        request: Request object.
    """
    if self.limiting_enabled:
        requester = self._get_ipaddr(request)
        if requester in self.limiter:
            del self.limiter[requester]
Functions
rate_limit_requests(day_limit: Optional[int] = None, minute_limit: Optional[int] = None) -> Callable[..., Any]

Decorator to handle exceptions in the API.

Parameters:

Name Type Description Default
day_limit Optional[int]

Number of requests allowed per day.

None
minute_limit Optional[int]

Number of requests allowed per minute.

None

Returns:

Type Description
Callable[..., Any]

Decorated function.

Source code in src/zenml/zen_server/rate_limit.py
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def rate_limit_requests(
    day_limit: Optional[int] = None,
    minute_limit: Optional[int] = None,
) -> Callable[..., Any]:
    """Decorator to handle exceptions in the API.

    Args:
        day_limit: Number of requests allowed per day.
        minute_limit: Number of requests allowed per minute.

    Returns:
        Decorated function.
    """
    limiter = RequestLimiter(day_limit=day_limit, minute_limit=minute_limit)

    def decorator(func: F) -> F:
        request_arg, request_kwarg = None, None
        parameters = inspect.signature(func).parameters
        for arg_num, arg_name in enumerate(parameters):
            if parameters[arg_name].annotation == Request:
                request_arg = arg_num
                request_kwarg = arg_name
                break
        if request_arg is None or request_kwarg is None:
            raise ValueError(
                "Rate limiting APIs must have argument of `Request` type."
            )

        @wraps(func)
        def decorated(
            *args: Any,
            **kwargs: Any,
        ) -> Any:
            if request_kwarg in kwargs:
                request = kwargs[request_kwarg]
            else:
                request = args[request_arg]
            with limiter.limit_failed_requests(request):
                return func(*args, **kwargs)

        return cast(F, decorated)

    return decorator

rbac

RBAC definitions.

Modules
endpoint_utils

High-level helper functions to write endpoints with RBAC.

Classes Functions
verify_permissions_and_batch_create_entity(batch: List[AnyRequest], create_method: Callable[[List[AnyRequest]], List[AnyResponse]]) -> List[AnyResponse]

Verify permissions and create a batch of entities if authorized.

Parameters:

Name Type Description Default
batch List[AnyRequest]

The batch to create.

required
create_method Callable[[List[AnyRequest]], List[AnyResponse]]

The method to create the entities.

required

Raises:

Type Description
RuntimeError

If the resource type is usage-tracked.

Returns:

Type Description
List[AnyResponse]

The created entities.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def verify_permissions_and_batch_create_entity(
    batch: List[AnyRequest],
    create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
) -> List[AnyResponse]:
    """Verify permissions and create a batch of entities if authorized.

    Args:
        batch: The batch to create.
        create_method: The method to create the entities.

    Raises:
        RuntimeError: If the resource type is usage-tracked.

    Returns:
        The created entities.
    """
    auth_context = get_auth_context()
    assert auth_context

    resource_types = set()
    for request_model in batch:
        resource_type = get_resource_type_for_model(request_model)
        if resource_type:
            resource_types.add(resource_type)

        if isinstance(request_model, UserScopedRequest):
            # Ignore the user field set in the request model, if any, and set it
            # to the current user's ID instead. This is just a precaution, given
            # that the SQLZenStore also does this same validation on all request
            # models.
            request_model.user = auth_context.user.id

    batch_verify_permissions_for_models(models=batch, action=Action.CREATE)

    if resource_types & set(server_config().reportable_resources):
        raise RuntimeError(
            "Batch requests are currently not possible with usage-tracked "
            "features."
        )

    created = create_method(batch)
    return dehydrate_response_model_batch(created)
verify_permissions_and_create_entity(request_model: AnyRequest, create_method: Callable[[AnyRequest], AnyResponse], surrogate_models: Optional[List[AnyOtherResponse]] = None, skip_entitlements: bool = False) -> AnyResponse

Verify permissions and create the entity if authorized.

Parameters:

Name Type Description Default
request_model AnyRequest

The entity request model.

required
create_method Callable[[AnyRequest], AnyResponse]

The method to create the entity.

required
surrogate_models Optional[List[AnyOtherResponse]]

Optional list of surrogate models to verify UPDATE permissions for instead of verifying CREATE permissions for the request model.

None
skip_entitlements bool

Whether to skip the entitlement check and usage increment.

False

Returns:

Type Description
AnyResponse

A model of the created entity.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def verify_permissions_and_create_entity(
    request_model: AnyRequest,
    create_method: Callable[[AnyRequest], AnyResponse],
    surrogate_models: Optional[List[AnyOtherResponse]] = None,
    skip_entitlements: bool = False,
) -> AnyResponse:
    """Verify permissions and create the entity if authorized.

    Args:
        request_model: The entity request model.
        create_method: The method to create the entity.
        surrogate_models: Optional list of surrogate models to verify
            UPDATE permissions for instead of verifying CREATE permissions for
            the request model.
        skip_entitlements: Whether to skip the entitlement check and usage
            increment.

    Returns:
        A model of the created entity.
    """
    if isinstance(request_model, UserScopedRequest):
        auth_context = get_auth_context()
        assert auth_context

        # Ignore the user field set in the request model, if any, and set it to
        # the current user's ID instead. This is just a precaution, given that
        # the SQLZenStore also does this same validation on all request models.
        request_model.user = auth_context.user.id

    if surrogate_models:
        batch_verify_permissions_for_models(
            models=surrogate_models, action=Action.UPDATE
        )
    else:
        verify_permission_for_model(model=request_model, action=Action.CREATE)

    resource_type = get_resource_type_for_model(request_model)

    if resource_type:
        needs_usage_increment = (
            not skip_entitlements
            and resource_type in server_config().reportable_resources
        )
        if needs_usage_increment:
            check_entitlement(resource_type)

    created = create_method(request_model)

    if resource_type and needs_usage_increment:
        report_usage(resource_type, resource_id=created.id)

    return dehydrate_response_model(created)
verify_permissions_and_delete_entity(id: UUIDOrStr, get_method: Callable[[UUIDOrStr, bool], AnyResponse], delete_method: Callable[[UUIDOrStr], None], **delete_method_kwargs: Any) -> AnyResponse

Verify permissions and delete an entity.

Parameters:

Name Type Description Default
id UUIDOrStr

The ID of the entity to delete.

required
get_method Callable[[UUIDOrStr, bool], AnyResponse]

The method to fetch the entity.

required
delete_method Callable[[UUIDOrStr], None]

The method to delete the entity.

required
delete_method_kwargs Any

Keyword arguments to pass to the delete method.

{}

Returns:

Type Description
AnyResponse

The deleted entity.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def verify_permissions_and_delete_entity(
    id: UUIDOrStr,
    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
    delete_method: Callable[[UUIDOrStr], None],
    **delete_method_kwargs: Any,
) -> AnyResponse:
    """Verify permissions and delete an entity.

    Args:
        id: The ID of the entity to delete.
        get_method: The method to fetch the entity.
        delete_method: The method to delete the entity.
        delete_method_kwargs: Keyword arguments to pass to the delete method.

    Returns:
        The deleted entity.
    """
    model = get_method(id, True)
    verify_permission_for_model(model, action=Action.DELETE)
    delete_method(model.id, **delete_method_kwargs)
    delete_model_resource(model)

    return model
verify_permissions_and_get_entity(id: UUIDOrStr, get_method: Callable[[UUIDOrStr], AnyResponse], **get_method_kwargs: Any) -> AnyResponse

Verify permissions and fetch an entity.

Parameters:

Name Type Description Default
id UUIDOrStr

The ID of the entity to fetch.

required
get_method Callable[[UUIDOrStr], AnyResponse]

The method to fetch the entity.

required
get_method_kwargs Any

Keyword arguments to pass to the get method.

{}

Returns:

Type Description
AnyResponse

A model of the fetched entity.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def verify_permissions_and_get_entity(
    id: UUIDOrStr,
    get_method: Callable[[UUIDOrStr], AnyResponse],
    **get_method_kwargs: Any,
) -> AnyResponse:
    """Verify permissions and fetch an entity.

    Args:
        id: The ID of the entity to fetch.
        get_method: The method to fetch the entity.
        get_method_kwargs: Keyword arguments to pass to the get method.

    Returns:
        A model of the fetched entity.
    """
    model = get_method(id, **get_method_kwargs)
    verify_permission_for_model(model, action=Action.READ)
    return dehydrate_response_model(model)
verify_permissions_and_get_or_create_entity(request_model: AnyRequest, get_or_create_method: Callable[[AnyRequest, Optional[Callable[[], None]]], Tuple[AnyResponse, bool]]) -> Tuple[AnyResponse, bool]

Verify permissions and create the entity if authorized.

Parameters:

Name Type Description Default
request_model AnyRequest

The entity request model.

required
get_or_create_method Callable[[AnyRequest, Optional[Callable[[], None]]], Tuple[AnyResponse, bool]]

The method to get or create the entity.

required

Returns:

Type Description
Tuple[AnyResponse, bool]

The entity and a boolean indicating whether the entity was created.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def verify_permissions_and_get_or_create_entity(
    request_model: AnyRequest,
    get_or_create_method: Callable[
        [AnyRequest, Optional[Callable[[], None]]], Tuple[AnyResponse, bool]
    ],
) -> Tuple[AnyResponse, bool]:
    """Verify permissions and create the entity if authorized.

    Args:
        request_model: The entity request model.
        get_or_create_method: The method to get or create the entity.

    Returns:
        The entity and a boolean indicating whether the entity was created.
    """
    if isinstance(request_model, UserScopedRequest):
        auth_context = get_auth_context()
        assert auth_context

        # Ignore the user field set in the request model, if any, and set it to
        # the current user's ID instead. This is just a precaution, given that
        # the SQLZenStore also does this same validation on all request models.
        request_model.user = auth_context.user.id

    resource_type = get_resource_type_for_model(request_model)
    needs_usage_increment = (
        resource_type and resource_type in server_config().reportable_resources
    )

    def _pre_creation_hook() -> None:
        verify_permission_for_model(model=request_model, action=Action.CREATE)
        if resource_type and needs_usage_increment:
            check_entitlement(resource_type=resource_type)

    model, created = get_or_create_method(request_model, _pre_creation_hook)

    if not created:
        verify_permission_for_model(model=model, action=Action.READ)
    elif resource_type and needs_usage_increment:
        report_usage(resource_type, resource_id=model.id)

    return dehydrate_response_model(model), created
verify_permissions_and_list_entities(filter_model: AnyFilter, resource_type: ResourceType, list_method: Callable[[AnyFilter], Page[AnyResponse]], **list_method_kwargs: Any) -> Page[AnyResponse]

Verify permissions and list entities.

Parameters:

Name Type Description Default
filter_model AnyFilter

The entity filter model.

required
resource_type ResourceType

The resource type of the entities to list.

required
list_method Callable[[AnyFilter], Page[AnyResponse]]

The method to list the entities.

required
list_method_kwargs Any

Keyword arguments to pass to the list method.

{}

Returns:

Type Description
Page[AnyResponse]

A page of entity models.

Raises:

Type Description
ValueError

If the filter's project scope is not set or is not a UUID.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
def verify_permissions_and_list_entities(
    filter_model: AnyFilter,
    resource_type: ResourceType,
    list_method: Callable[[AnyFilter], Page[AnyResponse]],
    **list_method_kwargs: Any,
) -> Page[AnyResponse]:
    """Verify permissions and list entities.

    Args:
        filter_model: The entity filter model.
        resource_type: The resource type of the entities to list.
        list_method: The method to list the entities.
        list_method_kwargs: Keyword arguments to pass to the list method.

    Returns:
        A page of entity models.

    Raises:
        ValueError: If the filter's project scope is not set or is not a UUID.
    """
    auth_context = get_auth_context()
    assert auth_context

    project_id: Optional[UUID] = None
    if isinstance(filter_model, ProjectScopedFilter):
        # A project scoped filter must always be scoped to a specific
        # project. This is required for the RBAC check to work.
        set_filter_project_scope(filter_model)
        if not filter_model.project or not isinstance(
            filter_model.project, UUID
        ):
            raise ValueError(
                "Project scope must be a UUID, got "
                f"{type(filter_model.project)}."
            )
        project_id = filter_model.project

    allowed_ids = get_allowed_resource_ids(
        resource_type=resource_type, project_id=project_id
    )
    filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id, id=allowed_ids
    )
    page = list_method(filter_model, **list_method_kwargs)
    return dehydrate_page(page)
verify_permissions_and_prune_entities(resource_type: ResourceType, prune_method: Callable[..., None], project_id: Optional[UUID] = None, **kwargs: Any) -> None

Verify permissions and prune entities of certain type.

Parameters:

Name Type Description Default
resource_type ResourceType

The resource type of the entities to prune.

required
prune_method Callable[..., None]

The method to prune the entities.

required
project_id Optional[UUID]

The project ID to prune the entities for.

None
kwargs Any

Keyword arguments to pass to the prune method.

{}
Source code in src/zenml/zen_server/rbac/endpoint_utils.py
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
def verify_permissions_and_prune_entities(
    resource_type: ResourceType,
    prune_method: Callable[..., None],
    project_id: Optional[UUID] = None,
    **kwargs: Any,
) -> None:
    """Verify permissions and prune entities of certain type.

    Args:
        resource_type: The resource type of the entities to prune.
        prune_method: The method to prune the entities.
        project_id: The project ID to prune the entities for.
        kwargs: Keyword arguments to pass to the prune method.
    """
    verify_permission(
        resource_type=resource_type,
        action=Action.PRUNE,
        project_id=project_id,
    )
    prune_method(**kwargs)
verify_permissions_and_update_entity(id: UUIDOrStr, update_model: AnyUpdate, get_method: Callable[[UUIDOrStr, bool], AnyResponse], update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse], **update_method_kwargs: Any) -> AnyResponse

Verify permissions and update an entity.

Parameters:

Name Type Description Default
id UUIDOrStr

The ID of the entity to update.

required
update_model AnyUpdate

The entity update model.

required
get_method Callable[[UUIDOrStr, bool], AnyResponse]

The method to fetch the entity.

required
update_method Callable[[UUIDOrStr, AnyUpdate], AnyResponse]

The method to update the entity.

required
update_method_kwargs Any

Keyword arguments to pass to the update method.

{}

Returns:

Type Description
AnyResponse

A model of the updated entity.

Source code in src/zenml/zen_server/rbac/endpoint_utils.py
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def verify_permissions_and_update_entity(
    id: UUIDOrStr,
    update_model: AnyUpdate,
    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
    update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
    **update_method_kwargs: Any,
) -> AnyResponse:
    """Verify permissions and update an entity.

    Args:
        id: The ID of the entity to update.
        update_model: The entity update model.
        get_method: The method to fetch the entity.
        update_method: The method to update the entity.
        update_method_kwargs: Keyword arguments to pass to the update method.

    Returns:
        A model of the updated entity.
    """
    # 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, **update_method_kwargs
    )
    return dehydrate_response_model(updated_model)
models

RBAC model classes.

Classes
Action

Bases: StrEnum

RBAC actions.

Resource

Bases: BaseModel

RBAC resource model.

Functions
parse(resource: str) -> Resource classmethod

Parse an RBAC resource string into a Resource object.

Parameters:

Name Type Description Default
resource str

The resource to convert.

required

Returns:

Type Description
Resource

The converted resource.

Source code in src/zenml/zen_server/rbac/models.py
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
@classmethod
def parse(cls, resource: str) -> "Resource":
    """Parse an RBAC resource string into a Resource object.

    Args:
        resource: The resource to convert.

    Returns:
        The converted resource.
    """
    project_id: Optional[str] = None
    if ":" in resource:
        (
            project_id,
            resource_type_and_id,
        ) = resource.split(":", maxsplit=1)
    else:
        project_id = None
        resource_type_and_id = resource

    resource_id: Optional[str] = None
    if "/" in resource_type_and_id:
        resource_type, resource_id = resource_type_and_id.split("/")
    else:
        resource_type = resource_type_and_id

    return Resource(
        type=resource_type, id=resource_id, project_id=project_id
    )
validate_project_id() -> Resource

Validate that project_id is set in combination with project-scoped resource types.

Raises:

Type Description
ValueError

If project_id is not set for a project-scoped resource or set for an unscoped resource.

Returns:

Type Description
Resource

The validated resource.

Source code in src/zenml/zen_server/rbac/models.py
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@model_validator(mode="after")
def validate_project_id(self) -> "Resource":
    """Validate that project_id is set in combination with project-scoped resource types.

    Raises:
        ValueError: If project_id is not set for a project-scoped
            resource or set for an unscoped resource.

    Returns:
        The validated resource.
    """
    resource_type = ResourceType(self.type)

    if resource_type.is_project_scoped() and not self.project_id:
        raise ValueError(
            "project_id must be set for project-scoped resource type "
            f"'{self.type}'"
        )

    if not resource_type.is_project_scoped() and self.project_id:
        raise ValueError(
            "project_id must not be set for global resource type "
            f"'{self.type}'"
        )

    return self
ResourceType

Bases: StrEnum

Resource types of the server API.

Functions
is_project_scoped() -> bool

Check if a resource type is project scoped.

Returns:

Type Description
bool

Whether the resource type is project scoped.

Source code in src/zenml/zen_server/rbac/models.py
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def is_project_scoped(self) -> bool:
    """Check if a resource type is project scoped.

    Returns:
        Whether the resource type is project scoped.
    """
    return self not in [
        self.FLAVOR,
        self.SECRET,
        self.SERVICE_CONNECTOR,
        self.STACK,
        self.STACK_COMPONENT,
        self.TAG,
        self.SERVICE_ACCOUNT,
        self.PROJECT,
        # Deactivated for now
        # self.USER,
    ]
rbac_interface

RBAC interface definition.

Classes
RBACInterface

Bases: ABC

RBAC interface definition.

Functions
check_permissions(user: UserResponse, resources: Set[Resource], action: Action) -> Dict[Resource, bool] abstractmethod

Checks if a user has permissions to perform an action on resources.

Parameters:

Name Type Description Default
user UserResponse

User which wants to access a resource.

required
resources Set[Resource]

The resources the user wants to access.

required
action Action

The action that the user wants to perform on the resources.

required

Returns:

Type Description
Dict[Resource, bool]

A dictionary mapping resources to a boolean which indicates whether

Dict[Resource, bool]

the user has permissions to perform the action on that resource.

Source code in src/zenml/zen_server/rbac/rbac_interface.py
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@abstractmethod
def check_permissions(
    self, user: "UserResponse", resources: Set[Resource], action: Action
) -> Dict[Resource, bool]:
    """Checks if a user has permissions to perform an action on resources.

    Args:
        user: User which wants to access a resource.
        resources: The resources the user wants to access.
        action: The action that the user wants to perform on the resources.

    Returns:
        A dictionary mapping resources to a boolean which indicates whether
        the user has permissions to perform the action on that resource.
    """
delete_resources(resources: List[Resource]) -> None abstractmethod

Delete resource membership information for a list of resources.

Parameters:

Name Type Description Default
resources List[Resource]

The resources for which to delete the resource membership information.

required
Source code in src/zenml/zen_server/rbac/rbac_interface.py
77
78
79
80
81
82
83
84
@abstractmethod
def delete_resources(self, resources: List[Resource]) -> None:
    """Delete resource membership information for a list of resources.

    Args:
        resources: The resources for which to delete the resource membership
            information.
    """
list_allowed_resource_ids(user: UserResponse, resource: Resource, action: Action) -> Tuple[bool, List[str]] abstractmethod

Lists all resource IDs of a resource type that a user can access.

Parameters:

Name Type Description Default
user UserResponse

User which wants to access a resource.

required
resource Resource

The resource the user wants to access.

required
action Action

The action that the user wants to perform on the resource.

required

Returns:

Type Description
bool

A tuple (full_resource_access, resource_ids).

List[str]

full_resource_access will be True if the user can perform the

Tuple[bool, List[str]]

given action on any instance of the given resource type, False

Tuple[bool, List[str]]

otherwise. If full_resource_access is False, resource_ids

Tuple[bool, List[str]]

will contain the list of instance IDs that the user can perform

Tuple[bool, List[str]]

the action on.

Source code in src/zenml/zen_server/rbac/rbac_interface.py
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
@abstractmethod
def list_allowed_resource_ids(
    self, user: "UserResponse", resource: Resource, action: Action
) -> Tuple[bool, List[str]]:
    """Lists all resource IDs of a resource type that a user can access.

    Args:
        user: User which wants to access a resource.
        resource: The resource the user wants to access.
        action: The action that the user wants to perform on the resource.

    Returns:
        A tuple (full_resource_access, resource_ids).
        `full_resource_access` will be `True` if the user can perform the
        given action on any instance of the given resource type, `False`
        otherwise. If `full_resource_access` is `False`, `resource_ids`
        will contain the list of instance IDs that the user can perform
        the action on.
    """
update_resource_membership(user: UserResponse, resource: Resource, actions: List[Action]) -> None abstractmethod

Update the resource membership of a user.

Parameters:

Name Type Description Default
user UserResponse

User for which the resource membership should be updated.

required
resource Resource

The resource.

required
actions List[Action]

The actions that the user should be able to perform on the resource.

required
Source code in src/zenml/zen_server/rbac/rbac_interface.py
64
65
66
67
68
69
70
71
72
73
74
75
@abstractmethod
def update_resource_membership(
    self, user: "UserResponse", resource: Resource, actions: List[Action]
) -> None:
    """Update the resource membership of a user.

    Args:
        user: User for which the resource membership should be updated.
        resource: The resource.
        actions: The actions that the user should be able to perform on the
            resource.
    """
rbac_sql_zen_store

RBAC SQL Zen Store implementation.

Classes
RBACSqlZenStore(skip_default_registrations: bool = False, **kwargs: Any)

Bases: SqlZenStore

Wrapper around the SQLZenStore that implements RBAC functionality.

Source code in src/zenml/zen_stores/base_zen_store.py
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def __init__(
    self,
    skip_default_registrations: bool = False,
    **kwargs: Any,
) -> None:
    """Create and initialize a store.

    Args:
        skip_default_registrations: If `True`, the creation of the default
            stack and user in the store will be skipped.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.
    """
    super().__init__(**kwargs)

    self._initialize()

    if not skip_default_registrations:
        logger.debug("Initializing database")
        self._initialize_database()
    else:
        logger.debug("Skipping database initialization")
Functions
utils

RBAC utility functions.

Classes Functions
batch_verify_permissions(resources: Set[Resource], action: Action) -> None

Batch permission verification.

Parameters:

Name Type Description Default
resources Set[Resource]

The resources the user wants to perform the action on.

required
action Action

The action the user wants to perform.

required

Raises:

Type Description
IllegalOperationError

If the user is not allowed to perform the action.

RuntimeError

If the permission verification failed unexpectedly.

Source code in src/zenml/zen_server/rbac/utils.py
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
def batch_verify_permissions(
    resources: Set[Resource],
    action: Action,
) -> None:
    """Batch permission verification.

    Args:
        resources: The resources the user wants to perform the action on.
        action: The action the user wants to perform.

    Raises:
        IllegalOperationError: If the user is not allowed to perform the action.
        RuntimeError: If the permission verification failed unexpectedly.
    """
    if not server_config().rbac_enabled:
        return

    auth_context = get_auth_context()
    assert auth_context

    permissions = rbac().check_permissions(
        user=auth_context.user, resources=resources, action=action
    )

    for resource in resources:
        if resource not in permissions:
            # This should never happen if the RBAC implementation is working
            # correctly
            raise RuntimeError(
                f"Failed to verify permissions to {action.upper()} resource "
                f"'{resource}'."
            )

        if not permissions[resource]:
            raise IllegalOperationError(
                message=f"Insufficient permissions to {action.upper()} "
                f"resource '{resource}'.",
            )
batch_verify_permissions_for_models(models: Sequence[AnyModel], action: Action) -> None

Batch permission verification for models.

Parameters:

Name Type Description Default
models Sequence[AnyModel]

The models the user wants to perform the action on.

required
action Action

The action the user wants to perform.

required
Source code in src/zenml/zen_server/rbac/utils.py
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def batch_verify_permissions_for_models(
    models: Sequence[AnyModel],
    action: Action,
) -> None:
    """Batch permission verification for models.

    Args:
        models: The models the user wants to perform the action on.
        action: The action the user wants to perform.
    """
    if not server_config().rbac_enabled:
        return

    resources = set()
    for model in models:
        if is_owned_by_authenticated_user(model):
            # The model owner always has permissions
            continue

        permission_model = get_surrogate_permission_model_for_model(
            model, action=action
        )

        if resource := get_resource_for_model(permission_model):
            resources.add(resource)

    batch_verify_permissions(resources=resources, action=action)
dehydrate_page(page: Page[AnyResponse]) -> Page[AnyResponse]

Dehydrate all items of a page.

Parameters:

Name Type Description Default
page Page[AnyResponse]

The page to dehydrate.

required

Returns:

Type Description
Page[AnyResponse]

The page with (potentially) dehydrated items.

Source code in src/zenml/zen_server/rbac/utils.py
51
52
53
54
55
56
57
58
59
60
61
def dehydrate_page(page: Page[AnyResponse]) -> Page[AnyResponse]:
    """Dehydrate all items of a page.

    Args:
        page: The page to dehydrate.

    Returns:
        The page with (potentially) dehydrated items.
    """
    new_items = dehydrate_response_model_batch(page.items)
    return page.model_copy(update={"items": new_items})
dehydrate_response_model(model: AnyModel, permissions: Optional[Dict[Resource, bool]] = None) -> AnyModel

Dehydrate a model if necessary.

Parameters:

Name Type Description Default
model AnyModel

The model to dehydrate.

required
permissions Optional[Dict[Resource, bool]]

Prefetched permissions that will be used to check whether sub-models will be included in the model or not. If a sub-model refers to a resource which is not included in this dictionary, the permissions will be checked with the RBAC component.

None

Returns:

Type Description
AnyModel

The (potentially) dehydrated model.

Source code in src/zenml/zen_server/rbac/utils.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def dehydrate_response_model(
    model: AnyModel, permissions: Optional[Dict[Resource, bool]] = None
) -> AnyModel:
    """Dehydrate a model if necessary.

    Args:
        model: The model to dehydrate.
        permissions: Prefetched permissions that will be used to check whether
            sub-models will be included in the model or not. If a sub-model
            refers to a resource which is not included in this dictionary, the
            permissions will be checked with the RBAC component.

    Returns:
        The (potentially) dehydrated model.
    """
    if not server_config().rbac_enabled:
        return model

    if not permissions:
        auth_context = get_auth_context()
        assert auth_context

        resources = get_subresources_for_model(model)
        permissions = rbac().check_permissions(
            user=auth_context.user, resources=resources, action=Action.READ
        )

    dehydrated_values = {}
    # See `get_subresources_for_model(...)` for a detailed explanation why we
    # need to use `model.__iter__()` here
    for key, value in model.__iter__():
        dehydrated_values[key] = _dehydrate_value(
            value, permissions=permissions
        )

    return type(model).model_validate(dehydrated_values)
dehydrate_response_model_batch(batch: List[AnyResponse]) -> List[AnyResponse]

Dehydrate all items of a batch.

Parameters:

Name Type Description Default
batch List[AnyResponse]

The batch to dehydrate.

required

Returns:

Type Description
List[AnyResponse]

The batch with (potentially) dehydrated items.

Source code in src/zenml/zen_server/rbac/utils.py
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def dehydrate_response_model_batch(
    batch: List[AnyResponse],
) -> List[AnyResponse]:
    """Dehydrate all items of a batch.

    Args:
        batch: The batch to dehydrate.

    Returns:
        The batch with (potentially) dehydrated items.
    """
    if not server_config().rbac_enabled:
        return batch

    auth_context = get_auth_context()
    assert auth_context

    resource_list = [get_subresources_for_model(item) for item in batch]
    resources = set.union(*resource_list) if resource_list else set()
    permissions = rbac().check_permissions(
        user=auth_context.user, resources=resources, action=Action.READ
    )

    new_batch = [
        dehydrate_response_model(item, permissions=permissions)
        for item in batch
    ]

    return new_batch
delete_model_resource(model: AnyModel) -> None

Delete resource membership information for a model.

Parameters:

Name Type Description Default
model AnyModel

The model for which to delete the resource membership information.

required
Source code in src/zenml/zen_server/rbac/utils.py
709
710
711
712
713
714
715
def delete_model_resource(model: AnyModel) -> None:
    """Delete resource membership information for a model.

    Args:
        model: The model for which to delete the resource membership information.
    """
    delete_model_resources(models=[model])
delete_model_resources(models: List[AnyModel]) -> None

Delete resource membership information for a list of models.

Parameters:

Name Type Description Default
models List[AnyModel]

The models for which to delete the resource membership information.

required
Source code in src/zenml/zen_server/rbac/utils.py
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
def delete_model_resources(models: List[AnyModel]) -> None:
    """Delete resource membership information for a list of models.

    Args:
        models: The models for which to delete the resource membership information.
    """
    if not server_config().rbac_enabled:
        return

    resources = set()
    for model in models:
        if resource := get_resource_for_model(model):
            resources.add(resource)

    delete_resources(resources=list(resources))
delete_resources(resources: List[Resource]) -> None

Delete resource membership information for a list of resources.

Parameters:

Name Type Description Default
resources List[Resource]

The resources for which to delete the resource membership information.

required
Source code in src/zenml/zen_server/rbac/utils.py
735
736
737
738
739
740
741
742
743
744
745
def delete_resources(resources: List[Resource]) -> None:
    """Delete resource membership information for a list of resources.

    Args:
        resources: The resources for which to delete the resource membership
            information.
    """
    if not server_config().rbac_enabled:
        return

    rbac().delete_resources(resources=resources)
get_allowed_resource_ids(resource_type: str, action: Action = Action.READ, project_id: Optional[UUID] = None) -> Optional[Set[UUID]]

Get all resource IDs of a resource type that a user can access.

Parameters:

Name Type Description Default
resource_type str

The resource type.

required
action Action

The action the user wants to perform on the resource.

READ
project_id Optional[UUID]

Optional project ID to filter the resources by. Required for project scoped resources.

None

Returns:

Type Description
Optional[Set[UUID]]

A list of resource IDs or None if the user has full access to the

Optional[Set[UUID]]

all instances of the resource.

Source code in src/zenml/zen_server/rbac/utils.py
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def get_allowed_resource_ids(
    resource_type: str,
    action: Action = Action.READ,
    project_id: Optional[UUID] = None,
) -> Optional[Set[UUID]]:
    """Get all resource IDs of a resource type that a user can access.

    Args:
        resource_type: The resource type.
        action: The action the user wants to perform on the resource.
        project_id: Optional project ID to filter the resources by.
            Required for project scoped resources.

    Returns:
        A list of resource IDs or `None` if the user has full access to the
        all instances of the resource.
    """
    if not server_config().rbac_enabled:
        return None

    auth_context = get_auth_context()
    assert auth_context

    (
        has_full_resource_access,
        allowed_ids,
    ) = rbac().list_allowed_resource_ids(
        user=auth_context.user,
        resource=Resource(type=resource_type, project_id=project_id),
        action=action,
    )

    if has_full_resource_access:
        return None

    return {UUID(id) for id in allowed_ids}
get_permission_denied_model(model: AnyResponse) -> AnyResponse

Get a model to return in case of missing read permissions.

Parameters:

Name Type Description Default
model AnyResponse

The original model.

required

Returns:

Type Description
AnyResponse

The permission denied model.

Source code in src/zenml/zen_server/rbac/utils.py
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def get_permission_denied_model(model: AnyResponse) -> AnyResponse:
    """Get a model to return in case of missing read permissions.

    Args:
        model: The original model.

    Returns:
        The permission denied model.
    """
    return model.model_copy(
        update={
            "body": None,
            "metadata": None,
            "resources": None,
            "permission_denied": True,
        }
    )
get_resource_for_model(model: AnyModel) -> Optional[Resource]

Get the resource associated with a model object.

Parameters:

Name Type Description Default
model AnyModel

The model for which to get the resource.

required

Returns:

Type Description
Optional[Resource]

The resource associated with the model, or None if the model

Optional[Resource]

is not associated with any resource type.

Source code in src/zenml/zen_server/rbac/utils.py
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def get_resource_for_model(model: AnyModel) -> Optional[Resource]:
    """Get the resource associated with a model object.

    Args:
        model: The model for which to get the resource.

    Returns:
        The resource associated with the model, or `None` if the model
        is not associated with any resource type.
    """
    resource_type = get_resource_type_for_model(model)
    if not resource_type:
        # This model is not tied to any RBAC resource type
        return None

    project_id: Optional[UUID] = None
    if isinstance(model, ProjectScopedResponse):
        # A project scoped response is always scoped to a specific project
        project_id = model.project.id
    elif isinstance(model, ProjectScopedRequest):
        # A project scoped request is always scoped to a specific project
        project_id = model.project

    resource_id: Optional[UUID] = None
    if isinstance(model, BaseIdentifiedResponse):
        resource_id = model.id

    return Resource(type=resource_type, id=resource_id, project_id=project_id)
get_resource_type_for_model(model: AnyModel) -> Optional[ResourceType]

Get the resource type associated with a model object.

Parameters:

Name Type Description Default
model AnyModel

The model for which to get the resource type.

required

Returns:

Type Description
Optional[ResourceType]

The resource type associated with the model, or None if the model

Optional[ResourceType]

is not associated with any resource type.

Source code in src/zenml/zen_server/rbac/utils.py
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
def get_resource_type_for_model(
    model: AnyModel,
) -> Optional[ResourceType]:
    """Get the resource type associated with a model object.

    Args:
        model: The model for which to get the resource type.

    Returns:
        The resource type associated with the model, or `None` if the model
        is not associated with any resource type.
    """
    from zenml.models import (
        ActionRequest,
        ActionResponse,
        ArtifactRequest,
        ArtifactResponse,
        ArtifactVersionRequest,
        ArtifactVersionResponse,
        CodeRepositoryRequest,
        CodeRepositoryResponse,
        ComponentRequest,
        ComponentResponse,
        EventSourceRequest,
        EventSourceResponse,
        FlavorRequest,
        FlavorResponse,
        ModelRequest,
        ModelResponse,
        ModelVersionRequest,
        ModelVersionResponse,
        PipelineBuildRequest,
        PipelineBuildResponse,
        PipelineDeploymentRequest,
        PipelineDeploymentResponse,
        PipelineRequest,
        PipelineResponse,
        PipelineRunRequest,
        PipelineRunResponse,
        ProjectRequest,
        ProjectResponse,
        RunMetadataRequest,
        RunTemplateRequest,
        RunTemplateResponse,
        SecretRequest,
        SecretResponse,
        ServiceAccountRequest,
        ServiceAccountResponse,
        ServiceConnectorRequest,
        ServiceConnectorResponse,
        ServiceRequest,
        ServiceResponse,
        StackRequest,
        StackResponse,
        TagRequest,
        TagResponse,
        TriggerExecutionRequest,
        TriggerExecutionResponse,
        TriggerRequest,
        TriggerResponse,
    )

    mapping: Dict[
        Any,
        ResourceType,
    ] = {
        ActionRequest: ResourceType.ACTION,
        ActionResponse: ResourceType.ACTION,
        ArtifactRequest: ResourceType.ARTIFACT,
        ArtifactResponse: ResourceType.ARTIFACT,
        ArtifactVersionRequest: ResourceType.ARTIFACT_VERSION,
        ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
        CodeRepositoryRequest: ResourceType.CODE_REPOSITORY,
        CodeRepositoryResponse: ResourceType.CODE_REPOSITORY,
        ComponentRequest: ResourceType.STACK_COMPONENT,
        ComponentResponse: ResourceType.STACK_COMPONENT,
        EventSourceRequest: ResourceType.EVENT_SOURCE,
        EventSourceResponse: ResourceType.EVENT_SOURCE,
        FlavorRequest: ResourceType.FLAVOR,
        FlavorResponse: ResourceType.FLAVOR,
        ModelRequest: ResourceType.MODEL,
        ModelResponse: ResourceType.MODEL,
        ModelVersionRequest: ResourceType.MODEL_VERSION,
        ModelVersionResponse: ResourceType.MODEL_VERSION,
        PipelineBuildRequest: ResourceType.PIPELINE_BUILD,
        PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
        PipelineDeploymentRequest: ResourceType.PIPELINE_DEPLOYMENT,
        PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
        PipelineRequest: ResourceType.PIPELINE,
        PipelineResponse: ResourceType.PIPELINE,
        PipelineRunRequest: ResourceType.PIPELINE_RUN,
        PipelineRunResponse: ResourceType.PIPELINE_RUN,
        RunMetadataRequest: ResourceType.RUN_METADATA,
        RunTemplateRequest: ResourceType.RUN_TEMPLATE,
        RunTemplateResponse: ResourceType.RUN_TEMPLATE,
        SecretRequest: ResourceType.SECRET,
        SecretResponse: ResourceType.SECRET,
        ServiceAccountRequest: ResourceType.SERVICE_ACCOUNT,
        ServiceAccountResponse: ResourceType.SERVICE_ACCOUNT,
        ServiceConnectorRequest: ResourceType.SERVICE_CONNECTOR,
        ServiceConnectorResponse: ResourceType.SERVICE_CONNECTOR,
        ServiceRequest: ResourceType.SERVICE,
        ServiceResponse: ResourceType.SERVICE,
        StackRequest: ResourceType.STACK,
        StackResponse: ResourceType.STACK,
        TagRequest: ResourceType.TAG,
        TagResponse: ResourceType.TAG,
        TriggerRequest: ResourceType.TRIGGER,
        TriggerResponse: ResourceType.TRIGGER,
        TriggerExecutionRequest: ResourceType.TRIGGER_EXECUTION,
        TriggerExecutionResponse: ResourceType.TRIGGER_EXECUTION,
        ProjectResponse: ResourceType.PROJECT,
        ProjectRequest: ResourceType.PROJECT,
        # UserResponse: ResourceType.USER,
    }

    return mapping.get(type(model))
get_schema_for_resource_type(resource_type: ResourceType) -> Type[BaseSchema]

Get the database schema for a resource type.

Parameters:

Name Type Description Default
resource_type ResourceType

The resource type for which to get the database schema.

required

Returns:

Type Description
Type[BaseSchema]

The database schema.

Source code in src/zenml/zen_server/rbac/utils.py
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
def get_schema_for_resource_type(
    resource_type: ResourceType,
) -> Type["BaseSchema"]:
    """Get the database schema for a resource type.

    Args:
        resource_type: The resource type for which to get the database schema.

    Returns:
        The database schema.
    """
    from zenml.zen_stores.schemas import (
        ActionSchema,
        ArtifactSchema,
        ArtifactVersionSchema,
        CodeRepositorySchema,
        EventSourceSchema,
        FlavorSchema,
        ModelSchema,
        ModelVersionSchema,
        PipelineBuildSchema,
        PipelineDeploymentSchema,
        PipelineRunSchema,
        PipelineSchema,
        RunMetadataSchema,
        RunTemplateSchema,
        SecretSchema,
        ServiceConnectorSchema,
        ServiceSchema,
        StackComponentSchema,
        StackSchema,
        TagSchema,
        TriggerExecutionSchema,
        TriggerSchema,
        UserSchema,
    )

    mapping: Dict[ResourceType, Type["BaseSchema"]] = {
        ResourceType.STACK: StackSchema,
        ResourceType.FLAVOR: FlavorSchema,
        ResourceType.STACK_COMPONENT: StackComponentSchema,
        ResourceType.PIPELINE: PipelineSchema,
        ResourceType.CODE_REPOSITORY: CodeRepositorySchema,
        ResourceType.MODEL: ModelSchema,
        ResourceType.MODEL_VERSION: ModelVersionSchema,
        ResourceType.SERVICE_CONNECTOR: ServiceConnectorSchema,
        ResourceType.ARTIFACT: ArtifactSchema,
        ResourceType.ARTIFACT_VERSION: ArtifactVersionSchema,
        ResourceType.SECRET: SecretSchema,
        ResourceType.SERVICE: ServiceSchema,
        ResourceType.TAG: TagSchema,
        ResourceType.SERVICE_ACCOUNT: UserSchema,
        # ResourceType.PROJECT: ProjectSchema,
        ResourceType.PIPELINE_RUN: PipelineRunSchema,
        ResourceType.PIPELINE_DEPLOYMENT: PipelineDeploymentSchema,
        ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
        ResourceType.RUN_TEMPLATE: RunTemplateSchema,
        ResourceType.RUN_METADATA: RunMetadataSchema,
        # ResourceType.USER: UserSchema,
        ResourceType.ACTION: ActionSchema,
        ResourceType.EVENT_SOURCE: EventSourceSchema,
        ResourceType.TRIGGER: TriggerSchema,
        ResourceType.TRIGGER_EXECUTION: TriggerExecutionSchema,
    }

    return mapping[resource_type]
get_subresources_for_model(model: AnyModel) -> Set[Resource]

Get all sub-resources of a model which need permission verification.

Parameters:

Name Type Description Default
model AnyModel

The model for which to get all the resources.

required

Returns:

Type Description
Set[Resource]

All resources of a model which need permission verification.

Source code in src/zenml/zen_server/rbac/utils.py
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
def get_subresources_for_model(
    model: AnyModel,
) -> Set[Resource]:
    """Get all sub-resources of a model which need permission verification.

    Args:
        model: The model for which to get all the resources.

    Returns:
        All resources of a model which need permission verification.
    """
    resources = set()

    # We don't want to use `model.model_dump()` here as that recursively
    # converts models to dicts, but we want to preserve those classes for
    # the recursive `_get_subresources_for_value` calls.
    # We previously used `dict(model)` here, but that lead to issues with
    # models overwriting `__getattr__`, this `model.__iter__()` has the same
    # results though.
    if isinstance(model, Page):
        for item in model:
            resources.update(_get_subresources_for_value(item))
    else:
        for _, value in model.__iter__():
            resources.update(_get_subresources_for_value(value))

    return resources
get_surrogate_permission_model_for_model(model: BaseModel, action: str) -> BaseModel

Get a surrogate permission model for a model.

In some cases a different model instead of the original model is used to verify permissions. For example, a parent container model might be used to verify permissions for all its children.

Parameters:

Name Type Description Default
model BaseModel

The original model.

required
action str

The action that the user wants to perform on the model.

required

Returns:

Type Description
BaseModel

A surrogate model or the original.

Source code in src/zenml/zen_server/rbac/utils.py
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
def get_surrogate_permission_model_for_model(
    model: BaseModel, action: str
) -> BaseModel:
    """Get a surrogate permission model for a model.

    In some cases a different model instead of the original model is used to
    verify permissions. For example, a parent container model might be used
    to verify permissions for all its children.

    Args:
        model: The original model.
        action: The action that the user wants to perform on the model.

    Returns:
        A surrogate model or the original.
    """
    from zenml.models import ArtifactVersionResponse, ModelVersionResponse

    # Permissions to read entities that represent versions of another entity
    # are checked on the parent entity
    if action == Action.READ:
        if isinstance(model, ModelVersionResponse):
            return model.model
        elif isinstance(model, ArtifactVersionResponse):
            return model.artifact

    return model
has_permissions_for_model(model: AnyModel, action: Action) -> bool

If the active user has permissions to perform the action on the model.

Parameters:

Name Type Description Default
model AnyModel

The model the user wants to perform the action on.

required
action Action

The action the user wants to perform.

required

Returns:

Type Description
bool

If the active user has permissions to perform the action on the model.

Source code in src/zenml/zen_server/rbac/utils.py
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def has_permissions_for_model(model: AnyModel, action: Action) -> bool:
    """If the active user has permissions to perform the action on the model.

    Args:
        model: The model the user wants to perform the action on.
        action: The action the user wants to perform.

    Returns:
        If the active user has permissions to perform the action on the model.
    """
    if is_owned_by_authenticated_user(model):
        return True

    try:
        verify_permission_for_model(model=model, action=action)
        return True
    except IllegalOperationError:
        return False
is_owned_by_authenticated_user(model: AnyModel) -> bool

Returns whether the currently authenticated user owns the model.

Parameters:

Name Type Description Default
model AnyModel

The model for which to check the ownership.

required

Returns:

Type Description
bool

Whether the currently authenticated user owns the model.

Source code in src/zenml/zen_server/rbac/utils.py
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
def is_owned_by_authenticated_user(model: AnyModel) -> bool:
    """Returns whether the currently authenticated user owns the model.

    Args:
        model: The model for which to check the ownership.

    Returns:
        Whether the currently authenticated user owns the model.
    """
    auth_context = get_auth_context()
    assert auth_context

    if isinstance(model, UserScopedResponse):
        if model.user:
            return model.user.id == auth_context.user.id
        else:
            # The model is server-owned and for RBAC purposes we consider
            # every user to be the owner of it
            return True

    return False
update_resource_membership(user: UserResponse, resource: Resource, actions: List[Action]) -> None

Update the resource membership of a user.

Parameters:

Name Type Description Default
user UserResponse

User for which the resource membership should be updated.

required
resource Resource

The resource.

required
actions List[Action]

The actions that the user should be able to perform on the resource.

required
Source code in src/zenml/zen_server/rbac/utils.py
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
def update_resource_membership(
    user: UserResponse, resource: Resource, actions: List[Action]
) -> None:
    """Update the resource membership of a user.

    Args:
        user: User for which the resource membership should be updated.
        resource: The resource.
        actions: The actions that the user should be able to perform on the
            resource.
    """
    if not server_config().rbac_enabled:
        return

    rbac().update_resource_membership(
        user=user, resource=resource, actions=actions
    )
verify_permission(resource_type: str, action: Action, resource_id: Optional[UUID] = None, project_id: Optional[UUID] = None) -> None

Verifies if a user has permission to perform an action on a resource.

Parameters:

Name Type Description Default
resource_type str

The type of resource that the user wants to perform the action on.

required
action Action

The action the user wants to perform.

required
resource_id Optional[UUID]

ID of the resource the user wants to perform the action on.

None
project_id Optional[UUID]

ID of the project the user wants to perform the action on. Only used for project scoped resources.

None
Source code in src/zenml/zen_server/rbac/utils.py
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
def verify_permission(
    resource_type: str,
    action: Action,
    resource_id: Optional[UUID] = None,
    project_id: Optional[UUID] = None,
) -> None:
    """Verifies if a user has permission to perform an action on a resource.

    Args:
        resource_type: The type of resource that the user wants to perform the
            action on.
        action: The action the user wants to perform.
        resource_id: ID of the resource the user wants to perform the action on.
        project_id: ID of the project the user wants to perform the action
            on. Only used for project scoped resources.
    """
    resource = Resource(
        type=resource_type, id=resource_id, project_id=project_id
    )
    batch_verify_permissions(resources={resource}, action=action)
verify_permission_for_model(model: AnyModel, action: Action) -> None

Verifies if a user has permission to perform an action on a model.

Parameters:

Name Type Description Default
model AnyModel

The model the user wants to perform the action on.

required
action Action

The action the user wants to perform.

required
Source code in src/zenml/zen_server/rbac/utils.py
249
250
251
252
253
254
255
256
def verify_permission_for_model(model: AnyModel, action: Action) -> None:
    """Verifies if a user has permission to perform an action on a model.

    Args:
        model: The model the user wants to perform the action on.
        action: The action the user wants to perform.
    """
    batch_verify_permissions_for_models(models=[model], action=action)
zenml_cloud_rbac

Cloud RBAC implementation.

Classes
ZenMLCloudRBAC()

Bases: RBACInterface

RBAC implementation that uses the ZenML Pro Management Plane as a backend.

Initialize the object.

Source code in src/zenml/zen_server/rbac/zenml_cloud_rbac.py
35
36
37
def __init__(self) -> None:
    """Initialize the object."""
    self._connection = cloud_connection()
Functions
check_permissions(user: UserResponse, resources: Set[Resource], action: Action) -> Dict[Resource, bool]

Checks if a user has permissions to perform an action on resources.

Parameters:

Name Type Description Default
user UserResponse

User which wants to access a resource.

required
resources Set[Resource]

The resources the user wants to access.

required
action Action

The action that the user wants to perform on the resources.

required

Returns:

Type Description
Dict[Resource, bool]

A dictionary mapping resources to a boolean which indicates whether

Dict[Resource, bool]

the user has permissions to perform the action on that resource.

Source code in src/zenml/zen_server/rbac/zenml_cloud_rbac.py
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def check_permissions(
    self, user: "UserResponse", resources: Set[Resource], action: Action
) -> Dict[Resource, bool]:
    """Checks if a user has permissions to perform an action on resources.

    Args:
        user: User which wants to access a resource.
        resources: The resources the user wants to access.
        action: The action that the user wants to perform on the resources.

    Returns:
        A dictionary mapping resources to a boolean which indicates whether
        the user has permissions to perform the action on that resource.
    """
    if not resources:
        # No need to send a request if there are no resources
        return {}

    if user.is_service_account:
        # Service accounts have full permissions for now
        return {resource: True for resource in resources}

    # At this point it's a regular user, which in a ZenML Pro with RBAC
    # enabled is always authenticated using external authentication
    assert user.external_user_id

    params = {
        "user_id": str(user.external_user_id),
        "resources": resources,
        "action": str(action),
    }
    response = self._connection.get(
        endpoint=PERMISSIONS_ENDPOINT, params=params
    )
    value = response.json()

    assert isinstance(value, dict)
    return {Resource.parse(k): v for k, v in value.items()}
delete_resources(resources: List[Resource]) -> None

Delete resource membership information for a list of resources.

Parameters:

Name Type Description Default
resources List[Resource]

The resources for which to delete the resource membership information.

required
Source code in src/zenml/zen_server/rbac/zenml_cloud_rbac.py
141
142
143
144
145
146
147
148
149
150
151
def delete_resources(self, resources: List[Resource]) -> None:
    """Delete resource membership information for a list of resources.

    Args:
        resources: The resources for which to delete the resource membership
            information.
    """
    params = {
        "resources": [str(resource) for resource in resources],
    }
    self._connection.delete(endpoint=RESOURCES_ENDPOINT, params=params)
list_allowed_resource_ids(user: UserResponse, resource: Resource, action: Action) -> Tuple[bool, List[str]]

Lists all resource IDs of a resource type that a user can access.

Parameters:

Name Type Description Default
user UserResponse

User which wants to access a resource.

required
resource Resource

The resource the user wants to access.

required
action Action

The action that the user wants to perform on the resource.

required

Returns:

Type Description
bool

A tuple (full_resource_access, resource_ids).

List[str]

full_resource_access will be True if the user can perform the

Tuple[bool, List[str]]

given action on any instance of the given resource type, False

Tuple[bool, List[str]]

otherwise. If full_resource_access is False, resource_ids

Tuple[bool, List[str]]

will contain the list of instance IDs that the user can perform

Tuple[bool, List[str]]

the action on.

Source code in src/zenml/zen_server/rbac/zenml_cloud_rbac.py
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def list_allowed_resource_ids(
    self, user: "UserResponse", resource: Resource, action: Action
) -> Tuple[bool, List[str]]:
    """Lists all resource IDs of a resource type that a user can access.

    Args:
        user: User which wants to access a resource.
        resource: The resource the user wants to access.
        action: The action that the user wants to perform on the resource.

    Returns:
        A tuple (full_resource_access, resource_ids).
        `full_resource_access` will be `True` if the user can perform the
        given action on any instance of the given resource type, `False`
        otherwise. If `full_resource_access` is `False`, `resource_ids`
        will contain the list of instance IDs that the user can perform
        the action on.
    """
    assert not resource.id
    if user.is_service_account:
        # Service accounts have full permissions for now
        return True, []

    # At this point it's a regular user, which in the ZenML Pro with RBAC
    # enabled is always authenticated using external authentication
    assert user.external_user_id
    params = {
        "user_id": str(user.external_user_id),
        "resource": str(resource),
        "action": str(action),
    }
    response = self._connection.get(
        endpoint=ALLOWED_RESOURCE_IDS_ENDPOINT, params=params
    )
    response_json = response.json()

    full_resource_access: bool = response_json["full_access"]
    allowed_ids: List[str] = response_json["ids"]

    return full_resource_access, allowed_ids
update_resource_membership(user: UserResponse, resource: Resource, actions: List[Action]) -> None

Update the resource membership of a user.

Parameters:

Name Type Description Default
user UserResponse

User for which the resource membership should be updated.

required
resource Resource

The resource.

required
actions List[Action]

The actions that the user should be able to perform on the resource.

required
Source code in src/zenml/zen_server/rbac/zenml_cloud_rbac.py
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def update_resource_membership(
    self, user: "UserResponse", resource: Resource, actions: List[Action]
) -> None:
    """Update the resource membership of a user.

    Args:
        user: User for which the resource membership should be updated.
        resource: The resource.
        actions: The actions that the user should be able to perform on the
            resource.
    """
    if user.is_service_account:
        # Service accounts have full permissions for now
        return

    data = {
        "user_id": str(user.external_user_id),
        "resource": str(resource),
        "actions": [str(action) for action in actions],
    }
    self._connection.post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)
Functions

routers

Endpoint definitions.

Modules
actions_endpoints

Endpoint definitions for actions.

Classes Functions
create_action(action: ActionRequest, _: AuthContext = Security(authorize)) -> ActionResponse

Creates an action.

Parameters:

Name Type Description Default
action ActionRequest

Action to create.

required

Raises:

Type Description
ValueError

If the action handler for flavor/type is not valid.

Returns:

Type Description
ActionResponse

The created action.

Source code in src/zenml/zen_server/routers/actions_endpoints.py
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_action(
    action: ActionRequest,
    _: AuthContext = Security(authorize),
) -> ActionResponse:
    """Creates an action.

    Args:
        action: Action to create.

    Raises:
        ValueError: If the action handler for flavor/type is not valid.

    Returns:
        The created action.
    """
    service_account = zen_store().get_service_account(
        service_account_name_or_id=action.service_account_id
    )
    verify_permission_for_model(service_account, action=Action.READ)

    action_handler = plugin_flavor_registry().get_plugin(
        name=action.flavor,
        _type=PluginType.ACTION,
        subtype=action.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an action
    # handler implementation
    if not isinstance(action_handler, BaseActionHandler):
        raise ValueError(
            f"Action handler plugin {action.plugin_subtype} "
            f"for flavor {action.flavor} is not a valid action "
            "handler plugin."
        )

    return verify_permissions_and_create_entity(
        request_model=action,
        create_method=action_handler.create_action,
    )
delete_action(action_id: UUID, force: bool = False, _: AuthContext = Security(authorize)) -> None

Delete an action.

Parameters:

Name Type Description Default
action_id UUID

ID of the action.

required
force bool

Flag deciding whether to force delete the action.

False

Raises:

Type Description
ValueError

If the action handler for flavor/type is not valid.

Source code in src/zenml/zen_server/routers/actions_endpoints.py
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
@router.delete(
    "/{action_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_action(
    action_id: UUID,
    force: bool = False,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete an action.

    Args:
        action_id: ID of the action.
        force: Flag deciding whether to force delete the action.

    Raises:
        ValueError: If the action handler for flavor/type is not valid.
    """
    action = zen_store().get_action(action_id=action_id)

    verify_permission_for_model(action, action=Action.DELETE)

    action_handler = plugin_flavor_registry().get_plugin(
        name=action.flavor,
        _type=PluginType.ACTION,
        subtype=action.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an action
    # handler implementation
    if not isinstance(action_handler, BaseActionHandler):
        raise ValueError(
            f"Action handler plugin {action.plugin_subtype} "
            f"for flavor {action.flavor} is not a valid action "
            "handler plugin."
        )

    action_handler.delete_action(
        action=action,
        force=force,
    )

    delete_model_resource(action)
get_action(action_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ActionResponse

Returns the requested action.

Parameters:

Name Type Description Default
action_id UUID

ID of the action.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Raises:

Type Description
ValueError

If the action handler for flavor/type is not valid.

Returns:

Type Description
ActionResponse

The requested action.

Source code in src/zenml/zen_server/routers/actions_endpoints.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
@router.get(
    "/{action_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_action(
    action_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ActionResponse:
    """Returns the requested action.

    Args:
        action_id: ID of the action.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Raises:
        ValueError: If the action handler for flavor/type is not valid.

    Returns:
        The requested action.
    """
    action = zen_store().get_action(action_id=action_id, hydrate=hydrate)

    verify_permission_for_model(action, action=Action.READ)

    action_handler = plugin_flavor_registry().get_plugin(
        name=action.flavor,
        _type=PluginType.ACTION,
        subtype=action.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an action
    # handler implementation
    if not isinstance(action_handler, BaseActionHandler):
        raise ValueError(
            f"Action handler plugin {action.plugin_subtype} "
            f"for flavor {action.flavor} is not a valid action "
            "handler plugin."
        )

    action = action_handler.get_action(action, hydrate=hydrate)

    return dehydrate_response_model(action)
list_actions(action_filter_model: ActionFilter = Depends(make_dependable(ActionFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ActionResponse]

List actions.

Parameters:

Name Type Description Default
action_filter_model ActionFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ActionFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ActionResponse]

Page of actions.

Source code in src/zenml/zen_server/routers/actions_endpoints.py
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_actions(
    action_filter_model: ActionFilter = Depends(make_dependable(ActionFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ActionResponse]:
    """List actions.

    Args:
        action_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page of actions.
    """

    def list_actions_fn(
        filter_model: ActionFilter,
    ) -> Page[ActionResponse]:
        """List actions through their associated plugins.

        Args:
            filter_model: Filter model used for pagination, sorting,
                filtering.

        Raises:
            ValueError: If the action handler for flavor/type is not valid.

        Returns:
            All actions.
        """
        actions = zen_store().list_actions(
            action_filter_model=filter_model, hydrate=hydrate
        )

        # Process the actions through their associated plugins
        for idx, action in enumerate(actions.items):
            action_handler = plugin_flavor_registry().get_plugin(
                name=action.flavor,
                _type=PluginType.ACTION,
                subtype=action.plugin_subtype,
            )

            # Validate that the flavor and plugin_type correspond to an action
            # handler implementation
            if not isinstance(action_handler, BaseActionHandler):
                raise ValueError(
                    f"Action handler plugin {action.plugin_subtype} "
                    f"for flavor {action.flavor} is not a valid action "
                    "handler plugin."
                )

            actions.items[idx] = action_handler.get_action(
                action, hydrate=hydrate
            )

        return actions

    return verify_permissions_and_list_entities(
        filter_model=action_filter_model,
        resource_type=ResourceType.ACTION,
        list_method=list_actions_fn,
    )
update_action(action_id: UUID, action_update: ActionUpdate, _: AuthContext = Security(authorize)) -> ActionResponse

Update an action.

Parameters:

Name Type Description Default
action_id UUID

ID of the action to update.

required
action_update ActionUpdate

The action update.

required

Raises:

Type Description
ValueError

If the action handler for flavor/type is not valid.

Returns:

Type Description
ActionResponse

The updated action.

Source code in src/zenml/zen_server/routers/actions_endpoints.py
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
@router.put(
    "/{action_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_action(
    action_id: UUID,
    action_update: ActionUpdate,
    _: AuthContext = Security(authorize),
) -> ActionResponse:
    """Update an action.

    Args:
        action_id: ID of the action to update.
        action_update: The action update.

    Raises:
        ValueError: If the action handler for flavor/type is not valid.

    Returns:
        The updated action.
    """
    action = zen_store().get_action(action_id=action_id)

    verify_permission_for_model(action, action=Action.UPDATE)

    if action_update.service_account_id:
        service_account = zen_store().get_service_account(
            service_account_name_or_id=action_update.service_account_id
        )
        verify_permission_for_model(service_account, action=Action.READ)

    action_handler = plugin_flavor_registry().get_plugin(
        name=action.flavor,
        _type=PluginType.ACTION,
        subtype=action.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an action
    # handler implementation
    if not isinstance(action_handler, BaseActionHandler):
        raise ValueError(
            f"Action handler plugin {action.plugin_subtype} "
            f"for flavor {action.flavor} is not a valid action "
            "handler plugin."
        )

    updated_action = action_handler.update_action(
        action=action,
        action_update=action_update,
    )

    return dehydrate_response_model(updated_action)
artifact_endpoint

Endpoint definitions for artifacts.

Classes Functions
create_artifact(artifact: ArtifactRequest, _: AuthContext = Security(authorize)) -> ArtifactResponse

Create a new artifact.

Parameters:

Name Type Description Default
artifact ArtifactRequest

The artifact to create.

required

Returns:

Type Description
ArtifactResponse

The created artifact.

Source code in src/zenml/zen_server/routers/artifact_endpoint.py
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
@artifact_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_artifact(
    artifact: ArtifactRequest,
    _: AuthContext = Security(authorize),
) -> ArtifactResponse:
    """Create a new artifact.

    Args:
        artifact: The artifact to create.

    Returns:
        The created artifact.
    """
    return verify_permissions_and_create_entity(
        request_model=artifact,
        create_method=zen_store().create_artifact,
    )
delete_artifact(artifact_id: UUID, _: AuthContext = Security(authorize)) -> None

Delete an artifact by ID.

Parameters:

Name Type Description Default
artifact_id UUID

The ID of the artifact to delete.

required
Source code in src/zenml/zen_server/routers/artifact_endpoint.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@artifact_router.delete(
    "/{artifact_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_artifact(
    artifact_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete an artifact by ID.

    Args:
        artifact_id: The ID of the artifact to delete.
    """
    verify_permissions_and_delete_entity(
        id=artifact_id,
        get_method=zen_store().get_artifact,
        delete_method=zen_store().delete_artifact,
    )
get_artifact(artifact_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ArtifactResponse

Get an artifact by ID.

Parameters:

Name Type Description Default
artifact_id UUID

The ID of the artifact to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ArtifactResponse

The artifact with the given ID.

Source code in src/zenml/zen_server/routers/artifact_endpoint.py
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
@artifact_router.get(
    "/{artifact_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_artifact(
    artifact_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ArtifactResponse:
    """Get an artifact by ID.

    Args:
        artifact_id: The ID of the artifact to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifact with the given ID.
    """
    return verify_permissions_and_get_entity(
        id=artifact_id,
        get_method=zen_store().get_artifact,
        hydrate=hydrate,
    )
list_artifacts(artifact_filter_model: ArtifactFilter = Depends(make_dependable(ArtifactFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ArtifactResponse]

Get artifacts according to query filters.

Parameters:

Name Type Description Default
artifact_filter_model ArtifactFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ArtifactFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ArtifactResponse]

The artifacts according to query filters.

Source code in src/zenml/zen_server/routers/artifact_endpoint.py
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@artifact_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_artifacts(
    artifact_filter_model: ArtifactFilter = Depends(
        make_dependable(ArtifactFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ArtifactResponse]:
    """Get artifacts according to query filters.

    Args:
        artifact_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifacts according to query filters.
    """
    return verify_permissions_and_list_entities(
        filter_model=artifact_filter_model,
        resource_type=ResourceType.ARTIFACT,
        list_method=zen_store().list_artifacts,
        hydrate=hydrate,
    )
update_artifact(artifact_id: UUID, artifact_update: ArtifactUpdate, _: AuthContext = Security(authorize)) -> ArtifactResponse

Update an artifact by ID.

Parameters:

Name Type Description Default
artifact_id UUID

The ID of the artifact to update.

required
artifact_update ArtifactUpdate

The update to apply to the artifact.

required

Returns:

Type Description
ArtifactResponse

The updated artifact.

Source code in src/zenml/zen_server/routers/artifact_endpoint.py
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
@artifact_router.put(
    "/{artifact_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_artifact(
    artifact_id: UUID,
    artifact_update: ArtifactUpdate,
    _: AuthContext = Security(authorize),
) -> ArtifactResponse:
    """Update an artifact by ID.

    Args:
        artifact_id: The ID of the artifact to update.
        artifact_update: The update to apply to the artifact.

    Returns:
        The updated artifact.
    """
    return verify_permissions_and_update_entity(
        id=artifact_id,
        update_model=artifact_update,
        get_method=zen_store().get_artifact,
        update_method=zen_store().update_artifact,
    )
artifact_version_endpoints

Endpoint definitions for artifact versions.

Classes Functions
batch_create_artifact_version(artifact_versions: List[ArtifactVersionRequest], _: AuthContext = Security(authorize)) -> List[ArtifactVersionResponse]

Create a batch of artifact versions.

Parameters:

Name Type Description Default
artifact_versions List[ArtifactVersionRequest]

The artifact versions to create.

required

Returns:

Type Description
List[ArtifactVersionResponse]

The created artifact versions.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@artifact_version_router.post(
    BATCH,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def batch_create_artifact_version(
    artifact_versions: List[ArtifactVersionRequest],
    _: AuthContext = Security(authorize),
) -> List[ArtifactVersionResponse]:
    """Create a batch of artifact versions.

    Args:
        artifact_versions: The artifact versions to create.

    Returns:
        The created artifact versions.
    """
    return verify_permissions_and_batch_create_entity(
        batch=artifact_versions,
        create_method=zen_store().batch_create_artifact_versions,
    )
create_artifact_version(artifact_version: ArtifactVersionRequest, _: AuthContext = Security(authorize)) -> ArtifactVersionResponse

Create a new artifact version.

Parameters:

Name Type Description Default
artifact_version ArtifactVersionRequest

The artifact version to create.

required

Returns:

Type Description
ArtifactVersionResponse

The created artifact version.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@artifact_version_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_artifact_version(
    artifact_version: ArtifactVersionRequest,
    _: AuthContext = Security(authorize),
) -> ArtifactVersionResponse:
    """Create a new artifact version.

    Args:
        artifact_version: The artifact version to create.

    Returns:
        The created artifact version.
    """
    return verify_permissions_and_create_entity(
        request_model=artifact_version,
        create_method=zen_store().create_artifact_version,
    )
delete_artifact_version(artifact_version_id: UUID, _: AuthContext = Security(authorize)) -> None

Delete an artifact version by ID.

Parameters:

Name Type Description Default
artifact_version_id UUID

The ID of the artifact version to delete.

required
Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
@artifact_version_router.delete(
    "/{artifact_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_artifact_version(
    artifact_version_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete an artifact version by ID.

    Args:
        artifact_version_id: The ID of the artifact version to delete.
    """
    verify_permissions_and_delete_entity(
        id=artifact_version_id,
        get_method=zen_store().get_artifact_version,
        delete_method=zen_store().delete_artifact_version,
    )
get_artifact_version(artifact_version_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ArtifactVersionResponse

Get an artifact version by ID.

Parameters:

Name Type Description Default
artifact_version_id UUID

The ID of the artifact version to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ArtifactVersionResponse

The artifact version with the given ID.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@artifact_version_router.get(
    "/{artifact_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_artifact_version(
    artifact_version_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ArtifactVersionResponse:
    """Get an artifact version by ID.

    Args:
        artifact_version_id: The ID of the artifact version to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifact version with the given ID.
    """
    return verify_permissions_and_get_entity(
        id=artifact_version_id,
        get_method=zen_store().get_artifact_version,
        hydrate=hydrate,
    )
get_artifact_visualization(artifact_version_id: UUID, index: int = 0, _: AuthContext = Security(authorize)) -> LoadedVisualization

Get the visualization of an artifact.

Parameters:

Name Type Description Default
artifact_version_id UUID

ID of the artifact version for which to get the visualization.

required
index int

Index of the visualization to get (if there are multiple).

0

Returns:

Type Description
LoadedVisualization

The visualization of the artifact version.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
@artifact_version_router.get(
    "/{artifact_version_id}" + VISUALIZE,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_artifact_visualization(
    artifact_version_id: UUID,
    index: int = 0,
    _: AuthContext = Security(authorize),
) -> LoadedVisualization:
    """Get the visualization of an artifact.

    Args:
        artifact_version_id: ID of the artifact version for which to get the visualization.
        index: Index of the visualization to get (if there are multiple).

    Returns:
        The visualization of the artifact version.
    """
    store = zen_store()
    artifact = verify_permissions_and_get_entity(
        id=artifact_version_id, get_method=store.get_artifact_version
    )
    return load_artifact_visualization(
        artifact=artifact, index=index, zen_store=store, encode_image=True
    )
list_artifact_versions(artifact_version_filter_model: ArtifactVersionFilter = Depends(make_dependable(ArtifactVersionFilter)), hydrate: bool = False, auth_context: AuthContext = Security(authorize)) -> Page[ArtifactVersionResponse]

Get artifact versions according to query filters.

Parameters:

Name Type Description Default
artifact_version_filter_model ArtifactVersionFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ArtifactVersionFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
Page[ArtifactVersionResponse]

The artifact versions according to query filters.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@artifact_version_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_artifact_versions(
    artifact_version_filter_model: ArtifactVersionFilter = Depends(
        make_dependable(ArtifactVersionFilter)
    ),
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[ArtifactVersionResponse]:
    """Get artifact versions according to query filters.

    Args:
        artifact_version_filter_model: Filter model used for pagination,
            sorting, filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: The authentication context.

    Returns:
        The artifact versions according to query filters.
    """
    # A project scoped request must always be scoped to a specific
    # project. This is required for the RBAC check to work.
    set_filter_project_scope(artifact_version_filter_model)
    assert isinstance(artifact_version_filter_model.project, UUID)

    allowed_artifact_ids = get_allowed_resource_ids(
        resource_type=ResourceType.ARTIFACT,
        project_id=artifact_version_filter_model.project,
    )
    artifact_version_filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id,
        artifact_id=allowed_artifact_ids,
    )
    artifact_versions = zen_store().list_artifact_versions(
        artifact_version_filter_model=artifact_version_filter_model,
        hydrate=hydrate,
    )
    return dehydrate_page(artifact_versions)
prune_artifact_versions(project_name_or_id: Union[str, UUID], only_versions: bool = True, _: AuthContext = Security(authorize)) -> None

Prunes unused artifact versions and their artifacts.

Parameters:

Name Type Description Default
project_name_or_id Union[str, UUID]

The project name or ID to prune artifact versions for.

required
only_versions bool

Only delete artifact versions, keeping artifacts

True
Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
@artifact_version_router.delete(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def prune_artifact_versions(
    project_name_or_id: Union[str, UUID],
    only_versions: bool = True,
    _: AuthContext = Security(authorize),
) -> None:
    """Prunes unused artifact versions and their artifacts.

    Args:
        project_name_or_id: The project name or ID to prune artifact
            versions for.
        only_versions: Only delete artifact versions, keeping artifacts
    """
    project_id = zen_store().get_project(project_name_or_id).id

    verify_permissions_and_prune_entities(
        resource_type=ResourceType.ARTIFACT_VERSION,
        prune_method=zen_store().prune_artifact_versions,
        only_versions=only_versions,
        project_id=project_id,
    )
update_artifact_version(artifact_version_id: UUID, artifact_version_update: ArtifactVersionUpdate, _: AuthContext = Security(authorize)) -> ArtifactVersionResponse

Update an artifact by ID.

Parameters:

Name Type Description Default
artifact_version_id UUID

The ID of the artifact version to update.

required
artifact_version_update ArtifactVersionUpdate

The update to apply to the artifact version.

required

Returns:

Type Description
ArtifactVersionResponse

The updated artifact.

Source code in src/zenml/zen_server/routers/artifact_version_endpoints.py
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
@artifact_version_router.put(
    "/{artifact_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_artifact_version(
    artifact_version_id: UUID,
    artifact_version_update: ArtifactVersionUpdate,
    _: AuthContext = Security(authorize),
) -> ArtifactVersionResponse:
    """Update an artifact by ID.

    Args:
        artifact_version_id: The ID of the artifact version to update.
        artifact_version_update: The update to apply to the artifact version.

    Returns:
        The updated artifact.
    """
    return verify_permissions_and_update_entity(
        id=artifact_version_id,
        update_model=artifact_version_update,
        get_method=zen_store().get_artifact_version,
        update_method=zen_store().update_artifact_version,
    )
auth_endpoints

Endpoint definitions for authentication (login).

Classes
OAuthLoginRequestForm(grant_type: Optional[str] = Form(None), username: Optional[str] = Form(None), password: Optional[str] = Form(None), client_id: Optional[str] = Form(None), device_code: Optional[str] = Form(None))

OAuth2 grant type request form.

This form allows multiple grant types to be used with the same endpoint: * standard OAuth2 password grant type * standard OAuth2 device authorization grant type * ZenML service account + API key grant type (proprietary) * ZenML External Authenticator grant type (proprietary)

Initializes the form.

Parameters:

Name Type Description Default
grant_type Optional[str]

The grant type.

Form(None)
username Optional[str]

The username. Only used for the password grant type.

Form(None)
password Optional[str]

The password. Only used for the password grant type.

Form(None)
client_id Optional[str]

The client ID.

Form(None)
device_code Optional[str]

The device code. Only used for the device authorization grant type.

Form(None)

Raises:

Type Description
HTTPException

If the request is invalid.

Source code in src/zenml/zen_server/routers/auth_endpoints.py
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def __init__(
    self,
    grant_type: Optional[str] = Form(None),
    username: Optional[str] = Form(None),
    password: Optional[str] = Form(None),
    client_id: Optional[str] = Form(None),
    device_code: Optional[str] = Form(None),
):
    """Initializes the form.

    Args:
        grant_type: The grant type.
        username: The username. Only used for the password grant type.
        password: The password. Only used for the password grant type.
        client_id: The client ID.
        device_code: The device code. Only used for the device authorization
            grant type.

    Raises:
        HTTPException: If the request is invalid.
    """
    config = server_config()

    if not grant_type:
        # Detect the grant type from the form data
        if username is not None:
            self.grant_type = OAuthGrantTypes.OAUTH_PASSWORD
        elif password:
            self.grant_type = OAuthGrantTypes.ZENML_API_KEY
        elif device_code:
            self.grant_type = OAuthGrantTypes.OAUTH_DEVICE_CODE
        elif config.auth_scheme == AuthScheme.EXTERNAL:
            self.grant_type = OAuthGrantTypes.ZENML_EXTERNAL
        elif config.auth_scheme in [
            AuthScheme.OAUTH2_PASSWORD_BEARER,
            AuthScheme.NO_AUTH,
            AuthScheme.HTTP_BASIC,
        ]:
            # For no auth and basic HTTP auth schemes, we also allow the
            # password grant type to be used for backwards compatibility
            self.grant_type = OAuthGrantTypes.OAUTH_PASSWORD
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Invalid request: grant type is required.",
            )
    else:
        if grant_type not in OAuthGrantTypes.values():
            logger.info(
                f"Request with unsupported grant type: {grant_type}"
            )
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unsupported grant type: {grant_type}",
            )
        self.grant_type = OAuthGrantTypes(grant_type)

    if self.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
        # For the no auth and basic HTTP auth schemes, we also allow the
        # password grant type to be used for compatibility with other
        # auth schemes
        if config.auth_scheme not in [
            AuthScheme.OAUTH2_PASSWORD_BEARER,
            AuthScheme.NO_AUTH,
            AuthScheme.HTTP_BASIC,
        ]:
            logger.info(
                f"Request with unsupported grant type: {self.grant_type}"
            )
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unsupported grant type: {self.grant_type}.",
            )
        if not username:
            logger.info("Request with missing username")
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Invalid request: username is required.",
            )
        self.username = username
        self.password = password or ""
    elif self.grant_type == OAuthGrantTypes.OAUTH_DEVICE_CODE:
        if config.auth_scheme in [AuthScheme.NO_AUTH, AuthScheme.EXTERNAL]:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Device authorization is not supported.",
            )

        if not device_code or not client_id:
            logger.info("Request with missing device code or client ID")
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Invalid request: device code and client ID are "
                "required.",
            )
        try:
            self.client_id = UUID(client_id)
        except ValueError:
            logger.info(f"Request with invalid client ID: {client_id}")
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Invalid request: invalid client ID.",
            )
        self.device_code = device_code
    elif self.grant_type == OAuthGrantTypes.ZENML_API_KEY:
        if not password:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="API key is required.",
            )
        self.api_key = password
    elif self.grant_type == OAuthGrantTypes.ZENML_EXTERNAL:
        if config.auth_scheme != AuthScheme.EXTERNAL:
            logger.info(
                f"Request with unsupported grant type: {self.grant_type}"
            )
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unsupported grant type: {self.grant_type}.",
            )
Functions
Functions
api_token(token_type: APITokenType = APITokenType.GENERIC, expires_in: Optional[int] = None, schedule_id: Optional[UUID] = None, pipeline_run_id: Optional[UUID] = None, step_run_id: Optional[UUID] = None, auth_context: AuthContext = Security(authorize)) -> str

Generate an API token for the current user.

Use this endpoint to generate an API token for the current user. Two types of API tokens are supported:

  • Generic API token: This token is short-lived and can be used for generic automation tasks. The expiration can be set by the user, but the server will impose a maximum expiration time.
  • Workload API token: This token is scoped to a specific pipeline run, step run or schedule and is used by pipeline workloads to authenticate with the server. A pipeline run ID, step run ID or schedule ID must be provided and the generated token will only be valid for the indicated pipeline run, step run or schedule. No time limit is imposed on the validity of the token. A workload API token can be used to authenticate and generate another workload API token, but only for the same schedule, pipeline run ID or step run ID, in that order.

Parameters:

Name Type Description Default
token_type APITokenType

The type of API token to generate.

GENERIC
expires_in Optional[int]

The expiration time of the generic API token in seconds. If not set, the server will use the default expiration time for generic API tokens. The server also imposes a maximum expiration time.

None
schedule_id Optional[UUID]

The ID of the schedule to scope the workload API token to.

None
pipeline_run_id Optional[UUID]

The ID of the pipeline run to scope the workload API token to.

None
step_run_id Optional[UUID]

The ID of the step run to scope the workload API token to.

None
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
str

The API token.

Raises:

Type Description
AuthorizationException

If not authorized to generate the API token.

ValueError

If the request is invalid.

Source code in src/zenml/zen_server/routers/auth_endpoints.py
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
@router.get(
    API_TOKEN,
    response_model=str,
)
@handle_exceptions
def api_token(
    token_type: APITokenType = APITokenType.GENERIC,
    expires_in: Optional[int] = None,
    schedule_id: Optional[UUID] = None,
    pipeline_run_id: Optional[UUID] = None,
    step_run_id: Optional[UUID] = None,
    auth_context: AuthContext = Security(authorize),
) -> str:
    """Generate an API token for the current user.

    Use this endpoint to generate an API token for the current user. Two types
    of API tokens are supported:

    * Generic API token: This token is short-lived and can be used for
    generic automation tasks. The expiration can be set by the user, but the
    server will impose a maximum expiration time.
    * Workload API token: This token is scoped to a specific pipeline run, step
    run or schedule and is used by pipeline workloads to authenticate with the
    server. A pipeline run ID, step run ID or schedule ID must be provided and
    the generated token will only be valid for the indicated pipeline run, step
    run or schedule. No time limit is imposed on the validity of the token.
    A workload API token can be used to authenticate and generate another
    workload API token, but only for the same schedule, pipeline run ID or step
    run ID, in that order.

    Args:
        token_type: The type of API token to generate.
        expires_in: The expiration time of the generic API token in seconds.
            If not set, the server will use the default expiration time for
            generic API tokens. The server also imposes a maximum expiration
            time.
        schedule_id: The ID of the schedule to scope the workload API token to.
        pipeline_run_id: The ID of the pipeline run to scope the workload API
            token to.
        step_run_id: The ID of the step run to scope the workload API token to.
        auth_context: The authentication context.

    Returns:
        The API token.

    Raises:
        AuthorizationException: If not authorized to generate the API token.
        ValueError: If the request is invalid.
    """
    token = auth_context.access_token
    if not token or not auth_context.encoded_access_token:
        # Should not happen
        raise AuthorizationException("Not authenticated.")

    if token_type == APITokenType.GENERIC:
        if schedule_id or pipeline_run_id or step_run_id:
            raise ValueError(
                "Generic API tokens cannot be scoped to a schedule, pipeline "
                "run or step run."
            )

        config = server_config()

        if not expires_in:
            expires_in = config.generic_api_token_lifetime

        if expires_in > config.generic_api_token_max_lifetime:
            raise ValueError(
                f"The maximum expiration time for generic API tokens allowed "
                f"by this server is {config.generic_api_token_max_lifetime} "
                "seconds."
            )

        return generate_access_token(
            user_id=token.user_id,
            expires_in=expires_in,
            # Don't include the access token as a cookie in the response
            response=None,
        ).access_token

    schedule_id = schedule_id or token.schedule_id
    pipeline_run_id = pipeline_run_id or token.pipeline_run_id
    step_run_id = step_run_id or token.step_run_id

    if not pipeline_run_id and not schedule_id and not step_run_id:
        raise ValueError(
            "Workload API tokens must be scoped to a schedule, pipeline run "
            "or step run."
        )

    if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
        raise AuthorizationException(
            f"Unable to scope API token to schedule {schedule_id}. The "
            f"token used to authorize this request is already scoped to "
            f"schedule {token.schedule_id}."
        )

    if (
        pipeline_run_id
        and token.pipeline_run_id
        and pipeline_run_id != token.pipeline_run_id
    ):
        raise AuthorizationException(
            f"Unable to scope API token to pipeline run {pipeline_run_id}. The "
            f"token used to authorize this request is already scoped to "
            f"pipeline run {token.pipeline_run_id}."
        )

    if step_run_id and token.step_run_id and step_run_id != token.step_run_id:
        raise AuthorizationException(
            f"Unable to scope API token to step run {step_run_id}. The "
            f"token used to authorize this request is already scoped to "
            f"step run {token.step_run_id}."
        )

    project_id: Optional[UUID] = None

    if schedule_id:
        # The schedule must exist
        try:
            schedule = zen_store().get_schedule(schedule_id, hydrate=True)
        except KeyError:
            raise ValueError(
                f"Schedule {schedule_id} does not exist and API tokens cannot "
                "be generated for non-existent schedules for security reasons."
            )
        project_id = schedule.project.id

        if not schedule.active:
            raise ValueError(
                f"Schedule {schedule_id} is not active and API tokens cannot "
                "be generated for inactive schedules for security reasons."
            )

    if pipeline_run_id:
        # The pipeline run must exist and the run must not be concluded
        try:
            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=True)
        except KeyError:
            raise ValueError(
                f"Pipeline run {pipeline_run_id} does not exist and API tokens "
                "cannot be generated for non-existent pipeline runs for "
                "security reasons."
            )

        verify_permission_for_model(model=pipeline_run, action=Action.READ)

        project_id = pipeline_run.project.id

        if pipeline_run.status.is_finished:
            raise ValueError(
                f"The execution of pipeline run {pipeline_run_id} has already "
                "concluded and API tokens can no longer be generated for it "
                "for security reasons."
            )

    if step_run_id:
        # The step run must exist and the step must not be concluded
        try:
            step_run = zen_store().get_run_step(step_run_id, hydrate=False)
        except KeyError:
            raise ValueError(
                f"Step run {step_run_id} does not exist and API tokens cannot "
                "be generated for non-existent step runs for security reasons."
            )

        project_id = step_run.project.id

        if step_run.status.is_finished:
            raise ValueError(
                f"The execution of step run {step_run_id} has already "
                "concluded and API tokens can no longer be generated for it "
                "for security reasons."
            )

    assert project_id is not None
    verify_permission(
        resource_type=ResourceType.PIPELINE_RUN,
        action=Action.CREATE,
        project_id=project_id,
    )

    return generate_access_token(
        user_id=token.user_id,
        # Keep the original API key and device token scopes
        api_key=auth_context.api_key,
        device=auth_context.device,
        schedule_id=schedule_id,
        pipeline_run_id=pipeline_run_id,
        step_run_id=step_run_id,
        # Don't include the access token as a cookie in the response
        response=None,
        # Never expire the token
        expires_in=0,
    ).access_token
device_authorization(request: Request, client_id: UUID = Form(...)) -> OAuthDeviceAuthorizationResponse

OAuth2 device authorization endpoint.

This endpoint implements the OAuth2 device authorization grant flow as defined in https://tools.ietf.org/html/rfc8628. It is called to initiate the device authorization flow by requesting a device and user code for a given client ID.

For a new client ID, a new OAuth device is created, stored in the DB and returned to the client along with a pair of newly generated device and user codes. If a device for the given client ID already exists, the existing DB entry is reused and new device and user codes are generated.

Parameters:

Name Type Description Default
request Request

The request object.

required
client_id UUID

The client ID.

Form(...)

Returns:

Type Description
OAuthDeviceAuthorizationResponse

The device authorization response.

Raises:

Type Description
HTTPException

If the device authorization is not supported.

Source code in src/zenml/zen_server/routers/auth_endpoints.py
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
@router.post(
    DEVICE_AUTHORIZATION,
    response_model=OAuthDeviceAuthorizationResponse,
)
def device_authorization(
    request: Request,
    client_id: UUID = Form(...),
) -> OAuthDeviceAuthorizationResponse:
    """OAuth2 device authorization endpoint.

    This endpoint implements the OAuth2 device authorization grant flow as
    defined in https://tools.ietf.org/html/rfc8628. It is called to initiate
    the device authorization flow by requesting a device and user code for a
    given client ID.

    For a new client ID, a new OAuth device is created, stored in the DB and
    returned to the client along with a pair of newly generated device and user
    codes. If a device for the given client ID already exists, the existing
    DB entry is reused and new device and user codes are generated.

    Args:
        request: The request object.
        client_id: The client ID.

    Returns:
        The device authorization response.

    Raises:
        HTTPException: If the device authorization is not supported.
    """
    config = server_config()

    if config.auth_scheme in [AuthScheme.NO_AUTH, AuthScheme.EXTERNAL]:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Device authorization is not supported.",
        )

    store = zen_store()

    try:
        # Use this opportunity to delete expired devices
        store.delete_expired_authorized_devices()
    except Exception:
        logger.exception("Failed to delete expired devices")

    # Fetch additional details about the client from the user-agent header
    user_agent_header = request.headers.get("User-Agent")
    if user_agent_header:
        device_details = OAuthDeviceUserAgentHeader.decode(user_agent_header)
    else:
        device_details = OAuthDeviceUserAgentHeader()

    # Fetch the IP address of the client
    ip_address: str = ""
    city, region, country = "", "", ""
    forwarded = request.headers.get("X-Forwarded-For")

    if forwarded:
        ip_address = forwarded.split(",")[0].strip()
    elif request.client and request.client.host:
        ip_address = request.client.host

    if ip_address:
        city, region, country = get_ip_location(ip_address)

    # Check if a device is already registered for the same client ID.
    try:
        device_model = store.get_internal_authorized_device(
            client_id=client_id
        )
    except KeyError:
        device_model = store.create_authorized_device(
            OAuthDeviceInternalRequest(
                client_id=client_id,
                expires_in=config.device_auth_timeout,
                ip_address=ip_address,
                city=city,
                region=region,
                country=country,
                **device_details.model_dump(exclude_none=True),
            )
        )
    else:
        # Put the device into pending state and generate new codes. This
        # effectively invalidates the old codes and the device cannot be used
        # for authentication anymore.
        device_model = store.update_internal_authorized_device(
            device_id=device_model.id,
            update=OAuthDeviceInternalUpdate(
                trusted_device=False,
                expires_in=config.device_auth_timeout,
                status=OAuthDeviceStatus.PENDING,
                failed_auth_attempts=0,
                generate_new_codes=True,
                ip_address=ip_address,
                city=city,
                region=region,
                country=country,
                **device_details.model_dump(exclude_none=True),
            ),
        )

    dashboard_url = config.dashboard_url or config.server_url

    if dashboard_url:
        verification_uri = dashboard_url.lstrip("/") + DEVICES + DEVICE_VERIFY
    else:
        verification_uri = DEVICES + DEVICE_VERIFY

    verification_uri_complete = (
        verification_uri
        + "?"
        + urlencode(
            dict(
                device_id=str(device_model.id),
                user_code=str(device_model.user_code),
            )
        )
    )
    return OAuthDeviceAuthorizationResponse(
        device_code=device_model.device_code,
        user_code=device_model.user_code,
        expires_in=config.device_auth_timeout,
        interval=config.device_auth_polling_interval,
        verification_uri=verification_uri,
        verification_uri_complete=verification_uri_complete,
    )
logout(response: Response) -> None

Logs out the user.

Parameters:

Name Type Description Default
response Response

The response object.

required
Source code in src/zenml/zen_server/routers/auth_endpoints.py
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
@router.get(
    LOGOUT,
)
def logout(
    response: Response,
) -> None:
    """Logs out the user.

    Args:
        response: The response object.
    """
    config = server_config()

    # Remove the HTTP only cookie even if it does not exist
    response.delete_cookie(
        key=config.get_auth_cookie_name(),
        httponly=True,
        samesite="lax",
        domain=config.auth_cookie_domain,
    )
token(request: Request, response: Response, auth_form_data: OAuthLoginRequestForm = Depends()) -> Union[OAuthTokenResponse, OAuthRedirectResponse]

OAuth2 token endpoint.

Parameters:

Name Type Description Default
request Request

The request object.

required
response Response

The response object.

required
auth_form_data OAuthLoginRequestForm

The OAuth 2.0 authentication form data.

Depends()

Returns:

Type Description
Union[OAuthTokenResponse, OAuthRedirectResponse]

An access token or a redirect response.

Raises:

Type Description
ValueError

If the grant type is invalid.

Source code in src/zenml/zen_server/routers/auth_endpoints.py
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
@router.post(
    LOGIN,
    response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
)
@rate_limit_requests(
    day_limit=server_config().login_rate_limit_day,
    minute_limit=server_config().login_rate_limit_minute,
)
@handle_exceptions
def token(
    request: Request,
    response: Response,
    auth_form_data: OAuthLoginRequestForm = Depends(),
) -> Union[OAuthTokenResponse, OAuthRedirectResponse]:
    """OAuth2 token endpoint.

    Args:
        request: The request object.
        response: The response object.
        auth_form_data: The OAuth 2.0 authentication form data.

    Returns:
        An access token or a redirect response.

    Raises:
        ValueError: If the grant type is invalid.
    """
    config = server_config()
    cookie_response: Optional[Response] = response

    if auth_form_data.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
        auth_context = authenticate_credentials(
            user_name_or_id=auth_form_data.username,
            password=auth_form_data.password,
        )

    elif auth_form_data.grant_type == OAuthGrantTypes.OAUTH_DEVICE_CODE:
        auth_context = authenticate_device(
            client_id=auth_form_data.client_id,
            device_code=auth_form_data.device_code,
        )
        # API tokens for authorized device are only meant for non-web clients
        # and should not be stored as cookies
        cookie_response = None

    elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_API_KEY:
        auth_context = authenticate_api_key(
            api_key=auth_form_data.api_key,
        )
        # API tokens for API keys are only meant for non-web clients
        # and should not be stored as cookies
        cookie_response = None

    elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_EXTERNAL:
        assert config.external_login_url is not None

        authorization_url = config.external_login_url

        # Try to get the external session token or authorization token from the
        # authorization header
        authorization_header = request.headers.get("Authorization")
        if authorization_header:
            scheme, _, token = authorization_header.partition(" ")
            if token and scheme.lower() == "bearer":
                external_access_token = token

        if not authorization_header:
            logger.info(
                "External session or authorization token not found. Redirecting to "
                "external authenticator."
            )

            # Redirect the user to the external authentication login endpoint
            return OAuthRedirectResponse(authorization_url=authorization_url)

        auth_context = authenticate_external_user(
            external_access_token=external_access_token,
            request=request,
        )

        # TODO: no easy way to detect which type of client issued this request
        # (web or non-web) in order to decide whether to store the access token
        # as a cookie in the response or not. For now, we always assume a web
        # client.
    else:
        # Shouldn't happen, because we verify all grants in the form data
        raise ValueError("Invalid grant type.")

    return generate_access_token(
        user_id=auth_context.user.id,
        response=cookie_response,
        request=request,
        device=auth_context.device,
        api_key=auth_context.api_key,
    )
code_repositories_endpoints

Endpoint definitions for code repositories.

Classes Functions
create_code_repository(code_repository: CodeRepositoryRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> CodeRepositoryResponse

Creates a code repository.

Parameters:

Name Type Description Default
code_repository CodeRepositoryRequest

Code repository to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
CodeRepositoryResponse

The created code repository.

Source code in src/zenml/zen_server/routers/code_repositories_endpoints.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + CODE_REPOSITORIES,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["code_repositories"],
)
@handle_exceptions
def create_code_repository(
    code_repository: CodeRepositoryRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> CodeRepositoryResponse:
    """Creates a code repository.

    Args:
        code_repository: Code repository to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created code repository.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        code_repository.project = project.id

    return verify_permissions_and_create_entity(
        request_model=code_repository,
        create_method=zen_store().create_code_repository,
    )
delete_code_repository(code_repository_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific code repository.

Parameters:

Name Type Description Default
code_repository_id UUID

The ID of the code repository to delete.

required
Source code in src/zenml/zen_server/routers/code_repositories_endpoints.py
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@router.delete(
    "/{code_repository_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_code_repository(
    code_repository_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific code repository.

    Args:
        code_repository_id: The ID of the code repository to delete.
    """
    verify_permissions_and_delete_entity(
        id=code_repository_id,
        get_method=zen_store().get_code_repository,
        delete_method=zen_store().delete_code_repository,
    )
get_code_repository(code_repository_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> CodeRepositoryResponse

Gets a specific code repository using its unique ID.

Parameters:

Name Type Description Default
code_repository_id UUID

The ID of the code repository to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
CodeRepositoryResponse

A specific code repository object.

Source code in src/zenml/zen_server/routers/code_repositories_endpoints.py
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@router.get(
    "/{code_repository_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_code_repository(
    code_repository_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> CodeRepositoryResponse:
    """Gets a specific code repository using its unique ID.

    Args:
        code_repository_id: The ID of the code repository to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific code repository object.
    """
    return verify_permissions_and_get_entity(
        id=code_repository_id,
        get_method=zen_store().get_code_repository,
        hydrate=hydrate,
    )
list_code_repositories(filter_model: CodeRepositoryFilter = Depends(make_dependable(CodeRepositoryFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[CodeRepositoryResponse]

Gets a page of code repositories.

Parameters:

Name Type Description Default
filter_model CodeRepositoryFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(CodeRepositoryFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[CodeRepositoryResponse]

Page of code repository objects.

Source code in src/zenml/zen_server/routers/code_repositories_endpoints.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + CODE_REPOSITORIES,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["code_repositories"],
)
@handle_exceptions
def list_code_repositories(
    filter_model: CodeRepositoryFilter = Depends(
        make_dependable(CodeRepositoryFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[CodeRepositoryResponse]:
    """Gets a page of code repositories.

    Args:
        filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page of code repository objects.
    """
    if project_name_or_id:
        filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=filter_model,
        resource_type=ResourceType.CODE_REPOSITORY,
        list_method=zen_store().list_code_repositories,
        hydrate=hydrate,
    )
update_code_repository(code_repository_id: UUID, update: CodeRepositoryUpdate, _: AuthContext = Security(authorize)) -> CodeRepositoryResponse

Updates a code repository.

Parameters:

Name Type Description Default
code_repository_id UUID

The ID of the code repository to update.

required
update CodeRepositoryUpdate

The model containing the attributes to update.

required

Returns:

Type Description
CodeRepositoryResponse

The updated code repository object.

Source code in src/zenml/zen_server/routers/code_repositories_endpoints.py
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
@router.put(
    "/{code_repository_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_code_repository(
    code_repository_id: UUID,
    update: CodeRepositoryUpdate,
    _: AuthContext = Security(authorize),
) -> CodeRepositoryResponse:
    """Updates a code repository.

    Args:
        code_repository_id: The ID of the code repository to update.
        update: The model containing the attributes to update.

    Returns:
        The updated code repository object.
    """
    return verify_permissions_and_update_entity(
        id=code_repository_id,
        update_model=update,
        get_method=zen_store().get_code_repository,
        update_method=zen_store().update_code_repository,
    )
devices_endpoints

Endpoint definitions for code repositories.

Classes Functions
delete_authorized_device(device_id: UUID, auth_context: AuthContext = Security(authorize)) -> None

Deletes a specific OAuth2 authorized device using its unique ID.

Parameters:

Name Type Description Default
device_id UUID

The ID of the OAuth2 authorized device to delete.

required
auth_context AuthContext

The current auth context.

Security(authorize)

Raises:

Type Description
KeyError

If the device with the given ID does not exist or does not belong to the current user.

Source code in src/zenml/zen_server/routers/devices_endpoints.py
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
@router.delete(
    "/{device_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_authorized_device(
    device_id: UUID,
    auth_context: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific OAuth2 authorized device using its unique ID.

    Args:
        device_id: The ID of the OAuth2 authorized device to delete.
        auth_context: The current auth context.

    Raises:
        KeyError: If the device with the given ID does not exist or does not
            belong to the current user.
    """
    device = zen_store().get_authorized_device(device_id=device_id)
    if not device.user or device.user.id != auth_context.user.id:
        raise KeyError(
            f"Unable to get device with ID {device_id}: No device with "
            "this ID found."
        )

    zen_store().delete_authorized_device(device_id=device_id)
get_authorization_device(device_id: UUID, user_code: Optional[str] = None, hydrate: bool = True, auth_context: AuthContext = Security(authorize)) -> OAuthDeviceResponse

Gets a specific OAuth2 authorized device using its unique ID.

Parameters:

Name Type Description Default
device_id UUID

The ID of the OAuth2 authorized device to get.

required
user_code Optional[str]

The user code of the OAuth2 authorized device to get. Needs to be specified with devices that have not been verified yet.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True
auth_context AuthContext

The current auth context.

Security(authorize)

Returns:

Type Description
OAuthDeviceResponse

A specific OAuth2 authorized device object.

Raises:

Type Description
KeyError

If the device with the given ID does not exist, does not belong to the current user or could not be verified using the given user code.

Source code in src/zenml/zen_server/routers/devices_endpoints.py
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@router.get(
    "/{device_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_authorization_device(
    device_id: UUID,
    user_code: Optional[str] = None,
    hydrate: bool = True,
    auth_context: AuthContext = Security(authorize),
) -> OAuthDeviceResponse:
    """Gets a specific OAuth2 authorized device using its unique ID.

    Args:
        device_id: The ID of the OAuth2 authorized device to get.
        user_code: The user code of the OAuth2 authorized device to get. Needs
            to be specified with devices that have not been verified yet.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: The current auth context.

    Returns:
        A specific OAuth2 authorized device object.

    Raises:
        KeyError: If the device with the given ID does not exist, does not
            belong to the current user or could not be verified using the
            given user code.
    """
    device = zen_store().get_authorized_device(
        device_id=device_id, hydrate=hydrate
    )
    if not device.user:
        # A device that hasn't been verified and associated with a user yet
        # can only be retrieved if the user code is specified and valid.
        if user_code:
            internal_device = zen_store().get_internal_authorized_device(
                device_id=device_id, hydrate=hydrate
            )
            if internal_device.verify_user_code(user_code=user_code):
                return device
    elif device.user.id == auth_context.user.id:
        return device

    raise KeyError(
        f"Unable to get device with ID {device_id}: No device with "
        "this ID found."
    )
list_authorized_devices(filter_model: OAuthDeviceFilter = Depends(make_dependable(OAuthDeviceFilter)), hydrate: bool = False, auth_context: AuthContext = Security(authorize)) -> Page[OAuthDeviceResponse]

Gets a page of OAuth2 authorized devices belonging to the current user.

Parameters:

Name Type Description Default
filter_model OAuthDeviceFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(OAuthDeviceFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
auth_context AuthContext

The current auth context.

Security(authorize)

Returns:

Type Description
Page[OAuthDeviceResponse]

Page of OAuth2 authorized device objects.

Source code in src/zenml/zen_server/routers/devices_endpoints.py
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_authorized_devices(
    filter_model: OAuthDeviceFilter = Depends(
        make_dependable(OAuthDeviceFilter)
    ),
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[OAuthDeviceResponse]:
    """Gets a page of OAuth2 authorized devices belonging to the current user.

    Args:
        filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: The current auth context.

    Returns:
        Page of OAuth2 authorized device objects.
    """
    filter_model.set_scope_user(auth_context.user.id)
    return zen_store().list_authorized_devices(
        filter_model=filter_model, hydrate=hydrate
    )
update_authorized_device(device_id: UUID, update: OAuthDeviceUpdate, auth_context: AuthContext = Security(authorize)) -> OAuthDeviceResponse

Updates a specific OAuth2 authorized device using its unique ID.

Parameters:

Name Type Description Default
device_id UUID

The ID of the OAuth2 authorized device to update.

required
update OAuthDeviceUpdate

The model containing the attributes to update.

required
auth_context AuthContext

The current auth context.

Security(authorize)

Returns:

Type Description
OAuthDeviceResponse

The updated OAuth2 authorized device object.

Raises:

Type Description
KeyError

If the device with the given ID does not exist or does not belong to the current user.

Source code in src/zenml/zen_server/routers/devices_endpoints.py
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@router.put(
    "/{device_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_authorized_device(
    device_id: UUID,
    update: OAuthDeviceUpdate,
    auth_context: AuthContext = Security(authorize),
) -> OAuthDeviceResponse:
    """Updates a specific OAuth2 authorized device using its unique ID.

    Args:
        device_id: The ID of the OAuth2 authorized device to update.
        update: The model containing the attributes to update.
        auth_context: The current auth context.

    Returns:
        The updated OAuth2 authorized device object.

    Raises:
        KeyError: If the device with the given ID does not exist or does not
            belong to the current user.
    """
    device = zen_store().get_authorized_device(device_id=device_id)
    if not device.user or device.user.id != auth_context.user.id:
        raise KeyError(
            f"Unable to get device with ID {device_id}: No device with "
            "this ID found."
        )

    return zen_store().update_authorized_device(
        device_id=device_id, update=update
    )
verify_authorized_device(device_id: UUID, request: OAuthDeviceVerificationRequest, auth_context: AuthContext = Security(authorize)) -> OAuthDeviceResponse

Verifies a specific OAuth2 authorized device using its unique ID.

This endpoint implements the OAuth2 device authorization grant flow as defined in https://tools.ietf.org/html/rfc8628. It is called to verify the user code for a given device ID.

If the user code is valid, the device is marked as verified and associated with the user that authorized the device. This association is required to be able to issue access tokens or revoke the device later on.

Parameters:

Name Type Description Default
device_id UUID

The ID of the OAuth2 authorized device to update.

required
request OAuthDeviceVerificationRequest

The model containing the verification request.

required
auth_context AuthContext

The current auth context.

Security(authorize)

Returns:

Type Description
OAuthDeviceResponse

The updated OAuth2 authorized device object.

Raises:

Type Description
ValueError

If the device verification request fails.

Source code in src/zenml/zen_server/routers/devices_endpoints.py
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
@router.put(
    "/{device_id}" + DEVICE_VERIFY,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def verify_authorized_device(
    device_id: UUID,
    request: OAuthDeviceVerificationRequest,
    auth_context: AuthContext = Security(authorize),
) -> OAuthDeviceResponse:
    """Verifies a specific OAuth2 authorized device using its unique ID.

    This endpoint implements the OAuth2 device authorization grant flow as
    defined in https://tools.ietf.org/html/rfc8628. It is called to verify
    the user code for a given device ID.

    If the user code is valid, the device is marked as verified and associated
    with the user that authorized the device. This association is required to
    be able to issue access tokens or revoke the device later on.

    Args:
        device_id: The ID of the OAuth2 authorized device to update.
        request: The model containing the verification request.
        auth_context: The current auth context.

    Returns:
        The updated OAuth2 authorized device object.

    Raises:
        ValueError: If the device verification request fails.
    """
    config = server_config()
    store = zen_store()

    # Check if a device is registered for the ID
    device_model = store.get_internal_authorized_device(
        device_id=device_id,
    )

    with track_handler(event=AnalyticsEvent.DEVICE_VERIFIED):
        # Check if the device is in a state that allows verification.
        if device_model.status != OAuthDeviceStatus.PENDING:
            raise ValueError(
                "Invalid request: device not pending verification.",
            )

        # Check if the device verification has expired.
        if device_model.expires and device_model.expires < utc_now():
            raise ValueError(
                "Invalid request: device verification expired.",
            )

        # Check if the device already has a user associated with it. If so, the
        # current user and the user associated with the device must be the same.
        if device_model.user and device_model.user.id != auth_context.user.id:
            raise ValueError(
                "Invalid request: this device is associated with another user.",
            )

        # Check if the device code is valid.
        if not device_model.verify_user_code(request.user_code):
            # If the device code is invalid, increment the failed auth attempts
            # counter and lock the device if the maximum number of failed auth
            # attempts has been reached.
            failed_auth_attempts = device_model.failed_auth_attempts + 1
            update = OAuthDeviceInternalUpdate(
                failed_auth_attempts=failed_auth_attempts
            )
            if failed_auth_attempts >= config.max_failed_device_auth_attempts:
                update.locked = True
            store.update_internal_authorized_device(
                device_id=device_model.id,
                update=update,
            )
            if failed_auth_attempts >= config.max_failed_device_auth_attempts:
                raise ValueError(
                    "Invalid request: device locked due to too many failed "
                    "authentication attempts.",
                )
            raise ValueError(
                "Invalid request: invalid user code.",
            )

        # If the device code is valid, associate the device with the current user.
        # We don't reset the expiration date yet, because we want to make sure
        # that the client has received the access token before we do so, to avoid
        # brute force attacks on the device code.
        update = OAuthDeviceInternalUpdate(
            status=OAuthDeviceStatus.VERIFIED,
            user_id=auth_context.user.id,
            failed_auth_attempts=0,
            trusted_device=request.trusted_device,
        )
        device_model = store.update_internal_authorized_device(
            device_id=device_model.id,
            update=update,
        )

    store.update_onboarding_state(
        completed_steps={OnboardingStep.DEVICE_VERIFIED}
    )
    return device_model
event_source_endpoints

Endpoint definitions for event sources.

Classes Functions
create_event_source(event_source: EventSourceRequest, _: AuthContext = Security(authorize)) -> EventSourceResponse

Creates an event source.

Parameters:

Name Type Description Default
event_source EventSourceRequest

EventSource to register.

required

Returns:

Type Description
EventSourceResponse

The created event source.

Raises:

Type Description
ValueError

If the plugin for an event source is not a valid event source plugin.

Source code in src/zenml/zen_server/routers/event_source_endpoints.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
@event_source_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_event_source(
    event_source: EventSourceRequest,
    _: AuthContext = Security(authorize),
) -> EventSourceResponse:
    """Creates an event source.

    Args:
        event_source: EventSource to register.

    Returns:
        The created event source.

    Raises:
        ValueError: If the plugin for an event source is not a valid event
            source plugin.
    """
    event_source_handler = plugin_flavor_registry().get_plugin(
        name=event_source.flavor,
        _type=PluginType.EVENT_SOURCE,
        subtype=event_source.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an event source
    # implementation
    if not isinstance(event_source_handler, BaseEventSourceHandler):
        raise ValueError(
            f"Event source plugin {event_source.plugin_subtype} "
            f"for flavor {event_source.flavor} is not a valid event source "
            "handler implementation."
        )

    return verify_permissions_and_create_entity(
        request_model=event_source,
        create_method=event_source_handler.create_event_source,
    )
delete_event_source(event_source_id: UUID, force: bool = False, _: AuthContext = Security(authorize)) -> None

Deletes a event_source.

Parameters:

Name Type Description Default
event_source_id UUID

Name of the event_source.

required
force bool

Flag deciding whether to force delete the event source.

False

Raises:

Type Description
ValueError

If the plugin for an event source is not a valid event source plugin.

Source code in src/zenml/zen_server/routers/event_source_endpoints.py
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
@event_source_router.delete(
    "/{event_source_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_event_source(
    event_source_id: UUID,
    force: bool = False,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a event_source.

    Args:
        event_source_id: Name of the event_source.
        force: Flag deciding whether to force delete the event source.

    Raises:
        ValueError: If the plugin for an event source is not a valid event
            source plugin.
    """
    event_source = zen_store().get_event_source(
        event_source_id=event_source_id
    )

    verify_permission_for_model(event_source, action=Action.DELETE)

    event_source_handler = plugin_flavor_registry().get_plugin(
        name=event_source.flavor,
        _type=PluginType.EVENT_SOURCE,
        subtype=event_source.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an event source
    # implementation
    if not isinstance(event_source_handler, BaseEventSourceHandler):
        raise ValueError(
            f"Event source plugin {event_source.plugin_subtype} "
            f"for flavor {event_source.flavor} is not a valid event source "
            "handler implementation."
        )

    event_source_handler.delete_event_source(
        event_source=event_source,
        force=force,
    )
get_event_source(event_source_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> EventSourceResponse

Returns the requested event_source.

Parameters:

Name Type Description Default
event_source_id UUID

ID of the event_source.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
EventSourceResponse

The requested event_source.

Raises:

Type Description
ValueError

If the plugin for an event source is not a valid event source plugin.

Source code in src/zenml/zen_server/routers/event_source_endpoints.py
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
@event_source_router.get(
    "/{event_source_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_event_source(
    event_source_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> EventSourceResponse:
    """Returns the requested event_source.

    Args:
        event_source_id: ID of the event_source.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested event_source.

    Raises:
        ValueError: If the plugin for an event source is not a valid event
            source plugin.
    """
    event_source = zen_store().get_event_source(
        event_source_id=event_source_id, hydrate=hydrate
    )

    verify_permission_for_model(event_source, action=Action.READ)

    event_source_handler = plugin_flavor_registry().get_plugin(
        name=event_source.flavor,
        _type=PluginType.EVENT_SOURCE,
        subtype=event_source.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an event source
    # implementation
    if not isinstance(event_source_handler, BaseEventSourceHandler):
        raise ValueError(
            f"Event source plugin {event_source.plugin_subtype} "
            f"for flavor {event_source.flavor} is not a valid event source "
            "handler implementation."
        )

    event_source = event_source_handler.get_event_source(
        event_source, hydrate=hydrate
    )

    return dehydrate_response_model(event_source)
list_event_sources(event_source_filter_model: EventSourceFilter = Depends(make_dependable(EventSourceFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[EventSourceResponse]

Returns all event_sources.

Parameters:

Name Type Description Default
event_source_filter_model EventSourceFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(EventSourceFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[EventSourceResponse]

All event_sources.

Source code in src/zenml/zen_server/routers/event_source_endpoints.py
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
@event_source_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_event_sources(
    event_source_filter_model: EventSourceFilter = Depends(
        make_dependable(EventSourceFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[EventSourceResponse]:
    """Returns all event_sources.

    Args:
        event_source_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        All event_sources.
    """

    def list_event_sources_fn(
        filter_model: EventSourceFilter,
    ) -> Page[EventSourceResponse]:
        """List event sources through their associated plugins.

        Args:
            filter_model: Filter model used for pagination, sorting,
                filtering.

        Returns:
            All event sources.

        Raises:
            ValueError: If the plugin for an event source is not a valid event
                source plugin.
        """
        event_sources = zen_store().list_event_sources(
            event_source_filter_model=filter_model, hydrate=hydrate
        )

        # Process the event sources through their associated plugins
        for idx, event_source in enumerate(event_sources.items):
            event_source_handler = plugin_flavor_registry().get_plugin(
                name=event_source.flavor,
                _type=PluginType.EVENT_SOURCE,
                subtype=event_source.plugin_subtype,
            )

            # Validate that the flavor and plugin_type correspond to an event
            # source implementation
            if not isinstance(event_source_handler, BaseEventSourceHandler):
                raise ValueError(
                    f"Event source plugin {event_source.plugin_subtype} "
                    f"for flavor {event_source.flavor} is not a valid event "
                    "source handler implementation."
                )

            event_sources.items[idx] = event_source_handler.get_event_source(
                event_source, hydrate=hydrate
            )

        return event_sources

    return verify_permissions_and_list_entities(
        filter_model=event_source_filter_model,
        resource_type=ResourceType.EVENT_SOURCE,
        list_method=list_event_sources_fn,
    )
update_event_source(event_source_id: UUID, event_source_update: EventSourceUpdate, _: AuthContext = Security(authorize)) -> EventSourceResponse

Updates an event_source.

Parameters:

Name Type Description Default
event_source_id UUID

Name of the event_source.

required
event_source_update EventSourceUpdate

EventSource to use for the update.

required

Returns:

Type Description
EventSourceResponse

The updated event_source.

Raises:

Type Description
ValueError

If the plugin for an event source is not a valid event source plugin.

Source code in src/zenml/zen_server/routers/event_source_endpoints.py
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
@event_source_router.put(
    "/{event_source_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_event_source(
    event_source_id: UUID,
    event_source_update: EventSourceUpdate,
    _: AuthContext = Security(authorize),
) -> EventSourceResponse:
    """Updates an event_source.

    Args:
        event_source_id: Name of the event_source.
        event_source_update: EventSource to use for the update.

    Returns:
        The updated event_source.

    Raises:
        ValueError: If the plugin for an event source is not a valid event
            source plugin.
    """
    event_source = zen_store().get_event_source(
        event_source_id=event_source_id
    )

    verify_permission_for_model(event_source, action=Action.UPDATE)

    event_source_handler = plugin_flavor_registry().get_plugin(
        name=event_source.flavor,
        _type=PluginType.EVENT_SOURCE,
        subtype=event_source.plugin_subtype,
    )

    # Validate that the flavor and plugin_type correspond to an event source
    # implementation
    if not isinstance(event_source_handler, BaseEventSourceHandler):
        raise ValueError(
            f"Event source plugin {event_source.plugin_subtype} "
            f"for flavor {event_source.flavor} is not a valid event source "
            "handler implementation."
        )

    updated_event_source = event_source_handler.update_event_source(
        event_source=event_source,
        event_source_update=event_source_update,
    )

    return dehydrate_response_model(updated_event_source)
flavors_endpoints

Endpoint definitions for flavors.

Classes Functions
create_flavor(flavor: FlavorRequest, _: AuthContext = Security(authorize)) -> FlavorResponse

Creates a stack component flavor.

Parameters:

Name Type Description Default
flavor FlavorRequest

Stack component flavor to register.

required

Returns:

Type Description
FlavorResponse

The created stack component flavor.

Source code in src/zenml/zen_server/routers/flavors_endpoints.py
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_flavor(
    flavor: FlavorRequest,
    _: AuthContext = Security(authorize),
) -> FlavorResponse:
    """Creates a stack component flavor.

    Args:
        flavor: Stack component flavor to register.

    Returns:
        The created stack component flavor.
    """
    return verify_permissions_and_create_entity(
        request_model=flavor,
        create_method=zen_store().create_flavor,
    )
delete_flavor(flavor_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a flavor.

Parameters:

Name Type Description Default
flavor_id UUID

ID of the flavor.

required
Source code in src/zenml/zen_server/routers/flavors_endpoints.py
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@router.delete(
    "/{flavor_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_flavor(
    flavor_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a flavor.

    Args:
        flavor_id: ID of the flavor.
    """
    verify_permissions_and_delete_entity(
        id=flavor_id,
        get_method=zen_store().get_flavor,
        delete_method=zen_store().delete_flavor,
    )
get_flavor(flavor_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> FlavorResponse

Returns the requested flavor.

Parameters:

Name Type Description Default
flavor_id UUID

ID of the flavor.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
FlavorResponse

The requested stack.

Source code in src/zenml/zen_server/routers/flavors_endpoints.py
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
@router.get(
    "/{flavor_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_flavor(
    flavor_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> FlavorResponse:
    """Returns the requested flavor.

    Args:
        flavor_id: ID of the flavor.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested stack.
    """
    return verify_permissions_and_get_entity(
        id=flavor_id, get_method=zen_store().get_flavor, hydrate=hydrate
    )
list_flavors(flavor_filter_model: FlavorFilter = Depends(make_dependable(FlavorFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[FlavorResponse]

Returns all flavors.

Parameters:

Name Type Description Default
flavor_filter_model FlavorFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(FlavorFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[FlavorResponse]

All flavors.

Source code in src/zenml/zen_server/routers/flavors_endpoints.py
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_flavors(
    flavor_filter_model: FlavorFilter = Depends(make_dependable(FlavorFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[FlavorResponse]:
    """Returns all flavors.

    Args:
        flavor_filter_model: Filter model used for pagination, sorting,
                             filtering
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        All flavors.
    """
    return verify_permissions_and_list_entities(
        filter_model=flavor_filter_model,
        resource_type=ResourceType.FLAVOR,
        list_method=zen_store().list_flavors,
        hydrate=hydrate,
    )
sync_flavors(_: AuthContext = Security(authorize)) -> None

Purge all in-built and integration flavors from the DB and sync.

Returns:

Type Description
None

None if successful. Raises an exception otherwise.

Source code in src/zenml/zen_server/routers/flavors_endpoints.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
@router.patch(
    "/sync",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def sync_flavors(
    _: AuthContext = Security(authorize),
) -> None:
    """Purge all in-built and integration flavors from the DB and sync.

    Returns:
        None if successful. Raises an exception otherwise.
    """
    verify_permission(resource_type=ResourceType.FLAVOR, action=Action.UPDATE)
    return zen_store()._sync_flavors()
update_flavor(flavor_id: UUID, flavor_update: FlavorUpdate, _: AuthContext = Security(authorize)) -> FlavorResponse

Updates a flavor.

noqa: DAR401

Parameters:

Name Type Description Default
flavor_id UUID

ID of the flavor to update.

required
flavor_update FlavorUpdate

Flavor update.

required

Returns:

Type Description
FlavorResponse

The updated flavor.

Source code in src/zenml/zen_server/routers/flavors_endpoints.py
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@router.put(
    "/{flavor_id}",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def update_flavor(
    flavor_id: UUID,
    flavor_update: FlavorUpdate,
    _: AuthContext = Security(authorize),
) -> FlavorResponse:
    """Updates a flavor.

    # noqa: DAR401

    Args:
        flavor_id: ID of the flavor to update.
        flavor_update: Flavor update.

    Returns:
        The updated flavor.
    """
    return verify_permissions_and_update_entity(
        id=flavor_id,
        update_model=flavor_update,
        get_method=zen_store().get_flavor,
        update_method=zen_store().update_flavor,
    )
logs_endpoints

Endpoint definitions for logs.

Classes Functions
get_logs(logs_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> LogsResponse

Returns the requested logs.

Parameters:

Name Type Description Default
logs_id UUID

ID of the logs.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
LogsResponse

The requested logs.

Source code in src/zenml/zen_server/routers/logs_endpoints.py
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@router.get(
    "/{logs_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_logs(
    logs_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> LogsResponse:
    """Returns the requested logs.

    Args:
        logs_id: ID of the logs.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested logs.
    """
    return verify_permissions_and_get_entity(
        id=logs_id, get_method=zen_store().get_logs, hydrate=hydrate
    )
model_versions_endpoints

Endpoint definitions for models.

Classes Functions
create_model_version(model_version: ModelVersionRequest, model_id: Optional[UUID] = None, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> ModelVersionResponse

Creates a model version.

Parameters:

Name Type Description Default
model_version ModelVersionRequest

Model version to create.

required
model_id Optional[UUID]

Optional ID of the model.

None
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
ModelVersionResponse

The created model version.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + MODELS + "/{model_id}" + MODEL_VERSIONS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["model_versions"],
)
@handle_exceptions
def create_model_version(
    model_version: ModelVersionRequest,
    model_id: Optional[UUID] = None,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Creates a model version.

    Args:
        model_version: Model version to create.
        model_id: Optional ID of the model.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created model version.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        model_version.project = project.id

    if model_id:
        model_version.model = model_id

    return verify_permissions_and_create_entity(
        request_model=model_version,
        create_method=zen_store().create_model_version,
    )
create_model_version_artifact_link(model_version_artifact_link: ModelVersionArtifactRequest, _: AuthContext = Security(authorize)) -> ModelVersionArtifactResponse

Create a new model version to artifact link.

Parameters:

Name Type Description Default
model_version_artifact_link ModelVersionArtifactRequest

The model version to artifact link to create.

required

Returns:

Type Description
ModelVersionArtifactResponse

The created model version to artifact link.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
@model_version_artifacts_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_model_version_artifact_link(
    model_version_artifact_link: ModelVersionArtifactRequest,
    _: AuthContext = Security(authorize),
) -> ModelVersionArtifactResponse:
    """Create a new model version to artifact link.

    Args:
        model_version_artifact_link: The model version to artifact link to create.

    Returns:
        The created model version to artifact link.
    """
    model_version = zen_store().get_model_version(
        model_version_artifact_link.model_version
    )

    return verify_permissions_and_create_entity(
        request_model=model_version_artifact_link,
        create_method=zen_store().create_model_version_artifact_link,
        # Check for UPDATE permissions on the model version instead of the
        # model version artifact link
        surrogate_models=[model_version],
    )
create_model_version_pipeline_run_link(model_version_pipeline_run_link: ModelVersionPipelineRunRequest, _: AuthContext = Security(authorize)) -> ModelVersionPipelineRunResponse

Create a new model version to pipeline run link.

Parameters:

Name Type Description Default
model_version_pipeline_run_link ModelVersionPipelineRunRequest

The model version to pipeline run link to create.

required

Returns:

Type Description
ModelVersionPipelineRunResponse
  • If Model Version to Pipeline Run Link already exists - returns the existing link.
ModelVersionPipelineRunResponse
  • Otherwise, returns the newly created model version to pipeline run link.
Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
@model_version_pipeline_runs_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_model_version_pipeline_run_link(
    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
    _: AuthContext = Security(authorize),
) -> ModelVersionPipelineRunResponse:
    """Create a new model version to pipeline run link.

    Args:
        model_version_pipeline_run_link: The model version to pipeline run link to create.

    Returns:
        - If Model Version to Pipeline Run Link already exists - returns the existing link.
        - Otherwise, returns the newly created model version to pipeline run link.
    """
    model_version = zen_store().get_model_version(
        model_version_pipeline_run_link.model_version, hydrate=False
    )

    return verify_permissions_and_create_entity(
        request_model=model_version_pipeline_run_link,
        create_method=zen_store().create_model_version_pipeline_run_link,
        # Check for UPDATE permissions on the model version instead of the
        # model version pipeline run link
        surrogate_models=[model_version],
    )
delete_all_model_version_artifact_links(model_version_id: UUID, only_links: bool = True, _: AuthContext = Security(authorize)) -> None

Deletes all model version to artifact links.

Parameters:

Name Type Description Default
model_version_id UUID

ID of the model version containing links.

required
only_links bool

Whether to only delete the link to the artifact.

True
Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
@router.delete(
    "/{model_version_id}" + ARTIFACTS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_all_model_version_artifact_links(
    model_version_id: UUID,
    only_links: bool = True,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes all model version to artifact links.

    Args:
        model_version_id: ID of the model version containing links.
        only_links: Whether to only delete the link to the artifact.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_all_model_version_artifact_links(
        model_version_id, only_links
    )
delete_model_version(model_version_id: UUID, _: AuthContext = Security(authorize)) -> None

Delete a model by name or ID.

Parameters:

Name Type Description Default
model_version_id UUID

The name or ID of the model version to delete.

required
Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
@router.delete(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_model_version(
    model_version_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a model by name or ID.

    Args:
        model_version_id: The name or ID of the model version to delete.
    """
    verify_permissions_and_delete_entity(
        id=model_version_id,
        get_method=zen_store().get_model_version,
        delete_method=zen_store().delete_model_version,
    )
delete_model_version_artifact_link(model_version_id: UUID, model_version_artifact_link_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Deletes a model version to artifact link.

Parameters:

Name Type Description Default
model_version_id UUID

ID of the model version containing the link.

required
model_version_artifact_link_name_or_id Union[str, UUID]

name or ID of the model version to artifact link to be deleted.

required
Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
@router.delete(
    "/{model_version_id}"
    + ARTIFACTS
    + "/{model_version_artifact_link_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_model_version_artifact_link(
    model_version_id: UUID,
    model_version_artifact_link_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a model version to artifact link.

    Args:
        model_version_id: ID of the model version containing the link.
        model_version_artifact_link_name_or_id: name or ID of the model
            version to artifact link to be deleted.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_model_version_artifact_link(
        model_version_id,
        model_version_artifact_link_name_or_id,
    )
delete_model_version_pipeline_run_link(model_version_id: UUID, model_version_pipeline_run_link_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Deletes a model version link.

Parameters:

Name Type Description Default
model_version_id UUID

name or ID of the model version containing the link.

required
model_version_pipeline_run_link_name_or_id Union[str, UUID]

name or ID of the model version link to be deleted.

required
Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
@router.delete(
    "/{model_version_id}"
    + RUNS
    + "/{model_version_pipeline_run_link_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_model_version_pipeline_run_link(
    model_version_id: UUID,
    model_version_pipeline_run_link_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a model version link.

    Args:
        model_version_id: name or ID of the model version containing the link.
        model_version_pipeline_run_link_name_or_id: name or ID of the model
            version link to be deleted.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_model_version_pipeline_run_link(
        model_version_id=model_version_id,
        model_version_pipeline_run_link_name_or_id=model_version_pipeline_run_link_name_or_id,
    )
get_model_version(model_version_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ModelVersionResponse

Get a model version by ID.

Parameters:

Name Type Description Default
model_version_id UUID

id of the model version to be retrieved.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ModelVersionResponse

The model version with the given name or ID.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@router.get(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_model_version(
    model_version_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Get a model version by ID.

    Args:
        model_version_id: id of the model version to be retrieved.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version with the given name or ID.
    """
    model_version = zen_store().get_model_version(
        model_version_id=model_version_id,
        hydrate=hydrate,
    )
    verify_permission_for_model(model_version.model, action=Action.READ)
    return dehydrate_response_model(model_version)
list_model_version_artifact_links(model_version_artifact_link_filter_model: ModelVersionArtifactFilter = Depends(make_dependable(ModelVersionArtifactFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ModelVersionArtifactResponse]

Get model version to artifact links according to query filters.

Parameters:

Name Type Description Default
model_version_artifact_link_filter_model ModelVersionArtifactFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ModelVersionArtifactFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ModelVersionArtifactResponse]

The model version to artifact links according to query filters.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
@model_version_artifacts_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_model_version_artifact_links(
    model_version_artifact_link_filter_model: ModelVersionArtifactFilter = Depends(
        make_dependable(ModelVersionArtifactFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ModelVersionArtifactResponse]:
    """Get model version to artifact links according to query filters.

    Args:
        model_version_artifact_link_filter_model: Filter model used for
            pagination, sorting, filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version to artifact links according to query filters.
    """
    return zen_store().list_model_version_artifact_links(
        model_version_artifact_link_filter_model=model_version_artifact_link_filter_model,
        hydrate=hydrate,
    )
list_model_version_pipeline_run_links(model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilter = Depends(make_dependable(ModelVersionPipelineRunFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ModelVersionPipelineRunResponse]

Get model version to pipeline run links according to query filters.

Parameters:

Name Type Description Default
model_version_pipeline_run_link_filter_model ModelVersionPipelineRunFilter

Filter model used for pagination, sorting, and filtering.

Depends(make_dependable(ModelVersionPipelineRunFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ModelVersionPipelineRunResponse]

The model version to pipeline run links according to query filters.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
@model_version_pipeline_runs_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_model_version_pipeline_run_links(
    model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilter = Depends(
        make_dependable(ModelVersionPipelineRunFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ModelVersionPipelineRunResponse]:
    """Get model version to pipeline run links according to query filters.

    Args:
        model_version_pipeline_run_link_filter_model: Filter model used for
            pagination, sorting, and filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version to pipeline run links according to query filters.
    """
    return zen_store().list_model_version_pipeline_run_links(
        model_version_pipeline_run_link_filter_model=model_version_pipeline_run_link_filter_model,
        hydrate=hydrate,
    )
list_model_versions(model_version_filter_model: ModelVersionFilter = Depends(make_dependable(ModelVersionFilter)), model_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, auth_context: AuthContext = Security(authorize)) -> Page[ModelVersionResponse]

Get model versions according to query filters.

Parameters:

Name Type Description Default
model_version_filter_model ModelVersionFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ModelVersionFilter))
model_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the model.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
Page[ModelVersionResponse]

The model versions according to query filters.

Raises:

Type Description
ValueError

If the model is missing from the filter.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
@model_router.get(
    "/{model_name_or_id}" + MODEL_VERSIONS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_model_versions(
    model_version_filter_model: ModelVersionFilter = Depends(
        make_dependable(ModelVersionFilter)
    ),
    model_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[ModelVersionResponse]:
    """Get model versions according to query filters.

    Args:
        model_version_filter_model: Filter model used for pagination, sorting,
            filtering.
        model_name_or_id: Optional name or ID of the model.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: The authentication context.

    Returns:
        The model versions according to query filters.

    Raises:
        ValueError: If the model is missing from the filter.
    """
    if model_name_or_id:
        model_version_filter_model.model = model_name_or_id

    if not model_version_filter_model.model:
        raise ValueError("Model missing from the filter")

    # A project scoped request must always be scoped to a specific
    # project. This is required for the RBAC check to work.
    set_filter_project_scope(model_version_filter_model)
    assert isinstance(model_version_filter_model.project, UUID)

    model = zen_store().get_model_by_name_or_id(
        model_version_filter_model.model,
        project=model_version_filter_model.project,
    )

    # Check read permissions on the model
    verify_permission_for_model(model, action=Action.READ)

    model_version_filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id
    )

    model_versions = zen_store().list_model_versions(
        model_version_filter_model=model_version_filter_model,
        hydrate=hydrate,
    )
    return dehydrate_page(model_versions)
update_model_version(model_version_id: UUID, model_version_update_model: ModelVersionUpdate, _: AuthContext = Security(authorize)) -> ModelVersionResponse

Get all model versions by filter.

Parameters:

Name Type Description Default
model_version_id UUID

The ID of model version to be updated.

required
model_version_update_model ModelVersionUpdate

The model version to be updated.

required

Returns:

Type Description
ModelVersionResponse

An updated model version.

Source code in src/zenml/zen_server/routers/model_versions_endpoints.py
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
@router.put(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_model_version(
    model_version_id: UUID,
    model_version_update_model: ModelVersionUpdate,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Get all model versions by filter.

    Args:
        model_version_id: The ID of model version to be updated.
        model_version_update_model: The model version to be updated.

    Returns:
        An updated model version.
    """
    model_version = zen_store().get_model_version(model_version_id)

    if model_version_update_model.stage:
        # Make sure the user has permissions to promote the model
        verify_permission_for_model(model_version.model, action=Action.PROMOTE)

    verify_permission_for_model(model_version, action=Action.UPDATE)
    updated_model_version = zen_store().update_model_version(
        model_version_id=model_version_id,
        model_version_update_model=model_version_update_model,
    )

    return dehydrate_response_model(updated_model_version)
models_endpoints

Endpoint definitions for models.

Classes Functions
create_model(model: ModelRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> ModelResponse

Creates a model.

Parameters:

Name Type Description Default
model ModelRequest

Model to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
ModelResponse

The created model.

Source code in src/zenml/zen_server/routers/models_endpoints.py
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + MODELS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["models"],
)
@handle_exceptions
def create_model(
    model: ModelRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ModelResponse:
    """Creates a model.

    Args:
        model: Model to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created model.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        model.project = project.id

    return verify_permissions_and_create_entity(
        request_model=model,
        create_method=zen_store().create_model,
    )
delete_model(model_id: UUID, _: AuthContext = Security(authorize)) -> None

Delete a model by ID.

Parameters:

Name Type Description Default
model_id UUID

The ID of the model to delete.

required
Source code in src/zenml/zen_server/routers/models_endpoints.py
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
@router.delete(
    "/{model_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_model(
    model_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a model by ID.

    Args:
        model_id: The ID of the model to delete.
    """
    model = verify_permissions_and_delete_entity(
        id=model_id,
        get_method=zen_store().get_model,
        delete_method=zen_store().delete_model,
    )

    if server_config().feature_gate_enabled:
        if ResourceType.MODEL in server_config().reportable_resources:
            report_decrement(ResourceType.MODEL, resource_id=model.id)
get_model(model_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ModelResponse

Get a model by name or ID.

Parameters:

Name Type Description Default
model_id UUID

The ID of the model to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ModelResponse

The model with the given name or ID.

Source code in src/zenml/zen_server/routers/models_endpoints.py
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
@router.get(
    "/{model_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_model(
    model_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ModelResponse:
    """Get a model by name or ID.

    Args:
        model_id: The ID of the model to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model with the given name or ID.
    """
    return verify_permissions_and_get_entity(
        id=model_id, get_method=zen_store().get_model, hydrate=hydrate
    )
list_models(model_filter_model: ModelFilter = Depends(make_dependable(ModelFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ModelResponse]

Get models according to query filters.

Parameters:

Name Type Description Default
model_filter_model ModelFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ModelFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ModelResponse]

The models according to query filters.

Source code in src/zenml/zen_server/routers/models_endpoints.py
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_models(
    model_filter_model: ModelFilter = Depends(make_dependable(ModelFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ModelResponse]:
    """Get models according to query filters.

    Args:
        model_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The models according to query filters.
    """
    return verify_permissions_and_list_entities(
        filter_model=model_filter_model,
        resource_type=ResourceType.MODEL,
        list_method=zen_store().list_models,
        hydrate=hydrate,
    )
update_model(model_id: UUID, model_update: ModelUpdate, _: AuthContext = Security(authorize)) -> ModelResponse

Updates a model.

Parameters:

Name Type Description Default
model_id UUID

Name of the stack.

required
model_update ModelUpdate

Stack to use for the update.

required

Returns:

Type Description
ModelResponse

The updated model.

Source code in src/zenml/zen_server/routers/models_endpoints.py
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@router.put(
    "/{model_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_model(
    model_id: UUID,
    model_update: ModelUpdate,
    _: AuthContext = Security(authorize),
) -> ModelResponse:
    """Updates a model.

    Args:
        model_id: Name of the stack.
        model_update: Stack to use for the update.

    Returns:
        The updated model.
    """
    return verify_permissions_and_update_entity(
        id=model_id,
        update_model=model_update,
        get_method=zen_store().get_model,
        update_method=zen_store().update_model,
    )
pipeline_builds_endpoints

Endpoint definitions for builds.

Classes Functions
create_build(build: PipelineBuildRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> PipelineBuildResponse

Creates a build, optionally in a specific project.

Parameters:

Name Type Description Default
build PipelineBuildRequest

Build to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
PipelineBuildResponse

The created build.

Source code in src/zenml/zen_server/routers/pipeline_builds_endpoints.py
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + PIPELINE_BUILDS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["builds"],
)
@handle_exceptions
def create_build(
    build: PipelineBuildRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> PipelineBuildResponse:
    """Creates a build, optionally in a specific project.

    Args:
        build: Build to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created build.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        build.project = project.id

    return verify_permissions_and_create_entity(
        request_model=build,
        create_method=zen_store().create_build,
    )
delete_build(build_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific build.

Parameters:

Name Type Description Default
build_id UUID

ID of the build to delete.

required
Source code in src/zenml/zen_server/routers/pipeline_builds_endpoints.py
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
@router.delete(
    "/{build_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_build(
    build_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific build.

    Args:
        build_id: ID of the build to delete.
    """
    verify_permissions_and_delete_entity(
        id=build_id,
        get_method=zen_store().get_build,
        delete_method=zen_store().delete_build,
    )
get_build(build_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> PipelineBuildResponse

Gets a specific build using its unique id.

Parameters:

Name Type Description Default
build_id UUID

ID of the build to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
PipelineBuildResponse

A specific build object.

Source code in src/zenml/zen_server/routers/pipeline_builds_endpoints.py
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
@router.get(
    "/{build_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_build(
    build_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> PipelineBuildResponse:
    """Gets a specific build using its unique id.

    Args:
        build_id: ID of the build to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific build object.
    """
    return verify_permissions_and_get_entity(
        id=build_id, get_method=zen_store().get_build, hydrate=hydrate
    )
list_builds(build_filter_model: PipelineBuildFilter = Depends(make_dependable(PipelineBuildFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[PipelineBuildResponse]

Gets a list of builds.

Parameters:

Name Type Description Default
build_filter_model PipelineBuildFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(PipelineBuildFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[PipelineBuildResponse]

List of build objects matching the filter criteria.

Source code in src/zenml/zen_server/routers/pipeline_builds_endpoints.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + PIPELINE_BUILDS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["builds"],
)
@handle_exceptions
def list_builds(
    build_filter_model: PipelineBuildFilter = Depends(
        make_dependable(PipelineBuildFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[PipelineBuildResponse]:
    """Gets a list of builds.

    Args:
        build_filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project to filter by.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of build objects matching the filter criteria.
    """
    if project_name_or_id:
        build_filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=build_filter_model,
        resource_type=ResourceType.PIPELINE_BUILD,
        list_method=zen_store().list_builds,
        hydrate=hydrate,
    )
pipeline_deployments_endpoints

Endpoint definitions for deployments.

Classes Functions
create_deployment(deployment: PipelineDeploymentRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> PipelineDeploymentResponse

Creates a deployment.

Parameters:

Name Type Description Default
deployment PipelineDeploymentRequest

Deployment to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
PipelineDeploymentResponse

The created deployment.

Source code in src/zenml/zen_server/routers/pipeline_deployments_endpoints.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + PIPELINE_DEPLOYMENTS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["deployments"],
)
@handle_exceptions
def create_deployment(
    deployment: PipelineDeploymentRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> PipelineDeploymentResponse:
    """Creates a deployment.

    Args:
        deployment: Deployment to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created deployment.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        deployment.project = project.id

    return verify_permissions_and_create_entity(
        request_model=deployment,
        create_method=zen_store().create_deployment,
    )
delete_deployment(deployment_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific deployment.

Parameters:

Name Type Description Default
deployment_id UUID

ID of the deployment to delete.

required
Source code in src/zenml/zen_server/routers/pipeline_deployments_endpoints.py
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
@router.delete(
    "/{deployment_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_deployment(
    deployment_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific deployment.

    Args:
        deployment_id: ID of the deployment to delete.
    """
    verify_permissions_and_delete_entity(
        id=deployment_id,
        get_method=zen_store().get_deployment,
        delete_method=zen_store().delete_deployment,
    )
deployment_logs(deployment_id: UUID, _: AuthContext = Security(authorize)) -> str

Get deployment logs.

Parameters:

Name Type Description Default
deployment_id UUID

ID of the deployment.

required

Returns:

Type Description
str

The deployment logs.

Source code in src/zenml/zen_server/routers/pipeline_deployments_endpoints.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
@router.get(
    "/{deployment_id}/logs",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def deployment_logs(
    deployment_id: UUID,
    _: AuthContext = Security(authorize),
) -> str:
    """Get deployment logs.

    Args:
        deployment_id: ID of the deployment.

    Returns:
        The deployment logs.
    """
    deployment = verify_permissions_and_get_entity(
        id=deployment_id,
        get_method=zen_store().get_deployment,
        hydrate=True,
    )

    return workload_manager().get_logs(workload_id=deployment.id)
get_deployment(deployment_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> PipelineDeploymentResponse

Gets a specific deployment using its unique id.

Parameters:

Name Type Description Default
deployment_id UUID

ID of the deployment to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
PipelineDeploymentResponse

A specific deployment object.

Source code in src/zenml/zen_server/routers/pipeline_deployments_endpoints.py
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@router.get(
    "/{deployment_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_deployment(
    deployment_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> PipelineDeploymentResponse:
    """Gets a specific deployment using its unique id.

    Args:
        deployment_id: ID of the deployment to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific deployment object.
    """
    return verify_permissions_and_get_entity(
        id=deployment_id,
        get_method=zen_store().get_deployment,
        hydrate=hydrate,
    )
list_deployments(deployment_filter_model: PipelineDeploymentFilter = Depends(make_dependable(PipelineDeploymentFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[PipelineDeploymentResponse]

Gets a list of deployments.

Parameters:

Name Type Description Default
deployment_filter_model PipelineDeploymentFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(PipelineDeploymentFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[PipelineDeploymentResponse]

List of deployment objects matching the filter criteria.

Source code in src/zenml/zen_server/routers/pipeline_deployments_endpoints.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + PIPELINE_DEPLOYMENTS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["deployments"],
)
@handle_exceptions
def list_deployments(
    deployment_filter_model: PipelineDeploymentFilter = Depends(
        make_dependable(PipelineDeploymentFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[PipelineDeploymentResponse]:
    """Gets a list of deployments.

    Args:
        deployment_filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project to filter by.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of deployment objects matching the filter criteria.
    """
    if project_name_or_id:
        deployment_filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=deployment_filter_model,
        resource_type=ResourceType.PIPELINE_DEPLOYMENT,
        list_method=zen_store().list_deployments,
        hydrate=hydrate,
    )
pipelines_endpoints

Endpoint definitions for pipelines.

Classes Functions
create_pipeline(pipeline: PipelineRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> PipelineResponse

Creates a pipeline.

Parameters:

Name Type Description Default
pipeline PipelineRequest

Pipeline to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
PipelineResponse

The created pipeline.

Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + PIPELINES,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["pipelines"],
)
@handle_exceptions
def create_pipeline(
    pipeline: PipelineRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> PipelineResponse:
    """Creates a pipeline.

    Args:
        pipeline: Pipeline to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created pipeline.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        pipeline.project = project.id

    # We limit pipeline namespaces, not pipeline versions
    skip_entitlements = (
        zen_store().count_pipelines(
            PipelineFilter(name=pipeline.name, project=pipeline.project)
        )
        > 0
    )

    return verify_permissions_and_create_entity(
        request_model=pipeline,
        create_method=zen_store().create_pipeline,
        skip_entitlements=skip_entitlements,
    )
delete_pipeline(pipeline_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific pipeline.

Parameters:

Name Type Description Default
pipeline_id UUID

ID of the pipeline to delete.

required
Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
@router.delete(
    "/{pipeline_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_pipeline(
    pipeline_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific pipeline.

    Args:
        pipeline_id: ID of the pipeline to delete.
    """
    pipeline = verify_permissions_and_delete_entity(
        id=pipeline_id,
        get_method=zen_store().get_pipeline,
        delete_method=zen_store().delete_pipeline,
    )

    should_decrement = (
        ResourceType.PIPELINE in server_config().reportable_resources
        and zen_store().count_pipelines(
            PipelineFilter(name=pipeline.name, project=pipeline.project.id)
        )
        == 0
    )
    if should_decrement:
        report_decrement(ResourceType.PIPELINE, resource_id=pipeline_id)
get_pipeline(pipeline_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> PipelineResponse

Gets a specific pipeline using its unique id.

Parameters:

Name Type Description Default
pipeline_id UUID

ID of the pipeline to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
PipelineResponse

A specific pipeline object.

Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@router.get(
    "/{pipeline_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_pipeline(
    pipeline_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> PipelineResponse:
    """Gets a specific pipeline using its unique id.

    Args:
        pipeline_id: ID of the pipeline to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific pipeline object.
    """
    return verify_permissions_and_get_entity(
        id=pipeline_id, get_method=zen_store().get_pipeline, hydrate=hydrate
    )
list_pipeline_runs(pipeline_run_filter_model: PipelineRunFilter = Depends(make_dependable(PipelineRunFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[PipelineRunResponse]

Get pipeline runs according to query filters.

Parameters:

Name Type Description Default
pipeline_run_filter_model PipelineRunFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(PipelineRunFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[PipelineRunResponse]

The pipeline runs according to query filters.

Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
@router.get(
    "/{pipeline_id}" + RUNS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_pipeline_runs(
    pipeline_run_filter_model: PipelineRunFilter = Depends(
        make_dependable(PipelineRunFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[PipelineRunResponse]:
    """Get pipeline runs according to query filters.

    Args:
        pipeline_run_filter_model: Filter model used for pagination, sorting,
            filtering
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The pipeline runs according to query filters.
    """
    return zen_store().list_runs(pipeline_run_filter_model, hydrate=hydrate)
list_pipelines(pipeline_filter_model: PipelineFilter = Depends(make_dependable(PipelineFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[PipelineResponse]

Gets a list of pipelines.

Parameters:

Name Type Description Default
pipeline_filter_model PipelineFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(PipelineFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[PipelineResponse]

List of pipeline objects matching the filter criteria.

Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + PIPELINES,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["pipelines"],
)
@handle_exceptions
def list_pipelines(
    pipeline_filter_model: PipelineFilter = Depends(
        make_dependable(PipelineFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[PipelineResponse]:
    """Gets a list of pipelines.

    Args:
        pipeline_filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project to filter by.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of pipeline objects matching the filter criteria.
    """
    if project_name_or_id:
        pipeline_filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=pipeline_filter_model,
        resource_type=ResourceType.PIPELINE,
        list_method=zen_store().list_pipelines,
        hydrate=hydrate,
    )
update_pipeline(pipeline_id: UUID, pipeline_update: PipelineUpdate, _: AuthContext = Security(authorize)) -> PipelineResponse

Updates the attribute on a specific pipeline using its unique id.

Parameters:

Name Type Description Default
pipeline_id UUID

ID of the pipeline to get.

required
pipeline_update PipelineUpdate

the model containing the attributes to update.

required

Returns:

Type Description
PipelineResponse

The updated pipeline object.

Source code in src/zenml/zen_server/routers/pipelines_endpoints.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
@router.put(
    "/{pipeline_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_pipeline(
    pipeline_id: UUID,
    pipeline_update: PipelineUpdate,
    _: AuthContext = Security(authorize),
) -> PipelineResponse:
    """Updates the attribute on a specific pipeline using its unique id.

    Args:
        pipeline_id: ID of the pipeline to get.
        pipeline_update: the model containing the attributes to update.

    Returns:
        The updated pipeline object.
    """
    return verify_permissions_and_update_entity(
        id=pipeline_id,
        update_model=pipeline_update,
        get_method=zen_store().get_pipeline,
        update_method=zen_store().update_pipeline,
    )
plugin_endpoints

Endpoint definitions for plugin flavors.

Classes Functions
get_flavor(name: str, type: PluginType = Query(..., alias='type'), subtype: PluginSubType = Query(..., alias='subtype'), _: AuthContext = Security(authorize)) -> BasePluginFlavorResponse

Returns the requested flavor.

Parameters:

Name Type Description Default
name str

Name of the flavor.

required
type PluginType

Type of Plugin

Query(..., alias='type')
subtype PluginSubType

Subtype of Plugin

Query(..., alias='subtype')

Returns:

Type Description
BasePluginFlavorResponse

The requested flavor response.

Source code in src/zenml/zen_server/routers/plugin_endpoints.py
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@plugin_router.get(
    "/{name}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_flavor(
    name: str,
    type: PluginType = Query(..., alias="type"),
    subtype: PluginSubType = Query(..., alias="subtype"),
    _: AuthContext = Security(authorize),
) -> BasePluginFlavorResponse:  # type: ignore[type-arg]
    """Returns the requested flavor.

    Args:
        name: Name of the flavor.
        type: Type of Plugin
        subtype: Subtype of Plugin

    Returns:
        The requested flavor response.
    """
    plugin_flavor = plugin_flavor_registry().get_flavor_class(
        name=name, _type=type, subtype=subtype
    )
    return plugin_flavor.get_flavor_response_model(hydrate=True)
list_flavors(type: PluginType, subtype: PluginSubType, page: int = PAGINATION_STARTING_PAGE, size: int = PAGE_SIZE_DEFAULT, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[BasePluginFlavorResponse]

Returns all event flavors.

Parameters:

Name Type Description Default
type PluginType

The type of Plugin

required
subtype PluginSubType

The subtype of the plugin

required
page int

Page for pagination (offset +1)

PAGINATION_STARTING_PAGE
size int

Page size for pagination

PAGE_SIZE_DEFAULT
hydrate bool

Whether to hydrate the response bodies

False

Returns:

Type Description
Page[BasePluginFlavorResponse]

A page of flavors.

Source code in src/zenml/zen_server/routers/plugin_endpoints.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
@plugin_router.get(
    "",
    response_model=Page[BasePluginFlavorResponse],  # type: ignore[type-arg]
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_flavors(
    type: PluginType,
    subtype: PluginSubType,
    page: int = PAGINATION_STARTING_PAGE,
    size: int = PAGE_SIZE_DEFAULT,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[BasePluginFlavorResponse]:  # type: ignore[type-arg]
    """Returns all event flavors.

    Args:
        type: The type of Plugin
        subtype: The subtype of the plugin
        page: Page for pagination (offset +1)
        size: Page size for pagination
        hydrate: Whether to hydrate the response bodies

    Returns:
        A page of flavors.
    """
    flavors = plugin_flavor_registry().list_available_flavor_responses_for_type_and_subtype(
        _type=type, subtype=subtype, page=page, size=size, hydrate=hydrate
    )
    return flavors
projects_endpoints

Endpoint definitions for projects.

Classes Functions
create_project(project_request: ProjectRequest, _: AuthContext = Security(authorize)) -> ProjectResponse

Creates a project based on the requestBody.

noqa: DAR401

Parameters:

Name Type Description Default
project_request ProjectRequest

Project to create.

required

Returns:

Type Description
ProjectResponse

The created project.

Source code in src/zenml/zen_server/routers/projects_endpoints.py
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@workspace_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
)
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_project(
    project_request: ProjectRequest,
    _: AuthContext = Security(authorize),
) -> ProjectResponse:
    """Creates a project based on the requestBody.

    # noqa: DAR401

    Args:
        project_request: Project to create.

    Returns:
        The created project.
    """
    return verify_permissions_and_create_entity(
        request_model=project_request,
        create_method=zen_store().create_project,
    )
delete_project(project_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Deletes a project.

Parameters:

Name Type Description Default
project_name_or_id Union[str, UUID]

Name or ID of the project.

required
Source code in src/zenml/zen_server/routers/projects_endpoints.py
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
@workspace_router.delete(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
)
@router.delete(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_project(
    project_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a project.

    Args:
        project_name_or_id: Name or ID of the project.
    """
    project = verify_permissions_and_delete_entity(
        id=project_name_or_id,
        get_method=zen_store().get_project,
        delete_method=zen_store().delete_project,
    )
    if server_config().feature_gate_enabled:
        if ResourceType.PROJECT in server_config().reportable_resources:
            report_decrement(ResourceType.PROJECT, resource_id=project.id)
get_project(project_name_or_id: Union[str, UUID], hydrate: bool = True, _: AuthContext = Security(authorize)) -> ProjectResponse

Get a project for given name.

noqa: DAR401

Parameters:

Name Type Description Default
project_name_or_id Union[str, UUID]

Name or ID of the project.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ProjectResponse

The requested project.

Source code in src/zenml/zen_server/routers/projects_endpoints.py
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@workspace_router.get(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
)
@router.get(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_project(
    project_name_or_id: Union[str, UUID],
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ProjectResponse:
    """Get a project for given name.

    # noqa: DAR401

    Args:
        project_name_or_id: Name or ID of the project.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested project.
    """
    return verify_permissions_and_get_entity(
        id=project_name_or_id,
        get_method=zen_store().get_project,
        hydrate=hydrate,
    )
get_project_statistics(project_name_or_id: Union[str, UUID], auth_context: AuthContext = Security(authorize)) -> ProjectStatistics

Gets statistics of a project.

noqa: DAR401

Parameters:

Name Type Description Default
project_name_or_id Union[str, UUID]

Name or ID of the project to get statistics for.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
ProjectStatistics

Project statistics.

Source code in src/zenml/zen_server/routers/projects_endpoints.py
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
@workspace_router.get(
    "/{project_name_or_id}" + STATISTICS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
)
@router.get(
    "/{project_name_or_id}" + STATISTICS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_project_statistics(
    project_name_or_id: Union[str, UUID],
    auth_context: AuthContext = Security(authorize),
) -> ProjectStatistics:
    """Gets statistics of a project.

    # noqa: DAR401

    Args:
        project_name_or_id: Name or ID of the project to get statistics for.
        auth_context: Authentication context.

    Returns:
        Project statistics.
    """
    project = verify_permissions_and_get_entity(
        id=project_name_or_id,
        get_method=zen_store().get_project,
    )

    user_id = auth_context.user.id

    run_filter = PipelineRunFilter(project=project.id)
    run_filter.configure_rbac(
        authenticated_user_id=user_id,
        id=get_allowed_resource_ids(
            resource_type=ResourceType.PIPELINE_RUN, project_id=project.id
        ),
    )

    pipeline_filter = PipelineFilter(project=project.id)
    pipeline_filter.configure_rbac(
        authenticated_user_id=user_id,
        id=get_allowed_resource_ids(
            resource_type=ResourceType.PIPELINE, project_id=project.id
        ),
    )

    return ProjectStatistics(
        pipelines=zen_store().count_pipelines(filter_model=pipeline_filter),
        runs=zen_store().count_runs(filter_model=run_filter),
    )
list_projects(project_filter_model: ProjectFilter = Depends(make_dependable(ProjectFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ProjectResponse]

Lists all projects in the organization.

Parameters:

Name Type Description Default
project_filter_model ProjectFilter

Filter model used for pagination, sorting, filtering,

Depends(make_dependable(ProjectFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ProjectResponse]

A list of projects.

Source code in src/zenml/zen_server/routers/projects_endpoints.py
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
@workspace_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
)
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_projects(
    project_filter_model: ProjectFilter = Depends(
        make_dependable(ProjectFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ProjectResponse]:
    """Lists all projects in the organization.

    Args:
        project_filter_model: Filter model used for pagination, sorting,
            filtering,
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of projects.
    """
    return verify_permissions_and_list_entities(
        filter_model=project_filter_model,
        resource_type=ResourceType.PROJECT,
        list_method=zen_store().list_projects,
        hydrate=hydrate,
    )
update_project(project_name_or_id: UUID, project_update: ProjectUpdate, _: AuthContext = Security(authorize)) -> ProjectResponse

Get a project for given name.

noqa: DAR401

Parameters:

Name Type Description Default
project_name_or_id UUID

Name or ID of the project to update.

required
project_update ProjectUpdate

the project to use to update

required

Returns:

Type Description
ProjectResponse

The updated project.

Source code in src/zenml/zen_server/routers/projects_endpoints.py
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@workspace_router.put(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
)
@router.put(
    "/{project_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_project(
    project_name_or_id: UUID,
    project_update: ProjectUpdate,
    _: AuthContext = Security(authorize),
) -> ProjectResponse:
    """Get a project for given name.

    # noqa: DAR401

    Args:
        project_name_or_id: Name or ID of the project to update.
        project_update: the project to use to update

    Returns:
        The updated project.
    """
    return verify_permissions_and_update_entity(
        id=project_name_or_id,
        update_model=project_update,
        get_method=zen_store().get_project,
        update_method=zen_store().update_project,
    )
run_metadata_endpoints

Endpoint definitions for run metadata.

Classes Functions
create_run_metadata(run_metadata: RunMetadataRequest, project_name_or_id: Optional[Union[str, UUID]] = None, auth_context: AuthContext = Security(authorize)) -> None

Creates run metadata.

Parameters:

Name Type Description Default
run_metadata RunMetadataRequest

The run metadata to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
auth_context AuthContext

Authentication context.

Security(authorize)

Raises:

Type Description
RuntimeError

If the resource type is not supported.

Source code in src/zenml/zen_server/routers/run_metadata_endpoints.py
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + RUN_METADATA,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["run_metadata"],
)
@handle_exceptions
def create_run_metadata(
    run_metadata: RunMetadataRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    auth_context: AuthContext = Security(authorize),
) -> None:
    """Creates run metadata.

    Args:
        run_metadata: The run metadata to create.
        project_name_or_id: Optional name or ID of the project.
        auth_context: Authentication context.

    Raises:
        RuntimeError: If the resource type is not supported.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        run_metadata.project = project.id

    run_metadata.user = auth_context.user.id

    verify_models: List[Any] = []
    for resource in run_metadata.resources:
        if resource.type == MetadataResourceTypes.PIPELINE_RUN:
            verify_models.append(zen_store().get_run(resource.id))
        elif resource.type == MetadataResourceTypes.STEP_RUN:
            verify_models.append(zen_store().get_run_step(resource.id))
        elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
            verify_models.append(zen_store().get_artifact_version(resource.id))
        elif resource.type == MetadataResourceTypes.MODEL_VERSION:
            verify_models.append(zen_store().get_model_version(resource.id))
        elif resource.type == MetadataResourceTypes.SCHEDULE:
            verify_models.append(zen_store().get_schedule(resource.id))
        else:
            raise RuntimeError(f"Unknown resource type: {resource.type}")

    batch_verify_permissions_for_models(
        models=verify_models,
        action=Action.UPDATE,
    )

    verify_permission_for_model(model=run_metadata, action=Action.CREATE)

    zen_store().create_run_metadata(run_metadata)
run_templates_endpoints

Endpoint definitions for run templates.

Classes Functions
create_run_template(run_template: RunTemplateRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> RunTemplateResponse

Create a run template.

Parameters:

Name Type Description Default
run_template RunTemplateRequest

Run template to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
RunTemplateResponse

The created run template.

Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + RUN_TEMPLATES,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["run_templates"],
)
@handle_exceptions
def create_run_template(
    run_template: RunTemplateRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> RunTemplateResponse:
    """Create a run template.

    Args:
        run_template: Run template to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created run template.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        run_template.project = project.id

    return verify_permissions_and_create_entity(
        request_model=run_template,
        create_method=zen_store().create_run_template,
    )
create_template_run(template_id: UUID, background_tasks: BackgroundTasks, config: Optional[PipelineRunConfiguration] = None, auth_context: AuthContext = Security(authorize)) -> PipelineRunResponse

Run a pipeline from a template.

Parameters:

Name Type Description Default
template_id UUID

The ID of the template.

required
background_tasks BackgroundTasks

Background tasks.

required
config Optional[PipelineRunConfiguration]

Configuration for the pipeline run.

None
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
PipelineRunResponse

The created pipeline run.

Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
@router.post(
    "/{template_id}/runs",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def create_template_run(
    template_id: UUID,
    background_tasks: BackgroundTasks,
    config: Optional[PipelineRunConfiguration] = None,
    auth_context: AuthContext = Security(authorize),
) -> PipelineRunResponse:
    """Run a pipeline from a template.

    Args:
        template_id: The ID of the template.
        background_tasks: Background tasks.
        config: Configuration for the pipeline run.
        auth_context: Authentication context.

    Returns:
        The created pipeline run.
    """
    from zenml.zen_server.template_execution.utils import run_template

    with track_handler(
        event=AnalyticsEvent.EXECUTED_RUN_TEMPLATE,
    ) as analytics_handler:
        template = verify_permissions_and_get_entity(
            id=template_id,
            get_method=zen_store().get_run_template,
            hydrate=True,
        )
        analytics_handler.metadata = {
            "project_id": template.project.id,
        }

        verify_permission(
            resource_type=ResourceType.PIPELINE_DEPLOYMENT,
            action=Action.CREATE,
            project_id=template.project.id,
        )
        verify_permission(
            resource_type=ResourceType.PIPELINE_RUN,
            action=Action.CREATE,
            project_id=template.project.id,
        )

        return run_template(
            template=template,
            auth_context=auth_context,
            background_tasks=background_tasks,
            run_config=config,
        )
delete_run_template(template_id: UUID, _: AuthContext = Security(authorize)) -> None

Delete a run template.

Parameters:

Name Type Description Default
template_id UUID

ID of the run template to delete.

required
Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
@router.delete(
    "/{template_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_run_template(
    template_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a run template.

    Args:
        template_id: ID of the run template to delete.
    """
    verify_permissions_and_delete_entity(
        id=template_id,
        get_method=zen_store().get_run_template,
        delete_method=zen_store().delete_run_template,
    )
get_run_template(template_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> RunTemplateResponse

Get a run template.

Parameters:

Name Type Description Default
template_id UUID

ID of the run template to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
RunTemplateResponse

The run template.

Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@router.get(
    "/{template_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_run_template(
    template_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> RunTemplateResponse:
    """Get a run template.

    Args:
        template_id: ID of the run template to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The run template.
    """
    return verify_permissions_and_get_entity(
        id=template_id,
        get_method=zen_store().get_run_template,
        hydrate=hydrate,
    )
list_run_templates(filter_model: RunTemplateFilter = Depends(make_dependable(RunTemplateFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[RunTemplateResponse]

Get a page of run templates.

Parameters:

Name Type Description Default
filter_model RunTemplateFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(RunTemplateFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[RunTemplateResponse]

Page of run templates.

Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + RUN_TEMPLATES,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["run_templates"],
)
@handle_exceptions
def list_run_templates(
    filter_model: RunTemplateFilter = Depends(
        make_dependable(RunTemplateFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[RunTemplateResponse]:
    """Get a page of run templates.

    Args:
        filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page of run templates.
    """
    if project_name_or_id:
        filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=filter_model,
        resource_type=ResourceType.RUN_TEMPLATE,
        list_method=zen_store().list_run_templates,
        hydrate=hydrate,
    )
update_run_template(template_id: UUID, update: RunTemplateUpdate, _: AuthContext = Security(authorize)) -> RunTemplateResponse

Update a run template.

Parameters:

Name Type Description Default
template_id UUID

ID of the run template to get.

required
update RunTemplateUpdate

The updates to apply.

required

Returns:

Type Description
RunTemplateResponse

The updated run template.

Source code in src/zenml/zen_server/routers/run_templates_endpoints.py
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@router.put(
    "/{template_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_run_template(
    template_id: UUID,
    update: RunTemplateUpdate,
    _: AuthContext = Security(authorize),
) -> RunTemplateResponse:
    """Update a run template.

    Args:
        template_id: ID of the run template to get.
        update: The updates to apply.

    Returns:
        The updated run template.
    """
    return verify_permissions_and_update_entity(
        id=template_id,
        update_model=update,
        get_method=zen_store().get_run_template,
        update_method=zen_store().update_run_template,
    )
runs_endpoints

Endpoint definitions for pipeline runs.

Classes Functions
delete_run(run_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a run.

Parameters:

Name Type Description Default
run_id UUID

ID of the run.

required
Source code in src/zenml/zen_server/routers/runs_endpoints.py
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
@router.delete(
    "/{run_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_run(
    run_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a run.

    Args:
        run_id: ID of the run.
    """
    verify_permissions_and_delete_entity(
        id=run_id,
        get_method=zen_store().get_run,
        delete_method=zen_store().delete_run,
    )
get_or_create_pipeline_run(pipeline_run: PipelineRunRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> Tuple[PipelineRunResponse, bool]

Get or create a pipeline run.

Parameters:

Name Type Description Default
pipeline_run PipelineRunRequest

Pipeline run to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
PipelineRunResponse

The pipeline run and a boolean indicating whether the run was created

bool

or not.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + RUNS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["runs"],
)
@handle_exceptions
def get_or_create_pipeline_run(
    pipeline_run: PipelineRunRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> Tuple[PipelineRunResponse, bool]:
    """Get or create a pipeline run.

    Args:
        pipeline_run: Pipeline run to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The pipeline run and a boolean indicating whether the run was created
        or not.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        pipeline_run.project = project.id

    return verify_permissions_and_get_or_create_entity(
        request_model=pipeline_run,
        get_or_create_method=zen_store().get_or_create_run,
    )
get_pipeline_configuration(run_id: UUID, _: AuthContext = Security(authorize)) -> Dict[str, Any]

Get the pipeline configuration of a specific pipeline run using its ID.

Parameters:

Name Type Description Default
run_id UUID

ID of the pipeline run to get.

required

Returns:

Type Description
Dict[str, Any]

The pipeline configuration of the pipeline run.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
@router.get(
    "/{run_id}" + PIPELINE_CONFIGURATION,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_pipeline_configuration(
    run_id: UUID,
    _: AuthContext = Security(authorize),
) -> Dict[str, Any]:
    """Get the pipeline configuration of a specific pipeline run using its ID.

    Args:
        run_id: ID of the pipeline run to get.

    Returns:
        The pipeline configuration of the pipeline run.
    """
    run = verify_permissions_and_get_entity(
        id=run_id, get_method=zen_store().get_run, hydrate=True
    )
    return run.config.model_dump()
get_run(run_id: UUID, hydrate: bool = True, refresh_status: bool = False, _: AuthContext = Security(authorize)) -> PipelineRunResponse

Get a specific pipeline run using its ID.

Parameters:

Name Type Description Default
run_id UUID

ID of the pipeline run to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True
refresh_status bool

Flag deciding whether we should try to refresh the status of the pipeline run using its orchestrator.

False

Returns:

Type Description
PipelineRunResponse

The pipeline run.

Raises:

Type Description
RuntimeError

If the stack or the orchestrator of the run is deleted.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@router.get(
    "/{run_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_run(
    run_id: UUID,
    hydrate: bool = True,
    refresh_status: bool = False,
    _: AuthContext = Security(authorize),
) -> PipelineRunResponse:
    """Get a specific pipeline run using its ID.

    Args:
        run_id: ID of the pipeline run to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        refresh_status: Flag deciding whether we should try to refresh
            the status of the pipeline run using its orchestrator.

    Returns:
        The pipeline run.

    Raises:
        RuntimeError: If the stack or the orchestrator of the run is deleted.
    """
    run = verify_permissions_and_get_entity(
        id=run_id, get_method=zen_store().get_run, hydrate=hydrate
    )
    if refresh_status:
        try:
            # Check the stack and its orchestrator
            if run.stack is not None:
                orchestrators = run.stack.components.get(
                    StackComponentType.ORCHESTRATOR, []
                )
                if orchestrators:
                    verify_permission_for_model(
                        model=orchestrators[0], action=Action.READ
                    )
                else:
                    raise RuntimeError(
                        f"The orchestrator, the run '{run.id}' was executed "
                        "with, is deleted."
                    )
            else:
                raise RuntimeError(
                    f"The stack, the run '{run.id}' was executed on, is deleted."
                )

            run = run.refresh_run_status()

        except Exception as e:
            logger.warning(
                "An error occurred while refreshing the status of the "
                f"pipeline run: {e}"
            )
    return run
get_run_status(run_id: UUID, _: AuthContext = Security(authorize)) -> ExecutionStatus

Get the status of a specific pipeline run.

Parameters:

Name Type Description Default
run_id UUID

ID of the pipeline run for which to get the status.

required

Returns:

Type Description
ExecutionStatus

The status of the pipeline run.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
@router.get(
    "/{run_id}" + STATUS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_run_status(
    run_id: UUID,
    _: AuthContext = Security(authorize),
) -> ExecutionStatus:
    """Get the status of a specific pipeline run.

    Args:
        run_id: ID of the pipeline run for which to get the status.

    Returns:
        The status of the pipeline run.
    """
    run = verify_permissions_and_get_entity(
        id=run_id, get_method=zen_store().get_run, hydrate=False
    )
    return run.status
get_run_steps(run_id: UUID, step_run_filter_model: StepRunFilter = Depends(make_dependable(StepRunFilter)), _: AuthContext = Security(authorize)) -> Page[StepRunResponse]

Get all steps for a given pipeline run.

Parameters:

Name Type Description Default
run_id UUID

ID of the pipeline run.

required
step_run_filter_model StepRunFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(StepRunFilter))

Returns:

Type Description
Page[StepRunResponse]

The steps for a given pipeline run.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
@router.get(
    "/{run_id}" + STEPS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_run_steps(
    run_id: UUID,
    step_run_filter_model: StepRunFilter = Depends(
        make_dependable(StepRunFilter)
    ),
    _: AuthContext = Security(authorize),
) -> Page[StepRunResponse]:
    """Get all steps for a given pipeline run.

    Args:
        run_id: ID of the pipeline run.
        step_run_filter_model: Filter model used for pagination, sorting,
            filtering

    Returns:
        The steps for a given pipeline run.
    """
    verify_permissions_and_get_entity(
        id=run_id, get_method=zen_store().get_run, hydrate=False
    )
    step_run_filter_model.pipeline_run_id = run_id
    return zen_store().list_run_steps(step_run_filter_model)
list_runs(runs_filter_model: PipelineRunFilter = Depends(make_dependable(PipelineRunFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[PipelineRunResponse]

Get pipeline runs according to query filters.

Parameters:

Name Type Description Default
runs_filter_model PipelineRunFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(PipelineRunFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[PipelineRunResponse]

The pipeline runs according to query filters.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + RUNS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["runs"],
)
@handle_exceptions
def list_runs(
    runs_filter_model: PipelineRunFilter = Depends(
        make_dependable(PipelineRunFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[PipelineRunResponse]:
    """Get pipeline runs according to query filters.

    Args:
        runs_filter_model: Filter model used for pagination, sorting, filtering.
        project_name_or_id: Optional name or ID of the project.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The pipeline runs according to query filters.
    """
    if project_name_or_id:
        runs_filter_model.project = project_name_or_id

    return verify_permissions_and_list_entities(
        filter_model=runs_filter_model,
        resource_type=ResourceType.PIPELINE_RUN,
        list_method=zen_store().list_runs,
        hydrate=hydrate,
    )
refresh_run_status(run_id: UUID, _: AuthContext = Security(authorize)) -> None

Refreshes the status of a specific pipeline run.

Parameters:

Name Type Description Default
run_id UUID

ID of the pipeline run to refresh.

required

Raises:

Type Description
RuntimeError

If the stack or the orchestrator of the run is deleted.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
@router.get(
    "/{run_id}" + REFRESH,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def refresh_run_status(
    run_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Refreshes the status of a specific pipeline run.

    Args:
        run_id: ID of the pipeline run to refresh.

    Raises:
        RuntimeError: If the stack or the orchestrator of the run is deleted.
    """
    # Verify access to the run
    run = verify_permissions_and_get_entity(
        id=run_id,
        get_method=zen_store().get_run,
        hydrate=True,
    )

    # Check the stack and its orchestrator
    if run.stack is not None:
        orchestrators = run.stack.components.get(
            StackComponentType.ORCHESTRATOR, []
        )
        if orchestrators:
            verify_permission_for_model(
                model=orchestrators[0], action=Action.READ
            )
        else:
            raise RuntimeError(
                f"The orchestrator, the run '{run.id}' was executed with, is "
                "deleted."
            )
    else:
        raise RuntimeError(
            f"The stack, the run '{run.id}' was executed on, is deleted."
        )
    run.refresh_run_status()
update_run(run_id: UUID, run_model: PipelineRunUpdate, _: AuthContext = Security(authorize)) -> PipelineRunResponse

Updates a run.

Parameters:

Name Type Description Default
run_id UUID

ID of the run.

required
run_model PipelineRunUpdate

Run model to use for the update.

required

Returns:

Type Description
PipelineRunResponse

The updated run model.

Source code in src/zenml/zen_server/routers/runs_endpoints.py
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
@router.put(
    "/{run_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_run(
    run_id: UUID,
    run_model: PipelineRunUpdate,
    _: AuthContext = Security(authorize),
) -> PipelineRunResponse:
    """Updates a run.

    Args:
        run_id: ID of the run.
        run_model: Run model to use for the update.

    Returns:
        The updated run model.
    """
    return verify_permissions_and_update_entity(
        id=run_id,
        update_model=run_model,
        get_method=zen_store().get_run,
        update_method=zen_store().update_run,
    )
schedule_endpoints

Endpoint definitions for pipeline run schedules.

Classes Functions
create_schedule(schedule: ScheduleRequest, project_name_or_id: Optional[Union[str, UUID]] = None, auth_context: AuthContext = Security(authorize)) -> ScheduleResponse

Creates a schedule.

Parameters:

Name Type Description Default
schedule ScheduleRequest

Schedule to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
ScheduleResponse

The created schedule.

Source code in src/zenml/zen_server/routers/schedule_endpoints.py
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + SCHEDULES,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["schedules"],
)
@handle_exceptions
def create_schedule(
    schedule: ScheduleRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    auth_context: AuthContext = Security(authorize),
) -> ScheduleResponse:
    """Creates a schedule.

    Args:
        schedule: Schedule to create.
        project_name_or_id: Optional name or ID of the project.
        auth_context: Authentication context.

    Returns:
        The created schedule.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        schedule.project = project.id

    # NOTE: no RBAC is enforced currently for schedules, but we're
    # keeping the RBAC checks here for consistency
    return verify_permissions_and_create_entity(
        request_model=schedule,
        create_method=zen_store().create_schedule,
    )
delete_schedule(schedule_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific schedule using its unique id.

Parameters:

Name Type Description Default
schedule_id UUID

ID of the schedule to delete.

required
Source code in src/zenml/zen_server/routers/schedule_endpoints.py
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
@router.delete(
    "/{schedule_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_schedule(
    schedule_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific schedule using its unique id.

    Args:
        schedule_id: ID of the schedule to delete.
    """
    zen_store().delete_schedule(schedule_id=schedule_id)
get_schedule(schedule_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ScheduleResponse

Gets a specific schedule using its unique id.

Parameters:

Name Type Description Default
schedule_id UUID

ID of the schedule to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ScheduleResponse

A specific schedule object.

Source code in src/zenml/zen_server/routers/schedule_endpoints.py
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
@router.get(
    "/{schedule_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_schedule(
    schedule_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ScheduleResponse:
    """Gets a specific schedule using its unique id.

    Args:
        schedule_id: ID of the schedule to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific schedule object.
    """
    return zen_store().get_schedule(
        schedule_id=schedule_id,
        hydrate=hydrate,
    )
list_schedules(schedule_filter_model: ScheduleFilter = Depends(make_dependable(ScheduleFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ScheduleResponse]

Gets a list of schedules.

Parameters:

Name Type Description Default
schedule_filter_model ScheduleFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(ScheduleFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ScheduleResponse]

List of schedule objects.

Source code in src/zenml/zen_server/routers/schedule_endpoints.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + SCHEDULES,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["schedules"],
)
@handle_exceptions
def list_schedules(
    schedule_filter_model: ScheduleFilter = Depends(
        make_dependable(ScheduleFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ScheduleResponse]:
    """Gets a list of schedules.

    Args:
        schedule_filter_model: Filter model used for pagination, sorting,
            filtering
        project_name_or_id: Optional name or ID of the project.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of schedule objects.
    """
    if project_name_or_id:
        schedule_filter_model.project = project_name_or_id

    return zen_store().list_schedules(
        schedule_filter_model=schedule_filter_model,
        hydrate=hydrate,
    )
update_schedule(schedule_id: UUID, schedule_update: ScheduleUpdate, _: AuthContext = Security(authorize)) -> ScheduleResponse

Updates the attribute on a specific schedule using its unique id.

Parameters:

Name Type Description Default
schedule_id UUID

ID of the schedule to get.

required
schedule_update ScheduleUpdate

the model containing the attributes to update.

required

Returns:

Type Description
ScheduleResponse

The updated schedule object.

Source code in src/zenml/zen_server/routers/schedule_endpoints.py
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@router.put(
    "/{schedule_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_schedule(
    schedule_id: UUID,
    schedule_update: ScheduleUpdate,
    _: AuthContext = Security(authorize),
) -> ScheduleResponse:
    """Updates the attribute on a specific schedule using its unique id.

    Args:
        schedule_id: ID of the schedule to get.
        schedule_update: the model containing the attributes to update.

    Returns:
        The updated schedule object.
    """
    return zen_store().update_schedule(
        schedule_id=schedule_id,
        schedule_update=schedule_update,
    )
secrets_endpoints

Endpoint definitions for pipeline run secrets.

Classes Functions
backup_secrets(ignore_errors: bool = True, delete_secrets: bool = False, _: AuthContext = Security(authorize)) -> None

Backs up all secrets in the secrets store to the backup secrets store.

Parameters:

Name Type Description Default
ignore_errors bool

Whether to ignore individual errors when backing up secrets and continue with the backup operation until all secrets have been backed up.

True
delete_secrets bool

Whether to delete the secrets that have been successfully backed up from the primary secrets store. Setting this flag effectively moves all secrets from the primary secrets store to the backup secrets store.

False
Source code in src/zenml/zen_server/routers/secrets_endpoints.py
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
@op_router.put(
    SECRETS_BACKUP,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def backup_secrets(
    ignore_errors: bool = True,
    delete_secrets: bool = False,
    _: AuthContext = Security(authorize),
) -> None:
    """Backs up all secrets in the secrets store to the backup secrets store.

    Args:
        ignore_errors: Whether to ignore individual errors when backing up
            secrets and continue with the backup operation until all secrets
            have been backed up.
        delete_secrets: Whether to delete the secrets that have been
            successfully backed up from the primary secrets store. Setting
            this flag effectively moves all secrets from the primary secrets
            store to the backup secrets store.
    """
    verify_permission(
        resource_type=ResourceType.SECRET, action=Action.BACKUP_RESTORE
    )

    zen_store().backup_secrets(
        ignore_errors=ignore_errors, delete_secrets=delete_secrets
    )
create_secret(secret: SecretRequest, workspace_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> SecretResponse

Creates a secret.

Parameters:

Name Type Description Default
secret SecretRequest

Secret to create.

required
workspace_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the workspace.

None

Returns:

Type Description
SecretResponse

The created secret.

Source code in src/zenml/zen_server/routers/secrets_endpoints.py
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{workspace_name_or_id}" + SECRETS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["secrets"],
)
@handle_exceptions
def create_secret(
    secret: SecretRequest,
    workspace_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> SecretResponse:
    """Creates a secret.

    Args:
        secret: Secret to create.
        workspace_name_or_id: Optional name or ID of the workspace.

    Returns:
        The created secret.
    """
    return verify_permissions_and_create_entity(
        request_model=secret,
        create_method=zen_store().create_secret,
    )
delete_secret(secret_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific secret using its unique id.

Parameters:

Name Type Description Default
secret_id UUID

ID of the secret to delete.

required
Source code in src/zenml/zen_server/routers/secrets_endpoints.py
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
@router.delete(
    "/{secret_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_secret(
    secret_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific secret using its unique id.

    Args:
        secret_id: ID of the secret to delete.
    """
    verify_permissions_and_delete_entity(
        id=secret_id,
        get_method=zen_store().get_secret,
        delete_method=zen_store().delete_secret,
    )
get_secret(secret_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> SecretResponse

Gets a specific secret using its unique id.

Parameters:

Name Type Description Default
secret_id UUID

ID of the secret to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
SecretResponse

A specific secret object.

Source code in src/zenml/zen_server/routers/secrets_endpoints.py
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
@router.get(
    "/{secret_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_secret(
    secret_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> SecretResponse:
    """Gets a specific secret using its unique id.

    Args:
        secret_id: ID of the secret to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific secret object.
    """
    secret = verify_permissions_and_get_entity(
        id=secret_id,
        get_method=zen_store().get_secret,
        hydrate=hydrate,
    )

    if not has_permissions_for_model(secret, action=Action.READ_SECRET_VALUE):
        secret.remove_secrets()

    return secret
list_secrets(secret_filter_model: SecretFilter = Depends(make_dependable(SecretFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[SecretResponse]

Gets a list of secrets.

Parameters:

Name Type Description Default
secret_filter_model SecretFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(SecretFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[SecretResponse]

List of secret objects.

Source code in src/zenml/zen_server/routers/secrets_endpoints.py
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_secrets(
    secret_filter_model: SecretFilter = Depends(make_dependable(SecretFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[SecretResponse]:
    """Gets a list of secrets.

    Args:
        secret_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of secret objects.
    """
    secrets = verify_permissions_and_list_entities(
        filter_model=secret_filter_model,
        resource_type=ResourceType.SECRET,
        list_method=zen_store().list_secrets,
        hydrate=hydrate,
    )

    # This will be `None` if the user is allowed to read secret values
    # for all secrets
    allowed_ids = get_allowed_resource_ids(
        resource_type=ResourceType.SECRET,
        action=Action.READ_SECRET_VALUE,
    )

    if allowed_ids is not None:
        for secret in secrets.items:
            if secret.id in allowed_ids or is_owned_by_authenticated_user(
                secret
            ):
                continue

            secret.remove_secrets()

    return secrets
restore_secrets(ignore_errors: bool = False, delete_secrets: bool = False, _: AuthContext = Security(authorize)) -> None

Restores all secrets from the backup secrets store into the main secrets store.

Parameters:

Name Type Description Default
ignore_errors bool

Whether to ignore individual errors when restoring secrets and continue with the restore operation until all secrets have been restored.

False
delete_secrets bool

Whether to delete the secrets that have been successfully restored from the backup secrets store. Setting this flag effectively moves all secrets from the backup secrets store to the primary secrets store.

False
Source code in src/zenml/zen_server/routers/secrets_endpoints.py
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
@op_router.put(
    SECRETS_RESTORE,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def restore_secrets(
    ignore_errors: bool = False,
    delete_secrets: bool = False,
    _: AuthContext = Security(authorize),
) -> None:
    """Restores all secrets from the backup secrets store into the main secrets store.

    Args:
        ignore_errors: Whether to ignore individual errors when restoring
            secrets and continue with the restore operation until all secrets
            have been restored.
        delete_secrets: Whether to delete the secrets that have been
            successfully restored from the backup secrets store. Setting
            this flag effectively moves all secrets from the backup secrets
            store to the primary secrets store.
    """
    verify_permission(
        resource_type=ResourceType.SECRET,
        action=Action.BACKUP_RESTORE,
    )

    zen_store().restore_secrets(
        ignore_errors=ignore_errors, delete_secrets=delete_secrets
    )
update_secret(secret_id: UUID, secret_update: SecretUpdate, patch_values: Optional[bool] = False, _: AuthContext = Security(authorize)) -> SecretResponse

Updates the attribute on a specific secret using its unique id.

Parameters:

Name Type Description Default
secret_id UUID

ID of the secret to get.

required
secret_update SecretUpdate

the model containing the attributes to update.

required
patch_values Optional[bool]

Whether to patch the secret values or replace them.

False

Returns:

Type Description
SecretResponse

The updated secret object.

Source code in src/zenml/zen_server/routers/secrets_endpoints.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
@router.put(
    "/{secret_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_secret(
    secret_id: UUID,
    secret_update: SecretUpdate,
    patch_values: Optional[bool] = False,
    _: AuthContext = Security(authorize),
) -> SecretResponse:
    """Updates the attribute on a specific secret using its unique id.

    Args:
        secret_id: ID of the secret to get.
        secret_update: the model containing the attributes to update.
        patch_values: Whether to patch the secret values or replace them.

    Returns:
        The updated secret object.
    """
    if not patch_values:
        # If patch_values is False, interpret the update values as a complete
        # replacement of the existing secret values. The only adjustment we
        # need to make is to set the value of any keys that are not present in
        # the update to None, so that they are deleted.
        secret = zen_store().get_secret(secret_id=secret_id)
        for key in secret.values.keys():
            if secret_update.values is not None:
                if key not in secret_update.values:
                    secret_update.values[key] = None

    return verify_permissions_and_update_entity(
        id=secret_id,
        update_model=secret_update,
        get_method=zen_store().get_secret,
        update_method=zen_store().update_secret,
    )
server_endpoints

Endpoint definitions for authentication (login).

Classes Functions
activate_server(activate_request: ServerActivationRequest) -> Optional[UserResponse]

Updates a stack.

Parameters:

Name Type Description Default
activate_request ServerActivationRequest

The request to activate the server.

required

Returns:

Type Description
Optional[UserResponse]

The default admin user that was created during activation, if any.

Source code in src/zenml/zen_server/routers/server_endpoints.py
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
@router.put(
    ACTIVATE,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def activate_server(
    activate_request: ServerActivationRequest,
) -> Optional[UserResponse]:
    """Updates a stack.

    Args:
        activate_request: The request to activate the server.

    Returns:
        The default admin user that was created during activation, if any.
    """
    return zen_store().activate_server(activate_request)
get_onboarding_state(_: AuthContext = Security(authorize)) -> List[str]

Get the onboarding state of the server.

Returns:

Type Description
List[str]

The onboarding state of the server.

Source code in src/zenml/zen_server/routers/server_endpoints.py
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
@router.get(
    ONBOARDING_STATE,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def get_onboarding_state(
    _: AuthContext = Security(authorize),
) -> List[str]:
    """Get the onboarding state of the server.

    Returns:
        The onboarding state of the server.
    """
    return zen_store().get_onboarding_state()
get_server_statistics(auth_context: AuthContext = Security(authorize)) -> ServerStatistics

Gets server statistics.

Parameters:

Name Type Description Default
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
ServerStatistics

Statistics of the server.

Source code in src/zenml/zen_server/routers/server_endpoints.py
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
@router.get(
    STATISTICS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_server_statistics(
    auth_context: AuthContext = Security(authorize),
) -> ServerStatistics:
    """Gets server statistics.

    Args:
        auth_context: Authentication context.

    Returns:
        Statistics of the server.
    """
    user_id = auth_context.user.id
    component_filter = ComponentFilter()
    component_filter.configure_rbac(
        authenticated_user_id=user_id,
        id=get_allowed_resource_ids(
            resource_type=ResourceType.STACK_COMPONENT
        ),
    )

    project_filter = ProjectFilter()
    project_filter.configure_rbac(
        authenticated_user_id=user_id,
        id=get_allowed_resource_ids(resource_type=ResourceType.PROJECT),
    )

    stack_filter = StackFilter()
    stack_filter.configure_rbac(
        authenticated_user_id=user_id,
        id=get_allowed_resource_ids(resource_type=ResourceType.STACK),
    )

    return ServerStatistics(
        stacks=zen_store().count_stacks(filter_model=stack_filter),
        components=zen_store().count_stack_components(
            filter_model=component_filter
        ),
        projects=zen_store().count_projects(filter_model=project_filter),
    )
get_settings(_: AuthContext = Security(authorize), hydrate: bool = True) -> ServerSettingsResponse

Get settings of the server.

Parameters:

Name Type Description Default
hydrate bool

Whether to hydrate the response.

True

Returns:

Type Description
ServerSettingsResponse

Settings of the server.

Source code in src/zenml/zen_server/routers/server_endpoints.py
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
@router.get(
    SERVER_SETTINGS,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def get_settings(
    _: AuthContext = Security(authorize),
    hydrate: bool = True,
) -> ServerSettingsResponse:
    """Get settings of the server.

    Args:
        hydrate: Whether to hydrate the response.

    Returns:
        Settings of the server.
    """
    return zen_store().get_server_settings(hydrate=hydrate)
server_info() -> ServerModel

Get information about the server.

Returns:

Type Description
ServerModel

Information about the server.

Source code in src/zenml/zen_server/routers/server_endpoints.py
68
69
70
71
72
73
74
75
76
77
78
79
@router.get(
    INFO,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def server_info() -> ServerModel:
    """Get information about the server.

    Returns:
        Information about the server.
    """
    return zen_store().get_store_info()
server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo

Get information about the server load.

Returns:

Type Description
ServerLoadInfo

Information about the server load.

Source code in src/zenml/zen_server/routers/server_endpoints.py
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@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,
    )
update_server_settings(settings_update: ServerSettingsUpdate, auth_context: AuthContext = Security(authorize)) -> ServerSettingsResponse

Updates the settings of the server.

Parameters:

Name Type Description Default
settings_update ServerSettingsUpdate

Settings update.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Raises:

Type Description
IllegalOperationError

If trying to update admin properties without admin permissions.

Returns:

Type Description
ServerSettingsResponse

The updated settings.

Source code in src/zenml/zen_server/routers/server_endpoints.py
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
@router.put(
    SERVER_SETTINGS,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def update_server_settings(
    settings_update: ServerSettingsUpdate,
    auth_context: AuthContext = Security(authorize),
) -> ServerSettingsResponse:
    """Updates the settings of the server.

    Args:
        settings_update: Settings update.
        auth_context: Authentication context.

    Raises:
        IllegalOperationError: If trying to update admin properties without
            admin permissions.

    Returns:
        The updated settings.
    """
    if not server_config().rbac_enabled:
        will_update_admin_properties = bool(
            settings_update.model_dump(
                exclude_none=True, exclude={"onboarding_state"}
            )
        )

        if not auth_context.user.is_admin and will_update_admin_properties:
            raise IllegalOperationError(
                "Only admins can update server settings."
            )

    return zen_store().update_server_settings(settings_update)
version() -> str

Get version of the server.

Returns:

Type Description
str

String representing the version of the server.

Source code in src/zenml/zen_server/routers/server_endpoints.py
58
59
60
61
62
63
64
65
@router.get("/version")
def version() -> str:
    """Get version of the server.

    Returns:
        String representing the version of the server.
    """
    return zenml.__version__
service_accounts_endpoints

Endpoint definitions for API keys.

Classes Functions
create_api_key(service_account_id: UUID, api_key: APIKeyRequest, _: AuthContext = Security(authorize)) -> APIKeyResponse

Creates an API key for a service account.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account for which to create the API key.

required
api_key APIKeyRequest

API key to create.

required

Returns:

Type Description
APIKeyResponse

The created API key.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
@router.post(
    "/{service_account_id}" + API_KEYS,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_api_key(
    service_account_id: UUID,
    api_key: APIKeyRequest,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Creates an API key for a service account.

    Args:
        service_account_id: ID of the service account for which to create the
            API key.
        api_key: API key to create.

    Returns:
        The created API key.
    """

    def create_api_key_wrapper(
        api_key: APIKeyRequest,
    ) -> APIKeyResponse:
        return zen_store().create_api_key(
            service_account_id=service_account_id,
            api_key=api_key,
        )

    service_account = zen_store().get_service_account(service_account_id)

    return verify_permissions_and_create_entity(
        request_model=api_key,
        create_method=create_api_key_wrapper,
        surrogate_models=[service_account],
    )
create_service_account(service_account: ServiceAccountRequest, _: AuthContext = Security(authorize)) -> ServiceAccountResponse

Creates a service account.

Parameters:

Name Type Description Default
service_account ServiceAccountRequest

Service account to create.

required

Returns:

Type Description
ServiceAccountResponse

The created service account.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
@router.post(
    "",
    responses={
        401: error_response,
        409: error_response,
        422: error_response,
    },
)
@handle_exceptions
def create_service_account(
    service_account: ServiceAccountRequest,
    _: AuthContext = Security(authorize),
) -> ServiceAccountResponse:
    """Creates a service account.

    Args:
        service_account: Service account to create.

    Returns:
        The created service account.
    """
    return verify_permissions_and_create_entity(
        request_model=service_account,
        create_method=zen_store().create_service_account,
    )
delete_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Deletes an API key.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account to which the API key belongs.

required
api_key_name_or_id Union[str, UUID]

Name or ID of the API key to delete.

required
Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
@router.delete(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes an API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to delete.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.UPDATE)
    zen_store().delete_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
    )
delete_service_account(service_account_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Delete a specific service account.

Parameters:

Name Type Description Default
service_account_name_or_id Union[str, UUID]

Name or ID of the service account.

required
Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@router.delete(
    "/{service_account_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_service_account(
    service_account_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.
    """
    verify_permissions_and_delete_entity(
        id=service_account_name_or_id,
        get_method=zen_store().get_service_account,
        delete_method=zen_store().delete_service_account,
    )
get_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], hydrate: bool = True, _: AuthContext = Security(authorize)) -> APIKeyResponse

Returns the requested API key.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account to which the API key belongs.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True
api_key_name_or_id Union[str, UUID]

Name or ID of the API key to return.

required

Returns:

Type Description
APIKeyResponse

The requested API key.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
@router.get(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Returns the requested API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        api_key_name_or_id: Name or ID of the API key to return.

    Returns:
        The requested API key.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.READ)
    api_key = zen_store().get_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        hydrate=hydrate,
    )
    return api_key
get_service_account(service_account_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize), hydrate: bool = True) -> ServiceAccountResponse

Returns a specific service account.

Parameters:

Name Type Description Default
service_account_name_or_id Union[str, UUID]

Name or ID of the service account.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ServiceAccountResponse

The service account matching the given name or ID.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
@router.get(
    "/{service_account_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_service_account(
    service_account_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
    hydrate: bool = True,
) -> ServiceAccountResponse:
    """Returns a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The service account matching the given name or ID.
    """
    return verify_permissions_and_get_entity(
        id=service_account_name_or_id,
        get_method=zen_store().get_service_account,
        hydrate=hydrate,
    )
list_api_keys(service_account_id: UUID, filter_model: APIKeyFilter = Depends(make_dependable(APIKeyFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[APIKeyResponse]

List API keys associated with a service account.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account to which the API keys belong.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
filter_model APIKeyFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(APIKeyFilter))

Returns:

Type Description
Page[APIKeyResponse]

All API keys matching the filter and associated with the supplied

Page[APIKeyResponse]

service account.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
@router.get(
    "/{service_account_id}" + API_KEYS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_api_keys(
    service_account_id: UUID,
    filter_model: APIKeyFilter = Depends(make_dependable(APIKeyFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[APIKeyResponse]:
    """List API keys associated with a service account.

    Args:
        service_account_id: ID of the service account to which the API keys
            belong.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        filter_model: Filter model used for pagination, sorting,
            filtering

    Returns:
        All API keys matching the filter and associated with the supplied
        service account.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.READ)
    return zen_store().list_api_keys(
        service_account_id=service_account_id,
        filter_model=filter_model,
        hydrate=hydrate,
    )
list_service_accounts(filter_model: ServiceAccountFilter = Depends(make_dependable(ServiceAccountFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ServiceAccountResponse]

Returns a list of service accounts.

Parameters:

Name Type Description Default
filter_model ServiceAccountFilter

Model that takes care of filtering, sorting and pagination.

Depends(make_dependable(ServiceAccountFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ServiceAccountResponse]

A list of service accounts matching the filter.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_service_accounts(
    filter_model: ServiceAccountFilter = Depends(
        make_dependable(ServiceAccountFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ServiceAccountResponse]:
    """Returns a list of service accounts.

    Args:
        filter_model: Model that takes care of filtering, sorting and
            pagination.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of service accounts matching the filter.
    """
    return verify_permissions_and_list_entities(
        filter_model=filter_model,
        resource_type=ResourceType.SERVICE_ACCOUNT,
        list_method=zen_store().list_service_accounts,
        hydrate=hydrate,
    )
rotate_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], rotate_request: APIKeyRotateRequest, _: AuthContext = Security(authorize)) -> APIKeyResponse

Rotate an API key.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account to which the API key belongs.

required
api_key_name_or_id Union[str, UUID]

Name or ID of the API key to rotate.

required
rotate_request APIKeyRotateRequest

API key rotation request.

required

Returns:

Type Description
APIKeyResponse

The updated API key.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
@router.put(
    "/{service_account_id}"
    + API_KEYS
    + "/{api_key_name_or_id}"
    + API_KEY_ROTATE,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def rotate_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    rotate_request: APIKeyRotateRequest,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Rotate an API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to rotate.
        rotate_request: API key rotation request.

    Returns:
        The updated API key.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.UPDATE)
    return zen_store().rotate_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        rotate_request=rotate_request,
    )
update_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], api_key_update: APIKeyUpdate, _: AuthContext = Security(authorize)) -> APIKeyResponse

Updates an API key for a service account.

Parameters:

Name Type Description Default
service_account_id UUID

ID of the service account to which the API key belongs.

required
api_key_name_or_id Union[str, UUID]

Name or ID of the API key to update.

required
api_key_update APIKeyUpdate

API key update.

required

Returns:

Type Description
APIKeyResponse

The updated API key.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
@router.put(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def update_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    api_key_update: APIKeyUpdate,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Updates an API key for a service account.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to update.
        api_key_update: API key update.

    Returns:
        The updated API key.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.UPDATE)
    return zen_store().update_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        api_key_update=api_key_update,
    )
update_service_account(service_account_name_or_id: Union[str, UUID], service_account_update: ServiceAccountUpdate, _: AuthContext = Security(authorize)) -> ServiceAccountResponse

Updates a specific service account.

Parameters:

Name Type Description Default
service_account_name_or_id Union[str, UUID]

Name or ID of the service account.

required
service_account_update ServiceAccountUpdate

the service account to use for the update.

required

Returns:

Type Description
ServiceAccountResponse

The updated service account.

Source code in src/zenml/zen_server/routers/service_accounts_endpoints.py
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
@router.put(
    "/{service_account_name_or_id}",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def update_service_account(
    service_account_name_or_id: Union[str, UUID],
    service_account_update: ServiceAccountUpdate,
    _: AuthContext = Security(authorize),
) -> ServiceAccountResponse:
    """Updates a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.
        service_account_update: the service account to use for the update.

    Returns:
        The updated service account.
    """
    return verify_permissions_and_update_entity(
        id=service_account_name_or_id,
        update_model=service_account_update,
        get_method=zen_store().get_service_account,
        update_method=zen_store().update_service_account,
    )
service_connectors_endpoints

Endpoint definitions for service connectors.

Classes Functions
create_service_connector(connector: ServiceConnectorRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> ServiceConnectorResponse

Creates a service connector.

Parameters:

Name Type Description Default
connector ServiceConnectorRequest

Service connector to register.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
ServiceConnectorResponse

The created service connector.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + SERVICE_CONNECTORS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["service_connectors"],
)
@handle_exceptions
def create_service_connector(
    connector: ServiceConnectorRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResponse:
    """Creates a service connector.

    Args:
        connector: Service connector to register.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created service connector.
    """
    return verify_permissions_and_create_entity(
        request_model=connector,
        create_method=zen_store().create_service_connector,
    )
delete_service_connector(connector_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a service connector.

Parameters:

Name Type Description Default
connector_id UUID

ID of the service connector.

required
Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
@router.delete(
    "/{connector_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_service_connector(
    connector_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a service connector.

    Args:
        connector_id: ID of the service connector.
    """
    verify_permissions_and_delete_entity(
        id=connector_id,
        get_method=zen_store().get_service_connector,
        delete_method=zen_store().delete_service_connector,
    )
get_resources_based_on_service_connector_info(connector_info: Optional[ServiceConnectorInfo] = None, connector_uuid: Optional[UUID] = None, _: AuthContext = Security(authorize)) -> ServiceConnectorResourcesInfo

Gets the list of resources that a service connector can access.

Parameters:

Name Type Description Default
connector_info Optional[ServiceConnectorInfo]

The service connector info.

None
connector_uuid Optional[UUID]

The service connector uuid.

None

Returns:

Type Description
ServiceConnectorResourcesInfo

The list of resources that the service connector configuration has

ServiceConnectorResourcesInfo

access to and consumable from UI/CLI.

Raises:

Type Description
ValueError

If both connector_info and connector_uuid are provided.

ValueError

If neither connector_info nor connector_uuid are provided.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
@router.post(
    SERVICE_CONNECTOR_FULL_STACK,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def get_resources_based_on_service_connector_info(
    connector_info: Optional[ServiceConnectorInfo] = None,
    connector_uuid: Optional[UUID] = None,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResourcesInfo:
    """Gets the list of resources that a service connector can access.

    Args:
        connector_info: The service connector info.
        connector_uuid: The service connector uuid.

    Returns:
        The list of resources that the service connector configuration has
        access to and consumable from UI/CLI.

    Raises:
        ValueError: If both connector_info and connector_uuid are provided.
        ValueError: If neither connector_info nor connector_uuid are provided.
    """
    if connector_info is not None and connector_uuid is not None:
        raise ValueError(
            "Only one of connector_info or connector_uuid must be provided."
        )
    if connector_info is None and connector_uuid is None:
        raise ValueError(
            "Either connector_info or connector_uuid must be provided."
        )

    if connector_info is not None:
        verify_permission(
            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
        )
    elif connector_uuid is not None:
        verify_permission(
            resource_type=ResourceType.SERVICE_CONNECTOR,
            action=Action.READ,
            resource_id=connector_uuid,
        )

    return get_resources_options_from_resource_model_for_full_stack(
        connector_details=connector_info or connector_uuid  # type: ignore[arg-type]
    )
get_service_connector(connector_id: UUID, expand_secrets: bool = True, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ServiceConnectorResponse

Returns the requested service connector.

Parameters:

Name Type Description Default
connector_id UUID

ID of the service connector.

required
expand_secrets bool

Whether to expand secrets or not.

True
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ServiceConnectorResponse

The requested service connector.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
@router.get(
    "/{connector_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_service_connector(
    connector_id: UUID,
    expand_secrets: bool = True,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResponse:
    """Returns the requested service connector.

    Args:
        connector_id: ID of the service connector.
        expand_secrets: Whether to expand secrets or not.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested service connector.
    """
    connector = zen_store().get_service_connector(
        connector_id, hydrate=hydrate
    )
    verify_permission_for_model(connector, action=Action.READ)

    if (
        expand_secrets
        and connector.secret_id
        and has_permissions_for_model(
            connector, action=Action.READ_SECRET_VALUE
        )
    ):
        secret = zen_store().get_secret(secret_id=connector.secret_id)

        # Update the connector configuration with the secret.
        connector.configuration.update(secret.secret_values)

    return dehydrate_response_model(connector)
get_service_connector_client(connector_id: UUID, resource_type: Optional[str] = None, resource_id: Optional[str] = None, _: AuthContext = Security(authorize)) -> ServiceConnectorResponse

Get a service connector client for a service connector and given resource.

This requires the service connector implementation to be installed on the ZenML server, otherwise a 501 Not Implemented error will be returned.

Parameters:

Name Type Description Default
connector_id UUID

ID of the service connector.

required
resource_type Optional[str]

Type of the resource to list.

None
resource_id Optional[str]

ID of the resource to list.

None

Returns:

Type Description
ServiceConnectorResponse

A service connector client that can be used to access the given

ServiceConnectorResponse

resource.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
@router.get(
    "/{connector_id}" + SERVICE_CONNECTOR_CLIENT,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_service_connector_client(
    connector_id: UUID,
    resource_type: Optional[str] = None,
    resource_id: Optional[str] = None,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResponse:
    """Get a service connector client for a service connector and given resource.

    This requires the service connector implementation to be installed
    on the ZenML server, otherwise a 501 Not Implemented error will be
    returned.

    Args:
        connector_id: ID of the service connector.
        resource_type: Type of the resource to list.
        resource_id: ID of the resource to list.

    Returns:
        A service connector client that can be used to access the given
        resource.
    """
    connector = zen_store().get_service_connector(connector_id)
    verify_permission_for_model(model=connector, action=Action.READ)
    verify_permission_for_model(model=connector, action=Action.CLIENT)

    return zen_store().get_service_connector_client(
        service_connector_id=connector_id,
        resource_type=resource_type,
        resource_id=resource_id,
    )
get_service_connector_type(connector_type: str, _: AuthContext = Security(authorize)) -> ServiceConnectorTypeModel

Returns the requested service connector type.

Parameters:

Name Type Description Default
connector_type str

the service connector type identifier.

required

Returns:

Type Description
ServiceConnectorTypeModel

The requested service connector type.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
@types_router.get(
    "/{connector_type}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_service_connector_type(
    connector_type: str,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorTypeModel:
    """Returns the requested service connector type.

    Args:
        connector_type: the service connector type identifier.

    Returns:
        The requested service connector type.
    """
    return zen_store().get_service_connector_type(connector_type)
list_service_connector_resources(filter_model: ServiceConnectorFilter = Depends(make_dependable(ServiceConnectorFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, auth_context: AuthContext = Security(authorize)) -> List[ServiceConnectorResourcesModel]

List resources that can be accessed by service connectors.

Parameters:

Name Type Description Default
filter_model ServiceConnectorFilter

The filter model to use when fetching service connectors.

Depends(make_dependable(ServiceConnectorFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
List[ServiceConnectorResourcesModel]

The matching list of resources that available service

List[ServiceConnectorResourcesModel]

connectors have access to.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@router.get(
    SERVICE_CONNECTOR_RESOURCES,
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + SERVICE_CONNECTOR_RESOURCES,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["service_connectors"],
)
@handle_exceptions
def list_service_connector_resources(
    filter_model: ServiceConnectorFilter = Depends(
        make_dependable(ServiceConnectorFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    auth_context: AuthContext = Security(authorize),
) -> List[ServiceConnectorResourcesModel]:
    """List resources that can be accessed by service connectors.

    Args:
        filter_model: The filter model to use when fetching service
            connectors.
        project_name_or_id: Optional name or ID of the project.
        auth_context: Authentication context.

    Returns:
        The matching list of resources that available service
        connectors have access to.
    """
    allowed_ids = get_allowed_resource_ids(
        resource_type=ResourceType.SERVICE_CONNECTOR
    )
    filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id, id=allowed_ids
    )

    return zen_store().list_service_connector_resources(
        filter_model=filter_model,
    )
list_service_connector_types(connector_type: Optional[str] = None, resource_type: Optional[str] = None, auth_method: Optional[str] = None, _: AuthContext = Security(authorize)) -> List[ServiceConnectorTypeModel]

Get a list of service connector types.

Parameters:

Name Type Description Default
connector_type Optional[str]

Filter by connector type.

None
resource_type Optional[str]

Filter by resource type.

None
auth_method Optional[str]

Filter by auth method.

None

Returns:

Type Description
List[ServiceConnectorTypeModel]

List of service connector types.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
@types_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_service_connector_types(
    connector_type: Optional[str] = None,
    resource_type: Optional[str] = None,
    auth_method: Optional[str] = None,
    _: AuthContext = Security(authorize),
) -> List[ServiceConnectorTypeModel]:
    """Get a list of service connector types.

    Args:
        connector_type: Filter by connector type.
        resource_type: Filter by resource type.
        auth_method: Filter by auth method.

    Returns:
        List of service connector types.
    """
    connector_types = zen_store().list_service_connector_types(
        connector_type=connector_type,
        resource_type=resource_type,
        auth_method=auth_method,
    )

    return connector_types
list_service_connectors(connector_filter_model: ServiceConnectorFilter = Depends(make_dependable(ServiceConnectorFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, expand_secrets: bool = True, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ServiceConnectorResponse]

Get a list of all service connectors.

Parameters:

Name Type Description Default
connector_filter_model ServiceConnectorFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(ServiceConnectorFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
expand_secrets bool

Whether to expand secrets or not.

True
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ServiceConnectorResponse]

Page with list of service connectors matching the filter criteria.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + SERVICE_CONNECTORS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["service_connectors"],
)
@handle_exceptions
def list_service_connectors(
    connector_filter_model: ServiceConnectorFilter = Depends(
        make_dependable(ServiceConnectorFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    expand_secrets: bool = True,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ServiceConnectorResponse]:
    """Get a list of all service connectors.

    Args:
        connector_filter_model: Filter model used for pagination, sorting,
            filtering
        project_name_or_id: Optional name or ID of the project to filter by.
        expand_secrets: Whether to expand secrets or not.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page with list of service connectors matching the filter criteria.
    """
    connectors = verify_permissions_and_list_entities(
        filter_model=connector_filter_model,
        resource_type=ResourceType.SERVICE_CONNECTOR,
        list_method=zen_store().list_service_connectors,
        hydrate=hydrate,
    )

    if expand_secrets:
        # This will be `None` if the user is allowed to read secret values
        # for all service connectors
        allowed_ids = get_allowed_resource_ids(
            resource_type=ResourceType.SERVICE_CONNECTOR,
            action=Action.READ_SECRET_VALUE,
        )

        for connector in connectors.items:
            if not connector.secret_id:
                continue

            if allowed_ids is None or is_owned_by_authenticated_user(
                connector
            ):
                # The user either owns the connector or has permissions to
                # read secret values for all service connectors
                pass
            elif connector.id not in allowed_ids:
                # The user is not allowed to read secret values for this
                # connector. We don't raise an exception here but don't include
                # the secret values
                continue

            secret = zen_store().get_secret(secret_id=connector.secret_id)

            # Update the connector configuration with the secret.
            connector.configuration.update(secret.secret_values)

    return connectors
update_service_connector(connector_id: UUID, connector_update: ServiceConnectorUpdate, _: AuthContext = Security(authorize)) -> ServiceConnectorResponse

Updates a service connector.

Parameters:

Name Type Description Default
connector_id UUID

ID of the service connector.

required
connector_update ServiceConnectorUpdate

Service connector to use to update.

required

Returns:

Type Description
ServiceConnectorResponse

Updated service connector.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
@router.put(
    "/{connector_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_service_connector(
    connector_id: UUID,
    connector_update: ServiceConnectorUpdate,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResponse:
    """Updates a service connector.

    Args:
        connector_id: ID of the service connector.
        connector_update: Service connector to use to update.

    Returns:
        Updated service connector.
    """
    return verify_permissions_and_update_entity(
        id=connector_id,
        update_model=connector_update,
        get_method=zen_store().get_service_connector,
        update_method=zen_store().update_service_connector,
    )
validate_and_verify_service_connector(connector_id: UUID, resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, _: AuthContext = Security(authorize)) -> ServiceConnectorResourcesModel

Verifies if a service connector instance has access to one or more resources.

This requires the service connector implementation to be installed on the ZenML server, otherwise a 501 Not Implemented error will be returned.

Parameters:

Name Type Description Default
connector_id UUID

The ID of the service connector to verify.

required
resource_type Optional[str]

The type of resource to verify access to.

None
resource_id Optional[str]

The ID of the resource to verify access to.

None
list_resources bool

If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned.

True

Returns:

Type Description
ServiceConnectorResourcesModel

The list of resources that the service connector has access to, scoped

ServiceConnectorResourcesModel

to the supplied resource type and ID, if provided.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
@router.put(
    "/{connector_id}" + SERVICE_CONNECTOR_VERIFY,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def validate_and_verify_service_connector(
    connector_id: UUID,
    resource_type: Optional[str] = None,
    resource_id: Optional[str] = None,
    list_resources: bool = True,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResourcesModel:
    """Verifies if a service connector instance has access to one or more resources.

    This requires the service connector implementation to be installed
    on the ZenML server, otherwise a 501 Not Implemented error will be
    returned.

    Args:
        connector_id: The ID of the service connector to verify.
        resource_type: The type of resource to verify access to.
        resource_id: The ID of the resource to verify access to.
        list_resources: If True, the list of all resources accessible
            through the service connector and matching the supplied resource
            type and ID are returned.

    Returns:
        The list of resources that the service connector has access to, scoped
        to the supplied resource type and ID, if provided.
    """
    connector = zen_store().get_service_connector(connector_id)
    verify_permission_for_model(model=connector, action=Action.READ)

    return zen_store().verify_service_connector(
        service_connector_id=connector_id,
        resource_type=resource_type,
        resource_id=resource_id,
        list_resources=list_resources,
    )
validate_and_verify_service_connector_config(connector: ServiceConnectorRequest, list_resources: bool = True, _: AuthContext = Security(authorize)) -> ServiceConnectorResourcesModel

Verifies if a service connector configuration has access to resources.

This requires the service connector implementation to be installed on the ZenML server, otherwise a 501 Not Implemented error will be returned.

Parameters:

Name Type Description Default
connector ServiceConnectorRequest

The service connector configuration to verify.

required
list_resources bool

If True, the list of all resources accessible through the service connector is returned.

True

Returns:

Type Description
ServiceConnectorResourcesModel

The list of resources that the service connector configuration has

ServiceConnectorResourcesModel

access to.

Source code in src/zenml/zen_server/routers/service_connectors_endpoints.py
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
@router.post(
    SERVICE_CONNECTOR_VERIFY,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def validate_and_verify_service_connector_config(
    connector: ServiceConnectorRequest,
    list_resources: bool = True,
    _: AuthContext = Security(authorize),
) -> ServiceConnectorResourcesModel:
    """Verifies if a service connector configuration has access to resources.

    This requires the service connector implementation to be installed
    on the ZenML server, otherwise a 501 Not Implemented error will be
    returned.

    Args:
        connector: The service connector configuration to verify.
        list_resources: If True, the list of all resources accessible
            through the service connector is returned.

    Returns:
        The list of resources that the service connector configuration has
        access to.
    """
    verify_permission_for_model(model=connector, action=Action.CREATE)

    return zen_store().verify_service_connector_config(
        service_connector=connector,
        list_resources=list_resources,
    )
service_endpoints

Endpoint definitions for services.

Classes Functions
create_service(service: ServiceRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> ServiceResponse

Creates a new service.

Parameters:

Name Type Description Default
service ServiceRequest

The service to create.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
ServiceResponse

The created service.

Source code in src/zenml/zen_server/routers/service_endpoints.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@router.post(
    "",
    responses={401: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + SERVICES,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["services"],
)
@handle_exceptions
def create_service(
    service: ServiceRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ServiceResponse:
    """Creates a new service.

    Args:
        service: The service to create.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created service.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        service.project = project.id

    return verify_permissions_and_create_entity(
        request_model=service,
        create_method=zen_store().create_service,
    )
delete_service(service_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a specific service.

Parameters:

Name Type Description Default
service_id UUID

The ID of the service to delete.

required
Source code in src/zenml/zen_server/routers/service_endpoints.py
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@router.delete(
    "/{service_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_service(
    service_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific service.

    Args:
        service_id: The ID of the service to delete.
    """
    verify_permissions_and_delete_entity(
        id=service_id,
        get_method=zen_store().get_service,
        delete_method=zen_store().delete_service,
    )
get_service(service_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ServiceResponse

Gets a specific service using its unique ID.

Parameters:

Name Type Description Default
service_id UUID

The ID of the service to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ServiceResponse

A specific service object.

Source code in src/zenml/zen_server/routers/service_endpoints.py
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
@router.get(
    "/{service_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_service(
    service_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ServiceResponse:
    """Gets a specific service using its unique ID.

    Args:
        service_id: The ID of the service to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A specific service object.
    """
    return verify_permissions_and_get_entity(
        id=service_id,
        get_method=zen_store().get_service,
        hydrate=hydrate,
    )
list_services(filter_model: ServiceFilter = Depends(make_dependable(ServiceFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ServiceResponse]

Gets a page of service objects.

Parameters:

Name Type Description Default
filter_model ServiceFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ServiceFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ServiceResponse]

Page of service objects.

Source code in src/zenml/zen_server/routers/service_endpoints.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_services(
    filter_model: ServiceFilter = Depends(make_dependable(ServiceFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ServiceResponse]:
    """Gets a page of service objects.

    Args:
        filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page of service objects.
    """
    return verify_permissions_and_list_entities(
        filter_model=filter_model,
        resource_type=ResourceType.SERVICE,
        list_method=zen_store().list_services,
        hydrate=hydrate,
    )
update_service(service_id: UUID, update: ServiceUpdate, _: AuthContext = Security(authorize)) -> ServiceResponse

Updates a service.

Parameters:

Name Type Description Default
service_id UUID

The ID of the service to update.

required
update ServiceUpdate

The model containing the attributes to update.

required

Returns:

Type Description
ServiceResponse

The updated service object.

Source code in src/zenml/zen_server/routers/service_endpoints.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
@router.put(
    "/{service_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_service(
    service_id: UUID,
    update: ServiceUpdate,
    _: AuthContext = Security(authorize),
) -> ServiceResponse:
    """Updates a service.

    Args:
        service_id: The ID of the service to update.
        update: The model containing the attributes to update.

    Returns:
        The updated service object.
    """
    return verify_permissions_and_update_entity(
        id=service_id,
        update_model=update,
        get_method=zen_store().get_service,
        update_method=zen_store().update_service,
    )
stack_components_endpoints

Endpoint definitions for stack components.

Classes Functions
create_stack_component(component: ComponentRequest, project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize)) -> ComponentResponse

Creates a stack component.

Parameters:

Name Type Description Default
component ComponentRequest

Stack component to register.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None

Returns:

Type Description
ComponentResponse

The created stack component.

Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + STACK_COMPONENTS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["stack_components"],
)
@handle_exceptions
def create_stack_component(
    component: ComponentRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ComponentResponse:
    """Creates a stack component.

    Args:
        component: Stack component to register.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created stack component.
    """
    if component.connector:
        service_connector = zen_store().get_service_connector(
            component.connector
        )
        verify_permission_for_model(service_connector, action=Action.READ)

    from zenml.stack.utils import validate_stack_component_config

    validate_stack_component_config(
        configuration_dict=component.configuration,
        flavor=component.flavor,
        component_type=component.type,
        zen_store=zen_store(),
        # We allow custom flavors to fail import on the server side.
        validate_custom_flavors=False,
    )

    return verify_permissions_and_create_entity(
        request_model=component,
        create_method=zen_store().create_stack_component,
    )
deregister_stack_component(component_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a stack component.

Parameters:

Name Type Description Default
component_id UUID

ID of the stack component.

required
Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
@router.delete(
    "/{component_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def deregister_stack_component(
    component_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a stack component.

    Args:
        component_id: ID of the stack component.
    """
    verify_permissions_and_delete_entity(
        id=component_id,
        get_method=zen_store().get_stack_component,
        delete_method=zen_store().delete_stack_component,
    )
get_stack_component(component_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> ComponentResponse

Returns the requested stack component.

Parameters:

Name Type Description Default
component_id UUID

ID of the stack component.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
ComponentResponse

The requested stack component.

Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@router.get(
    "/{component_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_stack_component(
    component_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ComponentResponse:
    """Returns the requested stack component.

    Args:
        component_id: ID of the stack component.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested stack component.
    """
    return verify_permissions_and_get_entity(
        id=component_id,
        get_method=zen_store().get_stack_component,
        hydrate=hydrate,
    )
get_stack_component_types(_: AuthContext = Security(authorize)) -> List[str]

Get a list of all stack component types.

Returns:

Type Description
List[str]

List of stack components.

Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
246
247
248
249
250
251
252
253
254
255
256
257
258
259
@types_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_stack_component_types(
    _: AuthContext = Security(authorize),
) -> List[str]:
    """Get a list of all stack component types.

    Returns:
        List of stack components.
    """
    return StackComponentType.values()
list_stack_components(component_filter_model: ComponentFilter = Depends(make_dependable(ComponentFilter)), project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[ComponentResponse]

Get a list of all stack components.

Parameters:

Name Type Description Default
component_filter_model ComponentFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(ComponentFilter))
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[ComponentResponse]

List of stack components matching the filter criteria.

Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + STACK_COMPONENTS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["stack_components"],
)
@handle_exceptions
def list_stack_components(
    component_filter_model: ComponentFilter = Depends(
        make_dependable(ComponentFilter)
    ),
    project_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ComponentResponse]:
    """Get a list of all stack components.

    Args:
        component_filter_model: Filter model used for pagination, sorting,
            filtering.
        project_name_or_id: Optional name or ID of the project to filter by.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of stack components matching the filter criteria.
    """
    return verify_permissions_and_list_entities(
        filter_model=component_filter_model,
        resource_type=ResourceType.STACK_COMPONENT,
        list_method=zen_store().list_stack_components,
        hydrate=hydrate,
    )
update_stack_component(component_id: UUID, component_update: ComponentUpdate, _: AuthContext = Security(authorize)) -> ComponentResponse

Updates a stack component.

Parameters:

Name Type Description Default
component_id UUID

ID of the stack component.

required
component_update ComponentUpdate

Stack component to use to update.

required

Returns:

Type Description
ComponentResponse

Updated stack component.

Source code in src/zenml/zen_server/routers/stack_components_endpoints.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
@router.put(
    "/{component_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_stack_component(
    component_id: UUID,
    component_update: ComponentUpdate,
    _: AuthContext = Security(authorize),
) -> ComponentResponse:
    """Updates a stack component.

    Args:
        component_id: ID of the stack component.
        component_update: Stack component to use to update.

    Returns:
        Updated stack component.
    """
    if component_update.configuration:
        from zenml.stack.utils import validate_stack_component_config

        existing_component = zen_store().get_stack_component(component_id)
        validate_stack_component_config(
            configuration_dict=component_update.configuration,
            flavor=existing_component.flavor_name,
            component_type=existing_component.type,
            zen_store=zen_store(),
            # We allow custom flavors to fail import on the server side.
            validate_custom_flavors=False,
        )

    if component_update.connector:
        service_connector = zen_store().get_service_connector(
            component_update.connector
        )
        verify_permission_for_model(service_connector, action=Action.READ)

    return verify_permissions_and_update_entity(
        id=component_id,
        update_model=component_update,
        get_method=zen_store().get_stack_component,
        update_method=zen_store().update_stack_component,
    )
stack_deployment_endpoints

Endpoint definitions for stack deployments.

Classes Functions
get_deployed_stack(provider: StackDeploymentProvider, stack_name: str, location: Optional[str] = None, date_start: Optional[datetime.datetime] = None, terraform: bool = False, _: AuthContext = Security(authorize)) -> Optional[DeployedStack]

Return a matching ZenML stack that was deployed and registered.

Parameters:

Name Type Description Default
provider StackDeploymentProvider

The stack deployment provider.

required
stack_name str

The name of the stack.

required
location Optional[str]

The location where the stack should be deployed.

None
date_start Optional[datetime]

The date when the deployment started.

None
terraform bool

Whether the stack was deployed using Terraform.

False

Returns:

Type Description
Optional[DeployedStack]

The ZenML stack that was deployed and registered or None if the stack

Optional[DeployedStack]

was not found.

Source code in src/zenml/zen_server/routers/stack_deployment_endpoints.py
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
@router.get(
    STACK,
)
@handle_exceptions
def get_deployed_stack(
    provider: StackDeploymentProvider,
    stack_name: str,
    location: Optional[str] = None,
    date_start: Optional[datetime.datetime] = None,
    terraform: bool = False,
    _: AuthContext = Security(authorize),
) -> Optional[DeployedStack]:
    """Return a matching ZenML stack that was deployed and registered.

    Args:
        provider: The stack deployment provider.
        stack_name: The name of the stack.
        location: The location where the stack should be deployed.
        date_start: The date when the deployment started.
        terraform: Whether the stack was deployed using Terraform.

    Returns:
        The ZenML stack that was deployed and registered or None if the stack
        was not found.
    """
    stack_deployment_class = get_stack_deployment_class(provider)
    return stack_deployment_class(
        terraform=terraform,
        stack_name=stack_name,
        location=location,
        # These fields are not needed for this operation
        zenml_server_url="",
        zenml_server_api_token="",
    ).get_stack(date_start=date_start)
get_stack_deployment_config(request: Request, provider: StackDeploymentProvider, stack_name: str, location: Optional[str] = None, terraform: bool = False, auth_context: AuthContext = Security(authorize)) -> StackDeploymentConfig

Return the URL to deploy the ZenML stack to the specified cloud provider.

Parameters:

Name Type Description Default
request Request

The FastAPI request object.

required
provider StackDeploymentProvider

The stack deployment provider.

required
stack_name str

The name of the stack.

required
location Optional[str]

The location where the stack should be deployed.

None
terraform bool

Whether the stack should be deployed using Terraform.

False
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
StackDeploymentConfig

The cloud provider console URL where the stack will be deployed and

StackDeploymentConfig

the configuration for the stack deployment.

Source code in src/zenml/zen_server/routers/stack_deployment_endpoints.py
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@router.get(
    CONFIG,
)
@handle_exceptions
def get_stack_deployment_config(
    request: Request,
    provider: StackDeploymentProvider,
    stack_name: str,
    location: Optional[str] = None,
    terraform: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> StackDeploymentConfig:
    """Return the URL to deploy the ZenML stack to the specified cloud provider.

    Args:
        request: The FastAPI request object.
        provider: The stack deployment provider.
        stack_name: The name of the stack.
        location: The location where the stack should be deployed.
        terraform: Whether the stack should be deployed using Terraform.
        auth_context: The authentication context.

    Returns:
        The cloud provider console URL where the stack will be deployed and
        the configuration for the stack deployment.
    """
    verify_permission(
        resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
    )
    verify_permission(
        resource_type=ResourceType.STACK_COMPONENT,
        action=Action.CREATE,
    )
    verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE)

    stack_deployment_class = get_stack_deployment_class(provider)

    config = server_config()
    if config.server_url:
        url = config.server_url
    else:
        # Get the base server URL used to call this FastAPI endpoint
        url = str(
            request.url.replace(path="")
            .replace(query="")
            .replace(scheme="https")
        )

    token = auth_context.access_token
    assert token is not None

    # A new API token is generated for the stack deployment
    api_token = generate_access_token(
        user_id=token.user_id,
        expires_in=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION * 60,
    ).access_token

    return stack_deployment_class(
        terraform=terraform,
        stack_name=stack_name,
        location=location,
        zenml_server_url=str(url),
        zenml_server_api_token=api_token,
    ).get_deployment_config()
get_stack_deployment_info(provider: StackDeploymentProvider, _: AuthContext = Security(authorize)) -> StackDeploymentInfo

Get information about a stack deployment provider.

Parameters:

Name Type Description Default
provider StackDeploymentProvider

The stack deployment provider.

required

Returns:

Type Description
StackDeploymentInfo

Information about the stack deployment provider.

Source code in src/zenml/zen_server/routers/stack_deployment_endpoints.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@router.get(
    INFO,
)
@handle_exceptions
def get_stack_deployment_info(
    provider: StackDeploymentProvider,
    _: AuthContext = Security(authorize),
) -> StackDeploymentInfo:
    """Get information about a stack deployment provider.

    Args:
        provider: The stack deployment provider.

    Returns:
        Information about the stack deployment provider.
    """
    stack_deployment_class = get_stack_deployment_class(provider)
    return stack_deployment_class.get_deployment_info()
stacks_endpoints

Endpoint definitions for stacks.

Classes Functions
create_stack(stack: StackRequest, project_name_or_id: Optional[Union[str, UUID]] = None, auth_context: AuthContext = Security(authorize)) -> StackResponse

Creates a stack.

Parameters:

Name Type Description Default
stack StackRequest

Stack to register.

required
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project.

None
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
StackResponse

The created stack.

Source code in src/zenml/zen_server/routers/stacks_endpoints.py
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + STACKS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["stacks"],
)
@handle_exceptions
def create_stack(
    stack: StackRequest,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    auth_context: AuthContext = Security(authorize),
) -> StackResponse:
    """Creates a stack.

    Args:
        stack: Stack to register.
        project_name_or_id: Optional name or ID of the project.
        auth_context: Authentication context.

    Returns:
        The created stack.
    """
    # Check the service connector creation
    is_connector_create_needed = False
    for connector_id_or_info in stack.service_connectors:
        if isinstance(connector_id_or_info, UUID):
            service_connector = zen_store().get_service_connector(
                connector_id_or_info, hydrate=False
            )
            verify_permission_for_model(
                model=service_connector, action=Action.READ
            )
        else:
            is_connector_create_needed = True

    # Check the component creation
    if is_connector_create_needed:
        verify_permission(
            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
        )
    is_component_create_needed = False
    for components in stack.components.values():
        for component_id_or_info in components:
            if isinstance(component_id_or_info, UUID):
                component = zen_store().get_stack_component(
                    component_id_or_info, hydrate=False
                )
                verify_permission_for_model(
                    model=component, action=Action.READ
                )
            else:
                is_component_create_needed = True
    if is_component_create_needed:
        verify_permission(
            resource_type=ResourceType.STACK_COMPONENT,
            action=Action.CREATE,
        )

    # Check the stack creation
    verify_permission_for_model(model=stack, action=Action.CREATE)

    return zen_store().create_stack(stack)
delete_stack(stack_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a stack.

Parameters:

Name Type Description Default
stack_id UUID

Name of the stack.

required
Source code in src/zenml/zen_server/routers/stacks_endpoints.py
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@router.delete(
    "/{stack_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_stack(
    stack_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a stack.

    Args:
        stack_id: Name of the stack.
    """
    verify_permissions_and_delete_entity(
        id=stack_id,
        get_method=zen_store().get_stack,
        delete_method=zen_store().delete_stack,
    )
get_stack(stack_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> StackResponse

Returns the requested stack.

Parameters:

Name Type Description Default
stack_id UUID

ID of the stack.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
StackResponse

The requested stack.

Source code in src/zenml/zen_server/routers/stacks_endpoints.py
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
@router.get(
    "/{stack_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_stack(
    stack_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> StackResponse:
    """Returns the requested stack.

    Args:
        stack_id: ID of the stack.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested stack.
    """
    return verify_permissions_and_get_entity(
        id=stack_id, get_method=zen_store().get_stack, hydrate=hydrate
    )
list_stacks(project_name_or_id: Optional[Union[str, UUID]] = None, stack_filter_model: StackFilter = Depends(make_dependable(StackFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[StackResponse]

Returns all stacks.

Parameters:

Name Type Description Default
project_name_or_id Optional[Union[str, UUID]]

Optional name or ID of the project to filter by.

None
stack_filter_model StackFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(StackFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[StackResponse]

All stacks matching the filter criteria.

Source code in src/zenml/zen_server/routers/stacks_endpoints.py
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.get(
    "/{project_name_or_id}" + STACKS,
    responses={401: error_response, 404: error_response, 422: error_response},
    deprecated=True,
    tags=["stacks"],
)
@handle_exceptions
def list_stacks(
    project_name_or_id: Optional[Union[str, UUID]] = None,
    stack_filter_model: StackFilter = Depends(make_dependable(StackFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[StackResponse]:
    """Returns all stacks.

    Args:
        project_name_or_id: Optional name or ID of the project to filter by.
        stack_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        All stacks matching the filter criteria.
    """
    return verify_permissions_and_list_entities(
        filter_model=stack_filter_model,
        resource_type=ResourceType.STACK,
        list_method=zen_store().list_stacks,
        hydrate=hydrate,
    )
update_stack(stack_id: UUID, stack_update: StackUpdate, _: AuthContext = Security(authorize)) -> StackResponse

Updates a stack.

Parameters:

Name Type Description Default
stack_id UUID

Name of the stack.

required
stack_update StackUpdate

Stack to use for the update.

required

Returns:

Type Description
StackResponse

The updated stack.

Source code in src/zenml/zen_server/routers/stacks_endpoints.py
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@router.put(
    "/{stack_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_stack(
    stack_id: UUID,
    stack_update: StackUpdate,
    _: AuthContext = Security(authorize),
) -> StackResponse:
    """Updates a stack.

    Args:
        stack_id: Name of the stack.
        stack_update: Stack to use for the update.

    Returns:
        The updated stack.
    """
    if stack_update.components:
        updated_components = [
            zen_store().get_stack_component(id)
            for ids in stack_update.components.values()
            for id in ids
        ]

        batch_verify_permissions_for_models(
            updated_components, action=Action.READ
        )

    return verify_permissions_and_update_entity(
        id=stack_id,
        update_model=stack_update,
        get_method=zen_store().get_stack,
        update_method=zen_store().update_stack,
    )
steps_endpoints

Endpoint definitions for steps (and artifacts) of pipeline runs.

Classes Functions
create_run_step(step: StepRunRequest, _: AuthContext = Security(authorize)) -> StepRunResponse

Create a run step.

Parameters:

Name Type Description Default
step StepRunRequest

The run step to create.

required
_ AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
StepRunResponse

The created run step.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_run_step(
    step: StepRunRequest,
    _: AuthContext = Security(authorize),
) -> StepRunResponse:
    """Create a run step.

    Args:
        step: The run step to create.
        _: Authentication context.

    Returns:
        The created run step.
    """
    pipeline_run = zen_store().get_run(step.pipeline_run_id)

    return verify_permissions_and_create_entity(
        request_model=step,
        create_method=zen_store().create_run_step,
        surrogate_models=[pipeline_run],
    )
get_step(step_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> StepRunResponse

Get one specific step.

Parameters:

Name Type Description Default
step_id UUID

ID of the step to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
StepRunResponse

The step.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
@router.get(
    "/{step_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_step(
    step_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> StepRunResponse:
    """Get one specific step.

    Args:
        step_id: ID of the step to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The step.
    """
    # We always fetch the step hydrated because we need the pipeline_run_id
    # for the permission checks. If the user requested an unhydrated response,
    # we later remove the metadata
    step = zen_store().get_run_step(step_id, hydrate=True)
    pipeline_run = zen_store().get_run(step.pipeline_run_id)
    verify_permission_for_model(pipeline_run, action=Action.READ)

    if hydrate is False:
        step.metadata = None

    return dehydrate_response_model(step)
get_step_configuration(step_id: UUID, _: AuthContext = Security(authorize)) -> Dict[str, Any]

Get the configuration of a specific step.

Parameters:

Name Type Description Default
step_id UUID

ID of the step to get.

required

Returns:

Type Description
Dict[str, Any]

The step configuration.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
@router.get(
    "/{step_id}" + STEP_CONFIGURATION,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_step_configuration(
    step_id: UUID,
    _: AuthContext = Security(authorize),
) -> Dict[str, Any]:
    """Get the configuration of a specific step.

    Args:
        step_id: ID of the step to get.

    Returns:
        The step configuration.
    """
    step = zen_store().get_run_step(step_id, hydrate=True)
    pipeline_run = zen_store().get_run(step.pipeline_run_id)
    verify_permission_for_model(pipeline_run, action=Action.READ)

    return step.config.model_dump()
get_step_logs(step_id: UUID, offset: int = 0, length: int = 1024 * 1024 * 16, _: AuthContext = Security(authorize)) -> str

Get the logs of a specific step.

Parameters:

Name Type Description Default
step_id UUID

ID of the step for which to get the logs.

required
offset int

The offset from which to start reading.

0
length int

The amount of bytes that should be read.

1024 * 1024 * 16

Returns:

Type Description
str

The logs of the step.

Raises:

Type Description
HTTPException

If no logs are available for this step.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
@router.get(
    "/{step_id}" + LOGS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_step_logs(
    step_id: UUID,
    offset: int = 0,
    length: int = 1024 * 1024 * 16,  # Default to 16MiB of data
    _: AuthContext = Security(authorize),
) -> str:
    """Get the logs of a specific step.

    Args:
        step_id: ID of the step for which to get the logs.
        offset: The offset from which to start reading.
        length: The amount of bytes that should be read.

    Returns:
        The logs of the step.

    Raises:
        HTTPException: If no logs are available for this step.
    """
    step = zen_store().get_run_step(step_id, hydrate=True)
    pipeline_run = zen_store().get_run(step.pipeline_run_id)
    verify_permission_for_model(pipeline_run, action=Action.READ)

    store = zen_store()
    logs = step.logs
    if logs is None:
        raise HTTPException(
            status_code=404, detail="No logs available for this step"
        )
    return fetch_logs(
        zen_store=store,
        artifact_store_id=logs.artifact_store_id,
        logs_uri=logs.uri,
        offset=offset,
        length=length,
    )
get_step_status(step_id: UUID, _: AuthContext = Security(authorize)) -> ExecutionStatus

Get the status of a specific step.

Parameters:

Name Type Description Default
step_id UUID

ID of the step for which to get the status.

required

Returns:

Type Description
ExecutionStatus

The status of the step.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
@router.get(
    "/{step_id}" + STATUS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_step_status(
    step_id: UUID,
    _: AuthContext = Security(authorize),
) -> ExecutionStatus:
    """Get the status of a specific step.

    Args:
        step_id: ID of the step for which to get the status.

    Returns:
        The status of the step.
    """
    step = zen_store().get_run_step(step_id, hydrate=True)
    pipeline_run = zen_store().get_run(step.pipeline_run_id)
    verify_permission_for_model(pipeline_run, action=Action.READ)

    return step.status
list_run_steps(step_run_filter_model: StepRunFilter = Depends(make_dependable(StepRunFilter)), hydrate: bool = False, auth_context: AuthContext = Security(authorize)) -> Page[StepRunResponse]

Get run steps according to query filters.

Parameters:

Name Type Description Default
step_run_filter_model StepRunFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(StepRunFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
Page[StepRunResponse]

The run steps according to query filters.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_run_steps(
    step_run_filter_model: StepRunFilter = Depends(
        make_dependable(StepRunFilter)
    ),
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[StepRunResponse]:
    """Get run steps according to query filters.

    Args:
        step_run_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: Authentication context.

    Returns:
        The run steps according to query filters.
    """
    # A project scoped request must always be scoped to a specific
    # project. This is required for the RBAC check to work.
    set_filter_project_scope(step_run_filter_model)
    assert isinstance(step_run_filter_model.project, UUID)

    allowed_pipeline_run_ids = get_allowed_resource_ids(
        resource_type=ResourceType.PIPELINE_RUN,
        project_id=step_run_filter_model.project,
    )
    step_run_filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id,
        pipeline_run_id=allowed_pipeline_run_ids,
    )

    page = zen_store().list_run_steps(
        step_run_filter_model=step_run_filter_model, hydrate=hydrate
    )
    return dehydrate_page(page)
update_step(step_id: UUID, step_model: StepRunUpdate, _: AuthContext = Security(authorize)) -> StepRunResponse

Updates a step.

Parameters:

Name Type Description Default
step_id UUID

ID of the step.

required
step_model StepRunUpdate

Step model to use for the update.

required

Returns:

Type Description
StepRunResponse

The updated step model.

Source code in src/zenml/zen_server/routers/steps_endpoints.py
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
@router.put(
    "/{step_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_step(
    step_id: UUID,
    step_model: StepRunUpdate,
    _: AuthContext = Security(authorize),
) -> StepRunResponse:
    """Updates a step.

    Args:
        step_id: ID of the step.
        step_model: Step model to use for the update.

    Returns:
        The updated step model.
    """
    step = zen_store().get_run_step(step_id, hydrate=True)
    pipeline_run = zen_store().get_run(step.pipeline_run_id)
    verify_permission_for_model(pipeline_run, action=Action.UPDATE)

    updated_step = zen_store().update_run_step(
        step_run_id=step_id, step_run_update=step_model
    )
    return dehydrate_response_model(updated_step)
tag_resource_endpoints

Endpoint definitions for the link between tags and resources.

Classes Functions
batch_create_tag_resource(tag_resources: List[TagResourceRequest], _: AuthContext = Security(authorize)) -> List[TagResourceResponse]

Attach different tags to different resources.

Parameters:

Name Type Description Default
tag_resources List[TagResourceRequest]

A list of tag resource requests.

required

Returns:

Type Description
List[TagResourceResponse]

A list of tag resource responses.

Source code in src/zenml/zen_server/routers/tag_resource_endpoints.py
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
@router.post(
    BATCH,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def batch_create_tag_resource(
    tag_resources: List[TagResourceRequest],
    _: AuthContext = Security(authorize),
) -> List[TagResourceResponse]:
    """Attach different tags to different resources.

    Args:
        tag_resources: A list of tag resource requests.

    Returns:
        A list of tag resource responses.
    """
    return [
        zen_store().create_tag_resource(tag_resource=tag_resource)
        for tag_resource in tag_resources
    ]
batch_delete_tag_resource(tag_resources: List[TagResourceRequest], _: AuthContext = Security(authorize)) -> None

Detach different tags from different resources.

Parameters:

Name Type Description Default
tag_resources List[TagResourceRequest]

A list of tag resource requests.

required
Source code in src/zenml/zen_server/routers/tag_resource_endpoints.py
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@router.delete(
    BATCH,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def batch_delete_tag_resource(
    tag_resources: List[TagResourceRequest],
    _: AuthContext = Security(authorize),
) -> None:
    """Detach different tags from different resources.

    Args:
        tag_resources: A list of tag resource requests.
    """
    zen_store().batch_delete_tag_resource(tag_resources=tag_resources)
create_tag_resource(tag_resource: TagResourceRequest, _: AuthContext = Security(authorize)) -> TagResourceResponse

Attach different tags to different resources.

Parameters:

Name Type Description Default
tag_resource TagResourceRequest

A tag resource request.

required

Returns:

Type Description
TagResourceResponse

A tag resource response.

Source code in src/zenml/zen_server/routers/tag_resource_endpoints.py
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@router.post(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def create_tag_resource(
    tag_resource: TagResourceRequest,
    _: AuthContext = Security(authorize),
) -> TagResourceResponse:
    """Attach different tags to different resources.

    Args:
        tag_resource: A tag resource request.

    Returns:
        A tag resource response.
    """
    return zen_store().create_tag_resource(tag_resource=tag_resource)
delete_tag_resource(tag_resource: TagResourceRequest, _: AuthContext = Security(authorize)) -> None

Detach a tag from a resource.

Parameters:

Name Type Description Default
tag_resource TagResourceRequest

The tag resource relationship to delete.

required
Source code in src/zenml/zen_server/routers/tag_resource_endpoints.py
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@router.delete(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_tag_resource(
    tag_resource: TagResourceRequest,
    _: AuthContext = Security(authorize),
) -> None:
    """Detach a tag from a resource.

    Args:
        tag_resource: The tag resource relationship to delete.
    """
    zen_store().delete_tag_resource(tag_resource=tag_resource)
tags_endpoints

Endpoint definitions for tags.

Classes Functions
create_tag(tag: TagRequest, _: AuthContext = Security(authorize)) -> TagResponse

Create a new tag.

Parameters:

Name Type Description Default
tag TagRequest

The tag to create.

required

Returns:

Type Description
TagResponse

The created tag.

Source code in src/zenml/zen_server/routers/tags_endpoints.py
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_tag(
    tag: TagRequest,
    _: AuthContext = Security(authorize),
) -> TagResponse:
    """Create a new tag.

    Args:
        tag: The tag to create.

    Returns:
        The created tag.
    """
    return zen_store().create_tag(tag=tag)
delete_tag(tag_name_or_id: Union[str, UUID], _: AuthContext = Security(authorize)) -> None

Delete a tag by name or ID.

Parameters:

Name Type Description Default
tag_name_or_id Union[str, UUID]

The name or ID of the tag to delete.

required
Source code in src/zenml/zen_server/routers/tags_endpoints.py
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@router.delete(
    "/{tag_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_tag(
    tag_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a tag by name or ID.

    Args:
        tag_name_or_id: The name or ID of the tag to delete.
    """
    verify_permissions_and_delete_entity(
        id=tag_name_or_id,
        get_method=zen_store().get_tag,
        delete_method=zen_store().delete_tag,
    )
get_tag(tag_name_or_id: Union[str, UUID], hydrate: bool = True, _: AuthContext = Security(authorize)) -> TagResponse

Get a tag by name or ID.

Parameters:

Name Type Description Default
tag_name_or_id Union[str, UUID]

The name or ID of the tag to get.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
TagResponse

The tag with the given name or ID.

Source code in src/zenml/zen_server/routers/tags_endpoints.py
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
@router.get(
    "/{tag_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_tag(
    tag_name_or_id: Union[str, UUID],
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> TagResponse:
    """Get a tag by name or ID.

    Args:
        tag_name_or_id: The name or ID of the tag to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The tag with the given name or ID.
    """
    return verify_permissions_and_get_entity(
        id=tag_name_or_id,
        get_method=zen_store().get_tag,
        hydrate=hydrate,
    )
list_tags(tag_filter_model: TagFilter = Depends(make_dependable(TagFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[TagResponse]

Get tags according to query filters.

Parameters:

Name Type Description Default
tag_filter_model TagFilter

Filter model used for pagination, sorting, filtering

Depends(make_dependable(TagFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[TagResponse]

The tags according to query filters.

Source code in src/zenml/zen_server/routers/tags_endpoints.py
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_tags(
    tag_filter_model: TagFilter = Depends(make_dependable(TagFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[TagResponse]:
    """Get tags according to query filters.

    Args:
        tag_filter_model: Filter model used for pagination, sorting,
            filtering
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The tags according to query filters.
    """
    return zen_store().list_tags(
        tag_filter_model=tag_filter_model,
        hydrate=hydrate,
    )
update_tag(tag_id: UUID, tag_update_model: TagUpdate, _: AuthContext = Security(authorize)) -> TagResponse

Updates a tag.

Parameters:

Name Type Description Default
tag_id UUID

Id or name of the tag.

required
tag_update_model TagUpdate

Tag to use for the update.

required

Returns:

Type Description
TagResponse

The updated tag.

Source code in src/zenml/zen_server/routers/tags_endpoints.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@router.put(
    "/{tag_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_tag(
    tag_id: UUID,
    tag_update_model: TagUpdate,
    _: AuthContext = Security(authorize),
) -> TagResponse:
    """Updates a tag.

    Args:
        tag_id: Id or name of the tag.
        tag_update_model: Tag to use for the update.

    Returns:
        The updated tag.
    """
    return verify_permissions_and_update_entity(
        id=tag_id,
        update_model=tag_update_model,
        get_method=zen_store().get_tag,
        update_method=zen_store().update_tag,
    )
triggers_endpoints

Endpoint definitions for triggers.

Classes Functions
create_trigger(trigger: TriggerRequest, _: AuthContext = Security(authorize)) -> TriggerResponse

Creates a trigger.

Parameters:

Name Type Description Default
trigger TriggerRequest

Trigger to register.

required

Returns:

Type Description
TriggerResponse

The created trigger.

Raises:

Type Description
ValueError

If the action flavor/subtype combination is not actually a webhook event source

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@handle_exceptions
def create_trigger(
    trigger: TriggerRequest,
    _: AuthContext = Security(authorize),
) -> TriggerResponse:
    """Creates a trigger.

    Args:
        trigger: Trigger to register.

    Returns:
        The created trigger.

    Raises:
        ValueError: If the action flavor/subtype combination is not actually a webhook event source
    """
    if trigger.event_source_id and trigger.event_filter:
        event_source = zen_store().get_event_source(
            event_source_id=trigger.event_source_id
        )

        event_source_handler = plugin_flavor_registry().get_plugin(
            name=event_source.flavor,
            _type=PluginType.EVENT_SOURCE,
            subtype=event_source.plugin_subtype,
        )

        # Validate that the flavor and plugin_type correspond to an event source
        # implementation
        if not isinstance(event_source_handler, BaseEventSourceHandler):
            raise ValueError(
                f"Event source plugin {event_source.plugin_subtype} "
                f"for flavor {event_source.flavor} is not a valid event source "
                "handler implementation."
            )

        # Validate the trigger event filter
        event_source_handler.validate_event_filter_configuration(
            trigger.event_filter
        )

    return verify_permissions_and_create_entity(
        request_model=trigger,
        create_method=zen_store().create_trigger,
    )
delete_trigger(trigger_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a trigger.

Parameters:

Name Type Description Default
trigger_id UUID

Name of the trigger.

required
Source code in src/zenml/zen_server/routers/triggers_endpoints.py
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@router.delete(
    "/{trigger_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_trigger(
    trigger_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a trigger.

    Args:
        trigger_id: Name of the trigger.
    """
    verify_permissions_and_delete_entity(
        id=trigger_id,
        get_method=zen_store().get_trigger,
        delete_method=zen_store().delete_trigger,
    )
delete_trigger_execution(trigger_execution_id: UUID, _: AuthContext = Security(authorize)) -> None

Deletes a trigger execution.

Parameters:

Name Type Description Default
trigger_execution_id UUID

ID of the trigger execution.

required
Source code in src/zenml/zen_server/routers/triggers_endpoints.py
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
@executions_router.delete(
    "/{trigger_execution_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def delete_trigger_execution(
    trigger_execution_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a trigger execution.

    Args:
        trigger_execution_id: ID of the trigger execution.
    """
    verify_permissions_and_delete_entity(
        id=trigger_execution_id,
        get_method=zen_store().get_trigger_execution,
        delete_method=zen_store().delete_trigger_execution,
    )
get_trigger(trigger_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> TriggerResponse

Returns the requested trigger.

Parameters:

Name Type Description Default
trigger_id UUID

ID of the trigger.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
TriggerResponse

The requested trigger.

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
@router.get(
    "/{trigger_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_trigger(
    trigger_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> TriggerResponse:
    """Returns the requested trigger.

    Args:
        trigger_id: ID of the trigger.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested trigger.
    """
    trigger = zen_store().get_trigger(trigger_id=trigger_id, hydrate=hydrate)
    verify_permission_for_model(trigger, action=Action.READ)
    return dehydrate_response_model(trigger)
get_trigger_execution(trigger_execution_id: UUID, hydrate: bool = True, _: AuthContext = Security(authorize)) -> TriggerExecutionResponse

Returns the requested trigger execution.

Parameters:

Name Type Description Default
trigger_execution_id UUID

ID of the trigger execution.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True

Returns:

Type Description
TriggerExecutionResponse

The requested trigger execution.

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
@executions_router.get(
    "/{trigger_execution_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_trigger_execution(
    trigger_execution_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> TriggerExecutionResponse:
    """Returns the requested trigger execution.

    Args:
        trigger_execution_id: ID of the trigger execution.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested trigger execution.
    """
    return verify_permissions_and_get_entity(
        id=trigger_execution_id,
        get_method=zen_store().get_trigger_execution,
        hydrate=hydrate,
    )
list_trigger_executions(trigger_execution_filter_model: TriggerExecutionFilter = Depends(make_dependable(TriggerExecutionFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[TriggerExecutionResponse]

List trigger executions.

Parameters:

Name Type Description Default
trigger_execution_filter_model TriggerExecutionFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(TriggerExecutionFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[TriggerExecutionResponse]

Page of trigger executions.

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
@executions_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_trigger_executions(
    trigger_execution_filter_model: TriggerExecutionFilter = Depends(
        make_dependable(TriggerExecutionFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[TriggerExecutionResponse]:
    """List trigger executions.

    Args:
        trigger_execution_filter_model: Filter model used for pagination,
            sorting, filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        Page of trigger executions.
    """
    return verify_permissions_and_list_entities(
        filter_model=trigger_execution_filter_model,
        resource_type=ResourceType.TRIGGER_EXECUTION,
        list_method=zen_store().list_trigger_executions,
        hydrate=hydrate,
    )
list_triggers(trigger_filter_model: TriggerFilter = Depends(make_dependable(TriggerFilter)), hydrate: bool = False, _: AuthContext = Security(authorize)) -> Page[TriggerResponse]

Returns all triggers.

Parameters:

Name Type Description Default
trigger_filter_model TriggerFilter

Filter model used for pagination, sorting, filtering.

Depends(make_dependable(TriggerFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False

Returns:

Type Description
Page[TriggerResponse]

All triggers.

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_triggers(
    trigger_filter_model: TriggerFilter = Depends(
        make_dependable(TriggerFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[TriggerResponse]:
    """Returns all triggers.

    Args:
        trigger_filter_model: Filter model used for pagination, sorting,
            filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        All triggers.
    """
    return verify_permissions_and_list_entities(
        filter_model=trigger_filter_model,
        resource_type=ResourceType.TRIGGER,
        list_method=zen_store().list_triggers,
        hydrate=hydrate,
    )
update_trigger(trigger_id: UUID, trigger_update: TriggerUpdate, _: AuthContext = Security(authorize)) -> TriggerResponse

Updates a trigger.

Parameters:

Name Type Description Default
trigger_id UUID

Name of the trigger.

required
trigger_update TriggerUpdate

Trigger to use for the update.

required

Returns:

Type Description
TriggerResponse

The updated trigger.

Raises:

Type Description
ValueError

If the action flavor/subtype combination is not actually a webhook event source

Source code in src/zenml/zen_server/routers/triggers_endpoints.py
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@router.put(
    "/{trigger_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def update_trigger(
    trigger_id: UUID,
    trigger_update: TriggerUpdate,
    _: AuthContext = Security(authorize),
) -> TriggerResponse:
    """Updates a trigger.

    Args:
        trigger_id: Name of the trigger.
        trigger_update: Trigger to use for the update.

    Returns:
        The updated trigger.

    Raises:
        ValueError: If the action flavor/subtype combination is not actually a webhook event source
    """
    trigger = zen_store().get_trigger(trigger_id=trigger_id)

    if trigger_update.event_filter:
        if not trigger.event_source:
            raise ValueError(
                "Trying to set event filter for trigger without event source."
            )

        event_source = zen_store().get_event_source(
            event_source_id=trigger.event_source.id
        )

        event_source_handler = plugin_flavor_registry().get_plugin(
            name=event_source.flavor,
            _type=PluginType.EVENT_SOURCE,
            subtype=event_source.plugin_subtype,
        )

        # Validate that the flavor and plugin_type correspond to an event source
        # implementation
        if not isinstance(event_source_handler, BaseEventSourceHandler):
            raise ValueError(
                f"Event source plugin {event_source.plugin_subtype} "
                f"for flavor {event_source.flavor} is not a valid event source "
                "handler implementation."
            )

        # Validate the trigger event filter
        event_source_handler.validate_event_filter_configuration(
            trigger_update.event_filter
        )

    verify_permission_for_model(trigger, action=Action.UPDATE)

    updated_trigger = zen_store().update_trigger(
        trigger_id=trigger_id, trigger_update=trigger_update
    )

    return dehydrate_response_model(updated_trigger)
users_endpoints

Endpoint definitions for users.

Classes Functions
activate_user(user_name_or_id: Union[str, UUID], user_update: UserUpdate) -> UserResponse

Activates a specific user.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
user_update UserUpdate

the user to use for the update.

required

Returns:

Type Description
UserResponse

The updated user.

Source code in src/zenml/zen_server/routers/users_endpoints.py
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
@activation_router.put(
    "/{user_name_or_id}" + ACTIVATE,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def activate_user(
    user_name_or_id: Union[str, UUID],
    user_update: UserUpdate,
) -> UserResponse:
    """Activates a specific user.

    Args:
        user_name_or_id: Name or ID of the user.
        user_update: the user to use for the update.

    Returns:
        The updated user.
    """
    user = zen_store().get_user(user_name_or_id)

    # Use a separate object to compute the update that will be applied to
    # the user to avoid giving the API requester direct control over the
    # user attributes that are updated.
    #
    # Exclude attributes that cannot be updated through this endpoint:
    #
    # - activation_token
    # - external_user_id
    # - is_admin
    # - active
    # - old_password
    #
    safe_user_update = user_update.create_copy(
        exclude={
            "activation_token",
            "external_user_id",
            "is_admin",
            "active",
            "old_password",
        },
    )

    # NOTE: if the activation token is not set, this will raise an
    # exception
    authenticate_credentials(
        user_name_or_id=user_name_or_id,
        activation_token=user_update.activation_token,
    )

    # Activate the user: set active to True and clear the activation token
    safe_user_update.active = True
    safe_user_update.activation_token = None
    return zen_store().update_user(
        user_id=user.id, user_update=safe_user_update
    )
create_user(user: UserRequest, auth_context: AuthContext = Security(authorize)) -> UserResponse

Creates a user.

noqa: DAR401

Parameters:

Name Type Description Default
user UserRequest

User to create.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

The created user.

Source code in src/zenml/zen_server/routers/users_endpoints.py
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
@router.post(
    "",
    responses={
        401: error_response,
        409: error_response,
        422: error_response,
    },
)
@handle_exceptions
def create_user(
    user: UserRequest,
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Creates a user.

    # noqa: DAR401

    Args:
        user: User to create.
        auth_context: Authentication context.

    Returns:
        The created user.
    """
    # Two ways of creating a new user:
    # 1. Create a new user with a password and have it immediately active
    # 2. Create a new user without a password and have it activated at a
    # later time with an activation token

    token: Optional[str] = None
    if user.password is None:
        user.active = False
        token = user.generate_activation_token()
    else:
        user.active = True

    verify_admin_status_if_no_rbac(
        auth_context.user.is_admin, "create user"
    )

    # new_user = verify_permissions_and_create_entity(
    #     request_model=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
    if token:
        new_user.get_body().activation_token = token
    return new_user
deactivate_user(user_name_or_id: Union[str, UUID], auth_context: AuthContext = Security(authorize)) -> UserResponse

Deactivates a user and generates a new activation token for it.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

The generated activation token.

Raises:

Type Description
IllegalOperationError

if the user is trying to deactivate themselves.

Source code in src/zenml/zen_server/routers/users_endpoints.py
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
@router.put(
    "/{user_name_or_id}" + DEACTIVATE,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def deactivate_user(
    user_name_or_id: Union[str, UUID],
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Deactivates a user and generates a new activation token for it.

    Args:
        user_name_or_id: Name or ID of the user.
        auth_context: Authentication context.

    Returns:
        The generated activation token.

    Raises:
        IllegalOperationError: if the user is trying to deactivate
            themselves.
    """
    user = zen_store().get_user(user_name_or_id)
    if user.id == auth_context.user.id:
        raise IllegalOperationError("Cannot deactivate yourself.")
    verify_admin_status_if_no_rbac(
        auth_context.user.is_admin, "deactivate user"
    )
    # verify_permission_for_model(
    #     user,
    #     action=Action.UPDATE,
    # )

    user_update = UserUpdate(
        active=False,
    )
    token = user_update.generate_activation_token()
    user = zen_store().update_user(
        user_id=user.id, user_update=user_update
    )
    # add back the original unhashed activation token
    user.get_body().activation_token = token
    return dehydrate_response_model(user)
delete_user(user_name_or_id: Union[str, UUID], auth_context: AuthContext = Security(authorize)) -> None

Deletes a specific user.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
auth_context AuthContext

The authentication context.

Security(authorize)

Raises:

Type Description
IllegalOperationError

If the user is not authorized to delete the user.

Source code in src/zenml/zen_server/routers/users_endpoints.py
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
@router.delete(
    "/{user_name_or_id}",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def delete_user(
    user_name_or_id: Union[str, UUID],
    auth_context: AuthContext = Security(authorize),
) -> None:
    """Deletes a specific user.

    Args:
        user_name_or_id: Name or ID of the user.
        auth_context: The authentication context.

    Raises:
        IllegalOperationError: If the user is not authorized to delete the user.
    """
    user = zen_store().get_user(user_name_or_id)

    if auth_context.user.id == user.id:
        raise IllegalOperationError(
            "You cannot delete the user account currently used to authenticate "
            "to the ZenML server. If you wish to delete this account, "
            "please authenticate with another account or contact your ZenML "
            "administrator."
        )
    else:
        verify_admin_status_if_no_rbac(
            auth_context.user.is_admin, "delete user"
        )
        # verify_permission_for_model(
        #     user,
        #     action=Action.DELETE,
        # )

    zen_store().delete_user(user_name_or_id=user_name_or_id)
email_opt_in_response(user_name_or_id: Union[str, UUID], user_response: UserUpdate, auth_context: AuthContext = Security(authorize)) -> UserResponse

Sets the response of the user to the email prompt.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
user_response UserUpdate

User Response to email prompt

required
auth_context AuthContext

The authentication context of the user

Security(authorize)

Returns:

Type Description
UserResponse

The updated user.

Raises:

Type Description
AuthorizationException

if the user does not have the required permissions

Source code in src/zenml/zen_server/routers/users_endpoints.py
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
@router.put(
    "/{user_name_or_id}" + EMAIL_ANALYTICS,
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def email_opt_in_response(
    user_name_or_id: Union[str, UUID],
    user_response: UserUpdate,
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Sets the response of the user to the email prompt.

    Args:
        user_name_or_id: Name or ID of the user.
        user_response: User Response to email prompt
        auth_context: The authentication context of the user

    Returns:
        The updated user.

    Raises:
        AuthorizationException: if the user does not have the required
            permissions
    """
    user = zen_store().get_user(user_name_or_id)

    if str(auth_context.user.id) == str(user_name_or_id):
        user_update = UserUpdate(
            email=user_response.email,
            email_opted_in=user_response.email_opted_in,
        )

        if user_response.email_opted_in is not None:
            email_opt_int(
                opted_in=user_response.email_opted_in,
                email=user_response.email,
                source="zenml server",
            )
        updated_user = zen_store().update_user(
            user_id=user.id, user_update=user_update
        )
        return dehydrate_response_model(updated_user)
    else:
        raise AuthorizationException(
            "Users can not opt in on behalf of another user."
        )
get_current_user(auth_context: AuthContext = Security(authorize)) -> UserResponse

Returns the model of the authenticated user.

Parameters:

Name Type Description Default
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

The model of the authenticated user.

Source code in src/zenml/zen_server/routers/users_endpoints.py
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
@current_user_router.get(
    "/current-user",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_current_user(
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Returns the model of the authenticated user.

    Args:
        auth_context: The authentication context.

    Returns:
        The model of the authenticated user.
    """
    return dehydrate_response_model(auth_context.user)
get_user(user_name_or_id: Union[str, UUID], hydrate: bool = True, auth_context: AuthContext = Security(authorize)) -> UserResponse

Returns a specific user.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

True
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

A specific user.

Source code in src/zenml/zen_server/routers/users_endpoints.py
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
@router.get(
    "/{user_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def get_user(
    user_name_or_id: Union[str, UUID],
    hydrate: bool = True,
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Returns a specific user.

    Args:
        user_name_or_id: Name or ID of the user.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: Authentication context.

    Returns:
        A specific user.
    """
    user = zen_store().get_user(
        user_name_or_id=user_name_or_id, hydrate=hydrate
    )
    if user.id != auth_context.user.id:
        verify_admin_status_if_no_rbac(
            auth_context.user.is_admin, "get other user"
        )
        # verify_permission_for_model(
        #     user,
        #     action=Action.READ,
        # )

    return dehydrate_response_model(user)
list_users(user_filter_model: UserFilter = Depends(make_dependable(UserFilter)), hydrate: bool = False, auth_context: AuthContext = Security(authorize)) -> Page[UserResponse]

Returns a list of all users.

Parameters:

Name Type Description Default
user_filter_model UserFilter

Model that takes care of filtering, sorting and pagination.

Depends(make_dependable(UserFilter))
hydrate bool

Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.

False
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
Page[UserResponse]

A list of all users.

Source code in src/zenml/zen_server/routers/users_endpoints.py
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@handle_exceptions
def list_users(
    user_filter_model: UserFilter = Depends(make_dependable(UserFilter)),
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[UserResponse]:
    """Returns a list of all users.

    Args:
        user_filter_model: Model that takes care of filtering, sorting and
            pagination.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: Authentication context.

    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}
    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
    )
    return dehydrate_page(page)
update_myself(user: UserUpdate, request: Request, auth_context: AuthContext = Security(authorize)) -> UserResponse

Updates a specific user.

Parameters:

Name Type Description Default
user UserUpdate

the user to use for the update.

required
request Request

The request object.

required
auth_context AuthContext

The authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

The updated user.

Raises:

Type Description
IllegalOperationError

if the current password is not supplied when changing the password or if the current password is incorrect.

Source code in src/zenml/zen_server/routers/users_endpoints.py
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
@current_user_router.put(
    "/current-user",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def update_myself(
    user: UserUpdate,
    request: Request,
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Updates a specific user.

    Args:
        user: the user to use for the update.
        request: The request object.
        auth_context: The authentication context.

    Returns:
        The updated user.

    Raises:
        IllegalOperationError: if the current password is not supplied when
            changing the password or if the current password is incorrect.
    """
    # Use a separate object to compute the update that will be applied to
    # the user to avoid giving the API requester direct control over the
    # user attributes that are updated.
    #
    # Exclude attributes that cannot be updated through this endpoint:
    #
    # - activation_token
    # - external_user_id
    # - admin
    # - is_active
    # - old_password
    #
    safe_user_update = user.create_copy(
        exclude={
            "activation_token",
            "external_user_id",
            "is_admin",
            "active",
            "old_password",
        },
    )

    # Validate a password change
    if user.password is not None:
        # If the user is updating their password, we need to verify
        # the old password
        if user.old_password is None:
            raise IllegalOperationError(
                "The current password must be supplied when changing the "
                "password."
            )
        with pass_change_limiter.limit_failed_requests(request):
            auth_user = zen_store().get_auth_user(auth_context.user.id)
            if not UserAuthModel.verify_password(
                user.old_password, auth_user
            ):
                raise IllegalOperationError(
                    "The current password is incorrect."
                )

        # Accept the password update
        safe_user_update.password = user.password

    updated_user = zen_store().update_user(
        user_id=auth_context.user.id, user_update=safe_user_update
    )
    return dehydrate_response_model(updated_user)
update_user(user_name_or_id: Union[str, UUID], user_update: UserUpdate, request: Request, auth_context: AuthContext = Security(authorize)) -> UserResponse

Updates a specific user.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
user_update UserUpdate

the user to use for the update.

required
request Request

The request object.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Returns:

Type Description
UserResponse

The updated user.

Raises:

Type Description
IllegalOperationError

if the user tries change admin status, while not an admin, if the user tries to change the password of another user, or if the user tries to change their own password without providing the old password or providing an incorrect old password.

Source code in src/zenml/zen_server/routers/users_endpoints.py
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
@router.put(
    "/{user_name_or_id}",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def update_user(
    user_name_or_id: Union[str, UUID],
    user_update: UserUpdate,
    request: Request,
    auth_context: AuthContext = Security(authorize),
) -> UserResponse:
    """Updates a specific user.

    Args:
        user_name_or_id: Name or ID of the user.
        user_update: the user to use for the update.
        request: The request object.
        auth_context: Authentication context.

    Returns:
        The updated user.

    Raises:
        IllegalOperationError: if the user tries change admin status,
            while not an admin, if the user tries to change the password
            of another user, or if the user tries to change their own
            password without providing the old password or providing
            an incorrect old password.
    """
    if server_config().auth_scheme == AuthScheme.EXTERNAL:
        # For external auth, we only allow updating the default project
        user_update = UserUpdate(
            default_project_id=user_update.default_project_id,
        )

    user = zen_store().get_user(user_name_or_id)

    # Use a separate object to compute the update that will be applied to
    # the user to avoid giving the API requester direct control over the
    # user attributes that are updated.
    #
    # Exclude attributes that cannot be updated through this endpoint:
    #
    # - activation_token
    # - external_user_id
    # - old_password
    #
    # Exclude things that are not always safe to update and need to be
    # validated first:
    #
    # - admin
    # - active
    # - password
    # - email_opted_in + email
    #
    safe_user_update = user_update.create_copy(
        exclude={
            "activation_token",
            "external_user_id",
            "is_admin",
            "active",
            "password",
            "old_password",
            "email_opted_in",
            "email",
        },
    )

    if user.id != auth_context.user.id:
        verify_admin_status_if_no_rbac(
            auth_context.user.is_admin, "update other user account"
        )
        # verify_permission_for_model(
        #     user,
        #     action=Action.UPDATE,
        # )

    # Validate a password change
    if user_update.password is not None:
        if user.id != auth_context.user.id:
            raise IllegalOperationError(
                "Users cannot change the password of other users. Use the "
                "account deactivation and activation flow instead."
            )

        # If the user is updating their own password, we need to verify
        # the old password
        if user_update.old_password is None:
            raise IllegalOperationError(
                "The current password must be supplied when changing the "
                "password."
            )

        with pass_change_limiter.limit_failed_requests(request):
            auth_user = zen_store().get_auth_user(user_name_or_id)
            if not UserAuthModel.verify_password(
                user_update.old_password, auth_user
            ):
                raise IllegalOperationError(
                    "The current password is incorrect."
                )

        # Accept the password update
        safe_user_update.password = user_update.password

    # Validate an admin status change
    if (
        user_update.is_admin is not None
        and user.is_admin != user_update.is_admin
    ):
        if user.id == auth_context.user.id:
            raise IllegalOperationError(
                "Cannot change the admin status of your own user account."
            )

        if user.id != auth_context.user.id and not auth_context.user.is_admin:
            raise IllegalOperationError(
                "Only admins are allowed to change the admin status of "
                "other user accounts."
            )

        # Accept the admin status update
        safe_user_update.is_admin = user_update.is_admin

    # Validate an active status change
    if user_update.active is not None and user.active != user_update.active:
        if user.id == auth_context.user.id:
            raise IllegalOperationError(
                "Cannot change the active status of your own user account."
            )

        if user.id != auth_context.user.id and not auth_context.user.is_admin:
            raise IllegalOperationError(
                "Only admins are allowed to change the active status of "
                "other user accounts."
            )

        # Accept the admin status update
        safe_user_update.is_admin = user_update.is_admin

    # Validate changes to private user account information
    if user_update.email_opted_in is not None or user_update.email is not None:
        if user.id != auth_context.user.id:
            raise IllegalOperationError(
                "Cannot change the private user account information for "
                "another user account."
            )

        # Accept the private user account information update
        if safe_user_update.email_opted_in is not None:
            safe_user_update.email_opted_in = user_update.email_opted_in
            safe_user_update.email = user_update.email

    updated_user = zen_store().update_user(
        user_id=user.id,
        user_update=safe_user_update,
    )
    return dehydrate_response_model(updated_user)
update_user_resource_membership(user_name_or_id: Union[str, UUID], resource_type: str, resource_id: UUID, actions: List[str], auth_context: AuthContext = Security(authorize)) -> None

Updates resource memberships of a user.

Parameters:

Name Type Description Default
user_name_or_id Union[str, UUID]

Name or ID of the user.

required
resource_type str

Type of the resource for which to update the membership.

required
resource_id UUID

ID of the resource for which to update the membership.

required
actions List[str]

List of actions that the user should be able to perform on the resource. If the user currently has permissions to perform actions which are not passed in this list, the permissions will be removed.

required
auth_context AuthContext

Authentication context.

Security(authorize)

Raises:

Type Description
ValueError

If a user tries to update their own membership.

KeyError

If no resource with the given type and ID exists.

Source code in src/zenml/zen_server/routers/users_endpoints.py
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
@router.post(
    "/{user_name_or_id}/resource_membership",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@handle_exceptions
def update_user_resource_membership(
    user_name_or_id: Union[str, UUID],
    resource_type: str,
    resource_id: UUID,
    actions: List[str],
    auth_context: AuthContext = Security(authorize),
) -> None:
    """Updates resource memberships of a user.

    Args:
        user_name_or_id: Name or ID of the user.
        resource_type: Type of the resource for which to update the
            membership.
        resource_id: ID of the resource for which to update the membership.
        actions: List of actions that the user should be able to perform on
            the resource. If the user currently has permissions to perform
            actions which are not passed in this list, the permissions will
            be removed.
        auth_context: Authentication context.

    Raises:
        ValueError: If a user tries to update their own membership.
        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)

    if user.id == auth_context.user.id:
        raise ValueError(
            "Not allowed to call endpoint with the authenticated user."
        )

    resource_type = ResourceType(resource_type)
    schema_class = get_schema_for_resource_type(resource_type)
    model = zen_store().get_entity_by_id(
        entity_id=resource_id, schema_class=schema_class
    )

    if not model:
        raise KeyError(
            f"Resource of type {resource_type} with ID {resource_id} does "
            "not exist."
        )

    project_id = None
    if isinstance(model, ProjectScopedResponse):
        project_id = model.project.id

    resource = Resource(
        type=resource_type, id=resource_id, project_id=project_id
    )

    verify_permission_for_model(model=model, action=Action.SHARE)
    for action in actions:
        # Make sure users aren't able to share permissions they don't have
        # themselves
        verify_permission_for_model(model=model, action=Action(action))

    update_resource_membership(
        user=user,
        resource=resource,
        actions=[Action(action) for action in actions],
    )
webhook_endpoints

Endpoint definitions for webhooks.

Classes Functions
get_body(request: Request) -> bytes async

Get access to the raw body.

Parameters:

Name Type Description Default
request Request

The request

required

Returns:

Type Description
bytes

The raw request body.

Source code in src/zenml/zen_server/routers/webhook_endpoints.py
42
43
44
45
46
47
48
49
50
51
async def get_body(request: Request) -> bytes:
    """Get access to the raw body.

    Args:
        request: The request

    Returns:
        The raw request body.
    """
    return await request.body()
webhook(event_source_id: UUID, request: Request, background_tasks: BackgroundTasks, raw_body: bytes = Depends(get_body)) -> Dict[str, str]

Webhook to receive events from external event sources.

Parameters:

Name Type Description Default
event_source_id UUID

The event_source_id

required
request Request

The request object

required
background_tasks BackgroundTasks

Background task handler

required
raw_body bytes

The raw request body

Depends(get_body)

Returns:

Type Description
Dict[str, str]

Static dict stating that event is received.

Raises:

Type Description
AuthorizationException

If the Event Source does not exist.

KeyError

If no appropriate Plugin found in the plugin registry

ValueError

If the id of the Event Source is not actually a webhook event source

WebhookInactiveError

In case this webhook has been deactivated

Source code in src/zenml/zen_server/routers/webhook_endpoints.py
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@router.post(
    "/{event_source_id}",
    response_model=Dict[str, str],
)
@handle_exceptions
def webhook(
    event_source_id: UUID,
    request: Request,
    background_tasks: BackgroundTasks,
    raw_body: bytes = Depends(get_body),
) -> Dict[str, str]:
    """Webhook to receive events from external event sources.

    Args:
        event_source_id: The event_source_id
        request: The request object
        background_tasks: Background task handler
        raw_body: The raw request body

    Returns:
        Static dict stating that event is received.

    Raises:
        AuthorizationException: If the Event Source does not exist.
        KeyError: If no appropriate Plugin found in the plugin registry
        ValueError: If the id of the Event Source is not actually a webhook event source
        WebhookInactiveError: In case this webhook has been deactivated
    """
    # Get the Event Source
    try:
        event_source = zen_store().get_event_source(event_source_id)
    except KeyError:
        logger.error(
            f"Webhook HTTP request received for unknown event source "
            f"'{event_source_id}'."
        )
        raise AuthorizationException(  # TODO: Are we sure about this error message?
            f"No webhook is registered at '{router.prefix}/{event_source_id}'"
        )
    if not event_source.is_active:
        raise WebhookInactiveError(f"Webhook {event_source_id} is inactive.")

    flavor = event_source.flavor
    try:
        plugin = plugin_flavor_registry().get_plugin(
            name=flavor,
            _type=PluginType.EVENT_SOURCE,
            subtype=PluginSubType.WEBHOOK,
        )
    except KeyError:
        logger.error(
            f"Webhook HTTP request received for event source "
            f"'{event_source_id}' and flavor {flavor} but no matching "
            f"plugin was found."
        )
        raise KeyError(
            f"No listener plugin found for event source {event_source_id}."
        )

    if not isinstance(plugin, BaseWebhookEventSourceHandler):
        raise ValueError(
            f"Event Source {event_source_id} is not a valid Webhook event "
            "source!"
        )

    # Pass the raw event and headers to the plugin
    background_tasks.add_task(
        plugin.process_webhook_event,
        event_source=event_source,
        raw_body=raw_body,
        headers=dict(request.headers.items()),
    )

    return {"status": "Event Received."}

secure_headers

Secure headers for the ZenML Server.

Functions
initialize_secure_headers() -> None

Initialize the secure headers component.

Source code in src/zenml/zen_server/secure_headers.py
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def initialize_secure_headers() -> None:
    """Initialize the secure headers component."""
    global _secure_headers

    config = server_config()

    # For each of the secure headers supported by the `secure` library, we
    # check if the corresponding configuration is set in the server
    # configuration:
    #
    # - if set to `True`, we use the default value for the header
    # - if set to a string, we use the string as the value for the header
    # - if set to `False`, we don't set the header

    server: Optional[secure.Server] = None
    if config.secure_headers_server:
        server = secure.Server()
        if isinstance(config.secure_headers_server, str):
            server.set(config.secure_headers_server)
        else:
            server.set(str(config.deployment_id))

    hsts: Optional[secure.StrictTransportSecurity] = None
    if config.secure_headers_hsts:
        hsts = secure.StrictTransportSecurity()
        if isinstance(config.secure_headers_hsts, str):
            hsts.set(config.secure_headers_hsts)

    xfo: Optional[secure.XFrameOptions] = None
    if config.secure_headers_xfo:
        xfo = secure.XFrameOptions()
        if isinstance(config.secure_headers_xfo, str):
            xfo.set(config.secure_headers_xfo)

    xxp: Optional[secure.XXSSProtection] = None
    if config.secure_headers_xxp:
        xxp = secure.XXSSProtection()
        if isinstance(config.secure_headers_xxp, str):
            xxp.set(config.secure_headers_xxp)

    csp: Optional[secure.ContentSecurityPolicy] = None
    if config.secure_headers_csp:
        csp = secure.ContentSecurityPolicy()
        if isinstance(config.secure_headers_csp, str):
            csp.set(config.secure_headers_csp)

    content: Optional[secure.XContentTypeOptions] = None
    if config.secure_headers_content:
        content = secure.XContentTypeOptions()
        if isinstance(config.secure_headers_content, str):
            content.set(config.secure_headers_content)

    referrer: Optional[secure.ReferrerPolicy] = None
    if config.secure_headers_referrer:
        referrer = secure.ReferrerPolicy()
        if isinstance(config.secure_headers_referrer, str):
            referrer.set(config.secure_headers_referrer)

    cache: Optional[secure.CacheControl] = None
    if config.secure_headers_cache:
        cache = secure.CacheControl()
        if isinstance(config.secure_headers_cache, str):
            cache.set(config.secure_headers_cache)

    permissions: Optional[secure.PermissionsPolicy] = None
    if config.secure_headers_permissions:
        permissions = secure.PermissionsPolicy()
        if isinstance(config.secure_headers_permissions, str):
            permissions.value = config.secure_headers_permissions

    _secure_headers = secure.Secure(
        server=server,
        hsts=hsts,
        xfo=xfo,
        xxp=xxp,
        csp=csp,
        content=content,
        referrer=referrer,
        cache=cache,
        permissions=permissions,
    )
secure_headers() -> secure.Secure

Return the secure headers component.

Returns:

Type Description
Secure

The secure headers component.

Raises:

Type Description
RuntimeError

If the secure headers component is not initialized.

Source code in src/zenml/zen_server/secure_headers.py
25
26
27
28
29
30
31
32
33
34
35
36
37
def secure_headers() -> secure.Secure:
    """Return the secure headers component.

    Returns:
        The secure headers component.

    Raises:
        RuntimeError: If the secure headers component is not initialized.
    """
    global _secure_headers
    if _secure_headers is None:
        raise RuntimeError("Secure headers component not initialized")
    return _secure_headers

template_execution

Modules
runner_entrypoint_configuration

Runner entrypoint configuration.

Classes
RunnerEntrypointConfiguration(arguments: List[str])

Bases: BaseEntrypointConfiguration

Runner entrypoint configuration.

Source code in src/zenml/entrypoints/base_entrypoint_configuration.py
60
61
62
63
64
65
66
def __init__(self, arguments: List[str]):
    """Initializes the entrypoint configuration.

    Args:
        arguments: Command line arguments to configure this object.
    """
    self.entrypoint_args = self._parse_arguments(arguments)
Functions
run() -> None

Run the entrypoint configuration.

This method runs the pipeline defined by the deployment given as input to the entrypoint configuration.

Source code in src/zenml/zen_server/template_execution/runner_entrypoint_configuration.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def run(self) -> None:
    """Run the entrypoint configuration.

    This method runs the pipeline defined by the deployment given as input
    to the entrypoint configuration.
    """
    deployment = self.load_deployment()

    stack = Client().active_stack
    assert deployment.stack and stack.id == deployment.stack.id

    placeholder_run = get_placeholder_run(deployment_id=deployment.id)
    deploy_pipeline(
        deployment=deployment,
        stack=stack,
        placeholder_run=placeholder_run,
    )
Functions
utils

Utility functions to run a pipeline from the server.

Classes Functions
deployment_request_from_template(template: RunTemplateResponse, config: PipelineRunConfiguration, user_id: UUID) -> PipelineDeploymentRequest

Generate a deployment request from a template.

Parameters:

Name Type Description Default
template RunTemplateResponse

The template from which to create the deployment request.

required
config PipelineRunConfiguration

The run configuration.

required
user_id UUID

ID of the user that is trying to run the template.

required

Raises:

Type Description
ValueError

If there are missing/extra step parameters in the run configuration.

Returns:

Type Description
PipelineDeploymentRequest

The generated deployment request.

Source code in src/zenml/zen_server/template_execution/utils.py
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
def deployment_request_from_template(
    template: RunTemplateResponse,
    config: PipelineRunConfiguration,
    user_id: UUID,
) -> "PipelineDeploymentRequest":
    """Generate a deployment request from a template.

    Args:
        template: The template from which to create the deployment request.
        config: The run configuration.
        user_id: ID of the user that is trying to run the template.

    Raises:
        ValueError: If there are missing/extra step parameters in the run
            configuration.

    Returns:
        The generated deployment request.
    """
    deployment = template.source_deployment
    assert deployment
    pipeline_configuration = PipelineConfiguration(
        **config.model_dump(
            include=set(PipelineConfiguration.model_fields),
            exclude={"name", "parameters"},
        ),
        name=deployment.pipeline_configuration.name,
        parameters=deployment.pipeline_configuration.parameters,
    )

    step_config_dict_base = pipeline_configuration.model_dump(
        exclude={"name", "parameters", "tags"}
    )
    steps = {}
    for invocation_id, step in deployment.step_configurations.items():
        step_config_dict = {
            **copy.deepcopy(step_config_dict_base),
            **step.config.model_dump(
                # TODO: Maybe we need to make some of these configurable via
                # yaml as well, e.g. the lazy loaders?
                include={
                    "name",
                    "caching_parameters",
                    "external_input_artifacts",
                    "model_artifacts_or_metadata",
                    "client_lazy_loaders",
                    "substitutions",
                    "outputs",
                }
            ),
        }

        required_parameters = set(step.config.parameters)
        configured_parameters = set()

        if update := config.steps.get(invocation_id):
            update_dict = update.model_dump()
            # Get rid of deprecated name to prevent overriding the step name
            # with `None`.
            update_dict.pop("name", None)
            configured_parameters = set(update.parameters)
            step_config_dict = dict_utils.recursive_update(
                step_config_dict, update=update_dict
            )

        unknown_parameters = configured_parameters - required_parameters
        if unknown_parameters:
            raise ValueError(
                "Run configuration contains the following unknown "
                f"parameters for step {step.config.name}: {unknown_parameters}."
            )

        missing_parameters = required_parameters - configured_parameters
        if missing_parameters:
            raise ValueError(
                "Run configuration is missing the following required "
                f"parameters for step {step.config.name}: {missing_parameters}."
            )

        step_config = StepConfiguration.model_validate(step_config_dict)
        steps[invocation_id] = Step(spec=step.spec, config=step_config)

    code_reference_request = None
    if deployment.code_reference:
        code_reference_request = CodeReferenceRequest(
            commit=deployment.code_reference.commit,
            subdirectory=deployment.code_reference.subdirectory,
            code_repository=deployment.code_reference.code_repository.id,
        )

    zenml_version = zen_store().get_store_info().version
    assert deployment.stack
    assert deployment.build
    deployment_request = PipelineDeploymentRequest(
        project=deployment.project.id,
        run_name_template=config.run_name
        or get_default_run_name(pipeline_name=pipeline_configuration.name),
        pipeline_configuration=pipeline_configuration,
        step_configurations=steps,
        client_environment={},
        client_version=zenml_version,
        server_version=zenml_version,
        stack=deployment.stack.id,
        pipeline=deployment.pipeline.id if deployment.pipeline else None,
        build=deployment.build.id,
        schedule=None,
        code_reference=code_reference_request,
        code_path=deployment.code_path,
        template=template.id,
        pipeline_version_hash=deployment.pipeline_version_hash,
        pipeline_spec=deployment.pipeline_spec,
    )

    return deployment_request
ensure_async_orchestrator(deployment: PipelineDeploymentRequest, stack: StackResponse) -> None

Ensures the orchestrator is configured to run async.

Parameters:

Name Type Description Default
deployment PipelineDeploymentRequest

Deployment request in which the orchestrator configuration should be updated to ensure the orchestrator is running async.

required
stack StackResponse

The stack on which the deployment will run.

required
Source code in src/zenml/zen_server/template_execution/utils.py
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def ensure_async_orchestrator(
    deployment: PipelineDeploymentRequest, stack: StackResponse
) -> None:
    """Ensures the orchestrator is configured to run async.

    Args:
        deployment: Deployment request in which the orchestrator
            configuration should be updated to ensure the orchestrator is
            running async.
        stack: The stack on which the deployment will run.
    """
    orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
    flavors = zen_store().list_flavors(
        FlavorFilter(name=orchestrator.flavor_name, type=orchestrator.type)
    )
    flavor = Flavor.from_model(flavors[0])

    if "synchronous" in flavor.config_class.model_fields:
        key = settings_utils.get_flavor_setting_key(flavor)

        if settings := deployment.pipeline_configuration.settings.get(key):
            settings_dict = settings.model_dump()
        else:
            settings_dict = {}

        settings_dict["synchronous"] = False
        deployment.pipeline_configuration.settings[key] = (
            BaseSettings.model_validate(settings_dict)
        )
generate_dockerfile(pypi_requirements: List[str], apt_packages: List[str], zenml_version: str, python_version: str) -> str

Generate a Dockerfile that installs the requirements.

Parameters:

Name Type Description Default
pypi_requirements List[str]

The PyPI requirements to install.

required
apt_packages List[str]

The APT packages to install.

required
zenml_version str

The ZenML version to use as parent image.

required
python_version str

The Python version to use as parent image.

required

Returns:

Type Description
str

The Dockerfile.

Source code in src/zenml/zen_server/template_execution/utils.py
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
def generate_dockerfile(
    pypi_requirements: List[str],
    apt_packages: List[str],
    zenml_version: str,
    python_version: str,
) -> str:
    """Generate a Dockerfile that installs the requirements.

    Args:
        pypi_requirements: The PyPI requirements to install.
        apt_packages: The APT packages to install.
        zenml_version: The ZenML version to use as parent image.
        python_version: The Python version to use as parent image.

    Returns:
        The Dockerfile.
    """
    parent_image = f"zenmldocker/zenml:{zenml_version}-py{python_version}"

    lines = [f"FROM {parent_image}"]
    if apt_packages:
        apt_packages_string = " ".join(f"'{p}'" for p in apt_packages)
        lines.append(
            "RUN apt-get update && apt-get install -y "
            f"--no-install-recommends {apt_packages_string}"
        )

    if pypi_requirements:
        pypi_requirements_string = " ".join(
            [f"'{r}'" for r in pypi_requirements]
        )

        if handle_bool_env_var(
            ENV_ZENML_RUNNER_IMAGE_DISABLE_UV, default=False
        ):
            lines.append(
                f"RUN pip install --default-timeout=60 --no-cache-dir "
                f"{pypi_requirements_string}"
            )
        else:
            lines.append("RUN pip install uv")
            lines.append(
                f"RUN uv pip install --no-cache-dir {pypi_requirements_string}"
            )

    return "\n".join(lines)
generate_image_hash(dockerfile: str) -> str

Generate a hash of the Dockerfile.

Parameters:

Name Type Description Default
dockerfile str

The Dockerfile for which to generate the hash.

required

Returns:

Type Description
str

The hash of the Dockerfile.

Source code in src/zenml/zen_server/template_execution/utils.py
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def generate_image_hash(dockerfile: str) -> str:
    """Generate a hash of the Dockerfile.

    Args:
        dockerfile: The Dockerfile for which to generate the hash.

    Returns:
        The hash of the Dockerfile.
    """
    hash_ = hashlib.md5()  # nosec
    # Uncomment this line when developing to guarantee a new docker image gets
    # built after restarting the server
    # hash_.update(f"{os.getpid()}".encode())
    hash_.update(dockerfile.encode())
    return hash_.hexdigest()
get_pipeline_run_analytics_metadata(deployment: PipelineDeploymentResponse, stack: StackResponse, template_id: UUID, run_id: UUID) -> Dict[str, Any]

Get metadata for the pipeline run analytics event.

Parameters:

Name Type Description Default
deployment PipelineDeploymentResponse

The deployment of the run.

required
stack StackResponse

The stack on which the run will happen.

required
template_id UUID

ID of the template from which the run was started.

required
run_id UUID

ID of the run.

required

Returns:

Type Description
Dict[str, Any]

The analytics metadata.

Source code in src/zenml/zen_server/template_execution/utils.py
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
def get_pipeline_run_analytics_metadata(
    deployment: "PipelineDeploymentResponse",
    stack: StackResponse,
    template_id: UUID,
    run_id: UUID,
) -> Dict[str, Any]:
    """Get metadata for the pipeline run analytics event.

    Args:
        deployment: The deployment of the run.
        stack: The stack on which the run will happen.
        template_id: ID of the template from which the run was started.
        run_id: ID of the run.

    Returns:
        The analytics metadata.
    """
    custom_materializer = False
    for step in deployment.step_configurations.values():
        for output in step.config.outputs.values():
            for source in output.materializer_source:
                if not source.is_internal:
                    custom_materializer = True

    assert deployment.user
    stack_creator = stack.user
    own_stack = stack_creator and stack_creator.id == deployment.user.id

    stack_metadata = {
        component_type.value: component_list[0].flavor_name
        for component_type, component_list in stack.components.items()
    }

    return {
        "project_id": deployment.project.id,
        "store_type": "rest",  # This method is called from within a REST endpoint
        **stack_metadata,
        "total_steps": len(deployment.step_configurations),
        "schedule": deployment.schedule is not None,
        "custom_materializer": custom_materializer,
        "own_stack": own_stack,
        "pipeline_run_id": str(run_id),
        "template_id": str(template_id),
    }
run_template(template: RunTemplateResponse, auth_context: AuthContext, background_tasks: Optional[BackgroundTasks] = None, run_config: Optional[PipelineRunConfiguration] = None) -> PipelineRunResponse

Run a pipeline from a template.

Parameters:

Name Type Description Default
template RunTemplateResponse

The template to run.

required
auth_context AuthContext

Authentication context.

required
background_tasks Optional[BackgroundTasks]

Background tasks.

None
run_config Optional[PipelineRunConfiguration]

The run configuration.

None

Raises:

Type Description
ValueError

If the template can not be run.

RuntimeError

If the server URL is not set in the server configuration.

Returns:

Type Description
PipelineRunResponse

ID of the new pipeline run.

Source code in src/zenml/zen_server/template_execution/utils.py
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def run_template(
    template: RunTemplateResponse,
    auth_context: AuthContext,
    background_tasks: Optional[BackgroundTasks] = None,
    run_config: Optional[PipelineRunConfiguration] = None,
) -> PipelineRunResponse:
    """Run a pipeline from a template.

    Args:
        template: The template to run.
        auth_context: Authentication context.
        background_tasks: Background tasks.
        run_config: The run configuration.

    Raises:
        ValueError: If the template can not be run.
        RuntimeError: If the server URL is not set in the server configuration.

    Returns:
        ID of the new pipeline run.
    """
    if not template.runnable:
        raise ValueError(
            "This template can not be run because its associated deployment, "
            "stack or build have been deleted."
        )

    # Guaranteed by the `runnable` check above
    build = template.build
    assert build
    stack = build.stack
    assert stack

    if build.stack_checksum and build.stack_checksum != compute_stack_checksum(
        stack=stack
    ):
        raise ValueError(
            f"The stack {stack.name} has been updated since it was used for "
            "the run that is the base for this template. This means the Docker "
            "images associated with this template most likely do not contain "
            "the necessary requirements. Please create a new template from a "
            "recent run on this stack."
        )

    validate_stack_is_runnable_from_server(zen_store=zen_store(), stack=stack)
    if run_config:
        validate_run_config_is_runnable_from_server(run_config)

    deployment_request = deployment_request_from_template(
        template=template,
        config=run_config or PipelineRunConfiguration(),
        user_id=auth_context.user.id,
    )

    ensure_async_orchestrator(deployment=deployment_request, stack=stack)

    new_deployment = zen_store().create_deployment(deployment_request)

    server_url = server_config().server_url
    if not server_url:
        raise RuntimeError(
            "The server URL is not set in the server configuration."
        )
    assert build.zenml_version
    zenml_version = build.zenml_version

    placeholder_run = create_placeholder_run(deployment=new_deployment)
    assert placeholder_run

    # We create an API token scoped to the pipeline run that never expires
    api_token = generate_access_token(
        user_id=auth_context.user.id,
        pipeline_run_id=placeholder_run.id,
        # Keep the original API key or device scopes, if any
        api_key=auth_context.api_key,
        device=auth_context.device,
        # Never expire the token
        expires_in=0,
    ).access_token

    environment = {
        ENV_ZENML_ACTIVE_PROJECT_ID: str(new_deployment.project.id),
        ENV_ZENML_ACTIVE_STACK_ID: str(stack.id),
        "ZENML_VERSION": zenml_version,
        "ZENML_STORE_URL": server_url,
        "ZENML_STORE_TYPE": StoreType.REST.value,
        "ZENML_STORE_API_TOKEN": api_token,
        "ZENML_STORE_VERIFY_SSL": "True",
    }

    command = RunnerEntrypointConfiguration.get_entrypoint_command()
    args = RunnerEntrypointConfiguration.get_entrypoint_arguments(
        deployment_id=new_deployment.id
    )

    if build.python_version:
        version_info = version.parse(build.python_version)
        python_version = f"{version_info.major}.{version_info.minor}"
    else:
        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"

    (
        pypi_requirements,
        apt_packages,
    ) = requirements_utils.get_requirements_for_stack(
        stack=stack, python_version=python_version
    )

    dockerfile = generate_dockerfile(
        pypi_requirements=pypi_requirements,
        apt_packages=apt_packages,
        zenml_version=zenml_version,
        python_version=python_version,
    )

    # Building a docker image with requirements and apt packages from the
    # stack only (no code). Ideally, only orchestrator requirements should
    # be added to the docker image, but we have to instantiate the entire
    # stack to get the orchestrator to run pipelines.
    image_hash = generate_image_hash(dockerfile=dockerfile)
    logger.info(
        "Building runner image %s for dockerfile:\n%s", image_hash, dockerfile
    )

    def _task() -> None:
        runner_image = workload_manager().build_and_push_image(
            workload_id=new_deployment.id,
            dockerfile=dockerfile,
            image_name=f"{RUNNER_IMAGE_REPOSITORY}:{image_hash}",
            sync=True,
        )

        workload_manager().log(
            workload_id=new_deployment.id,
            message="Starting pipeline run.",
        )

        # could do this same thing with a step operator, but we need some
        # minor changes to the abstract interface to support that.
        workload_manager().run(
            workload_id=new_deployment.id,
            image=runner_image,
            command=command,
            arguments=args,
            environment=environment,
            timeout_in_seconds=30,
            sync=True,
        )
        workload_manager().log(
            workload_id=new_deployment.id,
            message="Pipeline run started successfully.",
        )

    def _task_with_analytics_and_error_handling() -> None:
        with track_handler(
            event=AnalyticsEvent.RUN_PIPELINE
        ) as analytics_handler:
            analytics_handler.metadata = get_pipeline_run_analytics_metadata(
                deployment=new_deployment,
                stack=stack,
                template_id=template.id,
                run_id=placeholder_run.id,
            )

            try:
                _task()
            except Exception:
                logger.exception(
                    "Failed to run template %s, run ID: %s",
                    str(template.id),
                    str(placeholder_run.id),
                )
                zen_store().update_run(
                    run_id=placeholder_run.id,
                    run_update=PipelineRunUpdate(
                        status=ExecutionStatus.FAILED
                    ),
                )
                raise

    if background_tasks:
        background_tasks.add_task(_task_with_analytics_and_error_handling)
    else:
        # Run synchronously if no background tasks were passed. This is probably
        # when coming from a trigger which itself is already running in the
        # background
        _task_with_analytics_and_error_handling()

    return placeholder_run
Modules
workload_manager_interface

Workload manager interface definition.

Classes
WorkloadManagerInterface

Bases: ABC

Workload manager interface.

Functions
build_and_push_image(workload_id: UUID, dockerfile: str, image_name: str, sync: bool = True, timeout_in_seconds: int = 0) -> str abstractmethod

Build and push a Docker image.

Parameters:

Name Type Description Default
workload_id UUID

Workload ID.

required
dockerfile str

The dockerfile content to build the image.

required
image_name str

The image repository and tag.

required
sync bool

If True, will wait until the build finished before returning.

True
timeout_in_seconds int

Timeout in seconds to wait before cancelling the container. If set to 0 the container will run until it fails or finishes.

0

Returns:

Type Description
str

The full image name including container registry.

Source code in src/zenml/zen_server/template_execution/workload_manager_interface.py
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@abstractmethod
def build_and_push_image(
    self,
    workload_id: UUID,
    dockerfile: str,
    image_name: str,
    sync: bool = True,
    timeout_in_seconds: int = 0,
) -> str:
    """Build and push a Docker image.

    Args:
        workload_id: Workload ID.
        dockerfile: The dockerfile content to build the image.
        image_name: The image repository and tag.
        sync: If True, will wait until the build finished before returning.
        timeout_in_seconds: Timeout in seconds to wait before cancelling
            the container. If set to 0 the container will run until it
            fails or finishes.

    Returns:
        The full image name including container registry.
    """
    pass
delete_workload(workload_id: UUID) -> None abstractmethod

Delete a workload.

Parameters:

Name Type Description Default
workload_id UUID

Workload ID.

required
Source code in src/zenml/zen_server/template_execution/workload_manager_interface.py
63
64
65
66
67
68
69
70
@abstractmethod
def delete_workload(self, workload_id: UUID) -> None:
    """Delete a workload.

    Args:
        workload_id: Workload ID.
    """
    pass
get_logs(workload_id: UUID) -> str abstractmethod

Get logs for a workload.

Parameters:

Name Type Description Default
workload_id UUID

Workload ID.

required

Returns:

Type Description
str

The stored logs.

Source code in src/zenml/zen_server/template_execution/workload_manager_interface.py
72
73
74
75
76
77
78
79
80
81
82
@abstractmethod
def get_logs(self, workload_id: UUID) -> str:
    """Get logs for a workload.

    Args:
        workload_id: Workload ID.

    Returns:
        The stored logs.
    """
    pass
log(workload_id: UUID, message: str) -> None abstractmethod

Log a message.

Parameters:

Name Type Description Default
workload_id UUID

Workload ID.

required
message str

The message to log.

required
Source code in src/zenml/zen_server/template_execution/workload_manager_interface.py
84
85
86
87
88
89
90
91
92
@abstractmethod
def log(self, workload_id: UUID, message: str) -> None:
    """Log a message.

    Args:
        workload_id: Workload ID.
        message: The message to log.
    """
    pass
run(workload_id: UUID, image: str, command: List[str], arguments: List[str], environment: Optional[Dict[str, str]] = None, sync: bool = True, timeout_in_seconds: int = 0) -> None abstractmethod

Run a Docker container.

Parameters:

Name Type Description Default
workload_id UUID

Workload ID.

required
image str

The Docker image to run.

required
command List[str]

The command to run in the container.

required
arguments List[str]

The arguments for the command.

required
environment Optional[Dict[str, str]]

The environment to set in the container.

None
sync bool

If True, will wait until the container finished running before returning.

True
timeout_in_seconds int

Timeout in seconds to wait before cancelling the container. If set to 0 the container will run until it fails or finishes.

0
Source code in src/zenml/zen_server/template_execution/workload_manager_interface.py
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
@abstractmethod
def run(
    self,
    workload_id: UUID,
    image: str,
    command: List[str],
    arguments: List[str],
    environment: Optional[Dict[str, str]] = None,
    sync: bool = True,
    timeout_in_seconds: int = 0,
) -> None:
    """Run a Docker container.

    Args:
        workload_id: Workload ID.
        image: The Docker image to run.
        command: The command to run in the container.
        arguments: The arguments for the command.
        environment: The environment to set in the container.
        sync: If True, will wait until the container finished running before
            returning.
        timeout_in_seconds: Timeout in seconds to wait before cancelling
            the container. If set to 0 the container will run until it
            fails or finishes.
    """
    pass

utils

Util functions for the ZenML Server.

Classes
Functions
feature_gate() -> FeatureGateInterface

Return the initialized Feature Gate component.

Raises:

Type Description
RuntimeError

If the RBAC component is not initialized.

Returns:

Type Description
FeatureGateInterface

The RBAC component.

Source code in src/zenml/zen_server/utils.py
131
132
133
134
135
136
137
138
139
140
141
142
143
def feature_gate() -> FeatureGateInterface:
    """Return the initialized Feature Gate component.

    Raises:
        RuntimeError: If the RBAC component is not initialized.

    Returns:
        The RBAC component.
    """
    global _feature_gate
    if _feature_gate is None:
        raise RuntimeError("Feature gate component not initialized.")
    return _feature_gate
get_ip_location(ip_address: str) -> Tuple[str, str, str]

Get the location of the given IP address.

Parameters:

Name Type Description Default
ip_address str

The IP address to get the location for.

required

Returns:

Type Description
Tuple[str, str, str]

A tuple of city, region, country.

Source code in src/zenml/zen_server/utils.py
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
def get_ip_location(ip_address: str) -> Tuple[str, str, str]:
    """Get the location of the given IP address.

    Args:
        ip_address: The IP address to get the location for.

    Returns:
        A tuple of city, region, country.
    """
    import ipinfo  # type: ignore[import-untyped]

    try:
        handler = ipinfo.getHandler()
        details = handler.getDetails(ip_address)
        return (
            details.city,
            details.region,
            details.country_name,
        )
    except Exception:
        logger.exception(f"Could not get IP location for {ip_address}.")
        return "", "", ""
get_zenml_headers() -> Dict[str, str]

Get the ZenML specific headers to be included in requests made by the server.

Returns:

Type Description
Dict[str, str]

The ZenML specific headers.

Source code in src/zenml/zen_server/utils.py
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
def get_zenml_headers() -> Dict[str, str]:
    """Get the ZenML specific headers to be included in requests made by the server.

    Returns:
        The ZenML specific headers.
    """
    config = server_config()
    headers = {
        "zenml-server-id": str(config.get_external_server_id()),
        "zenml-server-version": zenml_version,
    }
    if config.server_url:
        headers["zenml-server-url"] = config.server_url

    return headers
handle_exceptions(func: F) -> F

Decorator to handle exceptions in the API.

Parameters:

Name Type Description Default
func F

Function to decorate.

required

Returns:

Type Description
F

Decorated function.

Source code in src/zenml/zen_server/utils.py
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def handle_exceptions(func: F) -> F:
    """Decorator to handle exceptions in the API.

    Args:
        func: Function to decorate.

    Returns:
        Decorated function.
    """

    @wraps(func)
    def decorated(*args: Any, **kwargs: Any) -> Any:
        # These imports can't happen at module level as this module is also
        # used by the CLI when installed without the `server` extra
        from fastapi import HTTPException
        from fastapi.responses import JSONResponse

        from zenml.zen_server.auth import AuthContext, set_auth_context

        for arg in args:
            if isinstance(arg, AuthContext):
                set_auth_context(arg)
                break
        else:
            for _, arg in kwargs.items():
                if isinstance(arg, AuthContext):
                    set_auth_context(arg)
                    break

        try:
            return func(*args, **kwargs)
        except OAuthError as error:
            # The OAuthError is special because it needs to have a JSON response
            return JSONResponse(
                status_code=error.status_code,
                content=error.to_dict(),
            )
        except HTTPException:
            raise
        except Exception as error:
            logger.exception("API error")
            http_exception = http_exception_from_error(error)
            raise http_exception

    return cast(F, decorated)
initialize_feature_gate() -> None

Initialize the Feature Gate component.

Source code in src/zenml/zen_server/utils.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def initialize_feature_gate() -> None:
    """Initialize the Feature Gate component."""
    global _feature_gate

    if (
        feature_gate_source
        := server_config().feature_gate_implementation_source
    ):
        from zenml.utils import source_utils

        implementation_class = source_utils.load_and_validate_class(
            feature_gate_source, expected_class=FeatureGateInterface
        )
        _feature_gate = implementation_class()
initialize_memcache(max_capacity: int, default_expiry: int) -> None

Initialize the memory cache.

Parameters:

Name Type Description Default
max_capacity int

The maximum capacity of the cache.

required
default_expiry int

The default expiry time in seconds.

required
Source code in src/zenml/zen_server/utils.py
228
229
230
231
232
233
234
235
236
def initialize_memcache(max_capacity: int, default_expiry: int) -> None:
    """Initialize the memory cache.

    Args:
        max_capacity: The maximum capacity of the cache.
        default_expiry: The default expiry time in seconds.
    """
    global _memcache
    _memcache = MemoryCache(max_capacity, default_expiry)
initialize_plugins() -> None

Initialize the event plugins registry.

Source code in src/zenml/zen_server/utils.py
199
200
201
def initialize_plugins() -> None:
    """Initialize the event plugins registry."""
    plugin_flavor_registry().initialize_plugins()
initialize_rbac() -> None

Initialize the RBAC component.

Source code in src/zenml/zen_server/utils.py
118
119
120
121
122
123
124
125
126
127
128
def initialize_rbac() -> None:
    """Initialize the RBAC component."""
    global _rbac

    if rbac_source := server_config().rbac_implementation_source:
        from zenml.utils import source_utils

        implementation_class = source_utils.load_and_validate_class(
            rbac_source, expected_class=RBACInterface
        )
        _rbac = implementation_class()
initialize_workload_manager() -> None

Initialize the workload manager component.

This does not fail if the source can't be loaded but only logs a warning.

Source code in src/zenml/zen_server/utils.py
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def initialize_workload_manager() -> None:
    """Initialize the workload manager component.

    This does not fail if the source can't be loaded but only logs a warning.
    """
    global _workload_manager

    if source := server_config().workload_manager_implementation_source:
        from zenml.utils import source_utils

        try:
            workload_manager_class: Type[WorkloadManagerInterface] = (
                source_utils.load_and_validate_class(
                    source=source, expected_class=WorkloadManagerInterface
                )
            )
        except (ModuleNotFoundError, KeyError):
            logger.warning("Unable to load workload manager source.")
        else:
            _workload_manager = workload_manager_class()
initialize_zen_store() -> None

Initialize the ZenML Store.

Raises:

Type Description
ValueError

If the ZenML Store is using a REST back-end.

Source code in src/zenml/zen_server/utils.py
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def initialize_zen_store() -> None:
    """Initialize the ZenML Store.

    Raises:
        ValueError: If the ZenML Store is using a REST back-end.
    """
    logger.debug("Initializing ZenML Store for FastAPI...")

    # Use an environment variable to flag the instance as a server
    os.environ[ENV_ZENML_SERVER] = "true"

    zen_store_ = GlobalConfiguration().zen_store

    if not isinstance(zen_store_, SqlZenStore):
        raise ValueError(
            "Server cannot be started with a REST store type. Make sure you "
            "configure ZenML to use a non-networked store backend "
            "when trying to start the ZenML Server."
        )

    global _zen_store
    _zen_store = zen_store_
is_same_or_subdomain(source_domain: str, target_domain: str) -> bool

Check if the source domain is the same or a subdomain of the target domain.

Examples:

is_same_or_subdomain("example.com", "example.com") -> True is_same_or_subdomain("alpha.example.com", "example.com") -> True is_same_or_subdomain("alpha.example.com", ".example.com") -> True is_same_or_subdomain("example.com", "alpha.example.com") -> False is_same_or_subdomain("alpha.beta.example.com", "beta.example.com") -> True is_same_or_subdomain("alpha.beta.example.com", "alpha.example.com") -> False is_same_or_subdomain("alphabeta.gamma.example", "beta.gamma.example") -> False

Parameters:

Name Type Description Default
source_domain str

The source domain to check.

required
target_domain str

The target domain to compare against.

required

Returns:

Type Description
bool

True if the source domain is the same or a subdomain of the target

bool

domain, False otherwise.

Source code in src/zenml/zen_server/utils.py
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
def is_same_or_subdomain(source_domain: str, target_domain: str) -> bool:
    """Check if the source domain is the same or a subdomain of the target domain.

    Examples:
        is_same_or_subdomain("example.com", "example.com") -> True
        is_same_or_subdomain("alpha.example.com", "example.com") -> True
        is_same_or_subdomain("alpha.example.com", ".example.com") -> True
        is_same_or_subdomain("example.com", "alpha.example.com") -> False
        is_same_or_subdomain("alpha.beta.example.com", "beta.example.com") -> True
        is_same_or_subdomain("alpha.beta.example.com", "alpha.example.com") -> False
        is_same_or_subdomain("alphabeta.gamma.example", "beta.gamma.example") -> False

    Args:
        source_domain: The source domain to check.
        target_domain: The target domain to compare against.

    Returns:
        True if the source domain is the same or a subdomain of the target
        domain, False otherwise.
    """
    import tldextract

    # Extract the registered domain and suffix for both
    src_parts = tldextract.extract(source_domain)
    tgt_parts = tldextract.extract(target_domain)

    if src_parts == tgt_parts:
        return True  # Same domain

    # Reconstruct the base domains (e.g., example.com)
    src_base_domain = f"{src_parts.domain}.{src_parts.suffix}"
    tgt_base_domain = f"{tgt_parts.domain}.{tgt_parts.suffix}"

    if src_base_domain != tgt_base_domain:
        return False  # Different base domains

    if tgt_parts.subdomain == "":
        return True  # Subdomain

    if src_parts.subdomain.endswith(f".{tgt_parts.subdomain.lstrip('.')}"):
        return True  # Subdomain of subdomain

    return False
is_user_request(request: Request) -> bool

Determine if the incoming request is a user request.

This function checks various aspects of the request to determine if it's a user-initiated request or a system request.

Parameters:

Name Type Description Default
request Request

The incoming FastAPI request object.

required

Returns:

Type Description
bool

True if it's a user request, False otherwise.

Source code in src/zenml/zen_server/utils.py
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
def is_user_request(request: "Request") -> bool:
    """Determine if the incoming request is a user request.

    This function checks various aspects of the request to determine
    if it's a user-initiated request or a system request.

    Args:
        request: The incoming FastAPI request object.

    Returns:
        True if it's a user request, False otherwise.
    """
    # Define system paths that should be excluded
    system_paths: List[str] = [
        "/health",
        "/metrics",
        "/system",
        "/docs",
        "/redoc",
        "/openapi.json",
    ]

    user_prefix = f"{API}{VERSION_1}"
    excluded_user_apis = [INFO]
    # Check if this is not an excluded endpoint
    if request.url.path in [
        user_prefix + suffix for suffix in excluded_user_apis
    ]:
        return False

    # Check if this is other user request
    if request.url.path.startswith(user_prefix):
        return True

    # Exclude system paths
    if any(request.url.path.startswith(path) for path in system_paths):
        return False

    # Exclude requests with specific headers
    if request.headers.get("X-System-Request") == "true":
        return False

    # Exclude requests from certain user agents (e.g., monitoring tools)
    user_agent = request.headers.get("User-Agent", "").lower()
    system_agents = ["prometheus", "datadog", "newrelic", "pingdom"]
    if any(agent in user_agent for agent in system_agents):
        return False

    # Check for internal IP addresses
    client_host = request.client.host if request.client else None
    if client_host and (
        client_host.startswith("10.") or client_host.startswith("192.168.")
    ):
        return False

    # Exclude OPTIONS requests (often used for CORS preflight)
    if request.method == "OPTIONS":
        return False

    # Exclude specific query parameters that might indicate system requests
    if request.query_params.get("system_check"):
        return False

    # If none of the above conditions are met, consider it a user request
    return True
make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]

This function makes a pydantic model usable for fastapi query parameters.

Additionally, it converts InternalServerErrors that would happen due to pydantic.ValidationError into 422 responses that signal an invalid request.

Check out https://github.com/tiangolo/fastapi/issues/1474 for context.

Usage

def f(model: Model = Depends(make_dependable(Model))): ...

UPDATE: Function from above mentioned Github issue was extended to support multi-input parameters, e.g. tags: List[str]. It needs a default set to Query(), rather just plain .

Parameters:

Name Type Description Default
cls Type[BaseModel]

The model class.

required

Returns:

Type Description
Callable[..., Any]

Function to use in FastAPI Depends.

Source code in src/zenml/zen_server/utils.py
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
    """This function makes a pydantic model usable for fastapi query parameters.

    Additionally, it converts `InternalServerError`s that would happen due to
    `pydantic.ValidationError` into 422 responses that signal an invalid
    request.

    Check out https://github.com/tiangolo/fastapi/issues/1474 for context.

    Usage:
        def f(model: Model = Depends(make_dependable(Model))):
            ...

    UPDATE: Function from above mentioned Github issue was extended to support
    multi-input parameters, e.g. tags: List[str]. It needs a default set to Query(<default>),
    rather just plain <default>.

    Args:
        cls: The model class.

    Returns:
        Function to use in FastAPI `Depends`.
    """
    from fastapi import Query

    from zenml.zen_server.exceptions import error_detail

    def init_cls_and_handle_errors(*args: Any, **kwargs: Any) -> BaseModel:
        from fastapi import HTTPException

        try:
            inspect.signature(init_cls_and_handle_errors).bind(*args, **kwargs)
            return cls(*args, **kwargs)
        except ValidationError as e:
            detail = error_detail(e, exception_type=ValueError)
            raise HTTPException(422, detail=detail)

    params = {v.name: v for v in inspect.signature(cls).parameters.values()}
    query_params = getattr(cls, "API_MULTI_INPUT_PARAMS", [])
    for qp in query_params:
        if qp in params:
            params[qp] = inspect.Parameter(
                name=params[qp].name,
                default=Query(params[qp].default),
                kind=params[qp].kind,
                annotation=params[qp].annotation,
            )

    init_cls_and_handle_errors.__signature__ = inspect.Signature(  # type: ignore[attr-defined]
        parameters=[v for v in params.values()]
    )

    return init_cls_and_handle_errors
memcache() -> MemoryCache

Return the memory cache.

Returns:

Type Description
MemoryCache

The memory cache.

Raises:

Type Description
RuntimeError

If the memory cache is not initialized.

Source code in src/zenml/zen_server/utils.py
239
240
241
242
243
244
245
246
247
248
249
250
def memcache() -> MemoryCache:
    """Return the memory cache.

    Returns:
        The memory cache.

    Raises:
        RuntimeError: If the memory cache is not initialized.
    """
    if _memcache is None:
        raise RuntimeError("Memory cache not initialized")
    return _memcache
plugin_flavor_registry() -> PluginFlavorRegistry

Get the plugin flavor registry.

Returns:

Type Description
PluginFlavorRegistry

The plugin flavor registry.

Source code in src/zenml/zen_server/utils.py
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def plugin_flavor_registry() -> PluginFlavorRegistry:
    """Get the plugin flavor registry.

    Returns:
        The plugin flavor registry.
    """
    global _plugin_flavor_registry

    if _plugin_flavor_registry is None:
        _plugin_flavor_registry = PluginFlavorRegistry()
        _plugin_flavor_registry.initialize_plugins()
    return _plugin_flavor_registry
rbac() -> RBACInterface

Return the initialized RBAC component.

Raises:

Type Description
RuntimeError

If the RBAC component is not initialized.

Returns:

Type Description
RBACInterface

The RBAC component.

Source code in src/zenml/zen_server/utils.py
103
104
105
106
107
108
109
110
111
112
113
114
115
def rbac() -> RBACInterface:
    """Return the initialized RBAC component.

    Raises:
        RuntimeError: If the RBAC component is not initialized.

    Returns:
        The RBAC component.
    """
    global _rbac
    if _rbac is None:
        raise RuntimeError("RBAC component not initialized")
    return _rbac
server_config() -> ServerConfiguration

Returns the ZenML Server configuration.

Returns:

Type Description
ServerConfiguration

The ZenML Server configuration.

Source code in src/zenml/zen_server/utils.py
256
257
258
259
260
261
262
263
264
265
def server_config() -> ServerConfiguration:
    """Returns the ZenML Server configuration.

    Returns:
        The ZenML Server configuration.
    """
    global _server_config
    if _server_config is None:
        _server_config = ServerConfiguration.get_server_config()
    return _server_config
set_filter_project_scope(filter_model: ProjectScopedFilter, project_name_or_id: Optional[Union[UUID, str]] = None) -> None

Set the project scope of the filter model.

Parameters:

Name Type Description Default
filter_model ProjectScopedFilter

The filter model to set the scope for.

required
project_name_or_id Optional[Union[UUID, str]]

The project to set the scope for. If not provided, the project scope is determined from the request project filter or the default project, in that order.

None
Source code in src/zenml/zen_server/utils.py
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
def set_filter_project_scope(
    filter_model: ProjectScopedFilter,
    project_name_or_id: Optional[Union[UUID, str]] = None,
) -> None:
    """Set the project scope of the filter model.

    Args:
        filter_model: The filter model to set the scope for.
        project_name_or_id: The project to set the scope for. If not
            provided, the project scope is determined from the request
            project filter or the default project, in that order.
    """
    zen_store().set_filter_project_id(
        filter_model=filter_model,
        project_name_or_id=project_name_or_id,
    )
verify_admin_status_if_no_rbac(admin_status: Optional[bool], action: Optional[str] = None) -> None

Validate the admin status for sensitive requests.

Only add this check in endpoints meant for admin use only.

Parameters:

Name Type Description Default
admin_status Optional[bool]

Whether the user is an admin or not. This is only used if explicitly specified in the call and even if passed will be ignored, if RBAC is enabled.

required
action Optional[str]

The action that is being performed, used for output only.

None

Raises:

Type Description
IllegalOperationError

If the admin status is not valid.

Source code in src/zenml/zen_server/utils.py
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
def verify_admin_status_if_no_rbac(
    admin_status: Optional[bool],
    action: Optional[str] = None,
) -> None:
    """Validate the admin status for sensitive requests.

    Only add this check in endpoints meant for admin use only.

    Args:
        admin_status: Whether the user is an admin or not. This is only used
            if explicitly specified in the call and even if passed will be
            ignored, if RBAC is enabled.
        action: The action that is being performed, used for output only.

    Raises:
        IllegalOperationError: If the admin status is not valid.
    """
    if not server_config().rbac_enabled:
        if not action:
            action = "this action"
        else:
            action = f"`{action.strip('`')}`"

        if admin_status is False:
            raise IllegalOperationError(
                message=f"Only admin users can perform {action} "
                "without RBAC enabled.",
            )
    return
workload_manager() -> WorkloadManagerInterface

Return the initialized workload manager component.

Raises:

Type Description
RuntimeError

If the workload manager component is not initialized.

Returns:

Type Description
WorkloadManagerInterface

The workload manager component.

Source code in src/zenml/zen_server/utils.py
162
163
164
165
166
167
168
169
170
171
172
173
174
def workload_manager() -> WorkloadManagerInterface:
    """Return the initialized workload manager component.

    Raises:
        RuntimeError: If the workload manager component is not initialized.

    Returns:
        The workload manager component.
    """
    global _workload_manager
    if _workload_manager is None:
        raise RuntimeError("Workload manager component not initialized")
    return _workload_manager
zen_store() -> SqlZenStore

Initialize the ZenML Store.

Returns:

Type Description
SqlZenStore

The ZenML Store.

Raises:

Type Description
RuntimeError

If the ZenML Store has not been initialized.

Source code in src/zenml/zen_server/utils.py
74
75
76
77
78
79
80
81
82
83
84
85
86
def zen_store() -> "SqlZenStore":
    """Initialize the ZenML Store.

    Returns:
        The ZenML Store.

    Raises:
        RuntimeError: If the ZenML Store has not been initialized.
    """
    global _zen_store
    if _zen_store is None:
        raise RuntimeError("ZenML Store not initialized")
    return _zen_store

zen_server_api

Zen Server API.

To run this file locally, execute:

```
uvicorn zenml.zen_server.zen_server_api:app --reload
```
Classes
RequestBodyLimit(app: ASGIApp, max_bytes: int)

Bases: BaseHTTPMiddleware

Limits the size of the request body.

Limits the size of the request body.

Parameters:

Name Type Description Default
app ASGIApp

The FastAPI app.

required
max_bytes int

The maximum size of the request body.

required
Source code in src/zenml/zen_server/zen_server_api.py
162
163
164
165
166
167
168
169
170
def __init__(self, app: ASGIApp, max_bytes: int) -> None:
    """Limits the size of the request body.

    Args:
        app: The FastAPI app.
        max_bytes: The maximum size of the request body.
    """
    super().__init__(app)
    self.max_bytes = max_bytes
Functions
dispatch(request: Request, call_next: RequestResponseEndpoint) -> Response async

Limits the size of the request body.

Parameters:

Name Type Description Default
request Request

The incoming request.

required
call_next RequestResponseEndpoint

The next function to be called.

required

Returns:

Type Description
Response

The response to the request.

Source code in src/zenml/zen_server/zen_server_api.py
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
async def dispatch(
    self, request: Request, call_next: RequestResponseEndpoint
) -> Response:
    """Limits the size of the request body.

    Args:
        request: The incoming request.
        call_next: The next function to be called.

    Returns:
        The response to the request.
    """
    if content_length := request.headers.get("content-length"):
        if int(content_length) > self.max_bytes:
            return Response(status_code=413)  # Request Entity Too Large

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )
RestrictFileUploadsMiddleware(app: FastAPI, allowed_paths: Set[str])

Bases: BaseHTTPMiddleware

Restrict file uploads to certain paths.

Restrict file uploads to certain paths.

Parameters:

Name Type Description Default
app FastAPI

The FastAPI app.

required
allowed_paths Set[str]

The allowed paths.

required
Source code in src/zenml/zen_server/zen_server_api.py
201
202
203
204
205
206
207
208
209
def __init__(self, app: FastAPI, allowed_paths: Set[str]):
    """Restrict file uploads to certain paths.

    Args:
        app: The FastAPI app.
        allowed_paths: The allowed paths.
    """
    super().__init__(app)
    self.allowed_paths = allowed_paths
Functions
dispatch(request: Request, call_next: RequestResponseEndpoint) -> Response async

Restrict file uploads to certain paths.

Parameters:

Name Type Description Default
request Request

The incoming request.

required
call_next RequestResponseEndpoint

The next function to be called.

required

Returns:

Type Description
Response

The response to the request.

Source code in src/zenml/zen_server/zen_server_api.py
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
async def dispatch(
    self, request: Request, call_next: RequestResponseEndpoint
) -> Response:
    """Restrict file uploads to certain paths.

    Args:
        request: The incoming request.
        call_next: The next function to be called.

    Returns:
        The response to the request.
    """
    if request.method == "POST":
        content_type = request.headers.get("content-type", "")
        if (
            "multipart/form-data" in content_type
            and request.url.path not in self.allowed_paths
        ):
            return JSONResponse(
                status_code=403,
                content={
                    "detail": "File uploads are not allowed on this endpoint."
                },
            )

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )
Functions
catch_all(request: Request, file_path: str) -> Any async

Dashboard endpoint.

Parameters:

Name Type Description Default
request Request

Request object.

required
file_path str

Path to a file in the dashboard root folder.

required

Returns:

Type Description
Any

The ZenML dashboard.

Source code in src/zenml/zen_server/zen_server_api.py
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
@app.get("/{file_path:path}", include_in_schema=False)
async def catch_all(request: Request, file_path: str) -> Any:
    """Dashboard endpoint.

    Args:
        request: Request object.
        file_path: Path to a file in the dashboard root folder.

    Returns:
        The ZenML dashboard.
    """
    if DASHBOARD_REDIRECT_URL:
        return RedirectResponse(url=DASHBOARD_REDIRECT_URL)
    # some static files need to be served directly from the root dashboard
    # directory
    if file_path and file_path in root_static_files:
        logger.debug(f"Returning static file: {file_path}")
        full_path = os.path.join(relative_path(DASHBOARD_DIRECTORY), file_path)
        return FileResponse(full_path)

    # everything else is directed to the index.html file that hosts the
    # single-page application
    return templates.TemplateResponse("index.html", {"request": request})
dashboard(request: Request) -> Any async

Dashboard endpoint.

Parameters:

Name Type Description Default
request Request

Request object.

required

Returns:

Type Description
Any

The ZenML dashboard.

Raises:

Type Description
HTTPException

If the dashboard files are not included.

Source code in src/zenml/zen_server/zen_server_api.py
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
@app.get("/", include_in_schema=False)
async def dashboard(request: Request) -> Any:
    """Dashboard endpoint.

    Args:
        request: Request object.

    Returns:
        The ZenML dashboard.

    Raises:
        HTTPException: If the dashboard files are not included.
    """
    if DASHBOARD_REDIRECT_URL:
        return RedirectResponse(url=DASHBOARD_REDIRECT_URL)

    if not os.path.isfile(
        os.path.join(relative_path(DASHBOARD_DIRECTORY), "index.html")
    ):
        raise HTTPException(status_code=404)
    return templates.TemplateResponse("index.html", {"request": request})
get_root_static_files() -> List[str]

Get the list of static files in the root dashboard directory.

These files are static files that are not in the /static subdirectory that need to be served as static files under the root URL path.

Returns:

Type Description
List[str]

List of static files in the root directory.

Source code in src/zenml/zen_server/zen_server_api.py
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
def get_root_static_files() -> List[str]:
    """Get the list of static files in the root dashboard directory.

    These files are static files that are not in the /static subdirectory
    that need to be served as static files under the root URL path.

    Returns:
        List of static files in the root directory.
    """
    root_path = relative_path(DASHBOARD_DIRECTORY)
    if not os.path.isdir(root_path):
        return []
    files = []
    for file in os.listdir(root_path):
        if file == "index.html":
            # this is served separately
            continue
        if isfile(os.path.join(root_path, file)):
            files.append(file)
    return files
health() -> str async

Get health status of the server.

Returns:

Type Description
str

String representing the health status of the server.

Source code in src/zenml/zen_server/zen_server_api.py
422
423
424
425
426
427
428
429
430
@app.head(HEALTH, include_in_schema=False)
@app.get(HEALTH)
async def health() -> str:
    """Get health status of the server.

    Returns:
        String representing the health status of the server.
    """
    return "OK"
infer_source_context(request: Request, call_next: Any) -> Any async

A middleware to track the source of an event.

It extracts the source context from the header of incoming requests and applies it to the ZenML source context on the API side. This way, the outgoing analytics request can append it as an additional field.

Parameters:

Name Type Description Default
request Request

the incoming request object.

required
call_next Any

a function that will receive the request as a parameter and pass it to the corresponding path operation.

required

Returns:

Type Description
Any

the response to the request.

Source code in src/zenml/zen_server/zen_server_api.py
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
@app.middleware("http")
async def infer_source_context(request: Request, call_next: Any) -> Any:
    """A middleware to track the source of an event.

    It extracts the source context from the header of incoming requests
    and applies it to the ZenML source context on the API side. This way, the
    outgoing analytics request can append it as an additional field.

    Args:
        request: the incoming request object.
        call_next: a function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        the response to the request.
    """
    try:
        s = request.headers.get(
            source_context.name,
            default=SourceContextTypes.API.value,
        )
        source_context.set(SourceContextTypes(s))
    except Exception as e:
        logger.warning(
            f"An unexpected error occurred while getting the source "
            f"context: {e}"
        )
        source_context.set(SourceContextTypes.API)

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )
initialize() -> None

Initialize the ZenML server.

Source code in src/zenml/zen_server/zen_server_api.py
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
@app.on_event("startup")
def initialize() -> None:
    """Initialize the ZenML server."""
    cfg = server_config()
    # Set the maximum number of worker threads
    to_thread.current_default_thread_limiter().total_tokens = (
        cfg.thread_pool_size
    )
    # IMPORTANT: these need to be run before the fastapi app starts, to avoid
    # race conditions
    initialize_zen_store()
    initialize_rbac()
    initialize_feature_gate()
    initialize_workload_manager()
    initialize_plugins()
    initialize_secure_headers()
    initialize_memcache(cfg.memcache_max_capacity, cfg.memcache_default_expiry)
    if cfg.deployment_type == ServerDeploymentType.CLOUD:
        # Send a workspace status update to the Cloud API to indicate that the
        # ZenML server is running or to update the version and server URL.
        send_pro_workspace_status_update()
invalid_api(invalid_api_path: str) -> None async

Invalid API endpoint.

All API endpoints that are not defined in the API routers will be redirected to this endpoint and will return a 404 error.

Parameters:

Name Type Description Default
invalid_api_path str

Invalid API path.

required

Raises:

Type Description
HTTPException

404 error.

Source code in src/zenml/zen_server/zen_server_api.py
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
@app.get(
    API + "/{invalid_api_path:path}", status_code=404, include_in_schema=False
)
async def invalid_api(invalid_api_path: str) -> None:
    """Invalid API endpoint.

    All API endpoints that are not defined in the API routers will be
    redirected to this endpoint and will return a 404 error.

    Args:
        invalid_api_path: Invalid API path.

    Raises:
        HTTPException: 404 error.
    """
    logger.debug(f"Invalid API path requested: {invalid_api_path}")
    raise HTTPException(status_code=404)
relative_path(rel: str) -> str

Get the absolute path of a path relative to the ZenML server module.

Parameters:

Name Type Description Default
rel str

Relative path.

required

Returns:

Type Description
str

Absolute path.

Source code in src/zenml/zen_server/zen_server_api.py
115
116
117
118
119
120
121
122
123
124
def relative_path(rel: str) -> str:
    """Get the absolute path of a path relative to the ZenML server module.

    Args:
        rel: Relative path.

    Returns:
        Absolute path.
    """
    return os.path.join(os.path.dirname(__file__), rel)
set_secure_headers(request: Request, call_next: Any) -> Any async

Middleware to set secure headers.

Parameters:

Name Type Description Default
request Request

The incoming request.

required
call_next Any

The next function to be called.

required

Returns:

Type Description
Any

The response with secure headers set.

Source code in src/zenml/zen_server/zen_server_api.py
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
@app.middleware("http")
async def set_secure_headers(request: Request, call_next: Any) -> Any:
    """Middleware to set secure headers.

    Args:
        request: The incoming request.
        call_next: The next function to be called.

    Returns:
        The response with secure headers set.
    """
    try:
        response = await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        response = JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )

    # If the request is for the openAPI docs, don't set secure headers
    if request.url.path.startswith("/docs") or request.url.path.startswith(
        "/redoc"
    ):
        return response

    secure_headers().framework.fastapi(response)
    return response
track_last_user_activity(request: Request, call_next: Any) -> Any async

A middleware to track last user activity.

This middleware checks if the incoming request is a user request and updates the last activity timestamp if it is.

Parameters:

Name Type Description Default
request Request

The incoming request object.

required
call_next Any

A function that will receive the request as a parameter and pass it to the corresponding path operation.

required

Returns:

Type Description
Any

The response to the request.

Source code in src/zenml/zen_server/zen_server_api.py
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
@app.middleware("http")
async def track_last_user_activity(request: Request, call_next: Any) -> Any:
    """A middleware to track last user activity.

    This middleware checks if the incoming request is a user request and
    updates the last activity timestamp if it is.

    Args:
        request: The incoming request object.
        call_next: A function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        The response to the request.
    """
    global last_user_activity
    global last_user_activity_reported

    now = utc_now()

    try:
        if is_user_request(request):
            last_user_activity = now
    except Exception as e:
        logger.debug(
            f"An unexpected error occurred while checking user activity: {e}"
        )
    if (
        (now - last_user_activity_reported).total_seconds()
        > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
    ):
        last_user_activity_reported = now
        zen_store()._update_last_user_activity_timestamp(
            last_user_activity=last_user_activity
        )

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )
validation_exception_handler(request: Any, exc: Exception) -> ORJSONResponse

Custom validation exception handler.

Parameters:

Name Type Description Default
request Any

The request.

required
exc Exception

The exception.

required

Returns:

Type Description
ORJSONResponse

The error response formatted using the ZenML API conventions.

Source code in src/zenml/zen_server/zen_server_api.py
143
144
145
146
147
148
149
150
151
152
153
154
155
156
@app.exception_handler(RequestValidationError)
def validation_exception_handler(
    request: Any, exc: Exception
) -> ORJSONResponse:
    """Custom validation exception handler.

    Args:
        request: The request.
        exc: The exception.

    Returns:
        The error response formatted using the ZenML API conventions.
    """
    return ORJSONResponse(error_detail(exc, ValueError), status_code=422)
Modules