Skip to main content

Async Connectors and Service Architecture

Connector lifecycle at a glance​

A connector-backed task carries a provider-defined identity from submission through polling:

Task + AsyncConnectorExecutorMixin
│ task_type, task_type_version
▼
ConnectorRegistry.get_connector(...)
│
├─ connector.create(TaskTemplate, inputs, output_prefix, ...)
│ └─ ResourceMeta.encode() ──► gRPC CreateTaskResponse
└─ connector.get(ResourceMeta.decode(...)) ──► Resource
└─ terminal phase + outputs/log links

AsyncConnector defines the asynchronous provider contract. ConnectorRegistry resolves an instance by (task_type_name, task_type_version). The local executor and gRPC service use the same registry. The public exports are available from flyte.connectors; implementation is in _connector.py and _server.py. Built-in BigQueryTask, DatabricksFunctionTask, and Snowflake integrations inherit AsyncConnectorExecutorMixin.

Implement a connector​

Start with a dataclass containing the external system's durable job identity. The custom connector in examples/apps/custom_connector_app/my_connector/connector.py stores a job ID and creation time:

import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Optional

from flyteidl2.core.execution_pb2 import TaskExecution

from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta


@dataclass
class BatchJobMetadata(ResourceMeta):
job_id: str
created_at: float


class BatchJobConnector(AsyncConnector):
"""Simulates an external batch job service."""

name = "Batch Job Connector"
task_type_name = "batch_job"
metadata_type = BatchJobMetadata

async def create(
self,
task_template,
inputs: Optional[Dict[str, Any]] = None,
**kwargs,
) -> BatchJobMetadata:
job_id = str(uuid.uuid4())[:8]
return BatchJobMetadata(job_id=job_id, created_at=time.time())

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

async def delete(self, resource_meta: BatchJobMetadata, **kwargs):
return None


ConnectorRegistry.register(BatchJobConnector())

The three required methods are asynchronous: create(task_template, output_prefix, inputs, task_execution_metadata, **kwargs) submits work and returns ResourceMeta; get(resource_meta, **kwargs) reports a Resource and may include outputs; and delete(resource_meta, **kwargs) deletes the external job. The delete docstring specifies an idempotent operation that raises if deletion fails, but the framework does not enforce idempotency.

Set task_type_version when a connector needs a version other than the default 0, and set metadata_type to the concrete metadata dataclass. The base get_metrics and get_logs methods raise NotImplementedError, so implement them when clients will request those endpoints.

Metadata and result envelopes​

ResourceMeta.encode() serializes the dataclass with json.dumps(asdict(...)) and UTF-8 encoding. decode() parses the JSON and reconstructs the concrete class with dataclass_from_dict:

create → provider metadata object → encode() → bytes in response/request
get/delete ← decode() ← bytes

Changing persisted metadata field names or types can make older bytes undecodable. Resource requires a Flyte TaskExecution.Phase and can carry a message, log links, native output dictionary, and custom information. get_resource_proto converts native outputs to Flyte literals, log links to protobuf task logs, and custom information to a protobuf Struct.

Use connectors.utils.convert_to_flyte_phase() for provider strings. It maps failed/timeout/canceled/skipped to FAILED, internal_error to RETRYABLE_FAILED, done/succeeded/success/completed to SUCCEEDED, running/terminating to RUNNING, and pending to INITIALIZING, case-insensitively. Unknown states raise ValueError.

Registration and discovery​

Register an instance after defining the connector:

ConnectorRegistry.register(BatchJobConnector())

ConnectorRegistry stores class-level _REGISTRY and _METADATA dictionaries. The lookup key is the frozen ConnectorRegistryKey(task_type_name, task_type_version). A duplicate key raises ValueError unless you explicitly pass override=True:

ConnectorRegistry.register(connector, override=True)

Metadata is grouped by display name; each registration appends a TaskCategory containing the task type and version. get_connector() raises FlyteConnectorNotFound for an unknown task type/version. _list_connectors() and _get_connector_metadata() support metadata discovery.

The registry is populated through import side effects. _load_connectors() in connectors.utils loads installed flyte.connectors entry points, then imports explicit modules with importlib.import_module. Entry-point failures are logged as warnings and startup continues, so a failed import can leave a service without the expected connector. Registry state is class-global; tests reusing a key commonly pass override=True.

gRPC service and server startup​

Use the public ConnectorService.run() facade. The c0 command in src/flyte/_bin/connect.py calls it with the gRPC port, Prometheus port, worker count, optional timeout, and additional modules:

from flyte.connectors import ConnectorService

ConnectorService.run(8000, 9090, 10, None, None)

ConnectorService.run is a syncify wrapper around asynchronous startup. It adds the current working directory to sys.path when necessary, then delegates to _start_grpc_server.

CLI defaults are --port 8000, --prometheus_port 9090, and --worker 10; --timeout is intended for testing, and --modules may be repeated. _start_grpc_server() loads connectors, starts the Prometheus endpoint, prints metadata, creates grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker)), registers AsyncConnectorService, ConnectorMetadataService, and health checking, binds [::]:{port} insecurely, starts, and waits for termination. TLS and authentication are not configured by this package.

RPC mapping​

AsyncConnectorService implements the generated flyteidl2.connector.service_pb2_grpc service:

  • CreateTask resolves template.type and template.task_type_version, builds a NativeInterface from template input types, converts protobuf inputs to native Python values, calls create, and returns encoded metadata. It forwards the output prefix, task execution metadata, and connection-derived keyword arguments.
  • GetTask resolves the TaskCategory, decodes bytes through the connector's metadata_type, calls get, and converts the Resource to a protobuf response.
  • DeleteTask uses the same category-and-metadata decode path and calls delete.
  • GetTaskMetrics and GetTaskLogs delegate to the connector methods; the base methods raise NotImplementedError rather than returning empty responses.

ConnectorMetadataService.GetConnector and ListConnectors expose registry metadata. The service test registers DummyConnector, verifies encoded DummyMetadata, invokes create/get/delete, checks metrics and logs, and verifies that unknown metadata raises FlyteConnectorNotFound.

_grpc.py is an import-order wrapper. Before importing the gRPC extension, it sets defaults only when GRPC_VERBOSITY is absent: GRPC_VERBOSITY=ERROR, GRPC_CPP_MIN_LOG_LEVEL=ERROR, GRPC_ENABLE_FORK_SUPPORT=0, GLOG_minloglevel=2, and ABSL_LOG=0.

Local execution through the mixin​

When a connector-backed task runs locally, AsyncConnectorExecutorMixin.execute() resolves the connector using the task's type and version:

connector = ConnectorRegistry.get_connector(task.task_type, task.task_type_version)

It requires a task context and raises RuntimeError("Task context is not set.") otherwise. Remote-mode async functions call AsyncFunctionTaskTemplate.execute directly; local async functions with a non-remote raw data path call TaskTemplate.execute directly.

For remote-storage async functions, the mixin builds a dry-run code bundle, uploads it under a generated prefix, serializes a rendered task template, converts native inputs, and uploads inputs.pb. Other connector tasks serialize using the context code bundle and compiled image cache. Custom secrets are interpreted as environment-variable names; a missing variable raises ValueError("Secret ... not found in environment."), and resolved values are passed to create and each get call.

The polling algorithm in _connector.py is:

resource_meta = await connector.create(
task_template=tt, output_prefix=ctx.raw_data.path, inputs=kwargs, **secrets
)

while not is_terminal_phase(resource.phase):
resource = await connector.get(resource_meta=resource_meta, **secrets)
await asyncio.sleep(3)

The actual method initializes resource with TaskExecution.RUNNING, records each returned log link, and then raises if the final phase is not SUCCEEDED. is_terminal_phase() recognizes only SUCCEEDED, ABORTED, and FAILED; RETRYABLE_FAILED continues polling. There is no built-in timeout or retry policy, and a never-terminal provider can wait indefinitely. ABORTED also raises the generic failure because abort support is marked TODO.

After success, remote async functions load and convert outputs from the remote output path when present. Otherwise, no outputs returns None, while a native output dictionary becomes tuple(resource.outputs.values()).

Rendering task templates​

_render_task_template(tt, file_prefix) mutates container arguments in place. It replaces input, output, raw-output, checkpoint, previous-checkpoint, run-name, and action-name placeholders, then inserts --run-base-dir, --org, --project, and --domain. The inserted org, project, and domain are currently the literal values test-org, test-project, and test-domain. A template without a container is returned unchanged. The replacement loop derives each replacement from the original args[i], so this is a narrow replacement routine rather than a general templating engine.

Deploy and load a custom connector​

ConnectorEnvironment packages the service command for an app. Its default port is 8080 (the h2c Port value), and generated arguments use Prometheus port 9092. Included paths become module names: my_connector/connector.py becomes my_connector.connector.

The deployment example uses an included package and installs the connector extra:

image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("flyte[connector]")

connector = flyte.app.ConnectorEnvironment(
name="batch-job-connector",
image=image,
resources=flyte.Resources(cpu="1", memory="1Gi"),
include=["my_connector"],
)

With no custom args, container_args() builds c0 --port <port> --prometheus_port 9092 and appends --modules plus included module names. A custom command may be a string, split with shlex.split, or a list. The base flyte dependency set does not install connector server dependencies; install flyte[connector] for gRPC, health checking, HTTP, and Prometheus support.

Provider implementations compared​

ProviderMetadata and external interactionConfiguration and outputs
BigQueryUses BigQueryMetadata, google-cloud-bigquery, and BigQuery job polling.Uses service-account credentials and settings such as ProjectID and Location; successful destination tables become a Flyte DataFrame output.
DatabricksUses DatabricksJobMetadata with workspace instance and run ID; submits and polls with aiohttp.databricksInstance may come from task configuration or fallback FLYTE_DATABRICKS_INSTANCE; a bearer token is a secret.
SnowflakeUses Snowflake connector tooling, cryptography, and async-lru.Account, database, schema, warehouse, user, and private-key secret names are provider configuration. Batch expansion requires one VALUES (...) clause, named %(key)s placeholders, and equal-length non-empty lists.

All three use the same registry and metadata lifecycle, while provider metadata, state normalization, authentication, and output handling differ. Snowflake tests reject empty or mismatched lists, missing or multiple VALUES clauses, and unsupported placeholders.

Test the lifecycle​

Use tests/flyte/connector/test_connector_service.py as the service pattern: define concrete metadata, implement create/get/delete plus metrics and logs, register with override=True, and invoke AsyncConnectorService with a mocked grpc.ServicerContext. This verifies registry behavior and the protobuf boundary without an external service.

Provider tests cover provider-specific behavior, including Snowflake batch expansion. Databricks tests verify that AsyncConnectorExecutorMixin is in the task MRO. ConnectorEnvironment tests should assert module conversion and generated ports/arguments. No direct tests were found for _load_connectors, _start_grpc_server, print_metadata, or _render_task_template; test those startup and rendering paths separately when changing service wiring.