在接触 Nix 大约半年后，我整理了相关的资料并先后写了 Nix 和 NixOps 的入门介绍，并在字节内部进行了分享，到现在已经有大几百人看过了。 虽然这究竟转化了多少人并让他们开始使用 Nix 不得而知，但还是收获了一个意外之喜，那就是用 Nix 帮助解决了组内的一些实际问题。
首先介绍一下我们的产品，我们提供了一个机器学习平台，其中有一个主要功能是用户指定他们要运行的任务的代码、入口命令、镜像、机器配置以及机器数量等配置，平台基于 K8S 和自定义调度器为用户启动对应的容器，并执行用户指定的任务。 我们还为 TensorFlowPS / PytorchDDP / MPI 等场景提供了自动环境变量配置的能力。
找上门来的问题来自一个客户。 他们虽然使用 GPU，但并不通过基于容器的方式来管理任务，而是通过 HPC 领域常用的 Slurm 。 Slurm 的整体架构与 K8S 大相径庭，但是客户已经在 Slurm 的基础上做了很多工作，为了让客户能够迁移到我们的平台，我们决定对 Slurm 的调度模式做一定的兼容。 Slurm 并没有提供基于容器的环境管理方式，只能提交原始的文件和入口命令到调度起来的机器上去执行。 而我们的兼容方式就是重新实现 Slurm 的一部分常用命令行 API（包括 sbatch 和 srun ），每次调用时，在背后解析对应的 Slurm 配置，动态调度起一个相等规格的 Slurm 集群，并执行用户的指令。 这里的 Slurm 集群包括了一个 master 节点和 n 个 worker 节点，各自都有我们预制的镜像。 支持这些能力之后客户的迁移和后续使用都一切正常。
接下来就是有趣的地方了，有一天用户突然说 Slurm 的环境管理确实不好使，想要我们支持自定义镜像的能力。 当然这个需求很离谱，毕竟 Slurm 本身都没法支持这个能力，不过结合我们的实际情况也不是做不了——这要求我们给用户的镜像注入启动 Slurm 集群需要的依赖，之后用之前的方式启动集群并执行任务就好了。 依赖注入是整个需求的难点，也正好撞在了 Nix 的刀刃上。 在我们的预制镜像中，基础镜像是 debian ，我们可以直接通过 apt 来安装 Slurm 、 environment-modules 等超算集群中的常见预置工具；但如果是用户的自定义镜像，我们没有办法做任何假设，它可能是某个稀奇古怪的系统，可能没有 apt 工具，可能用的是某个上古版本的 glibc …… 那么我们要怎么做到依赖注入呢？ 这里至少有一点可以确定： 我们运行的是一个 x86-64 架构的 Linux 系统。 结合 Nix 对软件依赖的完全捕获和特立独行的 /nix 目录机制，我们可以做到让 Nix 构建出来的软件在任意镜像上正常运行，而且不会跟用户现有环境产生任何冲突。 就算用户用的也是 Nix 镜像，/nix目录下的已有软件包也不会跟我们拷贝过去的有碰撞，参见 Nix 的哈希机制。
这样一来思路就很明显了，构建一个 Nix 环境，把之前我们在预制镜像中通过 apt 安装的包全部安装一遍，然后把这个环境连同它的所有依赖（也叫闭包）全部拷贝到用户镜像的同个位置，并把这个环境的目录添加到 PATH，这样用户就能指定任意镜像，同时平台还可以用这个镜像启动一个 Slurm 集群，还能提供一些预置工具了。
对应的 flake.nix 文件也很简单：
{
inputs = {
nixpkgs.url = "github:nixos/nixpkgs/nixos-21.11";
};
outputs = { self, nixpkgs }:
let
pkgs = nixpkgs.legacyPackages."x86_64-linux"; in
{
defaultPackage."x86_64-linux" = pkgs.buildEnv {
name = "slurm-pack";
paths = with pkgs; [
slurm
munge
tcl
self.packages."environment-modules"
];
pathsToLink = [ "/share" "/bin" ];
};
packages."environment-modules" = ... # 见附录
};
}
我们通过 buildEnv 构建了一个「环境」，指定了这个环境依赖的软件包以及需要聚合的路径，关于 buildEnv 的详细解释可以参考源码中的注释。 在这之后，把这个闭包整体拷贝到目标镜像的 /nix目录下就可以了，实现这一步的样例 Dockerfile ：
FROM nixos/nix AS nix
WORKDIR /root
# enable flake commands and tuna mirror
COPY nix.conf /etc/nix/nix.conf
COPY flake.nix flake.lock .
# build the environment, pick the closure and its path out
RUN nix build . && \
mkdir -p /tmp/nix-env-closure && \
cp -R $(nix-store -qR result) /tmp/nix-env-closure && \
readlink -f result > /tmp/env-path
FROM debian:11.2
COPY --from=nix /tmp/nix-env-closure /nix/store
COPY --from=nix /tmp/env-path /tmp/env-path
# add link to the environment and add it to path
RUN mkdir -p /nix/var && \
ln -s $(cat /tmp/env-path) /nix/var/ml-platform-env
ENV PATH="/nix/var/ml-platform-env/bin:${PATH}"
这里还有一些细节，比如 nix build 当前目录下的 result 指向构建出的软件包，因此我们通过 readlink -f result 获取了实际的绝对路径； 又因为 Docker 不支持我们把这个路径作为一个变量传递给两阶段构建中的下一个镜像，因此只能把这一信息存放在文件里拷贝给下一个镜像。
在这个需求完美解决之后，我们也改造了一些之前的实现，比如平台还支持运行一个基于用户镜像的开发机，本质上也是启动一个 K8S Pod ，在背后给它装上 sshd 、 code-server 等工具，方便用户做开发使用。 这里也是在向用户镜像注入我们的软件包。
最后还有一个有意思的细节。 我们知道现在很少有二进制分发的工具是完全静态链接的，至少也会依赖系统里的 glibc，原因可以看这里。 多嘴一句，这也导致在服务器上删除 glibc 会有严重的后果，例如这篇 如何拯救一台 glibc 被干掉的 Linux 服务器? 。 但如果是动态链接的话， Nix 要如何确保它分发的二进制使用的是 /nix 目录下对应版本的 glibc 而非用户系统里的某个 glibc 呢？ 毕竟版本对不上的话二进制文件基本是没有办法运行的。 在这之前我只知道 LD_LIBRARY_PATH 可以指定运行可执行文件时寻找动态链接库的路径，但是我们可以让分发的二进制固定加载某个路径下的动态链接库吗？ 这是可行的， 在 ELF 可执行文件标准中规定了DT_RPATH 属性，可以在可执行文件中声明要寻找动态链接库的路径，且优先级高于 LD_LIBRARY_PATH 。 在使用 nix 之前我们也有解决注入的 bash 、 sshd 等工具需要指定使用我们的 glibc 的问题，是通过 patchelf 工具修改了对应的二进制文件属性：
patchelf --set-interpreter /<our-root>/ld bash
patchelf --set-rpath "/<our-root>/lib/" bash
在整理的时候发现就连 patchelf 这个工具也是 Nix 社区开发的，为了方便修改一些已经编译好的二进制文件，把它们纳入到 Nix 的生态中。 当然，在改用 Nix 之后，这一过程就不再需要我们自己完成了。
附录
对 environment modules 的打包：
packages."environment-modules" = with pkgs; stdenv.mkDerivation rec {
pname = "modules";
version = "5.1.0";
src = fetchurl {
url = "https://github.com/cea-hpc/modules/releases/download/v${version}/modules-${version}.tar.gz";
sha256 = "1ab1e859b9c8bca8a8d332945366567fae4cf8dd7e312a689daaff46e7ffa949";
};
buildInputs = [ coreutils dejagnu tcl ps less ];
postPatch = "patchShebangs .";
configureFlags = [ "--with-bin-search-path=$PATH" "--with-tcl=${tcl}/lib/tclConfig.sh" ];
meta = {
homepage = "http://modules.sourceforge.net/";
description = "Dynamic modification of a user's environment via modulefiles";
};
};
};