Skip to content

README.md

Latest commit

 

History

History
94 lines (75 loc) · 4.16 KB

README.md

File metadata and controls

94 lines (75 loc) · 4.16 KB
Raw

Dask-on-Lambda — Terraform infrastructure

Runs a Dask cluster where workers are Lambda functions, not EC2. The long-lived scheduler runs on Fargate; a scheduled agent Lambda watches the scheduler and fans out worker Lambdas on demand. Modeled on the Fused kubernetes_agent + rt_scaling design (agent invokes ephemeral compute; scale-down is implicit as workers self-retire).

Architecture

        EventBridge Scheduler (rate(1 min))
                    │ invokes
                    ▼
        ┌──────────────────────┐        reads desired-vs-current
        │   scaling-agent λ    │──────────────┐  (dashboard :8787)
        └──────────┬───────────┘              │
                   │ async InvokeFunction     ▼
                   │              ┌────────────────────────┐
                   ├────────────► │  Dask scheduler        │
                   │              │  (ECS Fargate, :8786)  │
                   ▼              │  scheduler.<proj>.      │
        ┌──────────────────┐     │  internal via Cloud Map│
        │  worker λ  × N    │────►└────────────────────────┘
        │  `dask worker`    │  connect back over 8786
        │  --lifetime 840s  │
        └──────────────────┘
  • Scheduler — can't be Lambda (needs to stay up), so it's one Fargate task. Cloud Map gives it a stable DNS name scheduler.<project>.internal that workers dial regardless of the task's current IP.
  • Worker — container-image Lambda running dask worker. Each invocation is one worker for up to worker_timeout (900s hard cap). --lifetime 840s makes it retire gracefully before AWS kills it, so tasks aren't lost mid-flight.
  • Scaling agent — scheduled Lambda. Polls the scheduler, computes how many more workers are wanted (up to max_workers), and InvokeFunctions the worker Lambda that many times asynchronously. No scale-down logic — workers die on their own at --lifetime, so idle capacity drains automatically.

Files

File Purpose
vpc.tf VPC, single NAT, S3 gateway + ECR/logs/STS interface endpoints
security_groups.tf Scheduler ingress on 8786/8787 from VPC; worker egress
ecr.tf Repos for scheduler / worker / scaling-agent images
scheduler.tf ECS cluster, Fargate service, Cloud Map service discovery
lambda_worker.tf The worker Lambda (in-VPC, container image)
scaling_agent.tf The agent Lambda + EventBridge schedule
iam.tf Roles for ECS, both Lambdas, and the scheduler-invoke path
outputs.tf Scheduler address, ECR URLs, function ARNs

Deploy

  1. Build & push the three images to the ECR repos (created by ecr.tf — apply once first, or split repos into a bootstrap stage):

    terraform init && terraform apply -target=aws_ecr_repository.scheduler \
  -target=aws_ecr_repository.worker -target=aws_ecr_repository.scaling_agent

# then docker build/push each; see the app repos for the images

  2. Full apply:

    cp terraform.tfvars.example terraform.tfvars   # edit as needed
terraform apply

  3. Point your Dask client at the scheduler (from inside the VPC / over VPN):

    from dask.distributed import Client
client = Client("tcp://scheduler.dask-lambda.internal:8786")

Knobs

See variables.tf. The ones you'll actually touch: worker_memory, worker_lifetime_seconds (keep < worker_timeout), max_workers, and scaling_interval.

Not included (app code lives in the repo root, not here)

  • The scheduler image (dask scheduler)
  • The worker image entrypoint that reads DASK_SCHEDULER_ADDRESS / WORKER_LIFETIME and execs dask worker
  • The scaling-agent handler that reads SCHEDULER_DASHBOARD and invokes WORKER_FUNCTION_ARN

The Terraform passes everything those need as env vars.