Add support for all local containerized orchestrators

Author: schustmi

src/zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py

@@ -119,6 +119,15 @@ class AirflowOrchestratorConfig(
 
     local: bool = True
 
+    @property
+    def is_local(self) -> bool:
+        """Checks if this stack component is running locally.
+
+        Returns:
+            True if this config is for a local component, False otherwise.
+        """
+        return self.local
+
     @property
     def is_schedulable(self) -> bool:
         """Whether the orchestrator is schedulable or not.

src/zenml/orchestrators/containerized_orchestrator.py

@@ -14,9 +14,11 @@
 """Containerized orchestrator class."""
 
 from abc import ABC
-from typing import List, Optional
+from typing import List, Optional, Set
 
+import zenml
 from zenml.config.build_configuration import BuildConfiguration
+from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import ORCHESTRATOR_DOCKER_IMAGE_KEY
 from zenml.models import PipelineSnapshotBase, PipelineSnapshotResponse
 from zenml.orchestrators import BaseOrchestrator
@@ -25,6 +27,22 @@
 class ContainerizedOrchestrator(BaseOrchestrator, ABC):
     """Base class for containerized orchestrators."""
 
+    @property
+    def requirements(self) -> Set[str]:
+        """Set of PyPI requirements for the component.
+
+        Returns:
+            A set of PyPI requirements for the component.
+        """
+        requirements = super().requirements
+
+        if self.config.is_local and GlobalConfiguration().uses_sql_store:
+            # If we're directly connected to a DB, we need to install the
+            # `local` extra in the Docker image to include the DB dependencies.
+            requirements.add(f"'zenml[local]=={zenml.__version__}'")
+
+        return requirements
+
     @staticmethod
     def get_image(
         snapshot: "PipelineSnapshotResponse",

src/zenml/orchestrators/local_docker/local_docker_orchestrator.py

@@ -17,12 +17,11 @@
 import os
 import sys
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Type, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
 from uuid import uuid4
 
 from docker.errors import ContainerError
 
-import zenml
 from zenml.config.base_settings import BaseSettings
 from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import (
@@ -84,20 +83,6 @@ def validator(self) -> Optional[StackValidator]:
             required_components={StackComponentType.IMAGE_BUILDER}
         )
 
-    @property
-    def requirements(self) -> Set[str]:
-        """Set of PyPI requirements for the component.
-
-        Returns:
-            A set of PyPI requirements for the component.
-        """
-        if GlobalConfiguration().uses_sql_store:
-            # If we're directly connected to a DB, we need to install the
-            # `local` extra in the Docker image to include the DB dependencies.
-            return {f'"zenml[local]=={zenml.__version__}"'}
-
-        return set()
-
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.