wip: qa7920 sidecar + rag stream timeout
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
sj
+282
-23
@@ -25,6 +25,7 @@ from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
import hashlib
|
||||
import os
|
||||
import shlex
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(PROJECT_ROOT) not in sys.path:
|
||||
@@ -231,8 +232,17 @@ RNASEQ_BIN = str(rnaseq_bin())
|
||||
prepend_path(RNASEQ_BIN)
|
||||
|
||||
|
||||
MAX_TOOL_THREADS = 8
|
||||
|
||||
|
||||
def resolve_thread_plan(cpu: int, sample_parallel_jobs: int = 1, threads_per_sample: Optional[int] = None) -> Dict[str, int]:
|
||||
"""Resolve conservative RNA-seq sample/thread settings from CLI args."""
|
||||
"""Resolve conservative RNA-seq sample/thread settings from CLI args.
|
||||
|
||||
Per-tool thread count is capped at MAX_TOOL_THREADS because fastp (and most
|
||||
QC tools) are single-threaded for their core algorithm; oversized --thread
|
||||
values only add gzip-thread contention without speeding up filtering. Real
|
||||
parallelism is achieved via chunked splitting (see run_fastp).
|
||||
"""
|
||||
total_cpu = max(1, int(cpu or 1))
|
||||
jobs = max(1, int(sample_parallel_jobs or 1))
|
||||
if threads_per_sample is None or int(threads_per_sample) < 1:
|
||||
@@ -241,6 +251,7 @@ def resolve_thread_plan(cpu: int, sample_parallel_jobs: int = 1, threads_per_sam
|
||||
threads = max(1, int(threads_per_sample))
|
||||
if jobs * threads > total_cpu:
|
||||
threads = max(1, total_cpu // jobs)
|
||||
threads = min(threads, MAX_TOOL_THREADS)
|
||||
return {"cpu": total_cpu, "sample_parallel_jobs": jobs, "threads_per_sample": threads}
|
||||
|
||||
|
||||
@@ -378,6 +389,253 @@ def run_cutadapt(
|
||||
raise
|
||||
|
||||
|
||||
def _build_fastp_cmd(
|
||||
fastp_bin: Path,
|
||||
r1: Path,
|
||||
r2: Path,
|
||||
out_r1: Path,
|
||||
out_r2: Path,
|
||||
html_report: Path,
|
||||
json_report: Path,
|
||||
cpu: int,
|
||||
detect_adapter_for_pe: bool,
|
||||
cut_front: bool,
|
||||
cut_tail: bool,
|
||||
cut_window_size: int,
|
||||
cut_mean_quality: int,
|
||||
qualified_quality_phred: int,
|
||||
unqualified_percent_limit: int,
|
||||
n_base_limit: int,
|
||||
length_required: int,
|
||||
) -> List[str]:
|
||||
"""Build a fastp command line with a capped --thread value."""
|
||||
cmd = [
|
||||
str(fastp_bin),
|
||||
"-i", str(r1),
|
||||
"-I", str(r2),
|
||||
"-o", str(out_r1),
|
||||
"-O", str(out_r2),
|
||||
"--html", str(html_report),
|
||||
"--json", str(json_report),
|
||||
"--thread", str(max(1, min(MAX_TOOL_THREADS, int(cpu or 1)))),
|
||||
"--cut_window_size", str(cut_window_size),
|
||||
"--cut_mean_quality", str(cut_mean_quality),
|
||||
"--qualified_quality_phred", str(qualified_quality_phred),
|
||||
"--unqualified_percent_limit", str(unqualified_percent_limit),
|
||||
"--n_base_limit", str(n_base_limit),
|
||||
"--length_required", str(length_required),
|
||||
]
|
||||
if detect_adapter_for_pe:
|
||||
cmd.append("--detect_adapter_for_pe")
|
||||
if cut_front:
|
||||
cmd.append("--cut_front")
|
||||
if cut_tail:
|
||||
cmd.append("--cut_tail")
|
||||
return cmd
|
||||
|
||||
|
||||
def _count_fastq_reads(fastq_gz: Path) -> int:
|
||||
"""Count the number of reads (4 lines each) in a gzipped FASTQ using pigz."""
|
||||
proc = subprocess.run(
|
||||
["bash", "-c", f"pigz -dc {shlex.quote(str(fastq_gz))} | wc -l"],
|
||||
capture_output=True, text=True, check=True,
|
||||
)
|
||||
return int((proc.stdout or "0").strip() or 0) // 4
|
||||
|
||||
|
||||
def _split_fastq_gz(fastq_gz: Path, work_dir: Path, tag: str, n_chunks: int) -> List[Path]:
|
||||
"""Split a gzipped FASTQ into n_chunks contiguous, order-preserving gzipped chunks.
|
||||
|
||||
Uses pigz for parallel (de)compression. Read boundaries (4 lines per read) are
|
||||
preserved so the downstream merge stays in order and R1/R2 stay paired.
|
||||
"""
|
||||
work_dir.mkdir(parents=True, exist_ok=True)
|
||||
total_reads = _count_fastq_reads(fastq_gz)
|
||||
if total_reads <= 0:
|
||||
raise RuntimeError(f"fastq is empty or unreadable: {fastq_gz}")
|
||||
reads_per_chunk = -(-total_reads // n_chunks) # ceil division
|
||||
lines_per_chunk = max(4, reads_per_chunk * 4)
|
||||
raw_prefix = work_dir / f"{tag}.raw"
|
||||
subprocess.run(
|
||||
["bash", "-c",
|
||||
f"pigz -dc {shlex.quote(str(fastq_gz))} | split -l {lines_per_chunk} -d -a 4 - {shlex.quote(str(raw_prefix))}"],
|
||||
check=True,
|
||||
)
|
||||
raw_chunks = sorted(work_dir.glob(f"{tag}.raw[0-9]*"))
|
||||
chunk_paths: List[Path] = []
|
||||
for i, raw in enumerate(raw_chunks):
|
||||
out = work_dir / f"{tag}.part{i:04d}.fastq.gz"
|
||||
with open(out, "wb") as fout:
|
||||
subprocess.run(["pigz", "-p", "4", "-c", str(raw)], stdout=fout, check=True)
|
||||
raw.unlink(missing_ok=True)
|
||||
chunk_paths.append(out)
|
||||
if not chunk_paths:
|
||||
raise RuntimeError(f"failed to split fastq: {fastq_gz}")
|
||||
return chunk_paths
|
||||
|
||||
|
||||
def _merge_fastp_json(chunk_jsons: List[Path]) -> Dict[str, Any]:
|
||||
"""Merge per-chunk fastp JSON summaries into one combined summary.
|
||||
|
||||
Read/base counts are summed; rates, GC content and mean lengths are recomputed
|
||||
as (base|read)-weighted averages so merged stats match a single full run.
|
||||
"""
|
||||
summary: Dict[str, Any] = {"before_filtering": {}, "after_filtering": {}}
|
||||
for section in ("before_filtering", "after_filtering"):
|
||||
total_reads = 0
|
||||
total_bases = 0
|
||||
q20_bases = 0
|
||||
q30_bases = 0
|
||||
gc_sum = 0.0
|
||||
r1_len_sum = 0.0
|
||||
r2_len_sum = 0.0
|
||||
for jf in chunk_jsons:
|
||||
with open(jf, encoding="utf-8") as f:
|
||||
s = json.load(f).get("summary", {}).get(section, {})
|
||||
tr = int(s.get("total_reads", 0) or 0)
|
||||
tb = int(s.get("total_bases", 0) or 0)
|
||||
total_reads += tr
|
||||
total_bases += tb
|
||||
q20_bases += int(s.get("q20_bases", 0) or 0)
|
||||
q30_bases += int(s.get("q30_bases", 0) or 0)
|
||||
gc_sum += float(s.get("gc_content", 0) or 0) * tb
|
||||
r1_len_sum += float(s.get("read1_mean_length", 0) or 0) * tr
|
||||
r2_len_sum += float(s.get("read2_mean_length", 0) or 0) * tr
|
||||
summary[section] = {
|
||||
"total_reads": total_reads,
|
||||
"total_bases": total_bases,
|
||||
"q20_bases": q20_bases,
|
||||
"q30_bases": q30_bases,
|
||||
"q20_rate": (q20_bases / total_bases) if total_bases else 0.0,
|
||||
"q30_rate": (q30_bases / total_bases) if total_bases else 0.0,
|
||||
"gc_content": (gc_sum / total_bases) if total_bases else 0.0,
|
||||
"read1_mean_length": (r1_len_sum / total_reads) if total_reads else 0.0,
|
||||
"read2_mean_length": (r2_len_sum / total_reads) if total_reads else 0.0,
|
||||
}
|
||||
return {"summary": summary}
|
||||
|
||||
|
||||
def _write_fastp_summary_html(merged: Dict[str, Any], html_report: Path) -> None:
|
||||
"""Write a lightweight self-contained HTML report from the merged summary."""
|
||||
before = merged["summary"]["before_filtering"]
|
||||
after = merged["summary"]["after_filtering"]
|
||||
rows = [
|
||||
("Total reads", before.get("total_reads", 0), after.get("total_reads", 0)),
|
||||
("Q20 rate", f"{before.get('q20_rate', 0):.4f}", f"{after.get('q20_rate', 0):.4f}"),
|
||||
("Q30 rate", f"{before.get('q30_rate', 0):.4f}", f"{after.get('q30_rate', 0):.4f}"),
|
||||
("GC content", f"{before.get('gc_content', 0):.4f}", f"{after.get('gc_content', 0):.4f}"),
|
||||
]
|
||||
trs = "".join(f"<tr><td>{name}</td><td>{b}</td><td>{a}</td></tr>" for name, b, a in rows)
|
||||
html = (
|
||||
"<!DOCTYPE html><html><head><meta charset='utf-8'>"
|
||||
"<title>fastp QC Report (chunked)</title></head><body>"
|
||||
"<h1>fastp QC Report</h1><p>Generated from chunked parallel fastp runs.</p>"
|
||||
"<table border='1' cellpadding='6' cellspacing='0'>"
|
||||
"<tr><th>Metric</th><th>Before filtering</th><th>After filtering</th></tr>"
|
||||
f"{trs}</table></body></html>"
|
||||
)
|
||||
html_report.write_text(html, encoding="utf-8")
|
||||
|
||||
|
||||
def _run_fastp_chunked(
|
||||
fastp_bin: Path,
|
||||
r1: Path,
|
||||
r2: Path,
|
||||
out_r1: Path,
|
||||
out_r2: Path,
|
||||
html_report: Path,
|
||||
json_report: Path,
|
||||
log_file: Path,
|
||||
cpu: int,
|
||||
chunks: int,
|
||||
chunk_work_dir: Path,
|
||||
detect_adapter_for_pe: bool,
|
||||
cut_front: bool,
|
||||
cut_tail: bool,
|
||||
cut_window_size: int,
|
||||
cut_mean_quality: int,
|
||||
qualified_quality_phred: int,
|
||||
unqualified_percent_limit: int,
|
||||
n_base_limit: int,
|
||||
length_required: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run fastp in chunked-parallel mode: split FASTQ, parallel fastp, merge outputs."""
|
||||
chunk_work_dir.mkdir(parents=True, exist_ok=True)
|
||||
print(f" [fastp] Splitting into {chunks} chunks (parallel mode)...", flush=True)
|
||||
start_ts = time.time()
|
||||
|
||||
r1_chunks = _split_fastq_gz(r1, chunk_work_dir, "r1", chunks)
|
||||
r2_chunks = _split_fastq_gz(r2, chunk_work_dir, "r2", chunks)
|
||||
n = max(len(r1_chunks), len(r2_chunks))
|
||||
if len(r1_chunks) != len(r2_chunks):
|
||||
print(
|
||||
f" [fastp] WARNING: R1/R2 chunk counts differ "
|
||||
f"({len(r1_chunks)} vs {len(r2_chunks)}); using {n}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
tasks = []
|
||||
for i in range(n):
|
||||
tasks.append({
|
||||
"idx": i,
|
||||
"r1": r1_chunks[i],
|
||||
"r2": r2_chunks[i],
|
||||
"o1": chunk_work_dir / f"out{i:04d}.r1.fastq.gz",
|
||||
"o2": chunk_work_dir / f"out{i:04d}.r2.fastq.gz",
|
||||
"html": chunk_work_dir / f"out{i:04d}.fastp.html",
|
||||
"json": chunk_work_dir / f"out{i:04d}.fastp.json",
|
||||
})
|
||||
|
||||
def _run_one(t: Dict[str, Any]):
|
||||
cmd = _build_fastp_cmd(
|
||||
fastp_bin, t["r1"], t["r2"], t["o1"], t["o2"], t["html"], t["json"], cpu,
|
||||
detect_adapter_for_pe, cut_front, cut_tail, cut_window_size, cut_mean_quality,
|
||||
qualified_quality_phred, unqualified_percent_limit, n_base_limit, length_required,
|
||||
)
|
||||
res = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
||||
if res.returncode != 0:
|
||||
return t["idx"], None, res.stderr
|
||||
for p in (t["o1"], t["o2"], t["json"]):
|
||||
if not p.exists() or p.stat().st_size == 0:
|
||||
return t["idx"], None, f"empty output: {p}"
|
||||
return t["idx"], t, None
|
||||
|
||||
results: Dict[int, Dict[str, Any]] = {}
|
||||
errors: List[str] = []
|
||||
with ThreadPoolExecutor(max_workers=max(1, int(chunks))) as ex:
|
||||
futures = [ex.submit(_run_one, t) for t in tasks]
|
||||
for fut in as_completed(futures):
|
||||
idx, t, err = fut.result()
|
||||
if err:
|
||||
errors.append(f"chunk {idx}: {err}")
|
||||
else:
|
||||
results[idx] = t
|
||||
if errors:
|
||||
raise RuntimeError(f"fastp chunked failed: {'; '.join(errors[:3])}")
|
||||
|
||||
# Merge clean reads in chunk order (preserves read order/pairing)
|
||||
for out_path, key in ((out_r1, "o1"), (out_r2, "o2")):
|
||||
with open(out_path, "wb") as fout:
|
||||
for i in range(n):
|
||||
with open(results[i][key], "rb") as fin:
|
||||
shutil.copyfileobj(fin, fout)
|
||||
|
||||
# Merge JSON + generate HTML summary
|
||||
merged = _merge_fastp_json([results[i]["json"] for i in range(n)])
|
||||
json_report.write_text(json.dumps(merged), encoding="utf-8")
|
||||
_write_fastp_summary_html(merged, html_report)
|
||||
|
||||
elapsed = time.time() - start_ts
|
||||
log_file.write_text(
|
||||
f"fastp chunked parallel mode: {n} chunks x --thread {max(1, min(MAX_TOOL_THREADS, int(cpu or 1)))}\n"
|
||||
f"elapsed: {elapsed:.1f}s\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
shutil.rmtree(chunk_work_dir, ignore_errors=True)
|
||||
print(f" [fastp] Completed (chunked x{n}) in {elapsed:.1f}s", flush=True)
|
||||
return {"success": True, "time_sec": elapsed}
|
||||
|
||||
|
||||
def run_fastp(
|
||||
fastp_bin: Path,
|
||||
r1: Path,
|
||||
@@ -397,30 +655,22 @@ def run_fastp(
|
||||
unqualified_percent_limit: int = 30,
|
||||
n_base_limit: int = 5,
|
||||
length_required: int = 30,
|
||||
chunks: int = 1,
|
||||
chunk_work_dir: Optional[Path] = None,
|
||||
) -> Dict[str, Any]:
|
||||
cmd = [
|
||||
str(fastp_bin),
|
||||
"-i", str(r1),
|
||||
"-I", str(r2),
|
||||
"-o", str(out_r1),
|
||||
"-O", str(out_r2),
|
||||
"--html", str(html_report),
|
||||
"--json", str(json_report),
|
||||
"--thread", str(max(1, int(cpu or 1))),
|
||||
"--cut_window_size", str(cut_window_size),
|
||||
"--cut_mean_quality", str(cut_mean_quality),
|
||||
"--qualified_quality_phred", str(qualified_quality_phred),
|
||||
"--unqualified_percent_limit", str(unqualified_percent_limit),
|
||||
"--n_base_limit", str(n_base_limit),
|
||||
"--length_required", str(length_required),
|
||||
]
|
||||
if detect_adapter_for_pe:
|
||||
cmd.append("--detect_adapter_for_pe")
|
||||
if cut_front:
|
||||
cmd.append("--cut_front")
|
||||
if cut_tail:
|
||||
cmd.append("--cut_tail")
|
||||
if chunks and int(chunks) > 1 and chunk_work_dir is not None:
|
||||
return _run_fastp_chunked(
|
||||
fastp_bin, r1, r2, out_r1, out_r2, html_report, json_report, log_file,
|
||||
cpu, int(chunks), chunk_work_dir,
|
||||
detect_adapter_for_pe, cut_front, cut_tail, cut_window_size, cut_mean_quality,
|
||||
qualified_quality_phred, unqualified_percent_limit, n_base_limit, length_required,
|
||||
)
|
||||
|
||||
cmd = _build_fastp_cmd(
|
||||
fastp_bin, r1, r2, out_r1, out_r2, html_report, json_report, cpu,
|
||||
detect_adapter_for_pe, cut_front, cut_tail, cut_window_size, cut_mean_quality,
|
||||
qualified_quality_phred, unqualified_percent_limit, n_base_limit, length_required,
|
||||
)
|
||||
print(f" [fastp] Starting: {r1.name} + {r2.name}...", flush=True)
|
||||
start_ts = time.time()
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
||||
@@ -767,6 +1017,7 @@ def main():
|
||||
parser.add_argument("-t", "--cpu", type=int, default=4, help="Total CPU threads for this stage")
|
||||
parser.add_argument("--sample-parallel-jobs", type=int, default=1, help="Number of samples to process concurrently")
|
||||
parser.add_argument("--threads-per-sample", type=int, help="Threads used by each per-sample tool invocation")
|
||||
parser.add_argument("--chunks-per-sample", type=int, default=0, help="Split each sample into N chunks for parallel fastp. 0 = auto (cpu // threads-per-sample).")
|
||||
parser.add_argument("--resume-existing", action="store_true", help="Skip samples whose FastQC, clean reads, and cutadapt outputs already exist in the output directory")
|
||||
parser.add_argument("-p", "--prefix", default="", help="Output file prefix")
|
||||
parser.add_argument("--adapter-r1", default="AGATCGGAAGAGC", help="Adapter sequence for read1")
|
||||
@@ -786,6 +1037,12 @@ def main():
|
||||
args = parser.parse_args()
|
||||
thread_plan = resolve_thread_plan(args.cpu, args.sample_parallel_jobs, args.threads_per_sample)
|
||||
threads_per_sample = thread_plan["threads_per_sample"]
|
||||
|
||||
# Chunked parallelism: split each sample into N chunks and run fastp in
|
||||
# parallel (the real way to saturate many cores, since fastp is single-threaded).
|
||||
chunks_per_sample = int(args.chunks_per_sample or 0)
|
||||
if chunks_per_sample < 1:
|
||||
chunks_per_sample = max(1, int(thread_plan["cpu"]) // max(1, int(threads_per_sample)))
|
||||
|
||||
start_time = datetime.now()
|
||||
input_dir = Path(args.input)
|
||||
@@ -878,6 +1135,8 @@ def main():
|
||||
json_report=analysis_dir / f"{sample_name}_fastp.json",
|
||||
log_file=analysis_dir / f"{sample_name}_fastp.log",
|
||||
cpu=threads_per_sample,
|
||||
chunks=chunks_per_sample,
|
||||
chunk_work_dir=analysis_dir / f".{sample_name}.chunks",
|
||||
**fastp_params,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
|
||||
Reference in New Issue
Block a user