feat(server): compose AlphaFold CPU and GPU stages

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Server:
+  - workflow Composer: ordered, resumable task stages reuse existing Docker/SLURM jobs with independently snapshotted resource policies; AlphaFold MSA/features now run CPU-only before GPU model/relax.
   - version-2 scientific workspaces: modular task-selected RFdiffusion modes, Mol* residue selection (with result controls hidden by default), declarative linked result views, bounded table pages, and EASIFA table-to-structure mapping. Breaking: every task type must declare `input_workspace` explicitly — startup fails closed when a custom registry omits it.
   - runner: patch RFdiffusion's checkpoint-override parsing (`bool("false")` is True), so binder runs respect `preprocess.sidechain_input=false` instead of crashing in the broken upstream sidechain path.
   - runner: new alphafold family (official google-deepmind/alphafold @ c77e5d2a) — monomer/pTM/multimer presets, full_dbs MSA, Amber relaxation (best by default); DBs and the 2022-12-06 params release ro-mounted from /mnt/db.
@@ -56,6 +57,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Server:
+  - AlphaFold multimer full-database runs now pass the required UniRef30 database path.
+  - prepared SLURM deploys build staged SIFs from the matching `:next` runner image instead of silently reusing `:latest`.
+  - create-task parameters: choice controls now serialize their selected value, preserving AlphaFold multimer and every other non-default select option.
   - Result polling/viewers: terminal task states reload once without overlapping polls; pending result pages keep polling; Mol* teardown completes before iframe removal, stale teardown continuations cannot replace newer previews, and preview loaders remain visible.
   - AlphaFold stages: drain the stderr translator before wrapper exit and preserve both process statuses so final stage markers cannot be lost.
   - SLURM stages: `srun -u`, unbuffered AlphaFold stderr, and a Python translator stream wrapper phases live, so `run_stage` records intermediates before job exit.

server/README.md

@@ -449,6 +449,27 @@ and a watchdog kills work that exceeds the snapshotted runtime. Invalid fields
 fail closed in the admin API and again at submission/launch rather than being
 silently discarded.
 
+### Ordered workflows
+
+A task type may declare an ordered `workflow` whose stages reuse the same
+runtime family and immutable task snapshot. The worker acts as a lightweight
+Composer: it submits one existing `Job` at a time, validates that allocation's
+outputs through the runner contract, persists the stage/job state, and releases
+the next stage only after success. Each stage has an independently snapshotted
+resource policy in the configuration UI.
+
+```text
+Submission -> Composer -> [CPU: MSA + features] -> [GPU: model + relax] -> Results
+                           | features.pkl          | ranked_0.pdb
+                           +---- persisted --------+
+```
+
+AlphaFold is the first composed task. `alphafold.features` runs MSA and feature
+construction without GPU GRES or Apptainer `--nv`; after `features.pkl` is
+validated, `alphafold.model` receives the GPU allocation for inference and
+optional relaxation. Restarts cancel only the active allocation and resume at
+the first incomplete stage.
+
 ## 5. Authentication
 
 The server uses Bearer-token authentication (replaces the old HTTP Basic Auth + `users.txt` model).
@@ -1187,7 +1208,8 @@ job is enqueued.
 The allocation wrapper writes `REVODESIGN_JOB_ID=<numeric-id>` as its first
 stdout line. The worker stores that real SLURM ID in the `slurm_job_id` column
 of the tasks table so cancellation and restart recovery can address the
-allocation directly.
+allocation directly. Composed tasks additionally persist all stage handles and
+states in `workflow_state`; `slurm_job_id` remains the currently active handle.
 
 ### 13.8 Live Output
 

server/config/task_types.yaml

@@ -754,6 +754,17 @@ task_types:
       featuring: "Building features"
       modeling: "Folding models"
       relaxing: "Amber relaxation"
+    workflow:
+      - name: features
+        display_name: "MSA and features"
+        requires_gpu: false
+        runner_args: ["-s", "features"]
+        stage_markers: ["msa_searching", "featuring"]
+      - name: model
+        display_name: "Model and relax"
+        requires_gpu: true
+        runner_args: ["-s", "model"]
+        stage_markers: ["modeling", "relaxing"]
     params:
       - name: "model_preset"
         type: "str"

server/docker/runners/alphafold/Dockerfile

@@ -4,9 +4,11 @@
 FROM alpine/git:2.47.2 AS alphafold-source
 ARG ALPHAFOLD_REPO=https://github.com/google-deepmind/alphafold.git
 ARG ALPHAFOLD_REF=c77e5d2a8961d1a353632c462914ff0a32a950f6
+COPY ./docker/runners/alphafold/staged_pipeline.patch /tmp/staged_pipeline.patch
 RUN git init /opt/alphafold && git -C /opt/alphafold remote add origin ${ALPHAFOLD_REPO} && \
     git -C /opt/alphafold fetch --depth 1 origin ${ALPHAFOLD_REF} && git -C /opt/alphafold checkout --detach FETCH_HEAD && \
-    rm -rf /opt/alphafold/.git
+    git -C /opt/alphafold apply --check /tmp/staged_pipeline.patch && \
+    git -C /opt/alphafold apply /tmp/staged_pipeline.patch && rm -rf /opt/alphafold/.git
 
 FROM --platform=linux/amd64 python:3.11-slim
 

server/docker/runners/alphafold/run.sh

@@ -4,9 +4,11 @@ set -e
 task_context_src="${TASK_CONTEXT_SRC:-/app/revocompute/task_context.sh}"
 [[ -f "$task_context_src" ]] && source "$task_context_src"
 
-usage() { echo "Usage: $0 -i <task.json> -o <output_dir>"; exit 1; }
-while getopts ":i:o:" opt; do case "${opt}" in i) input_file=$OPTARG ;; o) output_dir=$OPTARG ;; ?) usage ;; esac; done
+usage() { echo "Usage: $0 -i <task.json> -o <output_dir> [-s all|features|model]"; exit 1; }
+run_stage=all
+while getopts ":i:o:s:" opt; do case "${opt}" in i) input_file=$OPTARG ;; o) output_dir=$OPTARG ;; s) run_stage=$OPTARG ;; ?) usage ;; esac; done
 [[ -z "${input_file:-}" || -z "${output_dir:-}" ]] && usage
+[[ "$run_stage" =~ ^(all|features|model)$ ]] || usage
 input_file=$(readlink -f "$input_file"); output_dir=$(readlink -f "$output_dir")
 [[ ! -f "$input_file" ]] && { echo "Task manifest not found: $input_file"; exit 1; }
 mkdir -p "$output_dir"
@@ -18,6 +20,15 @@ NUM_MULTIMER=$(_parse_param num_multimer_predictions_per_model 1)
 MODELS_TO_RELAX=$(_parse_param models_to_relax best)
 BENCHMARK=$(_parse_param benchmark false)
 fasta_path=$(primary_input)
+fasta_name=$(basename "$fasta_path")
+fasta_name=${fasta_name%.*}
+features_path="${output_dir}/${fasta_name}/features.pkl"
+features_marker="${output_dir}/.alphafold-features-complete"
+use_gpu_relax=true
+[[ "$run_stage" == features ]] && use_gpu_relax=false
+if [[ "$run_stage" == model ]]; then
+  [[ -s "$features_path" && -f "$features_marker" ]] || { echo "Validated AlphaFold features are missing" >&2; exit 1; }
+fi
 
 # A100 memory behaviour: let JAX overcommit via unified memory.
 export TF_FORCE_UNIFIED_MEMORY=1
@@ -34,13 +45,15 @@ af_args=(
   "--db_preset=${DB_PRESET}"
   "--model_preset=${MODEL_PRESET}"
   "--models_to_relax=${MODELS_TO_RELAX}"
-  "--use_gpu_relax=true"
+  "--use_gpu_relax=${use_gpu_relax}"
+  "--run_stage=${run_stage}"
   "--benchmark=${BENCHMARK}"
   "--bfd_database_path=${DB}/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt"
   "--mgnify_database_path=${DB}/mgnify/mgy_clusters.fa"
   "--template_mmcif_dir=${DB}/pdb_mmcif/mmcif_files"
   "--obsolete_pdbs_path=${DB}/pdb_mmcif/obsolete.dat"
   "--uniref90_database_path=${DB}/uniref90/uniref90.fasta"
+  "--uniref30_database_path=${DB}/uniref30_uc30/UniRef30_2022_02/UniRef30_2022_02"
 )
 if [[ "$MODEL_PRESET" == "multimer" ]]; then
   af_args+=(
@@ -51,7 +64,6 @@ if [[ "$MODEL_PRESET" == "multimer" ]]; then
 else
   af_args+=(
     "--pdb70_database_path=${DB}/pdb70/pdb70"
-    "--uniref30_database_path=${DB}/uniref30_uc30/UniRef30_2022_02/UniRef30_2022_02"
   )
 fi
 
@@ -90,6 +102,13 @@ if [[ ${alphafold_status} -ne 0 || ${translator_status} -ne 0 ]]; then
   exit "${translator_status}"
 fi
 
+if [[ "$run_stage" == features ]]; then
+  [[ -s "$features_path" ]] || { echo "AlphaFold produced no features.pkl" >&2; exit 1; }
+  touch "$features_marker"
+  echo "AlphaFold feature construction complete."
+  exit 0
+fi
+
 [[ -n "$(ls "${output_dir}"/*/ranked_0.pdb 2>/dev/null || true)" ]] || {
   echo "AlphaFold produced no ranked_0.pdb" >&2; exit 1; }
 touch "${output_dir}/task_finished"

server/docker/runners/alphafold/staged_pipeline.patch

@@ -0,0 +1,116 @@
+diff --git a/run_alphafold.py b/run_alphafold.py
+index 3d8a4f4..a41fb00 100644
+--- a/run_alphafold.py
++++ b/run_alphafold.py
+@@ -231,6 +231,12 @@ flags.DEFINE_boolean(
+     'recommended to enable if possible. GPUs must be available'
+     ' if this setting is enabled.',
+ )
++flags.DEFINE_enum(
++    'run_stage',
++    'all',
++    ['all', 'features', 'model'],
++    'Run the complete pipeline, feature construction only, or modeling only.',
++)
+ flags.DEFINE_integer(
+     'jackhmmer_n_cpu',
+     # Unfortunately, os.process_cpu_count() is only available in Python 3.13+.
+@@ -364,17 +370,22 @@ def predict_structure(
+   if not os.path.exists(msa_output_dir):
+     os.makedirs(msa_output_dir)
+
+-  # Get features.
+-  t_0 = time.time()
+-  feature_dict = data_pipeline.process(
+-      input_fasta_path=fasta_path, msa_output_dir=msa_output_dir
+-  )
+-  timings['features'] = time.time() - t_0
+-
+-  # Write out features as a pickled dictionary.
+   features_output_path = os.path.join(output_dir, 'features.pkl')
+-  with open(features_output_path, 'wb') as f:
+-    pickle.dump(feature_dict, f, protocol=4)
++  if FLAGS.run_stage == 'model':
++    logging.info('Reading precomputed features from %s', features_output_path)
++    with open(features_output_path, 'rb') as f:
++      feature_dict = pickle.load(f)
++  else:
++    t_0 = time.time()
++    feature_dict = data_pipeline.process(
++        input_fasta_path=fasta_path, msa_output_dir=msa_output_dir
++    )
++    timings['features'] = time.time() - t_0
++    with open(features_output_path, 'wb') as f:
++      pickle.dump(feature_dict, f, protocol=4)
++    if FLAGS.run_stage == 'features':
++      logging.info('Feature construction complete for %s', fasta_name)
++      return
+
+   unrelaxed_pdbs = {}
+   unrelaxed_proteins = {}
+@@ -667,36 +678,40 @@ def main(argv):
+     data_pipeline = monomer_data_pipeline
+
+   model_runners = {}
+-  model_names = config.MODEL_PRESETS[FLAGS.model_preset]
+-  for model_name in model_names:
+-    model_config = config.model_config(model_name)
+-    if run_multimer_system:
+-      model_config.model.num_ensemble_eval = num_ensemble
+-    else:
+-      model_config.data.eval.num_ensemble = num_ensemble
+-    model_params = data.get_model_haiku_params(
+-        model_name=model_name, data_dir=FLAGS.data_dir
+-    )
+-    model_runner = model.RunModel(model_config, model_params)
+-    for i in range(num_predictions_per_model):
+-      model_runners[f'{model_name}_pred_{i}'] = model_runner
++  amber_relaxer = None
++  if FLAGS.run_stage != 'features':
++    model_names = config.MODEL_PRESETS[FLAGS.model_preset]
++    for model_name in model_names:
++      model_config = config.model_config(model_name)
++      if run_multimer_system:
++        model_config.model.num_ensemble_eval = num_ensemble
++      else:
++        model_config.data.eval.num_ensemble = num_ensemble
++      model_params = data.get_model_haiku_params(
++          model_name=model_name, data_dir=FLAGS.data_dir
++      )
++      model_runner = model.RunModel(model_config, model_params)
++      for i in range(num_predictions_per_model):
++        model_runners[f'{model_name}_pred_{i}'] = model_runner
+
+-  logging.info(
+-      'Have %d models: %s', len(model_runners), list(model_runners.keys())
+-  )
++    logging.info(
++        'Have %d models: %s', len(model_runners), list(model_runners.keys())
++    )
+
+-  amber_relaxer = relax.AmberRelaxation(
+-      max_iterations=RELAX_MAX_ITERATIONS,
+-      tolerance=RELAX_ENERGY_TOLERANCE,
+-      stiffness=RELAX_STIFFNESS,
+-      exclude_residues=RELAX_EXCLUDE_RESIDUES,
+-      max_outer_iterations=RELAX_MAX_OUTER_ITERATIONS,
+-      use_gpu=FLAGS.use_gpu_relax,
+-  )
++    amber_relaxer = relax.AmberRelaxation(
++        max_iterations=RELAX_MAX_ITERATIONS,
++        tolerance=RELAX_ENERGY_TOLERANCE,
++        stiffness=RELAX_STIFFNESS,
++        exclude_residues=RELAX_EXCLUDE_RESIDUES,
++        max_outer_iterations=RELAX_MAX_OUTER_ITERATIONS,
++        use_gpu=FLAGS.use_gpu_relax,
++    )
+
+   random_seed = FLAGS.random_seed
+   if random_seed is None:
+-    random_seed = random.randrange(sys.maxsize // len(model_runners))
++    random_seed = random.randrange(
++        sys.maxsize // max(len(model_runners), len(config.MODEL_PRESETS[FLAGS.model_preset]))
++    )
+   logging.info('Using random seed %d for the data pipeline', random_seed)
+
+   # Predict structure for each of the sequences.

server/revocompute/app.py

@@ -198,6 +198,20 @@ def _add_security_headers(response):
     if manage_db.task_type_get(_tt.name) is None:
         manage_db.task_type_upsert(_tt.name, enabled=True)
         _log.info("Seeded task_type_config for %r (enabled=true)", _tt.name)
+    _parent_resources = manage_db.task_type_get(_tt.name) or {}
+    for _stage in _tt.workflow:
+        if manage_db.task_type_get(_stage.name) is None:
+            _initial = {"enabled": True}
+            if _stage.requires_gpu:
+                _initial.update(
+                    {
+                        key: value
+                        for key, value in _parent_resources.items()
+                        if key not in {"tool", "enabled"} and value is not None
+                    }
+                )
+            manage_db.task_type_upsert(_stage.name, **_initial)
+            _log.info("Seeded workflow resource profile %r", _stage.name)
 
 
 def _is_binary_file(path: str) -> bool:

server/revocompute/db.py

@@ -87,6 +87,7 @@ def __init__(self, path: str):
             Column("input_form", Text),
             Column("slurm_job_id", String),
             Column("container_id", String),
+            Column("workflow_state", Text),
         )
         Index("idx_tasks_uploaded_at", self.tasks_table.c.uploaded_at)
         self._initialize()
@@ -106,7 +107,7 @@ def _initialize(self) -> None:
             # create_all does not add columns to existing tables — backfill
             # ones added after a table first shipped (idempotent).
             existing = {row[1] for row in conn.exec_driver_sql("PRAGMA table_info(tasks)")}
-            for column in ("container_id",):
+            for column in ("container_id", "workflow_state"):
                 if column not in existing:
                     conn.exec_driver_sql(f"ALTER TABLE tasks ADD COLUMN {column} VARCHAR")
 
@@ -160,9 +161,9 @@ def upsert_task(self, md5sum: str, **fields) -> None:
         with self.engine.begin() as conn:
             conn.execute(stmt)
 
-    def update_task(self, md5sum: str, **fields) -> None:
+    def update_task(self, md5sum: str, **fields) -> bool:
         if not fields:
-            return
+            return False
         status = fields.get("status")
         if status:
             self._ensure_status(status)
@@ -174,7 +175,35 @@ def update_task(self, md5sum: str, **fields) -> None:
         if status is None or (not self._is_deleted_status(status)):
             stmt = stmt.where(self.tasks_table.c.status.notin_(tuple(self.TERMINAL_STATUSES)))
         with self.engine.begin() as conn:
-            conn.execute(stmt)
+            return conn.execute(stmt).rowcount == 1
+
+    def claim_task_recovery(self, md5sum: str, *, expected_status: str) -> bool:
+        """Atomically move one orphaned active task out of recovery scans."""
+        if expected_status not in {"queued", "running"}:
+            return False
+        stmt = (
+            update(self.tasks_table)
+            .where(
+                self.tasks_table.c.md5sum == md5sum,
+                self.tasks_table.c.status == expected_status,
+            )
+            .values(status="pending")
+        )
+        with self.engine.begin() as conn:
+            return conn.execute(stmt).rowcount == 1
+
+    def claim_task_cancellation(self, md5sum: str, **fields) -> bool:
+        """Atomically cancel a task only while it remains active."""
+        stmt = (
+            update(self.tasks_table)
+            .where(
+                self.tasks_table.c.md5sum == md5sum,
+                self.tasks_table.c.status.in_(("pending", "queued", "running")),
+            )
+            .values(status="cancelled", **fields)
+        )
+        with self.engine.begin() as conn:
+            return conn.execute(stmt).rowcount == 1
 
     def claim_task_cleanup(
         self,