Update sidv2, yt_prod_v3 test data and re-enable the test.

PiperOrigin-RevId: 802622705

Author: hejiang0116

export/orbax/export/protos/oex_orchestration.proto

@@ -22,24 +22,24 @@ message Pipeline {
   // nth pre-processor serves as the input to the (n+1)th pre-processor. while
   // multi-stage tf preprocessing is doable, we recommend consolidating the
   // steps into a single module to optimize for latency.
-  repeated string pre_processor_names = 20;
+  repeated string pre_processor_names = 3;
 
   // The model functions to be executed in the pipeline. The order of the
   // model functions determines the order of execution. The output of the nth
   // model function serves as the input to the (n+1)th model function.
   // Currently, we ONLY support AT MOST ONE model function (thus, having no
   // model functions is allowed). See b/421976850.
-  repeated BoundModelFunction model_functions = 30;
+  repeated BoundModelFunction model_functions = 4;
 
   // The names (in manifest) of the post-processor functions. For multi-stage
   // postprocessing, the postprocessors are exected in sequence. The output of
   // the nth post-processor serves as the input to the (n+1)th post-processor.
   // while multi-stage tf postprocessing is doable, we recommend consolidating
   // the steps into a single module to optimize for latency.
-  repeated string post_processor_names = 50;
+  repeated string post_processor_names = 5;
 
   // The batch options for the model. If not set, the model will not be batched.
-  optional BatchOptions batch_options = 60;
+  optional BatchOptions batch_options = 6;
 }
 
 // A bound model function.
@@ -53,7 +53,7 @@ message BoundModelFunction {
 
   // The name (in manifest) of the value representing the weights PyTree.
   // This PyTree will be given as the 1st arg to the model function.
-  optional string weights_name = 10;
+  optional string weights_name = 2;
 }
 
 message BatchOptions {

model/orbax/experimental/model/core/protos/manifest.proto

@@ -7,12 +7,12 @@ import "orbax/experimental/model/core/protos/type.proto";
 message Manifest {
   map<string, TopLevelObject> objects = 1;
   map<string, UnstructuredData> supplemental_info = 2;
-  optional DeviceAssignmentByCoords device_assignment_by_coords = 4;
+  optional DeviceAssignmentByCoords device_assignment_by_coords = 3;
 
   // A list of paths frequently used in this manifest. Once a path is added to
   // this list, use sites can refer to it by its index here instead of spelling
   // out the full string.
-  repeated string frequent_paths = 5;
+  repeated string frequent_paths = 4;
 }
 
 // Copied from /third_party/australis/google/ifrt/ifrt_australis.proto.
@@ -32,7 +32,7 @@ message TopLevelObject {
   oneof case {
     Function function = 1;
     Value value = 2;
-    PolymorphicFunction poly_fn = 30;
+    PolymorphicFunction poly_fn = 3;
   }
 }
 
@@ -41,10 +41,10 @@ message Function {
   FunctionBody body = 2;
   // The names of the dependent data. Each of this corresponds to a key in the
   // checkpoint metadata.
-  repeated string data_names = 5;
+  repeated string data_names = 3;
 
-  Visibility visibility = 3;
-  optional string gradient_function_name = 4;
+  Visibility visibility = 4;
+  optional string gradient_function_name = 5;
 }
 
 message Value {

model/orbax/experimental/model/core/python/manifest_constants.py

@@ -14,8 +14,18 @@
 
 """Manifest model format constants."""
 
-MANIFEST_FILENAME = 'manifest.pb'
+MANIFEST_VERSION_FILENAME = 'orbax_model_version.txt'
 
-MANIFEST_VERSION_FILENAME = 'manifest_version.txt'
+# The file path of the manifest proto file
+MANIFEST_FILE_PATH_KEY = 'manifest_file_path'
+# TODO(b/439870164): Update the `MANIFEST_FILENAME` to be `MANIFEST_FILE_PATH`
+# and treat it as a configurable path
+MANIFEST_FILENAME = 'manifest.pb'
 
+# The version of the manifest
+VERSION_KEY = 'version'
 MANIFEST_VERSION = '0.0.1'
+
+# The mime type of the manifest proto file
+MIME_TYPE_KEY = 'mime_type'
+MANIFEST_MIME_TYPE = 'application/protobuf; type=orbax_model_manifest.Manifest'

model/orbax/experimental/model/core/python/manifest_util.py

@@ -18,6 +18,7 @@
 from collections.abc import Mapping, Sequence
 from absl import logging
 from orbax.experimental.model.core.protos import manifest_pb2
+from orbax.experimental.model.core.python import manifest_constants
 from orbax.experimental.model.core.python import unstructured_data
 from orbax.experimental.model.core.python.device_assignment import DeviceAssignment
 from orbax.experimental.model.core.python.function import Function
@@ -28,7 +29,6 @@
 from orbax.experimental.model.core.python.unstructured_data import UnstructuredData
 from orbax.experimental.model.core.python.value import ExternalValue
 
-
 def _build_function(
     fn: Function,
     path: str,
@@ -115,6 +115,28 @@ def _is_seq_of_functions(obj: Saveable) -> bool:
   )
 
 
+def build_manifest_version_file() -> str:
+  """Builds a manifest version file content."""
+
+  # TODO(b/365967674): Remove this check once the manifest filename is
+  # configurable by the manifest version file. Currently, the manifest filename
+  # is hardcoded to "manifest.pb" in OBM & JSV codebase and that needs to be
+  # updated first.
+  if manifest_constants.MANIFEST_FILENAME != "manifest.pb":
+    raise ValueError(
+        "Currently, only manifest.pb is supported as the manifest filename."
+    )
+
+  return (
+      f"{manifest_constants.MANIFEST_FILE_PATH_KEY}:"
+      f' "{manifest_constants.MANIFEST_FILENAME}"\n'
+      f"{manifest_constants.VERSION_KEY}:"
+      f' "{manifest_constants.MANIFEST_VERSION}"\n'
+      f"{manifest_constants.MIME_TYPE_KEY}:"
+      f' "{manifest_constants.MANIFEST_MIME_TYPE}"\n'
+  )
+
+
 def build_manifest_proto(
     obm_module: dict[str, Saveable],
     path: str,

model/orbax/experimental/model/core/python/manifest_util_test.py

@@ -64,5 +64,15 @@ def test_build_device_assignment_by_coords_proto(self):
       self.assertEqual(device.core_on_chip, 0)  # Proto default
 
 
+def test_build_manifest_version_file_content(self):
+  content = manifest_util.build_manifest_version_file()
+  expected_content = (
+      'manifest_file_path: "manifest.pb"\n'
+      'version: "0.0.1"\n'
+      'mime_type: "application/protobuf; type=orbax_model_manifest.Manifest"\n'
+  )
+  self.assertEqual(content, expected_content)
+
+
 if __name__ == '__main__':
   absltest.main()

model/orbax/experimental/model/core/python/save_lib.py

@@ -29,6 +29,7 @@
 from orbax.experimental.model.core.python import unstructured_data
 from orbax.experimental.model.core.python.device_assignment import DeviceAssignment
 from orbax.experimental.model.core.python.manifest_util import build_manifest_proto
+from orbax.experimental.model.core.python.manifest_util import build_manifest_version_file
 from orbax.experimental.model.core.python.saveable import Saveable
 from orbax.experimental.model.core.python.unstructured_data import UnstructuredData
 
@@ -137,8 +138,11 @@ def save(
       os.path.join(path, manifest_constants.MANIFEST_FILENAME), 'wb'
   ) as f:
     f.write(manifest_proto.SerializeToString())
-  # Write the manifest version.
+
+  # Write the main metadata to detect and parse an Orbax Model. The version file
+  # should be THE LAST file to be written. It is used to validate the export and
+  # identify an Orbax Model.
   with file_utils.open_file(
       os.path.join(path, manifest_constants.MANIFEST_VERSION_FILENAME), 'w'
   ) as f:
-    f.write(manifest_constants.MANIFEST_VERSION)
+    f.write(build_manifest_version_file())

model/orbax/experimental/model/test_utils/simple_orchestration.proto

@@ -5,14 +5,14 @@ package orbax_model_simple_orchestration;
 // An orchestration pipeline.
 message Pipeline {
   // The name (in manifest) of the pre-processor function.
-  optional string pre_processor_name = 20;
+  optional string pre_processor_name = 1;
 
   // The name (in manifest) of the model function.
-  optional string model_function_name = 30;
+  optional string model_function_name = 2;
 
   // The name (in manifest) of the value representing the weights.
-  optional string weights_name = 40;
+  optional string weights_name = 3;
 
   // The name (in manifest) of the post-processor function.
-  optional string post_processor_name = 50;
+  optional string post_processor_name = 4;
 }

model/orbax/experimental/model/tf2obm/tf_concrete_function_handle.proto

@@ -14,8 +14,8 @@ message TfConcreteFunctionHandle {
   // SavedModel) but we want to allow calling the function with positional
   // arguments.
   //
-  repeated string input_names = 20;
+  repeated string input_names = 2;
 
   // Similar to `input_names`, but for outputs.
-  repeated string output_names = 30;
+  repeated string output_names = 3;
 }