Building Custom Connectors

This guide walks you through building and deploying custom connectors within flyte-sdk. By the end, you will understand how to define a custom connector, integrate it with Flyte tasks, and configure its deployment environment.

Prerequisites

To follow this tutorial, you need a basic understanding of Flyte tasks and Python asynchronous programming. Ensure you have the flyte-sdk installed.

Core Connector Components

Custom connectors in flyte-sdk are built around three core components: AsyncConnector, ResourceMeta, and Resource.

AsyncConnector

The flyte.connectors.AsyncConnector is an abstract base class that defines the interface for all asynchronous connectors. Any custom connector you build must inherit from this class and implement its abstract methods. The connector service uses this interface to interact with your custom logic for creating, monitoring, and managing external jobs.

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional

from flyteidl2.core.tasks_pb2 import TaskTemplate as ProtoTaskTemplate

from flyte.connectors._connector import Resource, ResourceMeta


class AsyncConnector(ABC):
    name: str
    task_type_name: str
    task_type_version: int = 0
    metadata_type: ResourceMeta

    @abstractmethod
    async def create(
        self,
        task_template: ProtoTaskTemplate,
        output_prefix: str,
        inputs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> ResourceMeta:
        raise NotImplementedError

    @abstractmethod
    async def get(self, resource_meta: ResourceMeta, **kwargs) -> Resource:
        raise NotImplementedError

    @abstractmethod
    async def delete(self, resource_meta: ResourceMeta, **kwargs):
        raise NotImplementedError

    async def get_metrics(self, resource_meta: ResourceMeta, **kwargs):
        raise NotImplementedError

    async def get_logs(self, resource_meta: ResourceMeta, **kwargs):
        raise NotImplementedError

Key attributes of AsyncConnector:

• name: A human-readable name for the connector.
• task_type_name: A string identifier for the task type this connector handles.
• task_type_version: An integer version for the task type.
• metadata_type: A reference to a ResourceMeta subclass that defines the metadata structure for jobs managed by this connector.

The essential methods you must implement are:

• create: Initiates an external job and returns a ResourceMeta instance containing its metadata.
• get: Retrieves the current status and, optionally, outputs of the external job, returning a Resource instance.
• delete: Terminates or cleans up the external job. This method should be idempotent.
• get_metrics and get_logs: (Optional) Provide metrics and logs for the external job.

ResourceMeta

The flyte.connectors.ResourceMeta is a dataclass used to store metadata about an external job. This metadata is specific to your connector and helps identify and manage the job. It provides encode() and decode() methods for serialization and deserialization, allowing the Flyte system to persist and retrieve job-specific information.

from dataclasses import dataclass, asdict
import json

# Assuming dataclass_from_dict and other necessary imports are available

@dataclass
class ResourceMeta:
    def encode(self) -> bytes:
        return json.dumps(asdict(self)).encode('utf-8')

    @classmethod
    def decode(cls, data: bytes) -> "ResourceMeta":
        # Simplified for example, actual implementation uses dataclass_from_dict
        return cls(**json.loads(data.decode('utf-8')))

When creating your custom connector, you will define a subclass of ResourceMeta to hold the unique identifiers and state necessary to interact with your external service.

Resource

The flyte.connectors.Resource dataclass represents the current status and potential outputs of an external job managed by a connector. The get method of your AsyncConnector implementation will return an instance of this class to update Flyte on the job