Repository: pytorch
Issue: Segmentation fault in PyTorch dataloader
Labels: needs reproduction,module: dataloader,triaged

## 🐛 Bug

I encountered a segmentation fault error when running the program [Vamb ](https://github.com/RasmussenLab/vamb) and the developer suggested it was related to the C or C++ implementation of pytorch's dataloader.

This error arose when running the Snakemake workflow for Vamb, specifically in the rule that runs the Vamb program itself. I have attached the error log file here along with the Snakefile and the original error message. Full documentation for Vamb can be found on the GitHub linked above.

Many thanks in advance!

## To Reproduce

Steps to reproduce the behavior:

Follow instructions on the vamb/workflow GitHub to set up input files for Vamb. Then run

snakemake --cores 20 --configfile config.json --snakefile vamb.snake.conda.py

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

Snakemake error:
[Thu Mar 25 10:45:02 2021]
rule vamb:
input: jgi_matrix/jgi.abundance.dat, contigs.flt.fna.gz
output: vamb/clusters.tsv, vamb/latent.npz, vamb/lengths.npz, vamb/log.txt, vamb/model.pt, vamb/mask.npz, vamb/tnf.npz
log: log/vamb/vamb.log
jobid: 250
threads: 10

[Thu Mar 25 10:49:09 2021]
Error in rule vamb:
jobid: 250
output: vamb/clusters.tsv, vamb/latent.npz, vamb/lengths.npz, vamb/log.txt, vamb/model.pt, vamb/mask.npz, vamb/tnf.npz
log: log/vamb/vamb.log (check log file(s) for error message)
shell:
rm -rf vamb;vamb --outdir vamb --fasta contigs.flt.fna.gz --jgi jgi_matrix/jgi.abundance.dat -o C -m 2000 --minfasta 500000 --outdir vamb 2>log/vamb/vamb.log
(one of the commands exited with non-zero exit code; note that snakemake uses bash strict mode!)

Snakefile:

#!/usr/bin/env python

# set configurations
INDEX_SIZE = config.get("index_size", "12G")
MM_MEM = config.get("minimap_mem", "35gb")
MM_PPN = config.get("minimap_ppn", "10")
VAMB_MEM = config.get("vamb_mem", "20gb")
VAMB_PPN = config.get("vamb_ppn", "10")
CHECKM_MEM = config.get("checkm_mem", "25gb")
CHECKM_PPN = config.get("checkm_ppn", "10")
SAMPLE_DATA = config.get("sample_data", "samples2data.txt")
CONTIGS = config.get("contigs", "contigs.txt")
VAMB_PARAMS = config.get("vamb_params", "-o C -m 2000 --minfasta 500000")
VAMB_PRELOAD = config.get("vamb_preload", "")

# parse if GPUs is needed #
VAMB_split = VAMB_PPN.split(":") 
VAMB_threads = VAMB_split[0]

## read in sample information ##

# read in sample2path
IDS = []
sample2path = {}
fh_in = open(SAMPLE_DATA, 'r')
for line in fh_in:
    line = line.rstrip()
    fields = line.split('\t')
    IDS.append(fields[0])
    sample2path[fields[0]] = [fields[1], fields[2]]

# read in list of per-sample assemblies
contigs_list = []
fh_in = open(CONTI
GS, 'r')
for line in fh_in:
    line = line.rstrip()
    contigs_list.append(line)


## start of snakemake rules ##

# targets
rule all:
    input:
        "jgi_matrix/jgi.abundance.dat",
        "vamb/clusters.tsv",
        "vamb/checkm.results"

rule cat_contigs:
    input:
        contigs_list
    output:
        "contigs.flt.fna.gz"
    params:
        walltime="864000", nodes="1", ppn="1", mem="10gb"
    threads:
        int(1)
    log:
        "log/contigs/catcontigs.log"
    conda:
        "envs/vamb.yaml"
    shell:
        "concatenate.py {output} {input} -m 2000"

rule index:
    input:
        contigs = "contigs.flt.fna.gz"
    output:
        mmi = "contigs.flt.mmi"
    params:
        walltime="864000", nodes="1", ppn="1", mem="90gb"
    threads:
        int(1)
    log:
        "log/contigs/index.log"
    conda: 
        "envs/minimap2.yaml"
    shell:
        "minimap2 -I {INDEX_SIZE} -d {output} {input} 2> {log}"

rule dict:
    input:
        contigs = "contigs.flt.fna.gz"
    output:
        dict = "contigs.flt.dict"
    params:
        walltime="864000", nodes="1", ppn="1", mem="10gb"
    threads:
        int(1)
    log:
        "log/contigs/dict.log"
    conda:
        "envs/samtools.yaml"
    shell:
        "samtools dict {input} | cut -f1-3 > {output} 2> {log}"

rule minimap:
    input:
        fq = lambda wildcards: sample2path[wildcards.sample],
        mmi = "contigs.flt.mmi",
        dict = "contigs.flt.dict"
    output:
        bam = temp("mapped/{sample}.bam")
    params:
        walltime="864000", nodes="1", ppn=MM_PPN, mem=MM_MEM
    threads:
        int(MM_PPN)
    log:
        "log/map/{sample}.minimap.log"
    conda:
        "envs/minimap2.yaml"
    shell:
        '''minimap2 -t {threads} -ax sr {input.mmi} {input.fq} | grep -v "^@" | cat {input.dict} - | samtools view -F 3584 -b - > {output.bam} 2>{log}'''

rule sort:
    input:
        "mapped/{sample}.bam"
    output:
        temp("mapped/{sample}.sort.bam")
    params:
        walltime="864000", nodes="1", ppn="2", mem="15gb",
        prefix="mapped/tmp.{sample}"
    threads:
        int(2)
    log:
        "log/map/{sample}.sort.log"
    conda:
        "envs/samtools.yaml"
    shell:
        "samtools sort {input} -T {params.prefix} --threads 1 -m 3G -o {output} 2>{log}"

rule jgi:
    input:
        bam = "mapped/{sample}.sort.bam"
    output:
        jgi = temp("jgi/{sample}.raw.jgi")
    params:
        walltime="864000", nodes="1", ppn="1", mem="10gb"
    threads:
        int(1)
    log:
        "log/jgi/{sample}.jgi"
    conda:
        "envs/metabat2.yaml"
    shell:
        "jgi_summarize_bam_contig_depths --noIntraDepthVariance --outputDepth {output} {input} 2>{log}"

rule cut_column1to3: 
    input:
        "jgi/%s.raw.jgi" % IDS[0] 
    output:
        "jgi/jgi.column1to3"
    params:
        walltime="86400", nodes="1", ppn="1", mem="1gb"
    log:
        "log/jgi/column1to3"
    shell: 
        "cut -f1-3 {input} > {output} 2>{log}"

rule cut_column4to5:
    input:
        "jgi/{sample}.raw.jgi"
    output:
        "jgi/{sample}.cut.jgi"
    params:
        walltime="86400", nodes="1", ppn="1", mem="1gb"
    log:
        "log/jgi/{sample}.cut.log"
    shell: 
        "cut -f1-3 --complement {input} > {output} 2>{log}"

rule paste_abundances:
    input:
        column1to3="jgi/jgi.column1to3",
        data=expand("jgi/{sample}.cut.jgi", sample=IDS)
    output:
        "jgi_matrix/jgi.abundance.dat" 
    params:
        walltime="86400", nodes="1", ppn="1", mem="1gb"
    log:
        "log/jgi/paste_abundances.log"
    shell: 
        "paste {input.column1to3} {input.data} > {output} 2>{log}" 

rule vamb:
    input:
        jgi = "jgi_matrix/jgi.abundance.dat",
        contigs = "contigs.flt.fna.gz"
    output:
        "vamb/clusters.tsv",
        "vamb/latent.npz",
        "vamb/lengths.npz",
        "vamb/log.txt",
        "vamb/model.pt",
        "vamb/mask.npz",
        "vamb/tnf.npz"
    params:
        walltime="86400", nodes="1", ppn=VAMB_PPN, mem=VAMB_MEM
    log:
        "log/vamb/vamb.log"
    threads:
        int(VAMB_threads)
    conda:
        "envs/vamb.yaml"
    shell:
        "{VAMB_PRELOAD}"
        "rm -rf vamb;"
        "vamb --outdir vamb --fasta {input.contigs} --jgi {input.jgi} {VAMB_PARAMS} 2>{log}"

rule checkm:
    input:
        "vamb/clusters.tsv"
    output:
        "vamb/checkm.results"
    params:
        walltime="86400", nodes="1", ppn=CHECKM_PPN, mem=CHECKM_MEM,
        bins = "vamb/bins",
        outdir = "vamb/checkm.outdir"
    log:
        "log/vamb/checkm.log"
    threads:
        int(CHECKM_PPN)
    conda:
        "envs/checkm.yaml"
    shell:
        "checkm lineage_wf -f {output} -t {threads} -x fna {params.bins} {params.outdir} 2>{log}"

Error log:

ERROR: Unexpected segmentation fault encountered in worker.
Traceback (most recent call last):
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 986, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/multiprocessing/queues.py", line 104, in get
    if not self._poll(timeout):
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/multiprocessing/connection.py", line 257, in poll
    return self._poll(timeout)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/multiprocessing/connection.py", line 414, in _poll
    r = wait([self], timeout)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/multiprocessing/connection.py", line 921, in wait
    ready = selector.select(timeout)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/selectors.py", line 415, in select
    fd_event_list = self._selector.poll(timeout)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/_utils/signal_handling.py", line 66, in handler
    _error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 25630) is killed by signal: Segmentation fault. 

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/bin/vamb", line 8, in <module>
    sys.exit(main())
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/vamb/__main__.py", line 528, in main
    logfile=logfile)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/vamb/__main__.py", line 251, in run
    dropout, cuda, batchsize, nepochs, lrate, batchsteps, logfile)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/vamb/__main__.py", line 150, in trainvae
    logfile=logfile, modelfile=modelpath)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/vamb/encode.py", line 492, in trainmodel
    dataloader = self.trainepoch(dataloader, epoch, optimizer, batchsteps_set, logfile)
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/vamb/encode.py", line 299, in trainepoch
    for depths_in, tnf_in in data_loader:
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 517, in __next__
    data = self._next_data()
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1182, in _next_data
    idx, data = self._get_data()
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1148, in _get_data
    success, data = self._try_get_data()
  File "/home/cns.local/nastassia.patin/.conda/envs/vamb/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 999, in _try_get_data
    raise RuntimeError('DataLoader worker (pid(s) {}) exited unexpectedly'.format(pids_str)) from e
RuntimeError: DataLoader worker (pid(s) 25630) exited unexpectedly

## Expected behavior

Snakemake workflow completes without errors and Vamb successfully trains autoencoder and clusters contigs into bins.

## Environment

PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.7 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.1
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              h8f6ccaa_8    conda-forge
[conda] mkl                       2020.2                      256  
[conda] numpy                     1.20.1           py37haa41c4c_0    conda-forge
[conda] pytorch                   1.8.0           py3.7_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchvision               0.9.0                    pypi_0    pypi

## Additional context

snakemake version 6.0.5


cc @SsnL @VitalyFedyunin @ejguan