如何优雅设置DeepVariant和PEPPER使用的线程数？

DeepVariant和PEPPER介绍

DeepVariant是基于深度学习模型来检测遗传变异（SNP和INDEL）的一款软件，PEPPER是在DeepVariant基础上进行了改进优化，二者都支持三代测序数据（HiFi和ONT测序）。

DeepVariant官网：https://github.com/google/deepvariant

PEPPER官网：https://github.com/kishwarshafin/pepper

由于PEPPER是基于DeepVariant，DeepVariant存在的问题也存在于PEPPER中，因此先介绍一下DeepVariant的基本工作流程，再介绍二者在线程使用上存在的问题。

DeepVariant有以下三个步骤：

1. make_examples: 将测序数据（经过比对的bam文件）转化为下一步的输入文件，在这里存储为TFRecord文件，这是TensorFlow的一种数据存储格式。这一步借用parallel软件实现了多进程，即每个进程仅处理一小批的数据，DeepVariant所指定的线程数就是这一步所用的进程数。
2. call_variants: 这一步是调用已经训练好的深度学习模型来进行推理（也叫预测，或者深度学习中的前向传播），这是一个已经训练好的准确率很高的模型，存放于DeepVariant和PEPPER的镜像中。这一步会接收上一步制作好的TFRecord文件，将其传入模型得到预测结果即SNP和INDEL。这一步不涉及到线程的概念，但是由于它需要调用TensorFlow进行前向传播，而TensorFlow在计算时会默认使用CPU所有的线程（核心），所以这一步其实是多线程计算。这样的设计存在问题，后面会解释。
3. postprocess_variants: 对第二步得到的variants进行处理，排序并合并多等位基因的结果，最后输出VCF文件。该步骤是单线程执行，不需要设置线程数。

DeepVariant在其HiFi-case-study中给的命令如下：

ulimit -u 10000 # https://stackoverflow.com/questions/52026652/openblas-blas-thread-init-pthread-create-resource-temporarily-unavailable/54746150#54746150
BIN_VERSION="1.5.0"
mkdir -p deepvariant_output

singularity exec --bind /usr/lib/locale/ \
  docker://google/deepvariant:${BIN_VERSION} \
    /opt/deepvariant/bin/run_deepvariant \
    --model_type PACBIO \
    --ref reference/GRCh38_no_alt_analysis_set.fasta \
    --reads input/HG003.GRCh38.chr20.pFDA_truthv2.bam \
    --output_vcf deepvariant_output/output.vcf.gz \
    --num_shards $(nproc) \
    --regions chr20 \
    --intermediate_results_dir outputs/intermediate_results # 这个参数是我自己加的，目的是让Deepvariant输出中间文件，方便后续debug

其中--num_shards表示将数据切分为多少份再传入make_examples.py，因此这个参数设置的是第一步的线程数，第二步（call_variants）未给出设置线程的方法，第三步（postprocess_variants）则是单线程工作，无需设置线程。

潜在问题与解决方案

潜在问题

前面提到DeepVariant提供的参数只能设置第一步（make_examples）的线程，而第三步（postprocess_variants）属于单线程任务，无需设置，因此可能出现问题的地方就在于第二步（call_variants）。

官网介绍这一步是多线程步骤，该步骤会调用TensorFlow进行神经网络的前向传播，而TensorFlow会默认占用CPU的所有计算核心，因此当一块CPU上跑了多个DeepVariant任务且都跑到第二步时，CPU就会非常繁忙，因为每个任务都会占用所有计算核心，但总的资源是固定的，CPU只能一会执行这个程序，一会又去执行另一个程序，这样下来会导致CPU出现高负载，每个任务的执行效率会下降，长此以往可能会损坏CPU。

下面介绍如何查看这一步实际使用的线程数，再介绍可行的解决方案。

当将上面的HiHi-case-study执行完毕以后，在--intermediate_results_dir参数指定的目录下面会出现每个步骤的结果文件，如第一个步骤的结果文件为make_examples.tfrecord-000xx-of-000xx.gz ，其中xx根据--num_shards 的结果确定。

因为DeepVariant在执行这三个步骤前都会将要执行的命令打印输出，这样就很方便我们进行debug。可以从它的log文件中找到第二步所要执行的命令如下

# 在log文件中搜索Running即可找到下面命令
***** Running the command:*****
( time /opt/deepvariant/bin/call_variants --outfile "outputs/intermediate_results/call_variants_output.tfrecord.gz" --examples "outputs/intermediate_results/make_examples.tfrecord@10.gz" --checkpoint "/opt/models/pacbio/model.ckpt" )

可以发现该步骤会读取outputs/intermediate_results/make_examples.tfrecord@10.gz这里指定的10个文件，并读取预先训练好的模型/opt/models/pacbio/model.ckpt进行预测，将结果输出在outputs/intermediate_results/call_variants_output.tfrecord.gz文件中。这一步只有这三个参数，没有指定线程的参数。

当在终端单独运行这一步时，可以使用top命令观察到该程序的%CPU占用，该占比除以100即为所用线程数，可以发现该程序的%CPU超过了100%，说明该程序并不使用一个线程。

解决方案

由于DeepVariant并未提供指定TensorFlow线程的参数，因此只能通过修改特定环境变量的形式来限制其使用的线程数。经过搜索，在该issue中找到了可行方案，现将方案解释如下：

TensorFlow所使用的线程主要受三个环境变量控制，分别是OMP_NUM_THREADS（表示OpenMP并行库的线程数）, TF_NUM_INTEROP_THREADS和TF_NUM_INTRAOP_THREADS（表示TensorFlow内部进行计算所使用的线程数），因此只需要同时设置这三个环境变量即可限制TensorFlow使用的线程。

实际测试发现call_variants这一步所用的参数会比这三个环境变量指定的线程数多2个左右，因此若拟让该步骤使用10个线程，那么环境变量需要设置为8，即比计划线程数小2。

为了方便使用，将这一步骤封装如下：

set_TF_threads(){
	### 设置TensorFlow所用线程数
    threads=$1
    threads_to_set=$((threads-2)) # 测试发现实际使用的线程比设置的线程数多2个左右，因此设置线程数要减去2

    ### 通过下面三个环境变量来控制TensorFlow使用的线程数目
    ### 这三个环境变量必须全部设置，否则无法很好的限制TensorFlow使用的线程数
    export OMP_NUM_THREADS=${threads_to_set}
    export TF_NUM_INTEROP_THREADS=${threads_to_set}
    export TF_NUM_INTRAOP_THREADS=${threads_to_set}
    echo "Have set threads of TensorFlow to $threads"
}

num_threads=8 # 拟使用8个线程
set_TF_threads $num_threads

因此一个更完整的使用DeepVariant的示例如下：

set_TF_threads(){
		### 设置TensorFlow所用线程数
    threads=$1
    threads_to_set=$((threads-2)) # 测试发现实际使用的线程比设置的线程数多2个左右，因此设置线程数要减去2

    ### 通过下面三个环境变量来控制TensorFlow使用的线程数目
    ### 这三个环境变量必须全部设置，否则无法很好的限制TensorFlow使用的线程数
    export OMP_NUM_THREADS=${threads_to_set}
    export TF_NUM_INTEROP_THREADS=${threads_to_set}
    export TF_NUM_INTRAOP_THREADS=${threads_to_set}
    echo "Have set threads of TensorFlow to $threads"
}

intermediate_results_dir=outputs/intermediate_results; mkdir -p ${intermediate_results_dir}

threads=10
region=chr20
set_TF_threads $num_threads

ulimit -u 10000 # https://stackoverflow.com/questions/52026652/openblas-blas-thread-init-pthread-create-resource-temporarily-unavailable/54746150#54746150
BIN_VERSION="1.5.0"
mkdir -p deepvariant_output

singularity exec --bind /usr/lib/locale/ \
  docker://google/deepvariant:${BIN_VERSION} \
    /opt/deepvariant/bin/run_deepvariant \
    --model_type PACBIO \
    --ref reference/GRCh38_no_alt_analysis_set.fasta \
    --reads input/HG003.GRCh38.chr20.pFDA_truthv2.bam \
    --output_vcf deepvariant_output/output.vcf.gz \
    --num_shards ${threads} \
    --regions "${region}"\
    --intermediate_results_dir outputs/intermediate_results # 这个参数是我自己加的，目的是让Deepvariant输出中间文件，方便后续debug

这样将可以较好的设置DeepVariant在第二个步骤中使用的线程数和申请的线程数一致。

至于PEPPER也是一样的道理，在运行PEPPER之前调用set_TF_threads函数来设置线程数即可。