Size STARsolo memory to index + sort buffer; make it configurable

AGENTS.md

@@ -337,7 +337,8 @@ build the strings with `str.format()` for consistency (Snakemake 7's tokenizer r
 f-string contents; fixed in 8+), and emit each message once via `_diag`
 (params lambdas run once per job). `_get_config` stays silent for keys that are
 legitimately optional (`_OPTIONAL_KEYS` / `_OPTIONAL_SUFFIXES`: `cutadapt_*`, `bc_cut`,
-`bc_length1`, `*_extra_args`, `*_polyA_filter`, `*_priming_offset`, `bam_*`); any other
+`bc_length1`, `*_extra_args`, `*_polyA_filter`, `*_priming_offset`, `bam_*`,
+`STAR_mem_mb`); any other
 missing key is a real misconfiguration and is reported once.
 
 ---
@@ -355,7 +356,8 @@ missing key is a real misconfiguration and is reported once.
 - `DISCOVERY`: Optional de novo RT priming site discovery block (see `docs/discovery.md`)
 - `DEFAULTS`: Lowest-priority values for any per-sample key: `platform`, `chemistry`,
   `alignments`, and the mode switches `R2_polyA_filter`, `R2_priming_offset`,
-  `bam_mate`, `bam_strand`, `bam_polyA_filter`, `bam_priming_offset`
+  `bam_mate`, `bam_strand`, `bam_polyA_filter`, `bam_priming_offset`, and the optional
+  `STAR_mem_mb`
 - `SAMPLES`: Per-sample configuration (`basename` or `bam`, `platform`, `chemistry`,
   `alignments`, `discovery_read`, the `bam_*`/`R2_*` switches, and any
   chemistry/platform key such as `STAR_R2` or `STAR_R2_extra_args` as an override)
@@ -562,7 +564,13 @@ above).
 - Missing conda dependencies → check `scraps_conda.yml`
 - YAML syntax errors → validate with `snakemake -np`
 - Missing input files → check `DATA` path in config.yaml
-- Resource exhaustion → adjust `mem_mb` or `threads` in rules
+- Resource exhaustion → adjust `mem_mb` or `threads` in rules. STARsolo is the
+  exception: its `mem_mb` comes from `_star_mem_mb` (Snakefile). That is `STAR_mem_mb`
+  from the usual lookup, or by default the index size on disk (`STAR_INDEX_MB`) + 16 GB
+  sort buffer + 4 GB headroom, times the attempt number. `--limitBAMsortRAM` is derived
+  from the job's final `mem_mb` (`_star_sort_ram`), so never hardcode it: STAR holds the
+  index in RAM in addition to the sort buffer, and a hardcoded buffer equal to `mem_mb`
+  is what got starsolo jobs OOM-killed. The SLURM profile sets `retries: 1`.
 
 **Debugging Snakemake**:
 ```bash

README.md

@@ -73,6 +73,8 @@ See [profiles/slurm/README.md](profiles/slurm/README.md) and the [SLURM executor
 
 Notes: each STAR alignment job (one per sample × alignment mode) uses `total_impact=5`; `--resources total_impact=5*N` allows N concurrent STAR jobs.
 
+STAR memory: each STARsolo job requests the genome index size (read from `STAR_INDEX`) plus a 16 GB BAM-sort buffer and 4 GB headroom, and `--limitBAMsortRAM` is derived from that request, so sorting stays inside the allocation. For large libraries set `STAR_mem_mb` (MB) in `DEFAULTS`, a platform block, or a sample block, or override `mem_mb` with `set-resources`/`--set-resources`; the sort buffer follows either way. The SLURM profile retries a failed job once, and STAR jobs get twice the memory on the retry.
+
 ---
 
 ## Configuration

Snakefile

@@ -18,6 +18,7 @@ READS = ["R1", "R2", "paired"]
 GENOME_FASTA = config.get("GENOME_FASTA")
 DISCOVERY = config.get("DISCOVERY", {}) or {}
 
+import os
 import sys
 import yaml
 from snakemake.logging import logger
@@ -106,7 +107,7 @@ _validate_modes()
 
 # keys that are legitimately absent (rules treat "" as "off"); no diagnostic
 _OPTIONAL_KEYS = {"cutadapt_R1", "cutadapt_paired", "bc_cut", "bc_length1",
-                  "bam_strand", "bam_mate"}
+                  "bam_strand", "bam_mate", "STAR_mem_mb"}
 _OPTIONAL_SUFFIXES = ("_extra_args", "_polyA_filter", "_priming_offset")
 
 def _get_config(sample, item):
@@ -144,6 +145,42 @@ def _get_config(sample, item):
           "string.".format(item, sample, chemistry, platform))
   return ""
 
+# --- STARsolo memory ----------------------------------------------------------
+# STAR holds the whole genome index in RAM and sorts the BAM in a buffer on top
+# of it, so a job needs index + sort buffer + headroom. STAR_INDEX_MB is the
+# index size on disk (32000, a human index, when STAR_INDEX is absent, e.g. a
+# CI dry-run). STAR_mem_mb (sample/platform/chemistry/DEFAULTS) sets the job
+# memory; unset, it is index + 16 GB sort buffer + headroom. The starsolo rules
+# multiply it by the attempt number (profile `retries`) and derive
+# --limitBAMsortRAM from the job's final mem_mb (_star_sort_ram), so a
+# set-resources override keeps the two in sync.
+_STAR_HEADROOM_MB = 4000
+
+def _dir_size_mb(path):
+  try:
+    files = [os.path.join(path, f) for f in os.listdir(path)]
+  except OSError:
+    return None
+  return sum(os.path.getsize(f) for f in files if os.path.isfile(f)) // 1000000
+
+STAR_INDEX_MB = _dir_size_mb(STAR_INDEX) or 32000
+
+def _star_mem_mb(wildcards, attempt):
+  mem = _get_config(wildcards.sample, "STAR_mem_mb")
+  if mem in ("", None):
+    mem = STAR_INDEX_MB + 16000 + _STAR_HEADROOM_MB
+  return int(mem) * attempt
+
+def _star_sort_ram(wildcards, resources):
+  """--limitBAMsortRAM in bytes: job memory minus index and headroom."""
+  sort_mb = resources.mem_mb - STAR_INDEX_MB - _STAR_HEADROOM_MB
+  if sort_mb < 2000:
+    _diag(("star_mem", wildcards.sample),
+          "Warning: sample '{}': STARsolo mem_mb {} leaves less than 2 GB for "
+          "BAM sorting after the {} MB index; raise STAR_mem_mb.".format(
+            wildcards.sample, resources.mem_mb, STAR_INDEX_MB))
+  return max(sort_mb, 2000) * 1000000
+
 # --- per-mode behaviour switches ----------------------------------------------
 # Per-mode keys ({mode} = R2 or bam), settable per sample / platform / chemistry
 # with defaults in config.yaml DEFAULTS:

config.yaml

@@ -159,6 +159,10 @@ DEFAULTS:
   # proper pairs only, the other mate is dropped, positions are per fragment).
   # bam_priming_offset: bp the priming site lies beyond the counted read end.
   # ASSUMED, not enforced: full-length molecules, CB/UB tags, accurate strand.
+  # STAR_mem_mb: STARsolo job memory (MB). Unset = genome index size + 16 GB
+  # BAM-sort buffer + 4 GB headroom; --limitBAMsortRAM is derived from it.
+  # Set here, per sample, or per platform for large libraries, e.g.
+  # STAR_mem_mb: 96000
   bam_mate: 0
   bam_strand: same
   bam_polyA_filter: false

profiles/slurm/README.md

@@ -19,6 +19,15 @@ Edit `config.v8+.yaml`:
 Each rule's `threads` becomes `--cpus-per-task` and its `mem_mb` becomes
 `--mem`, so no per-rule CPU/memory settings are needed here.
 
+STARsolo memory is the genome index size (read from `STAR_INDEX`) plus a 16 GB
+BAM-sort buffer and 4 GB headroom, and `--limitBAMsortRAM` is derived from the
+job's `mem_mb`, so STAR never sorts past its allocation. For large libraries
+raise it with `STAR_mem_mb` in `config.yaml` (DEFAULTS, a platform, or one
+sample) or with `mem_mb` under `set-resources` here. `retries: 1` resubmits a
+failed job once, and STAR jobs get twice the memory on the retry. On clusters
+that cap memory per core (Alpine `acpu`: about 3.8 GB/core) a large `mem_mb`
+needs a matching partition (e.g. `amem`) or more CPUs.
+
 ## Run
 
 ```bash

profiles/slurm/config.v8+.yaml

@@ -9,6 +9,7 @@ latency-wait: 60             # seconds to wait for outputs on shared filesystems
 keep-going: true
 printshellcmds: true
 rerun-incomplete: true
+retries: 1                   # resubmit a failed job once; STAR jobs get 2x mem_mb
 
 default-resources:
   - slurm_account="CHANGE_ME"
@@ -25,6 +26,7 @@ resources:
 set-resources:
   starsolo_R1:
     runtime: 1440
+    # mem_mb: 96000          # overrides STAR_mem_mb; --limitBAMsortRAM follows
   starsolo_R2:
     runtime: 1440
   starsolo_paired:

rules/cutadapt_star.snake

@@ -158,12 +158,13 @@ rule starsolo_R1:
     params:
       chemistry = _get_chem_version_R1,
       out_dir = "{results}/{sample}/{sample}_R1_",
-      job_name = "star_R1"
+      job_name = "star_R1",
+      sort_ram = _star_sort_ram
     log: "{results}/logs/starsolo_{sample}_R1.txt"
     benchmark: "{results}/benchmarks/starsolo_{sample}_R1.tsv"
     resources:
       total_impact = 5,
-      mem_mb = 48000
+      mem_mb = _star_mem_mb
     threads: 14
     shell:
       """
@@ -173,7 +174,7 @@ rule starsolo_R1:
       --soloUMIfiltering MultiGeneUMI \
       --soloUMIdedup 1MM_Directional \
       --readFilesCommand gunzip -c \
-      --runThreadN 12 \
+      --runThreadN {threads} \
       --soloCBwhitelist {params.chemistry} \
       --soloBarcodeMate 1 \
       --soloStrand Reverse \
@@ -184,7 +185,7 @@ rule starsolo_R1:
       --soloCellFilter None \
       --outFileNamePrefix {params.out_dir} \
       --outSAMtype BAM SortedByCoordinate \
-      --limitBAMsortRAM 48000000000 \
+      --limitBAMsortRAM {params.sort_ram} \
       --outSAMattributes NH HI nM AS CR UR CB UB GX GN sS sQ sM \
       --readFilesIn {input.R1}
 
@@ -200,12 +201,13 @@ rule starsolo_R2:
     params:
       chemistry = _get_chem_version_R2,
       out_dir = "{results}/{sample}/{sample}_R2_", 
-      job_name = "star_R2"
+      job_name = "star_R2",
+      sort_ram = _star_sort_ram
     log: "{results}/logs/starsolo_{sample}_R2.txt"
     benchmark: "{results}/benchmarks/starsolo_{sample}_R2.tsv"
     resources:
       total_impact = 5,
-      mem_mb = 48000
+      mem_mb = _star_mem_mb
     threads: 14
     shell:
       """
@@ -215,7 +217,7 @@ rule starsolo_R2:
       --soloUMIfiltering MultiGeneUMI \
       --soloUMIdedup 1MM_Directional \
       --readFilesCommand gunzip -c \
-      --runThreadN 12 \
+      --runThreadN {threads} \
       --outFilterMultimapNmax 1 \
       --soloBarcodeReadLength 0 \
       --soloCBwhitelist {params.chemistry} \
@@ -225,7 +227,7 @@ rule starsolo_R2:
       --soloCellFilter None \
       --outFileNamePrefix {params.out_dir} \
       --outSAMtype BAM SortedByCoordinate \
-      --limitBAMsortRAM 48000000000 \
+      --limitBAMsortRAM {params.sort_ram} \
       --outSAMattributes NH HI nM AS CR UR CB UB GX GN sS sQ sM \
       --readFilesIn $(echo {input.R1} | sed 's/ /,/g;s/R1/R2/g') $(echo {input.R1} | sed 's/ /,/g')
 
@@ -242,12 +244,13 @@ rule starsolo_paired:
     params:
       chemistry = _get_chem_version_paired,
       out_dir = "{results}/{sample}/{sample}_paired_", 
-      job_name = "star_paired"
+      job_name = "star_paired",
+      sort_ram = _star_sort_ram
     log: "{results}/logs/starsolo_{sample}_paired.txt"
     benchmark: "{results}/benchmarks/starsolo_{sample}_paired.tsv"
     resources:
       total_impact = 5,
-      mem_mb = 48000
+      mem_mb = _star_mem_mb
     threads: 14
     shell:
       """
@@ -257,7 +260,7 @@ rule starsolo_paired:
       --soloUMIfiltering MultiGeneUMI \
       --soloUMIdedup 1MM_Directional \
       --readFilesCommand gunzip -c \
-      --runThreadN 12 \
+      --runThreadN {threads} \
       --soloCBwhitelist {params.chemistry} \
       --soloBarcodeMate 1 \
       --soloStrand Reverse \
@@ -268,7 +271,7 @@ rule starsolo_paired:
       --soloCellFilter None \
       --outFileNamePrefix {params.out_dir} \
       --outSAMtype BAM SortedByCoordinate \
-      --limitBAMsortRAM 48000000000 \
+      --limitBAMsortRAM {params.sort_ram} \
       --outSAMattributes NH HI nM AS CR UR CB UB GX GN sS sQ sM \
       --readFilesIn {input.R1} {input.R2}