fix(customizer): namespace backend schemas to fix OpenAPI name collision

packages/nmp_common/src/nmp/common/api/utils.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+from collections import defaultdict
 from copy import deepcopy
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
@@ -147,35 +148,58 @@ def _walk_spec(d: Dict, visitor: Callable[[str, Any, Dict], None]):
                     _walk_spec(item, visitor)
 
 
-def _normalize_refs_and_schema_keys(spec: Dict) -> Dict:
+def _normalize_refs_and_schema_keys(spec: Dict, *, strict_collisions: bool = False) -> Dict:
     """Normalize all ``$ref`` values and schema dictionary keys.
 
     Schema keys are renamed first so that ``$ref`` values can be rewritten
     consistently.  When two raw keys normalize to the same target and have
-    identical content the duplicate is silently dropped.  When the content
-    differs a warning is logged and the *existing* schema wins (the
-    duplicate is still dropped).
+    identical content the duplicate is silently dropped.
+
+    When the content *differs*, two distinct Pydantic models are fighting over
+    one schema name, and keeping one silently makes the other's ``$ref``\\ s point
+    at the wrong contract.  With ``strict_collisions`` this raises ``ValueError``
+    so the build fails loudly — used for self-contained plugin specs (e.g. the
+    merged ``/apis/customization`` app, where each backend must namespace its
+    own models).  Without it (the default, used for the platform/service specs)
+    it logs a warning and keeps the first-seen schema, preserving legacy
+    behaviour: the platform spec carries pre-existing such collisions that
+    predate this gate and are tracked separately.
     """
     schemas = spec["components"]["schemas"]
 
-    rename_map: Dict[str, str] = {}
-
-    for old_key, value in list(schemas.items()):
-        new_key = normalize_schema_name(old_key)
-        if new_key == old_key:
-            continue
-
-        if new_key in schemas and schemas[new_key] != value:
-            logger.warning(
-                "Schema key collision: '%s' renamed to '%s' which already exists "
-                "with different content. Keeping the existing schema. This likely "
-                "means two distinct Pydantic models from different modules share "
-                "the same class name.",
-                old_key,
-                new_key,
+    # Normalize each raw key once, then group by the name it maps to.  Two
+    # distinct models from different modules (e.g. each customization backend's
+    # own ``TrainingSpec``) produce module-qualified raw keys that normalize to
+    # the same bare name; collapsing them would make one silently steal the
+    # other's ``$ref``\\ s, shipping a wrong contract.
+    key_to_target: Dict[str, str] = {old_key: normalize_schema_name(old_key) for old_key in schemas}
+
+    by_target: Dict[str, List[str]] = defaultdict(list)
+    for old_key, target in key_to_target.items():
+        by_target[target].append(old_key)
+
+    for target, old_keys in by_target.items():
+        # Identical-content duplicates are harmless — they dedup to one schema.
+        # Keep one representative per distinct content shape; more than one means
+        # genuinely different models are colliding on a single name.
+        reps: List[str] = []
+        for key in old_keys:
+            if not any(schemas[key] == schemas[rep] for rep in reps):
+                reps.append(key)
+        if len(reps) > 1:
+            message = (
+                f"OpenAPI schema name collision: {sorted(old_keys)} all normalize to "
+                f"'{target}' with differing content. Two distinct Pydantic models share "
+                f"a class name across modules — namespace them (e.g. via a per-backend "
+                f"NamespacedModel base) so they emit distinct schema names."
             )
+            if strict_collisions:
+                raise ValueError(message)
+            logger.warning("%s Keeping the first-seen schema.", message)
 
-        rename_map[old_key] = new_key
+    # Iterating ``key_to_target`` preserves ``schemas`` order, so the first raw
+    # key that maps to a given target still wins the collapse below.
+    rename_map: Dict[str, str] = {old_key: target for old_key, target in key_to_target.items() if target != old_key}
 
     for old_key, new_key in rename_map.items():
         if new_key not in schemas:
@@ -331,9 +355,9 @@ def _sort_schemas(spec: Dict) -> Dict:
     return spec
 
 
-def tweak_spec(spec: Dict) -> Dict:
+def tweak_spec(spec: Dict, *, strict_collisions: bool = False) -> Dict:
     _walk_spec(spec, _anyof_null_visitor)
-    spec = _normalize_refs_and_schema_keys(spec)
+    spec = _normalize_refs_and_schema_keys(spec, strict_collisions=strict_collisions)
     spec = _split_input_output_schemas(spec)
     spec = _annotate_string_references(spec)
     spec = _sync_schema_titles(spec)

packages/nmp_common/tests/api/test_utils_openapi_spec.py

@@ -328,6 +328,93 @@ def test_tweak_spec_full_pipeline():
     assert result == expected
 
 
+def _collision_spec():
+    """Two distinct models sharing a class name across modules that normalize to
+    the same bare name with *differing* content."""
+    return {
+        "components": {
+            "schemas": {
+                "automodel__schema__TrainingSpec": {
+                    "type": "object",
+                    "properties": {"finetuning_type": {"type": "string"}},
+                },
+                "unsloth__schemas__TrainingSpec": {
+                    "type": "object",
+                    "properties": {"use_gradient_checkpointing": {"type": "string"}},
+                },
+            }
+        },
+        "paths": {
+            "/a": {
+                "post": {
+                    "requestBody": {
+                        "content": {"application/json": {"schema": {"$ref": REF + "automodel__schema__TrainingSpec"}}}
+                    }
+                }
+            }
+        },
+    }
+
+
+def test_tweak_spec_raises_on_collision_when_strict():
+    """With ``strict_collisions`` (plugin specs, e.g. the customization app) a
+    differing-content collision must fail the build loudly rather than silently
+    keeping one and mis-pointing the other's ``$ref``s."""
+    with pytest.raises(ValueError, match="schema name collision"):
+        tweak_spec(_collision_spec(), strict_collisions=True)
+
+
+def test_tweak_spec_warns_and_collapses_on_collision_by_default(caplog):
+    """Non-strict (platform/service specs) preserves legacy behaviour: warn and
+    keep the first-seen schema, so pre-existing platform collisions don't newly
+    break the build."""
+    import logging
+
+    with caplog.at_level(logging.WARNING, logger="nmp.common.api.utils"):
+        result = tweak_spec(_collision_spec())
+
+    assert "schema name collision" in caplog.text
+    schemas = result["components"]["schemas"]
+    assert set(schemas) == {"TrainingSpec"}
+    # First-seen (automodel) wins the collapse.
+    assert "finetuning_type" in schemas["TrainingSpec"]["properties"]
+    ref = result["paths"]["/a"]["post"]["requestBody"]["content"]["application/json"]["schema"]["$ref"]
+    assert ref == REF + "TrainingSpec"
+
+
+def test_tweak_spec_dedups_identical_content_module_qualified_collision():
+    """Two module-qualified keys with *identical* content dedup to one schema —
+    no raise (the fix must not over-trigger on genuinely-shared shapes)."""
+    spec = {
+        "components": {
+            "schemas": {
+                "pkg_a__schema__Shared": {
+                    "type": "object",
+                    "properties": {"value": {"type": "string"}},
+                },
+                "pkg_b__schema__Shared": {
+                    "type": "object",
+                    "properties": {"value": {"type": "string"}},
+                },
+            }
+        },
+        "paths": {
+            "/a": {
+                "post": {
+                    "requestBody": {
+                        "content": {"application/json": {"schema": {"$ref": REF + "pkg_a__schema__Shared"}}}
+                    }
+                }
+            }
+        },
+    }
+
+    result = tweak_spec(spec)
+    assert set(result["components"]["schemas"]) == {"Shared"}
+    ref = result["paths"]["/a"]["post"]["requestBody"]["content"]["application/json"]["schema"]["$ref"]
+    assert ref == REF + "Shared"
+
+
 def test_anyof_null_collapse_preserves_format_and_write_only():
     """Collapsing ``anyOf: [SecretStr, null]`` must keep ``format`` / ``writeOnly`` so
     SDK + docs treat the field as sensitive."""

packages/nmp_customization_common/src/nmp/customization_common/schema.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Namespaced Pydantic base for customization-backend schemas.
+
+Every backend (automodel, unsloth, …) is merged into one FastAPI app under
+``/apis/customization``, so two backends that each define e.g. ``TrainingSpec``
+collide in the generated OpenAPI. Pydantic freezes a model's JSON-schema name at
+class-creation time, so the prefix has to be applied there — this metaclass does
+it. ``class TrainingSpec(AutomodelSchema)`` is emitted as ``AutomodelTrainingSpec``.
+
+Usage: declare a per-backend subclass that sets ``__schema_namespace__`` and
+inherit *that* from every model the backend owns::
+
+    class AutomodelSchema(NamespacedModel):
+        __schema_namespace__ = "Automodel"
+
+    class TrainingSpec(AutomodelSchema):  # emitted as ``AutomodelTrainingSpec``
+        ...
+
+A model whose class name already starts with the prefix (e.g.
+``AutomodelJobInput``) is left unchanged, so top-level request/response names stay
+stable.
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from pydantic import BaseModel, ConfigDict
+
+# ``type(BaseModel)`` is pydantic's ``ModelMetaclass``; deriving from it this way
+# avoids importing pydantic internals.
+_ModelMeta = type(BaseModel)
+
+
+class _NamespacedMeta(_ModelMeta):
+    def __new__(mcs, name, bases, ns, **kw):
+        prefix = ns.get("__schema_namespace__") or next(
+            (p for b in bases if (p := getattr(b, "__schema_namespace__", None))), None
+        )
+        if prefix and not name.startswith(prefix):  # don't double-prefix e.g. AutomodelJobInput
+            name = f"{prefix}{name}"
+            ns["__qualname__"] = name
+        return super().__new__(mcs, name, bases, ns, **kw)
+
+
+class NamespacedModel(BaseModel, metaclass=_NamespacedMeta):
+    """Base for backend schemas.
+
+    Declare a per-backend subclass that sets ``__schema_namespace__`` and inherit
+    *that* from every model, so each backend's schemas emit distinct component
+    names in the merged ``/apis/customization`` OpenAPI spec.
+    """
+
+    # Dunder ClassVar → pydantic ignores it as a field; the metaclass reads it at
+    # class-creation time to compute the emitted schema name.
+    __schema_namespace__: ClassVar[str | None] = None
+    model_config = ConfigDict(extra="forbid")  # the config every backend model already sets

packages/nmp_customization_common/tests/test_schema.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for the ``NamespacedModel`` backend-schema base.
+
+These pin the behaviour the customization OpenAPI fix relies on: each backend's
+models emit a namespace-prefixed schema name at *class-creation time*, so two
+backends that each define e.g. ``TrainingSpec`` don't collapse into one component
+when merged into the single ``/apis/customization`` FastAPI app.
+"""
+
+import pytest
+from nmp.customization_common.schema import NamespacedModel
+from pydantic import Field, ValidationError
+
+
+class AutomodelSchema(NamespacedModel):
+    __schema_namespace__ = "Automodel"
+
+
+class UnslothSchema(NamespacedModel):
+    __schema_namespace__ = "Unsloth"
+
+
+# Two backends each define a leaf literally named ``TrainingSpec`` — exactly the
+# collision the metaclass exists to prevent. Capture each binding under an alias
+# before the next ``class TrainingSpec`` statement shadows the module name.
+class TrainingSpec(AutomodelSchema):
+    lr: float = 1e-4
+
+
+AutomodelTrainingSpec = TrainingSpec
+
+
+class TrainingSpec(UnslothSchema):  # noqa: F811 — intentional same-name redefinition
+    steps: int = 100
+
+
+UnslothTrainingSpec = TrainingSpec
+
+
+class AutomodelJobInput(AutomodelSchema):
+    training: AutomodelTrainingSpec = Field(default_factory=AutomodelTrainingSpec)
+
+
+def test_leaf_models_are_prefixed_at_class_creation():
+    assert AutomodelTrainingSpec.__name__ == "AutomodelTrainingSpec"
+    assert UnslothTrainingSpec.__name__ == "UnslothTrainingSpec"
+
+
+def test_same_leaf_name_across_backends_emits_distinct_component_names():
+    a_title = AutomodelTrainingSpec.model_json_schema()["title"]
+    u_title = UnslothTrainingSpec.model_json_schema()["title"]
+    assert a_title == "AutomodelTrainingSpec"
+    assert u_title == "UnslothTrainingSpec"
+    # The two normalize to different bare names, so the platform normalizer won't
+    # collapse them into one component.
+    assert a_title != u_title
+
+
+def test_nested_ref_points_at_the_owning_backends_schema():
+    schema = AutomodelJobInput.model_json_schema()
+    assert "AutomodelTrainingSpec" in schema["$defs"]
+    assert schema["properties"]["training"]["$ref"].endswith("/AutomodelTrainingSpec")
+
+
+def test_top_level_name_already_prefixed_is_not_double_prefixed():
+    assert AutomodelJobInput.__name__ == "AutomodelJobInput"
+    assert AutomodelJobInput.model_json_schema()["title"] == "AutomodelJobInput"
+
+
+def test_backend_base_class_itself_is_not_prefixed():
+    # The intermediate per-backend base keeps its own name (it is never emitted
+    # as a request/response schema, but must not be mangled either).
+    assert AutomodelSchema.__name__ == "AutomodelSchema"
+    assert UnslothSchema.__name__ == "UnslothSchema"
+
+
+def test_extra_forbid_is_inherited():
+    assert AutomodelJobInput.model_config.get("extra") == "forbid"
+    assert AutomodelJobInput.model_json_schema()["additionalProperties"] is False
+    with pytest.raises(ValidationError):
+        # Extra field rejected (validated via dict so the extra key isn't a
+        # static type error); training defaults, so ``bogus`` is the only fault.
+        AutomodelJobInput.model_validate({"training": {}, "bogus": 1})
+
+
+def test_schema_namespace_is_not_a_field():
+    assert "__schema_namespace__" not in AutomodelJobInput.model_fields
+    assert AutomodelSchema.__schema_namespace__ == "Automodel"
+
+
+def test_plain_namespaced_model_without_namespace_is_unprefixed():
+    class Bare(NamespacedModel):
+        x: int = 0
+
+    assert Bare.__name__ == "Bare"
+    assert Bare.model_json_schema()["title"] == "Bare"
+
+
+def test_subclass_with_own_model_config_still_inherits_extra_forbid():
+    """A subclass that declares its OWN ``model_config`` must still inherit
+    ``extra='forbid'`` from the base.
+
+    ``RlJobInput``, ``_TrainingBase`` and ``RlJobOutput`` each set
+    ``ConfigDict(protected_namespaces=())`` and rely on pydantic *merging* (not
+    replacing) the base config to keep rejecting unknown fields. If that ever
+    regressed, ``additionalProperties: false`` would silently vanish from those
+    request bodies and this suite would still pass without this guard.
+    """
+    from pydantic import ConfigDict
+
+    class WithOwnConfig(AutomodelSchema):
+        model_config = ConfigDict(protected_namespaces=())
+
+        value: int = 0
+
+    # Subclass's own key applied *and* the base's ``extra='forbid'`` preserved.
+    assert WithOwnConfig.model_config.get("protected_namespaces") == ()
+    assert WithOwnConfig.model_config.get("extra") == "forbid"
+    assert WithOwnConfig.model_json_schema()["additionalProperties"] is False
+    with pytest.raises(ValidationError):
+        WithOwnConfig.model_validate({"value": 1, "bogus": 2})