Building Custom Orchestrators

from zenml.orchestrators import BaseOrchestratorConfig
from pydantic import Field

class MyOrchestratorConfig(BaseOrchestratorConfig):
    """Configuration for my custom orchestrator."""
    
    region: str = Field(
        description="Cloud region for execution"
    )
    
    execution_role: str = Field(
        description="IAM role for pipeline execution"
    )
    
    @property
    def is_synchronous(self) -> bool:
        """Whether the orchestrator runs synchronously.
        
        Returns:
            False for cloud orchestrators that submit async jobs.
        """
        return False
    
    @property
    def is_schedulable(self) -> bool:
        """Whether the orchestrator supports scheduling.
        
        Returns:
            True if your backend supports cron schedules.
        """
        return True

import os
from typing import Dict, Optional
from zenml.orchestrators import ContainerizedOrchestrator
from zenml.models import PipelineRunResponse, PipelineSnapshotResponse
from zenml.stack import Stack

class MyOrchestrator(ContainerizedOrchestrator):
    """Custom orchestrator implementation."""
    
    @property
    def config(self) -> MyOrchestratorConfig:
        """Returns the orchestrator configuration."""
        return cast(MyOrchestratorConfig, self._config)
    
    def get_orchestrator_run_id(self) -> str:
        """Returns the run ID of the active orchestrator run.
        
        This is the most critical method to implement correctly.
        
        Requirements:
        - Must return the SAME value for all steps in a pipeline run
        - Must be UNIQUE across different runs
        - Limited to ~250 characters (MySQL column limit)
        
        Returns:
            The unique orchestrator run ID.
        """
        # For static pipelines: read from environment variable
        # set by your orchestrator when launching steps
        if "MY_ORCHESTRATOR_RUN_ID" in os.environ:
            return os.environ["MY_ORCHESTRATOR_RUN_ID"]
        
        # For dynamic pipelines: use container-specific identifier
        # This only needs to be unique within the orchestration container
        return os.environ.get("CONTAINER_ID", "unknown")
    
    def submit_pipeline(
        self,
        snapshot: PipelineSnapshotResponse,
        stack: Stack,
        base_environment: Dict[str, str],
        step_environments: Dict[str, Dict[str, str]],
        placeholder_run: Optional[PipelineRunResponse] = None,
    ) -> Optional[SubmissionResult]:
        """Submit a static pipeline to the orchestration backend.
        
        Args:
            snapshot: Pipeline snapshot with step configurations
            stack: The stack the pipeline will run on
            base_environment: Environment variables for all steps
            step_environments: Per-step environment variables
            placeholder_run: Optional placeholder for the pipeline run
            
        Returns:
            Optional submission result with wait callback.
        """
        from zenml.orchestrators import SubmissionResult
        
        # Generate unique run ID for this pipeline execution
        orchestrator_run_id = self._generate_run_id()
        
        # Prepare steps for execution
        for invocation_id, step in snapshot.step_configurations.items():
            # Get the Docker image for this step
            image = self.get_image(snapshot, step.config.name)
            
            # Get environment variables for this step
            env = step_environments[invocation_id].copy()
            env["MY_ORCHESTRATOR_RUN_ID"] = orchestrator_run_id
            
            # Get resource requirements
            resources = step.config.resource_settings
            
            # Submit step to your backend
            self._submit_step(
                step_name=step.config.name,
                image=image,
                environment=env,
                cpu_count=resources.cpu_count,
                memory=resources.memory,
                gpu_count=resources.gpu_count,
            )
        
        # Return result with optional wait callback
        return SubmissionResult(
            wait_for_completion=lambda: self._wait_for_run(orchestrator_run_id)
        )
    
    def submit_dynamic_pipeline(
        self,
        snapshot: PipelineSnapshotResponse,
        stack: Stack,
        environment: Dict[str, str],
        placeholder_run: Optional[PipelineRunResponse] = None,
    ) -> Optional[SubmissionResult]:
        """Submit a dynamic pipeline to the orchestration backend.
        
        For dynamic pipelines, submit an orchestration container that will
        determine and launch the actual pipeline steps at runtime.
        
        Args:
            snapshot: Pipeline snapshot
            stack: The stack the pipeline will run on
            environment: Environment variables for orchestration container
            placeholder_run: Optional placeholder for the pipeline run
            
        Returns:
            Optional submission result.
        """
        from zenml.orchestrators import SubmissionResult
        
        # Get image for orchestration container
        image = self.get_image(snapshot)
        
        # Submit orchestration container
        job_id = self._submit_orchestration_job(
            image=image,
            environment=environment,
        )
        
        return SubmissionResult(
            wait_for_completion=lambda: self._wait_for_job(job_id)
        )
    
    def _generate_run_id(self) -> str:
        """Generate a unique run ID for the pipeline."""
        import uuid
        return f"run-{uuid.uuid4().hex[:12]}"
    
    def _submit_step(
        self,
        step_name: str,
        image: str,
        environment: Dict[str, str],
        cpu_count: Optional[float],
        memory: Optional[str],
        gpu_count: Optional[int],
    ) -> None:
        """Submit a single step to your execution backend."""
        # Implementation depends on your backend
        # Examples: boto3 for AWS, google.cloud for GCP, kubernetes for K8s
        pass
    
    def _submit_orchestration_job(self, image: str, environment: Dict[str, str]) -> str:
        """Submit orchestration container for dynamic pipelines."""
        # Submit a job that will run the orchestration logic
        pass
    
    def _wait_for_run(self, run_id: str) -> None:
        """Wait for a pipeline run to complete."""
        # Poll your backend until the run finishes
        pass
    
    def _wait_for_job(self, job_id: str) -> None:
        """Wait for an orchestration job to complete."""
        pass

from typing import Type, TYPE_CHECKING
from zenml.orchestrators import BaseOrchestratorFlavor

if TYPE_CHECKING:
    from zenml.orchestrators import BaseOrchestrator

class MyOrchestratorFlavor(BaseOrchestratorFlavor):
    """Flavor for my custom orchestrator."""
    
    @property
    def name(self) -> str:
        """The flavor name."""
        return "my_orchestrator"
    
    @property
    def config_class(self) -> Type[MyOrchestratorConfig]:
        """The configuration class."""
        return MyOrchestratorConfig
    
    @property
    def implementation_class(self) -> Type["BaseOrchestrator"]:
        """The implementation class."""
        from my_integration.orchestrators import MyOrchestrator
        return MyOrchestrator

import socket
import os

def get_orchestrator_run_id(self) -> str:
    # Static pipelines: use pre-set environment variable
    try:
        return os.environ["ZENML_KUBERNETES_RUN_ID"]
    except KeyError:
        # Dynamic pipelines: use pod name (unique per container)
        return socket.gethostname()

import os
import json

def get_orchestrator_run_id(self) -> str:
    # Check environment variables
    for env_var in ["MY_PLATFORM_JOB_ID", "MY_PLATFORM_EXECUTION_ARN"]:
        if env_var in os.environ:
            return os.environ[env_var]
    
    # Fall back to config file if your platform provides one
    config_path = "/opt/platform/config.json"
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = json.load(f)
            return config["job_id"]
    
    raise RuntimeError("Unable to determine orchestrator run ID")

from zenml.config.resource_settings import ResourceSettings

def _apply_resources(self, step_config, backend_job_config):
    """Apply resource settings to backend job configuration."""
    resources: ResourceSettings = step_config.resource_settings
    
    # CPU configuration
    if resources.cpu_count:
        backend_job_config["cpu"] = resources.cpu_count
    
    # Memory configuration
    if resources.memory:
        # Convert to backend-specific format
        memory_mb = resources.get_memory(unit="MB")
        backend_job_config["memory_mb"] = int(memory_mb)
    
    # GPU configuration
    if resources.gpu_count:
        backend_job_config["gpu_count"] = resources.gpu_count
        # Some backends need GPU type specification
        backend_job_config["gpu_type"] = "nvidia-tesla-t4"

from zenml import pipeline, step

@step
def load_data() -> int:
    return 42

@step
def process_data(data: int) -> int:
    return data * 2

@pipeline
def test_pipeline():
    data = load_data()
    process_data(data)

# Run with your orchestrator
test_pipeline.with_options(
    orchestrator="my_orchestrator"
)()