diff --git a/README.md b/README.md index 0815dd8..b97b6e4 100644 --- a/README.md +++ b/README.md @@ -340,6 +340,26 @@ prediction_container: "/path/to/images/alphafold3-2.5.0.sif" +### Container temporary directories + +Errors such as `ptxas fatal: Could not open output file` for a `tmpxft_*` file can +mean the job's inherited `TMPDIR` points to scratch that is not mounted inside +the container. Python may silently choose another temporary directory, but +native tools such as XLA's `ptxas` use `TMPDIR` directly. + +Every container rule checks `TMPDIR` before starting work. A usable directory is +kept without logging; an unset or inaccessible path falls back to `/tmp`, with +one `[container-tmpdir]` message in the job log. To choose another fallback: + +```yaml +container_tmpdir: /path/visible/inside/container +``` + +The fallback must be writable and searchable inside the container. The guard +creates it when possible and stops with a clear error if it is unusable. This +setting does not add a host bind mount; configure site mounts separately if the +fallback is outside the container's existing mounts. + ### GPU compatibility The containers carry their own CUDA runtime (pip `nvidia-*` wheels), so GPU support depends on the diff --git a/config/config.yaml b/config/config.yaml index 4f63247..6980bbb 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -88,6 +88,9 @@ exclude_permutations: true ## ADVANCED -- optional features and resource tuning ################################################################################ +# Writable fallback inside containers when the inherited TMPDIR is unavailable. +container_tmpdir: /tmp + ## -- Batch small predictions ------------------------------------------------- # Run several predictions in one job to share model loading and queue time. # 1 keeps one prediction per job. batch_max_tokens limits the combined residue diff --git a/test/test_container_tmpdir.py b/test/test_container_tmpdir.py new file mode 100644 index 0000000..37c3427 --- /dev/null +++ b/test/test_container_tmpdir.py @@ -0,0 +1,224 @@ +"""Container shells repair an inherited TMPDIR before native tools write into it. + +Like the inference command-line tests, these render the real workflow with the +resident-batch fixtures and execute its shell prologue under bash strict mode. +No container image, GPU, database, or cluster allocation is needed. +""" + +from __future__ import annotations + +import importlib.machinery +import importlib.util +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +_REPOSITORY = Path(__file__).resolve().parents[1] +_FIXTURES = _REPOSITORY / "test" / "fixtures" / "resident_batch" +_COMMON = _REPOSITORY / "workflow" / "rules" / "common.smk" +_loader = importlib.machinery.SourceFileLoader("aps_common", str(_COMMON)) +_spec = importlib.util.spec_from_loader("aps_common", _loader) +common = importlib.util.module_from_spec(_spec) +_loader.exec_module(common) + + +def _dry_run(tmp_path, **extra): + snakemake = shutil.which("snakemake") + if snakemake is None: + pytest.skip("Snakemake executable is not available") + config = yaml.safe_load((_FIXTURES / "af3.yaml").read_text(encoding="utf-8")) + config.update( + output_directory=str(tmp_path / "output"), + structure_inference_unified_memory=False, + **extra, + ) + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + completed = subprocess.run( + [ + snakemake, + "--snakefile", str(_REPOSITORY / "workflow" / "Snakefile"), + "--configfile", str(config_path), + "--cores", "1", + "--rerun-triggers", "mtime", + "--dry-run", + "--printshellcmds", + ], + cwd=_REPOSITORY, + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "APPTAINER_BINDPATH": "", "SINGULARITY_BINDPATH": ""}, + ) + assert completed.returncode == 0, completed.stdout[-2000:] + completed.stderr[-2000:] + jobs = re.split(r"^(?=(?:local)?rule \w+:$)", completed.stdout, flags=re.MULTILINE) + return [job for job in jobs if re.match(r"(?:local)?rule \w+:$", job, re.MULTILINE)] + + +def test_inference_native_tempfile_write_recovers_from_unavailable_tmpdir(tmp_path): + """A native writer uses TMPDIR directly, unlike Python's tempfile fallback.""" + fallback = tmp_path / "container-tmp" + jobs = _dry_run(tmp_path, container_tmpdir=str(fallback)) + job, = [job for job in jobs if job.startswith("rule structure_inference:")] + prologue = re.split( + r'if \[ "(?:true|false)" = "true" \] && command -v run_structure_prediction_batch.py', + job.split("Shell command:", 1)[1], + maxsplit=1, + )[0] + # An inaccessible parent makes the inherited path unusable even when tests + # run as root; it stands in for a job-local directory absent in the container. + blocked = tmp_path / "unbound-scratch" + blocked.write_text("not a mounted directory", encoding="utf-8") + completed = subprocess.run( + ["bash", "-c", 'set -euo pipefail\n' + prologue + '\n: > "$TMPDIR/tmpxft_regression"\n'], + capture_output=True, + text=True, + timeout=30, + env={**os.environ, "TMPDIR": str(blocked / "job-tmp")}, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert (fallback / "tmpxft_regression").is_file() + assert completed.stderr.count("[container-tmpdir]") == 1 + + +def _run_guard(tmp_path, fallback, inherited=None): + env = dict(os.environ) + env.pop("TMPDIR", None) + if inherited is not None: + env["TMPDIR"] = str(inherited) + return subprocess.run( + [ + "bash", "-c", + 'set -euo pipefail\n' + common.container_tmpdir_guard(str(fallback)) + + '\n: > "$TMPDIR/tmpxft_regression"\nprintf "%s\\n" "$TMPDIR"\n', + ], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=30, + env=env, + ) + + +@pytest.mark.parametrize("inherited", [None, ""]) +def test_unset_or_empty_tmpdir_uses_fallback(tmp_path, inherited): + fallback = tmp_path / "fallback" + completed = _run_guard(tmp_path, fallback, inherited) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == str(fallback) + assert (fallback / "tmpxft_regression").is_file() + assert completed.stderr.count("[container-tmpdir]") == 1 + + +@pytest.mark.parametrize("exists", [True, False]) +def test_usable_inherited_tmpdir_is_preserved_without_logging(tmp_path, exists): + inherited = tmp_path / "job scratch's $HOME" + if exists: + inherited.mkdir() + fallback = tmp_path / "fallback" + completed = _run_guard(tmp_path, fallback, inherited) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == str(inherited) + assert (inherited / "tmpxft_regression").is_file() + assert completed.stderr == "" + assert not fallback.exists() + + +@pytest.mark.parametrize("mode", [0o555, 0o666]) +def test_read_only_or_unsearchable_tmpdir_uses_fallback(tmp_path, mode): + if os.geteuid() == 0: + pytest.skip("root bypasses directory permissions") + inherited = tmp_path / "inaccessible" + inherited.mkdir(mode=mode) + fallback = tmp_path / "fallback" + try: + completed = _run_guard(tmp_path, fallback, inherited) + finally: + inherited.chmod(0o755) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == str(fallback) + assert (fallback / "tmpxft_regression").is_file() + assert completed.stderr.count("[container-tmpdir]") == 1 + + +@pytest.mark.parametrize("kind", ["file", "read-only", "unsearchable"]) +def test_unusable_fallback_fails_before_native_work(tmp_path, kind): + fallback = tmp_path / "unusable" + if kind == "file": + fallback.write_text("not a directory", encoding="utf-8") + else: + if os.geteuid() == 0: + pytest.skip("root bypasses directory permissions") + fallback.mkdir(mode=0o555 if kind == "read-only" else 0o666) + try: + completed = _run_guard(tmp_path, fallback) + finally: + fallback.chmod(0o755) + assert completed.returncode != 0 + assert completed.stdout == "" + assert "[container-tmpdir] fallback" in completed.stderr + assert "not writable in the container" in completed.stderr + assert not (fallback / "tmpxft_regression").exists() + + +def test_fallback_path_is_literal_shell_data(tmp_path): + fallback = tmp_path / 'scratch space\'s "$HOME" `touch INJECTED` $(touch INJECTED) {cache}' + completed = _run_guard(tmp_path, fallback) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == str(fallback) + assert (fallback / "tmpxft_regression").is_file() + assert f"using {fallback}" in completed.stderr + assert not (tmp_path / "INJECTED").exists() + + +def test_every_container_rule_starts_with_the_rendered_guard(tmp_path): + """Both feature modes together render every container rule, including reports.""" + source = (_REPOSITORY / "workflow" / "Snakefile").read_text(encoding="utf-8") + rules = re.split(r"^[ \t]*rule (\w+):", source, flags=re.MULTILINE) + container_rules = { + name for name, body in zip(rules[1::2], rules[2::2]) + if re.search(r"^[ \t]*container:", body, re.MULTILINE) + } + rendered = set() + for local_mmseqs in (False, True): + case = tmp_path / str(local_mmseqs) + data = case / "output" / "data" + data.mkdir(parents=True) + for protein in ("A", "B"): + (data / f"{protein}.fasta").write_text(f">{protein}\nACDE\n", encoding="utf-8") + fallback = case / "scratch's $HOME {cache}" + extra = { + "feature_directory": [], + "enable_structure_analysis": True, + "generate_recursive_report": True, + "container_tmpdir": str(fallback), + "mmseqs2_features": {"enabled": local_mmseqs}, + } + if local_mmseqs: + database_dir = case / "databases" + database_dir.mkdir() + databases = {} + for name in ("uniref90", "mgnify", "small_bfd", "uniprot"): + (database_dir / f"{name}.index").write_text("0\t0\t4\n", encoding="utf-8") + databases[name] = {"path": str(database_dir / name), "identifier": f"test-{name}"} + extra["mmseqs2_features"].update( + temp_dir=str(case / "mmseqs-tmp"), + template_database_ids={"pdb_seqres": "test-pdb-seqres", "mmcif": "test-mmcif"}, + databases=databases, + ) + for job in _dry_run(case, **extra): + name = re.match(r"(?:local)?rule (\w+):", job).group(1) + if name not in container_rules or "Shell command:" not in job: + continue + shell = job.split("Shell command:", 1)[1].lstrip() + # The local-MMseqs create_features variant is a Python symlink rule. + if shell.splitlines()[0] == "None": + continue + assert shell.startswith(common.container_tmpdir_guard(str(fallback))), name + rendered.add(name) + assert rendered == container_rules diff --git a/workflow/Snakefile b/workflow/Snakefile index 1644cb3..c4ffb97 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -19,6 +19,7 @@ from alphapulldown_input_parser.parser import FoldDataset, generate_fold_specifi config["output_directory"] = abspath(config["output_directory"]) makedirs(config["output_directory"], exist_ok=True) +CONTAINER_TMPDIR = config.get("container_tmpdir", "/tmp") feature_directories = config.get("feature_directory", []) if isinstance(feature_directories, (str, Path)): @@ -703,6 +704,7 @@ if LOCAL_MMSEQS2.enabled: "{mmseqs_shard}.json", ), params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), msa_output_directory=LOCAL_MMSEQS2_MSA_CACHE, fasta_paths=lambda wildcards, input: ",".join( str(path) for path in input.fastas @@ -745,6 +747,7 @@ if LOCAL_MMSEQS2.enabled: config["prediction_container"], shell: r""" + {params.tmpdir_guard} create_batch_msas.py \ --fasta_paths="{params.fasta_paths}" \ --msa_output_dir={params.msa_output_directory} \ @@ -767,6 +770,7 @@ if LOCAL_MMSEQS2.enabled: wildcard_constraints: fasta_basename="[^/]+", params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), data_directory=config["databases_directory"], output_directory=LOCAL_MMSEQS2_CACHE, msa_output_directory=LOCAL_MMSEQS2_MSA_CACHE, @@ -814,6 +818,7 @@ if LOCAL_MMSEQS2.enabled: config["prediction_container"], shell: r""" + {params.tmpdir_guard} finalize_batch_features.py \ --fasta_paths={input.fasta} \ --msa_input_dir={params.msa_output_directory} \ @@ -868,6 +873,7 @@ else: wildcard_constraints: fasta_basename="[^/]+", params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), data_directory=config["databases_directory"], output_directory=join(config["output_directory"], "features"), cli_parameters=" ".join( @@ -902,6 +908,7 @@ else: config["prediction_container"], shell: r""" + {params.tmpdir_guard} create_individual_features.py \ --fasta_paths={input} \ --data_dir={params.data_directory} \ @@ -945,6 +952,7 @@ rule structure_inference: output: join(config["output_directory"], "predictions", "{batch}", "completed_fold.txt"), params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), data_directory=config["backend_weights_directory"], feature_directory=join(config["output_directory"], "features"), # Parallel space-separated lists (one entry per fold in the batch) consumed by @@ -1011,6 +1019,7 @@ rule structure_inference: config["prediction_container"], shell: r""" + {params.tmpdir_guard} # Keep only JAX's own entries in the shared compile cache. XLA's per-fusion # autotune cache fails an atomic rename on shared filesystems and then segfaults # the compiler (jaxlib 0.9.1, the AF3 image). Newer AlphaPulldown sets this @@ -1100,6 +1109,7 @@ if ENABLE_STRUCTURE_ANALYSIS: qos=DEFAULT_SLURM_QOS, **linear_resources(mem=8000, runtime=1440), params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), prediction_dir=lambda wildcards: join( config["output_directory"], "predictions", wildcards.fold ), @@ -1111,6 +1121,7 @@ if ENABLE_STRUCTURE_ANALYSIS: config["analysis_container"], shell: r""" + {params.tmpdir_guard} alphajudge \ {params.prediction_dir} \ {params.cli_parameters} @@ -1152,6 +1163,7 @@ if ENABLE_STRUCTURE_ANALYSIS and GENERATE_RECURSIVE_REPORT: output: RECURSIVE_REPORT, params: + tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), prediction_dir=lambda wildcards: join(config["output_directory"], "predictions"), report_dir=join(config["output_directory"], "reports"), extra_args=lambda wildcards: " ".join( @@ -1164,6 +1176,7 @@ if ENABLE_STRUCTURE_ANALYSIS and GENERATE_RECURSIVE_REPORT: config["analysis_container"], shell: r""" + {params.tmpdir_guard} mkdir -p {params.report_dir} alphajudge \ {params.prediction_dir} \ diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 2889670..b888037 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -15,6 +15,7 @@ import json import lzma import os import re +import shlex import tempfile import time import urllib.error @@ -665,6 +666,29 @@ def _collect_roots(paths: Iterable[str | Path]) -> set[str]: return roots +def container_tmpdir_guard(fallback: str = "/tmp") -> str: + """Bash run inside a container before native tools use the inherited TMPDIR. + + Job-local scratch may not be mounted in the container. Keep a usable TMPDIR, + otherwise select a writable, searchable fallback or fail before starting work. + """ + quoted = shlex.quote(fallback) + return ( + 'if [ -z "${TMPDIR:-}" ] || ! { mkdir -p -- "$TMPDIR" 2>/dev/null ' + '&& [ -d "$TMPDIR" ] && [ -w "$TMPDIR" ] && [ -x "$TMPDIR" ]; }; then\n' + f' if ! {{ mkdir -p -- {quoted} 2>/dev/null && [ -d {quoted} ] ' + f'&& [ -w {quoted} ] && [ -x {quoted} ]; }}; then\n' + " printf '[container-tmpdir] fallback %s is not writable in the container\\n' " + f'{quoted} >&2\n' + ' exit 1\n' + ' fi\n' + " printf '[container-tmpdir] TMPDIR=%s is not writable in the container; using %s\\n' " + f'"${{TMPDIR:-}}" {quoted} >&2\n' + f' export TMPDIR={quoted}\n' + 'fi\n' + ) + + def prepare_container_binds( *, output_directory: str,