set CF_CUDA_ENABLED=True if compiler('cuda') is detected in meta.yaml

conda_smithy/configure_feedstock.py

@@ -2,6 +2,7 @@
 from itertools import product, chain
 import logging
 import os
+import re
 import subprocess
 import textwrap
 import yaml
@@ -588,6 +589,19 @@ def _render_ci_provider(
             if ver:
                 os.environ["DEFAULT_LINUX_VERSION"] = ver
 
+        # detect if `compiler('cuda')` is used in meta.yaml,
+        # and set appropriate environment variable
+        with open(
+            os.path.join(forge_dir, forge_config["recipe_dir"], "meta.yaml")
+        ) as f:
+            meta_lines = f.readlines()
+        # looking for `compiler('cuda')` with both quote variants;
+        # do not match if there is a `#` somewhere before on the line
+        pat = re.compile(r"^[^\#]*compiler\((\"cuda\"|\'cuda\')\).*")
+        for ml in meta_lines:
+            if pat.match(ml):
+                os.environ["CF_CUDA_ENABLED"] = "True"
+
         config = conda_build.config.get_or_merge_config(
             None,
             exclusive_config_file=forge_config["exclusive_config_file"],

tests/conftest.py

@@ -515,6 +515,37 @@ def choco_recipe(config_yaml, request):
     )
 
 
+@pytest.fixture(scope="function")
+def cuda_enabled_recipe(config_yaml, request):
+    with open(os.path.join(config_yaml, "recipe", "meta.yaml"), "w") as fh:
+        fh.write(
+            """
+package:
+    name: py-test
+    version: 1.0.0
+requirements:
+    build:
+        - {{ compiler('c') }}
+        - {{ compiler('cuda') }}
+    host:
+        - python
+    run:
+        - python
+about:
+    home: home
+    """
+        )
+    return RecipeConfigPair(
+        str(config_yaml),
+        _load_forge_config(
+            config_yaml,
+            exclusive_config_file=os.path.join(
+                config_yaml, "recipe", "default_config.yaml"
+            ),
+        ),
+    )
+
+
 @pytest.fixture(scope="function")
 def jinja_env(request):
     tmplt_dir = os.path.join(conda_forge_content, "templates")

tests/test_configure_feedstock.py

@@ -815,3 +815,26 @@ def test_cos7_env_render(py_recipe, jinja_env):
         else:
             if "DEFAULT_LINUX_VERSION" in os.environ:
                 del os.environ["DEFAULT_LINUX_VERSION"]
+
+
+def test_cuda_enabled_render(cuda_enabled_recipe, jinja_env):
+    has_env = "CF_CUDA_ENABLED" in os.environ
+    if has_env:
+        old_val = os.environ["CF_CUDA_ENABLED"]
+        del os.environ["CF_CUDA_ENABLED"]
+
+    try:
+        assert "CF_CUDA_ENABLED" not in os.environ
+        cnfgr_fdstk.render_azure(
+            jinja_env=jinja_env,
+            forge_config=cuda_enabled_recipe.config,
+            forge_dir=cuda_enabled_recipe.recipe,
+        )
+        assert os.environ["CF_CUDA_ENABLED"] == "True"
+
+    finally:
+        if has_env:
+            os.environ["CF_CUDA_ENABLED"] = old_val
+        else:
+            if "CF_CUDA_ENABLED" in os.environ:
+                del os.environ["CF_CUDA_ENABLED"]