用于泛基因组学的伪Shell

# 泛基因组学的假壳 来源：https://www.cs.cornell.edu/~asampson/blog/flash.html 我一直在开发一个高效的泛基因组学工具包，名为FlatGFA (https://www.cs.cornell.edu/~asampson/blog/flatgfa.html)\. 与其他泛基因组学工具如odgi (https://odgi.readthedocs.io/en/latest/)相比，FlatGFA只有一个窍门：零拷贝数据格式。内存中的数据格式与磁盘上的格式完全相同，因此FlatGFA可以跳过所有序列化和反序列化成本；打开一个文件只需执行一次mmap\(2\) (https://linux.die.net/man/2/mmap)。对于不公平挑选的工作负载 (https://www.cs.cornell.edu/~asampson/blog/flatgfa.html#speedup)，FlatGFA可以比odgi快数千倍。 现在到了困难的部分：我希望我的基因组学家同事们真正使用FlatGFA。我想编写一份高性能操作清单，让真正的科学家将它们组合成完整的工作流。 为了实现这种组合，有两个简单的选择：要么 (1) 制作一个命令行界面，暴露所有操作符，让科学家编写shell脚本进行组合，要么 (2) 设计一个Rust API，让科学家编写Rust代码。两个选择都不太令人满意： 1. CLI方法限制了组合的类型。所有中间结果必须通过文件或管道传递，这可能会变得笨拙，而且肯定会有一些开销。 2. 我们内部的Rust API，由于我们耍的各种数据结构技巧，采用了一种相当独特的风格。尽管我们的生物学家合作者是出色的Rust黑客，但我不能昧着良心说我们有一个他们乐于使用的*好*API。 本篇文章介绍的是我们最近构建的一个非常愚蠢的替代方案：一个*假壳*，它假装提供选项1，但近似于选项2的性能。 ## 关于Ousterhout二分法 很长一段时间里，我认为“打包”像FlatGFA这样的性能导向库的正确方式可能是标准的Ousterhout二分法 (https://web.stanford.edu/~ouster/cgi-bin/papers/scripting.pdf)。性能敏感的例程保留在Rust中，但我们为更高级的语言构建绑定，以将这些例程组合成完整的工作流。结果看起来会很像PyTorch (https://pytorch.org/)：对于机器学习工程师来说，Python并不快，因为超过99%的时间花在用C++和CUDA编写的优化内核例程上，所以这并不重要。 在现代，Python是Ousterhout二分法中“胶水语言”部分的自然选择。1 (https://www.cs.cornell.edu/~asampson/blog/flash.html#fn:tcl)于是我们开始使用优秀的PyO3 (https://pyo3.rs/)项目为FlatGFA构建Python绑定。我们让基本功能 (https://cucapra.github.io/pollen/flatgfa/)相当好用——例如，试试以下操作看看效果： ``` $ curl -LO https://raw.githubusercontent.com/pangenome/odgi/refs/heads/master/test/LPA.gfa $ uv run --with flatgfa python >>> import flatgfa >>> graph = flatgfa.parse("LPA.gfa") >>> [path.name for path in graph.paths] ``` 然而，Python绑定有一些严重的缺点： - 即使使用PyO3，高效编写绑定也很难。问题在于Rust的静态生命周期与Python的动态管理堆之间不匹配的固有复杂性。FlatGFA的性能优势来自于消除拷贝、分配和指针追逐——而这些在Rust/Python边界上都容易重新出现。 - 我们无法获得整个工作负载的全局视图。直接的Python绑定意味着我们只有*在每次调用库内部*才有机会跑得快，而无法跨多个调用进行优化。例如，一旦用户编写了一个Python`for`循环来迭代FlatGFA的数据结构，我们几乎肯定会输掉性能游戏。这也是PyTorch有独立的、可选的编译模式 (https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler.html#torch-compiler-overview)的相同根本原因。 - 事实证明，我们的生物学家合作者其实并不那么喜欢Python！传统上，组合泛基因组学管道的熟悉方式是Unix shell。就我个人而言，我已经太习惯Python是易用性的默认选择。当然，领域专家的偏好因环境而异。 重新考虑odgi (https://odgi.readthedocs.io/en/latest/)及其同类工具所使用的面向CLI的方法是合理的。 ## 重新思考Shell 让我们看看这个领域中基于shell组合的一个例子。odgi文档 (https://odgi.readthedocs.io/en/latest/)中的一篇教程 (https://odgi.readthedocs.io/en/latest/rst/tutorials/detect_complex_regions.html#obtain-the-depth-over-the-pangenome)展示了如何通过组合odgi和bedtools (https://bedtools.readthedocs.io/en/latest/)的操作符来查找人类8号染色体中的重复序列： ``` odgi depth -i chr8.pan.og -r chm13#chr8 | \ bedtools makewindows -b /dev/stdin -w 5000 > chm13.chr8.w5kbps.bed odgi depth -i chr8.pan.og -b chm13.chr8.w5kbps.bed --threads 2 | \ bedtools sort > chr8.pan.depth.w5kbps.bed ``` 可能有人会觉得将shell脚本与功能齐全的动态脚本语言相比有些奇怪，但像这样的shell脚本比Python有一些实质性的优势： - 通过管道进行流式I/O在适当的情况下非常适合大型数据集。 - 简单的管道并行性易于表达。 - 在文件中持久化中间结果非常直接。 - shell算是*终极*胶水语言：你可以组合分别开发、用不同语言编写的组件，而无需在绑定上费额外功夫。（唯一的“绑定”就是Unix用户态API。） 这个示例工作流使用了来自两个不同包的四位操作符，两个Unix管道和一个中间文件。我认为在这个例子中关系不大，但shell管道让两对命令可以并发运行，这很不错。借用Greenberg等人 (https://nikos.vasilak.is/p/pash:hotos:2021.pdf)的话来说，*shell实际上挺好的。*2 (https://www.cs.cornell.edu/~asampson/blog/flash.html#fn:hotos) 然而，有一个巨大的缺点：在操作之间交换数据的唯一方式是文件和管道。文件可能会导致不必要地将数据写入磁盘，即使所有字节都能舒适地放在内存中。管道可以避免磁盘I/O，并且非常适合流式操作符，但它们通常需要将所有内容序列化为文本，而且并非每个生产者-消费者关系都天然支持流式处理。例如，如果一个命令生成了一个新的变异图（一个新的GFA文件），下一个命令很可能需要先读取整个文件才能开始工作。 在我们关于泛基因组学资助的每周会议上 (https://news.cornell.edu/stories/2021/11/5m-grant-will-tackle-pangenomics-computing-challenge)，小组就这些基于shell组合的根本限制展开了略显激烈的讨论。也许操作系统的磁盘缓存可以大部分缓解文件I/O成本？能否通过将文件放在RAM磁盘中来强制实现？（当你`mmap`一个位于RAM磁盘上的文件时，到底会发生什么？）当数据集增长到超出主内存大小时，这些方法可能都不实用？如果我们把所有底层命令都改为使用零拷贝二进制文件格式而不是基于文本的流式处理，这些权衡会如何变化？这种方法会不会让脚本更加混乱，并失去流水线的所有好处？ 在讨论中，我意识到存在一个荒谬、不切实际但非常有趣的替代方案，可以避开所有这些缺点。 ## 题外话：向量化解释器 Graydon演讲中名为'19th Century Tech Calls'的幻灯片，附有一张火车图片 (https://venge.net/graydon/talks/VectorizedInterpretersTalk-2023-05-12.pdf)2023年，Graydon Hoare在UCSC做了一个关于“向量化解释器”的演讲 (https://venge.net/graydon/talks/VectorizedInterpretersTalk-2023-05-12.pdf)，给我留下了深刻印象。3 (https://www.cs.cornell.edu/~asampson/blog/flash.html#fn:graydon)他指出，本地代码编译器（尤其是JIT）是从代码中提取性能的一种极其复杂的方式。让我印象深刻的想法是，在编程模型的适当配合下，*批量*操作的解释器可以成为一种简单快速的替代方案。如果你的字节码中的每条指令都代表对大量数据的一次大型计算（而不是，比如，一个简单的标量整数加法），那么直接解释该字节码就足够高效了。例如，当99.99%的时间都花在执行那些粗粒度指令的实现上时，就无需担心字节码指令分发的成本。 在Graydon的演示中，PyTorch和NumPy都是向量化解释器的例子。但正如我上面提到的，它们复用了Python的程序表示和解释器——因此它们的可寻址“指令窗口”是有限的。 我一直以为，对于泛基因组学操作，一定可以通过一个专门构建的向量化解释器做得更好。而shell脚本工作流的问题正好提供了一个尝试的借口。 ## 一个假壳 想法是构建一个*假壳*：它支持一小部分POSIX shell语法，并在运行泛基因组学操作符时“作弊”。目标是运行未修改的shell脚本，这些脚本使用传统的CLI工具，如odgi (https://odgi.readthedocs.io/en/latest/)和bedtools (https://bedtools.readthedocs.io/en/latest/)，并通过管道和文件进行通信。我们将通过机会主义地切换到更快的实现并避免I/O，使相同的生物学家编写的shell脚本运行得更快。 这个shell叫做Flash（FlatGFA shell），如果你想一起玩，可以在我们的泛基因组学单体仓库 (https://github.com/cucapra/pollen/tree/main/flatgfa-sh)中找到它。使用`cargo run`获取交互式提示。 ### Shell基本功能 Flash首先可以像真正的shell一样运行普通命令。例如，这是可行的： ``` echo llenroc | rev > message.txt ; cat message.txt ``` 为了实现这一点，我借用了现有的shell语法解析器 (https://crates.io/crates/brush-parser)，它来自一个“用Rust重写它”的shell项目 (https://crates.io/crates/brush)。但是，Flash没有直接解释shell AST，而是首先将其翻译成基于指令的中间表示。上面那个小脚本翻译成三条指令，每条对应它所运行的一条命令。如果你给Flash加上`--pretend`（`-p`）标志，它可以美化打印IR： ``` $ flash -p -c 'echo llenroc | rev > message.txt ; cat message.txt' shell("echo", ["llenroc"], input=stdin) -> pipe-0 shell("rev", [], input=pipe-0) -> "message.txt" shell("cat", ["message.txt"], input=stdin) -> stdout ``` 到目前为止，我们只使用了`shell`指令，它实际上会fork一个子进程（就像真正的shell一样）。Flash的IR是围绕*资源*构建的：资源是指令的输入和输出。这个程序使用了`stdin`和`stdout`资源、一个Unix管道和一个文件。Flash的IR求值器 (https://github.com/cucapra/pollen/blob/main/flatgfa-sh/src/eval/mod.rs)负责为代表每条指令设置管道和打开文件。 ### 作弊 Flash之所以是*假*壳，是因为它特化处理了一组内置的已知泛基因组学CLI工具。例如，让我们借用前面脚本的一部分： ``` $ flash -p -c 'odgi depth -i chr8.pan.gfa -r chm13#chr8' parse-gfa("chr8.pan.gfa") -> gfa-store-0 path-depth(gfa-store-0, path="chm13#chr8") -> stdout ``` Flash识别了我们的`odgi depth`调用，并生成了它可以在*内部*运行的一些专门指令，而不是`shell`指令。`path-depth`指令通过直接调用FlatGFA库中的Rust函数工作，并且从不fork子进程。 上面看到的`shell`指令使用输入和输出*资源*。在上面的IR列表中，它们被命名为`pipe-0`和`"message.txt"`。对于`shell`，唯一允许的资源类型是字节流（即管道和文件）。然而，`parse-gfa`指令产生了一种新类型的资源：一个*GFA存储*，这是泛基因组学变异图的高效内存表示。这是一个普通的Rust值，存储在Flash解释器的环境中。当最终求值`path-depth`指令时，Flash解释器检索该值并将其输入到相关的库函数中。 这就是发明假壳的全部原因：以便它能作弊。Flash机会主义地避免运行真正命令，而是使用一个精心设计的内置库。 ### 完整示例 让我们看看更完整工作流的IR表示。我将这个脚本放在一个名为`wdepth.sh`的文件中： ``` #!/bin/sh odgi depth -i chr8.pan.gfa -r chm13#chr8 \ | bedtools makewindows -b /dev/stdin -w 5000 > chm13.chr8.w5kbps.bed odgi depth -i chr8.pan.gfa -b chm13.chr8.w5kbps.bed rm -f chm13.chr8.w5kbps.bed ``` 这是IR列表： ``` $ flash -p wdepth.sh parse-gfa("chr8.pan.gfa") -> gfa-store-0 path-depth(gfa-store-0, path="chm13#chr8") -> pipe-0 parse-bed(pipe-0) -> bed-store-0 make-windows(bed-store-0, size=5000) -> "chm13.chr8.w5kbps.bed" parse-gfa("chr8.pan.gfa") -> gfa-store-1 parse-bed("chm13.chr8.w5kbps.bed") -> bed-store-1 interval-depth(gfa-store-1, bed-store-1) -> stdout shell("rm", ["-f", "chm13.chr8.w5kbps.bed"], input=stdin) -> stdout ``` 我喜欢Flash设计的一点是，它天然支持混合搭配不同的资源类型。有些指令通过文件系统或Unix管道与外部世界交互；另一些则与Flash的内部数据结构交互。两种类型的指令可以共存并交换数据。 这个示例脚本既可以在普通`sh`下运行，也可以在Flash下运行。让我们比较一下性能：4 (https://www.cs.cornell.edu/~asampson/blog/flash.html#fn:setup) 一张条形图，比较了脚本在sh下（使用本地odgi输入文件）、Flash下（使用本地FlatGFA输入文件）和Flash下（使用文本GFA输入文件）的运行情况。 前两个条是重要的比较：odgi（通过普通`sh`）和FlatGFA（通过Flash）都分别以其高效的二进制文件格式读取输入图。最后一个条显示了Flash在还必须解析标准文本GFA格式作为输入时的性能。 这个特定的简单shell脚本在Flash上运行比使用“真正”的shell和底层odgi快28倍。我觉得这已经是一个令人满意的加速了！而且很有趣的是，即使从文本GFA格式开始，Flash也能超越odgi。（我省略了odgi读取文本GFA的条，因为它需要超过7分钟。我相信，尽管没有证据，FlatGFA拥有世界上最快的GFA解析器。） ### 优化 这一切都构成了一个有趣的语言实现爱好 (https://discuss.systems/@adrian/116518791774005898)，但所有这些设置的真实原因是进行优化。Flash基于指令的IR使得优化成为可行（我无法想象如何直接在shell AST上实现它们）。以下是我到目前为止实现的优化： - 当程序生成一个BED文件，然后又用`parse-bed`指令加载它时，避免通过字节的往返。直接生成内存中的BED资源并直接使用。 - 识别实际只需要碱基对数的`path-depth`使用，并将其替换为更便宜的`path-length`指令。（这个有点取巧：恰好`odgi depth -r`是获取路径长度的便捷CLI方式，尽管它 （注：原文在此处被截断，但根据上下文，后续内容应继续关于优化的讨论。由于用户提供的输入到此结束，翻译也相应停止。如果需要继续翻译剩余部分，请提供完整原文。）