zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__py3-none-any.whl

zenml/zen_server/utils.py

@@ -39,13 +39,14 @@ from zenml.constants import (
     INFO,
     VERSION_1,
 )
+from zenml.enums import StoreType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
-from zenml.zen_server.deploy.deployment import ServerDeployment
-from zenml.zen_server.deploy.local.local_zen_server import (
-    LocalServerDeploymentConfig,
+from zenml.zen_server.deploy.deployment import (
+    LocalServerDeployment,
 )
+from zenml.zen_server.deploy.exceptions import ServerDeploymentNotFoundError
 from zenml.zen_server.exceptions import http_exception_from_error
 from zenml.zen_server.feature_gate.feature_gate_interface import (
     FeatureGateInterface,
@@ -236,7 +237,7 @@ def server_config() -> ServerConfiguration:
     return _server_config
 
 
-def get_local_server() -> Optional["ServerDeployment"]:
+def get_local_server() -> Optional["LocalServerDeployment"]:
     """Get the active local server.
 
     Call this function to retrieve the local server deployed on this machine.
@@ -245,48 +246,75 @@ def get_local_server() -> Optional["ServerDeployment"]:
         The local server deployment or None, if no local server deployment was
         found.
     """
-    from zenml.zen_server.deploy.deployer import ServerDeployer
+    from zenml.zen_server.deploy.deployer import LocalServerDeployer
 
-    deployer = ServerDeployer()
-    servers = deployer.list_servers()
-    if not servers:
+    deployer = LocalServerDeployer()
+    try:
+        return deployer.get_server()
+    except ServerDeploymentNotFoundError:
         return None
 
-    return servers[0]
 
+def connected_to_local_server() -> bool:
+    """Check if the client is connected to a local server.
 
-def get_active_server_details() -> Tuple[str, Optional[int]]:
-    """Get the URL of the current ZenML Server.
+    Returns:
+        True if the client is connected to a local server, False otherwise.
+    """
+    from zenml.zen_server.deploy.deployer import LocalServerDeployer
 
-    When multiple servers are present, the following precedence is used to
-    determine which server to use:
-    - If the client is connected to a server, that server has precedence.
-    - If no server is connected, the local server is used, if present.
+    deployer = LocalServerDeployer()
+    return deployer.is_connected_to_server()
 
-    Returns:
-        The URL and port of the currently active server.
+
+def show_dashboard(
+    local: bool = False,
+    ngrok_token: Optional[str] = None,
+) -> None:
+    """Show the ZenML dashboard.
+
+    Args:
+        local: Whether to show the dashboard for the local server or the
+            one for the active server.
+        ngrok_token: An ngrok auth token to use for exposing the ZenML
+            dashboard on a public domain. Primarily used for accessing the
+            dashboard in Colab.
 
     Raises:
-        RuntimeError: If no server is active.
+        RuntimeError: If no server is connected.
     """
-    # Check for connected servers first
-    gc = GlobalConfiguration()
-    if not gc.uses_default_store():
-        logger.debug("Getting URL of connected server.")
-        parsed_url = urlparse(gc.store_configuration.url)
-        return f"{parsed_url.scheme}://{parsed_url.hostname}", parsed_url.port
-    # Else, check for local servers
-    server = get_local_server()
-    if server and server.status and server.status.url:
-        if isinstance(server.config, LocalServerDeploymentConfig):
-            return server.status.url, server.config.port
-        return server.status.url, None
-
-    raise RuntimeError(
-        "ZenML is not connected to any server right now. Please use "
-        "`zenml connect` to connect to a server or spin up a new local server "
-        "via `zenml up`."
-    )
+    from zenml.utils.dashboard_utils import show_dashboard
+    from zenml.utils.networking_utils import get_or_create_ngrok_tunnel
+
+    url: Optional[str] = None
+    if not local:
+        gc = GlobalConfiguration()
+        if gc.store_configuration.type == StoreType.REST:
+            url = gc.store_configuration.url
+
+    if not url:
+        # Else, check for local servers
+        server = get_local_server()
+        if server and server.status and server.status.url:
+            url = server.status.url
+
+    if not url:
+        raise RuntimeError(
+            "ZenML is not connected to any server right now. Please use "
+            "`zenml login` to connect to a server or spin up a new local server "
+            "via `zenml login --local`."
+        )
+
+    if ngrok_token:
+        parsed_url = urlparse(url)
+
+        ngrok_url = get_or_create_ngrok_tunnel(
+            ngrok_token=ngrok_token, port=parsed_url.port or 80
+        )
+        logger.debug(f"Tunneling dashboard from {url} to {ngrok_url}.")
+        url = ngrok_url
+
+    show_dashboard(url)
 
 
 F = TypeVar("F", bound=Callable[..., Any])

zenml/zen_stores/base_zen_store.py

@@ -24,11 +24,9 @@ from typing import (
     Type,
     Union,
 )
-from urllib.parse import urlparse
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, model_validator
-from requests import ConnectionError
 
 import zenml
 from zenml.config.global_config import GlobalConfiguration
@@ -39,13 +37,11 @@ from zenml.constants import (
     DEFAULT_WORKSPACE_NAME,
     ENV_ZENML_DEFAULT_WORKSPACE_NAME,
     IS_DEBUG_ENV,
-    ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT,
 )
 from zenml.enums import (
     SecretsStoreType,
     StoreType,
 )
-from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
     ServerDatabaseType,
@@ -57,9 +53,6 @@ from zenml.models import (
     WorkspaceResponse,
 )
 from zenml.utils.pydantic_utils import before_validator_handler
-from zenml.zen_stores.secrets_stores.sql_secrets_store import (
-    SqlSecretsStoreConfiguration,
-)
 from zenml.zen_stores.zen_store_interface import ZenStoreInterface
 
 logger = get_logger(__name__)
@@ -136,51 +129,10 @@ class BaseZenStore(
                 stack and user in the store will be skipped.
             **kwargs: Additional keyword arguments to pass to the Pydantic
                 constructor.
-
-        Raises:
-            RuntimeError: If the store cannot be initialized.
-            AuthorizationException: If the store cannot be initialized due to
-                authentication errors.
         """
         super().__init__(**kwargs)
 
-        try:
-            self._initialize()
-
-        # Handle cases where the ZenML server is not available
-        except ConnectionError as e:
-            error_message = (
-                "Cannot connect to the ZenML database because the ZenML server "
-                f"at {self.url} is not running."
-            )
-            if urlparse(self.url).hostname in ["localhost", "127.0.0.1"]:
-                recommendation = (
-                    "Please run `zenml down` and `zenml up` to restart the "
-                    "server."
-                )
-            else:
-                recommendation = (
-                    "Please run `zenml disconnect` and `zenml connect --url "
-                    f"{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 ".zenml.io" in self.url:
-                zenml_pro_extra = (
-                    ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT
-                )
-            raise RuntimeError(
-                f"Error initializing {self.type.value} store with URL "
-                f"'{self.url}': {str(e)}" + zenml_pro_extra
-            ) from e
+        self._initialize()
 
         if not skip_default_registrations:
             logger.debug("Initializing database")
@@ -293,6 +245,9 @@ class BaseZenStore(
         Returns:
             The default store configuration.
         """
+        from zenml.zen_stores.secrets_stores.sql_secrets_store import (
+            SqlSecretsStoreConfiguration,
+        )
         from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration
 
         config = SqlZenStoreConfiguration(

zenml/zen_stores/migrations/versions/0.68.1_release.py

@@ -0,0 +1,23 @@
+"""Release [0.68.1].
+
+Revision ID: 0.68.1
+Revises: 0.68.0
+Create Date: 2024-10-28 10:34:17.307663
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.68.1"
+down_revision = "0.68.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/c22561cbb3a9_add_artifact_unique_constraints.py

@@ -0,0 +1,86 @@
+"""Add artifact unique constraints [c22561cbb3a9].
+
+Revision ID: c22561cbb3a9
+Revises: 0.68.1
+Create Date: 2024-10-17 16:41:25.053677
+
+"""
+
+from collections import defaultdict
+from typing import Dict, Set
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "c22561cbb3a9"
+down_revision = "0.68.1"
+branch_labels = None
+depends_on = None
+
+
+def resolve_duplicate_versions() -> None:
+    """Resolve duplicate artifact versions."""
+    connection = op.get_bind()
+
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=op.get_bind(),
+        only=("artifact_version",),
+    )
+
+    artifact_version_table = sa.Table("artifact_version", meta)
+    query = sa.select(
+        artifact_version_table.c.id,
+        artifact_version_table.c.artifact_id,
+        artifact_version_table.c.version,
+    )
+
+    versions_per_artifact: Dict[str, Set[str]] = defaultdict(set)
+
+    for id, artifact_id, version in connection.execute(query).fetchall():
+        versions = versions_per_artifact[artifact_id]
+        if version in versions:
+            for suffix_length in range(4, len(id)):
+                new_version = f"{version}-{id[:suffix_length]}"
+                if new_version not in versions:
+                    version = new_version
+                    break
+
+            connection.execute(
+                sa.update(artifact_version_table)
+                .where(artifact_version_table.c.id == id)
+                .values(version=version)
+            )
+
+        versions.add(version)
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    resolve_duplicate_versions()
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_artifact_name", ["name"])
+
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_version_for_artifact_id", ["version", "artifact_id"]
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_version_for_artifact_id", type_="unique"
+        )
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_artifact_name", type_="unique")
+
+    # ### end Alembic commands ###