From 5716e910c57912c13921308eb0ef42f1d6141a3d Mon Sep 17 00:00:00 2001 From: Sam Oluwalana Date: Thu, 16 Jul 2026 10:52:00 -0600 Subject: [PATCH 1/3] fix(customizer): resolve metadata drift in model upload Signed-off-by: Sam Oluwalana --- .../tasks/file_io_metadata.py | 57 +++++++++------ .../tests/tasks/test_file_io_metadata.py | 71 +++++++++++++------ .../src/nmp/automodel/app/jobs/compiler.py | 33 +-------- services/rl/src/nmp/rl/app/jobs/compiler.py | 12 ++-- services/rl/tests/test_compiler.py | 6 +- .../src/nmp/unsloth/app/jobs/compiler.py | 13 ++-- .../tests/test_compiler_validation_path.py | 7 +- 7 files changed, 97 insertions(+), 102 deletions(-) diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py index f48c50c569..357ef6dc47 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py +++ b/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)) diff --git a/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py index 966c31f758..ea91cdd634 100644 --- a/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py +++ b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py @@ -3,32 +3,59 @@ """Tests for output fileset metadata helpers.""" -from nmp.customization_common.tasks.file_io_metadata import build_output_metadata +from types import SimpleNamespace +from nmp.customization_common.tasks.file_io_metadata import ( + build_model_fileset_metadata, + build_output_fileset_metadata_from_model_entity, + extract_tool_calling_metadata, +) -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", + +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"}}} + + 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", + assert build_output_fileset_metadata_from_model_entity(me) == { + "model": {"tool_calling": {"tool_call_parser": "hermes"}}, } diff --git a/services/automodel/src/nmp/automodel/app/jobs/compiler.py b/services/automodel/src/nmp/automodel/app/jobs/compiler.py index b086482fc3..4e0f7f78bc 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/compiler.py +++ b/services/automodel/src/nmp/automodel/app/jobs/compiler.py @@ -53,6 +53,7 @@ FileSetRef, UploadItem, ) +from nmp.customization_common.tasks.file_io_metadata import build_output_fileset_metadata_from_model_entity from nmp.customization_common.schemas.model_entity import ( DeploymentParameters as ModelEntityDeploymentParameters, ) @@ -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, @@ -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 diff --git a/services/rl/src/nmp/rl/app/jobs/compiler.py b/services/rl/src/nmp/rl/app/jobs/compiler.py index cc63bdc352..df8cd12a70 100644 --- a/services/rl/src/nmp/rl/app/jobs/compiler.py +++ b/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, diff --git a/services/rl/tests/test_compiler.py b/services/rl/tests/test_compiler.py index 41f0b46d18..f4c784b19c 100644 --- a/services/rl/tests/test_compiler.py +++ b/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 diff --git a/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py b/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py index 2743974f4a..63b1bffea6 100644 --- a/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py +++ b/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, diff --git a/services/unsloth/tests/test_compiler_validation_path.py b/services/unsloth/tests/test_compiler_validation_path.py index 51fe9fe0b8..bf85e6491e 100644 --- a/services/unsloth/tests/test_compiler_validation_path.py +++ b/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 From 12024374336e330270a56c9907bfe1f1330a1c13 Mon Sep 17 00:00:00 2001 From: Sam O Date: Fri, 17 Jul 2026 10:25:28 -0600 Subject: [PATCH 2/3] Update services/automodel/src/nmp/automodel/app/jobs/compiler.py Co-authored-by: Albert Cui Signed-off-by: Sam Oluwalana --- services/automodel/src/nmp/automodel/app/jobs/compiler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/services/automodel/src/nmp/automodel/app/jobs/compiler.py b/services/automodel/src/nmp/automodel/app/jobs/compiler.py index 4e0f7f78bc..c29a453873 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/compiler.py +++ b/services/automodel/src/nmp/automodel/app/jobs/compiler.py @@ -238,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): From f8ab18501bf41319570eb0bb57231994a80ae5e4 Mon Sep 17 00:00:00 2001 From: Sam Oluwalana Date: Fri, 17 Jul 2026 10:49:49 -0600 Subject: [PATCH 3/3] Fix test case Signed-off-by: Sam Oluwalana --- .../tests/tasks/test_file_io_metadata.py | 90 ++++++++++++++++++- .../src/nmp/automodel/app/jobs/compiler.py | 2 +- 2 files changed, 88 insertions(+), 4 deletions(-) diff --git a/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py index ea91cdd634..05e4375d5b 100644 --- a/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py +++ b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py @@ -5,11 +5,32 @@ 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 + + +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: @@ -17,6 +38,11 @@ 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 @@ -56,6 +82,64 @@ def test_builds_nested_model_metadata(self) -> None: ), ), ) - assert build_output_fileset_metadata_from_model_entity(me) == { - "model": {"tool_calling": {"tool_call_parser": "hermes"}}, - } + 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 diff --git a/services/automodel/src/nmp/automodel/app/jobs/compiler.py b/services/automodel/src/nmp/automodel/app/jobs/compiler.py index c29a453873..78a99751f5 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/compiler.py +++ b/services/automodel/src/nmp/automodel/app/jobs/compiler.py @@ -53,7 +53,6 @@ FileSetRef, UploadItem, ) -from nmp.customization_common.tasks.file_io_metadata import build_output_fileset_metadata_from_model_entity from nmp.customization_common.schemas.model_entity import ( DeploymentParameters as ModelEntityDeploymentParameters, ) @@ -64,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__)