Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 32 additions & 3 deletions ci/tools/fetch_ctk_redistrib.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,25 @@ def load_metadata(*, metadata_path: str | None, metadata_url: str | None) -> dic
raise ValueError("exactly one of --metadata-path or --metadata-url is required")

if metadata_path is not None:
return json.loads(Path(metadata_path).read_text(encoding="utf-8"))
return _as_metadata_object(json.loads(Path(metadata_path).read_text(encoding="utf-8")), metadata_path)

assert metadata_url is not None
metadata_url = validate_metadata_url(metadata_url)
with urllib.request.urlopen(metadata_url) as response: # noqa: S310 - scheme is restricted to https above
return json.load(response)
return _as_metadata_object(json.load(response), metadata_url)


def _as_metadata_object(metadata: Any, source: str) -> dict[str, Any]:
"""Reject JSON that parsed fine but is not a redistrib manifest.

The manifest is downloaded with ``curl -LSs`` (no ``--fail``), so an error
page or a redirect body lands in the file and may still be valid JSON --
just not an object. Without this the failure surfaces several frames later
as ``TypeError: argument of type 'NoneType' is not iterable``.
"""
if not isinstance(metadata, dict):
raise ValueError(f"CTK redistrib metadata from {source} must be a JSON object, got {type(metadata).__name__}")
return metadata


def resolve_component_name(metadata: dict[str, Any], component: str) -> str:
Expand All @@ -101,7 +114,11 @@ def filter_components(
skipped = []
for component in filter_static_components(split_components(components), host_platform, cuda_version):
resolved_component = resolve_component_name(metadata, component)
if ctk_subdir in metadata.get(resolved_component, {}):
# Guard the type: a top-level key such as "release_label" holds a
# string, and ``ctk_subdir in "13.0.0"`` is a substring test rather
# than the intended key lookup.
component_info = metadata.get(resolved_component)
if isinstance(component_info, dict) and ctk_subdir in component_info:
filtered.append(resolved_component)
else:
skipped.append(component)
Expand All @@ -114,10 +131,22 @@ def get_component_relative_path(metadata: dict[str, Any], *, host_platform: str,
component_info = metadata.get(component)
if component_info is None:
raise KeyError(f"unknown CTK component {component!r}")
if not isinstance(component_info, dict):
# Real manifests carry string-valued top-level keys ("release_date",
# "release_label", "release_product") alongside the component objects,
# so "present" is not the same as "is a component".
raise KeyError(
f"CTK metadata entry {component!r} is not a component object (got {type(component_info).__name__})"
)

subdir_info = component_info.get(ctk_subdir)
if subdir_info is None:
raise KeyError(f"CTK component {component!r} is not available for redistrib subdir {ctk_subdir!r}")
if not isinstance(subdir_info, dict):
raise KeyError(
f"CTK component {component!r} entry for redistrib subdir {ctk_subdir!r} "
f"is not an object (got {type(subdir_info).__name__})"
)

relative_path = subdir_info.get("relative_path")
if relative_path is None:
Expand Down
117 changes: 117 additions & 0 deletions ci/tools/tests/test_fetch_ctk_redistrib.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import json
import os
import sys

import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from fetch_ctk_redistrib import main

# Shaped like a real redistrib_*.json: string-valued release keys sit at the
# top level alongside the component objects.
METADATA = {
"release_date": "2026-01-01",
"release_label": "13.0.0",
"release_product": "cuda",
"cuda_nvcc": {
"linux-x86_64": {"relative_path": "cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64.tar.xz"},
},
}


def write_metadata(tmp_path, payload):
path = tmp_path / "redistrib.json"
path.write_text(json.dumps(payload), encoding="utf-8")
return str(path)


def relpath_argv(metadata_path, component):
return [
"component-relative-path",
"--host-platform",
"linux-64",
"--component",
component,
"--metadata-path",
metadata_path,
]


def filter_argv(metadata_path, components="cuda_nvcc"):
return [
"filter-components",
"--host-platform",
"linux-64",
"--cuda-version",
"13.0.0",
"--components",
components,
"--metadata-path",
metadata_path,
]


@pytest.mark.agent_authored(model="claude-opus-5")
def test_valid_component_is_resolved(tmp_path, capsys):
assert main(relpath_argv(write_metadata(tmp_path, METADATA), "cuda_nvcc")) == 0
assert capsys.readouterr().out.strip() == "cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64.tar.xz"


@pytest.mark.agent_authored(model="claude-opus-5")
@pytest.mark.parametrize("component", ["release_label", "release_date", "release_product"])
def test_string_valued_top_level_key_is_not_a_component(tmp_path, capsys, component):
"""`is None` only rejects an absent key, not a wrongly-typed one.

Every real manifest carries these string-valued keys next to the component
objects, so asking for one used to reach `component_info.get(...)` and die
with `AttributeError: 'str' object has no attribute 'get'` instead of the
tool's own diagnostic.
"""
assert main(relpath_argv(write_metadata(tmp_path, METADATA), component)) == 1
assert "ERROR:" in capsys.readouterr().err


@pytest.mark.agent_authored(model="claude-opus-5")
def test_absent_component_still_reports_cleanly(tmp_path, capsys):
assert main(relpath_argv(write_metadata(tmp_path, METADATA), "not_a_component")) == 1
assert "unknown CTK component" in capsys.readouterr().err


@pytest.mark.agent_authored(model="claude-opus-5")
def test_non_object_subdir_entry_is_reported(tmp_path, capsys):
metadata = {"cuda_nvcc": {"linux-x86_64": "cuda_nvcc/linux-x86_64/x.tar.xz"}}
assert main(relpath_argv(write_metadata(tmp_path, metadata), "cuda_nvcc")) == 1
assert "ERROR:" in capsys.readouterr().err


@pytest.mark.agent_authored(model="claude-opus-5")
@pytest.mark.parametrize(
"payload",
[
pytest.param(None, id="null"),
pytest.param([1, 2], id="array"),
pytest.param("13.0.0", id="string"),
],
)
@pytest.mark.parametrize("argv_builder", [relpath_argv, filter_argv], ids=["relative-path", "filter"])
def test_metadata_that_is_not_an_object_is_reported(tmp_path, capsys, payload, argv_builder):
"""The manifest is downloaded with `curl -LSs` (no --fail), so an error
page or redirect body can parse as valid JSON that is not an object."""
path = write_metadata(tmp_path, payload)
argv = argv_builder(path, "cuda_nvcc") if argv_builder is relpath_argv else argv_builder(path)

assert main(argv) == 1
assert "must be a JSON object" in capsys.readouterr().err


@pytest.mark.agent_authored(model="claude-opus-5")
def test_filter_skips_a_string_valued_top_level_key(tmp_path, capsys):
assert main(filter_argv(write_metadata(tmp_path, METADATA), "release_label")) == 0
captured = capsys.readouterr()
assert captured.out.strip() == ""
assert "Skipping unsupported CTK component 'release_label'" in captured.err
Loading