python3Packages.torch.tests: MPS tests for darwin

[mikatammi]

Things done

Depends on #351778. Add new test to pytorch for Darwin platform. It tests CPU without torch.compile decorator, and MPS if MPS device is present.

Modify pytorch tests so it also tests MPS during CPU tests on Darwin. I'm not sure is this the right way to do this but at least this is a good starting point for discussion. I was also unsure do the tests work for other platforms. MPS is built for macOS by default if #351778 gets merged in its current state, so it does not need specific torchWithMps build.

• Built on platform(s)
  • x86_64-linux
  • aarch64-linux
  • x86_64-darwin
  • aarch64-darwin
• For non-Linux: Is sandboxing enabled in nix.conf? (See Nix manual)
  • sandbox = relaxed
  • sandbox = true
• Tested, as applicable:
  • NixOS test(s) (look inside nixos/tests)
  • and/or package tests
  • or, for functions and "core" functionality, tests in lib/tests or pkgs/test
  • made sure NixOS tests are linked to the relevant packages
• Tested compilation of all packages that depend on this change using nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD". Note: all changes have to be committed, also see nixpkgs-review usage
• Tested basic functionality of all binary files (usually in ./result/bin/)
• 24.11 Release Notes (or backporting 23.11 and 24.05 Release notes)
  • (Package updates) Added a release notes entry if the change is major or breaking
  • (Module updates) Added a release notes entry if the change is significant
  • (Module addition) Added a release notes entry if adding a new NixOS module
• Fits CONTRIBUTING.md.

---

Add a 👍 reaction to pull requests you find important.

[mikatammi]

Also got bitten by nixfmt check, fixing that made the diff quite hard to read

[samuela]

a lot of the diff appears to be just formatting changes. let's try to keep the git blame clean

[mikatammi]

a lot of the diff appears to be just formatting changes. let's try to keep the git blame clean

I re force-pushed the whole thing and now the diff only shows the lines I changed, however I now get complaint from the nix formatting checks.

[mikatammi]

a lot of the diff appears to be just formatting changes. let's try to keep the git blame clean

I re force-pushed the whole thing and now the diff only shows the lines I changed, however I now get complaint from the nix formatting checks.

Made another commit on top to clearly show which part is formatting and which is not

[mikatammi]

Actually are you doing merge commits nowadays? Should I make separate PR for the .git-blame-ignore-revs in case the commit hash gets mangled

[SomeoneSerge]

Actually are you doing merge commits nowadays

Yes we do, I think it's ok to keep it in 1 PR

[mikatammi]

I changed this now so there's a separate test for Darwin. I think in the end it is cleaner this way

[SomeoneSerge]

This doesn't need writeGpuTestPython then (it adds a gpuCheck attribute with requiredSystemFeatures set to access CUDA in the sandbox).

FWIW you make this a simple

runCommand "test-torch-mps" { nativeBuildInputs = [ (python.withPackages ...) ]; } ''
  python << \EOF
  import torch
  ...
  EOF
''

This maybe has the undesired side-effect that getting a script to run outside the sandbox would take an extra step, but also I don't think anyone would actually need that...

[mikatammi]

What do you think about this version? I made it to use writeShellApplcation and simply put the python script to separate file because it doesn't need any generated things from Nix side

[SomeoneSerge]

This will not run the test as part of the derivation though:)

Good idea to move the script into a separate file

It's still a bit disappointing that the script is essentially the same as for cuda/rocm, the difference is just one line which we could simply parameterize (e.g. instead of the @ decorator syntax, we could f = compile(f) under an if guard)

[mikatammi]

Can we somehow also add this as part of the tests for the package?

[mikatammi]

Is this good to go in now?

[SomeoneSerge]

Can we somehow also add this as part of the tests for the package?

We need to make a derivation which runs this script during the build, and we need to put that derivation under passthru.tests