" (declare that task + its app authorization in # config). Override/extend via UPSUN_QUEUE_TASK_MAP (JSON # {"": " "}). self._queue_map: dict = json.loads(os.environ.get("UPSUN_QUEUE_TASK_MAP", "{}")) # The task container can't reach the api-server on localhost and has no # PLATFORM_ROUTES, so we resolve the public Execution API URL here (the # web app knows it) and pass it to the task. self._exec_api = self._execution_api_url() @staticmethod def _execution_api_url() -> str: base = os.environ.get("AIRFLOW__API__BASE_URL") if base: return base.rstrip("/") + "/execution/" routes = json.loads(base64.b64decode(os.environ["PLATFORM_ROUTES"]).decode()) for url, cfg in routes.items(): if cfg.get("primary"): return url.rstrip("/") + "/execution/" raise RuntimeError("cannot determine Execution API URL (no AIRFLOW__API__BASE_URL / primary route)") def _task_for_queue(self, queue: str | None) -> str: q = queue or "default" if q in self._queue_map: return self._queue_map[q] if q == "default": return self.task_name return f"{self.task_name}-{q}" # -- auth ---------------------------------------------------------------- def _token_value(self, force: bool = False) -> str: explicit = os.environ.get("UPSUN_API_TOKEN") if explicit: return explicit if not force and self._token and time.monotonic() < self._token_expiry - 30: return self._token data = urllib.parse.urlencode({"grant_type": "client_credentials"}).encode() req = urllib.request.Request(_TOKEN_URL, data=data, method="POST") req.add_header("Content-Type", "application/x-www-form-urlencoded") req.add_header("x-token-ttl", "900") # match our polling cadence with urllib.request.urlopen(req, timeout=15) as resp: payload = json.loads(resp.read().decode()) self._token = payload["access_token"] self._token_expiry = time.monotonic() + int(payload.get("expires_in", 900)) return self._token def _api(self, method: str, url: str, body: dict | None = None, _retry: bool = True) -> dict: data = json.dumps(body).encode() if body is not None else None req = urllib.request.Request(url, data=data, method=method) req.add_header("Authorization", f"Bearer {self._token_value()}") if data is not None: req.add_header("Content-Type", "application/json") try: with urllib.request.urlopen(req, timeout=30) as resp: raw = resp.read().decode() except urllib.error.HTTPError as e: if e.code == 401 and _retry: self._token_value(force=True) return self._api(method, url, body, _retry=False) raise RuntimeError(f"{method} {url} -> HTTP {e.code}: {e.read().decode(errors='replace')}") from e return json.loads(raw) if raw else {} # -- scheduler interface ------------------------------------------------- def _process_workloads(self, workloads_seq: Sequence["workloads.All"]) -> None: from airflow.executors.workloads import ExecuteTask for workload in workloads_seq: if not isinstance(workload, ExecuteTask): raise ValueError(f"{type(self).__name__} cannot handle {type(workload).__name__}") key = workload.ti.key task_name = self._task_for_queue(getattr(workload.ti, "queue", None)) try: activity_id = self._trigger(workload, task_name) except Exception as e: # noqa: BLE001 - report as a failed launch self.log.error("Failed to trigger Upsun task %r for %s: %s", task_name, key, e) self.queued_tasks.pop(key, None) self.event_buffer[key] = (TaskInstanceState.FAILED, None) continue self.queued_tasks.pop(key, None) self.running.add(key) self._activities[key] = activity_id self.event_buffer[key] = (TaskInstanceState.QUEUED, activity_id) self.log.info("Task %s (queue=%s) -> %s / Upsun activity %s", key, getattr(workload.ti, "queue", None), task_name, activity_id) def _trigger(self, workload: "workloads.ExecuteTask", task_name: str) -> str: body = {"variables": {"env": { "UPSUN_WORKLOAD": workload.model_dump_json(), "AIRFLOW__CORE__EXECUTION_API_SERVER_URL": self._exec_api, }}} resp = self._api("POST", f"{self._base}/tasks/{task_name}/run", body) emb = resp.get("_embedded", {}) acts = emb.get("activities") activity = acts[0] if isinstance(acts, list) and acts else (acts if isinstance(acts, dict) else resp) activity_id = activity.get("id") if not activity_id: raise RuntimeError(f"no activity id in run response: {resp}") return activity_id def sync(self) -> None: for key, activity_id in list(self._activities.items()): try: act = self._api("GET", f"{self._base}/activities/{activity_id}") except Exception as e: # noqa: BLE001 - transient API error, retry next heartbeat self.log.warning("Polling activity %s failed: %s", activity_id, e) continue state = (act.get("state") or act.get("status") or "").lower() if state not in _TERMINAL: continue result = (act.get("result") or "").lower() self._activities.pop(key, None) if state == "complete" and result != "failure": self.log.info("Upsun activity %s complete -> success for %s", activity_id, key) self.success(key) else: self.log.warning("Upsun activity %s ended state=%s result=%s -> fail for %s", activity_id, state, result, key) self.fail(key) def try_adopt_task_instances(self, tis: Sequence["TaskInstance"]) -> Sequence["TaskInstance"]: # We don't persist activity ids across scheduler restarts; let the # scheduler reschedule orphaned TIs rather than adopt them. return list(tis) def end(self) -> None: self.heartbeat() def terminate(self) -> None: # pragma: no cover - nothing long-lived to kill here pass ``` The task container's side is a short script that deserializes the workload and runs it through the same Task SDK path a Celery worker would use. The token in the workload authorizes it, so the container needs no database or broker access at all. The app declares an [authorization](/docs/configure-apps/image-properties/authorizations) to operate the task, and Upsun's in-container token service mints a scoped, credential-less token at runtime. The part I find genuinely useful is that you can define more than one class of task container. Route memory-hungry DAGs to a large one and everything else to a small one: ```yaml theme={null} tasks: task-runner: type: python:3.12 container_profile: HIGH_MEMORY source: root: /runtime run: command: LD_PRELOAD=/app/lib/no_nondumpable.so python task_runner.py timeout: 3600 ``` Because a task container only exists while a task runs, you pay for compute per task instead of for a worker that sits idle between DAG runs. The queue-to-task mapping in the executor lets you tune that per workload. Two caveats to plan around. A per-task container has a cold start of a few seconds, so this suits isolation-sensitive or long-running tasks better than a flood of sub-second ones. And task containers cap concurrency at a default of 3 parallel runs, with further triggers queuing behind them. That default is adjustable, but design for it rather than against it. ## Why go to the trouble At the end of the ladder you have Airflow running the way it's meant to, and you have it inside an Upsun environment. That second part is what makes the effort worth it. Branch an environment and Upsun clones its data, so a staging copy of your Airflow comes up with a real Postgres database instead of an empty one. Add [object storage](/docs/manage-resources/object-storage) and the buckets clone alongside it, so your DAGs run against realistic data in staging without touching production. You get scoped secrets, per-environment variables, and the same container hardening this whole article started with, applied to a data platform you can branch, test, and throw away. The 30 lines of C get you in the door. The isolation model is the choice worth making deliberately, and it's nice that the platform lets you make it at all. Task containers are in prerelease. Read [the task containers documentation](/docs/configure-apps/tasks) to see how they work and request access on your project.