diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 29482a3f..fed7667e 100644 --- a/.github/workflows/ci.yaml +++ b/.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' diff --git a/README.md b/README.md index 307a5920..29daeaf9 100644 --- a/README.md +++ b/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 . diff --git a/auto_update.py b/auto_update.py index 4a2cf4db..1e748d57 100644 --- a/auto_update.py +++ b/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: diff --git a/driver_config.yml b/driver_config.yml index 83486fd8..bbc6c6d9 100644 --- a/driver_config.yml +++ b/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" diff --git a/test/test_auto_update.py b/test/test_auto_update.py new file mode 100644 index 00000000..3d878343 --- /dev/null +++ b/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()