Store compiled plugins to disk

[drake7707]

I'm currently using nvdiffrast in pytorch/pytorch:2.4.0-cuda12.1-cudnn9-devel and everything works as expected. I've built a wheel and would like to use it in pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime as it's significantly smaller in image size. The installation of the wheel in the runtime image worked fine.

However the _get_plugin() still requires the nvcc compiler and development environment to compile the plugin. Instead of storing it memory, can it be stored in a cache folder as well (or at least give me an option to compile it ahead of time)? Then I can copy the those to the runtime environment and have no need for compilation in my production environment, I really don't like using development environments in production.

[s-laine]

Storing precompiled extensions in a wheel is a rather complex topic, and would require modifications to the plugin loader of nvdiffrast itself. The code doesn't currently support this.

A possible way to create a custom, small docker image would be by installing and compiling nvdiffrast directly in the Dockerfile. That is, base the image on the smaller runtime image, then using a single RUN instruction in the Dockerfile, install devel tools, install nvdiffrast, execute a simple script (e.g., the triangle example) that triggers the extension build, and uninstall the development tools. This should create a Docker layer where the only difference is that nvdiffrast has been installed and is ready to run.

However, it may be that PyTorch still requires the devel tools to be present even when the compiled plugin is in the cache. If so, you'd have to tweak its internals to bypass those parts.

[drake7707]

I'll try it out thanks, I thought that it built it every single time without a cache, but I might have missed that browsing through the code. I've used the "copy cache to runtime" trick before on Jetson nanos with some other project, compiling was adding a significant overhead and already having all binaries in place made starting everything much much quicker.

[s-laine]

Looking at the internals of PyTorch's cpp_extension, you might be able to load the compiled library from the cache directory by doing something similar to what _import_module_from_library() function here does. If that works, you should be able to modify nvdiffrast's _get_plugin() to just do that in the runtime environment, instead of invoking PyTorch's compilation machinery at all.