Refactor: fix update_json API and add dtype generalization field constants

Author: Dayuxiaoshui

graph_net/graph_net_json_file_util.py

@@ -1,11 +1,15 @@
 import json
 from pathlib import Path
-from typing import Union
 
 kDimensionGeneralizationPasses = "dimension_generalization_passes"
 kDataTypeGeneralizationPasses = "data_type_generalization_passes"
 kSymbolicDimensionReifier = "symbolic_dimension_reifier"
 
+# Fields for dtype generalization metadata
+kDtypeGeneralizationTargetDtype = "dtype_generalization_target_dtype"
+kDtypeGeneralizationPrecision = "dtype_generalization_precision"
+kDtypeGeneralizationGenerated = "dtype_generalization_generated"
+
 
 def read_json(model_path):
     """
@@ -21,44 +25,37 @@ def read_json(model_path):
     return json.loads(graph_net_json_file_path.read_text())
 
 
-def update_json(json_path: Union[str, Path], updates: dict) -> None:
+def update_json(model_path, field, value):
     """
-    Atomically update a JSON file with the given updates.
+    Update a single field in graph_net.json.
 
     Args:
-        json_path: Path to the JSON file
-        updates: Dictionary of key-value pairs to update
+        model_path: Path to model directory or graph_net.json file
+        field: Field name to update
+        value: Value to set
     """
-    json_path = Path(json_path)
+    if isinstance(model_path, (str, Path)):
+        model_path = Path(model_path)
+        # If it's a file path, use it directly; otherwise assume it's a directory
+        if model_path.suffix == ".json":
+            graph_net_json_file_path = model_path
+        else:
+            graph_net_json_file_path = model_path / "graph_net.json"
+    else:
+        graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
 
     # Read existing JSON
-    if json_path.exists():
-        with open(json_path, "r") as f:
-            metadata = json.load(f)
+    if graph_net_json_file_path.exists():
+        with open(graph_net_json_file_path, "r") as f:
+            graph_net_json = json.load(f)
     else:
-        metadata = {}
+        graph_net_json = {}
 
-    # Apply updates
-    metadata.update(updates)
+    # Update field
+    graph_net_json[field] = value
 
     # Atomic write: write to temp file then rename
-    temp_path = json_path.with_suffix(".json.tmp")
+    temp_path = graph_net_json_file_path.with_suffix(".json.tmp")
     with open(temp_path, "w") as f:
-        json.dump(metadata, f, indent=4)
-    temp_path.replace(json_path)
-
-
-# Backward compatibility: old interface using model_path, field, value
-def update_json_legacy(model_path, field, value):
-    """
-    Legacy interface for updating a single field in graph_net.json.
-
-    Args:
-        model_path: Path to model directory
-        field: Field name to update
-        value: Value to set
-    """
-    graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
-    graph_net_json = json.loads(graph_net_json_file_path.read_text())
-    graph_net_json[field] = value
-    graph_net_json_file_path.write_text(json.dumps(graph_net_json, indent=4))
+        json.dump(graph_net_json, f, indent=4)
+    temp_path.replace(graph_net_json_file_path)

graph_net/torch/dtype_generalizer.py

@@ -21,6 +21,9 @@
 
 from graph_net.graph_net_json_file_util import (
     kDataTypeGeneralizationPasses,
+    kDtypeGeneralizationTargetDtype,
+    kDtypeGeneralizationPrecision,
+    kDtypeGeneralizationGenerated,
     update_json,
 )
 from graph_net.torch.constraint_util import RunModelPredicator
@@ -185,10 +188,7 @@ def _save_dtype_pass_names(
             dtype_pass_names: List of working pass names
             model_path: Path to model directory
         """
-        graph_net_json_path = Path(model_path) / "graph_net.json"
-        update_json(
-            graph_net_json_path, {kDataTypeGeneralizationPasses: dtype_pass_names}
-        )
+        update_json(model_path, kDataTypeGeneralizationPasses, dtype_pass_names)
 
 
 class ApplyDataTypeGeneralizationPasses:
@@ -308,13 +308,25 @@ def _apply_pass_and_generate(
             Path to the generated sample directory
         """
         # Parse pass name to extract base name and dtype
-        # Format: "dtype_generalization_pass_float16"
+        # Format: "dtype_generalization_pass_float16" or "dtype_generalization_pass_bfloat16"
+        # The base name "dtype_generalization_pass" corresponds to the file
+        # dtype_generalization_pass.py, which contains the ConcretePass class.
         parts = pass_name.rsplit("_", 1)
         if len(parts) != 2:
-            raise ValueError(f"Invalid pass name format: {pass_name}")
+            raise ValueError(
+                f"Invalid pass name format: {pass_name}. "
+                f"Expected format: 'dtype_generalization_pass_<dtype>'"
+            )
 
         base_name, dtype = parts
 
+        # Validate base name
+        if base_name != "dtype_generalization_pass":
+            raise ValueError(
+                f"Unknown pass base name: {base_name}. "
+                f"Expected: 'dtype_generalization_pass'"
+            )
+
         # Load and apply the pass
         dtype_pass_class = get_dtype_generalization_pass(base_name)
         dtype_pass = dtype_pass_class(
@@ -363,14 +375,9 @@ def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None:
             dtype: Target dtype
         """
         graph_net_json_path = sample_dir / "graph_net.json"
-        update_json(
-            graph_net_json_path,
-            {
-                "dtype": dtype,
-                "precision": dtype,
-                "generated_from_dtype_generalization": True,
-            },
-        )
+        update_json(graph_net_json_path, kDtypeGeneralizationTargetDtype, dtype)
+        update_json(graph_net_json_path, kDtypeGeneralizationPrecision, dtype)
+        update_json(graph_net_json_path, kDtypeGeneralizationGenerated, True)
 
 
 class MultiDtypeFilter: