Resumable Pipeline & UX Improvements for Long-Running Jobs

[cytokineking]

First, thank you for releasing PXDesign.

After running several multi-hour jobs on multi-GPU setups, we encountered some pain points with the current pipeline:

• No resume capability. If a job crashes, gets preempted, or is terminated (common on shared HPC clusters), all progress is lost. For jobs generating the recommended 10,000+ designs that can run dozens of hours, this is a significant issue.
• Results only emitted at completion. CIF files are held in memory until the entire diffusion stage finishes. Users have no intermediate artifacts to inspect, and a late-stage failure means losing everything.
• Difficult progress monitoring. It's hard to tell how far along a job is, especially in distributed multi-GPU runs.
• Cluttered output structure. The current output layout mixes intermediate artifacts with final results, making it harder to find what matters.

We've implemented solutions for all of the above in this fork: https://github.com/cytokineking/PXDesign

Key changes:

• World-size agnostic resume: Deterministic design_id partitioning with durable state tracking—resume on any GPU count
• Incremental CIF streaming: Atomic writes after each diffusion chunk completes, so progress is never lost
• Heartbeat monitoring: Standardized status.json files for real-time progress visibility
• Versioned results output: Clean results/ snapshots that don't overwrite on re-ranking, with internal runs/ structure for reproducibility

All changes have been tested on single-GPU and multi-GPU configurations with resume scenarios at various pipeline stages.
Feel free to use any of this work if it's helpful—happy to discuss the approach or prepare a formal PR if there's interest.

[guanjq]

Thanks a lot for the detailed write-up and for sharing your fork — this is a substantial and thoughtful set of improvements.

You’re right about these pain points. In the current PXDesign inference pipeline, samples for a single task are kept in memory until the diffusion stage completes. With a large N_sample, this means no intermediate CIFs are written and a late-stage failure can indeed lose all samples for that task. On the evaluation side, we previously emitted intermediate JSONs with per-sample confidence metrics, but these were removed during output cleanup. In hindsight, that does make resume and partial recovery harder.

We agree that resume support, incremental persistence, and better progress visibility are important for long-running multi-GPU jobs. I took a quick look at your fork, and it seems these issues are addressed in a systematic way.

We’d be happy to review this as a formal PR. If you’re open to it, please feel free to open a PR against main, ideally with a brief design summary. Given the scope, we may iterate on some interfaces during review, but the overall direction makes sense.

Thanks again for the contribution — a PR would be very welcome!

[cytokineking]

Additional work is needed to address handling of the design/prediction metrics written to the .csv files. Those are also held in memory and written once. Unfortunately, those activities are performed by PXDesignBench and can't be addressed by a single PR in this repo.

[xinshi-chen]

Hi @cytokineking ,
Thanks for pointing that out! PXDesignBench is also maintained by our team, so we can coordinate the updates across both repositories. Please feel free to propose the order that makes the most sense to you, as you've implemented the changes and understand the dependencies very well. We look forward to supporting this integration, so just let us know how you'd like us to cooperate.

[cytokineking]

Thanks, glad to know you are open to altering PXDesignBench as well.

Here is a quick update on progress towards resume and UX feature enhancements. I now have a fully resumable path for generating the .csv outputs by reconstructing them from durable JSON artifacts emitted by the AlphaFold2 / Protenix prediction stages. The PXDesignBench computations themselves are fast, CPU-bound, and parallelizable (configurable worker count), so if a job is interrupted and later resumed, the CSVs can be regenerated deterministically without issue.

In addition, I extended parallelization beyond the diffusion/inference stage so that AF2/Protenix tasks are also distributed across GPU workers, rather than being serialized downstream.

A few other notable changes in the fork:
• Optional use of ColabFold’s MMseqs2 server for MSA generation
• Support for specifying a range of target lengths rather than a single fixed length
• An adapted version of PXDesignBench bundled directly in the repo to enable the above workflow immediately

I may not be able to package all of this cleanly into a PR in the near term, but please feel free to evaluate or incorporate any pieces of the code or logic from the fork as useful:
https://github.com/cytokineking/PXDesign