diff --git a/docs/book/stacks-and-components/auth-management/auth-management.md b/docs/book/stacks-and-components/auth-management/auth-management.md index b312970b7a7..f590278a448 100644 --- a/docs/book/stacks-and-components/auth-management/auth-management.md +++ b/docs/book/stacks-and-components/auth-management/auth-management.md @@ -101,6 +101,11 @@ zenml service-connector list-types ┃ │ │ 🌀 kubernetes-cluster │ service-account │ │ ┃ ┃ │ │ 🐳 docker-registry │ oauth2-token │ │ ┃ ┃ │ │ │ impersonation │ │ ┃ +┠──────────────────────────────┼───────────────┼───────────────────────┼──────────────────┼───────┼────────┨ +┃ HyperAI Service Connector │ 🤖 hyperai │ 🤖 hyperai-instance │ rsa-key │ ✅ │ ✅ ┃ +┃ │ │ │ dsa-key │ │ ┃ +┃ │ │ │ ecdsa-key │ │ ┃ +┃ │ │ │ ed25519-key │ │ ┃ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━┷━━━━━━━┷━━━━━━━━┛ ``` {% endcode %} diff --git a/docs/book/stacks-and-components/auth-management/hyperai-service-connector.md b/docs/book/stacks-and-components/auth-management/hyperai-service-connector.md new file mode 100644 index 00000000000..56458c1b64d --- /dev/null +++ b/docs/book/stacks-and-components/auth-management/hyperai-service-connector.md @@ -0,0 +1,57 @@ +--- +description: Configuring HyperAI Connectors to connect ZenML to HyperAI instances. +--- + +# HyperAI Service Connector + +The ZenML HyperAI Service Connector allows authenticating with a HyperAI instance for deployment of pipeline runs. This connector provides pre-authenticated Paramiko SSH clients to Stack Components that are linked to it. + +``` +$ zenml service-connector list-types --type hyperai +┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━┯━━━━━━━┯━━━━━━━━┓ +┃ NAME │ TYPE │ RESOURCE TYPES │ AUTH METHODS │ LOCAL │ REMOTE ┃ +┠───────────────────────────┼────────────┼────────────────────┼──────────────┼───────┼────────┨ +┃ HyperAI Service Connector │ 🤖 hyperai │ 🤖 hyperai-instance │ rsa-key │ ✅ │ ✅ ┃ +┃ │ │ │ dsa-key │ │ ┃ +┃ │ │ │ ecdsa-key │ │ ┃ +┃ │ │ │ ed25519-key │ │ ┃ +┗━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━┷━━━━━━━┷━━━━━━━━┛ +``` + +## Prerequisites +The HyperAI Service Connector is part of the HyperAI integration. It is necessary to install the integration in order to use this Service Connector: + +* `zenml integration install hyperai` installs the HyperAI integration + +## Resource Types +The HyperAI Service Connector supports HyperAI instances. + +## Authentication Methods +ZenML creates an SSH connection to the HyperAI instance in the background when using this Service Connector. It then provides these connections to stack components requiring them, such as the HyperAI Orchestrator. Multiple authentication methods are supported: + +1. RSA key based authentication. +2. DSA (DSS) key based authentication. +3. ECDSA key based authentication. +4. ED25519 key based authentication. + +{% hint style="warning" %} +SSH private keys configured in the connector will be distributed to all clients that use them to run pipelines with the HyperAI orchestrator. SSH keys are long-lived credentials that give unrestricted access to HyperAI instances. +{% endhint %} + +When configuring the Service Connector, it is required to provide at least one hostname via `hostnames` and the `username` with which to login. Optionally, it is possible to provide an `ssh_passphrase` if applicable. This way, it is possible to use the HyperAI service connector in multiple ways: + +1. Create one service connector per HyperAI instance with different SSH keys. +2. Configure a reused SSH key just once for multiple HyperAI instances, then select the individual instance when creating the HyperAI orchestrator component. + +## Auto-configuration + +{% hint style="info" %} +This Service Connector does not support auto-discovery and extraction of authentication credentials from HyperAI instances. If this feature is useful to you or your organization, please let us know by messaging us in [Slack](https://zenml.io/slack-invite) or [creating an issue on GitHub](https://github.com/zenml-io/zenml/issues). +{% endhint %} + +## Stack Components use + +The HyperAI Service Connector can be used by the HyperAI Orchestrator to deploy pipeline runs to HyperAI instances. + + +
ZenML Scarf
diff --git a/docs/book/stacks-and-components/component-guide/orchestrators/hyperai.md b/docs/book/stacks-and-components/component-guide/orchestrators/hyperai.md new file mode 100644 index 00000000000..d54dbdc41cc --- /dev/null +++ b/docs/book/stacks-and-components/component-guide/orchestrators/hyperai.md @@ -0,0 +1,90 @@ +--- +description: Orchestrating your pipelines to run on HyperAI.ai instances. +--- + +# HyperAI orchestrator +[HyperAI](https://www.hyperai.ai) is a cutting-edge cloud compute platform designed to make AI accessible for everyone. The HyperAI orchestrator is an [orchestrator](orchestrators.md) flavor that allows you to easily deploy your pipelines on HyperAI instances. + +{% hint style="warning" %} +This component is only meant to be used within the context of +a [remote ZenML deployment scenario](/docs/book/deploying-zenml/zenml-self-hosted/zenml-self-hosted.md). +Usage with a local ZenML deployment may lead to unexpected behavior! +{% endhint %} + +### When to use it + +You should use the HyperAI orchestrator if: + +* you're looking for a managed solution for running your pipelines. +* you're a HyperAI customer. + +### Prerequisites +You will need to do the following to start using the HyperAI orchestrator: + +* Have a running HyperAI instance. It must be accessible from the internet (or at least from the IP addresses of your ZenML users) and allow SSH key based access (passwords are not supported). +* Ensure that a recent version of Docker is installed. This version must include Docker Compose, meaning that the command `docker compose` works. +* Ensure that the appropriate [NVIDIA Driver](https://www.nvidia.com/en-us/drivers/unix/) is installed on the HyperAI instance (if not already installed by the HyperAI team). +* Ensure that the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) is installed and configured on the HyperAI instance. + +Note that it is possible to omit installing the NVIDIA Driver and NVIDIA Container Toolkit. However, you will then be unable to use the GPU from within your ZenML pipeline. Additionally, you will then need to disable GPU access within the container when configuring the Orchestrator component, or the pipeline will not start correctly. + +## How it works +The HyperAI orchestrator works with Docker Compose, which can be used to construct machine learning pipelines. +Under the hood, it creates a Docker Compose file which it then deploys and executes on the configured HyperAI instance. +For each ZenML pipeline step, it creates a service in this file. It uses the `service_completed_successfully` condition +to ensure that pipeline steps will only run if their connected upstream steps have successfully finished. + +If configured for it, the HyperAI orchestrator will connect the HyperAI instance to the stack's container registry to ensure +a smooth transfer of Docker images. + +### Scheduled pipelines + +[Scheduled pipelines](../../../user-guide/advanced-guide/pipelining-features/schedule-pipeline-runs.md) are supported by the HyperAI orchestrator. Currently, only cron expressions are supported via `cron_expression`. When pipeline runs are scheduled, they are added as a crontab entry +on the HyperAI instance. + +### How to deploy it +To use the HyperAI orchestrator, you must configure a HyperAI Service Connector in ZenML and link it to the HyperAI orchestrator +component. The service connector contains credentials with which ZenML connects to the HyperAI instance. + +Additionally, the HyperAI orchestrator must be used in a stack that contains a container registry and an image builder. + +### How to use it + +To use the HyperAI orchestrator, we must configure a HyperAI Service Connector first using one of its supported authentication +methods. For example, for authentication with an RSA-based key, create the service connector as follows: + +```shell +zenml service-connector register --type=hyperai --auth-method=rsa-key --base64_ssh_key= --hostnames=,,.., --username= +``` + +Hostnames are either DNS resolvable names or IP addresses. + +For example, if you have two servers - one at `1.2.3.4` and another at `4.3.2.1`, you could provide them as `--hostnames=1.2.3.4,4.3.2.1`. + +Optionally, it is possible to provide a passphrase for the key (`--ssh_passphrase`). + +Following registering the service connector, we can register the orchestrator and use it in our active stack: + +```shell +zenml orchestrator register --flavor=hyperai + +# Register and activate a stack with the new orchestrator +zenml stack register -o ... --set +``` + +You can now run any ZenML pipeline using the HyperAI orchestrator: + +```shell +python file_that_runs_a_zenml_pipeline.py +``` + +#### Enabling CUDA for GPU-backed hardware + +Note that if you wish to use this orchestrator to run steps on a GPU, you will need to +follow [the instructions on this page](/docs/book/user-guide/advanced-guide/infrastructure-management/scale-compute-to-the-cloud.md) to ensure +that it works. It requires adding some extra settings customization and is essential to enable CUDA for the GPU to +give its full acceleration. + + +
ZenML Scarf
+ diff --git a/docs/book/stacks-and-components/component-guide/orchestrators/orchestrators.md b/docs/book/stacks-and-components/component-guide/orchestrators/orchestrators.md index 9bbca83fcf4..139b1f83b31 100644 --- a/docs/book/stacks-and-components/component-guide/orchestrators/orchestrators.md +++ b/docs/book/stacks-and-components/component-guide/orchestrators/orchestrators.md @@ -38,6 +38,7 @@ Additional orchestrators are provided by integrations: | [SkypilotAWSOrchestrator](skypilot-vm.md) | `vm_aws` | `skypilot[aws]` | Runs your pipelines in AWS VMs using SkyPilot | | [SkypilotGCPOrchestrator](skypilot-vm.md) | `vm_gcp` | `skypilot[gcp]` | Runs your pipelines in GCP VMs using SkyPilot | | [SkypilotAzureOrchestrator](skypilot-vm.md) | `vm_azure` | `skypilot[azure]` | Runs your pipelines in Azure VMs using SkyPilot | +| [HyperAIOrchestrator](hyperai.md) | `hyperai` | `hyperai` | Runs your pipeline in HyperAI.ai instances. | [Custom Implementation](custom.md) | _custom_ | | Extend the orchestrator abstraction and provide your own implementation | If you would like to see the available flavors of orchestrators, you can use the command: diff --git a/docs/mocked_libs.json b/docs/mocked_libs.json index 07bb7c50427..0bd5108ae55 100644 --- a/docs/mocked_libs.json +++ b/docs/mocked_libs.json @@ -144,6 +144,7 @@ "neptune", "neuralprophet", "openai", + "paramiko", "polars", "pyarrow", "pyarrow.parquet", diff --git a/pyproject.toml b/pyproject.toml index 96e41dabfb8..0b7d39f1f45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,6 +152,7 @@ types-certifi = { version = "^2021.10.8.0", optional = true } types-croniter = { version = "^1.0.2", optional = true } types-futures = { version = "^3.3.1", optional = true } types-Markdown = { version = "^3.3.6", optional = true } +types-paramiko = { version = ">=3.4.0", optional = true } types-Pillow = { version = "^9.2.1", optional = true } types-protobuf = { version = "^3.18.0", optional = true } types-PyMySQL = { version = "^1.0.4", optional = true } @@ -232,6 +233,7 @@ dev = [ "types-croniter", "types-futures", "types-Markdown", + "types-paramiko", "types-Pillow", "types-protobuf", "types-PyMySQL", diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index 0049bc3029e..73fa3fb7013 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -915,12 +915,14 @@ def prompt_configuration( config_dict = {} for attr_name, attr_schema in config_schema.get("properties", {}).items(): title = attr_schema.get("title", attr_name) - attr_type = attr_schema.get("type", "string") + attr_type_name = attr_type = attr_schema.get("type", "string") + if attr_type == "array": + attr_type_name = "list (CSV or JSON)" title = f"[{attr_name}] {title}" required = attr_name in config_schema.get("required", []) hidden = attr_schema.get("format", "") == "password" subtitles: List[str] = [] - subtitles.append(attr_type) + subtitles.append(attr_type_name) if hidden: subtitles.append("secret") if required: @@ -938,6 +940,8 @@ def prompt_configuration( if hidden and not show_secrets: title += " is currently set to: [HIDDEN]" else: + if attr_type == "array": + existing_value = json.dumps(existing_value) title += f" is currently set to: '{existing_value}'" else: title += " is not currently set" diff --git a/src/zenml/integrations/__init__.py b/src/zenml/integrations/__init__.py index 247282a1b06..786e4d86f77 100644 --- a/src/zenml/integrations/__init__.py +++ b/src/zenml/integrations/__init__.py @@ -35,6 +35,7 @@ GreatExpectationsIntegration, ) from zenml.integrations.huggingface import HuggingfaceIntegration # noqa +from zenml.integrations.hyperai import HyperAIIntegration # noqa from zenml.integrations.kaniko import KanikoIntegration # noqa from zenml.integrations.kserve import KServeIntegration # noqa from zenml.integrations.kubeflow import KubeflowIntegration # noqa diff --git a/src/zenml/integrations/azure/__init__.py b/src/zenml/integrations/azure/__init__.py index 4d1ceac545d..1d7e2438ab9 100644 --- a/src/zenml/integrations/azure/__init__.py +++ b/src/zenml/integrations/azure/__init__.py @@ -43,7 +43,7 @@ class AzureIntegration(Integration): "azure-keyvault-keys", "azure-keyvault-secrets", "azure-identity==1.10.0", - "azureml-core==1.48.0", + "azureml-core==1.54.0.post1", "azure-mgmt-containerservice>=20.0.0", "azure-storage-blob==12.17.0", # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056 "kubernetes", diff --git a/src/zenml/integrations/constants.py b/src/zenml/integrations/constants.py index 8d654bee0a0..a833d08e2d3 100644 --- a/src/zenml/integrations/constants.py +++ b/src/zenml/integrations/constants.py @@ -30,6 +30,7 @@ GRAPHVIZ = "graphviz" KSERVE = "kserve" HUGGINGFACE = "huggingface" +HYPERAI = "hyperai" GREAT_EXPECTATIONS = "great_expectations" KANIKO = "kaniko" KUBEFLOW = "kubeflow" diff --git a/src/zenml/integrations/hyperai/__init__.py b/src/zenml/integrations/hyperai/__init__.py new file mode 100644 index 00000000000..936c1f78635 --- /dev/null +++ b/src/zenml/integrations/hyperai/__init__.py @@ -0,0 +1,53 @@ +# 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. +"""Initialization of the HyperAI integration.""" +from typing import List, Type + +from zenml.integrations.constants import HYPERAI +from zenml.integrations.integration import Integration +from zenml.stack import Flavor + +# Service connector constants +HYPERAI_CONNECTOR_TYPE = "hyperai" +HYPERAI_RESOURCE_TYPE = "hyperai-instance" + + +class HyperAIIntegration(Integration): + """Definition of HyperAI integration for ZenML.""" + + NAME = HYPERAI + REQUIREMENTS = [ + "paramiko>=3.4.0", + ] + + @classmethod + def activate(cls) -> None: + """Activates the integration.""" + from zenml.integrations.hyperai import service_connectors # noqa + + @classmethod + def flavors(cls) -> List[Type[Flavor]]: + """Declare the stack component flavors for the HyperAI integration. + + Returns: + List of stack component flavors for this integration. + """ + from zenml.integrations.hyperai.flavors import ( + HyperAIOrchestratorFlavor + ) + + return [HyperAIOrchestratorFlavor] + + +HyperAIIntegration.check_installation() diff --git a/src/zenml/integrations/hyperai/flavors/__init__.py b/src/zenml/integrations/hyperai/flavors/__init__.py new file mode 100644 index 00000000000..861bff7be1a --- /dev/null +++ b/src/zenml/integrations/hyperai/flavors/__init__.py @@ -0,0 +1,20 @@ +# 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. +"""Airflow integration flavors.""" + +from zenml.integrations.hyperai.flavors.hyperai_orchestrator_flavor import ( + HyperAIOrchestratorFlavor, +) + +__all__ = ["HyperAIOrchestratorFlavor"] diff --git a/src/zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py b/src/zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py new file mode 100644 index 00000000000..6acc64b839d --- /dev/null +++ b/src/zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py @@ -0,0 +1,161 @@ +# 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. +"""Implementation of the ZenML HyperAI orchestrator.""" + +from typing import TYPE_CHECKING, Dict, Optional, Type + +from zenml.config.base_settings import BaseSettings +from zenml.integrations.hyperai import HYPERAI_RESOURCE_TYPE +from zenml.logger import get_logger +from zenml.models import ServiceConnectorRequirements +from zenml.orchestrators import ( + BaseOrchestratorConfig, + BaseOrchestratorFlavor, +) + +if TYPE_CHECKING: + from zenml.integrations.hyperai.orchestrators import HyperAIOrchestrator + +logger = get_logger(__name__) + + +class HyperAIOrchestratorSettings(BaseSettings): + """HyperAI orchestrator settings. + + Attributes: + mounts_from_to: A dictionary mapping from paths on the HyperAI instance + to paths within the Docker container. This allows users to mount + directories from the HyperAI instance into the Docker container that runs + on it. + """ + + mounts_from_to: Dict[str, str] = {} + + +class HyperAIOrchestratorConfig( # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173 + BaseOrchestratorConfig, HyperAIOrchestratorSettings +): + """Configuration for the HyperAI orchestrator. + + Attributes: + container_registry_autologin: If True, the orchestrator will attempt to + automatically log in to the container registry specified in the stack + configuration on the HyperAI instance. This is useful if the container + registry requires authentication and the HyperAI instance has not been + manually logged in to the container registry. Defaults to `False`. + automatic_cleanup_pipeline_files: If True, the orchestrator will + automatically clean up old pipeline files that are on the HyperAI + instance. Pipeline files will be cleaned up if they are 7 days old or + older. Defaults to `True`. + gpu_enabled_in_container: If True, the orchestrator will enable GPU + support in the Docker container that runs on the HyperAI instance. + Defaults to `True`. + + """ + + container_registry_autologin: bool = False + automatic_cleanup_pipeline_files: bool = True + gpu_enabled_in_container: bool = True + + @property + def is_remote(self) -> bool: + """Checks if this stack component is running remotely. + + This designation is used to determine if the stack component can be + used with a local ZenML database or if it requires a remote ZenML + server. + + Returns: + True if this config is for a remote component, False otherwise. + """ + return True + + +class HyperAIOrchestratorFlavor(BaseOrchestratorFlavor): + """Flavor for the HyperAI orchestrator.""" + + @property + def name(self) -> str: + """Name of the orchestrator flavor. + + Returns: + Name of the orchestrator flavor. + """ + return "hyperai" + + @property + def service_connector_requirements( + self, + ) -> Optional[ServiceConnectorRequirements]: + """Service connector resource requirements for service connectors. + + Specifies resource requirements that are used to filter the available + service connector types that are compatible with this flavor. + + Returns: + Requirements for compatible service connectors, if a service + connector is required for this flavor. + """ + return ServiceConnectorRequirements( + resource_type=HYPERAI_RESOURCE_TYPE + ) + + @property + def docs_url(self) -> Optional[str]: + """A url to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A url to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A url to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/connectors/hyperai/hyperai.png" + + @property + def config_class(self) -> Type[BaseOrchestratorConfig]: + """Config class for the base orchestrator flavor. + + Returns: + The config class. + """ + return HyperAIOrchestratorConfig + + @property + def implementation_class(self) -> Type["HyperAIOrchestrator"]: + """Implementation class for this flavor. + + Returns: + Implementation class for this flavor. + """ + from zenml.integrations.hyperai.orchestrators import ( + HyperAIOrchestrator, + ) + + return HyperAIOrchestrator diff --git a/src/zenml/integrations/hyperai/orchestrators/__init__.py b/src/zenml/integrations/hyperai/orchestrators/__init__.py new file mode 100644 index 00000000000..26789a5331a --- /dev/null +++ b/src/zenml/integrations/hyperai/orchestrators/__init__.py @@ -0,0 +1,21 @@ +# 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 orchestrator.""" +from zenml.integrations.hyperai.orchestrators.hyperai_orchestrator import ( + HyperAIOrchestrator +) + +__all__ = [ + "HyperAIOrchestrator" +] diff --git a/src/zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py b/src/zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py new file mode 100644 index 00000000000..68993ce3ec9 --- /dev/null +++ b/src/zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py @@ -0,0 +1,418 @@ +# 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. +"""Implementation of the ZenML HyperAI orchestrator.""" + +import os +import re +import tempfile +from shlex import quote +from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast + +import paramiko +import yaml + +from zenml.config.base_settings import BaseSettings +from zenml.entrypoints import StepEntrypointConfiguration +from zenml.enums import StackComponentType +from zenml.integrations.hyperai.flavors.hyperai_orchestrator_flavor import ( + HyperAIOrchestratorConfig, + HyperAIOrchestratorSettings, +) +from zenml.logger import get_logger +from zenml.orchestrators import ( + ContainerizedOrchestrator, +) +from zenml.stack import Stack, StackValidator + +if TYPE_CHECKING: + from zenml.models import PipelineDeploymentResponse + +logger = get_logger(__name__) + +ENV_ZENML_HYPERAI_RUN_ID = "ZENML_HYPERAI_ORCHESTRATOR_RUN_ID" + + +class HyperAIOrchestrator(ContainerizedOrchestrator): + """Orchestrator responsible for running pipelines on HyperAI instances.""" + + @property + def config(self) -> HyperAIOrchestratorConfig: + """Returns the `HyperAIOrchestratorConfig` config. + + Returns: + The configuration. + """ + return cast(HyperAIOrchestratorConfig, self._config) + + @property + def settings_class(self) -> Optional[Type["BaseSettings"]]: + """Settings class for the HyperAI orchestrator. + + Returns: + The settings class. + """ + return HyperAIOrchestratorSettings + + @property + def validator(self) -> Optional[StackValidator]: + """Ensures there is an image builder in the stack. + + Returns: + A `StackValidator` instance. + """ + return StackValidator( + required_components={ + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, + } + ) + + def get_orchestrator_run_id(self) -> str: + """Returns the active orchestrator run id. + + Raises: + RuntimeError: If the environment variable specifying the run id + is not set. + + Returns: + The orchestrator run id. + """ + try: + return os.environ[ENV_ZENML_HYPERAI_RUN_ID] + except KeyError: + raise RuntimeError( + "Unable to read run id from environment variable " + f"{ENV_ZENML_HYPERAI_RUN_ID}." + ) + + def _validate_mount_path(self, path: str) -> str: + """Validates if a given string is in a valid path format. + + Args: + path: The path to be validated. + + Returns: + The path in a valid format. + + Raises: + RuntimeError: If the path is not in a valid format. + """ + # Define a regular expression pattern to match a valid path format + pattern = r'^(?:[a-zA-Z]:\\(\\[^\\/:*?"<>|]*)*$|^/([^\0]*)*$)' + + if bool(re.match(pattern, path)): + return path + else: + raise RuntimeError( + f"Path '{path}' is not in a valid format, so a mount cannot be established." + ) + + def _escape_shell_command(self, command: str) -> str: + """Escapes a shell command. + + Args: + command: The command to escape. + + Returns: + The escaped command. + """ + return quote(command) + + def prepare_or_run_pipeline( + self, + deployment: "PipelineDeploymentResponse", + stack: "Stack", + environment: Dict[str, str], + ) -> Any: + """Sequentially runs all pipeline steps in Docker containers. + + Assumes that: + - A HyperAI (hyperai.ai) instance is running on the configured IP address. + - The HyperAI instance has been configured to allow SSH connections from the + machine running the pipeline. + - Docker and Docker Compose are installed on the HyperAI instance. + - A key pair has been generated and the public key has been added to the + HyperAI instance's `authorized_keys` file. + - The private key is available in a HyperAI service connector linked to this + orchestrator. + + Args: + deployment: The pipeline deployment to prepare or run. + stack: The stack the pipeline will run on. + environment: Environment variables to set in the orchestration + environment. + + Raises: + RuntimeError: If a step fails. + """ + from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import ( + HyperAIServiceConnector, + ) + + # Basic Docker Compose definition + compose_definition: Dict[str, Any] = {"version": "3", "services": {}} + + # Get deployment id + deployment_id = deployment.id + + # Set environment + os.environ[ENV_ZENML_HYPERAI_RUN_ID] = str(deployment_id) + environment[ENV_ZENML_HYPERAI_RUN_ID] = str(deployment_id) + + # Add each step as a service to the Docker Compose definition + logger.info("Preparing pipeline steps for deployment.") + for step_name, step in deployment.step_configurations.items(): + # Get image + image = self.get_image(deployment=deployment, step_name=step_name) + + # Get settings + step_settings = cast( + HyperAIOrchestratorSettings, self.get_settings(step) + ) + + # Define container name as combination between deployment id and step name + container_name = f"{deployment_id}-{step_name}" + + # Make Compose service definition for step + compose_definition["services"][container_name] = { + "image": image, + "container_name": container_name, + "network_mode": "host", + "entrypoint": StepEntrypointConfiguration.get_entrypoint_command(), + "command": StepEntrypointConfiguration.get_entrypoint_arguments( + step_name=step_name, deployment_id=deployment.id + ), + "volumes": [ + "{}:{}".format( + self._validate_mount_path(mount_from), + self._validate_mount_path(mount_to), + ) + for mount_from, mount_to in step_settings.mounts_from_to.items() + ], + } + + # Depending on GPU setting, add GPU support to service definition + if self.config.gpu_enabled_in_container: + compose_definition["services"][container_name]["deploy"] = { + "resources": { + "reservations": { + "devices": [ + {"driver": "nvidia", "capabilities": ["gpu"]} + ] + } + } + } + + # Depending on whether it is a scheduled or a realtime pipeline, add + # potential .env file to service definition for deployment ID override. + if deployment.schedule: + # drop ZENML_HYPERAI_ORCHESTRATOR_RUN_ID from environment + del environment[ENV_ZENML_HYPERAI_RUN_ID] + compose_definition["services"][container_name]["env_file"] = [ + ".env" + ] + + compose_definition["services"][container_name][ + "environment" + ] = environment + + # Add dependency on upstream steps if applicable + upstream_steps = step.spec.upstream_steps + for upstream_step_name in upstream_steps: + upstream_container_name = ( + f"{deployment_id}-{upstream_step_name}" + ) + compose_definition["services"][container_name][ + "depends_on" + ] = { + upstream_container_name: { + "condition": "service_completed_successfully" + } + } + + # Convert into yaml + logger.info("Finalizing Docker Compose definition.") + compose_definition_yaml: str = yaml.dump(compose_definition) + + # Connect to configured HyperAI instance + logger.info( + "Connecting to HyperAI instance and placing Docker Compose file." + ) + paramiko_client: paramiko.SSHClient + if connector := self.get_connector(): + paramiko_client = connector.connect() + if paramiko_client is None: + raise RuntimeError( + "Expected to receive a `paramiko.SSHClient` object from the " + "linked connector, but got `None`. This likely originates from " + "a misconfigured service connector, typically caused by a wrong " + "SSH key type being selected. Please check your " + "`hyperai_orchestrator` configuration and make sure that the " + "`ssh_key_type` of its connected service connector is set to the " + "correct value." + ) + elif not isinstance(paramiko_client, paramiko.SSHClient): + raise RuntimeError( + f"Expected to receive a `paramiko.SSHClient` object from the " + f"linked connector, but got type `{type(paramiko_client)}`." + ) + else: + raise RuntimeError( + "You must link a HyperAI service connector to the orchestrator." + ) + + # Get container registry autologin setting + if self.config.container_registry_autologin: + logger.info( + "Attempting to automatically log in to container registry used by stack." + ) + + # Select stack container registry + container_registry = stack.container_registry + + # Raise error if no container registry is found + if not container_registry: + raise RuntimeError( + "Unable to find container registry in stack." + ) + + # Get container registry credentials from its config + credentials = container_registry.credentials + if credentials is None: + raise RuntimeError( + "The container registry in the active stack has no " + "credentials or service connector configured, but the " + "HyperAI orchestrator is set to autologin to the container " + "registry. Please configure the container registry with " + "credentials or turn off the `container_registry_autologin` " + "setting in the HyperAI orchestrator configuration." + ) + + container_registry_url = container_registry.config.uri + ( + container_registry_username, + container_registry_password, + ) = credentials + + # Escape inputs + container_registry_username = self._escape_shell_command( + container_registry_username + ) + container_registry_url = self._escape_shell_command( + container_registry_url + ) + + # Log in to container registry using --password-stdin + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"docker login -u {container_registry_username} " + f"--password-stdin {container_registry_url}" + ) + # Send the password to stdin + stdin.channel.send( + f"{container_registry_password}\n".encode("utf-8") + ) + stdin.channel.shutdown_write() + + # Log stdout + for line in stdout.readlines(): + logger.info(line) + + # Get username from connector + assert isinstance(connector, HyperAIServiceConnector) + username = connector.config.username + + # Set up pipeline-runs directory if it doesn't exist + nonscheduled_directory_name = self._escape_shell_command( + f"/home/{username}/pipeline-runs" + ) + directory_name = ( + nonscheduled_directory_name + if not deployment.schedule + else self._escape_shell_command( + f"/home/{username}/scheduled-pipeline-runs" + ) + ) + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"mkdir -p {directory_name}" + ) + + # Get pipeline run id and create directory for it + orchestrator_run_id = self.get_orchestrator_run_id() + directory_name = self._escape_shell_command( + f"{directory_name}/{orchestrator_run_id}" + ) + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"mkdir -p {directory_name}" + ) + + # Remove all folders from nonscheduled pipelines if they are 7 days old or older + if self.config.automatic_cleanup_pipeline_files: + logger.info( + "Cleaning up old pipeline files on HyperAI instance. This may take a while." + ) + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"find {nonscheduled_directory_name} -type d -ctime +7 -exec rm -rf {{}} +" + ) + + # Create temporary file and write Docker Compose file to it + with tempfile.NamedTemporaryFile(mode="w", delete=True) as f: + # Write Docker Compose file to temporary file + with f.file as f_: + f_.write(compose_definition_yaml) + + # Scp Docker Compose file to HyperAI instance + try: + scp_client = paramiko_client.open_sftp() + scp_client.put(f.name, f"{directory_name}/docker-compose.yaml") + scp_client.close() + except FileNotFoundError: + raise RuntimeError( + "Failed to write Docker Compose file to HyperAI instance. Does the user have permissions to write?" + ) + + # Run or schedule Docker Compose file depending on settings + if not deployment.schedule: + logger.info( + "Starting ZenML pipeline on HyperAI instance. Depending on the size of your container image, this may take a while..." + ) + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"cd {directory_name} && docker compose up -d" + ) + + # Log errors in case of failure + for line in stderr.readlines(): + logger.info(line) + else: + # Get cron expression for scheduled pipeline + cron_expression = deployment.schedule.cron_expression + if not cron_expression: + raise RuntimeError( + "A cron expression is required for scheduled pipelines." + ) + expected_cron_pattern = r"^(?:(?:[0-9]|[1-5][0-9]|60)(?:,(?:[0-9]|[1-5][0-9]|60))*|[*](?:\/[1-9][0-9]*)?)(?:[ \t]+(?:(?:[0-9]|[0-5][0-9]|60)(?:,(?:[0-9]|[0-5][0-9]|60))*|[*](?:\/[1-9][0-9]*)?)){4}$" + if not re.match(expected_cron_pattern, cron_expression): + raise RuntimeError( + f"The cron expression '{cron_expression}' is not in a valid format." + ) + + # Log about scheduling + logger.info("Scheduling ZenML pipeline on HyperAI instance.") + logger.info(f"Cron expression: {cron_expression}") + + # Create cron job for scheduled pipeline on HyperAI instance + stdin, stdout, stderr = paramiko_client.exec_command( # nosec + f"(crontab -l ; echo '{cron_expression} cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}=\"{deployment_id}_$(date +\%s)\" > .env && docker compose up -d') | crontab -" + ) + + logger.info("Pipeline scheduled successfully.") diff --git a/src/zenml/integrations/hyperai/service_connectors/__init__.py b/src/zenml/integrations/hyperai/service_connectors/__init__.py new file mode 100644 index 00000000000..1b641cb2970 --- /dev/null +++ b/src/zenml/integrations/hyperai/service_connectors/__init__.py @@ -0,0 +1,20 @@ +# 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.""" + +from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import ( + HyperAIServiceConnector, +) + +__all__ = ["HyperAIServiceConnector"] diff --git a/src/zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py b/src/zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py new file mode 100644 index 00000000000..12680850b78 --- /dev/null +++ b/src/zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py @@ -0,0 +1,373 @@ +# 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, SecretStr + +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 + +logger = get_logger(__name__) + + +class HyperAICredentials(AuthenticationConfig): + """HyperAI client authentication credentials.""" + + base64_ssh_key: SecretStr = Field( + title="SSH key (base64)", + ) + ssh_passphrase: Optional[SecretStr] = 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 diff --git a/src/zenml/models/v2/core/service_connector.py b/src/zenml/models/v2/core/service_connector.py index 08e189d2c9f..cbc6938d610 100644 --- a/src/zenml/models/v2/core/service_connector.py +++ b/src/zenml/models/v2/core/service_connector.py @@ -829,6 +829,7 @@ def _validate_and_configure_resources( "required", [] ) secret = attr_schema.get("format", "") == "password" + attr_type = attr_schema.get("type", "string") value = configuration.get(attr_name, secrets.get(attr_name)) if required: if value is None: @@ -846,6 +847,11 @@ def _validate_and_configure_resources( else: update_connector_metadata.secrets[attr_name] = SecretStr(value) else: + if attr_type == "array" and isinstance(value, str): + try: + value = json.loads(value) + except json.decoder.JSONDecodeError: + value = value.split(",") update_connector_metadata.configuration[attr_name] = value # Warn about attributes that are not part of the configuration schema diff --git a/src/zenml/service_connectors/service_connector_registry.py b/src/zenml/service_connectors/service_connector_registry.py index fdb580b0242..dd25bc969e7 100644 --- a/src/zenml/service_connectors/service_connector_registry.py +++ b/src/zenml/service_connectors/service_connector_registry.py @@ -245,5 +245,12 @@ def register_builtin_service_connectors(self) -> None: except ImportError as e: logger.warning(f"Could not import Docker service connector: {e}.") + try: + from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import ( # noqa + HyperAIServiceConnector, + ) + except ImportError as e: + logger.warning(f"Could not import HyperAI service connector: {e}.") + service_connector_registry = ServiceConnectorRegistry() diff --git a/tests/integration/integrations/hyperai/__init__.py b/tests/integration/integrations/hyperai/__init__.py new file mode 100644 index 00000000000..cd90a82cfc2 --- /dev/null +++ b/tests/integration/integrations/hyperai/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/integration/integrations/hyperai/orchestrators/__init__.py b/tests/integration/integrations/hyperai/orchestrators/__init__.py new file mode 100644 index 00000000000..cd90a82cfc2 --- /dev/null +++ b/tests/integration/integrations/hyperai/orchestrators/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/integration/integrations/hyperai/orchestrators/test_hyperai_orchestrator.py b/tests/integration/integrations/hyperai/orchestrators/test_hyperai_orchestrator.py new file mode 100644 index 00000000000..b0a7fa88e1f --- /dev/null +++ b/tests/integration/integrations/hyperai/orchestrators/test_hyperai_orchestrator.py @@ -0,0 +1,89 @@ +# 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. +from datetime import datetime +from uuid import uuid4 + +from zenml.enums import StackComponentType +from zenml.integrations.hyperai.flavors.hyperai_orchestrator_flavor import ( + HyperAIOrchestratorConfig, +) +from zenml.integrations.hyperai.orchestrators.hyperai_orchestrator import ( + HyperAIOrchestrator, +) + + +def test_hyperai_orchestrator_attributes(): + """Tests that the basic attributes of the HyperAI orchestrator are set correctly.""" + orchestrator = HyperAIOrchestrator( + name="", + id=uuid4(), + config=HyperAIOrchestratorConfig(), + flavor="hyperai", + type=StackComponentType.ORCHESTRATOR, + user=uuid4(), + workspace=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + + assert orchestrator.type == StackComponentType.ORCHESTRATOR + assert orchestrator.flavor == "hyperai" + assert orchestrator.config.is_remote is True + assert orchestrator.config.container_registry_autologin is False + assert orchestrator.config.automatic_cleanup_pipeline_files is True + assert orchestrator.config.gpu_enabled_in_container is True + + +def test_validate_mount_path(): + """Tests that only valid mount paths are accepted by the HyperAI orchestrator.""" + orchestrator = HyperAIOrchestrator( + name="", + id=uuid4(), + config=HyperAIOrchestratorConfig(), + flavor="hyperai", + type=StackComponentType.ORCHESTRATOR, + user=uuid4(), + workspace=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + + # Valid POSIX path + valid_posix_path = "/mnt/hello/there" + assert ( + orchestrator._validate_mount_path(valid_posix_path) == valid_posix_path + ) + + # Valid Windows path + valid_windows_path = r"C:\\Users\\user\\Documents" + assert ( + orchestrator._validate_mount_path(valid_windows_path) + == valid_windows_path + ) + + # Invalid POSIX path + invalid_posix_path = "echo '>something>' ; /mnt/hello/there/.." + try: + orchestrator._validate_mount_path(invalid_posix_path) + except RuntimeError: + pass + + # Invalid Windows path + invalid_windows_path = ( + "set SOMETHING=123; C:\\Users\\user\\Documents\\..\\file.txt" + ) + try: + orchestrator._validate_mount_path(invalid_windows_path) + except RuntimeError: + pass