Prima Mente scales DNA foundation-model pipelines past 250,000 tasks with pure Python orchestration

Prima Mente is an AI biology company building general-purpose biological foundation models, starting with the human brain. The company generates its own data, trains multi-modal foundation models on enormous volumes of molecular biology data, and translates the results into downstream research and clinical outcomes — from gene- and cell-level prediction tasks to diagnostics for neurological diseases like Alzheimer’s and Parkinson’s. Its team of AI researchers, experimentalists, clinicians, and operators works across London, San Francisco, and UAE.

Every model the company trains depends on a preprocessing layer that reprocesses raw sequencing data in-house. Each model needs its own pipeline and its own scaling profile, and the volumes are substantial: Prima Mente has roughly 5 petabytes of data to move through these pipelines before the resulting datasets are handed off to a SLURM cluster on Nebius for foundation-model pre-training.

Before Union.ai, that orchestration ran on Nextflow on top of AWS Batch — a setup that became harder to sustain as the team’s requirements grew.

Several constraints became more pronounced as the team pushed for scale:

• A scaling ceiling. Around 25,000 fanned-out tasks, the team began hitting naming and hash-ID collisions that were difficult to design around.
• Limited central visibility. Jobs could be submitted from anywhere, with little historical record of what had run or who had run it.
• A head-node model that required careful tuning. Every run required its own head node with allocated resources; under-provisioning could fail the run and affect other jobs sharing it.
• A separate DSL. Nextflow’s domain-specific language offered limited support from IDEs and coding agents, and required the team to work outside Python.
• Limited modularity. Running a single task in isolation typically required exporting the workflow entry point, creating placeholder files, and adding boilerplate code.
• No infrastructure team to absorb the complexity. Maintaining a Kubernetes cluster wasn’t realistic, so the team opted for the most accessible setup — which came with AWS Batch’s limited transparency and per-job instance churn.

Sensitive patient data further raised the bar for reproducibility and lineage. Before Union, run history was mostly stored in GCS paths and GitHub repositories, with only a bare-bones, in-house attempt to track what had been executed.

What the team needed was clear: an orchestrator that could scale far past 25,000 tasks, give them operational visibility and role-based access control, and let them author everything in plain Python — no new DSL to learn.

“I waited too long to switch to Flyte 2 — it’s just so much better. We can scale to 200,000–300,000 pods with the escalation logic baked right in, and the out-of-memory and scheduling headaches I used to fight are simply gone.”
— Jay Ganbat, Principal Bioinformatics Engineer, Prima Mente

Having seen the success others had with open source Flyte, Prima Mente skipped a formal tool evaluation and moved their preprocessing straight onto Union.ai 

A few capabilities did the most work:

Pure Python authoring

Where Nextflow relied on a separate DSL, Union lets the bioinformatics team write workflows in the language they already use. That removed an entire learning curve and made the team’s coding agents far more effective — a meaningful advantage for scientists who come from a computational-biology background rather than an infrastructure one.

Flyte 2 scale with escalation logic built in

The single most impactful change was Flyte 2. Workloads the team previously struggled to complete now run cleanly, and the team can scale to 200,000–300,000 pods with automatic resource escalation baked in. The OOMKilled and scheduling errors that used to dominate the day are gone.

A remote image builder, so no one touches a Dockerfile

Every Prima Mente workflow now uses Union’s remote image builder. The team maintains a base Docker image, and Union layers additional requirements on top at runtime — no rebuilding from scratch, no local image builds. Even for an engineer fully capable of writing Dockerfiles, it removes a recurring chore from daily work.

Operational visibility, RBAC, and true modularity

Union gave the team a centralized control plane it previously didn’t have: a clear record of inputs, outputs, and who ran what, with role-based access scoped to projects and domains. And any single task can be re-run straight from the UI by entering its inputs — no endpoint exports or throwaway files required.

The scale before-and-after is stark. Where the previous setup topped out around 25,000 tasks, Flyte 2 changed the picture entirely: the team's largest single run reached 278,000 concurrent actions executed across 150,000 reused Pods, without any of the reliability problems that capped the earlier stacks.

That scale is already doing real work. In the span of a week or two, roughly 500 terabytes of sequencing data moved through the Flyte 2 pipeline, part of the 5 petabytes the team reprocesses in-house to feed foundation-model pre-training.

The change shows up in the day-to-day, too. Prima Mente’s bioinformatics team — two engineers in London and four in San Francisco — works in Union daily, authoring workflows entirely in Python with no infrastructure translation layer in between. And when issues arise, support has kept pace:

“Support-wise it’s been nothing but excellent. Every request from our team gets addressed within a few hours.”
— Jay Ganbat, Principal Bioinformatics Engineer, Prima Mente

Looking ahead

For Prima Mente, the orchestration layer has gone from a recurring source of pain to something the team simply doesn't have to think about. As the team scales toward its next model, it will look to expand Union’s impact beyond preprocessing and data generation: configuring the Nebius data plane, exploring reusable containers and Spark-backed SQL for large-scale data access, and giving Prima Mente’s infrastructure engineer direct visibility into node pools. For a company building foundation models to understand and protect the brain, every cycle reclaimed from infrastructure is a cycle returned to science. With their infrastructure finally matching their ambition, Prima Mente is accelerating what’s possible.