torch.compile fails because it does not find ldconfig -p

[meditans]

Describe the bug

Trying to optimize a machine learning with torch.compile for a GPU creates an error because ldconfig -p is not found.
Here's the precise error:

torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
CalledProcessError: Command '['ldconfig', '-p']' returned non-zero exit status 1.

Steps To Reproduce

On a computer with pytorch and an enabled NVidia GPU, try to run the following python program:

import torch
import torch.nn as nn

torch.cuda.set_device(0)


class Simple(nn.Module):
    def __init__(self):
        super().__init__()
        self.ff = nn.Linear(8, 8)
        self.cc = nn.Conv2d(10, 10, kernel_size=3, padding="same")

    @torch.compile
    def forward(self, x):
        x = self.ff(x)
        x = self.cc(x)
        return x


simple = Simple().to("cuda")
x = torch.rand((1, 10, 8, 8)).to("cuda")
simple(x)

Expected behavior

The program should complete successfully.

Additional context

This reminds me of another issue I raised on not having ldconfig -p in a pytorch environment, #285307

Notify maintainers

@NixOS/cuda-maintainers @SomeoneSerge @Madouura

Metadata

Please run nix-shell -p nix-info --run "nix-info -m" and paste the result.

$ nix run nixpkgs#nix-info -- -m
 - system: `"x86_64-linux"`
 - host os: `Linux 5.10.215, NixOS, 24.05 (Uakari), 24.05.20240427.58a1abd`
 - multi-user?: `yes`
 - sandbox: `yes`
 - version: `nix-env (Nix) 2.18.2`
 - nixpkgs: `/nix/store/801l7gvdz7yaibhjsxqx82sc7zkakjbq-source`

---

Add a 👍 reaction to issues you find important.

[SomeoneSerge]

Just in case this is the same error propagating from triton, could you please test #320888? I probably won't get around to do this myself for another week

[meditans]

Hi @SomeoneSerge, I tested the PR you mentioned and it's surely a step in the right direction! I don't get the ldconfig -p related errors anymore.

I get this now, which is technically a different problem but still related to torch.compile (I didn't have much luck googling for it):

  File "/nix/store/r42c89mj865p7vjb62jnknxi9bsqlr47-python3-3.11.9-env/lib/python3.11/site-packages/torch/utils/_triton.py", line 37, in triton_backend_hash
    from triton.common.backend import get_backend, get_cuda_version_key
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:

ImportError: cannot import name 'get_cuda_version_key' from 'triton.common.backend' (/nix/store/r42c89mj865p7vjb62jnknxi9bsqlr47-python3-3.11.9-env/lib/python3.11/site-packages/triton/common/backend.py)

Do you know where this could stem from?

For the curious, here's the flake I used to test the PR:

{
  description = "Check SomeoneSerge's ldconfig -p fix";

  inputs.nixpkgs.url = "github:SomeoneSerge/nixpkgs/530d5769d57e88bde975c5fefb9ecfcddba89742";

  outputs =
    { self, nixpkgs }:
    let
      pkgs = import nixpkgs {
        system = "x86_64-linux";
        config.allowUnfree = true;
      };
    in
    {
      devShell.x86_64-linux =
        with pkgs;
        (mkShell.override { stdenv = gcc11Stdenv; }) {
          venvDir = "./.venv";
          buildInputs = [
            (pkgs.python3.withPackages (
              ps: with ps; [
                torch-bin
                openai-triton-bin
              ]
            ))
            pkgs.virtualenv
            pkgs.python3Packages.venvShellHook
          ];
          postVenvCreation = ''
            unset SOURCE_DATE_EPOCH
          '';
        };
    };
}

[SomeoneSerge]

ImportError: cannot import name 'get_cuda_version_key' from 'triton.common.backend' (/nix/store/r42c89mj865p7vjb62jnknxi9bsqlr47-python3-3.11.9-env/lib/python3.11/site-packages/triton/common/backend.py)

Maybe our torch and triton had diverged? We may be relaxing version constraints sometimes and we may not have included integration tests for torch.compile because we've so far been just relying on the upstream's test suite. Speaking of which, @meditans, do you think you could package your reproducers into passthru.tests for our torch? Even if they require cuda (it's desirable to also have cpu-only tests because ofborg builds them, but it's ok to start small) you can experiment with:

nixpkgs/pkgs/development/python-modules/torch/gpu-checks.nix

Lines 16 to 26 in efcebae

	cudaPackages.writeGpuTestPython
	{
	inherit feature;
	libraries = [ torch ];
	name = "${feature}Available";
	}
	''
	import torch
	message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
	assert torch.cuda.is_available() and torch.version.${versionAttr}, message
	print(message)

nixpkgs/pkgs/development/python-modules/pynvml/default.nix

Lines 55 to 60 in efcebae

	passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
	import pynvml
	from pynvml.smi import nvidia_smi # noqa: F401
	print(f"{pynvml.nvmlInit()=}")
	'';