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

zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py

@@ -0,0 +1,375 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""HyperAI Service Connector.
+
+The HyperAI Service Connector allows authenticating to HyperAI (hyperai.ai)
+GPU equipped instances.
+"""
+
+import base64
+import io
+from typing import Any, List, Optional, Type
+
+import paramiko
+from pydantic import Field
+
+from zenml.exceptions import AuthorizationException
+from zenml.integrations.hyperai import (
+    HYPERAI_CONNECTOR_TYPE,
+    HYPERAI_RESOURCE_TYPE,
+)
+from zenml.logger import get_logger
+from zenml.models import (
+    AuthenticationMethodModel,
+    ResourceTypeModel,
+    ServiceConnectorTypeModel,
+)
+from zenml.service_connectors.service_connector import (
+    AuthenticationConfig,
+    ServiceConnector,
+)
+from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
+
+logger = get_logger(__name__)
+
+
+class HyperAICredentials(AuthenticationConfig):
+    """HyperAI client authentication credentials."""
+
+    base64_ssh_key: PlainSerializedSecretStr = Field(
+        title="SSH key (base64)",
+    )
+    ssh_passphrase: Optional[PlainSerializedSecretStr] = Field(
+        default=None,
+        title="SSH key passphrase",
+    )
+
+
+class HyperAIConfiguration(HyperAICredentials):
+    """HyperAI client configuration."""
+
+    hostnames: List[str] = Field(
+        title="Hostnames of the supported HyperAI instances.",
+    )
+
+    username: str = Field(
+        title="Username to use to connect to the HyperAI instance.",
+    )
+
+
+class HyperAIAuthenticationMethods(StrEnum):
+    """HyperAI Authentication methods."""
+
+    RSA_KEY_OPTIONAL_PASSPHRASE = "rsa-key"
+    DSA_KEY_OPTIONAL_PASSPHRASE = "dsa-key"
+    ECDSA_KEY_OPTIONAL_PASSPHRASE = "ecdsa-key"
+    ED25519_KEY_OPTIONAL_PASSPHRASE = "ed25519-key"
+
+
+HYPERAI_SERVICE_CONNECTOR_TYPE_SPEC = ServiceConnectorTypeModel(
+    name="HyperAI Service Connector",
+    connector_type=HYPERAI_CONNECTOR_TYPE,
+    description="""
+The ZenML HyperAI Service Connector allows authenticating to HyperAI (hyperai.ai)
+GPU equipped instances.
+
+This connector provides an SSH connection to your HyperAI instance, which can be
+used to run ZenML pipelines.
+
+The instance must be configured to allow SSH connections from the ZenML server.
+Docker and Docker Compose must be installed on the HyperAI instance. If you want
+to use scheduled pipeline runs, also ensure that a working cron daemon is installed
+and running on the HyperAI instance.
+""",
+    logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/connectors/hyperai/hyperai.png",
+    emoji=":robot_face:",
+    auth_methods=[
+        AuthenticationMethodModel(
+            name="RSA key with optional passphrase",
+            auth_method=HyperAIAuthenticationMethods.RSA_KEY_OPTIONAL_PASSPHRASE,
+            description="""
+Use an RSA private key to authenticate with a HyperAI instance. The key may be
+encrypted with a passphrase. If the key is encrypted, the passphrase must be
+provided. Make sure to provide the key as a Base64 encoded string.
+""",
+            config_class=HyperAIConfiguration,
+        ),
+        AuthenticationMethodModel(
+            name="DSA/DSS key with optional passphrase",
+            auth_method=HyperAIAuthenticationMethods.DSA_KEY_OPTIONAL_PASSPHRASE,
+            description="""
+Use a DSA/DSS private key to authenticate with a HyperAI instance. The key may be
+encrypted with a passphrase. If the key is encrypted, the passphrase must be
+provided. Make sure to provide the key as a Base64 encoded string.
+""",
+            config_class=HyperAIConfiguration,
+        ),
+        AuthenticationMethodModel(
+            name="ECDSA key with optional passphrase",
+            auth_method=HyperAIAuthenticationMethods.ECDSA_KEY_OPTIONAL_PASSPHRASE,
+            description="""
+Use an ECDSA private key to authenticate with a HyperAI instance. The key may be
+encrypted with a passphrase. If the key is encrypted, the passphrase must be
+provided. Make sure to provide the key as a Base64 encoded string.
+""",
+            config_class=HyperAIConfiguration,
+        ),
+        AuthenticationMethodModel(
+            name="Ed25519 key with optional passphrase",
+            auth_method=HyperAIAuthenticationMethods.ED25519_KEY_OPTIONAL_PASSPHRASE,
+            description="""
+Use an Ed25519 private key to authenticate with a HyperAI instance. The key may be
+encrypted with a passphrase. If the key is encrypted, the passphrase must be
+provided. Make sure to provide the key as a Base64 encoded string.
+""",
+            config_class=HyperAIConfiguration,
+        ),
+    ],
+    resource_types=[
+        ResourceTypeModel(
+            name="HyperAI instance",
+            resource_type=HYPERAI_RESOURCE_TYPE,
+            description="""
+Allows users to access a HyperAI instance as a resource. When used by
+connector consumers, they are provided a pre-authenticated SSH client
+instance.
+""",
+            auth_methods=HyperAIAuthenticationMethods.values(),
+            supports_instances=True,
+            logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/connectors/hyperai/hyperai.png",
+            emoji=":robot_face:",
+        ),
+    ],
+)
+
+
+class HyperAIServiceConnector(ServiceConnector):
+    """HyperAI service connector."""
+
+    config: HyperAIConfiguration
+
+    @classmethod
+    def _get_connector_type(cls) -> ServiceConnectorTypeModel:
+        """Get the service connector specification.
+
+        Returns:
+            The service connector specification.
+        """
+        return HYPERAI_SERVICE_CONNECTOR_TYPE_SPEC
+
+    def _paramiko_key_type_given_auth_method(self) -> Type[paramiko.PKey]:
+        """Get the Paramiko key type given the authentication method.
+
+        Returns:
+            The Paramiko key type.
+
+        Raises:
+            ValueError: If the authentication method is invalid.
+        """
+        mapping = {
+            HyperAIAuthenticationMethods.RSA_KEY_OPTIONAL_PASSPHRASE: paramiko.RSAKey,
+            HyperAIAuthenticationMethods.DSA_KEY_OPTIONAL_PASSPHRASE: paramiko.DSSKey,
+            HyperAIAuthenticationMethods.ECDSA_KEY_OPTIONAL_PASSPHRASE: paramiko.ECDSAKey,
+            HyperAIAuthenticationMethods.ED25519_KEY_OPTIONAL_PASSPHRASE: paramiko.Ed25519Key,
+        }
+
+        try:
+            return mapping[HyperAIAuthenticationMethods(self.auth_method)]
+        except KeyError:
+            raise ValueError(
+                f"Invalid authentication method: {self.auth_method}"
+            )
+
+    def _create_paramiko_client(
+        self, hostname: str
+    ) -> paramiko.client.SSHClient:
+        """Create a Paramiko SSH client based on the configuration.
+
+        Args:
+            hostname: The hostname of the HyperAI instance.
+
+        Returns:
+            A Paramiko SSH client.
+
+        Raises:
+            AuthorizationException: If the client cannot be created.
+        """
+        if self.config.ssh_passphrase is None:
+            ssh_passphrase = None
+        else:
+            ssh_passphrase = self.config.ssh_passphrase.get_secret_value()
+
+        # Connect to the HyperAI instance
+        try:
+            # Convert the SSH key from base64 to string
+            base64_key_value = self.config.base64_ssh_key.get_secret_value()
+            ssh_key = base64.b64decode(base64_key_value).decode("utf-8")
+            paramiko_key = None
+
+            with io.StringIO(ssh_key) as f:
+                paramiko_key = self._paramiko_key_type_given_auth_method().from_private_key(
+                    f, password=ssh_passphrase
+                )
+
+            # Trim whitespace from the IP address
+            hostname = hostname.strip()
+
+            paramiko_client = paramiko.client.SSHClient()
+            paramiko_client.set_missing_host_key_policy(
+                paramiko.AutoAddPolicy()  # nosec
+            )
+            paramiko_client.connect(
+                hostname=hostname,
+                username=self.config.username,
+                pkey=paramiko_key,
+                timeout=30,
+            )
+
+            return paramiko_client
+
+        except paramiko.ssh_exception.BadHostKeyException as e:
+            logger.error("Bad host key: %s", e)
+        except paramiko.ssh_exception.AuthenticationException as e:
+            logger.error("Authentication failed: %s", e)
+        except paramiko.ssh_exception.SSHException as e:
+            logger.error(
+                "SSH error: %s. A common cause for this error is selection of the wrong key type in your service connector.",
+                e,
+            )
+        except Exception as e:
+            logger.error(
+                "Unknown error while connecting to HyperAI instance: %s. Please check your network connection, IP address, and authentication details.",
+                e,
+            )
+
+        raise AuthorizationException(
+            "Could not create SSH client for HyperAI instance."
+        )
+
+    def _authorize_client(self, hostname: str) -> None:
+        """Verify that the client can authenticate with the HyperAI instance.
+
+        Args:
+            hostname: The hostname of the HyperAI instance.
+        """
+        logger.info("Verifying connection to HyperAI instance...")
+
+        paramiko_client = self._create_paramiko_client(hostname)
+        paramiko_client.close()
+
+    def _connect_to_resource(
+        self,
+        **kwargs: Any,
+    ) -> Any:
+        """Connect to a HyperAI instance. Returns an authenticated SSH client.
+
+        Args:
+            kwargs: Additional implementation specific keyword arguments to pass
+                to the session or client constructor.
+
+        Returns:
+            An authenticated Paramiko SSH client.
+        """
+        logger.info("Connecting to HyperAI instance...")
+        assert self.resource_id is not None
+
+        paramiko_client = self._create_paramiko_client(self.resource_id)
+        return paramiko_client
+
+    def _configure_local_client(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        """There is no local client for the HyperAI connector, so it does nothing.
+
+        Args:
+            kwargs: Additional implementation specific keyword arguments to pass
+                to the session or client constructor.
+
+        Raises:
+            NotImplementedError: If there is no local client for the HyperAI
+                connector.
+        """
+        raise NotImplementedError(
+            "There is no local client for the HyperAI connector."
+        )
+
+    @classmethod
+    def _auto_configure(
+        cls,
+        auth_method: Optional[str] = None,
+        resource_type: Optional[str] = None,
+        resource_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> "HyperAIServiceConnector":
+        """Auto-configure the connector.
+
+        Not supported by the HyperAI connector.
+
+        Args:
+            auth_method: The particular authentication method to use. If not
+                specified, the connector implementation must decide which
+                authentication method to use or raise an exception.
+            resource_type: The type of resource to configure.
+            resource_id: The ID of the resource to configure. The
+                implementation may choose to either require or ignore this
+                parameter if it does not support or detect an resource type that
+                supports multiple instances.
+            kwargs: Additional implementation specific keyword arguments to use.
+
+        Raises:
+            NotImplementedError: If the connector auto-configuration fails or
+                is not supported.
+        """
+        raise NotImplementedError(
+            "Auto-configuration is not supported by the HyperAI connector."
+        )
+
+    def _verify(
+        self,
+        resource_type: Optional[str] = None,
+        resource_id: Optional[str] = None,
+    ) -> List[str]:
+        """Verify that a connection can be established to the HyperAI instance.
+
+        Args:
+            resource_type: The type of resource to verify. Must be set to the
+                Docker resource type.
+            resource_id: The HyperAI instance to verify.
+
+        Returns:
+            The resource ID if the connection can be established.
+
+        Raises:
+            ValueError: If the resource ID is not in the list of configured
+                hostnames.
+        """
+        if resource_id:
+            if resource_id not in self.config.hostnames:
+                raise ValueError(
+                    f"The supplied hostname '{resource_id}' is not in the list "
+                    f"of configured hostnames: {self.config.hostnames}. Please "
+                    f"check your configuration."
+                )
+            hostnames = [resource_id]
+        else:
+            hostnames = self.config.hostnames
+
+        resources = []
+        for hostname in hostnames:
+            self._authorize_client(hostname)
+            resources.append(hostname)
+
+        return resources

zenml/integrations/integration.py

@@ -13,15 +13,21 @@
 #  permissions and limitations under the License.
 """Base and meta classes for ZenML integrations."""
 
-from typing import Any, Dict, List, Optional, Tuple, Type, cast
+import re
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 
 import pkg_resources
+from pkg_resources import Requirement
 
 from zenml.integrations.registry import integration_registry
 from zenml.logger import get_logger
 from zenml.stack.flavor import Flavor
 from zenml.utils.integration_utils import parse_requirement
 
+if TYPE_CHECKING:
+    from zenml.plugins.base_plugin_flavor import BasePluginFlavor
+
+
 logger = get_logger(__name__)
 
 
@@ -54,6 +60,7 @@ class Integration(metaclass=IntegrationMeta):
 
     REQUIREMENTS: List[str] = []
     APT_PACKAGES: List[str] = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL: List[str] = []
 
     @classmethod
     def check_installation(cls) -> bool:
@@ -62,53 +69,68 @@ class Integration(metaclass=IntegrationMeta):
         Returns:
             True if all required packages are installed, False otherwise.
         """
-        try:
-            for r in cls.get_requirements():
-                name, extras = parse_requirement(r)
-                if name:
-                    dist = pkg_resources.get_distribution(name)
-                    # Check if extras are specified and installed
-                    if extras:
-                        extra_list = extras[1:-1].split(",")
-                        for extra in extra_list:
-                            try:
-                                requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
-                            except pkg_resources.UnknownExtra as e:
-                                logger.debug("Unknown extra: " + str(e))
-                                return False
-
-                            for ri in requirements:
-                                try:
-                                    pkg_resources.get_distribution(ri)
-                                except IndexError:
-                                    logger.debug(
-                                        f"Unable to find required extra '{ri.project_name}' for "
-                                        f"requirements '{name}' coming from integration '{cls.NAME}'."
-                                    )
-                                    return False
+        for r in cls.get_requirements():
+            try:
+                # First check if the base package is installed
+                dist = pkg_resources.get_distribution(r)
+
+                # Next, check if the dependencies (including extras) are
+                # installed
+                deps: List[Requirement] = []
+
+                _, extras = parse_requirement(r)
+                if extras:
+                    extra_list = extras[1:-1].split(",")
+                    for extra in extra_list:
+                        try:
+                            requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
+                        except pkg_resources.UnknownExtra as e:
+                            logger.debug(f"Unknown extra: {str(e)}")
+                            return False
+                        deps.extend(requirements)
                 else:
-                    logger.debug(
-                        f"Invalid requirement format '{r}' for integration {cls.NAME}."
-                    )
-                    return False
-
-            logger.debug(
-                f"Integration {cls.NAME} is installed correctly with "
-                f"requirements {cls.get_requirements()}."
-            )
-            return True
-        except pkg_resources.DistributionNotFound as e:
-            logger.debug(
-                f"Unable to find required package '{e.req}' for "
-                f"integration {cls.NAME}."
-            )
-            return False
-        except pkg_resources.VersionConflict as e:
-            logger.debug(
-                f"VersionConflict error when loading installation {cls.NAME}: "
-                f"{str(e)}"
-            )
-            return False
+                    deps = dist.requires()
+
+                for ri in deps:
+                    try:
+                        # Remove the "extra == ..." part from the requirement string
+                        cleaned_req = re.sub(
+                            r"; extra == \"\w+\"", "", str(ri)
+                        )
+                        pkg_resources.get_distribution(cleaned_req)
+                    except pkg_resources.DistributionNotFound as e:
+                        logger.debug(
+                            f"Unable to find required dependency "
+                            f"'{e.req}' for requirement '{r}' "
+                            f"necessary for integration '{cls.NAME}'."
+                        )
+                        return False
+                    except pkg_resources.VersionConflict as e:
+                        logger.debug(
+                            f"Package version '{e.dist}' does not match "
+                            f"version '{e.req}' required by '{r}' "
+                            f"necessary for integration '{cls.NAME}'."
+                        )
+                        return False
+
+            except pkg_resources.DistributionNotFound as e:
+                logger.debug(
+                    f"Unable to find required package '{e.req}' for "
+                    f"integration {cls.NAME}."
+                )
+                return False
+            except pkg_resources.VersionConflict as e:
+                logger.debug(
+                    f"Package version '{e.dist}' does not match version "
+                    f"'{e.req}' necessary for integration {cls.NAME}."
+                )
+                return False
+
+        logger.debug(
+            f"Integration {cls.NAME} is installed correctly with "
+            f"requirements {cls.get_requirements()}."
+        )
+        return True
 
     @classmethod
     def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
@@ -122,6 +144,29 @@ class Integration(metaclass=IntegrationMeta):
         """
         return cls.REQUIREMENTS
 
+    @classmethod
+    def get_uninstall_requirements(
+        cls, target_os: Optional[str] = None
+    ) -> List[str]:
+        """Method to get the uninstall requirements for the integration.
+
+        Args:
+            target_os: The target operating system to get the requirements for.
+
+        Returns:
+            A list of requirements.
+        """
+        ret = []
+        for each in cls.get_requirements(target_os=target_os):
+            is_ignored = False
+            for ignored in cls.REQUIREMENTS_IGNORED_ON_UNINSTALL:
+                if each.startswith(ignored):
+                    is_ignored = True
+                    break
+            if not is_ignored:
+                ret.append(each)
+        return ret
+
     @classmethod
     def activate(cls) -> None:
         """Abstract method to activate the integration."""
@@ -134,3 +179,12 @@ class Integration(metaclass=IntegrationMeta):
             A list of new stack component flavors.
         """
         return []
+
+    @classmethod
+    def plugin_flavors(cls) -> List[Type["BasePluginFlavor"]]:
+        """Abstract method to declare new plugin flavors.
+
+        Returns:
+            A list of new plugin flavors.
+        """
+        return []

zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py

@@ -13,10 +13,9 @@
 #  permissions and limitations under the License.
 """Kaniko image builder flavor."""
 
-import json
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
-from pydantic import validator
+from pydantic import PositiveInt
 
 from zenml.image_builders import BaseImageBuilderConfig, BaseImageBuilderFlavor
 from zenml.integrations.kaniko import KANIKO_IMAGE_BUILDER_FLAVOR
@@ -29,6 +28,7 @@ KANIKO_EXECUTOR_IMAGE_TAG = "v1.9.1"
 DEFAULT_KANIKO_EXECUTOR_IMAGE = (
     f"gcr.io/kaniko-project/executor:{KANIKO_EXECUTOR_IMAGE_TAG}"
 )
+DEFAULT_KANIKO_POD_RUNNING_TIMEOUT = 300
 
 
 class KanikoImageBuilderConfig(BaseImageBuilderConfig):
@@ -47,6 +47,8 @@ class KanikoImageBuilderConfig(BaseImageBuilderConfig):
             Kaniko pod. This namespace will not be created and must already
             exist.
         executor_image: The image of the Kaniko executor to use.
+        pod_running_timeout: The timeout to wait until the pod is running
+            in seconds. Defaults to `300`.
         env: `env` section of the Kubernetes container spec.
         env_from: `envFrom` section of the Kubernetes container spec.
         volume_mounts: `volumeMounts` section of the Kubernetes container spec.
@@ -67,6 +69,7 @@ class KanikoImageBuilderConfig(BaseImageBuilderConfig):
     kubernetes_context: str
     kubernetes_namespace: str = "zenml-kaniko"
     executor_image: str = DEFAULT_KANIKO_EXECUTOR_IMAGE
+    pod_running_timeout: PositiveInt = DEFAULT_KANIKO_POD_RUNNING_TIMEOUT
 
     env: List[Dict[str, Any]] = []
     env_from: List[Dict[str, Any]] = []
@@ -78,47 +81,6 @@ class KanikoImageBuilderConfig(BaseImageBuilderConfig):
 
     executor_args: List[str] = []
 
-    @validator(
-        "env",
-        "env_from",
-        "volume_mounts",
-        "volumes",
-        "executor_args",
-        pre=True,
-    )
-    def _convert_json_string(
-        cls, value: Union[None, str, List[Any]]
-    ) -> Optional[List[Any]]:
-        """Converts potential JSON strings passed via the CLI to lists.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `List` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a list.
-        """
-        if isinstance(value, str):
-            try:
-                list_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(list_, List):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a list."
-                )
-
-            return list_
-        elif isinstance(value, List) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a list.")
-
 
 class KanikoImageBuilderFlavor(BaseImageBuilderFlavor):
     """Kaniko image builder flavor."""