fix(customizer): resolve metadata drift in model upload

packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py

@@ -3,25 +3,40 @@
 
 """Helpers for stamping output fileset metadata during customization uploads."""
 
+from typing import Any
 
-def build_output_metadata(
-    *,
-    model: str,
-    finetuning_type: str,
-    output_type: str,
-    save_method: str | None = None,
-) -> dict:
-    """Build the metadata dict stamped onto the output fileset.
-
-    Captures the bits a downstream consumer (model-entity creation,
-    deployment) needs about this artefact without re-deriving them
-    from the training spec.
-    """
-    metadata: dict[str, str] = {
-        "model": model,
-        "finetuning_type": finetuning_type,
-        "output_type": output_type,
-    }
-    if save_method is not None:
-        metadata["save_method"] = save_method
-    return metadata
+
+def extract_tool_calling_metadata(model_entity: Any) -> dict | None:
+    """Extract tool_calling fields from a model entity for output fileset metadata."""
+    spec = getattr(model_entity, "spec", None)
+    if spec is None:
+        return None
+
+    tool_calling: dict[str, Any] = {}
+
+    chat_template = getattr(spec, "chat_template", None)
+    if chat_template:
+        tool_calling["chat_template"] = chat_template
+
+    tcc = getattr(spec, "tool_call_config", None)
+    if tcc is not None:
+        if getattr(tcc, "tool_call_parser", None):
+            tool_calling["tool_call_parser"] = tcc.tool_call_parser
+        if getattr(tcc, "tool_call_plugin", None):
+            tool_calling["tool_call_plugin"] = tcc.tool_call_plugin
+        if getattr(tcc, "auto_tool_choice", None) is not None:
+            tool_calling["auto_tool_choice"] = tcc.auto_tool_choice
+
+    return tool_calling or None
+
+
+def build_model_fileset_metadata(*, tool_calling: dict | None = None) -> dict | None:
+    """Build a ``FilesetMetadata`` dict for model-purpose filesets."""
+    if not tool_calling:
+        return None
+    return {"model": {"tool_calling": tool_calling}}
+
+
+def build_output_fileset_metadata_from_model_entity(model_entity: Any) -> dict | None:
+    """Build output fileset metadata by propagating tool_calling from a source model entity."""
+    return build_model_fileset_metadata(tool_calling=extract_tool_calling_metadata(model_entity))

packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py

@@ -3,32 +3,143 @@
 
 """Tests for output fileset metadata helpers."""
 
-from nmp.customization_common.tasks.file_io_metadata import build_output_metadata
+from types import SimpleNamespace
 
+import pytest
+from nemo_platform_plugin.files.metadata import (
+    FilesetMetadata,
+    ModelMetadataContent,
+    ToolCallingMetadataContent,
+)
+from nmp.customization_common.tasks.file_io_metadata import (
+    build_model_fileset_metadata,
+    build_output_fileset_metadata_from_model_entity,
+    extract_tool_calling_metadata,
+)
+from pydantic import ValidationError
 
-class TestBuildOutputMetadata:
-    def test_extracts_canonical_fields(self) -> None:
-        meta = build_output_metadata(
-            model="Qwen/Qwen3-0.6B",
-            finetuning_type="all_weights",
-            save_method="lora",
-            output_type="model",
+
+def assert_conforms_to_fileset_metadata(meta: dict) -> FilesetMetadata:
+    """Validate a produced metadata dict against the real ``FilesetMetadata`` schema.
+
+    Asserting on raw dicts (as this suite originally did) cannot catch structural
+    drift: a payload with a misplaced or misspelled key still "looks right" but is
+    silently dropped by the schema the platform actually enforces. Round-tripping
+    through ``FilesetMetadata`` (validate -> dump) guarantees every key the helper
+    emits lands in a real schema field, which is the guarantee the upload path needs.
+    """
+    validated = FilesetMetadata.model_validate(meta)
+    assert validated.model_dump(exclude_none=True, by_alias=True) == meta
+    return validated
+
+
+class TestBuildModelFilesetMetadata:
+    def test_wraps_tool_calling_under_model(self) -> None:
+        meta = build_model_fileset_metadata(tool_calling={"tool_call_parser": "llama3_json"})
+        assert meta == {"model": {"tool_calling": {"tool_call_parser": "llama3_json"}}}
+
+        validated = assert_conforms_to_fileset_metadata(meta)
+        assert validated.model is not None
+        assert validated.model.tool_calling is not None
+        assert validated.model.tool_calling.tool_call_parser == "llama3_json"
+
+    def test_returns_none_when_empty(self) -> None:
+        assert build_model_fileset_metadata(tool_calling=None) is None
+
+
+class TestExtractToolCallingMetadata:
+    def test_extracts_from_model_entity_spec(self) -> None:
+        me = SimpleNamespace(
+            spec=SimpleNamespace(
+                chat_template="{% for m in messages %}{{ m }}{% endfor %}",
+                tool_call_config=SimpleNamespace(
+                    tool_call_parser="llama3_json",
+                    tool_call_plugin="default/plugin-fs",
+                    auto_tool_choice=True,
+                ),
+            ),
         )
-        assert meta == {
-            "model": "Qwen/Qwen3-0.6B",
-            "finetuning_type": "all_weights",
-            "save_method": "lora",
-            "output_type": "model",
+        assert extract_tool_calling_metadata(me) == {
+            "chat_template": "{% for m in messages %}{{ m }}{% endfor %}",
+            "tool_call_parser": "llama3_json",
+            "tool_call_plugin": "default/plugin-fs",
+            "auto_tool_choice": True,
         }
 
-    def test_omits_save_method_when_not_provided(self) -> None:
-        meta = build_output_metadata(
-            model="default/base-model",
-            finetuning_type="all_weights",
-            output_type="model",
+    def test_returns_none_without_spec(self) -> None:
+        assert extract_tool_calling_metadata(SimpleNamespace(spec=None)) is None
+
+
+class TestBuildOutputFilesetMetadataFromModelEntity:
+    def test_builds_nested_model_metadata(self) -> None:
+        me = SimpleNamespace(
+            spec=SimpleNamespace(
+                chat_template=None,
+                tool_call_config=SimpleNamespace(
+                    tool_call_parser="hermes",
+                    tool_call_plugin=None,
+                    auto_tool_choice=None,
+                ),
+            ),
         )
-        assert meta == {
-            "model": "default/base-model",
-            "finetuning_type": "all_weights",
-            "output_type": "model",
-        }
+        meta = build_output_fileset_metadata_from_model_entity(me)
+        assert meta == {"model": {"tool_calling": {"tool_call_parser": "hermes"}}}
+
+        validated = assert_conforms_to_fileset_metadata(meta)
+        assert validated.model is not None
+        assert validated.model.tool_calling is not None
+        assert validated.model.tool_calling.tool_call_parser == "hermes"
+
+    def test_all_tool_calling_fields_map_onto_schema(self) -> None:
+        me = SimpleNamespace(
+            spec=SimpleNamespace(
+                chat_template="{% for m in messages %}{{ m }}{% endfor %}",
+                tool_call_config=SimpleNamespace(
+                    tool_call_parser="hermes",
+                    tool_call_plugin="default/plugin-fs",
+                    auto_tool_choice=True,
+                ),
+            ),
+        )
+        meta = build_output_fileset_metadata_from_model_entity(me)
+
+        validated = assert_conforms_to_fileset_metadata(meta)
+        assert validated == FilesetMetadata(
+            model=ModelMetadataContent(
+                tool_calling=ToolCallingMetadataContent(
+                    chat_template="{% for m in messages %}{{ m }}{% endfor %}",
+                    tool_call_parser="hermes",
+                    tool_call_plugin="default/plugin-fs",
+                    auto_tool_choice=True,
+                ),
+            ),
+        )
+
+
+class TestFilesetMetadataSchemaConformance:
+    """Regression coverage for the metadata-drift bug.
+
+    The original failure stamped an output fileset with metadata that did not match
+    the ``FilesetMetadata`` schema the platform validates against. These tests pin
+    the schema as the source of truth rather than trusting hand-written dicts.
+    """
+
+    def test_schema_rejects_string_metadata(self) -> None:
+        # Root cause of the original failure: metadata was set to a bare string
+        # instead of the tagged {model: {tool_calling: {...}}} structure.
+        with pytest.raises(ValidationError):
+            FilesetMetadata.model_validate("llama3_json")
+
+    def test_schema_rejects_string_tool_calling(self) -> None:
+        with pytest.raises(ValidationError):
+            FilesetMetadata.model_validate({"model": {"tool_calling": "llama3_json"}})
+
+    def test_round_trip_detects_misplaced_tool_calling(self) -> None:
+        # The pre-fix shape omitted the `model` wrapper. It validates (extra keys
+        # are ignored) but the tool_calling payload is silently dropped -- exactly
+        # the drift the round-trip assertion in assert_conforms_to_fileset_metadata
+        # is designed to catch.
+        drifted = {"tool_calling": {"tool_call_parser": "hermes"}}
+        validated = FilesetMetadata.model_validate(drifted)
+        assert validated.model is None
+        assert validated.model_dump(exclude_none=True, by_alias=True) != drifted

services/automodel/src/nmp/automodel/app/jobs/compiler.py

@@ -63,6 +63,7 @@
     PEFTConfig as ModelEntityPEFTConfig,
 )
 from nmp.customization_common.service.platform_client import fetch_model_entity
+from nmp.customization_common.tasks.file_io_metadata import build_output_fileset_metadata_from_model_entity
 
 logger = logging.getLogger(__name__)
 
@@ -180,36 +181,6 @@ def _build_file_download_config(
     return FileIOTaskConfig(download=downloads)
 
 
-def _build_output_fileset_metadata(me: ModelEntity) -> dict | None:
-    """Build tool_calling metadata to propagate to the output fileset.
-
-    Extracts chat_template and tool_call_config from the source model entity's spec
-    so the model-spec-runner will apply them to the output model entity.
-
-    Returns:
-        A dict like {"tool_calling": {...}} suitable for fileset metadata, or None
-        if there is nothing to propagate.
-    """
-    if me.spec is None:
-        return None
-
-    tool_calling: dict = {}
-
-    if me.spec.chat_template:
-        tool_calling["chat_template"] = me.spec.chat_template
-
-    if me.spec.tool_call_config:
-        tcc = me.spec.tool_call_config
-        if tcc.tool_call_parser:
-            tool_calling["tool_call_parser"] = tcc.tool_call_parser
-        if tcc.tool_call_plugin:
-            tool_calling["tool_call_plugin"] = tcc.tool_call_plugin
-        if tcc.auto_tool_choice is not None:
-            tool_calling["auto_tool_choice"] = tcc.auto_tool_choice
-
-    return {"tool_calling": tool_calling} if tool_calling else None
-
-
 def _build_file_upload_config(
     output_fileset_name: str,
     fileset_metadata: dict | None = None,
@@ -267,7 +238,7 @@ def _build_model_entity_config(
 
     # Only forward the user-supplied deployment_config from the job spec.
     # tool_call_config from the *source* model entity's spec is propagated
-    # separately via fileset metadata (see _build_output_fileset_metadata),
+    # separately via fileset metadata (see build_output_fileset_metadata_from_model_entity),
     # so we intentionally do not merge it here.
     deployment_config: str | ModelEntityDeploymentParameters | None = None
     if isinstance(job_spec.deployment_config, str):
@@ -441,7 +412,7 @@ async def platform_job_config_compiler(
     # These are propagated to:
     #   1. The training step config (chat_template takes highest priority in template resolution)
     #   2. The output fileset metadata (so the model-spec-runner sets them on the output model)
-    fileset_metadata = _build_output_fileset_metadata(me)
+    fileset_metadata = build_output_fileset_metadata_from_model_entity(me)
     file_io_upload_config = _build_file_upload_config(transformed_spec.output.fileset, fileset_metadata)
 
     # Build model_entity config for creating the model entity

services/rl/src/nmp/rl/app/jobs/compiler.py

@@ -50,7 +50,7 @@
 )
 from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig
 from nmp.customization_common.service.platform_client import fetch_model_entity
-from nmp.customization_common.tasks.file_io_metadata import build_output_metadata
+from nmp.customization_common.tasks.file_io_metadata import build_output_fileset_metadata_from_model_entity
 from nmp.rl.app.constants import (
     BASE_LOG_DIR_ENVVAR,
     DEFAULT_DATASET_PATH,
@@ -121,17 +121,13 @@ def _build_download_config(job_spec: RlJobOutput, me: ModelEntity, *, workspace:
     )
 
 
-def _build_upload_config(job_spec: RlJobOutput) -> FileIOTaskConfig:
+def _build_upload_config(job_spec: RlJobOutput, me) -> FileIOTaskConfig:
     return FileIOTaskConfig(
         upload=[
             UploadItem(
                 src=DEFAULT_OUTPUT_MODEL_PATH,
                 dest=FileSetRef(workspace=None, name=job_spec.output.fileset),
-                metadata=build_output_metadata(
-                    model=job_spec.model,
-                    finetuning_type=FinetuningType.ALL_WEIGHTS.value,
-                    output_type=str(job_spec.output.type),
-                ),
+                metadata=build_output_fileset_metadata_from_model_entity(me),
             ),
         ],
     )
@@ -375,7 +371,7 @@ def _cpu_task_step(
             _build_download_config(job_spec, me, workspace=workspace),
         ),
         _build_training_step(job_spec, base_env, trust_remote_code=trust_remote_code, profile=profile),
-        _cpu_task_step("model-upload", FILE_IO_TASK_COMMAND, _build_upload_config(job_spec)),
+        _cpu_task_step("model-upload", FILE_IO_TASK_COMMAND, _build_upload_config(job_spec, me)),
         _cpu_task_step(
             "model-entity-creation",
             MODEL_ENTITY_TASK_COMMAND,

services/rl/tests/test_compiler.py

@@ -218,11 +218,7 @@ async def test_compiler_emits_four_steps(monkeypatch: pytest.MonkeyPatch, mock_s
     ]
 
     upload_meta = steps[2]["config"]["upload"][0]["metadata"]
-    assert upload_meta == {
-        "model": "default/base-model",
-        "finetuning_type": "all_weights",
-        "output_type": "model",
-    }
+    assert upload_meta is None
 
 
 @pytest.mark.asyncio

services/unsloth/src/nmp/unsloth/app/jobs/compiler.py

@@ -41,7 +41,7 @@
 )
 from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig
 from nmp.customization_common.service.platform_client import fetch_model_entity
-from nmp.customization_common.tasks.file_io_metadata import build_output_metadata
+from nmp.customization_common.tasks.file_io_metadata import build_output_fileset_metadata_from_model_entity
 from nmp.unsloth.app.constants import (
     DEFAULT_DATASET_PATH,
     DEFAULT_MODEL_PATH,
@@ -171,7 +171,7 @@ def _build_file_download_config(
     return FileIOTaskConfig(download=downloads)
 
 
-def _build_file_upload_config(job_spec: UnslothJobOutput) -> FileIOTaskConfig:
+def _build_file_upload_config(job_spec: UnslothJobOutput, me: ModelEntity) -> FileIOTaskConfig:
     """Compile the upload step.
 
     ``workspace=None`` tells the file_io task to use the job's workspace
@@ -182,12 +182,7 @@ def _build_file_upload_config(job_spec: UnslothJobOutput) -> FileIOTaskConfig:
             UploadItem(
                 src=DEFAULT_OUTPUT_MODEL_PATH,
                 dest=FileSetRef(workspace=None, name=job_spec.output.fileset),
-                metadata=build_output_metadata(
-                    model=job_spec.model.name,
-                    finetuning_type=job_spec.training.finetuning_type,
-                    save_method=job_spec.output.save_method,
-                    output_type=job_spec.output.type,
-                ),
+                metadata=build_output_fileset_metadata_from_model_entity(me),
             ),
         ],
     )
@@ -242,7 +237,7 @@ async def platform_job_config_compiler(
 
     validation_dataset_path = _resolve_validation_dataset_path(job_spec, workspace=workspace)
     download_config = _build_file_download_config(job_spec, me, workspace=workspace)
-    upload_config = _build_file_upload_config(job_spec)
+    upload_config = _build_file_upload_config(job_spec, me)
     model_entity_config = _build_model_entity_config(
         workspace,
         job_spec,

services/unsloth/tests/test_compiler_validation_path.py

@@ -122,9 +122,4 @@ async def test_upload_step_stamps_output_metadata() -> None:
         compiler_mod.fetch_model_entity = original_fetch
 
     upload = next(s for s in job["steps"] if s["name"] == "model-upload")
-    assert upload["config"]["upload"][0]["metadata"] == {
-        "model": "default/qwen3-1.7b",
-        "finetuning_type": "lora",
-        "save_method": "lora",
-        "output_type": "adapter",
-    }
+    assert upload["config"]["upload"][0]["metadata"] is None