docker镜像创建时使用GPU环境

我想将一个使用tensorrt加载的推理服务打包为docker镜像，如果推理引擎每次都加载onnx会比较慢，导致容器启动后一段时间内无法提供服务，所以最好能预先序列化为tensorrt engine文件保存在镜像中，这样容器启动的速度就很快了。

问题来了，由于engine文件高度依赖于系统运行时环境，如cuda, tensorrt的版本等。我在开发机上预先建好的engine文件，很有可能打包到docker镜像中后无法启动，我想到的方法是在docker build阶段，运行一段代码，直接在当前的docker环境中加载onnx并生成engine文件保存。这需要镜像构建时能够使用GPU环境，否则无法加载tensorrt。

为了实现这个目的，目前只查到一种方法，就是在docker 配置文件/etc/docker/daemon.json 中指定默认的runtime：

{
    "runtimes": {
        "nvidia": {
            "path": "/usr/bin/nvidia-container-runtime",
            "runtimeArgs": []
         } 
    },
    "default-runtime": "nvidia"  ## 这一步需要添加
}

将默认runtime设定为nvidia后，docker build阶段会默认开启gpu，dockerfile中可以运行需要GPU环境的代码了。构建好的镜像启动容器时，也不需要再指定--gpu。

感觉这个方法不太灵活，毕竟大部分情况下并不需要nvidia的runtime，镜像也不一定要默认使用GPU。如果docker build能像docker run 指定 --gpu参数 那样灵活设置就好了，不知道是不是有这种方法。

Advanced topics · NVIDIA/nvidia-docker Wiki · GitHub
在这里提到了

Default runtime

The default runtime used by the Docker® Engine is runc, our runtime can become the default one by configuring the docker daemon with --default-runtime=nvidia. Doing so will remove the need to add the --runtime=nvidia argument to docker run. It is also the only way to have GPU access during docker build.

看起来像是docker build阶段使用gpu的唯一办法了。。。