zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250115__py3-none-any.whl

zenml/zen_stores/rest_zen_store.py

@@ -12,9 +12,11 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """REST Zen Store implementation."""
+
 import os
 import re
-from pathlib import Path, PurePath
+from datetime import datetime
+from pathlib import Path
 from typing import (
     Any,
     ClassVar,
@@ -31,14 +33,23 @@ from uuid import UUID
 
 import requests
 import urllib3
-from pydantic import BaseModel, root_validator, validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationError,
+    field_validator,
+    model_validator,
+)
 from requests.adapters import HTTPAdapter, Retry
 
 import zenml
 from zenml.analytics import source_context
 from zenml.config.global_config import GlobalConfiguration
+from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
+    ACTIONS,
     API,
     API_KEY_ROTATE,
     API_KEYS,
@@ -46,13 +57,17 @@ from zenml.constants import (
     ARTIFACT_VERSIONS,
     ARTIFACT_VISUALIZATIONS,
     ARTIFACTS,
+    BATCH,
     CODE_REFERENCES,
     CODE_REPOSITORIES,
+    CONFIG,
     CURRENT_USER,
+    DEACTIVATE,
     DEFAULT_HTTP_TIMEOUT,
     DEVICES,
     DISABLE_CLIENT_SERVER_MISMATCH_WARNING,
     ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING,
+    EVENT_SOURCES,
     FLAVORS,
     GET_OR_CREATE,
     INFO,
@@ -66,36 +81,58 @@ from zenml.constants import (
     PIPELINE_DEPLOYMENTS,
     PIPELINES,
     RUN_METADATA,
+    RUN_TEMPLATES,
     RUNS,
     SCHEDULES,
     SECRETS,
     SECRETS_BACKUP,
     SECRETS_OPERATIONS,
     SECRETS_RESTORE,
+    SERVER_SETTINGS,
     SERVICE_ACCOUNTS,
     SERVICE_CONNECTOR_CLIENT,
     SERVICE_CONNECTOR_RESOURCES,
     SERVICE_CONNECTOR_TYPES,
     SERVICE_CONNECTOR_VERIFY,
+    SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
     SERVICE_CONNECTORS,
+    SERVICES,
+    STACK,
     STACK_COMPONENTS,
+    STACK_DEPLOYMENT,
     STACKS,
     STEPS,
     TAGS,
+    TRIGGER_EXECUTIONS,
+    TRIGGERS,
     USERS,
     VERSION_1,
     WORKSPACES,
 )
 from zenml.enums import (
+    APITokenType,
     OAuthGrantTypes,
+    StackDeploymentProvider,
     StoreType,
 )
 from zenml.exceptions import (
     AuthorizationException,
+    CredentialsNotValid,
+    MethodNotAllowedError,
 )
 from zenml.io import fileio
 from zenml.logger import get_logger
+from zenml.login.credentials import APIToken
+from zenml.login.credentials_store import get_credentials_store
+from zenml.login.pro.utils import (
+    get_troubleshooting_instructions,
+    is_zenml_pro_server_url,
+)
 from zenml.models import (
+    ActionFilter,
+    ActionRequest,
+    ActionResponse,
+    ActionUpdate,
     APIKeyFilter,
     APIKeyRequest,
     APIKeyResponse,
@@ -111,8 +148,8 @@ from zenml.models import (
     ArtifactVersionUpdate,
     ArtifactVisualizationResponse,
     BaseFilter,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
     CodeReferenceResponse,
     CodeRepositoryFilter,
     CodeRepositoryRequest,
@@ -122,6 +159,11 @@ from zenml.models import (
     ComponentRequest,
     ComponentResponse,
     ComponentUpdate,
+    DeployedStack,
+    EventSourceFilter,
+    EventSourceRequest,
+    EventSourceResponse,
+    EventSourceUpdate,
     FlavorFilter,
     FlavorRequest,
     FlavorResponse,
@@ -144,6 +186,7 @@ from zenml.models import (
     OAuthDeviceFilter,
     OAuthDeviceResponse,
     OAuthDeviceUpdate,
+    OAuthTokenResponse,
     Page,
     PipelineBuildFilter,
     PipelineBuildRequest,
@@ -159,9 +202,11 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
+    RunTemplateFilter,
+    RunTemplateRequest,
+    RunTemplateResponse,
+    RunTemplateUpdate,
     ScheduleFilter,
     ScheduleRequest,
     ScheduleResponse,
@@ -171,6 +216,8 @@ from zenml.models import (
     SecretResponse,
     SecretUpdate,
     ServerModel,
+    ServerSettingsResponse,
+    ServerSettingsUpdate,
     ServiceAccountFilter,
     ServiceAccountRequest,
     ServiceAccountResponse,
@@ -181,6 +228,12 @@ from zenml.models import (
     ServiceConnectorResponse,
     ServiceConnectorTypeModel,
     ServiceConnectorUpdate,
+    ServiceFilter,
+    ServiceRequest,
+    ServiceResponse,
+    ServiceUpdate,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
     StackFilter,
     StackRequest,
     StackResponse,
@@ -193,6 +246,12 @@ from zenml.models import (
     TagRequest,
     TagResponse,
     TagUpdate,
+    TriggerExecutionFilter,
+    TriggerExecutionResponse,
+    TriggerFilter,
+    TriggerRequest,
+    TriggerResponse,
+    TriggerUpdate,
     UserFilter,
     UserRequest,
     UserResponse,
@@ -209,6 +268,7 @@ from zenml.service_connectors.service_connector_registry import (
 from zenml.utils.networking_utils import (
     replace_localhost_with_internal_hostname,
 )
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_server.exceptions import exception_from_response
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
@@ -219,7 +279,7 @@ Json = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
 
 
 AnyRequest = TypeVar("AnyRequest", bound=BaseRequest)
-AnyResponse = TypeVar("AnyResponse", bound=BaseResponse)  # type: ignore[type-arg]
+AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
 AnyWorkspaceScopedRequest = TypeVar(
     "AnyWorkspaceScopedRequest",
     bound=WorkspaceScopedRequest,
@@ -234,10 +294,10 @@ class RestZenStoreConfiguration(StoreConfiguration):
         username: The username to use to connect to the Zen server.
         password: The password to use to connect to the Zen server.
         api_key: The service account API key to use to connect to the Zen
-            server.
-        api_token: The API token to use to connect to the Zen server. Generated
-            by the client and stored in the configuration file on the first
-            login and every time the API key is refreshed.
+            server. This is only set if the API key is configured explicitly via
+            environment variables or the ZenML global configuration file. API
+            keys configured via the CLI are stored in the credentials store
+            instead.
         verify_ssl: Either a boolean, in which case it controls whether we
             verify the server's TLS certificate, or a string, in which case it
             must be a path to a CA bundle to use or the CA bundle value itself.
@@ -247,40 +307,13 @@ class RestZenStoreConfiguration(StoreConfiguration):
 
     type: StoreType = StoreType.REST
 
-    username: Optional[str] = None
-    password: Optional[str] = None
-    api_key: Optional[str] = None
-    api_token: Optional[str] = None
-    verify_ssl: Union[bool, str] = True
+    verify_ssl: Union[bool, str] = Field(
+        default=True, union_mode="left_to_right"
+    )
     http_timeout: int = DEFAULT_HTTP_TIMEOUT
 
-    @root_validator
-    def validate_credentials(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        """Validates the credentials provided in the values dictionary.
-
-        Args:
-            values: A dictionary containing the values to be validated.
-
-        Raises:
-            ValueError: If neither api_token nor username nor api_key is set.
-
-        Returns:
-            The values dictionary.
-        """
-        # Check if the values dictionary contains either an API token, an API
-        # key or a username as non-empty strings.
-        if (
-            values.get("api_token")
-            or values.get("username")
-            or values.get("api_key")
-        ):
-            return values
-        raise ValueError(
-            "Neither api_token nor username nor api_key is set in the "
-            "store config."
-        )
-
-    @validator("url")
+    @field_validator("url")
+    @classmethod
     def validate_url(cls, url: str) -> str:
         """Validates that the URL is a well-formed REST store URL.
 
@@ -308,7 +341,8 @@ class RestZenStoreConfiguration(StoreConfiguration):
 
         return url
 
-    @validator("verify_ssl")
+    @field_validator("verify_ssl")
+    @classmethod
     def validate_verify_ssl(
         cls, verify_ssl: Union[bool, str]
     ) -> Union[bool, str]:
@@ -335,9 +369,10 @@ class RestZenStoreConfiguration(StoreConfiguration):
 
         fileio.makedirs(str(secret_folder))
         file_path = Path(secret_folder, "ca_bundle.pem")
-        with open(file_path, "w") as f:
+        with os.fdopen(
+            os.open(file_path, flags=os.O_RDWR | os.O_CREAT, mode=0o600), "w"
+        ) as f:
             f.write(verify_ssl)
-        file_path.chmod(0o600)
         verify_ssl = str(file_path)
 
         return verify_ssl
@@ -363,54 +398,48 @@ class RestZenStoreConfiguration(StoreConfiguration):
             with open(self.verify_ssl, "r") as f:
                 self.verify_ssl = f.read()
 
+    @model_validator(mode="before")
     @classmethod
-    def copy_configuration(
-        cls,
-        config: "StoreConfiguration",
-        config_path: str,
-        load_config_path: Optional[PurePath] = None,
-    ) -> "StoreConfiguration":
-        """Create a copy of the store config using a different path.
-
-        This method is used to create a copy of the store configuration that can
-        be loaded using a different configuration path or in the context of a
-        new environment, such as a container image.
-
-        The configuration files accompanying the store configuration are also
-        copied to the new configuration path (e.g. certificates etc.).
-
-        Args:
-            config: The store configuration to copy.
-            config_path: new path where the configuration copy will be loaded
-                from.
-            load_config_path: absolute path that will be used to load the copied
-                configuration. This can be set to a value different from
-                `config_path` if the configuration copy will be loaded from
-                a different environment, e.g. when the configuration is copied
-                to a container image and loaded using a different absolute path.
-                This will be reflected in the paths and URLs encoded in the
-                copied configuration.
-
-        Returns:
-            A new store configuration object that reflects the new configuration
-            path.
-        """
-        assert isinstance(config, RestZenStoreConfiguration)
-        assert config.api_token is not None or config.api_key is not None
-        config = config.copy(exclude={"username", "password"}, deep=True)
-        # Load the certificate values back into the configuration
-        config.expand_certificates()
-        return config
+    @before_validator_handler
+    def _move_credentials(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Moves credentials (API keys, API tokens, passwords) from the config to the credentials store.
+
+        Args:
+            data: The values dict used to instantiate the model.
 
-    class Config:
-        """Pydantic configuration class."""
+        Returns:
+            The values dict without credentials.
+        """
+        url = data.get("url")
+        if not url:
+            return data
+
+        url = replace_localhost_with_internal_hostname(url)
+
+        if api_token := data.pop("api_token", None):
+            credentials_store = get_credentials_store()
+            credentials_store.set_bare_token(url, api_token)
 
+        username = data.pop("username", None)
+        password = data.pop("password", None)
+        if username is not None and password is not None:
+            credentials_store = get_credentials_store()
+            credentials_store.set_password(url, username, password)
+
+        if api_key := data.pop("api_key", None):
+            credentials_store = get_credentials_store()
+            credentials_store.set_api_key(url, api_key)
+
+        return data
+
+    model_config = ConfigDict(
         # Don't validate attributes when assigning them. This is necessary
         # because the `verify_ssl` attribute can be expanded to the contents
         # of the certificate file.
-        validate_assignment = False
-        # Forbid extra attributes set in the class.
-        extra = "forbid"
+        validate_assignment=False,
+        # Ignore extra attributes set in the class.
+        extra="ignore",
+    )
 
 
 class RestZenStore(BaseZenStore):
@@ -419,7 +448,7 @@ class RestZenStore(BaseZenStore):
     config: RestZenStoreConfiguration
     TYPE: ClassVar[StoreType] = StoreType.REST
     CONFIG_TYPE: ClassVar[Type[StoreConfiguration]] = RestZenStoreConfiguration
-    _api_token: Optional[str] = None
+    _api_token: Optional[APIToken] = None
     _session: Optional[requests.Session] = None
 
     # ====================================
@@ -431,9 +460,54 @@ class RestZenStore(BaseZenStore):
     # --------------------------------
 
     def _initialize(self) -> None:
-        """Initialize the REST store."""
-        client_version = zenml.__version__
-        server_version = self.get_store_info().version
+        """Initialize the REST store.
+
+        Raises:
+            RuntimeError: If the store cannot be initialized.
+            AuthorizationException: If the store cannot be initialized due to
+                authentication errors.
+        """
+        try:
+            client_version = zenml.__version__
+            server_version = self.get_store_info().version
+
+        # Handle cases where the ZenML server is not available
+        except ConnectionError as e:
+            error_message = (
+                f"Cannot connect to the ZenML server at {self.url}."
+            )
+            if urlparse(self.url).hostname in [
+                "localhost",
+                "127.0.0.1",
+                "host.docker.internal",
+            ]:
+                recommendation = (
+                    "Please run `zenml login --local --restart` to restart the "
+                    "server."
+                )
+            else:
+                recommendation = (
+                    f"Please run `zenml login {self.url}` to reconnect to the "
+                    "server."
+                )
+            raise RuntimeError(f"{error_message}\n{recommendation}") from e
+
+        except AuthorizationException as e:
+            raise AuthorizationException(
+                f"Authorization failed for store at '{self.url}'. Please check "
+                f"your credentials: {str(e)}"
+            )
+
+        except Exception as e:
+            zenml_pro_extra = ""
+            if is_zenml_pro_server_url(self.url):
+                zenml_pro_extra = (
+                    "\nHINT: " + get_troubleshooting_instructions(self.url)
+                )
+            raise RuntimeError(
+                f"Error connecting to URL "
+                f"'{self.url}': {str(e)}" + zenml_pro_extra
+            ) from e
 
         if not DISABLE_CLIENT_SERVER_MISMATCH_WARNING and (
             server_version != client_version
@@ -455,7 +529,7 @@ class RestZenStore(BaseZenStore):
             Information about the server.
         """
         body = self.get(INFO)
-        return ServerModel.parse_obj(body)
+        return ServerModel.model_validate(body)
 
     def get_deployment_id(self) -> UUID:
         """Get the ID of the deployment.
@@ -465,6 +539,131 @@ class RestZenStore(BaseZenStore):
         """
         return self.get_store_info().id
 
+    # -------------------- Server Settings --------------------
+
+    def get_server_settings(
+        self, hydrate: bool = True
+    ) -> ServerSettingsResponse:
+        """Get the server settings.
+
+        Args:
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The server settings.
+        """
+        response_body = self.get(SERVER_SETTINGS, params={"hydrate": hydrate})
+        return ServerSettingsResponse.model_validate(response_body)
+
+    def update_server_settings(
+        self, settings_update: ServerSettingsUpdate
+    ) -> ServerSettingsResponse:
+        """Update the server settings.
+
+        Args:
+            settings_update: The server settings update.
+
+        Returns:
+            The updated server settings.
+        """
+        response_body = self.put(SERVER_SETTINGS, body=settings_update)
+        return ServerSettingsResponse.model_validate(response_body)
+
+    # -------------------- Actions  --------------------
+
+    def create_action(self, action: ActionRequest) -> ActionResponse:
+        """Create an action.
+
+        Args:
+            action: The action to create.
+
+        Returns:
+            The created action.
+        """
+        return self._create_resource(
+            resource=action,
+            route=ACTIONS,
+            response_model=ActionResponse,
+        )
+
+    def get_action(
+        self,
+        action_id: UUID,
+        hydrate: bool = True,
+    ) -> ActionResponse:
+        """Get an action by ID.
+
+        Args:
+            action_id: The ID of the action to get.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The action.
+        """
+        return self._get_resource(
+            resource_id=action_id,
+            route=ACTIONS,
+            response_model=ActionResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_actions(
+        self,
+        action_filter_model: ActionFilter,
+        hydrate: bool = False,
+    ) -> Page[ActionResponse]:
+        """List all actions matching the given filter criteria.
+
+        Args:
+            action_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all actions matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=ACTIONS,
+            response_model=ActionResponse,
+            filter_model=action_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_action(
+        self,
+        action_id: UUID,
+        action_update: ActionUpdate,
+    ) -> ActionResponse:
+        """Update an existing action.
+
+        Args:
+            action_id: The ID of the action to update.
+            action_update: The update to be applied to the action.
+
+        Returns:
+            The updated action.
+        """
+        return self._update_resource(
+            resource_id=action_id,
+            resource_update=action_update,
+            route=ACTIONS,
+            response_model=ActionResponse,
+        )
+
+    def delete_action(self, action_id: UUID) -> None:
+        """Delete an action.
+
+        Args:
+            action_id: The ID of the action to delete.
+        """
+        self._delete_resource(
+            resource_id=action_id,
+            route=ACTIONS,
+        )
+
     # ----------------------------- API Keys -----------------------------
 
     def create_api_key(
@@ -511,16 +710,6 @@ class RestZenStore(BaseZenStore):
             params={"hydrate": hydrate},
         )
 
-    def set_api_key(self, api_key: str) -> None:
-        """Set the API key to use for authentication.
-
-        Args:
-            api_key: The API key to use for authentication.
-        """
-        self.config.api_key = api_key
-        self.clear_session()
-        GlobalConfiguration()._write_config()
-
     def list_api_keys(
         self,
         service_account_id: UUID,
@@ -592,7 +781,7 @@ class RestZenStore(BaseZenStore):
             f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}/{str(api_key_name_or_id)}{API_KEY_ROTATE}",
             body=rotate_request,
         )
-        return APIKeyResponse.parse_obj(response_body)
+        return APIKeyResponse.model_validate(response_body)
 
     def delete_api_key(
         self,
@@ -611,6 +800,93 @@ class RestZenStore(BaseZenStore):
             route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
         )
 
+    # ----------------------------- Services -----------------------------
+
+    def create_service(
+        self, service_request: ServiceRequest
+    ) -> ServiceResponse:
+        """Create a new service.
+
+        Args:
+            service_request: The service to create.
+
+        Returns:
+            The created service.
+        """
+        return self._create_resource(
+            resource=service_request,
+            response_model=ServiceResponse,
+            route=SERVICES,
+        )
+
+    def get_service(
+        self, service_id: UUID, hydrate: bool = True
+    ) -> ServiceResponse:
+        """Get a service.
+
+        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:
+            The service.
+        """
+        return self._get_resource(
+            resource_id=service_id,
+            route=SERVICES,
+            response_model=ServiceResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_services(
+        self, filter_model: ServiceFilter, hydrate: bool = False
+    ) -> Page[ServiceResponse]:
+        """List all services matching the given filter criteria.
+
+        Args:
+            filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all services matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=SERVICES,
+            response_model=ServiceResponse,
+            filter_model=filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_service(
+        self, service_id: UUID, update: ServiceUpdate
+    ) -> ServiceResponse:
+        """Update a service.
+
+        Args:
+            service_id: The ID of the service to update.
+            update: The update to be applied to the service.
+
+        Returns:
+            The updated service.
+        """
+        return self._update_resource(
+            resource_id=service_id,
+            resource_update=update,
+            response_model=ServiceResponse,
+            route=SERVICES,
+        )
+
+    def delete_service(self, service_id: UUID) -> None:
+        """Delete a service.
+
+        Args:
+            service_id: The ID of the service to delete.
+        """
+        self._delete_resource(resource_id=service_id, route=SERVICES)
+
     # ----------------------------- Artifacts -----------------------------
 
     def create_artifact(self, artifact: ArtifactRequest) -> ArtifactResponse:
@@ -715,6 +991,23 @@ class RestZenStore(BaseZenStore):
             route=ARTIFACT_VERSIONS,
         )
 
+    def batch_create_artifact_versions(
+        self, artifact_versions: List[ArtifactVersionRequest]
+    ) -> List[ArtifactVersionResponse]:
+        """Creates a batch of artifact versions.
+
+        Args:
+            artifact_versions: The artifact versions to create.
+
+        Returns:
+            The created artifact versions.
+        """
+        return self._batch_create_resources(
+            resources=artifact_versions,
+            response_model=ArtifactVersionResponse,
+            route=ARTIFACT_VERSIONS,
+        )
+
     def get_artifact_version(
         self, artifact_version_id: UUID, hydrate: bool = True
     ) -> ArtifactVersionResponse:
@@ -1308,10 +1601,6 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_BUILDS,
         )
 
-        # ----------------------
-        # Pipeline Deployments
-        # ----------------------
-
     # -------------------------- Pipeline Deployments --------------------------
 
     def create_deployment(
@@ -1386,6 +1675,228 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_DEPLOYMENTS,
         )
 
+    # -------------------- Run templates --------------------
+
+    def create_run_template(
+        self,
+        template: RunTemplateRequest,
+    ) -> RunTemplateResponse:
+        """Create a new run template.
+
+        Args:
+            template: The template to create.
+
+        Returns:
+            The newly created template.
+        """
+        return self._create_workspace_scoped_resource(
+            resource=template,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+        )
+
+    def get_run_template(
+        self, template_id: UUID, hydrate: bool = True
+    ) -> RunTemplateResponse:
+        """Get a run template with a given ID.
+
+        Args:
+            template_id: ID of the template.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The template.
+        """
+        return self._get_resource(
+            resource_id=template_id,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_run_templates(
+        self,
+        template_filter_model: RunTemplateFilter,
+        hydrate: bool = False,
+    ) -> Page[RunTemplateResponse]:
+        """List all run templates matching the given filter criteria.
+
+        Args:
+            template_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all templates matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+            filter_model=template_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_run_template(
+        self,
+        template_id: UUID,
+        template_update: RunTemplateUpdate,
+    ) -> RunTemplateResponse:
+        """Updates a run template.
+
+        Args:
+            template_id: The ID of the template to update.
+            template_update: The update to apply.
+
+        Returns:
+            The updated template.
+        """
+        return self._update_resource(
+            resource_id=template_id,
+            resource_update=template_update,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+        )
+
+    def delete_run_template(self, template_id: UUID) -> None:
+        """Delete a run template.
+
+        Args:
+            template_id: The ID of the template to delete.
+        """
+        self._delete_resource(
+            resource_id=template_id,
+            route=RUN_TEMPLATES,
+        )
+
+    def run_template(
+        self,
+        template_id: UUID,
+        run_configuration: Optional[PipelineRunConfiguration] = None,
+    ) -> PipelineRunResponse:
+        """Run a template.
+
+        Args:
+            template_id: The ID of the template to run.
+            run_configuration: Configuration for the run.
+
+        Raises:
+            RuntimeError: If the server does not support running a template.
+
+        Returns:
+            Model of the pipeline run.
+        """
+        run_configuration = run_configuration or PipelineRunConfiguration()
+
+        try:
+            response_body = self.post(
+                f"{RUN_TEMPLATES}/{template_id}/runs",
+                body=run_configuration,
+            )
+        except MethodNotAllowedError as e:
+            raise RuntimeError(
+                "Running a template is not supported for this server."
+            ) from e
+
+        return PipelineRunResponse.model_validate(response_body)
+
+    # -------------------- Event Sources  --------------------
+
+    def create_event_source(
+        self, event_source: EventSourceRequest
+    ) -> EventSourceResponse:
+        """Create an event_source.
+
+        Args:
+            event_source: The event_source to create.
+
+        Returns:
+            The created event_source.
+        """
+        return self._create_resource(
+            resource=event_source,
+            route=EVENT_SOURCES,
+            response_model=EventSourceResponse,
+        )
+
+    def get_event_source(
+        self,
+        event_source_id: UUID,
+        hydrate: bool = True,
+    ) -> EventSourceResponse:
+        """Get an event_source by ID.
+
+        Args:
+            event_source_id: The ID of the event_source to get.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The event_source.
+        """
+        return self._get_resource(
+            resource_id=event_source_id,
+            route=EVENT_SOURCES,
+            response_model=EventSourceResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_event_sources(
+        self,
+        event_source_filter_model: EventSourceFilter,
+        hydrate: bool = False,
+    ) -> Page[EventSourceResponse]:
+        """List all event_sources matching the given filter criteria.
+
+        Args:
+            event_source_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all event_sources matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=EVENT_SOURCES,
+            response_model=EventSourceResponse,
+            filter_model=event_source_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_event_source(
+        self,
+        event_source_id: UUID,
+        event_source_update: EventSourceUpdate,
+    ) -> EventSourceResponse:
+        """Update an existing event_source.
+
+        Args:
+            event_source_id: The ID of the event_source to update.
+            event_source_update: The update to be applied to the event_source.
+
+        Returns:
+            The updated event_source.
+        """
+        return self._update_resource(
+            resource_id=event_source_id,
+            resource_update=event_source_update,
+            route=EVENT_SOURCES,
+            response_model=EventSourceResponse,
+        )
+
+    def delete_event_source(self, event_source_id: UUID) -> None:
+        """Delete an event_source.
+
+        Args:
+            event_source_id: The ID of the event_source to delete.
+        """
+        self._delete_resource(
+            resource_id=event_source_id,
+            route=EVENT_SOURCES,
+        )
+
     # ----------------------------- Pipeline runs -----------------------------
 
     def create_run(
@@ -1502,9 +2013,7 @@ class RestZenStore(BaseZenStore):
 
     # ----------------------------- Run Metadata -----------------------------
 
-    def create_run_metadata(
-        self, run_metadata: RunMetadataRequest
-    ) -> List[RunMetadataResponse]:
+    def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
         """Creates run metadata.
 
         Args:
@@ -1514,55 +2023,8 @@ class RestZenStore(BaseZenStore):
             The created run metadata.
         """
         route = f"{WORKSPACES}/{str(run_metadata.workspace)}{RUN_METADATA}"
-        response_body = self.post(f"{route}", body=run_metadata)
-        result: List[RunMetadataResponse] = []
-        if isinstance(response_body, list):
-            for metadata in response_body or []:
-                result.append(RunMetadataResponse.parse_obj(metadata))
-        return result
-
-    def get_run_metadata(
-        self, run_metadata_id: UUID, hydrate: bool = True
-    ) -> RunMetadataResponse:
-        """Gets run metadata with the given ID.
-
-        Args:
-            run_metadata_id: The ID of the run metadata to get.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._get_resource(
-            resource_id=run_metadata_id,
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            params={"hydrate": hydrate},
-        )
-
-    def list_run_metadata(
-        self,
-        run_metadata_filter_model: RunMetadataFilter,
-        hydrate: bool = False,
-    ) -> Page[RunMetadataResponse]:
-        """List run metadata.
-
-        Args:
-            run_metadata_filter_model: All filter parameters including
-                pagination params.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._list_paginated_resources(
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            filter_model=run_metadata_filter_model,
-            params={"hydrate": hydrate},
-        )
+        self.post(f"{route}", body=run_metadata)
+        return None
 
     # ----------------------------- Schedules -----------------------------
 
@@ -2110,9 +2572,15 @@ class RestZenStore(BaseZenStore):
             f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}",
             body=service_connector,
             params={"list_resources": list_resources},
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
-        resources = ServiceConnectorResourcesModel.parse_obj(response_body)
+        resources = ServiceConnectorResourcesModel.model_validate(
+            response_body
+        )
         self._populate_connector_type(resources)
         return resources
 
@@ -2145,9 +2613,15 @@ class RestZenStore(BaseZenStore):
         response_body = self.put(
             f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}",
             params=params,
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
-        resources = ServiceConnectorResourcesModel.parse_obj(response_body)
+        resources = ServiceConnectorResourcesModel.model_validate(
+            response_body
+        )
         self._populate_connector_type(resources)
         return resources
 
@@ -2178,7 +2652,7 @@ class RestZenStore(BaseZenStore):
             params=params,
         )
 
-        connector = ServiceConnectorResponse.parse_obj(response_body)
+        connector = ServiceConnectorResponse.model_validate(response_body)
         self._populate_connector_type(connector)
         return connector
 
@@ -2211,11 +2685,15 @@ class RestZenStore(BaseZenStore):
         response_body = self.get(
             f"{WORKSPACES}/{workspace_name_or_id}{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_RESOURCES}",
             params=params,
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         assert isinstance(response_body, list)
         resource_list = [
-            ServiceConnectorResourcesModel.parse_obj(item)
+            ServiceConnectorResourcesModel.model_validate(item)
             for item in response_body
         ]
 
@@ -2248,7 +2726,7 @@ class RestZenStore(BaseZenStore):
                 )
             except (ValueError, AuthorizationException) as e:
                 logger.error(
-                    f'Failed to fetch {resource_type or "available"} '
+                    f"Failed to fetch {resource_type or 'available'} "
                     f"resources from service connector {connector.name}/"
                     f"{connector.id}: {e}"
                 )
@@ -2288,7 +2766,8 @@ class RestZenStore(BaseZenStore):
 
         assert isinstance(response_body, list)
         remote_connector_types = [
-            ServiceConnectorTypeModel.parse_obj(item) for item in response_body
+            ServiceConnectorTypeModel.model_validate(item)
+            for item in response_body
         ]
 
         # Mark the remote connector types as being only remotely available
@@ -2344,7 +2823,7 @@ class RestZenStore(BaseZenStore):
             response_body = self.get(
                 f"{SERVICE_CONNECTOR_TYPES}/{connector_type}",
             )
-            remote_connector_type = ServiceConnectorTypeModel.parse_obj(
+            remote_connector_type = ServiceConnectorTypeModel.model_validate(
                 response_body
             )
             if local_connector_type:
@@ -2376,10 +2855,12 @@ class RestZenStore(BaseZenStore):
         Returns:
             The registered stack.
         """
-        return self._create_workspace_scoped_resource(
+        assert stack.workspace is not None
+
+        return self._create_resource(
             resource=stack,
-            route=STACKS,
             response_model=StackResponse,
+            route=f"{WORKSPACES}/{str(stack.workspace)}{STACKS}",
         )
 
     def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse:
@@ -2451,6 +2932,88 @@ class RestZenStore(BaseZenStore):
             route=STACKS,
         )
 
+    # ---------------- Stack deployments-----------------
+
+    def get_stack_deployment_info(
+        self,
+        provider: StackDeploymentProvider,
+    ) -> StackDeploymentInfo:
+        """Get information about a stack deployment provider.
+
+        Args:
+            provider: The stack deployment provider.
+
+        Returns:
+            Information about the stack deployment provider.
+        """
+        body = self.get(
+            f"{STACK_DEPLOYMENT}{INFO}",
+            params={"provider": provider.value},
+        )
+        return StackDeploymentInfo.model_validate(body)
+
+    def get_stack_deployment_config(
+        self,
+        provider: StackDeploymentProvider,
+        stack_name: str,
+        location: Optional[str] = None,
+    ) -> StackDeploymentConfig:
+        """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.
+
+        Args:
+            provider: The stack deployment provider.
+            stack_name: The name of the stack.
+            location: The location where the stack should be deployed.
+
+        Returns:
+            The cloud provider console URL and configuration needed to deploy
+            the ZenML stack to the specified cloud provider.
+        """
+        params = {
+            "provider": provider.value,
+            "stack_name": stack_name,
+        }
+        if location:
+            params["location"] = location
+        body = self.get(f"{STACK_DEPLOYMENT}{CONFIG}", params=params)
+        return StackDeploymentConfig.model_validate(body)
+
+    def get_stack_deployment_stack(
+        self,
+        provider: StackDeploymentProvider,
+        stack_name: str,
+        location: Optional[str] = None,
+        date_start: Optional[datetime] = None,
+    ) -> 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.
+
+        Returns:
+            The ZenML stack that was deployed and registered or None if the
+            stack was not found.
+        """
+        params = {
+            "provider": provider.value,
+            "stack_name": stack_name,
+        }
+        if location:
+            params["location"] = location
+        if date_start:
+            params["date_start"] = str(date_start)
+        body = self.get(
+            f"{STACK_DEPLOYMENT}{STACK}",
+            params=params,
+        )
+        if body:
+            return DeployedStack.model_validate(body)
+
+        return None
+
     # ----------------------------- Step runs -----------------------------
 
     def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
@@ -2532,6 +3095,158 @@ class RestZenStore(BaseZenStore):
             route=STEPS,
         )
 
+    # -------------------- Triggers  --------------------
+
+    def create_trigger(self, trigger: TriggerRequest) -> TriggerResponse:
+        """Create an trigger.
+
+        Args:
+            trigger: The trigger to create.
+
+        Returns:
+            The created trigger.
+        """
+        return self._create_resource(
+            resource=trigger,
+            route=TRIGGERS,
+            response_model=TriggerResponse,
+        )
+
+    def get_trigger(
+        self,
+        trigger_id: UUID,
+        hydrate: bool = True,
+    ) -> TriggerResponse:
+        """Get a trigger by ID.
+
+        Args:
+            trigger_id: The ID of the trigger to get.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The trigger.
+        """
+        return self._get_resource(
+            resource_id=trigger_id,
+            route=TRIGGERS,
+            response_model=TriggerResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_triggers(
+        self,
+        trigger_filter_model: TriggerFilter,
+        hydrate: bool = False,
+    ) -> Page[TriggerResponse]:
+        """List all triggers matching the given filter criteria.
+
+        Args:
+            trigger_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all triggers matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=TRIGGERS,
+            response_model=TriggerResponse,
+            filter_model=trigger_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_trigger(
+        self,
+        trigger_id: UUID,
+        trigger_update: TriggerUpdate,
+    ) -> TriggerResponse:
+        """Update an existing trigger.
+
+        Args:
+            trigger_id: The ID of the trigger to update.
+            trigger_update: The update to be applied to the trigger.
+
+        Returns:
+            The updated trigger.
+        """
+        return self._update_resource(
+            resource_id=trigger_id,
+            resource_update=trigger_update,
+            route=TRIGGERS,
+            response_model=TriggerResponse,
+        )
+
+    def delete_trigger(self, trigger_id: UUID) -> None:
+        """Delete an trigger.
+
+        Args:
+            trigger_id: The ID of the trigger to delete.
+        """
+        self._delete_resource(
+            resource_id=trigger_id,
+            route=TRIGGERS,
+        )
+
+    # -------------------- Trigger Executions --------------------
+
+    def get_trigger_execution(
+        self,
+        trigger_execution_id: UUID,
+        hydrate: bool = True,
+    ) -> TriggerExecutionResponse:
+        """Get an trigger execution by ID.
+
+        Args:
+            trigger_execution_id: The ID of the trigger execution to get.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The trigger execution.
+        """
+        return self._get_resource(
+            resource_id=trigger_execution_id,
+            route=TRIGGER_EXECUTIONS,
+            response_model=TriggerExecutionResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_trigger_executions(
+        self,
+        trigger_execution_filter_model: TriggerExecutionFilter,
+        hydrate: bool = False,
+    ) -> Page[TriggerExecutionResponse]:
+        """List all trigger executions matching the given filter criteria.
+
+        Args:
+            trigger_execution_filter_model: All filter parameters including
+                pagination params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all trigger executions matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=TRIGGER_EXECUTIONS,
+            response_model=TriggerExecutionResponse,
+            filter_model=trigger_execution_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def delete_trigger_execution(self, trigger_execution_id: UUID) -> None:
+        """Delete a trigger execution.
+
+        Args:
+            trigger_execution_id: The ID of the trigger execution to delete.
+        """
+        self._delete_resource(
+            resource_id=trigger_execution_id,
+            route=TRIGGER_EXECUTIONS,
+        )
+
     # ----------------------------- Users -----------------------------
 
     def create_user(self, user: UserRequest) -> UserResponse:
@@ -2580,7 +3295,7 @@ class RestZenStore(BaseZenStore):
             )
         else:
             body = self.get(CURRENT_USER, params={"hydrate": hydrate})
-            return UserResponse.parse_obj(body)
+            return UserResponse.model_validate(body)
 
     def list_users(
         self,
@@ -2624,6 +3339,23 @@ class RestZenStore(BaseZenStore):
             response_model=UserResponse,
         )
 
+    def deactivate_user(
+        self, user_name_or_id: Union[str, UUID]
+    ) -> UserResponse:
+        """Deactivates a user.
+
+        Args:
+            user_name_or_id: The name or ID of the user to delete.
+
+        Returns:
+            The deactivated user containing the activation token.
+        """
+        response_body = self.put(
+            f"{USERS}/{str(user_name_or_id)}{DEACTIVATE}",
+        )
+
+        return UserResponse.model_validate(response_body)
+
     def delete_user(self, user_name_or_id: Union[str, UUID]) -> None:
         """Deletes a user.
 
@@ -2940,10 +3672,10 @@ class RestZenStore(BaseZenStore):
         Returns:
             The newly created model version to artifact link.
         """
-        return self._create_workspace_scoped_resource(
+        return self._create_resource(
             resource=model_version_artifact_link,
             response_model=ModelVersionArtifactResponse,
-            route=f"{MODEL_VERSIONS}/{model_version_artifact_link.model_version}{ARTIFACTS}",
+            route=MODEL_VERSION_ARTIFACTS,
         )
 
     def list_model_version_artifact_links(
@@ -3020,10 +3752,10 @@ class RestZenStore(BaseZenStore):
             - Otherwise, returns the newly created model version to pipeline
                 run link.
         """
-        return self._create_workspace_scoped_resource(
+        return self._create_resource(
             resource=model_version_pipeline_run_link,
             response_model=ModelVersionPipelineRunResponse,
-            route=f"{MODEL_VERSIONS}/{model_version_pipeline_run_link.model_version}{RUNS}",
+            route=MODEL_VERSION_PIPELINE_RUNS,
         )
 
     def list_model_version_pipeline_run_links(
@@ -3141,17 +3873,20 @@ class RestZenStore(BaseZenStore):
 
     def get_api_token(
         self,
-        pipeline_id: Optional[UUID] = None,
+        token_type: APITokenType = APITokenType.WORKLOAD,
+        expires_in: Optional[int] = None,
         schedule_id: Optional[UUID] = None,
-        expires_minutes: Optional[int] = None,
+        pipeline_run_id: Optional[UUID] = None,
+        step_run_id: Optional[UUID] = None,
     ) -> str:
-        """Get an API token for a workload.
+        """Get an API token.
 
         Args:
-            pipeline_id: The ID of the pipeline to get a token for.
+            token_type: The type of the token to get.
+            expires_in: The time in seconds until the token expires.
             schedule_id: The ID of the schedule to get a token for.
-            expires_minutes: The number of minutes for which the token should
-                be valid. If not provided, the token will be valid indefinitely.
+            pipeline_run_id: The ID of the pipeline run to get a token for.
+            step_run_id: The ID of the step run to get a token for.
 
         Returns:
             The API token.
@@ -3159,13 +3894,17 @@ class RestZenStore(BaseZenStore):
         Raises:
             ValueError: if the server response is not valid.
         """
-        params: Dict[str, Any] = {}
-        if pipeline_id:
-            params["pipeline_id"] = pipeline_id
+        params: Dict[str, Any] = {
+            "token_type": token_type.value,
+        }
+        if expires_in:
+            params["expires_in"] = expires_in
         if schedule_id:
             params["schedule_id"] = schedule_id
-        if expires_minutes:
-            params["expires_minutes"] = expires_minutes
+        if pipeline_run_id:
+            params["pipeline_run_id"] = pipeline_run_id
+        if step_run_id:
+            params["step_run_id"] = step_run_id
         response_body = self.get(API_TOKEN, params=params)
         if not isinstance(response_body, str):
             raise ValueError(
@@ -3272,104 +4011,282 @@ class RestZenStore(BaseZenStore):
     # Internal helper methods
     # =======================
 
-    def _get_auth_token(self) -> str:
-        """Get the authentication token for the REST store.
+    def get_or_generate_api_token(self) -> str:
+        """Get or generate an API token.
 
         Returns:
-            The authentication token.
+            The API token.
 
         Raises:
-            ValueError: if the response from the server isn't in the right
-                format.
-        """
-        if self._api_token is None:
-            # Check if the API token is already stored in the config
-            if self.config.api_token:
-                self._api_token = self.config.api_token
-            # Check if the username and password are provided in the config
-            elif (
-                self.config.username is not None
-                and self.config.password is not None
-                or self.config.api_key is not None
-            ):
-                data: Optional[Dict[str, str]] = None
-                if self.config.api_key is not None:
-                    data = {
-                        "grant_type": OAuthGrantTypes.ZENML_API_KEY.value,
-                        "password": self.config.api_key,
-                    }
-                elif (
-                    self.config.username is not None
-                    and self.config.password is not None
-                ):
-                    data = {
-                        "grant_type": OAuthGrantTypes.OAUTH_PASSWORD.value,
-                        "username": self.config.username,
-                        "password": self.config.password,
-                    }
-
-                response = self._handle_response(
-                    requests.post(
-                        self.url + API + VERSION_1 + LOGIN,
-                        data=data,
-                        verify=self.config.verify_ssl,
-                        timeout=self.config.http_timeout,
-                    )
+            CredentialsNotValid: if an API token cannot be fetched or
+                generated because the client credentials are not valid.
+        """
+        if self._api_token is None or self._api_token.expired:
+            # Check if a valid API token is already in the cache
+            credentials_store = get_credentials_store()
+            credentials = credentials_store.get_credentials(self.url)
+            token = credentials.api_token if credentials else None
+            if credentials and token and not token.expired:
+                self._api_token = token
+
+                # Populate the server info in the credentials store if it is
+                # not already present
+                if not credentials.server_id:
+                    try:
+                        server_info = self.get_store_info()
+                    except Exception as e:
+                        logger.warning(f"Failed to get server info: {e}.")
+                    else:
+                        credentials_store.update_server_info(
+                            self.url, server_info
+                        )
+
+                return self._api_token.access_token
+
+            # Token is expired or not found in the cache. Time to get a new one.
+
+            if not token:
+                logger.debug(f"Authenticating to {self.url}")
+            else:
+                logger.debug(
+                    f"Authentication token for {self.url} expired; refreshing..."
                 )
-                if (
-                    not isinstance(response, dict)
-                    or "access_token" not in response
-                ):
-                    raise ValueError(
-                        f"Bad API Response. Expected access token dict, got "
-                        f"{type(response)}"
+
+            data: Optional[Dict[str, str]] = None
+
+            # Use a custom user agent to identify the ZenML client in the server
+            # logs.
+            headers: Dict[str, str] = {
+                "User-Agent": "zenml/" + zenml.__version__,
+            }
+
+            # Check if an API key is configured
+            api_key = credentials_store.get_api_key(self.url)
+
+            # Check if username and password are configured
+            username, password = credentials_store.get_password(self.url)
+
+            api_key_hint = (
+                "\nHint: If you're getting this error in an automated, "
+                "non-interactive workload like a pipeline run or a CI/CD job, "
+                "you should use a service account API key to authenticate to "
+                "the server instead of temporary CLI login credentials. For "
+                "more information, see "
+                "https://docs.zenml.io/how-to/project-setup-and-management/connecting-to-zenml/connect-with-a-service-account"
+            )
+
+            if api_key is not None:
+                # An API key is configured. Use it as a password to
+                # authenticate.
+                data = {
+                    "grant_type": OAuthGrantTypes.ZENML_API_KEY.value,
+                    "password": api_key,
+                }
+            elif username is not None and password is not None:
+                # Username and password are configured. Use them to authenticate.
+                data = {
+                    "grant_type": OAuthGrantTypes.OAUTH_PASSWORD.value,
+                    "username": username,
+                    "password": password,
+                }
+            elif is_zenml_pro_server_url(self.url):
+                # ZenML Pro tenants use a proprietary authorization grant
+                # where the ZenML Pro API session token is exchanged for a
+                # regular ZenML server access token.
+
+                # Get the ZenML Pro API session token, if cached and valid
+                pro_token = credentials_store.get_pro_token(allow_expired=True)
+                if not pro_token:
+                    raise CredentialsNotValid(
+                        "You need to be logged in to ZenML Pro in order to "
+                        f"access the ZenML Pro server '{self.url}'. Please run "
+                        "'zenml login' to log in or choose a different server."
+                        + api_key_hint
+                    )
+
+                elif pro_token.expired:
+                    raise CredentialsNotValid(
+                        "Your ZenML Pro login session has expired. "
+                        "Please log in again using 'zenml login'."
+                        + api_key_hint
                     )
-                self._api_token = response["access_token"]
-                self.config.api_token = self._api_token
+
+                data = {
+                    "grant_type": OAuthGrantTypes.ZENML_EXTERNAL.value,
+                }
+                headers.update(
+                    {"Authorization": "Bearer " + pro_token.access_token}
+                )
             else:
-                raise ValueError(
-                    "No API token, API key or username/password provided. "
-                    "Please provide either an API token, an API key or a "
-                    "username and password in the ZenML config."
+                if not token:
+                    raise CredentialsNotValid(
+                        "No valid credentials found. Please run 'zenml login "
+                        f"--url {self.url}' to connect to the current server."
+                        + api_key_hint
+                    )
+                elif token.expired:
+                    raise CredentialsNotValid(
+                        "Your authentication to the current server has expired. "
+                        "Please log in again using 'zenml login --url "
+                        f"{self.url}'." + api_key_hint
+                    )
+
+            response = self._handle_response(
+                requests.post(
+                    self.url + API + VERSION_1 + LOGIN,
+                    data=data,
+                    verify=self.config.verify_ssl,
+                    timeout=self.config.http_timeout,
+                    headers=headers,
                 )
-        return self._api_token
+            )
+            try:
+                token_response = OAuthTokenResponse.model_validate(response)
+            except ValidationError as e:
+                raise CredentialsNotValid(
+                    "Unexpected response received while authenticating to "
+                    f"the server {e}"
+                ) from e
+
+            # Cache the token
+            self._api_token = credentials_store.set_token(
+                self.url, token_response
+            )
+
+            # Update the server info in the credentials store with the latest
+            # information from the server.
+            # NOTE: this is the best place to do this because we know that
+            # the token is valid and the server is reachable.
+            try:
+                server_info = self.get_store_info()
+            except Exception as e:
+                logger.warning(f"Failed to get server info: {e}.")
+            else:
+                credentials_store.update_server_info(self.url, server_info)
+
+        return self._api_token.access_token
 
     @property
     def session(self) -> requests.Session:
-        """Authenticate to the ZenML server.
+        """Initialize and return a requests session.
 
         Returns:
-            A requests session with the authentication token.
+            A requests session.
         """
         if self._session is None:
+            # We only need to initialize the session once over the lifetime
+            # of the client. We can swap the token out when it expires.
             if self.config.verify_ssl is False:
                 urllib3.disable_warnings(
                     urllib3.exceptions.InsecureRequestWarning
                 )
 
             self._session = requests.Session()
-            retries = Retry(backoff_factor=0.1, connect=5)
+            # Retries are triggered for idempotent HTTP methods (GET, HEAD, PUT,
+            # OPTIONS and DELETE) on specific HTTP status codes:
+            #
+            #     500: Internal Server Error.
+            #     502: Bad Gateway.
+            #     503: Service Unavailable.
+            #     504: Gateway Timeout.
+            #
+            # This also handles connection level errors, if a connection attempt
+            # fails due to transient issues like:
+            #
+            #     DNS resolution errors.
+            #     Connection timeouts.
+            #     Network disruptions.
+            #
+            # Additional errors retried:
+            #
+            #     Read Timeouts: If the server does not send a response within
+            #     the timeout period.
+            #     Connection Refused: If the server refuses the connection.
+            #
+            retries = Retry(
+                connect=5,
+                read=8,
+                redirect=3,
+                status=10,
+                allowed_methods=["HEAD", "GET", "PUT", "DELETE", "OPTIONS"],
+                status_forcelist=[
+                    408,  # Request Timeout
+                    429,  # Too Many Requests
+                    500,  # Internal Server Error
+                    502,  # Bad Gateway
+                    503,  # Service Unavailable
+                    504,  # Gateway Timeout
+                ],
+                other=3,
+                backoff_factor=0.5,
+            )
             self._session.mount("https://", HTTPAdapter(max_retries=retries))
             self._session.mount("http://", HTTPAdapter(max_retries=retries))
             self._session.verify = self.config.verify_ssl
-            token = self._get_auth_token()
-            self._session.headers.update({"Authorization": "Bearer " + token})
-            logger.debug("Authenticated to ZenML server.")
+            # Use a custom user agent to identify the ZenML client in the server
+            # logs.
+            self._session.headers.update(
+                {"User-Agent": "zenml/" + zenml.__version__}
+            )
+
+        # Note that we return an unauthenticated session here. An API token
+        # is only fetched and set in the authorization header when and if it is
+        # needed.
         return self._session
 
-    def clear_session(self) -> None:
-        """Clear the authentication session and any cached API tokens."""
-        self._session = None
+    def authenticate(self, force: bool = False) -> None:
+        """Authenticate or re-authenticate to the ZenML server.
+
+        Args:
+            force: If True, force a re-authentication even if a valid API token
+                is currently cached. This is useful when the current API token
+                is known to be invalid or expired.
+        """
+        # This is called to trigger an authentication flow, either because
+        # the current API token is expired or no longer valid, or because
+        # a configuration change has happened or merely because an
+        # authentication was never attempted before.
+        #
+        # 1. Drop the API token currently being used, if any.
+        # 2. If force=True, clear the current API token from the credentials
+        # store, if any, otherwise it will just be re-used on the next call.
+        # 3. Get a new API token
+
+        # The authentication token could have expired or invalidated through
+        # other means; refresh it and try again. This will clear any cached
+        # token and trigger a new authentication flow.
+        if self._api_token and not force:
+            if self._api_token.expired:
+                logger.info(
+                    "Authentication session expired; attempting to "
+                    "re-authenticate."
+                )
+            else:
+                logger.info(
+                    "Authentication session was invalidated by the server; "
+                    "This can happen for example if the user's permissions "
+                    "have been revoked or if the server has been restarted "
+                    "and lost its session state. Attempting to "
+                    "re-authenticate."
+                )
+        else:
+            if force:
+                # Clear the current API token from the credentials store, if
+                # any, to force a new authentication flow.
+                get_credentials_store().clear_token(self.url)
+            # Never authenticated since the client was created or the API token
+            # was explicitly cleared.
+            logger.debug(f"Authenticating to {self.url}...")
+
         self._api_token = None
-        # Clear the configured API token only if it's possible to fetch a new
-        # one from the server using other credentials (username/password or
-        # service account API key).
-        if (
-            self.config.username is not None
-            and self.config.password is not None
-            or self.config.api_key is not None
-        ):
-            self.config.api_token = None
+
+        new_api_token = self.get_or_generate_api_token()
+
+        # Set or refresh the authentication token
+        self.session.headers.update(
+            {"Authorization": "Bearer " + new_api_token}
+        )
+        logger.debug(f"Authenticated to {self.url}")
 
     @staticmethod
     def _handle_response(response: requests.Response) -> Json:
@@ -3417,6 +4334,7 @@ class RestZenStore(BaseZenStore):
         method: str,
         url: str,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a request to the REST API.
@@ -3425,14 +4343,15 @@ class RestZenStore(BaseZenStore):
             method: The HTTP method to use.
             url: The URL to request.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
             The parsed response.
 
         Raises:
-            AuthorizationException: if the request fails due to an expired
-                authentication token.
+            CredentialsNotValid: if the request fails due to invalid
+                client credentials.
         """
         params = {k: str(v) for k, v in params.items()} if params else {}
 
@@ -3440,49 +4359,88 @@ class RestZenStore(BaseZenStore):
             {source_context.name: source_context.get().value}
         )
 
-        try:
-            return self._handle_response(
-                self.session.request(
-                    method,
-                    url,
-                    params=params,
-                    verify=self.config.verify_ssl,
-                    timeout=self.config.http_timeout,
-                    **kwargs,
-                )
-            )
-        except AuthorizationException:
-            # The authentication token could have expired; refresh it and try
-            # again. This will clear any cached token and trigger a new
-            # authentication flow.
-            self.clear_session()
-            logger.info("Authentication token expired; refreshing...")
-
-        try:
-            return self._handle_response(
-                self.session.request(
-                    method,
-                    url,
-                    params=params,
-                    verify=self.config.verify_ssl,
-                    timeout=self.config.http_timeout,
-                    **kwargs,
+        # If the server replies with a credentials validation (401 Unauthorized)
+        # error, we (re-)authenticate and retry the request here in the
+        # following cases:
+        #
+        # 1. initial authentication: the last request was not authenticated
+        # with an API token.
+        # 2. re-authentication: the last request was authenticated with an API
+        # token that was rejected by the server. This is to cover the case
+        # of expired tokens that can be refreshed by the client automatically
+        # without user intervention from other sources (e.g. API keys).
+        #
+        # NOTE: it can happen that the same request is retried here for up to
+        # two times: once after initial authentication and once after
+        # re-authentication.
+        re_authenticated = False
+        while True:
+            try:
+                return self._handle_response(
+                    self.session.request(
+                        method,
+                        url,
+                        params=params,
+                        verify=self.config.verify_ssl,
+                        timeout=timeout or self.config.http_timeout,
+                        **kwargs,
+                    )
                 )
-            )
-        except AuthorizationException:
-            logger.info(
-                "Your authentication token has expired. Please re-authenticate."
-            )
-            raise
+            except CredentialsNotValid as e:
+                # NOTE: CredentialsNotValid is raised only when the server
+                # explicitly indicates that the credentials are not valid and
+                # they can be thrown away or when the request is not
+                # authenticated at all.
+
+                if self._api_token is None:
+                    # The last request was not authenticated with an API
+                    # token at all. We authenticate here and then try the
+                    # request again, this time with a valid API token in the
+                    # header.
+                    logger.debug(
+                        f"The last request was not authenticated: {e}\n"
+                        "Re-authenticating and retrying..."
+                    )
+                    self.authenticate()
+                elif not re_authenticated:
+                    # The last request was authenticated with an API token
+                    # that was rejected by the server. We attempt a
+                    # re-authentication here and then retry the request.
+                    logger.debug(
+                        "The last request was authenticated with an API token "
+                        f"that was rejected by the server: {e}\n"
+                        "Re-authenticating and retrying..."
+                    )
+                    re_authenticated = True
+                    self.authenticate(
+                        # Ignore the current token and force a re-authentication
+                        force=True
+                    )
+                else:
+                    # The last request was made after re-authenticating but
+                    # still failed. Bailing out.
+                    logger.debug(
+                        f"The last request failed after re-authenticating: {e}\n"
+                        "Bailing out..."
+                    )
+                    raise CredentialsNotValid(
+                        "The current credentials are no longer valid. Please "
+                        "log in again using 'zenml login'."
+                    ) from e
 
     def get(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a GET request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3490,17 +4448,26 @@ class RestZenStore(BaseZenStore):
         """
         logger.debug(f"Sending GET request to {path}...")
         return self._request(
-            "GET", self.url + API + VERSION_1 + path, params=params, **kwargs
+            "GET",
+            self.url + API + VERSION_1 + path,
+            params=params,
+            timeout=timeout,
+            **kwargs,
         )
 
     def delete(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a DELETE request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3511,6 +4478,7 @@ class RestZenStore(BaseZenStore):
             "DELETE",
             self.url + API + VERSION_1 + path,
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -3519,6 +4487,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: BaseModel,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a POST request to the given endpoint path.
@@ -3527,6 +4496,7 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3536,8 +4506,9 @@ class RestZenStore(BaseZenStore):
         return self._request(
             "POST",
             self.url + API + VERSION_1 + path,
-            data=body.json(),
+            json=body.model_dump(mode="json"),
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -3546,6 +4517,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: Optional[BaseModel] = None,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a PUT request to the given endpoint path.
@@ -3554,18 +4526,22 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
             The response body.
         """
         logger.debug(f"Sending PUT request to {path}...")
-        data = body.json(exclude_unset=True) if body else None
+        json = (
+            body.model_dump(mode="json", exclude_unset=True) if body else None
+        )
         return self._request(
             "PUT",
             self.url + API + VERSION_1 + path,
-            data=data,
+            json=json,
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -3589,7 +4565,42 @@ class RestZenStore(BaseZenStore):
             The created resource.
         """
         response_body = self.post(f"{route}", body=resource, params=params)
-        return response_model.parse_obj(response_body)
+
+        return response_model.model_validate(response_body)
+
+    def _batch_create_resources(
+        self,
+        resources: List[AnyRequest],
+        response_model: Type[AnyResponse],
+        route: str,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> List[AnyResponse]:
+        """Create a new batch of resources.
+
+        Args:
+            resources: The resources to create.
+            response_model: The response model of an individual resource.
+            route: The resource REST route to use.
+            params: Optional query parameters to pass to the endpoint.
+
+        Returns:
+            List of response models.
+        """
+        json_data = [
+            resource.model_dump(mode="json") for resource in resources
+        ]
+        response = self._request(
+            "POST",
+            self.url + API + VERSION_1 + route + BATCH,
+            json=json_data,
+            params=params,
+        )
+        assert isinstance(response, list)
+
+        return [
+            response_model.model_validate(model_data)
+            for model_data in response
+        ]
 
     def _create_workspace_scoped_resource(
         self,
@@ -3665,7 +4676,7 @@ class RestZenStore(BaseZenStore):
                 f"response from the {route}{GET_OR_CREATE} endpoint but got "
                 f"{type(was_created)} instead."
             )
-        return response_model.parse_obj(model_json), was_created
+        return response_model.model_validate(model_json), was_created
 
     def _get_or_create_workspace_scoped_resource(
         self,
@@ -3713,7 +4724,7 @@ class RestZenStore(BaseZenStore):
             The retrieved resource.
         """
         body = self.get(f"{route}/{str(resource_id)}", params=params)
-        return response_model.parse_obj(body)
+        return response_model.model_validate(body)
 
     def _list_paginated_resources(
         self,
@@ -3738,17 +4749,17 @@ class RestZenStore(BaseZenStore):
         """
         # leave out filter params that are not supplied
         params = params or {}
-        params.update(filter_model.dict(exclude_none=True))
+        params.update(filter_model.model_dump(exclude_none=True))
         body = self.get(f"{route}", params=params)
         if not isinstance(body, dict):
             raise ValueError(
                 f"Bad API Response. Expected list, got {type(body)}"
             )
         # The initial page of items will be of type BaseResponseModel
-        page_of_items: Page[AnyResponse] = Page.parse_obj(body)
+        page_of_items: Page[AnyResponse] = Page.model_validate(body)
         # So these items will be parsed into their correct types like here
         page_of_items.items = [
-            response_model.parse_obj(generic_item)
+            response_model.model_validate(generic_item)
             for generic_item in body["items"]
         ]
         return page_of_items
@@ -3779,7 +4790,7 @@ class RestZenStore(BaseZenStore):
             raise ValueError(
                 f"Bad API Response. Expected list, got {type(body)}"
             )
-        return [response_model.parse_obj(entry) for entry in body]
+        return [response_model.model_validate(entry) for entry in body]
 
     def _update_resource(
         self,
@@ -3806,7 +4817,7 @@ class RestZenStore(BaseZenStore):
             f"{route}/{str(resource_id)}", body=resource_update, params=params
         )
 
-        return response_model.parse_obj(response_body)
+        return response_model.model_validate(response_body)
 
     def _delete_resource(
         self, resource_id: Union[str, UUID], route: str