Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/source/configuration_reference.md
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ Configures the ColabFold MSA server integration.
- `server_user_agent` *(str)*: User agent string (default: `openfold`)
- `server_url` *(Url)*: ColabFold server URL (default: `https://api.colabfold.com`)
- `save_mappings` *(bool)*: Save sequence ID mappings (default: `true`)
- `msa_output_directory` *(Path)*: Directory for MSA outputs (default: temporary directory)
- `msa_output_directory` *(Path)*: Directory for MSA outputs (default: `temporary directory/of3-of-<user>/colabfold_msas`)
- `cleanup_msa_dir` *(bool)*: Delete MSAs after processing (default: `true`)

**Example**:
Expand Down
8 changes: 4 additions & 4 deletions examples/reference_full_config/full_config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ msa_computation_settings:
server_user_agent: openfold
server_url: https://api.colabfold.com/
save_mappings: true
msa_output_directory: <tmp-dir>/of3_colabfold_msas/
msa_output_directory: <tmp-dir>/of3-of-<user>/colabfold_msas/
cleanup_msa_dir: true

# TemplatePreprocessorSettings: https://github.com/aqlaboratory/openfold-3/blob/main/openfold3/core/data/pipelines/preprocessing/template.py#L1459
Expand All @@ -152,11 +152,11 @@ template_preprocessor_settings:
create_logs: false
n_processes: 1
chunksize: 1
structure_directory: <tmp-dir>/of3_template_data/template_structures
structure_directory: <tmp-dir>/of3-of-<user>/template_data/template_structures
structure_file_format: cif
output_directory: <tmp-dir>/of3_template_data
output_directory: <tmp-dir>/of3-of-<user>/template_data
precache_directory: null
structure_array_directory: null
cache_directory: <tmp-dir>/of3_template_data/template_cache
cache_directory: <tmp-dir>/of3-of-<user>/template_data/template_cache
log_directory: null
ccd_file_path: null
9 changes: 5 additions & 4 deletions openfold3/core/data/pipelines/preprocessing/template.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import os
import random
import re
import tempfile
import traceback
from datetime import datetime
from functools import wraps
Expand Down Expand Up @@ -1601,9 +1600,11 @@ def _prepare_output_directories(self) -> "TemplatePreprocessorSettings":
"pipeline."
)

self.output_directory = (
self.output_directory or Path(tempfile.gettempdir()) / "of3_template_data"
)
if self.output_directory is None:
from openfold3.core.data.tools.utils import get_of3_tmpdir

self.output_directory = get_of3_tmpdir("template_data")

base = self.output_directory

# only set these if the user did not give them explicitly
Expand Down
9 changes: 6 additions & 3 deletions openfold3/core/data/tools/colabfold_msa_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
import random
import shutil
import tarfile
import tempfile
import time
import warnings
from dataclasses import dataclass, field
Expand Down Expand Up @@ -997,13 +996,17 @@ class MsaComputationSettings(BaseModel):
server_user_agent: str = "openfold"
server_url: Url = Url("https://api.colabfold.com")
save_mappings: bool = True
msa_output_directory: Path = Path(tempfile.gettempdir()) / "of3_colabfold_msas"
msa_output_directory: Path | None = None
cleanup_msa_dir: bool = True

@model_validator(mode="after")
def create_dir(self) -> "MsaComputationSettings":
"""Creates the output directory if it does not exist."""
if not self.msa_output_directory.exists():
if self.msa_output_directory is None:
from openfold3.core.data.tools.utils import get_of3_tmpdir

self.msa_output_directory = get_of3_tmpdir("colabfold_msas")
elif not self.msa_output_directory.exists():
self.msa_output_directory.mkdir(parents=True, exist_ok=True)
return self

Expand Down
17 changes: 17 additions & 0 deletions openfold3/core/data/tools/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,27 @@

import contextlib
import datetime
import getpass
import logging
import shutil
import tempfile
import time
from pathlib import Path


def get_of3_tmpdir(subdir: str | None = None) -> Path:
"""Return a user-namespaced OpenFold3 temporary directory.

Follows the same convention as pytest (``/tmp/pytest-of-<user>/``):
``<tmpdir>/of3-of-<user>/<subdir>``. Respects ``$TMPDIR`` and other
platform conventions via :func:`tempfile.gettempdir`.

The returned directory is created on disk if it does not already exist.
"""
base = Path(tempfile.gettempdir()) / f"of3-of-{getpass.getuser()}"
path = base / subdir if subdir else base
path.mkdir(parents=True, exist_ok=True)
return path


@contextlib.contextmanager
Expand Down
77 changes: 77 additions & 0 deletions openfold3/tests/core/data/tools/test_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
"""Tests for ``openfold3.core.data.tools.utils``."""

import getpass
import tempfile
from pathlib import Path
from unittest.mock import patch

from openfold3.core.data.tools.utils import get_of3_tmpdir


class TestGetOf3Tmpdir:
"""Tests for get_of3_tmpdir."""

def test_returns_path(self):
result = get_of3_tmpdir("test_subdir")
assert isinstance(result, Path)

def test_directory_is_created(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
result = get_of3_tmpdir("mysubdir")

assert result.is_dir()

def test_contains_username(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
result = get_of3_tmpdir("subdir")

assert f"of3-of-{getpass.getuser()}" in str(result)

def test_subdir_appended(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
result = get_of3_tmpdir("colabfold_msas")

assert result.name == "colabfold_msas"
assert result.parent.name == f"of3-of-{getpass.getuser()}"

def test_no_subdir(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
result = get_of3_tmpdir()

assert result.name == f"of3-of-{getpass.getuser()}"
assert result.is_dir()

def test_respects_tmpdir_env(self, tmp_path, monkeypatch):
custom_tmp = tmp_path / "custom_tmp"
custom_tmp.mkdir()
monkeypatch.setenv("TMPDIR", str(custom_tmp))
# Force tempfile to re-evaluate TMPDIR
tempfile.tempdir = None

result = get_of3_tmpdir("subdir")

assert str(result).startswith(str(custom_tmp))

def test_idempotent(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
first = get_of3_tmpdir("subdir")
second = get_of3_tmpdir("subdir")

assert first == second

def test_different_subdirs_are_isolated(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
a = get_of3_tmpdir("aaa")
b = get_of3_tmpdir("bbb")

assert a != b
assert a.parent == b.parent

def test_different_users_are_isolated(self, tmp_path):
with patch.object(tempfile, "gettempdir", return_value=str(tmp_path)):
real = get_of3_tmpdir("data")
with patch.object(getpass, "getuser", return_value="other_user"):
other = get_of3_tmpdir("data")

assert real != other
assert "of3-of-other_user" in str(other)
1 change: 0 additions & 1 deletion openfold3/tests/test_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,6 @@ def write_confidence_scores(
def test_full_confidence_scores_written(
self, tmp_path, output_fmt, dummy_confidence_scores
):

self.write_confidence_scores(
tmp_path, output_fmt, True, dummy_confidence_scores
)
Expand Down
Loading