fix: keep A10 GRID on the v19 LTS branch

.github/workflows/ci.yaml

@@ -237,3 +237,7 @@ jobs:
           sudo apt-get install -y bats
       - name: Run install.sh unit tests
         run: bats test/install.bats
+      - name: Install GRID updater dependencies
+        run: python3 -m pip install -r grid_requirements.txt
+      - name: Test GRID branch selection
+        run: python3 -m unittest discover -s test -p 'test_*.py'

README.md

@@ -5,6 +5,34 @@ Kubernetes Nvidia GPU integration. Run it as a privileged container in the host
 It will enter the host mount namespace and install the nvidia drivers, container runtime, 
 and associated libraries on the host, validating their functionality
 
+## GRID driver branches
+
+The GRID images intentionally track different branches from the
+[Azure Linux GRID compatibility matrix](https://learn.microsoft.com/azure/virtual-machines/linux/n-series-driver-setup#supported-grid-drivers):
+
+| Configuration | Image | Driver | Azure compatibility |
+| --- | --- | --- | --- |
+| `grid` | `aks-gpu-grid` | vGPU 19.6 LTS / `580.178.04` | NVadsA10_v5 |
+| `grid_v20` | `aks-gpu-grid-v20` | vGPU 20.2 / `595.91.07` | NCv6 RTX PRO 6000 BSE, NCasT4_v3, NVadsA10_v5 |
+
+AgentBaker selects the standard GRID image for its A10 GRID SKUs and the v20 image
+for NCv6 RTX PRO 6000 BSE. CUDA driver selection is separate and is unchanged;
+the presence of a GRID option for a SKU does not require switching its CUDA
+workloads to GRID. The updater checks both upstream `Latest` and `Archive`
+releases, but only updates each image within its configured driver branch.
+
+Changing this producer does not update AKS nodes by itself. After the R580 image
+is published, AgentBaker must pin its actual MCR tag and constrain Renovate for
+`aks/aks-gpu-grid` to R580; existing R595 tags in that repository must not be
+selected for the A10 LTS path.
+
+Keeping the R580 option avoids forcing older applications onto R595. For example,
+[Isaac Sim 5.1 pins Kit SDK 107.3.3](https://github.com/isaac-sim/IsaacSim/blob/v5.1.0/deps/kit-sdk.packman.xml),
+and [upstream reports that this Kit version does not support 59x drivers](https://github.com/isaac-sim/IsaacSim/issues/794#issuecomment-5606456050).
+This is not a guarantee for every application: the separate MIG/vGPU identifier
+bug in that issue requires a newer Kit version, even with a supported driver.
+The exact application workload still needs validation on its target SKU.
+
 ## Build
 ```
 docker build -f Dockerfile  --build-arg DRIVER_VERSION=??? -t docker.io/alexeldeib/aks-gpu:latest .

auto_update.py

@@ -31,26 +31,41 @@ def validate_driver_url(url):
     return url
 
 
-def get_latest_grid_driver():
+def get_grid_driver_data():
     # URL of the JSON file containing driver information
     url = "https://raw.githubusercontent.com/Azure/azhpc-extensions/refs/heads/master/NvidiaGPU/Nvidia-GPU-Linux-Resources.json"
-    response = requests.get(url)
-    response.raise_for_status()  
-    data = response.json()
-
-    # Extract the latest GRID driver information
-    grid_versions = data['Latest']['Category']
-    grid_info = next((item for item in grid_versions if item["Name"] == "GRID"), None)
-
-    if grid_info:
-        latest_version_info = grid_info['Versions'][0]
-        latest_version = validate_driver_version(latest_version_info['DriverVersion'])
-        latest_url = validate_driver_url(latest_version_info['Driver'][0]['DirLink'])
-        return latest_version, latest_url
-
-    raise Exception("Could not find latest GRID driver version")
+    response = requests.get(url, timeout=30)
+    response.raise_for_status()
+    return response.json()
+
+
+def get_latest_grid_driver(data, driver_branch):
+    candidates = []
+    for section in ("Latest", "Archive"):
+        for release in data[section]:
+            for category in release["Category"]:
+                if category["Name"] != "GRID":
+                    continue
+                for version_info in category["Versions"]:
+                    version = validate_driver_version(version_info["DriverVersion"])
+                    if version.split(".")[0] == driver_branch:
+                        candidates.append(version_info)
+
+    if not candidates:
+        raise ValueError(f"Could not find GRID R{driver_branch} in upstream driver metadata")
+
+    latest = max(
+        candidates,
+        key=lambda item: tuple(int(part) for part in item["DriverVersion"].split(".")),
+    )
+    version = latest["DriverVersion"]
+    url = validate_driver_url(latest["Drivers"][0]["DirLink"])
+    expected_filename = f"NVIDIA-Linux-x86_64-{version}-grid-azure.run"
+    if os.path.basename(urlparse(url).path) != expected_filename:
+        raise ValueError(f"Driver URL does not match GRID version {version}: {url!r}")
+    return version, url
+
 
-# Add this at the end of your update_driver_config function
 def update_driver_config():
     yaml = YAML()
     yaml.preserve_quotes = True
@@ -62,12 +77,19 @@ def update_driver_config():
     with open("driver_config.yml", "r") as f:
         config = yaml.load(f)
 
-    # Get latest version and URL
-    latest_version, latest_url = get_latest_grid_driver()
-
-    # Update the grid section while preserving order
-    config['grid']['version'] = latest_version
-    config['grid']['url'] = latest_url
+    data = get_grid_driver_data()
+    for config_key in ("grid", "grid_v20"):
+        current_version = validate_driver_version(config[config_key]["version"])
+        latest_version, latest_url = get_latest_grid_driver(
+            data, current_version.split(".")[0]
+        )
+        if tuple(map(int, latest_version.split("."))) < tuple(map(int, current_version.split("."))):
+            raise ValueError(
+                f"Refusing to downgrade {config_key} from {current_version} to {latest_version}"
+            )
+        config[config_key]["version"] = latest_version
+        config[config_key]["url"] = latest_url
+        print(f"{config_key}: {current_version} -> {latest_version}")
 
     # Write back to file
     with open("driver_config.yml", "w") as f:

driver_config.yml

@@ -7,9 +7,10 @@ cuda_lts:
   version: "580.159.04"
 
 grid:
-  version: "595.91.07"
+  # vGPU 19.6 LTS for NVadsA10_v5; keep the R580 branch separate from GRID v20.
+  version: "580.178.04"
   # We do not support GRID drivers on ARM64 architecture.
-  url: "https://download.microsoft.com/download/a7cb6d36-3bbc-43d6-9e88-e0842e6f9ab9/NVIDIA-Linux-x86_64-595.91.07-grid-azure.run"
+  url: "https://download.microsoft.com/download/4efc1974-7b6f-431f-ab2d-b8d99bb8a512/NVIDIA-Linux-x86_64-580.178.04-grid-azure.run"
 
 grid_v20:
   version: "595.91.07"

test/test_auto_update.py

@@ -0,0 +1,155 @@
+import copy
+import os
+from pathlib import Path
+import tempfile
+import unittest
+from unittest.mock import Mock, patch
+
+from ruamel.yaml import YAML
+
+import auto_update
+
+
+def grid_version(version):
+    return {
+        "DriverVersion": version,
+        "Drivers": [{
+            "DirLink": (
+                "https://download.microsoft.com/download/test/"
+                f"NVIDIA-Linux-x86_64-{version}-grid-azure.run"
+            ),
+        }],
+    }
+
+
+def driver_data():
+    return {
+        "Latest": [{
+            "Category": [{
+                "Name": "GRID",
+                "Versions": [grid_version("595.91.07")],
+            }],
+        }],
+        "Archive": [{
+            "Category": [{
+                "Name": "GRID",
+                "Versions": [grid_version("580.178.04"), grid_version("570.237")],
+            }],
+        }],
+    }
+
+
+class GridDriverUpdateTests(unittest.TestCase):
+    def test_selects_lts_from_archive_and_v20_from_latest(self):
+        for branch, expected in (("580", "580.178.04"), ("595", "595.91.07")):
+            with self.subTest(branch=branch):
+                version, url = auto_update.get_latest_grid_driver(driver_data(), branch)
+                self.assertEqual(version, expected)
+                self.assertTrue(url.endswith(f"{expected}-grid-azure.run"))
+
+    def test_selects_highest_numeric_version_within_branch(self):
+        data = driver_data()
+        data["Archive"][0]["Category"][0]["Versions"].extend([
+            grid_version("580.99.01"),
+            grid_version("580.180.01"),
+        ])
+        version, _ = auto_update.get_latest_grid_driver(data, "580")
+        self.assertEqual(version, "580.180.01")
+
+    def test_missing_branch_fails_instead_of_switching_to_latest(self):
+        with self.assertRaisesRegex(ValueError, "Could not find GRID R580"):
+            auto_update.get_latest_grid_driver(
+                {"Latest": driver_data()["Latest"], "Archive": []}, "580"
+            )
+
+    def test_rejects_invalid_upstream_version(self):
+        data = driver_data()
+        data["Latest"][0]["Category"][0]["Versions"][0]["DriverVersion"] = "invalid"
+        with self.assertRaisesRegex(ValueError, "Unexpected driver version"):
+            auto_update.get_latest_grid_driver(data, "595")
+
+    def test_rejects_untrusted_driver_url(self):
+        data = driver_data()
+        data["Archive"][0]["Category"][0]["Versions"][0]["Drivers"][0]["DirLink"] = (
+            "https://example.com/driver.run"
+        )
+        with self.assertRaisesRegex(ValueError, "Driver URL host is not allowed"):
+            auto_update.get_latest_grid_driver(data, "580")
+
+    def test_rejects_installer_for_a_different_driver_version(self):
+        data = driver_data()
+        data["Archive"][0]["Category"][0]["Versions"][0]["Drivers"] = (
+            grid_version("595.91.07")["Drivers"]
+        )
+        with self.assertRaisesRegex(ValueError, "Driver URL does not match GRID version"):
+            auto_update.get_latest_grid_driver(data, "580")
+
+    def test_fetch_uses_timeout_and_propagates_http_errors(self):
+        response = Mock()
+        response.raise_for_status.side_effect = auto_update.requests.HTTPError("failed")
+        with patch.object(auto_update.requests, "get", return_value=response) as get:
+            with self.assertRaises(auto_update.requests.HTTPError):
+                auto_update.get_grid_driver_data()
+            self.assertEqual(get.call_args.kwargs["timeout"], 30)
+            response.json.assert_not_called()
+
+    def test_update_preserves_lts_v20_and_cuda_separation(self):
+        original = {
+            "cuda": {"version": "595.71.05"},
+            "cuda_lts": {"version": "580.159.04"},
+            "grid": {"version": "580.100.01", "url": "old"},
+            "grid_v20": {"version": "595.58.03", "url": "old"},
+        }
+        with tempfile.TemporaryDirectory() as directory:
+            config_path = Path(directory) / "driver_config.yml"
+            yaml = YAML()
+            with config_path.open("w") as config:
+                yaml.dump(copy.deepcopy(original), config)
+            previous_directory = os.getcwd()
+            try:
+                os.chdir(directory)
+                with patch.object(auto_update, "get_grid_driver_data", return_value=driver_data()):
+                    auto_update.update_driver_config()
+            finally:
+                os.chdir(previous_directory)
+            with config_path.open() as config:
+                updated = yaml.load(config)
+        self.assertEqual(updated["grid"]["version"], "580.178.04")
+        self.assertEqual(updated["grid_v20"]["version"], "595.91.07")
+        self.assertEqual(updated["cuda"], original["cuda"])
+        self.assertEqual(updated["cuda_lts"], original["cuda_lts"])
+
+    def test_update_refuses_to_downgrade_and_leaves_config_unchanged(self):
+        original = {
+            "grid": {"version": "580.180.01", "url": "old"},
+            "grid_v20": {"version": "595.91.07", "url": "old"},
+        }
+        with tempfile.TemporaryDirectory() as directory:
+            config_path = Path(directory) / "driver_config.yml"
+            yaml = YAML()
+            with config_path.open("w") as config:
+                yaml.dump(original, config)
+            before = config_path.read_text()
+            previous_directory = os.getcwd()
+            try:
+                os.chdir(directory)
+                with patch.object(auto_update, "get_grid_driver_data", return_value=driver_data()):
+                    with self.assertRaisesRegex(ValueError, "Refusing to downgrade grid"):
+                        auto_update.update_driver_config()
+            finally:
+                os.chdir(previous_directory)
+            self.assertEqual(config_path.read_text(), before)
+
+    def test_checked_in_config_keeps_grid_branch_and_installer_consistent(self):
+        path = Path(__file__).resolve().parents[1] / "driver_config.yml"
+        with path.open() as config:
+            data = YAML().load(config)
+        for key, branch in (("grid", "580"), ("grid_v20", "595")):
+            with self.subTest(key=key):
+                version = data[key]["version"]
+                self.assertEqual(version.split(".")[0], branch)
+                self.assertTrue(data[key]["url"].endswith(f"{version}-grid-azure.run"))
+
+
+if __name__ == "__main__":
+    unittest.main()