Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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
4 changes: 3 additions & 1 deletion openfold3/core/data/tools/jackhmmer.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,9 @@ def db_remote_chunk(db_idx):
return f"{self.database_path}.{db_idx}"

def db_local_chunk(db_idx):
return f"/tmp/ramdisk/{db_basename}.{db_idx}"
from openfold3.core.data.tools.utils import get_of3_tmpdir

return str(get_of3_tmpdir("ramdisk") / f"{db_basename}.{db_idx}")

# Remove existing files to prevent OOM
for f in glob.glob(db_local_chunk("[0-9]*")):
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