Develop a custom orchestrator

Learning how to develop a custom orchestrator.

Before diving into the specifics of this component type, it is beneficial to familiarize yourself with our general guide to writing custom component flavors in ZenML. This guide provides an essential understanding of ZenML's component flavor concepts.

Base Implementation

ZenML aims to enable orchestration with any orchestration tool. This is where the BaseOrchestrator comes into play. It abstracts away many of the ZenML-specific details from the actual implementation and exposes a simplified interface:

from abc import ABC, abstractmethod

from typing import Any, Dict, Type

from zenml.models import PipelineDeploymentResponseModel

from zenml.enums import StackComponentType

from zenml.stack import StackComponent, StackComponentConfig, Stack, Flavor

class BaseOrchestratorConfig(StackComponentConfig):

"""Base class for all ZenML orchestrator configurations."""

class BaseOrchestrator(StackComponent, ABC):

"""Base class for all ZenML orchestrators"""

@abstractmethod

def prepare_or_run_pipeline(

self,

deployment: PipelineDeploymentResponseModel,

stack: Stack,

environment: Dict[str, str],

) -> Any:

"""Prepares and runs the pipeline outright or returns an intermediate

pipeline representation that gets deployed.

"""

@abstractmethod

def get_orchestrator_run_id(self) -> str:

"""Returns the run id of the active orchestrator run.

Important: This needs to be a unique ID and return the same value for

all steps of a pipeline run.

Returns:

The orchestrator run id.

"""

class BaseOrchestratorFlavor(Flavor):

"""Base orchestrator for all ZenML orchestrator flavors."""

@property

@abstractmethod

def name(self):

"""Returns the name of the flavor."""

@property

def type(self) -> StackComponentType:

"""Returns the flavor type."""

return StackComponentType.ORCHESTRATOR

@property

def config_class(self) -> Type[BaseOrchestratorConfig]: