Dagster with AWS compute

This page adds Polars Nanook authentication and per-asset compute sizing to a Dagster pipeline. It uses typed ConfigurableResource subclasses.

Credentials

Credentials are referenced through EnvVar so the raw values are resolved at run launch and never displayed in the Dagster UI. Populate NANOOK_CLIENT_ID and NANOOK_CLIENT_SECRET in the Dagster deployment environment, either directly or by exporting from a secret manager (e.g. AWS Secrets Manager).

Compute context lifecycle

PolarsNanookCompute declares PolarsNanookAuth as a nested resource so Dagster runs PolarsNanookAuth.setup_for_execution, which calls authenticate(...), before any compute resource is used. The compute resource owns the ComputeContext lifecycle through yield_for_execution: the compute cluster is started before each asset that uses the resource runs and stopped on exit, including on failure. Asset bodies contain only Polars code.

Per-asset machine size

Per-asset compute sizing follows from the same pattern: one PolarsNanookCompute instance per machine size (cpus/memory) under distinct resource keys, and each asset takes the instance sized for it. Dagster manages the lifecycle of each instance independently, so the asset that uses the larger cluster does not pay the cost of spinning it up while the smaller-cluster assets are running.

Complete example

The complete module below defines two resources (2 vCPU/4 GB and 4 vCPU/16 GB), three assets, and the Definitions object. Replace the ... placeholders with your source paths and S3 destinations. Each resource keeps the context it started, and each asset passes its resource's context to .remote(). A process-wide default set with pn.set_compute_context would send every asset to the cluster started last when Dagster runs assets in one process.

from contextlib import contextmanager

import polars as pl
import polars_nanook as pn
from dagster import asset, ConfigurableResource, Definitions, EnvVar
from pydantic import PrivateAttr


class PolarsNanookAuth(ConfigurableResource):
    """Service account authentication for Polars Nanook."""

    client_id: str
    client_secret: str

    def setup_for_execution(self, context) -> None:
        pn.authenticate(client_id=self.client_id, client_secret=self.client_secret)


class PolarsNanookCompute(ConfigurableResource):
    """Polars Nanook compute context with auto-managed cluster lifecycle.

    Declares `PolarsNanookAuth` as a nested resource so Dagster runs `authenticate(...)`
    before this resource is used. `yield_for_execution(...)` starts a `ComputeContext`
    of the configured shape and the `finally` clause guarantees the cluster is stopped even
    when an asset raises.
    """

    auth: PolarsNanookAuth
    cpus: int
    memory: int
    _ctx: pn.ComputeContext = PrivateAttr()

    @property
    def ctx(self) -> pn.ComputeContext:
        return self._ctx

    @contextmanager
    def yield_for_execution(self, context):
        ctx = pn.ComputeContext(
            workspace="My workspace", cpus=self.cpus, memory=self.memory
        )
        ctx.start()
        try:
            self._ctx = ctx
            yield self
        finally:
            ctx.stop()


@asset
def dataset_1(small_vm: PolarsNanookCompute):
    pl.scan_csv(...).remote(small_vm.ctx).sink_parquet(...).await_result()


@asset
def dataset_2(small_vm: PolarsNanookCompute):
    pl.scan_ndjson(...).remote(small_vm.ctx).sink_parquet(...).await_result()


# Use a bigger machine for the join.
@asset(deps=[dataset_1, dataset_2])
def joined(large_vm: PolarsNanookCompute):
    pl.scan_parquet(...).remote(large_vm.ctx).sink_parquet(...).await_result()


# Share a single `PolarsNanookAuth` instance between the compute resources.
auth = PolarsNanookAuth(
    client_id=EnvVar("NANOOK_CLIENT_ID"),
    client_secret=EnvVar("NANOOK_CLIENT_SECRET"),
)

defs = Definitions(
    assets=[dataset_1, dataset_2, joined],
    resources={
        "small_vm": PolarsNanookCompute(auth=auth, cpus=2, memory=4),
        "large_vm": PolarsNanookCompute(auth=auth, cpus=4, memory=16),
    },
)

Materializing joined starts the 4 vCPU cluster, runs the join remotely, and stops the cluster in the finally clause even if the asset raises. A service account starts without a default workspace, so the ComputeContext names the workspace with workspace=.

dataset_1 and dataset_2 do not depend on each other, so Dagster runs them in parallel, each on its own 2 vCPU cluster.

Resource key and parameter name

The Dagster resource key (the dictionary key in Definitions.resources) must match the parameter name in each asset signature. The type annotation marks the parameter as a resource, but Dagster injects by key, not by type.