Proposal: add work pool support for Armada (K8s multi-cluster) batch system

[richscott]

Hello, I am working with a team that is supporting applications running on very large (100k+ nodes) configurations, and is interested in using Prefect for running batch workflows.

Background: conventional Kubernetes clusters typically can support up to 5000-7000 nodes total. To allow scaling of batch jobs on much larger configurations, the Armada system was developed (https://armadaproject.io and source code at https://github.com/armadaproject/armada) to provide a facility layer on top of vanilla Kubernetes to allow submission and management of jobs that may potentially use hundreds of thousands of nodes. It provides a GRPC-based submission mechanism, an intelligent scheduler, and manager, using Redis, Postgresql, Pulsar, and other components to manage jobs. (There are
alternatives/competitors to Armada, e.g. Google Cloud Platform's hosted K8S uses a proprietary alternative to K8s' etcd to provide similar high-scalability.)

I have been working on an initial implementation of a new "armada" Work Pool type, heavily based on the existing kubernetes work pool implementation (Armada and K8s are isomorphically quite similar). A new armada work pool can be created and configured via the Prefect web UI, and it does successfully submit and run basic Prefect jobs. I haven't yet done any serious load-testing or intensive regression testing on error-handling, etc.

Observations / Questions:

• Although the Armada user community is not large (admittedly, it's designed for very, very large clusters), it is a CNCF Sandbox project and is used in production environments right now. Would the Prefect project require a sizable user community to accept it as a new work pool integration?

• It was really a nice dev experience that the initial work to get the Armada work pool functioning went quickly, and that the changes were all new files in src/integrations/prefect-armada/*.

After I verified basic functionality, I worked on adding Armada to the "Add Work Pool" web UI page (I used AI assistance, as I am not a front-end dev), which went pretty quickly. Those changes are more interleaved with existing code, e.g. in ui-v2/src/components/work-pools/create/infrastructure-type-step/infrastructure-type-step.stories.tsx where some new clauses for armada were added. I know it would probably be a substantial amount of work, but would it ever be possible to isolate the
per-worker-type front-end config/data/logic in the same way that the back-end has src/integrations/perfect-<work-pool-type>/? If a work pool type was not directly integrated into the Prefect repo, then a simple additive overlay could be dropped onto a repo tree, and the new work pool type could be immediately used.

So far, the enhancements are in my repo fork at https://github.com/richscott/prefect within the armada-integration branch.
All suggestions/questions/ideas would be welcomed - thank you.

[desertaxle]

Thanks for sharing this! Armada sounds like a great fit for a standalone prefect-armada package with your Armada worker hosted in it. You’ve already put together a thoughtful implementation, especially the configurable Armada-native job template and the coverage of the job lifecycle.

We’d also be happy to explore making the UI extensible so third-party worker types like this can be visible and configurable without requiring changes to Prefect itself.