Add CI stack for ML packages

[adamjstewart]

Basic stack of ML packages we would like to test and generate binaries for in CI.

Closes #31551

@tgamblin @scottwittenburg I said this when I was working on #27798, but this is the least intuitive and least user friendly thing I've ever done in Spack. I basically copy-n-pasted a gigantic file containing a million settings that have absolutely no meaning to me in the hopes that it works and I didn't make any mistakes. If there is even a single setting that is wrong in this file, someone is going to have to help me fix it, because I'm unaware of any documentation explaining what all of these mysterious keywords mean. Does this have to be this complicated? Can we have a file of defaults somewhere and I only need to override the ones I care about? I don't see why this should be any more than a simple environment file like:

spack:
  specs:
    - py-tensorflow
    - py-torch

[adamjstewart]

TODO: add ml tag to all packages we decide to consider ML libraries and test in CI.

[adamjstewart]

@scottwittenburg any tips to get this pipeline to run?

[blue42u]

Advice from a fellow traveler on this road (#31598). See comments below.

[adamjstewart]

Still no luck.

[adamjstewart]

@spackbot run pipeline

[spackbot-app]

I've started that pipeline for you!

[adamjstewart]

I think it's working! Like there are still concretization issues I need to resolve, but it gets far enough that it tries to concretize the environment!

[adamjstewart]

@aweits can you investigate the build issues with JAX in CI?

[aweits]

@aweits can you investigate the build issues with JAX in CI?

replicated a build failure locally w/ bazel 5.2.0, builds fine with 4.1.0...

[adamjstewart]

On what OS? I was hoping that Bazel 5 build failures would only be on macOS arm64...

[aweits]

On what OS? I was hoping that Bazel 5 build failures would only be on macOS arm64...

linux-rhel7-skylake_avx512

I think it's this:

bazelbuild/bazel#13811

[adamjstewart]

Okay, so from reading that, it seems like the problem is at the package level, not the bazel level. So would the right approach be to lock down supported Bazel versions until those rules are updated?

[aweits]

Okay, so from reading that, it seems like the problem is at the package level, not the bazel level. So would the right approach be to lock down supported Bazel versions until those rules are updated?

That's quite probably the "right" way to do it - alternatively, the workaround in that thread seems to work:

diff --git a/var/spack/repos/builtin/packages/py-jaxlib/package.py b/var/spack/repos/builtin/packages/py-jaxlib/package.p
index eb543a2..3ad4399 100644
--- a/var/spack/repos/builtin/packages/py-jaxlib/package.py
+++ b/var/spack/repos/builtin/packages/py-jaxlib/package.py
@@ -67,3 +67,19 @@ def patch(self):
             "build/build_wheel.py",
             string=True,
         )
+
+        matchtext = '''load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")'''
+        newtext = '''
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+http_archive(
+    name = "build_bazel_rules_apple",
+    sha256 = "0052d452af7742c8f3a4e0929763388a66403de363775db7e90adecb2ba4944b",
+    urls = [
+        "https://github.com/bazelbuild/rules_apple/releases/download/0.31.3/rules_apple.0.31.3.tar.gz",
+    ],
+)
+'''
+        filter_file(matchtext,
+                    newtext,
+                    'WORKSPACE',
+                    string=True)

[adamjstewart]

@spackbot run pipeline

[spackbot-app]

I've started that pipeline for you!

[adamjstewart]

I think this PR is finally ready to merge! All new pipelines are passing. Some caveats:

1. Some existing pipelines are failing. I have no control over those. I can force merge if I need to.
2. PyTorch does not yet support ROCm, JAX does not yet build properly. These issues can be solved in future PRs.
3. There are 3 separate stacks for CPU, CUDA, and ROCm pipelines. 90% of the packages in those DAGs do not have +cuda/+rocm variants, and so are duplicated. I would love to be able to write a single pipeline with a matrix over all possible variant combinations, but that's not possible at the moment: Environments: matrix of optional variants #32893
4. We would like to someday repeat these 3 stacks on several OS/arch combinations. Much like 3, this is not yet supported without duplicating the entire pipeline dozens of times. Would love to see support for this someday: Environments: matrix of OS/arch #32894

Addressed by @adamjstewart

[tgamblin]

This looks really awesome to me -- thanks for all the work everyone. Once we merge, this should start getting tested on every PR. We may still have issues on develop until GitHub merge queue goes into broader beta later this year (then we can use it) but I think this goes a most of the way towards keeping these packages built!