Prefect with AWS compute

This page adds Polars Nanook authentication and per-task compute sizing to a Prefect flow. The credentials are kept in Prefect Secret blocks.

Credentials

The service-account client ID and secret are stored as Prefect Secret blocks so raw values are never hard-coded in flow code. Create them once from Python:

from prefect.blocks.system import Secret

Secret(value="<CLIENT_ID>").save("polars-nanook-client-id")
Secret(value="<CLIENT_SECRET>").save("polars-nanook-client-secret")

Cluster lifecycle

Authentication is performed once at the top of the flow so every task in the run shares the same session. The flow creates the contexts after that, because creating a ComputeContext contacts the control plane. Each task receives its context as an argument and passes it to .remote(ctx), and the flow owns the compute cluster lifecycle: start() is called before any task runs and stop() is guaranteed by the finally clause. The tasks set cache_policy=NO_CACHE because Prefect's default cache policy cannot hash a ComputeContext.

dataset_1 and dataset_2 share a single compute context (small_ctx) which handles both tasks sequentially (concurrently for Prefect, but sequentially in Polars Nanook). joined runs on a larger cluster (4 vCPU / 16 GB), which is started and stopped independently.

Complete example

The complete flow below starts both clusters, runs dataset_1 and dataset_2 on the 2 vCPU/4 GB cluster, then joined on the 4 vCPU/16 GB cluster. Replace ... with your paths.

import polars as pl
import polars_nanook as pn
from prefect import flow, task
from prefect.blocks.system import Secret
from prefect.cache_policies import NO_CACHE


@task(cache_policy=NO_CACHE)
def dataset_1(ctx):
    pl.scan_csv(...).remote(ctx).sink_parquet(...).await_result()


@task(cache_policy=NO_CACHE)
def dataset_2(ctx):
    pl.scan_ndjson(...).remote(ctx).sink_parquet(...).await_result()


# Use a bigger machine for the join.
@task(cache_policy=NO_CACHE)
def joined(ctx):
    pl.scan_parquet(...).remote(ctx).sink_parquet(...).await_result()


@flow(name="Report")
def report():
    pn.authenticate(
        client_id=Secret.load("polars-nanook-client-id").get(),
        client_secret=Secret.load("polars-nanook-client-secret").get(),
    )

    small_ctx = pn.ComputeContext(workspace="My workspace", cpus=2, memory=4)
    large_ctx = pn.ComputeContext(workspace="My workspace", cpus=4, memory=16)
    small_ctx.start()
    large_ctx.start()
    try:
        f1 = dataset_1.submit(small_ctx)
        f2 = dataset_2.submit(small_ctx)
        joined.submit(large_ctx, wait_for=[f1, f2]).wait()
    finally:
        small_ctx.stop()
        large_ctx.stop()


report()

Both clusters are stopped in the finally clause even when a task fails. The contexts pass workspace= because a service account starts without a default workspace.