fix: type edge application resources to match OEG schema (741fd5dd) · Commits · OOP / code / ai2

mcp_module/models/edge_application_models.py

+90 −8

Changes for mcp_module/models/edge_application_models.py: 90 added lines, 8 removed lines.
Original line number	Diff line number	Diff line
		from typing import Literal
		from typing import Literal, Optional

		from pydantic import BaseModel, Field, field_validator
		from pydantic import BaseModel, ConfigDict, Field, field_validator


		class AppRepoInput(BaseModel):
		@@ -49,6 +49,92 @@ class OperatingSystemInput(BaseModel):
		]


		# Resource contract mirrors OEG v0.1alpha1; only Kubernetes is supported.
		class CpuPoolTopology(BaseModel):
		model_config = ConfigDict(extra="forbid")

		minNumberOfNodes: int = Field(ge=1, le=1000)
		minNodeCpu: int = Field(ge=1, le=256)
		minNodeMemory: int = Field(ge=1, le=16384)


		class CpuPool(BaseModel):
		model_config = ConfigDict(extra="forbid")

		numCPU: int = Field(ge=1, le=256)
		memory: int = Field(ge=1, le=16384)
		topology: CpuPoolTopology


		class GpuPoolTopology(BaseModel):
		model_config = ConfigDict(extra="forbid")

		minNumberOfNodes: int = Field(ge=1, le=1000)
		minNodeCpu: int = Field(ge=1, le=256)
		minNodeMemory: int = Field(ge=1, le=16384)
		minNodeGpuMemory: int = Field(ge=1, le=16)


		class GpuPool(BaseModel):
		model_config = ConfigDict(extra="forbid")

		numCPU: int = Field(ge=1, le=1024)
		memory: int = Field(ge=1, le=16384)
		gpuMemory: int = Field(ge=1, le=16)
		topology: GpuPoolTopology


		class ApplicationResources(BaseModel):
		model_config = ConfigDict(extra="forbid")

		cpuPool: Optional[CpuPool] = None
		gpuPool: Optional[GpuPool] = None


		class K8sPrimaryNetwork(BaseModel):
		model_config = ConfigDict(extra="forbid")

		provider: Optional[str] = Field(default=None, max_length=64)
		version: Optional[str] = Field(default=None, max_length=64)


		class K8sAdditionalNetwork(BaseModel):
		model_config = ConfigDict(extra="forbid")

		name: Optional[str] = Field(default=None, max_length=64)
		interfaceType: Optional[Literal["netdevice", "vfio-pci", "interface"]] = None


		class K8sNetworking(BaseModel):
		model_config = ConfigDict(extra="forbid")

		primaryNetwork: K8sPrimaryNetwork
		additionalNetworks: Optional[list[K8sAdditionalNetwork]] = Field(default=None, max_length=100)


		K8sAddon = Literal["monitoring", "ingress"]


		class KubernetesResources(BaseModel):
		model_config = ConfigDict(extra="forbid")

		infraKind: Literal["kubernetes"]
		applicationResources: ApplicationResources
		isStandalone: bool
		additionalStorage: Optional[str] = Field(default=None, max_length=32, pattern=r"^\d+(GB\|MB)$")
		version: Optional[str] = Field(default=None, max_length=64)
		networking: Optional[K8sNetworking] = None
		addons: Optional[list[K8sAddon]] = Field(default=None, max_length=2)

		@field_validator("addons")
		@classmethod
		def _addons_unique(cls, v: Optional[list[str]]) -> Optional[list[str]]:
		if v is not None and len(set(v)) != len(v):
		raise ValueError("addons must be unique")
		return v



		class RegisterAppDefinitionInput(BaseModel):
		appId: str \| None = Field(
		default=None,
		@@ -72,13 +158,9 @@ class RegisterAppDefinitionInput(BaseModel):
		)
		operatingSystem: OperatingSystemInput \| None = None
		componentSpec: list[ComponentSpecInput] = Field(..., max_length=100)
		requiredResources: dict = Field(
		requiredResources: KubernetesResources = Field(
		...,
		description=(
		"OEG Kubernetes resources. Required fields: infraKind='kubernetes', "
		"applicationResources, and isStandalone. Optional fields: "
		"additionalStorage, version, networking, and addons."
		),
		description="Kubernetes resource requirements supported by OEG.",
		)

		@field_validator("packageType", mode="before")

mcp_module/tools/edge_application.py

+3 −1

Changes for mcp_module/tools/edge_application.py: 3 added lines, 1 removed line.
Original line number	Diff line number	Diff line
		@@ -124,7 +124,9 @@ async def register_app_definition(
		payload["version"] = inp.version
		payload["packageType"] = inp.packageType
		payload["appRepo"] = inp.appRepo.model_dump(exclude_none=True)
		payload["requiredResources"] = inp.requiredResources
		payload["requiredResources"] = inp.requiredResources.model_dump(
		mode="json", exclude_none=True
		)
		payload["componentSpec"] = [
		component.model_dump() for component in inp.componentSpec
		]

tests/mcp_module/test_edge_resources.py

0 → 100644

+75 −0

Changes for tests/mcp_module/test_edge_resources.py: 75 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		import json

		import pytest
		from pydantic import ValidationError

		from mcp_module.models.edge_application_models import (
		KubernetesResources,
		RegisterAppDefinitionInput,
		)


		def resources():
		return {
		"infraKind": "kubernetes",
		"applicationResources": {
		"cpuPool": {
		"numCPU": 2,
		"memory": 1024,
		"topology": {
		"minNumberOfNodes": 1,
		"minNodeCpu": 2,
		"minNodeMemory": 1024,
		},
		},
		},
		"isStandalone": True,
		}


		def test_resources_serialize_without_optional_nulls():
		payload = resources()
		model = KubernetesResources.model_validate(payload)
		assert model.model_dump(mode="json", exclude_none=True) == payload
		json.dumps(model.model_dump(mode="json", exclude_none=True))


		@pytest.mark.parametrize("level", ["root", "application", "cpu", "topology"])
		def test_unknown_resource_fields_are_rejected(level):
		payload = resources()
		targets = {
		"root": payload,
		"application": payload["applicationResources"],
		"cpu": payload["applicationResources"]["cpuPool"],
		"topology": payload["applicationResources"]["cpuPool"]["topology"],
		}
		targets[level]["guessedField"] = 1
		with pytest.raises(ValidationError, match="extra_forbidden"):
		KubernetesResources.model_validate(payload)


		@pytest.mark.parametrize("case", ["infra", "topology", "cpu", "storage", "addons"])
		def test_invalid_resource_constraints(case):
		payload = resources()
		if case == "infra":
		payload["infraKind"] = "vm"
		elif case == "topology":
		del payload["applicationResources"]["cpuPool"]["topology"]
		elif case == "cpu":
		payload["applicationResources"]["cpuPool"]["numCPU"] = 257
		elif case == "storage":
		payload["additionalStorage"] = "10Gi"
		else:
		payload["addons"] = ["ingress", "ingress"]
		with pytest.raises(ValidationError):
		KubernetesResources.model_validate(payload)


		def test_tool_schema_exposes_nested_resources():
		schema = RegisterAppDefinitionInput.model_json_schema()
		assert schema["properties"]["requiredResources"]["$ref"].endswith(
		"/KubernetesResources"
		)
		definitions = schema["$defs"]
		assert "topology" in definitions["CpuPool"]["required"]
		assert definitions["CpuPoolTopology"]["additionalProperties"] is False