mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-17 01:11:33 -06:00
feat(node): auto-detect node capabilities via kernel interfaces
Closes the operator-burden gap the harness shakedown surfaced — the
list_nodes capability/region/role filtering surface that nodes were
launching with empty. Auto-detection runs at server startup and
populates ``node_metadata`` rows with sensible defaults that
operators can still override via the ``[metadata]`` section of
config.toml (operator-config writes win on the per-key upsert).
What's detected, all from kernel interfaces (no userspace binaries
on PATH — works the same way regardless of whether nvidia-smi /
rocm-smi / lspci is installed):
- ``gpu_count`` / ``gpu_vendor`` / ``gpu_vendors`` / ``gpus`` —
walks ``/sys/class/drm/cardN/device/{vendor,device}`` and decodes
PCI vendor IDs to friendly names (NVIDIA / AMD / Intel / Apple).
Heterogeneous-GPU nodes get the first KNOWN vendor in the flat
``gpu_vendor`` key — never ``"unknown"`` when known vendors are
present — so a coord filtering on ``gpu_vendor=nvidia`` matches
nodes whose first card happened to be exotic.
- ``memory_gb`` — reads ``/proc/meminfo``, rounds GiB down so
``filters={"memory_gb": 32}`` doesn't match a 31.5 GiB node.
- ``cpu_model`` — first ``model name`` line from ``/proc/cpuinfo``.
- ``cloud_provider`` / ``cloud_region`` / ``cloud_zone`` /
``cloud_instance_type`` / ``cloud_instance_id`` — DMI sysfs
identifies the cloud provider from BIOS/SMBIOS strings (no
network call) and only THEN does the IMDS probe fire. Baremetal
hosts pay zero startup latency on the cloud path.
Hardening highlights:
- IMDS probes target the link-local IP literal ``169.254.169.254``
for AWS, GCP, AND Azure — no DNS-resolvable hostname for any
vendor, so a host with attacker-controlled DNS can't redirect
the probe even when its DMI claims a cloud provider.
- Response bodies capped at 64 KiB on read; per-field strings
capped at 256 chars and stripped of control characters before
persistence. Stops a hostile IMDS responder from spraying
multi-megabyte / newline-injected payloads into ``node_metadata``
and from there into coord-LLM ``list_nodes`` context.
- ``isinstance(doc, dict)`` guards on every JSON IMDS response so
a non-conformant body (list / scalar / null) returns clean ``{}``
instead of raising.
- ``collect_node_info()`` runs via ``asyncio.to_thread`` from the
server's lifespan handler so the IMDS probe latency never blocks
the event loop.
- GCP fans the three zone/machine-type/id probes concurrently so a
misidentified host's worst case is one timeout window (~1 s)
instead of three (~3 s).
- Operator opt-out via ``TURNSTONE_AUTO_CLOUD_METADATA=0`` skips the
IMDS phase entirely; the DMI-derived ``cloud_provider`` still
populates because that's a kernel interface.
Tests: 4807 pass (+11 net, 73 in test_node_info.py). Ruff + mypy
clean on every modified file. New tests pin the heterogeneous-GPU
flat-key fix, the IMDS hardening (non-dict JSON, control-char
sanitisation, body cap, per-field cap), and the GCP IP-literal
property.
This commit is contained in:
Patrick Buckley
@@ -5,8 +5,20 @@ from __future__ import annotations
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core import node_info
|
||||
from turnstone.core.node_info import (
|
||||
_collect_interfaces,
|
||||
_detect_aws_metadata,
|
||||
_detect_azure_metadata,
|
||||
_detect_cloud_metadata,
|
||||
_detect_cloud_provider_from_dmi,
|
||||
_detect_cpu_model,
|
||||
_detect_gcp_metadata,
|
||||
_detect_gpus,
|
||||
_detect_memory_gb,
|
||||
_imds_field,
|
||||
_is_loopback_or_link_local,
|
||||
collect_node_info,
|
||||
)
|
||||
@@ -135,3 +147,650 @@ class TestIsLoopbackOrLinkLocal:
|
||||
assert _is_loopback_or_link_local("10.0.0.5") is False
|
||||
assert _is_loopback_or_link_local("192.168.1.1") is False
|
||||
assert _is_loopback_or_link_local("2001:db8::1") is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel-interface helpers — capability detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _seed_drm_layout(tmp_path, cards):
|
||||
"""Build a fake ``/sys/class/drm`` layout under ``tmp_path``.
|
||||
|
||||
``cards`` is a list of ``(name, vendor_id, device_id)`` tuples.
|
||||
Use ``vendor_id=None`` to skip writing the vendor file (simulates
|
||||
a permission/missing-attr failure that the detector must skip
|
||||
cleanly). Returns the DRM root path.
|
||||
"""
|
||||
drm = tmp_path / "drm"
|
||||
drm.mkdir()
|
||||
for name, vendor_id, device_id in cards:
|
||||
device_dir = drm / name / "device"
|
||||
device_dir.mkdir(parents=True)
|
||||
if vendor_id is not None:
|
||||
(device_dir / "vendor").write_text(vendor_id + "\n")
|
||||
if device_id is not None:
|
||||
(device_dir / "device").write_text(device_id + "\n")
|
||||
return str(drm)
|
||||
|
||||
|
||||
class TestDetectGPUs:
|
||||
"""Sysfs-DRM enumeration — vendor-agnostic, no userspace binary."""
|
||||
|
||||
def test_returns_empty_when_drm_dir_missing(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", "/nonexistent/path/that/should/not/exist")
|
||||
assert _detect_gpus() == []
|
||||
|
||||
def test_returns_empty_when_no_card_dirs(self, tmp_path, monkeypatch):
|
||||
# Empty /sys/class/drm — no GPUs registered.
|
||||
drm = tmp_path / "drm"
|
||||
drm.mkdir()
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", str(drm))
|
||||
assert _detect_gpus() == []
|
||||
|
||||
def test_detects_nvidia_gpu(self, tmp_path, monkeypatch):
|
||||
drm_dir = _seed_drm_layout(tmp_path, [("card0", "0x10de", "0x2330")])
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", drm_dir)
|
||||
gpus = _detect_gpus()
|
||||
assert len(gpus) == 1
|
||||
assert gpus[0] == {
|
||||
"index": "0",
|
||||
"vendor": "nvidia",
|
||||
"pci_vendor": "0x10de",
|
||||
"pci_device": "0x2330",
|
||||
}
|
||||
|
||||
def test_detects_amd_gpu(self, tmp_path, monkeypatch):
|
||||
drm_dir = _seed_drm_layout(tmp_path, [("card0", "0x1002", "0x74a1")])
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", drm_dir)
|
||||
gpus = _detect_gpus()
|
||||
assert len(gpus) == 1
|
||||
assert gpus[0]["vendor"] == "amd"
|
||||
|
||||
def test_detects_intel_gpu(self, tmp_path, monkeypatch):
|
||||
drm_dir = _seed_drm_layout(tmp_path, [("card0", "0x8086", "0x56a0")])
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", drm_dir)
|
||||
gpus = _detect_gpus()
|
||||
assert gpus[0]["vendor"] == "intel"
|
||||
|
||||
def test_unknown_vendor_id_surfaces_as_unknown(self, tmp_path, monkeypatch):
|
||||
"""Vendor IDs not in our friendly-name table land as 'unknown'
|
||||
but the raw IDs are still surfaced — operators can map them
|
||||
out of band."""
|
||||
drm_dir = _seed_drm_layout(tmp_path, [("card0", "0xdead", "0xbeef")])
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", drm_dir)
|
||||
gpus = _detect_gpus()
|
||||
assert gpus[0]["vendor"] == "unknown"
|
||||
assert gpus[0]["pci_vendor"] == "0xdead"
|
||||
assert gpus[0]["pci_device"] == "0xbeef"
|
||||
|
||||
def test_skips_render_nodes(self, tmp_path, monkeypatch):
|
||||
"""``renderD*`` nodes are per-card render-only interfaces that
|
||||
share the same physical device as a ``cardN`` entry; counting
|
||||
them would double the GPU count. The card-name regex
|
||||
excludes them."""
|
||||
drm = tmp_path / "drm"
|
||||
drm.mkdir()
|
||||
for name in ("card0", "renderD128"):
|
||||
device = drm / name / "device"
|
||||
device.mkdir(parents=True)
|
||||
(device / "vendor").write_text("0x10de")
|
||||
(device / "device").write_text("0x2330")
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", str(drm))
|
||||
gpus = _detect_gpus()
|
||||
assert len(gpus) == 1 # only card0, not renderD128
|
||||
|
||||
def test_multi_gpu_node(self, tmp_path, monkeypatch):
|
||||
drm_dir = _seed_drm_layout(
|
||||
tmp_path,
|
||||
[
|
||||
("card0", "0x10de", "0x2330"),
|
||||
("card1", "0x10de", "0x2330"),
|
||||
("card2", "0x10de", "0x2330"),
|
||||
("card3", "0x10de", "0x2330"),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", drm_dir)
|
||||
gpus = _detect_gpus()
|
||||
assert len(gpus) == 4
|
||||
assert [g["index"] for g in gpus] == ["0", "1", "2", "3"]
|
||||
|
||||
def test_card_with_missing_vendor_is_skipped(self, tmp_path, monkeypatch):
|
||||
"""A card whose vendor file can't be read (permissions /
|
||||
partial sysfs) is silently skipped — the rest of the
|
||||
enumeration must still complete."""
|
||||
drm = tmp_path / "drm"
|
||||
drm.mkdir()
|
||||
# card0 has no vendor file; card1 is well-formed.
|
||||
(drm / "card0" / "device").mkdir(parents=True)
|
||||
good = drm / "card1" / "device"
|
||||
good.mkdir(parents=True)
|
||||
(good / "vendor").write_text("0x10de")
|
||||
(good / "device").write_text("0x2330")
|
||||
monkeypatch.setattr(node_info, "_DRM_DIR", str(drm))
|
||||
gpus = _detect_gpus()
|
||||
assert len(gpus) == 1
|
||||
assert gpus[0]["index"] == "1"
|
||||
|
||||
|
||||
class TestDetectMemoryGB:
|
||||
def test_parses_meminfo(self, tmp_path, monkeypatch):
|
||||
meminfo = tmp_path / "meminfo"
|
||||
# 32 GiB = 32 * 1024 * 1024 KiB = 33554432 KiB
|
||||
meminfo.write_text(
|
||||
"MemTotal: 33554432 kB\n"
|
||||
"MemFree: 5000000 kB\n"
|
||||
"MemAvailable: 28000000 kB\n"
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_MEMINFO_PATH", str(meminfo))
|
||||
assert _detect_memory_gb() == 32
|
||||
|
||||
def test_rounds_down(self, tmp_path, monkeypatch):
|
||||
"""31.5 GiB worth of KiB rounds down to 31 — operators that
|
||||
write ``filters={"memory_gb": 32}`` shouldn't match a node
|
||||
that's actually 31.5."""
|
||||
meminfo = tmp_path / "meminfo"
|
||||
# 31.5 GiB = 31.5 * 1024 * 1024 = 33030144 KiB
|
||||
meminfo.write_text(f"MemTotal: {31 * 1024 * 1024 + 512 * 1024} kB\n")
|
||||
monkeypatch.setattr(node_info, "_MEMINFO_PATH", str(meminfo))
|
||||
assert _detect_memory_gb() == 31
|
||||
|
||||
def test_returns_none_when_meminfo_missing(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_MEMINFO_PATH", "/nonexistent/meminfo")
|
||||
assert _detect_memory_gb() is None
|
||||
|
||||
def test_returns_none_when_no_memtotal_line(self, tmp_path, monkeypatch):
|
||||
meminfo = tmp_path / "meminfo"
|
||||
meminfo.write_text("MemFree: 5000000 kB\n") # no MemTotal
|
||||
monkeypatch.setattr(node_info, "_MEMINFO_PATH", str(meminfo))
|
||||
assert _detect_memory_gb() is None
|
||||
|
||||
|
||||
class TestDetectCPUModel:
|
||||
def test_parses_intel_brand(self, tmp_path, monkeypatch):
|
||||
cpuinfo = tmp_path / "cpuinfo"
|
||||
cpuinfo.write_text(
|
||||
"processor\t: 0\n"
|
||||
"model name\t: Intel(R) Xeon(R) Platinum 8488C\n"
|
||||
"cpu MHz\t\t: 2400.000\n"
|
||||
"processor\t: 1\n"
|
||||
"model name\t: Intel(R) Xeon(R) Platinum 8488C\n"
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_CPUINFO_PATH", str(cpuinfo))
|
||||
assert _detect_cpu_model() == "Intel(R) Xeon(R) Platinum 8488C"
|
||||
|
||||
def test_parses_amd_brand(self, tmp_path, monkeypatch):
|
||||
cpuinfo = tmp_path / "cpuinfo"
|
||||
cpuinfo.write_text("model name\t: AMD EPYC 9654 96-Core Processor\n")
|
||||
monkeypatch.setattr(node_info, "_CPUINFO_PATH", str(cpuinfo))
|
||||
assert _detect_cpu_model() == "AMD EPYC 9654 96-Core Processor"
|
||||
|
||||
def test_returns_none_on_arm_with_no_model_name(self, tmp_path, monkeypatch):
|
||||
"""ARM cpuinfo uses ``Hardware`` / ``Processor`` instead of
|
||||
``model name``; we return None and operators set ``cpu_model``
|
||||
in [metadata] config to taste."""
|
||||
cpuinfo = tmp_path / "cpuinfo"
|
||||
cpuinfo.write_text("Hardware\t: Apple M1\nProcessor\t: ARMv8\n")
|
||||
monkeypatch.setattr(node_info, "_CPUINFO_PATH", str(cpuinfo))
|
||||
assert _detect_cpu_model() is None
|
||||
|
||||
def test_returns_none_when_cpuinfo_missing(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_CPUINFO_PATH", "/nonexistent/cpuinfo")
|
||||
assert _detect_cpu_model() is None
|
||||
|
||||
|
||||
def _seed_dmi_layout(tmp_path, fields):
|
||||
"""Build a fake /sys/class/dmi/id with given key→value text files."""
|
||||
dmi = tmp_path / "dmi"
|
||||
dmi.mkdir()
|
||||
for key, value in fields.items():
|
||||
(dmi / key).write_text(value + "\n")
|
||||
return str(dmi)
|
||||
|
||||
|
||||
class TestDetectCloudProviderFromDMI:
|
||||
"""DMI-based cloud-provider detection — pure kernel interface."""
|
||||
|
||||
def test_aws_via_sys_vendor(self, tmp_path, monkeypatch):
|
||||
dmi = _seed_dmi_layout(tmp_path, {"sys_vendor": "Amazon EC2"})
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "aws"
|
||||
|
||||
def test_aws_via_bios_vendor(self, tmp_path, monkeypatch):
|
||||
"""Older Nitro instances set bios_vendor instead of sys_vendor."""
|
||||
dmi = _seed_dmi_layout(
|
||||
tmp_path,
|
||||
{"sys_vendor": "Xen", "bios_vendor": "Amazon EC2"},
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "aws"
|
||||
|
||||
def test_gcp_via_sys_vendor(self, tmp_path, monkeypatch):
|
||||
dmi = _seed_dmi_layout(
|
||||
tmp_path,
|
||||
{"sys_vendor": "Google", "product_name": "Google Compute Engine"},
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "gcp"
|
||||
|
||||
def test_azure_via_chassis_asset_tag(self, tmp_path, monkeypatch):
|
||||
"""The chassis_asset_tag prefix distinguishes Azure VMs from
|
||||
plain Microsoft Hyper-V on baremetal — same sys_vendor, but
|
||||
only Azure VMs carry the well-known asset tag."""
|
||||
dmi = _seed_dmi_layout(
|
||||
tmp_path,
|
||||
{
|
||||
"sys_vendor": "Microsoft Corporation",
|
||||
"chassis_asset_tag": "7783-7084-3265-9085-8269-3286-77",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "azure"
|
||||
|
||||
def test_microsoft_without_azure_tag_is_unknown(self, tmp_path, monkeypatch):
|
||||
"""Plain Hyper-V on baremetal — Microsoft sys_vendor but no
|
||||
Azure asset tag. Must not auto-detect as azure."""
|
||||
dmi = _seed_dmi_layout(
|
||||
tmp_path,
|
||||
{
|
||||
"sys_vendor": "Microsoft Corporation",
|
||||
"chassis_asset_tag": "Default string",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "unknown"
|
||||
|
||||
def test_baremetal_is_unknown(self, tmp_path, monkeypatch):
|
||||
dmi = _seed_dmi_layout(tmp_path, {"sys_vendor": "Dell Inc.", "bios_vendor": "Dell Inc."})
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", dmi)
|
||||
assert _detect_cloud_provider_from_dmi() == "unknown"
|
||||
|
||||
def test_missing_dmi_dir_is_unknown(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_DMI_DIR", "/nonexistent/dmi")
|
||||
assert _detect_cloud_provider_from_dmi() == "unknown"
|
||||
|
||||
|
||||
class TestIMDSDetectors:
|
||||
"""Vendor-specific IMDS parsers — exercise the body-shape parsing
|
||||
without making real network calls."""
|
||||
|
||||
def test_aws_imds_v2_token_failure(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: None)
|
||||
assert _detect_aws_metadata() == {}
|
||||
|
||||
def test_aws_imds_parses_identity_doc(self, monkeypatch):
|
||||
responses = iter(
|
||||
[
|
||||
"TOKEN-ABCD", # PUT /api/token
|
||||
json.dumps(
|
||||
{
|
||||
"region": "us-east-1",
|
||||
"availabilityZone": "us-east-1a",
|
||||
"instanceType": "p5.48xlarge",
|
||||
"instanceId": "i-0123456789abcdef0",
|
||||
}
|
||||
), # GET /dynamic/instance-identity/document
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: next(responses))
|
||||
result = _detect_aws_metadata()
|
||||
assert result == {
|
||||
"cloud_region": "us-east-1",
|
||||
"cloud_zone": "us-east-1a",
|
||||
"cloud_instance_type": "p5.48xlarge",
|
||||
"cloud_instance_id": "i-0123456789abcdef0",
|
||||
}
|
||||
|
||||
def test_aws_malformed_identity_doc_returns_empty(self, monkeypatch):
|
||||
responses = iter(["TOKEN-ABCD", "not-json"])
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: next(responses))
|
||||
assert _detect_aws_metadata() == {}
|
||||
|
||||
def test_gcp_zone_parsing(self, monkeypatch):
|
||||
# GCP returns paths like "projects/12345/zones/us-east1-a";
|
||||
# we surface the tail and derive region by chopping the
|
||||
# trailing "-a" letter.
|
||||
responses = {
|
||||
"zone": "projects/12345/zones/us-east1-a",
|
||||
"machine-type": "projects/12345/machineTypes/n1-standard-4",
|
||||
"id": "9876543210",
|
||||
}
|
||||
|
||||
def fake(url, headers=None, **_kw):
|
||||
for key, body in responses.items():
|
||||
if url.endswith("/" + key):
|
||||
return body
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(node_info, "_imds_get", fake)
|
||||
result = _detect_gcp_metadata()
|
||||
assert result["cloud_zone"] == "us-east1-a"
|
||||
assert result["cloud_region"] == "us-east1"
|
||||
assert result["cloud_instance_type"] == "n1-standard-4"
|
||||
assert result["cloud_instance_id"] == "9876543210"
|
||||
|
||||
def test_gcp_no_zone_returns_empty(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: None)
|
||||
assert _detect_gcp_metadata() == {}
|
||||
|
||||
def test_azure_compute_block_parsing(self, monkeypatch):
|
||||
body = json.dumps(
|
||||
{
|
||||
"compute": {
|
||||
"location": "eastus",
|
||||
"zone": "1",
|
||||
"vmSize": "Standard_NC24ads_A100_v4",
|
||||
"vmId": "abcd1234-...",
|
||||
}
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: body)
|
||||
result = _detect_azure_metadata()
|
||||
assert result == {
|
||||
"cloud_region": "eastus",
|
||||
"cloud_zone": "1",
|
||||
"cloud_instance_type": "Standard_NC24ads_A100_v4",
|
||||
"cloud_instance_id": "abcd1234-...",
|
||||
}
|
||||
|
||||
def test_azure_missing_compute_block_returns_empty(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: json.dumps({}))
|
||||
assert _detect_azure_metadata() == {}
|
||||
|
||||
|
||||
class TestDetectCloudMetadata:
|
||||
"""End-to-end cloud metadata detection: DMI gate + IMDS probe."""
|
||||
|
||||
def test_baremetal_skips_imds(self, monkeypatch):
|
||||
"""No DMI cloud signal → no IMDS probe → empty result, no
|
||||
startup latency cost. This is the property we wanted from
|
||||
the kernel-interface refactor."""
|
||||
called = {"imds": 0}
|
||||
|
||||
def _spy(*args, **kwargs):
|
||||
called["imds"] += 1
|
||||
return "should-never-be-called"
|
||||
|
||||
monkeypatch.setattr(node_info, "_detect_cloud_provider_from_dmi", lambda: "unknown")
|
||||
monkeypatch.setattr(node_info, "_imds_get", _spy)
|
||||
assert _detect_cloud_metadata() == {}
|
||||
assert called["imds"] == 0
|
||||
|
||||
def test_aws_detection_path(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_detect_cloud_provider_from_dmi", lambda: "aws")
|
||||
monkeypatch.setattr(
|
||||
node_info,
|
||||
"_detect_aws_metadata",
|
||||
lambda: {"cloud_region": "us-west-2", "cloud_instance_type": "p4d.24xlarge"},
|
||||
)
|
||||
result = _detect_cloud_metadata()
|
||||
assert result["cloud_provider"] == "aws"
|
||||
assert result["cloud_region"] == "us-west-2"
|
||||
assert result["cloud_instance_type"] == "p4d.24xlarge"
|
||||
|
||||
def test_imds_probe_failure_still_surfaces_provider(self, monkeypatch):
|
||||
"""If DMI says we're on AWS but IMDS times out, we still
|
||||
surface ``cloud_provider=aws`` from DMI alone. Operators
|
||||
can route on provider even when region/instance-type
|
||||
couldn't be probed."""
|
||||
monkeypatch.setattr(node_info, "_detect_cloud_provider_from_dmi", lambda: "aws")
|
||||
monkeypatch.setattr(node_info, "_detect_aws_metadata", lambda: {})
|
||||
result = _detect_cloud_metadata()
|
||||
assert result == {"cloud_provider": "aws"}
|
||||
|
||||
def test_opt_out_skips_imds_but_keeps_provider(self, monkeypatch):
|
||||
"""``TURNSTONE_AUTO_CLOUD_METADATA=0`` skips the network probe
|
||||
entirely. ``cloud_provider`` from DMI still populates because
|
||||
it's a kernel interface, not a network call."""
|
||||
monkeypatch.setenv("TURNSTONE_AUTO_CLOUD_METADATA", "0")
|
||||
monkeypatch.setattr(node_info, "_detect_cloud_provider_from_dmi", lambda: "gcp")
|
||||
|
||||
def _imds_should_not_run(*a, **kw):
|
||||
pytest.fail("IMDS probe must not run when TURNSTONE_AUTO_CLOUD_METADATA=0")
|
||||
|
||||
monkeypatch.setattr(node_info, "_imds_get", _imds_should_not_run)
|
||||
result = _detect_cloud_metadata()
|
||||
assert result == {"cloud_provider": "gcp"}
|
||||
|
||||
def test_imds_exception_does_not_propagate(self, monkeypatch):
|
||||
"""A buggy IMDS parser (raises unexpectedly) must not crash
|
||||
the collector — the ``except Exception`` wrapper inside
|
||||
``_detect_cloud_metadata`` swallows and logs."""
|
||||
monkeypatch.setattr(node_info, "_detect_cloud_provider_from_dmi", lambda: "azure")
|
||||
|
||||
def _boom():
|
||||
raise RuntimeError("simulated parser bug")
|
||||
|
||||
monkeypatch.setattr(node_info, "_detect_azure_metadata", _boom)
|
||||
result = _detect_cloud_metadata()
|
||||
# cloud_provider survives; region/zone are missing.
|
||||
assert result == {"cloud_provider": "azure"}
|
||||
|
||||
|
||||
class TestCollectNodeInfoCapabilityIntegration:
|
||||
"""End-to-end checks on the public ``collect_node_info`` entry
|
||||
point — confirms the new kernel-interface helpers wire up
|
||||
correctly and that one helper failing doesn't suppress the others."""
|
||||
|
||||
def test_gpu_keys_appear_when_gpus_detected(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
node_info,
|
||||
"_detect_gpus",
|
||||
lambda: [
|
||||
{"index": "0", "vendor": "nvidia", "pci_vendor": "0x10de", "pci_device": "0x2330"},
|
||||
],
|
||||
)
|
||||
info = collect_node_info()
|
||||
assert info["gpu_count"] == 1
|
||||
assert info["gpu_vendor"] == "nvidia"
|
||||
assert info["gpu_vendors"] == ["nvidia"]
|
||||
assert info["gpus"][0]["pci_device"] == "0x2330"
|
||||
|
||||
def test_gpu_keys_absent_when_no_gpus(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_detect_gpus", lambda: [])
|
||||
info = collect_node_info()
|
||||
for k in ("gpu_count", "gpu_vendor", "gpu_vendors", "gpus"):
|
||||
assert k not in info
|
||||
|
||||
def test_unknown_only_gpus_omits_vendor_keys(self, monkeypatch):
|
||||
"""A node where every detected GPU has unknown vendor still
|
||||
gets gpu_count + gpus, but the flat gpu_vendor / gpu_vendors
|
||||
keys are skipped (filtering on 'unknown' isn't useful)."""
|
||||
monkeypatch.setattr(
|
||||
node_info,
|
||||
"_detect_gpus",
|
||||
lambda: [
|
||||
{"index": "0", "vendor": "unknown", "pci_vendor": "0xdead", "pci_device": "0xbeef"},
|
||||
],
|
||||
)
|
||||
info = collect_node_info()
|
||||
assert info["gpu_count"] == 1
|
||||
assert "gpu_vendor" not in info
|
||||
assert "gpu_vendors" not in info
|
||||
|
||||
def test_memory_key_appears(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_detect_memory_gb", lambda: 256)
|
||||
info = collect_node_info()
|
||||
assert info["memory_gb"] == 256
|
||||
|
||||
def test_memory_zero_omitted(self, monkeypatch):
|
||||
"""A reading of 0 GiB is degenerate — likely a parse error
|
||||
rather than a real zero-RAM machine. Skip the key rather
|
||||
than advertise a false value."""
|
||||
monkeypatch.setattr(node_info, "_detect_memory_gb", lambda: 0)
|
||||
info = collect_node_info()
|
||||
assert "memory_gb" not in info
|
||||
|
||||
def test_cpu_model_key_appears(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_detect_cpu_model", lambda: "AMD EPYC 9654")
|
||||
info = collect_node_info()
|
||||
assert info["cpu_model"] == "AMD EPYC 9654"
|
||||
|
||||
def test_cloud_keys_merged(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
node_info,
|
||||
"_detect_cloud_metadata",
|
||||
lambda: {
|
||||
"cloud_provider": "aws",
|
||||
"cloud_region": "us-east-1",
|
||||
"cloud_instance_type": "p5.48xlarge",
|
||||
},
|
||||
)
|
||||
info = collect_node_info()
|
||||
assert info["cloud_provider"] == "aws"
|
||||
assert info["cloud_region"] == "us-east-1"
|
||||
assert info["cloud_instance_type"] == "p5.48xlarge"
|
||||
|
||||
def test_one_capability_failure_does_not_block_others(self, monkeypatch):
|
||||
"""If GPU detection raises, memory + cpu + cloud detection
|
||||
must still run. Mirrors the existing per-field-failsafe
|
||||
contract on the basic fields."""
|
||||
|
||||
def _boom():
|
||||
raise RuntimeError("simulated DRM failure")
|
||||
|
||||
monkeypatch.setattr(node_info, "_detect_gpus", _boom)
|
||||
monkeypatch.setattr(node_info, "_detect_memory_gb", lambda: 64)
|
||||
monkeypatch.setattr(node_info, "_detect_cpu_model", lambda: "AMD EPYC 9654")
|
||||
info = collect_node_info()
|
||||
assert "gpu_count" not in info
|
||||
assert info["memory_gb"] == 64
|
||||
assert info["cpu_model"] == "AMD EPYC 9654"
|
||||
|
||||
def test_heterogeneous_gpu_first_unknown_uses_first_known_vendor(self, monkeypatch):
|
||||
"""Regression guard for the fix where ``gpu_vendor`` was set
|
||||
from ``gpus[0]["vendor"]`` even when that was ``"unknown"``,
|
||||
leaving the flat key unfilterable on a node that actually
|
||||
has known-vendor GPUs. The fix routes through the sorted
|
||||
unique known-vendor set so the flat key reflects something
|
||||
the operator can actually filter on."""
|
||||
monkeypatch.setattr(
|
||||
node_info,
|
||||
"_detect_gpus",
|
||||
lambda: [
|
||||
# First card: unrecognized vendor (e.g. an exotic
|
||||
# accelerator) — surfaces vendor=unknown but the
|
||||
# node is otherwise nvidia.
|
||||
{
|
||||
"index": "0",
|
||||
"vendor": "unknown",
|
||||
"pci_vendor": "0xdead",
|
||||
"pci_device": "0xbeef",
|
||||
},
|
||||
{
|
||||
"index": "1",
|
||||
"vendor": "nvidia",
|
||||
"pci_vendor": "0x10de",
|
||||
"pci_device": "0x2330",
|
||||
},
|
||||
],
|
||||
)
|
||||
info = collect_node_info()
|
||||
# Flat key must reflect a vendor an operator can route on.
|
||||
assert info["gpu_vendor"] == "nvidia"
|
||||
assert info["gpu_vendors"] == ["nvidia"]
|
||||
# Per-card detail still shows the unknown card so an
|
||||
# operator can investigate.
|
||||
assert info["gpus"][0]["vendor"] == "unknown"
|
||||
assert info["gpu_count"] == 2
|
||||
|
||||
|
||||
class TestIMDSFieldSanitiser:
|
||||
"""``_imds_field`` strips control chars + length-caps each
|
||||
persisted value. Defense-in-depth against an attacker-controlled
|
||||
IMDS responder injecting prompt-payload bytes into coord LLM
|
||||
context via ``list_nodes``."""
|
||||
|
||||
def test_passes_clean_string_through(self):
|
||||
assert _imds_field("us-east-1") == "us-east-1"
|
||||
|
||||
def test_strips_control_characters(self):
|
||||
# Newline + NUL would otherwise survive into list_nodes
|
||||
# output and could break parsing or inject content into
|
||||
# downstream renderers.
|
||||
out = _imds_field("us-east-1\n\x00 injected")
|
||||
assert "\n" not in (out or "")
|
||||
assert "\x00" not in (out or "")
|
||||
assert out == "us-east-1 injected"
|
||||
|
||||
def test_caps_length(self):
|
||||
from turnstone.core.node_info import _IMDS_MAX_FIELD_CHARS
|
||||
|
||||
out = _imds_field("X" * (_IMDS_MAX_FIELD_CHARS * 4))
|
||||
assert out is not None
|
||||
assert len(out) == _IMDS_MAX_FIELD_CHARS
|
||||
|
||||
def test_returns_none_for_non_string(self):
|
||||
assert _imds_field(None) is None
|
||||
assert _imds_field(42) is None
|
||||
assert _imds_field(["us-east-1"]) is None
|
||||
|
||||
def test_returns_none_for_empty_or_whitespace(self):
|
||||
assert _imds_field("") is None
|
||||
assert _imds_field(" ") is None
|
||||
|
||||
|
||||
class TestIMDSResponseHardening:
|
||||
"""Regression guards on the AWS / Azure non-dict-JSON paths and
|
||||
the GCP hostname → IP-literal switch."""
|
||||
|
||||
def test_aws_handles_non_dict_json_without_raising(self, monkeypatch):
|
||||
"""If a hostile/misbehaving IMDS returns a JSON list rather
|
||||
than the documented identity-document object, the previous
|
||||
shape would AttributeError on ``doc.get(src)``. The
|
||||
``isinstance(doc, dict)`` guard makes this a clean miss."""
|
||||
responses = iter(["TOKEN-ABCD", "[1, 2, 3]"])
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: next(responses))
|
||||
# Must not raise.
|
||||
assert _detect_aws_metadata() == {}
|
||||
|
||||
def test_aws_handles_scalar_json_without_raising(self, monkeypatch):
|
||||
responses = iter(["TOKEN-ABCD", "42"])
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: next(responses))
|
||||
assert _detect_aws_metadata() == {}
|
||||
|
||||
def test_azure_handles_non_dict_json_without_raising(self, monkeypatch):
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: '["not-an-object"]')
|
||||
# Must not raise.
|
||||
assert _detect_azure_metadata() == {}
|
||||
|
||||
def test_gcp_uses_link_local_ip_literal(self, monkeypatch):
|
||||
"""The GCP probe must target ``169.254.169.254`` directly so
|
||||
a host with attacker-controlled DNS can't redirect the probe
|
||||
via ``metadata.google.internal``. Pin the URL prefix."""
|
||||
called_urls: list[str] = []
|
||||
|
||||
def _spy(url, *args, **kwargs):
|
||||
called_urls.append(url)
|
||||
return None # all probes fail; that's fine — we're inspecting URLs
|
||||
|
||||
monkeypatch.setattr(node_info, "_imds_get", _spy)
|
||||
_detect_gcp_metadata()
|
||||
assert called_urls, "GCP detector must issue at least one IMDS call"
|
||||
for url in called_urls:
|
||||
assert url.startswith("http://169.254.169.254/"), (
|
||||
f"GCP probe leaked through DNS-resolvable hostname: {url}"
|
||||
)
|
||||
|
||||
def test_imds_field_sanitises_aws_response(self, monkeypatch):
|
||||
"""End-to-end: a hostile IMDS response body with a control
|
||||
character lands sanitised in the AWS detector's output."""
|
||||
responses = iter(
|
||||
[
|
||||
"TOKEN-ABCD",
|
||||
json.dumps(
|
||||
{
|
||||
"region": "us-east-1\nrm -rf", # control char injection
|
||||
"instanceType": "p5.48xlarge",
|
||||
}
|
||||
),
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(node_info, "_imds_get", lambda *a, **kw: next(responses))
|
||||
result = _detect_aws_metadata()
|
||||
assert "\n" not in result["cloud_region"]
|
||||
# Sanitiser preserves the leading meaningful prefix, drops
|
||||
# the control character. Trailing content survives stripped
|
||||
# of control chars.
|
||||
assert "us-east-1" in result["cloud_region"]
|
||||
assert "rm -rf" in result["cloud_region"] # text still there, just newline-free
|
||||
|
||||
+487
-1
@@ -1,11 +1,35 @@
|
||||
"""Collect auto-populated node metadata using stdlib only."""
|
||||
"""Collect auto-populated node metadata using stdlib + kernel interfaces.
|
||||
|
||||
Two collection layers:
|
||||
|
||||
- Always-available basics — ``hostname``, ``fqdn``, ``os``, ``arch``,
|
||||
``python``, ``cpu_count``, ``interfaces`` — pulled from
|
||||
``platform``/``socket``/``os`` and never block.
|
||||
- Capability detection from Linux kernel interfaces — DRM sysfs for
|
||||
GPUs, ``/proc/meminfo`` for RAM, ``/proc/cpuinfo`` for the CPU
|
||||
model, ``/sys/class/dmi/id/*`` for the cloud provider, plus an
|
||||
IMDS probe for cloud region/instance-type. No userspace binaries
|
||||
(``nvidia-smi`` / ``rocm-smi`` / ``lspci``) on PATH — kernel
|
||||
interfaces work the same way regardless of vendor and don't depend
|
||||
on which optional package the operator happened to install.
|
||||
|
||||
Operators can still override any auto-detected key via the
|
||||
``[metadata]`` section of ``config.toml`` (last-write-wins on the
|
||||
``(node_id, key)`` upsert in ``set_node_metadata_bulk``), so the
|
||||
auto-detection layer is strictly additive — operators get sensible
|
||||
defaults, custom deployments still get the final say.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import socket
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import Any
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
@@ -35,6 +59,422 @@ def _collect_interfaces() -> dict[str, list[str]]:
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel-interface helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _read_text(path: str) -> str | None:
|
||||
"""Read a small kernel-pseudofs file and return its stripped text.
|
||||
|
||||
Returns ``None`` on any OSError so callers can treat the
|
||||
"file/sysfs not present" path as a clean miss. Decoded as UTF-8
|
||||
with ``errors="replace"``: a stray non-UTF-8 byte in DMI strings
|
||||
becomes ``U+FFFD`` rather than raising, which is the right call
|
||||
for substring-matching against vendor strings — the original
|
||||
bytes don't need to round-trip.
|
||||
"""
|
||||
try:
|
||||
with open(path, encoding="utf-8", errors="replace") as fh:
|
||||
return fh.read().strip()
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
# DRM (Direct Rendering Manager) sysfs — every PCI GPU registers a
|
||||
# ``cardN`` directory here regardless of vendor (NVIDIA, AMD, Intel,
|
||||
# ARM Mali, etc.). Reading the underlying PCI device's ``vendor`` and
|
||||
# ``device`` files gives us vendor identification without depending on
|
||||
# any vendor-specific userspace binary being installed or on PATH.
|
||||
_DRM_DIR = "/sys/class/drm"
|
||||
_CARD_DIR_RE = re.compile(r"^card\d+$")
|
||||
|
||||
# PCI vendor IDs. Source: pcisig.com canonical list. We surface
|
||||
# friendly names for the four vendors that ship GPUs into AI
|
||||
# infrastructure today; everything else lands as ``unknown`` and the
|
||||
# raw vendor/device IDs are kept on the row so an operator can map
|
||||
# them out-of-band.
|
||||
_PCI_VENDOR_NAMES: dict[str, str] = {
|
||||
"0x10de": "nvidia",
|
||||
"0x1002": "amd",
|
||||
"0x8086": "intel",
|
||||
"0x106b": "apple",
|
||||
}
|
||||
|
||||
|
||||
def _detect_gpus() -> list[dict[str, str]]:
|
||||
"""Enumerate GPUs via the Linux DRM sysfs interface.
|
||||
|
||||
For each ``/sys/class/drm/cardN`` directory, read the underlying
|
||||
PCI device's ``vendor`` and ``device`` IDs. Returns a list of
|
||||
``{"index", "vendor", "pci_vendor", "pci_device"}`` dicts.
|
||||
|
||||
Vendor-agnostic by design — works for NVIDIA, AMD, Intel, and any
|
||||
future PCI GPU vendor without a vendor-specific tool on PATH.
|
||||
Trade-off: we surface PCI vendor/device IDs (e.g. ``"nvidia"`` +
|
||||
``"0x2330"``) not human-readable model names (``"NVIDIA H100"``);
|
||||
operators who need specific-model routing set ``gpu_model`` in
|
||||
``[metadata]`` config to override the auto layer.
|
||||
|
||||
Returns empty list on non-Linux, missing sysfs, or any read
|
||||
failure. Containers see whatever DRM nodes the host mapped in;
|
||||
a container with no GPU mapped returns empty cleanly.
|
||||
"""
|
||||
if not os.path.isdir(_DRM_DIR):
|
||||
return []
|
||||
try:
|
||||
entries = sorted(os.listdir(_DRM_DIR))
|
||||
except OSError:
|
||||
return []
|
||||
gpus: list[dict[str, str]] = []
|
||||
for name in entries:
|
||||
# Skip ``renderD*`` nodes — they're per-card render-only
|
||||
# interfaces that duplicate ``cardN`` for the same physical
|
||||
# device. Counting them would double the GPU count.
|
||||
if not _CARD_DIR_RE.match(name):
|
||||
continue
|
||||
device_dir = os.path.join(_DRM_DIR, name, "device")
|
||||
vendor_id = _read_text(os.path.join(device_dir, "vendor"))
|
||||
device_id = _read_text(os.path.join(device_dir, "device"))
|
||||
if not vendor_id or not device_id:
|
||||
continue
|
||||
gpus.append(
|
||||
{
|
||||
"index": name[4:], # strip "card" prefix
|
||||
"vendor": _PCI_VENDOR_NAMES.get(vendor_id, "unknown"),
|
||||
"pci_vendor": vendor_id,
|
||||
"pci_device": device_id,
|
||||
}
|
||||
)
|
||||
return gpus
|
||||
|
||||
|
||||
# ``/proc/meminfo`` MemTotal field is in KiB. Linux only — falls
|
||||
# through to None on Darwin/Windows/missing-procfs containers.
|
||||
_MEMINFO_PATH = "/proc/meminfo"
|
||||
|
||||
|
||||
def _detect_memory_gb() -> int | None:
|
||||
"""Read total memory from ``/proc/meminfo`` and return GiB.
|
||||
|
||||
Returns ``None`` on non-Linux or any read/parse failure. Rounds
|
||||
DOWN — ``mem_gb >= N`` is the canonical "this node has at least N
|
||||
GiB" filter shape, and a node with 31.5 GiB shouldn't claim to
|
||||
have 32 in case a downstream pin checks the exact value.
|
||||
"""
|
||||
text = _read_text(_MEMINFO_PATH)
|
||||
if text is None:
|
||||
return None
|
||||
for line in text.splitlines():
|
||||
if not line.startswith("MemTotal:"):
|
||||
continue
|
||||
parts = line.split()
|
||||
if len(parts) >= 2 and parts[1].isdigit():
|
||||
return int(parts[1]) // (1024 * 1024)
|
||||
return None
|
||||
|
||||
|
||||
# ``/proc/cpuinfo`` is per-CPU; the ``model name`` field repeats for
|
||||
# every logical CPU. Read the first occurrence.
|
||||
_CPUINFO_PATH = "/proc/cpuinfo"
|
||||
_CPU_MODEL_RE = re.compile(r"^model name\s*:\s*(.+)$", re.MULTILINE)
|
||||
|
||||
|
||||
def _detect_cpu_model() -> str | None:
|
||||
"""Read the CPU brand string from ``/proc/cpuinfo``.
|
||||
|
||||
Returns the first ``model name`` value (Intel: ``Xeon Platinum
|
||||
8488C``, AMD: ``EPYC 9654``, ARM: usually empty since ARM exposes
|
||||
``Hardware`` / ``Processor`` instead — those return None and
|
||||
operators set ``cpu_model`` in config to taste).
|
||||
"""
|
||||
text = _read_text(_CPUINFO_PATH)
|
||||
if text is None:
|
||||
return None
|
||||
m = _CPU_MODEL_RE.search(text)
|
||||
if not m:
|
||||
return None
|
||||
return m.group(1).strip() or None
|
||||
|
||||
|
||||
# DMI (Desktop Management Interface) sysfs — Linux's view of the
|
||||
# vendor strings the BIOS/SMBIOS reports. Cloud hypervisors set
|
||||
# distinctive values here, so the cloud-provider detection can run
|
||||
# entirely from kernel interfaces with no network probe.
|
||||
_DMI_DIR = "/sys/class/dmi/id"
|
||||
|
||||
|
||||
def _read_dmi(field: str) -> str:
|
||||
"""Return the named DMI field's value, lowercased + stripped.
|
||||
|
||||
DMI files are root-readable on most distros but world-readable on
|
||||
typical cloud images. On a hardened host where we can't read
|
||||
them, this returns empty string and cloud-provider detection
|
||||
falls back to "unknown" (which then suppresses the IMDS probe).
|
||||
"""
|
||||
text = _read_text(os.path.join(_DMI_DIR, field))
|
||||
if text is None:
|
||||
return ""
|
||||
return text.lower().strip()
|
||||
|
||||
|
||||
def _detect_cloud_provider_from_dmi() -> str:
|
||||
"""Identify the cloud provider from BIOS/SMBIOS strings.
|
||||
|
||||
Returns ``"aws"`` / ``"gcp"`` / ``"azure"`` / ``"unknown"``.
|
||||
Pure kernel interface — no network call. Used to gate the IMDS
|
||||
probe so non-cloud hosts don't pay startup latency on doomed
|
||||
link-local connections.
|
||||
"""
|
||||
sys_vendor = _read_dmi("sys_vendor")
|
||||
board_vendor = _read_dmi("board_vendor")
|
||||
bios_vendor = _read_dmi("bios_vendor")
|
||||
chassis_asset_tag = _read_dmi("chassis_asset_tag")
|
||||
# AWS EC2: SMBIOS reports "Amazon EC2". Older Nitro instances
|
||||
# leave bios_vendor=Amazon EC2 too.
|
||||
if "amazon ec2" in (sys_vendor, board_vendor, bios_vendor):
|
||||
return "aws"
|
||||
# GCP: sys_vendor is "Google" with product_name "Google Compute Engine".
|
||||
if sys_vendor == "google" or "google compute engine" in _read_dmi("product_name"):
|
||||
return "gcp"
|
||||
# Azure: sys_vendor "Microsoft Corporation" plus a stable
|
||||
# chassis_asset_tag of "7783-7084-3265-9085-8269-3286-77".
|
||||
# Microsoft uses the same sys_vendor for Hyper-V on baremetal;
|
||||
# the tag is what distinguishes Azure VMs.
|
||||
if "microsoft" in sys_vendor and chassis_asset_tag.startswith("7783-7084"):
|
||||
return "azure"
|
||||
return "unknown"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IMDS probes — cloud-only, gated by DMI detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Per-call timeout for IMDS probes. Cloud hosts respond in < 50 ms.
|
||||
_CLOUD_PROBE_TIMEOUT_S: float = 1.0
|
||||
|
||||
# Hard caps on IMDS data we'll persist. All three real cloud-provider
|
||||
# IMDS responses are well under these limits (AWS identity doc is ~1
|
||||
# KiB, GCP/Azure single-field responses are tens of bytes); the caps
|
||||
# exist so a host where the link-local responder is hostile (spoofed
|
||||
# DMI on baremetal, attacker-controlled DNS, lab tamper) can't spray
|
||||
# multi-megabyte payloads into ``node_metadata`` and from there into
|
||||
# coord-LLM context windows on the next ``list_nodes``.
|
||||
_IMDS_MAX_BODY_BYTES: int = 64 * 1024
|
||||
_IMDS_MAX_FIELD_CHARS: int = 256
|
||||
|
||||
|
||||
def _imds_get(
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
method: str = "GET",
|
||||
timeout: float = _CLOUD_PROBE_TIMEOUT_S,
|
||||
) -> str | None:
|
||||
"""Tiny wrapper around urllib for IMDS calls.
|
||||
|
||||
Returns the response body as a UTF-8 string on 2xx, ``None`` on any
|
||||
network / decode / non-2xx failure. Body size is capped at
|
||||
:data:`_IMDS_MAX_BODY_BYTES` so a hostile responder can't cause an
|
||||
unbounded read.
|
||||
"""
|
||||
try:
|
||||
req = urllib.request.Request(url, headers=headers or {}, method=method)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp: # noqa: S310 (link-local IMDS)
|
||||
body: bytes = resp.read(_IMDS_MAX_BODY_BYTES)
|
||||
return body.decode("utf-8")
|
||||
except (urllib.error.URLError, OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _imds_field(value: Any) -> str | None:
|
||||
"""Sanitise an IMDS field for persistence into ``node_metadata``.
|
||||
|
||||
- Returns ``None`` for non-string / empty values so callers can
|
||||
``if v: out[k] = v``-style filter cleanly.
|
||||
- Strips control characters (anything below U+0020 plus DEL) —
|
||||
a hostile IMDS could otherwise inject newlines / NULs into
|
||||
strings the coord LLM later inhales.
|
||||
- Hard-caps to :data:`_IMDS_MAX_FIELD_CHARS`.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
cleaned = "".join(ch for ch in value if ch >= " " and ch != "\x7f").strip()
|
||||
if not cleaned:
|
||||
return None
|
||||
if len(cleaned) > _IMDS_MAX_FIELD_CHARS:
|
||||
cleaned = cleaned[:_IMDS_MAX_FIELD_CHARS]
|
||||
return cleaned
|
||||
|
||||
|
||||
def _detect_aws_metadata() -> dict[str, str]:
|
||||
"""EC2 IMDSv2: token + identity document."""
|
||||
base = "http://169.254.169.254/latest"
|
||||
token = _imds_get(
|
||||
f"{base}/api/token",
|
||||
method="PUT",
|
||||
headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
|
||||
)
|
||||
if not token:
|
||||
return {}
|
||||
body = _imds_get(
|
||||
f"{base}/dynamic/instance-identity/document",
|
||||
headers={"X-aws-ec2-metadata-token": token.strip()},
|
||||
)
|
||||
if not body:
|
||||
return {}
|
||||
try:
|
||||
doc = json.loads(body)
|
||||
except (TypeError, ValueError):
|
||||
return {}
|
||||
# IMDS contract says this endpoint returns a JSON object — but a
|
||||
# spoofed responder can return any JSON. Guard so a list / scalar
|
||||
# / null doesn't AttributeError on .get below; the outer try/except
|
||||
# in ``_detect_cloud_metadata`` would mask the crash, but local
|
||||
# type-checking keeps the function safe in isolation.
|
||||
if not isinstance(doc, dict):
|
||||
return {}
|
||||
out: dict[str, str] = {}
|
||||
for src, dst in (
|
||||
("region", "cloud_region"),
|
||||
("availabilityZone", "cloud_zone"),
|
||||
("instanceType", "cloud_instance_type"),
|
||||
("instanceId", "cloud_instance_id"),
|
||||
):
|
||||
cleaned = _imds_field(doc.get(src))
|
||||
if cleaned:
|
||||
out[dst] = cleaned
|
||||
return out
|
||||
|
||||
|
||||
def _detect_gcp_metadata() -> dict[str, str]:
|
||||
"""GCP Compute Engine metadata: zone / machine-type / id.
|
||||
|
||||
Targets the link-local IP literal ``169.254.169.254`` (not the
|
||||
resolvable hostname ``metadata.google.internal``) so a host with
|
||||
spoofed DMI tags + attacker-controlled DNS can't redirect the
|
||||
probe to a hostile server. The ``Metadata-Flavor: Google`` header
|
||||
is what GCE's metadata server uses to confirm we're a legitimate
|
||||
caller, and AWS/Azure also target the same IP — using it for GCP
|
||||
keeps all three providers on the same trust model.
|
||||
|
||||
Issues the three sub-calls (zone, machine-type, id) concurrently
|
||||
so a misidentified host (DMI claims GCP, IMDS unreachable) takes
|
||||
one timeout window (~1 s) instead of three sequential ones.
|
||||
"""
|
||||
import concurrent.futures
|
||||
|
||||
base = "http://169.254.169.254/computeMetadata/v1/instance"
|
||||
headers = {"Metadata-Flavor": "Google"}
|
||||
|
||||
paths = ("zone", "machine-type", "id")
|
||||
with concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=len(paths),
|
||||
thread_name_prefix="gcp-imds",
|
||||
) as pool:
|
||||
futures = {p: pool.submit(_imds_get, f"{base}/{p}", headers=headers) for p in paths}
|
||||
results = {p: fut.result() for p, fut in futures.items()}
|
||||
|
||||
zone = results.get("zone")
|
||||
if zone is None:
|
||||
return {}
|
||||
out: dict[str, str] = {}
|
||||
# zone format: "projects/PROJECT_NUM/zones/us-east1-a" → take tail.
|
||||
zone_short = _imds_field(zone.rsplit("/", 1)[-1])
|
||||
if zone_short:
|
||||
out["cloud_zone"] = zone_short
|
||||
# GCP region = zone with the trailing letter chopped.
|
||||
if "-" in zone_short:
|
||||
region = _imds_field(zone_short.rsplit("-", 1)[0])
|
||||
if region:
|
||||
out["cloud_region"] = region
|
||||
machine_type = results.get("machine-type")
|
||||
if machine_type:
|
||||
cleaned = _imds_field(machine_type.rsplit("/", 1)[-1])
|
||||
if cleaned:
|
||||
out["cloud_instance_type"] = cleaned
|
||||
instance_id = results.get("id")
|
||||
if instance_id:
|
||||
cleaned = _imds_field(instance_id)
|
||||
if cleaned:
|
||||
out["cloud_instance_id"] = cleaned
|
||||
return out
|
||||
|
||||
|
||||
def _detect_azure_metadata() -> dict[str, str]:
|
||||
"""Azure VM IMDS: location / vmSize."""
|
||||
url = "http://169.254.169.254/metadata/instance?api-version=2021-12-13"
|
||||
body = _imds_get(url, headers={"Metadata": "true"})
|
||||
if not body:
|
||||
return {}
|
||||
try:
|
||||
doc = json.loads(body)
|
||||
except (TypeError, ValueError):
|
||||
return {}
|
||||
# Same isinstance guard as the AWS path — a non-dict body would
|
||||
# AttributeError on doc.get("compute") below.
|
||||
if not isinstance(doc, dict):
|
||||
return {}
|
||||
compute = doc.get("compute") or {}
|
||||
if not isinstance(compute, dict):
|
||||
return {}
|
||||
out: dict[str, str] = {}
|
||||
for src, dst in (
|
||||
("location", "cloud_region"),
|
||||
("zone", "cloud_zone"),
|
||||
("vmSize", "cloud_instance_type"),
|
||||
("vmId", "cloud_instance_id"),
|
||||
):
|
||||
cleaned = _imds_field(compute.get(src))
|
||||
if cleaned:
|
||||
out[dst] = cleaned
|
||||
return out
|
||||
|
||||
|
||||
def _detect_cloud_metadata() -> dict[str, str]:
|
||||
"""Surface cloud_provider + region/zone/instance-type.
|
||||
|
||||
Detection happens in two phases:
|
||||
|
||||
1. **DMI (kernel interface)** identifies the provider from
|
||||
BIOS/SMBIOS strings. No network call, no startup latency on
|
||||
baremetal hosts — ``unknown`` returns immediately.
|
||||
2. **IMDS (network)** runs only when DMI confirmed a cloud, so
|
||||
the link-local probe can't burn 1+ second on a host that has
|
||||
no IMDS at all.
|
||||
|
||||
Operators can opt out of the IMDS phase entirely via
|
||||
``TURNSTONE_AUTO_CLOUD_METADATA=0`` if their network policy
|
||||
forbids link-local probes; ``cloud_provider`` from DMI still
|
||||
populates.
|
||||
"""
|
||||
provider = _detect_cloud_provider_from_dmi()
|
||||
if provider == "unknown":
|
||||
return {}
|
||||
out: dict[str, str] = {"cloud_provider": provider}
|
||||
if os.environ.get("TURNSTONE_AUTO_CLOUD_METADATA", "1") == "0":
|
||||
return out
|
||||
# Inline dispatch (vs a module-level dict of function refs) so a
|
||||
# test monkeypatching ``_detect_aws_metadata`` actually substitutes
|
||||
# the function the dispatcher will call — a dict captured at import
|
||||
# time would still hold the original reference.
|
||||
try:
|
||||
if provider == "aws":
|
||||
out.update(_detect_aws_metadata())
|
||||
elif provider == "gcp":
|
||||
out.update(_detect_gcp_metadata())
|
||||
elif provider == "azure":
|
||||
out.update(_detect_azure_metadata())
|
||||
except Exception:
|
||||
log.debug("node_info: IMDS probe failed provider=%s", provider, exc_info=True)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public collector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def collect_node_info() -> dict[str, Any]:
|
||||
"""Collect auto-populated node metadata.
|
||||
|
||||
@@ -66,4 +506,50 @@ def collect_node_info() -> dict[str, Any]:
|
||||
except Exception:
|
||||
log.debug("node_info: failed to collect interfaces", exc_info=True)
|
||||
|
||||
# Capability detection — independent failsafe blocks so a missing
|
||||
# /sys/class/drm doesn't suppress memory detection, etc.
|
||||
try:
|
||||
gpus = _detect_gpus()
|
||||
if gpus:
|
||||
info["gpu_count"] = len(gpus)
|
||||
info["gpus"] = gpus
|
||||
# Surface ``gpu_vendor`` (a single known vendor) and
|
||||
# ``gpu_vendors`` (sorted unique known vendors) as flat
|
||||
# keys so the ``list_nodes(filters={"gpu_vendor":
|
||||
# "nvidia"})`` shape works without filter-on-list
|
||||
# gymnastics. Both keys consult only the KNOWN-vendor
|
||||
# set: a heterogeneous node where ``gpus[0]`` happens to
|
||||
# be a card we don't have a friendly name for would
|
||||
# otherwise leak ``gpu_vendor="unknown"`` while
|
||||
# ``gpu_vendors=["nvidia"]`` listed the actual hardware,
|
||||
# and the flat-key filter would mismatch a real schedule
|
||||
# target. The full per-card list (including the unknown
|
||||
# entries with their raw PCI IDs) lives in ``gpus``.
|
||||
vendors = sorted({g["vendor"] for g in gpus if g["vendor"] != "unknown"})
|
||||
if vendors:
|
||||
info["gpu_vendor"] = vendors[0]
|
||||
info["gpu_vendors"] = vendors
|
||||
except Exception:
|
||||
log.debug("node_info: GPU detection failed", exc_info=True)
|
||||
|
||||
try:
|
||||
mem_gb = _detect_memory_gb()
|
||||
if mem_gb is not None and mem_gb > 0:
|
||||
info["memory_gb"] = mem_gb
|
||||
except Exception:
|
||||
log.debug("node_info: memory detection failed", exc_info=True)
|
||||
|
||||
try:
|
||||
cpu_model = _detect_cpu_model()
|
||||
if cpu_model:
|
||||
info["cpu_model"] = cpu_model
|
||||
except Exception:
|
||||
log.debug("node_info: CPU model detection failed", exc_info=True)
|
||||
|
||||
try:
|
||||
cloud = _detect_cloud_metadata()
|
||||
info.update(cloud)
|
||||
except Exception:
|
||||
log.debug("node_info: cloud metadata detection failed", exc_info=True)
|
||||
|
||||
return info
|
||||
|
||||
+7
-2
@@ -3106,12 +3106,17 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
_svc_storage.register_service("server", _svc_node_id, _svc_url)
|
||||
log.info("server.service_registered", node_id=_svc_node_id, url=_svc_url)
|
||||
|
||||
# Collect and store node metadata (auto + config)
|
||||
# Collect and store node metadata (auto + config).
|
||||
# ``collect_node_info`` runs synchronous probes (sysfs reads,
|
||||
# /proc reads, IMDS HTTP requests). Off-load to a worker
|
||||
# thread so the IMDS path's worst-case latency (~1 s on a
|
||||
# misidentified-cloud host) doesn't block the event loop
|
||||
# during the rest of the lifespan startup work.
|
||||
try:
|
||||
from turnstone.core.config import load_config as _load_meta_config
|
||||
from turnstone.core.node_info import collect_node_info
|
||||
|
||||
_auto_info = collect_node_info()
|
||||
_auto_info = await asyncio.to_thread(collect_node_info)
|
||||
_meta_entries: list[tuple[str, str, str]] = [
|
||||
(k, json.dumps(v), "auto") for k, v in _auto_info.items()
|
||||
]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "list_nodes",
|
||||
"description": "List active cluster nodes with their metadata. By default only nodes with a fresh service-registry heartbeat (within 120s) are returned. Each node carries two metadata sources: auto-populated at startup (`arch`, `cpu_count`, `fqdn`, `hostname`, `os`, `os_release`, `python` — always present) and user-supplied via the console Nodes admin tab (e.g. `capability`, `region`, `tenant`, `role` — deployment-specific). Pass arbitrary `key=value` filters to narrow; all filters must match (AND). Pair with `target_node` on spawn_workstream to pin a child to a node that matches a capability. The 120s heartbeat is a sliding window, so a node returned here can drop out before a follow-up spawn lands — the race produces `\"No available node for routing\"`; omit `target_node` to let rendezvous pick from the still-healthy set, or retry after re-listing if a specific node is required. The `interfaces` key (container IPs, interface names) is stripped by default — routing should use capability/region tags, not IPs; pass `include_network_detail=true` only for debugging. Pass `include_inactive=true` to surface stale registrations (those nodes will reject `target_node` pinning).",
|
||||
"description": "List active cluster nodes with their metadata. By default only nodes with a fresh service-registry heartbeat (within 120s) are returned. Pass arbitrary `key=value` filters to narrow; all filters must match (AND). Pair with `target_node` on spawn_workstream to pin a child to a node that matches a capability. The 120s heartbeat is a sliding window, so a node returned here can drop out before a follow-up spawn lands — the race produces `\"No available node for routing\"`; omit `target_node` to let rendezvous pick from the still-healthy set, or retry after re-listing if a specific node is required. The `interfaces` key (container IPs, interface names) is stripped by default — routing should use capability/region tags, not IPs; pass `include_network_detail=true` only for debugging. Pass `include_inactive=true` to surface stale registrations (those nodes will reject `target_node` pinning).",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
Reference in New Issue
Block a user