Enable Google Cloud Storage (gs://) path support in AlphaFold3 using etils.epath.

This change migrates path operations in AlphaFold3 from Python's standard `pathlib.Path` and `os.path` to the unified filesystem wrapper `etils.epath.Path`. This enables several core flags (`--json_path`, `--input_dir`, `--output_dir`, `--model_dir`, and `--pdb_database_path`) to natively support GCS cloud-hosted paths (`gs://...`).

Support for Google Cloud Storage paths is treated as an optional dependency. Open source users can opt into GCS path support by installing the `gcsfs` add-in, e.g., via: `uv sync --extra` , or adding `--build-arg UV_EXTRAS="--extra gcsfs"` to their docker build command.

PiperOrigin-RevId: 938058741
Change-Id: I3463fb159a082deca4e6752fceec7dae3d75c1d3

Author: DeepMind

docker/Dockerfile

@@ -77,8 +77,9 @@ WORKDIR /app/alphafold
 # --no-editable: install as a static package.
 # If using this as a recipe for local installation, we recommend removing the
 # --frozen and --no-editable flags.
-RUN --mount=type=cache,target=/root/.cache/uv \ 
-    UV_LINK_MODE=copy uv sync --frozen --all-groups --no-editable
+ARG UV_EXTRAS=""
+RUN --mount=type=cache,target=/root/.cache/uv \
+    UV_LINK_MODE=copy uv sync --frozen --all-groups --no-editable $UV_EXTRAS
 
 # Build chemical components database (this binary was installed by uv sync).
 RUN uv run build_data

docs/input.md

@@ -5,9 +5,20 @@
 You can provide inputs to `run_alphafold.py` in one of two ways:
 
 -   Single input file: Use the `--json_path` flag followed by the path to a
-    single JSON file.
+    single JSON file. This path can be either a local path or a Google Cloud
+    Storage path (starting with `gs://`), if enabled.
 -   Multiple input files: Use the `--input_dir` flag followed by the path to a
-    directory of JSON files.
+    directory of JSON files. This path can be either a local directory path or a
+    GCS path, if enabled.
+-   Within an input file, fields with names ending in `Path` (e.g.
+    `pairedMsaPath`) can be provided as Google Cloud Storage paths, if enabled.
+
+### Note: Google Cloud Storage paths
+
+Google Cloud Storage (`gs://`) paths are supported for certain flags and fields
+if the optional `gcsfs` dependency is installed. See the
+[installation instructions](installation.md#optional-enable-google-cloud-storage-path-support)
+for details and a complete list of supported paths.
 
 ## Input Format
 

docs/installation.md

@@ -17,8 +17,8 @@ should aid others with different setups. If you are installing locally outside
 of a Docker container, please ensure CUDA, cuDNN, and JAX are correctly
 installed; the
 [JAX installation documentation](https://jax.readthedocs.io/en/latest/installation.html#nvidia-gpu)
-is a useful reference for this case. Please note that the Docker container
-requires that the host machine has CUDA 12.6 installed.
+is a useful reference for this case. Note that the Docker container requires
+that the host machine has CUDA 12.6 installed.
 
 The instructions provided below describe how to:
 
@@ -368,6 +368,33 @@ docker run alphafold3 python run_alphafold.py --help
 
 for more information.
 
+:ledger: **Optional: Enable Google cloud storage path support**
+
+AlphaFold 3 supports reading and writing specific files directly from Google
+Cloud Storage (`gs://` paths). Since GCS is an optional feature, the required
+`gcsfs` dependency is not installed by default. To enable GCS support when
+building the Docker container, pass the `UV_EXTRAS` build argument when running
+`docker build`:
+
+```sh
+docker build --build-arg UV_EXTRAS="--extra gcsfs" -t alphafold3 -f docker/Dockerfile .
+```
+
+The following command-line flags support `gs://` paths:
+
+-   `--input_dir`
+-   `--json_path`
+-   `--output_dir`
+-   `--model_dir`
+-   `--pdb_database_path`
+-   In the JSON input file, fields with names ending in `Path` (e.g.
+    `pairedMsaPath`) can be provided as Google Cloud Storage paths, if enabled.
+
+:warning: All other database paths (e.g., `--db_dir`,
+`--small_bfd_database_path`, etc.) and binary paths (e.g.,
+`--jackhmmer_binary_path`) do **not** support `gs://` paths and must be local
+file paths.
+
 ## Running Using Singularity Instead of Docker
 
 You may prefer to run AlphaFold 3 within Singularity. You'll still need to

docs/output.md

@@ -9,6 +9,10 @@ AlphaFold 3 will write its outputs in a directory called `My_first_fold_TEST`
 append a timestamp to the directory name to avoid overwriting existing data
 unless `--force_output_dir` is passed.
 
+The output directory can also be a Google Cloud Storage path (`gs://`) if the
+`gcsfs` dependency is installed (see the
+[installation instructions](installation.md#optional-enable-google-cloud-storage-path-support)).
+
 The following structure is used within the output directory:
 
 *   Sub-directories with results for each sample and seed. There will be

pyproject.toml

@@ -17,6 +17,7 @@ license = {file = "LICENSE"}
 dependencies = [
     "absl-py>=2.3.1",
     "dm-haiku==0.0.16",
+    "etils[epath]",
     "jax==0.9.1",
     "jax[cuda12]==0.9.1",
     "numpy",
@@ -26,6 +27,12 @@ dependencies = [
     "zstandard",
 ]
 
+[project.optional-dependencies]
+gcsfs = [
+  # Use `uv sync --extra gcsfs` to enable GCS paths for inputs and outputs
+  "gcsfs",
+]
+
 [dependency-groups]
 dev = [
     "pytest>=6.0",

run_alphafold.py

@@ -34,7 +34,6 @@
 import datetime
 import functools
 import os
-import pathlib
 import shutil
 import string
 import textwrap
@@ -56,34 +55,34 @@
 from alphafold3.model import params
 from alphafold3.model import post_processing
 from alphafold3.model.components import utils
+from etils import epath
 import haiku as hk
 import jax
 from jax import numpy as jnp
 import numpy as np
 import tokamax
 
-_HOME_DIR = pathlib.Path.home()
+_HOME_DIR = epath.Path('~').expanduser()
 _DEFAULT_MODEL_DIR = _HOME_DIR / 'models'
 _DEFAULT_DB_DIR = _HOME_DIR / 'public_databases'
 
-
 # Input and output paths.
-_JSON_PATH = flags.DEFINE_string(
+_JSON_PATH = epath.DEFINE_path(
     'json_path',
     None,
     'Path to the input JSON file.',
 )
-_INPUT_DIR = flags.DEFINE_string(
+_INPUT_DIR = epath.DEFINE_path(
     'input_dir',
     None,
     'Path to the directory containing input JSON files.',
 )
-_OUTPUT_DIR = flags.DEFINE_string(
+_OUTPUT_DIR = epath.DEFINE_path(
     'output_dir',
     None,
     'Path to a directory where the results will be saved.',
 )
-MODEL_DIR = flags.DEFINE_string(
+MODEL_DIR = epath.DEFINE_path(
     'model_dir',
     _DEFAULT_MODEL_DIR.as_posix(),
     'Path to the model to use for inference.',
@@ -220,7 +219,7 @@
     ' calculation. Must be set for sharded databases.',
     lower_bound=0.0,
 )
-_PDB_DATABASE_PATH = flags.DEFINE_string(
+_PDB_DATABASE_PATH = epath.DEFINE_path(
     'pdb_database_path',
     '${DB_DIR}/mmcif_files',
     'PDB database directory with mmCIF files path, used for template search.',
@@ -420,11 +419,11 @@ def __init__(
       self,
       config: model.Model.Config,
       device: jax.Device,
-      model_dir: pathlib.Path,
+      model_dir: epath.PathLike,
   ):
     self._model_config = config
     self._device = device
-    self._model_dir = model_dir
+    self._model_dir = epath.Path(model_dir)
 
   @functools.cached_property
   def model_params(self) -> hk.Params:
@@ -605,19 +604,19 @@ def predict_structure(
 
 def write_fold_input_json(
     fold_input: folding_input.Input,
-    output_dir: os.PathLike[str] | str,
+    output_dir: epath.PathLike,
 ) -> None:
   """Writes the input JSON to the output directory."""
-  os.makedirs(output_dir, exist_ok=True)
-  path = os.path.join(output_dir, f'{fold_input.sanitised_name()}_data.json')
+  output_dir = epath.Path(output_dir)
+  output_dir.mkdir(parents=True, exist_ok=True)
+  path = output_dir / f'{fold_input.sanitised_name()}_data.json'
   print(f'Writing model input JSON to {path}')
-  with open(path, 'wt') as f:
-    f.write(fold_input.to_json())
+  path.write_text(fold_input.to_json())
 
 
 def write_outputs(
     all_inference_results: Sequence[ResultsForSeed],
-    output_dir: os.PathLike[str] | str,
+    output_dir: epath.PathLike,
     job_name: str,
     compress_large_output_files: bool = False,
 ) -> None:
@@ -627,15 +626,16 @@ def write_outputs(
   max_ranking_result = None
 
   output_terms = (
-      pathlib.Path(alphafold3.cpp.__file__).parent / 'OUTPUT_TERMS_OF_USE.md'
+      epath.Path(alphafold3.cpp.__file__).parent / 'OUTPUT_TERMS_OF_USE.md'
   ).read_text()
 
-  os.makedirs(output_dir, exist_ok=True)
+  output_dir = epath.Path(output_dir)
+  output_dir.mkdir(parents=True, exist_ok=True)
   for results_for_seed in all_inference_results:
     seed = results_for_seed.seed
     for sample_idx, result in enumerate(results_for_seed.inference_results):
-      sample_dir = os.path.join(output_dir, f'seed-{seed}_sample-{sample_idx}')
-      os.makedirs(sample_dir, exist_ok=True)
+      sample_dir = output_dir / f'seed-{seed}_sample-{sample_idx}'
+      sample_dir.mkdir(parents=True, exist_ok=True)
       post_processing.write_output(
           inference_result=result,
           output_dir=sample_dir,
@@ -649,21 +649,19 @@ def write_outputs(
         max_ranking_result = result
 
     if embeddings := results_for_seed.embeddings:
-      embeddings_dir = os.path.join(output_dir, f'seed-{seed}_embeddings')
-      os.makedirs(embeddings_dir, exist_ok=True)
+      embeddings_dir = output_dir / f'seed-{seed}_embeddings'
+      embeddings_dir.mkdir(parents=True, exist_ok=True)
       post_processing.write_embeddings(
           embeddings=embeddings,
           output_dir=embeddings_dir,
           name=f'{job_name}_seed-{seed}',
       )
 
     if (distogram := results_for_seed.distogram) is not None:
-      distogram_dir = os.path.join(output_dir, f'seed-{seed}_distogram')
-      os.makedirs(distogram_dir, exist_ok=True)
-      distogram_path = os.path.join(
-          distogram_dir, f'{job_name}_seed-{seed}_distogram.npz'
-      )
-      with open(distogram_path, 'wb') as f:
+      distogram_dir = output_dir / f'seed-{seed}_distogram'
+      distogram_dir.mkdir(parents=True, exist_ok=True)
+      distogram_path = distogram_dir / f'{job_name}_seed-{seed}_distogram.npz'
+      with distogram_path.open('wb') as f:
         np.savez_compressed(f, distogram=distogram.astype(np.float16))
 
   if max_ranking_result is not None:  # True iff ranking_scores non-empty.
@@ -677,29 +675,31 @@ def write_outputs(
     )
     # Save csv of ranking scores with seeds and sample indices, to allow easier
     # comparison of ranking scores across different runs.
-    with open(
-        os.path.join(output_dir, f'{job_name}_ranking_scores.csv'), 'wt'
-    ) as f:
+    ranking_scores_csv_path = output_dir / f'{job_name}_ranking_scores.csv'
+    with ranking_scores_csv_path.open('w') as f:
       writer = csv.writer(f)
       writer.writerow(['seed', 'sample', 'ranking_score'])
       writer.writerows(ranking_scores)
 
 
-def replace_db_dir(path_with_db_dir: str, db_dirs: Sequence[str]) -> str:
+def replace_db_dir(
+    path_with_db_dir: epath.PathLike, db_dirs: Sequence[str]
+) -> str:
   """Replaces the DB_DIR placeholder in a path with the given DB_DIR."""
+  path_with_db_dir = os.fspath(path_with_db_dir)
   template = string.Template(path_with_db_dir)
   if 'DB_DIR' in template.get_identifiers():
     for db_dir in db_dirs:
       path = template.substitute(DB_DIR=db_dir)
-      if os.path.exists(path):
+      if epath.Path(path).exists():
         return path
     raise FileNotFoundError(
         f'{path_with_db_dir} with ${{DB_DIR}} not found in any of {db_dirs}.'
     )
   if (sharded_paths := shards.get_sharded_paths(path_with_db_dir)) is not None:
-    db_exists = all(os.path.exists(p) for p in sharded_paths)
+    db_exists = all(epath.Path(p).exists() for p in sharded_paths)
   else:
-    db_exists = os.path.exists(path_with_db_dir)
+    db_exists = epath.Path(path_with_db_dir).exists()
   if not db_exists:
     raise FileNotFoundError(f'{path_with_db_dir} does not exist.')
   return path_with_db_dir
@@ -711,7 +711,7 @@ def process_fold_input(
     data_pipeline_config: pipeline.DataPipelineConfig | None,
     *,
     model_runner: None,
-    output_dir: os.PathLike[str] | str,
+    output_dir: epath.PathLike,
     buckets: Sequence[int] | None = None,
     ref_max_modified_date: datetime.date | None = None,
     conformer_max_iterations: int | None = None,
@@ -729,7 +729,7 @@ def process_fold_input(
     data_pipeline_config: pipeline.DataPipelineConfig | None,
     *,
     model_runner: ModelRunner,
-    output_dir: os.PathLike[str] | str,
+    output_dir: epath.PathLike,
     buckets: Sequence[int] | None = None,
     ref_max_modified_date: datetime.date | None = None,
     conformer_max_iterations: int | None = None,
@@ -746,7 +746,7 @@ def process_fold_input(
     data_pipeline_config: pipeline.DataPipelineConfig | None,
     *,
     model_runner: ModelRunner | None,
-    output_dir: os.PathLike[str] | str,
+    output_dir: epath.PathLike,
     buckets: Sequence[int] | None = None,
     ref_max_modified_date: datetime.date | None = None,
     conformer_max_iterations: int | None = None,
@@ -803,13 +803,11 @@ def process_fold_input(
   if not fold_input.chains:
     raise ValueError('Fold input has no chains.')
 
-  if (
-      not force_output_dir
-      and os.path.exists(output_dir)
-      and os.listdir(output_dir)
-  ):
+  output_dir = epath.Path(output_dir)
+  if not force_output_dir and output_dir.exists() and any(output_dir.iterdir()):
     new_output_dir = (
-        f'{output_dir}_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}'
+        output_dir.parent
+        / f'{output_dir.name}_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}'
     )
     print(
         f'Output will be written in {new_output_dir} since {output_dir} is'
@@ -874,13 +872,9 @@ def main(_):
     )
 
   if _INPUT_DIR.value is not None:
-    fold_inputs = folding_input.load_fold_inputs_from_dir(
-        pathlib.Path(_INPUT_DIR.value)
-    )
+    fold_inputs = folding_input.load_fold_inputs_from_dir(_INPUT_DIR.value)
   elif _JSON_PATH.value is not None:
-    fold_inputs = folding_input.load_fold_inputs_from_path(
-        pathlib.Path(_JSON_PATH.value)
-    )
+    fold_inputs = folding_input.load_fold_inputs_from_path(_JSON_PATH.value)
   else:
     raise AssertionError(
         'Exactly one of --json_path or --input_dir must be specified.'
@@ -891,7 +885,7 @@ def main(_):
 
   # Make sure we can create the output directory before running anything.
   try:
-    os.makedirs(_OUTPUT_DIR.value, exist_ok=True)
+    _OUTPUT_DIR.value.mkdir(parents=True, exist_ok=True)
   except OSError as e:
     print(f'Failed to create output directory {_OUTPUT_DIR.value}: {e}')
     raise
@@ -993,7 +987,7 @@ def main(_):
             return_distogram=_SAVE_DISTOGRAM.value,
         ),
         device=devices[_GPU_DEVICE.value],
-        model_dir=pathlib.Path(MODEL_DIR.value),
+        model_dir=MODEL_DIR.value,
     )
     # Check we can load the model parameters before launching anything.
     print('Checking that model parameters can be loaded...')
@@ -1010,7 +1004,7 @@ def main(_):
         fold_input=fold_input,
         data_pipeline_config=data_pipeline_config,
         model_runner=model_runner,
-        output_dir=os.path.join(_OUTPUT_DIR.value, fold_input.sanitised_name()),
+        output_dir=_OUTPUT_DIR.value / fold_input.sanitised_name(),
         buckets=tuple(int(bucket) for bucket in _BUCKETS.value),
         ref_max_modified_date=max_template_date,
         conformer_max_iterations=_CONFORMER_MAX_ITERATIONS.value,