HyperAI integration: orchestrator and service connector

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",

src/zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py

@@ -95,7 +95,7 @@ def get_orchestrator_run_id(self) -> str:
                 f"{ENV_ZENML_HYPERAI_RUN_ID}."
             )
 
-    def _validate_mount_path(self, path) -> str:
+    def _validate_mount_path(self, path: str) -> str:
         """Validates if a given string is in a valid path format.
 
         Args:
@@ -144,8 +144,12 @@ def prepare_or_run_pipeline(
         Raises:
             RuntimeError: If a step fails.
         """
+        from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import (
+            HyperAIServiceConnector,
+        )
+
         # Basic Docker Compose definition
-        compose_definition = {"version": "3", "services": {}}
+        compose_definition: Dict[str, Any] = {"version": "3", "services": {}}
 
         # Get deployment id
         deployment_id = deployment.id
@@ -226,7 +230,7 @@ def prepare_or_run_pipeline(
 
         # Convert into yaml
         logger.info("Finalizing Docker Compose definition.")
-        compose_definition = yaml.dump(compose_definition)
+        compose_definition_yaml = yaml.dump(compose_definition)
 
         # Connect to configured HyperAI instance
         logger.info(
@@ -271,22 +275,27 @@ def prepare_or_run_pipeline(
                 )
 
             # Get container registry credentials from its config
-            container_registry_url = container_registry.config.uri
-            (
-                container_registry_username,
-                container_registry_password,
-            ) = container_registry.credentials
-
-            # Log in to container registry using --password-stdin
-            stdin, stdout, stderr = paramiko_client.exec_command(
-                f"docker login -u {container_registry_username} --password-stdin {container_registry_url} <<< {container_registry_password}"
-            )
+            credentials = container_registry.credentials
+            if credentials is not None:
+                container_registry_url = container_registry.config.uri
+                (
+                    container_registry_username,
+                    container_registry_password,
+                ) = credentials
+
+                # Log in to container registry using --password-stdin
+                _, stdout, stderr = paramiko_client.exec_command(
+                    f"docker login -u {container_registry_username} "
+                    f"--password-stdin {container_registry_url} <<< "
+                    f"{container_registry_password}"
+                )
 
-            # Log stdout
-            for line in stdout.readlines():
-                logger.info(line)
+                # 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
@@ -320,7 +329,7 @@ def prepare_or_run_pipeline(
         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)
+                f_.write(compose_definition_yaml)
 
             # Scp Docker Compose file to HyperAI instance
             scp_client = paramiko_client.open_sftp()

src/zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py

@@ -18,7 +18,7 @@
 """
 import base64
 import io
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Type
 
 import paramiko
 from pydantic import Field, SecretStr
@@ -167,7 +167,7 @@ def _get_connector_type(cls) -> ServiceConnectorTypeModel:
         """
         return HYPERAI_SERVICE_CONNECTOR_TYPE_SPEC
 
-    def _paramiko_key_type_given_auth_method(self) -> paramiko.PKey:
+    def _paramiko_key_type_given_auth_method(self) -> Type[paramiko.PKey]:
         """Get the Paramiko key type given the authentication method.
 
         Returns: