Implementing a Custom Run Coordinator

[ldnicolasmay]

Context

I couldn't find an example of a Custom Run Coordinator anywhere, so I thought it might be helpful to share my solution with anyone else looking to implement their own custom run coordinator.

This may not be the most elegant solution, but it works.

If you have any trouble following this, please comment below and I'll revise when I find the time.

File Tree

my-dagster-project/
├── my-dagster-customizations/
│   ├── run_coordinator/
│   │   ├── __init__.py
│   │   └── custom_run_coordinator.py
│   └── pyproject.toml
│
├── my-user-code/
│   ├── dagster/
│   │   ├── __init__.py
│   │   └── definitions.py
│   └── pyproject.toml
│
├── helm/
│   └── dagster/
│       └── values.yaml
│
├── Dockerfile.dagster_customizations
└── Dockerfile.my_user_code

I assume you have a working Dagster project with your own user code, here exemplified as the my-user-code directory and Dockerfile.my_user_code Dockerfile.

I'll focus on the my-dagster-customizations directory, the Helm values.yaml file, and both Dockerfiles.

I suspect Kubernetes is the most common way to deploy Dagster instances, so my examples will assume you're using Kubernetes. This example can probably be easily adapted for different approaches.

Dockerfiles

You'll need to install your customizations into all the Docker images that are used in your deployment. That includes your user code image and the image used for both the Dagster daemon and webserver.

Daemon and Webserver

Here, we'll be installing the my-dagster-customizations package over Dagster's dagster-k8s image.

FROM --platform=$BUILDARCH dagster/dagster-k8s

ENV DAGSTER_DIR='/opt'

# Install my-dagster-customizations package
COPY my-dagster-customizations ${DAGSTER_DIR}/my-dagster-customizations
RUN pip install -U pip; \
    pip install -e ${DAGSTER_DIR}/my-dagster-customizations

We'll also need to install the package in the user code image.

FROM --platform=$BUILDARCH python

ENV DAGSTER_APP='/opt/dagster/app'

# Install my-dagster-customizations package
COPY my-dagster-customizations ${DAGSTER_APP}/my-dagster-customizations
RUN pip install -U pip; \
    pip install -e ${DAGSTER_APP}/my-dagster-customizations

# Install Python dependencies
COPY my-user-code/pyproject.toml pyproject.toml
RUN pip install -U pip; \
    pip install -e '.'

# Copy required user code files into image volume
WORKDIR ${DAGSTER_APP}
COPY my-user-code/ ./my-user-code/

Contents of my-dagster-customizations

• pyproject.toml

  [build-system]
  requires = ["setuptools"]
  build-backend = "setuptools.build_meta"
  
  [project]
  name = "my-dagster-customizations"
  version = "0.1.0"
  authors = [{ name = "L. D. Nicolas May" }]
  maintainers = [{ name = "L. D. Nicolas May" }]
  readme = "README.md"
  classifiers = ["Programming Language :: Python"]
  requires-python = ">=3.10,<3.12"
  
  dependencies = ["dagster"]
  
  [tool.setuptools]
  packages = ["run_coordinator"]

• run_coordinator/__init__.py

  Empty file.

• run_coordinator/custom_run_coordinator.py

  This module extends the Dagster's QueuedRunCoordinator class.

  from typing import Any, List, Mapping, NamedTuple, Optional, Sequence
  
  from typing_extensions import Self
  
  from dagster import _check as check
  from dagster._core.run_coordinator.base import SubmitRunContext
  from dagster._core.run_coordinator.queued_run_coordinator import QueuedRunCoordinator
  from dagster._core.storage.dagster_run import DagsterRun
  from dagster._serdes import ConfigurableClassData
  
  
  class RunQueueConfig(
      NamedTuple(
          "_RunQueueConfig",
          [
              ("max_concurrent_runs", int),
              ("tag_concurrency_limits", Sequence[Mapping[str, Any]]),
              ("max_user_code_failure_retries", int),
              ("user_code_failure_retry_delay", int),
              ("should_block_op_concurrency_limited_runs", bool),
              ("op_concurrency_slot_buffer", int),
          ],
      )
  ):
      def __new__(
          cls,
          max_concurrent_runs: int,
          tag_concurrency_limits: Optional[Sequence[Mapping[str, Any]]],
          max_user_code_failure_retries: int = 0,
          user_code_failure_retry_delay: int = 60,
          should_block_op_concurrency_limited_runs: bool = True,
          op_concurrency_slot_buffer: int = 0,
      ):
          return super().__new__(
              cls,
              check.int_param(max_concurrent_runs, "max_concurrent_runs"),
              check.opt_sequence_param(tag_concurrency_limits, "tag_concurrency_limits"),
              check.int_param(max_user_code_failure_retries, "max_user_code_failure_retries"),
              check.int_param(user_code_failure_retry_delay, "user_code_failure_retry_delay"),
              check.bool_param(should_block_op_concurrency_limited_runs, "should_block_op_concurrency_limited_runs"),
              check.int_param(op_concurrency_slot_buffer, "op_concurrency_slot_buffer"),
          )
  
      def with_concurrency_settings(
          self, concurrency_settings: Mapping[str, Any]
      ) -> "RunQueueConfig":
          run_settings = concurrency_settings.get("runs", {})
          pool_settings = concurrency_settings.get("pools", {})
          return RunQueueConfig(
              max_concurrent_runs=run_settings.get("max_concurrent_runs", self.max_concurrent_runs),
              tag_concurrency_limits=run_settings.get("tag_concurrency_limits", self.tag_concurrency_limits),
              max_user_code_failure_retries=self.max_user_code_failure_retries,
              user_code_failure_retry_delay=self.user_code_failure_retry_delay,
              should_block_op_concurrency_limited_runs=self.should_block_op_concurrency_limited_runs,
              op_concurrency_slot_buffer=pool_settings.get("op_granularity_run_buffer", self.op_concurrency_slot_buffer),
          )
  
  
  class CustomRunCoordinator(QueuedRunCoordinator):
      """
      Run coordinator that extends QueuedRunCoordinator with additional filtering
      capability based on tag values.
      """
  
      def __init__(
          self,
          max_concurrent_runs: Optional[int] = None,
          tag_concurrency_limits: Optional[List[Mapping[str, Any]]] = None,
          dequeue_interval_seconds: Optional[int] = None,
          dequeue_use_threads: Optional[bool] = None,
          dequeue_num_workers: Optional[int] = None,
          max_user_code_failure_retries: Optional[int] = None,
          user_code_failure_retry_delay: Optional[int] = None,
          block_op_concurrency_limited_runs: Optional[Mapping[str, Any]] = None,
          inst_data: Optional[ConfigurableClassData] = None,
      ):
          super().__init__(
              max_concurrent_runs=max_concurrent_runs,
              tag_concurrency_limits=tag_concurrency_limits,
              dequeue_interval_seconds=dequeue_interval_seconds,
              dequeue_use_threads=dequeue_use_threads,
              dequeue_num_workers=dequeue_num_workers,
              max_user_code_failure_retries=max_user_code_failure_retries,
              user_code_failure_retry_delay=user_code_failure_retry_delay,
              block_op_concurrency_limited_runs=block_op_concurrency_limited_runs,
              inst_data=inst_data,
          )
  
      def get_run_queue_config(self) -> RunQueueConfig:
          return RunQueueConfig(
              max_concurrent_runs=self._max_concurrent_runs,
              tag_concurrency_limits=self._tag_concurrency_limits,
              max_user_code_failure_retries=self._max_user_code_failure_retries,
              user_code_failure_retry_delay=self._user_code_failure_retry_delay,
              should_block_op_concurrency_limited_runs=self._should_block_op_concurrency_limited_runs,
              op_concurrency_slot_buffer=self._op_concurrency_slot_buffer,
          )
  
      @classmethod
      def config_type(cls):
          parent_config = super().config_type()
          return {**parent_config}  # You might want to expand this if adding class properties
  
      @classmethod
      def from_config_value(
          cls, inst_data: ConfigurableClassData, config_value: Mapping[str, Any]
      ) -> Self:
          return cls(
              inst_data=inst_data,
              max_concurrent_runs=config_value.get("max_concurrent_runs"),
              tag_concurrency_limits=config_value.get("tag_concurrency_limits"),
              dequeue_interval_seconds=config_value.get("dequeue_interval_seconds"),
              dequeue_use_threads=config_value.get("dequeue_use_threads"),
              dequeue_num_workers=config_value.get("dequeue_num_workers"),
              max_user_code_failure_retries=config_value.get("max_user_code_failure_retries"),
              user_code_failure_retry_delay=config_value.get("user_code_failure_retry_delay"),
              block_op_concurrency_limited_runs=config_value.get("block_op_concurrency_limited_runs"),
          )
  
      def _is_run_allowed(self, dagster_run: DagsterRun) -> bool:
          """
          Check if the run is allowed based on tag criteria.
          """
          run_tags = dagster_run.tags or {}
  
          # Example `run_tags`:
          # { 'foo': 'bar' }
          #
          # Example `self._tag_concurrency_limits`:
          # [
          #     {'key': 'foo', 'value': 'bar', limit: 1},
          # ]
  
          # If there are tag_concurrency_limits, the project must be included in order
          # for a job run to be allowed.
          if self._tag_concurrency_limits:
              # Below can be simplified to a set comprehension
              values = set()
              for limit in self._tag_concurrency_limits:
                  limit_key = limit["key"]
                  limit_value = limit["value"]
                  if limit_key == "foo":
                      values.add(limit_value)
  
              if "foo" in run_tags and run_tags.get("foo") in values:
                  return True
  
              return False
  
          # If there aren't any `tag_concurrency_limits` defined, no run should be allowed
          return False
  
      def submit_run(self, context: SubmitRunContext) -> DagsterRun:
          """
          Override submit_run to add tag-based filtering.
          """
          dagster_run: DagsterRun = context.dagster_run
          run_tags = dagster_run.tags or {}
  
          # Check tags
          if not self._is_run_allowed(dagster_run):
              # Unfortunately, "canceled" run jobs appear as red (like failed) in the UI.
              self._instance.report_run_canceled(
                  dagster_run,
                  message=(
                      "Run denied by CustomRunCoordinator: Failed tag permission check. "
                      f'Run tag `{{"foo": "{run_tags.get("foo")}"}}` '
                      "not in `tag_concurrency_limits`"
                  ),
              )
              return dagster_run
  
          # If we've passed all filters, submit run to the queue
          return super().submit_run(context)

Contents of helm/dagster/

• values.yaml

  The portion of interest is the values of dagsterDaemon.runCoordinator.

  dagsterDaemon:
    runCoordintor:
      enabled: true
      type: CustomRunCoordinator
      config:
        customRunCoordinator:
          module: run_coordinator.custom_run_coordinator
          class: CustomRunCoordinator
          config:
            tag_concurrency_limits:
              - { key: foo, value: bar, limit: 1}

Wrap-up

If you want to add a config field to the custom run coordinator (for example, alongside dagsterDaemon.runCoordinator.config.customRunCoordinator.config.tag_conccurreny_limits), you'll need to add that as a property to the CustomRunCoordinator and RunQueueConfig classes in my-dagster-customizations/run_coordinator/custom_run_coordinator.py, and then account for how the property changes submit_run method's triggering of DagsterRuns.

I hope this helps anyone looking to create their own run coordinator behavior get up to speed much faster than it took me.

If you're looking to customize Dagster's run launcher, here's a GitHub issue that should help: #15856. The process is pretty similar to customizing the run coordinator.

[AndreaGiardini]

@garethbrickman Probably worth moving it to a discussion?

@ldnicolasmay Thank you for the writeup!