From f81d7ee77a75f1859f939bb09b5ff84bc65cd075 Mon Sep 17 00:00:00 2001
From: Dom Heinzeller <dom.heinzeller@icloud.com>
Date: Tue, 5 Mar 2024 13:03:26 -0700
Subject: [PATCH] Add ai-env virtual package

---
 configs/templates/unified-dev/spack.yaml      |  1 +
 .../spack-stack/packages/ai-env/package.py    | 23 +++++++++++++++++++
 .../spack-stack/packages/ewok-env/package.py  |  5 ----
 .../packages/jedi-base-env/package.py         |  9 ++++----
 4 files changed, 28 insertions(+), 10 deletions(-)
 create mode 100644 spack-ext/repos/spack-stack/packages/ai-env/package.py

diff --git a/configs/templates/unified-dev/spack.yaml b/configs/templates/unified-dev/spack.yaml
index c6ada84bf..b451cc5be 100644
--- a/configs/templates/unified-dev/spack.yaml
+++ b/configs/templates/unified-dev/spack.yaml
@@ -9,6 +9,7 @@ spack:
   - compilers: ['%aocc', '%apple-clang', '%gcc', '%intel']
   - packages:
       - ewok-env +ecflow +cylc
+      - ai-env
       - geos-gcm-env
       - global-workflow-env
       - gmao-swell-env
diff --git a/spack-ext/repos/spack-stack/packages/ai-env/package.py b/spack-ext/repos/spack-stack/packages/ai-env/package.py
new file mode 100644
index 000000000..1d50e0a78
--- /dev/null
+++ b/spack-ext/repos/spack-stack/packages/ai-env/package.py
@@ -0,0 +1,23 @@
+# Copyright 2013-2022 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+from spack.package import *
+
+
+class AiEnv(BundlePackage):
+    """Development environment for AI/ML applications"""
+
+    # DH* TODO UPDATE FROM INTERNAL TO PUBLIC
+    homepage = "https://github.com/JCSDA/spack-stack"
+    git = "https://github.comJCSDA/spack-stack.git"
+
+    maintainers("climbfuji", "srherbener")
+
+    version("1.0.0")
+
+    depends_on("jedi-base-env", type="run")
+    depends_on("py-torch", type="run")
+
+    # There is no need for install() since there is no code.
diff --git a/spack-ext/repos/spack-stack/packages/ewok-env/package.py b/spack-ext/repos/spack-stack/packages/ewok-env/package.py
index 5090da78e..f4d494132 100644
--- a/spack-ext/repos/spack-stack/packages/ewok-env/package.py
+++ b/spack-ext/repos/spack-stack/packages/ewok-env/package.py
@@ -67,9 +67,4 @@ class EwokEnv(BundlePackage):
     depends_on("r2d2", when="+r2d2", type="run")
     depends_on("ewok", when="+ewok", type="run")
 
-    conflicts(
-        "%gcc platform=darwin",
-        msg="ewok-env does " + "not build with gcc on macOS, use apple-clang",
-    )
-
     # There is no need for install() since there is no code.
diff --git a/spack-ext/repos/spack-stack/packages/jedi-base-env/package.py b/spack-ext/repos/spack-stack/packages/jedi-base-env/package.py
index da9b0de69..01774834c 100644
--- a/spack-ext/repos/spack-stack/packages/jedi-base-env/package.py
+++ b/spack-ext/repos/spack-stack/packages/jedi-base-env/package.py
@@ -21,9 +21,6 @@ class JediBaseEnv(BundlePackage):
     variant("fftw", default=True, description="Build fftw")
     variant("hdf4", default=True, description="Build hdf4 library and python hdf module")
 
-    # Variants for AI packages
-    variant("py-torch", default=False, description="Enable support for py-torch")
-
     depends_on("base-env", type="run")
     depends_on("bison", type="run")
     depends_on("blas", type="run")
@@ -66,7 +63,9 @@ class JediBaseEnv(BundlePackage):
     depends_on("py-scipy", type="run")
     depends_on("py-xarray", type="run")
 
-    # AI packages
-    depends_on("py-torch", when="+py-torch", type="run")
+    conflicts(
+        "%gcc platform=darwin",
+        msg="jedi-base-env does not build with gcc on macOS, use apple-clang",
+    )
 
     # There is no need for install() since there is no code.