update protos to match TF 2.11 and pin the CI version to TF 2.11 (#6029)

* Motivation for features / changes
This PR syncs the Tensorboard Protos to match the Tensorflow 2.11
release candidate. The required a change in our update script which is
also in our master branch.

This also pins the TF version in the CI to 2.11.

Author: JamesHollyer

.github/workflows/ci.yml

@@ -27,7 +27,7 @@ env:
   BUILDTOOLS_VERSION: '3.0.0'
   BUILDIFIER_SHA256SUM: 'e92a6793c7134c5431c58fbc34700664f101e5c9b1c1fcd93b97978e8b7f88db'
   BUILDOZER_SHA256SUM: '3d58a0b6972e4535718cdd6c12778170ea7382de7c75bc3728f5719437ffb84d'
-  TENSORFLOW_VERSION: 'tf-nightly'
+  TENSORFLOW_VERSION: 'tensorflow==2.11.0rc2'
 
 jobs:
   build:

tensorboard/compat/proto/BUILD

@@ -134,6 +134,11 @@ tb_proto_library(
     ],
 )
 
+tb_proto_library(
+    name = "histogram",
+    srcs = ["histogram.proto"],
+)
+
 tb_proto_library(
     name = "meta_graph",
     srcs = ["meta_graph.proto"],
@@ -186,6 +191,11 @@ tb_proto_library(
     ],
 )
 
+tb_proto_library(
+    name = "rpc_options",
+    srcs = ["rpc_options.proto"],
+)
+
 tb_proto_library(
     name = "saved_object_graph",
     srcs = ["saved_object_graph.proto"],
@@ -227,6 +237,7 @@ tb_proto_library(
     name = "summary",
     srcs = ["summary.proto"],
     deps = [
+        ":histogram",
         ":tensor",
     ],
 )
@@ -310,11 +321,13 @@ tb_proto_library(
         ":full_type",
         ":function",
         ":graph",
+        ":histogram",
         ":meta_graph",
         ":node_def",
         ":op_def",
         ":resource_handle",
         ":rewriter_config",
+        ":rpc_options",
         ":saved_object_graph",
         ":saver",
         ":step_stats",

tensorboard/compat/proto/config.proto

@@ -123,6 +123,12 @@ message GPUOptions {
       // created with the default. If this field has values set, then the size
       // of this must match with the above memory_limit_mb.
       repeated int32 priority = 2;
+
+      // Virtual Device ordinal number determines the device ID of the device.
+      // A Virtual device with a lower ordinal number always receives the a
+      // smaller device id. The phyiscal device id and location in the
+      // virtual device list is used to break ties.
+      repeated int32 device_ordinal = 3;
     }
 
     // The multi virtual device settings. If empty (not set), it will create
@@ -132,15 +138,30 @@ message GPUOptions {
     // "visible_device_list" filtering if it is set), and the string represented
     // device names (e.g. /device:GPU:<id>) will refer to the virtual
     // devices and have the <id> field assigned sequentially starting from 0,
-    // according to the order they appear in this list and the "memory_limit"
-    // list inside each element. For example,
+    // according to the order of the virtual devices determined by
+    // device_ordinal and the location in the virtual device list.
+    //
+    // For example,
     //   visible_device_list = "1,0"
     //   virtual_devices { memory_limit: 1GB memory_limit: 2GB }
-    //   virtual_devices {}
-    // will create three virtual devices as:
+    //   virtual_devices { memory_limit: 3GB memory_limit: 4GB }
+    // will create 4 virtual devices as:
     //   /device:GPU:0 -> visible GPU 1 with 1GB memory
     //   /device:GPU:1 -> visible GPU 1 with 2GB memory
-    //   /device:GPU:2 -> visible GPU 0 with all available memory
+    //   /device:GPU:2 -> visible GPU 0 with 3GB memory
+    //   /device:GPU:3 -> visible GPU 0 with 4GB memory
+    //
+    // but
+    //   visible_device_list = "1,0"
+    //   virtual_devices { memory_limit: 1GB memory_limit: 2GB
+    //                     device_ordinal: 10 device_ordinal: 20}
+    //   virtual_devices { memory_limit: 3GB memory_limit: 4GB
+    //                     device_ordinal: 10 device_ordinal: 20}
+    // will create 4 virtual devices as:
+    //   /device:GPU:0 -> visible GPU 1 with 1GB memory  (ordinal 10)
+    //   /device:GPU:1 -> visible GPU 0 with 3GB memory  (ordinal 10)
+    //   /device:GPU:2 -> visible GPU 1 with 2GB memory  (ordinal 20)
+    //   /device:GPU:3 -> visible GPU 0 with 4GB memory  (ordinal 20)
     //
     // NOTE:
     // 1. It's invalid to set both this and "per_process_gpu_memory_fraction"
@@ -798,6 +819,9 @@ message RunMetadata {
   // level idea of what the built graph looks like (since the various graph
   // optimization passes might change the structure of the graph significantly).
   repeated FunctionGraphs function_graphs = 4;
+
+  // Metadata about the session.
+  SessionMetadata session_metadata = 5;
 }
 
 // Defines a connection between two tensors in a `GraphDef`.

tensorboard/compat/proto/coordination_config.proto

@@ -4,6 +4,13 @@ package tensorboard;
 
 option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
 
+// Represents a job type and the number of tasks under this job.
+// For example, ("worker", 20) implies that there will be 20 worker tasks.
+message CoordinatedJob {
+  string name = 1;
+  int32 num_tasks = 2;
+}
+
 // Coordination service configuration parameters.
 // The system picks appropriate values for fields that are not set.
 message CoordinationServiceConfig {
@@ -29,9 +36,9 @@ message CoordinationServiceConfig {
   // recording the state change and the agent stopping heartbeats.
   int64 heartbeat_timeout_in_ms = 5;
 
-  // The list of jobs that partipate in the coordination service. If empty, all
-  // jobs will be included in the coordination service by default.
-  repeated string coordinated_jobs = 6;
+  // The list of `CoordinatedJob`s that will register in coordination service.
+  reserved 6;
+  repeated CoordinatedJob coordinated_job_list = 10;
 
   // Denotes how long to wait for all coordination agents to reach the barriers
   // (after the first shutdown request) before disconnecting together. If
@@ -43,4 +50,10 @@ message CoordinationServiceConfig {
   // find out about the disconnecte agent via stale heartbeats. Used for
   // testing.
   bool agent_destruction_without_shutdown = 8;
+
+  // The list of jobs which are recoverable. If a task in this list fails,
+  // it will not propagate error to other tasks.
+  // If empty, no jobs will be recoverable and every task failure will cause
+  // error propagation to other tasks.
+  repeated string recoverable_jobs = 9;
 }

tensorboard/compat/proto/full_type.proto

@@ -174,6 +174,18 @@ enum FullTypeId {
   //   TFT_LITERAL[TFT_INT32]{1} is the compile-time constant 1.
   TFT_LITERAL = 1003;
 
+  // Encoding types describe a value of a certain type, encoded as a different
+  // type.
+  //
+  // Parametrization:
+  //   TFT_ENCODED[<encoded type>, <encoding type>]
+  //   * <encoded type> may be any type
+  //   * <encoding type> may be any type
+  //
+  // Examples:
+  //   TFT_ENCODING[TFT_INT32, TFT_STRING] is an integer encoded as string.
+  TFT_ENCODED = 1004;
+
   // Type attributes. These always appear in the parametrization of a type,
   // never alone. For example, there is no such thing as a "bool" TensorFlow
   // object (for now).
@@ -217,7 +229,7 @@ enum FullTypeId {
   // is expressed using TFT_PRODUCT.
   //
   //
-  // Parametrization: TFT_ARRAY[<element type>].
+  // Parametrization: TFT_DATASET[<element type>].
   //   * <element type> may be a concrete type or a type symbol. It represents
   //     the data type of the elements produced by the dataset.
   //
@@ -235,6 +247,15 @@ enum FullTypeId {
   // Parametrization: TFT_RAGGED[<element_type>].
   TFT_RAGGED = 10103;
 
+  // Iterators created by tf.data ops and APIs. Very similar to Datasets, except
+  // they are mutable.
+  //
+  //
+  // Parametrization: TFT_ITERATOR[<element type>].
+  //   * <element type> may be a concrete type or a type symbol. It represents
+  //     the data type of the elements produced by the dataset.
+  TFT_ITERATOR = 10104;
+
   // A mutex lock tensor, produced by tf.raw_ops.MutexLock.
   // Unlike strict execution models, where ownership of a lock is denoted by
   // "running after the lock has been acquired", in non-strict mode, lock
@@ -265,7 +286,7 @@ message FullTypeDef {
 
   repeated FullTypeDef args = 2;
 
-  // Literal values of this type object, if the the type admits one.
+  // Literal values of this type object, if the type admits one.
   // For example, a type variable admits a string attribute - its name.
   // Shape-related types may admit int attributes - their static shape values.
   // Fields for more data types to be added as needed.

tensorboard/compat/proto/histogram.proto

@@ -0,0 +1,26 @@
+syntax = "proto3";
+
+package tensorboard;
+
+option cc_enable_arenas = true;
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+option go_package = "github.com/google/tsl/tsl/go/core/protobuf/summary_go_proto";
+
+// Serialization format for histogram module in
+// tsl/lib/histogram/histogram.h
+message HistogramProto {
+  double min = 1;
+  double max = 2;
+  double num = 3;
+  double sum = 4;
+  double sum_squares = 5;
+
+  // Parallel arrays encoding the bucket boundaries and the bucket values.
+  // bucket(i) is the count for the bucket i.  The range for
+  // a bucket is:
+  //   i == 0:  -DBL_MAX .. bucket_limit(0)
+  //   i != 0:  bucket_limit(i-1) .. bucket_limit(i)
+  repeated double bucket_limit = 6 [packed = true];
+  repeated double bucket = 7 [packed = true];
+}

tensorboard/compat/proto/meta_graph.proto

@@ -17,9 +17,6 @@ option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
 option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
 
-// NOTE: This protocol buffer is evolving, and will go through revisions in the
-// coming months.
-//
 // Protocol buffer containing the following which are necessary to restart
 // training, run inference. It can be used to serialize/de-serialize memory
 // objects necessary for running computation in a graph when crossing the

tensorboard/compat/proto/proto_test.py

@@ -150,6 +150,7 @@
     ("tensorflow/core/profiler/", "tensorboard/compat/proto/"),
     ("tensorflow/python/framework/", "tensorboard/compat/proto/"),
     ("tensorflow/core/util/", "tensorboard/compat/proto/"),
+    ("tensorflow/tsl/protobuf/", "tensorboard/compat/proto/"),
     ('package: "tensorflow.tfprof"', 'package: "tensorboard"'),
     ('package: "tensorflow"', 'package: "tensorboard"'),
     ('type_name: ".tensorflow.tfprof', 'type_name: ".tensorboard'),

tensorboard/compat/proto/rewriter_config.proto

@@ -107,10 +107,17 @@ message RewriterConfig {
   // Note that this can change the numerical stability of the graph and may
   // require the use of loss scaling to maintain model convergence.
   Toggle auto_mixed_precision = 23;
-  // Optimize data types for MKL (default is OFF).
+  // Optimize data types for oneDNN (default is OFF).
   // This will try to use bfloat16 on CPUs, which is faster.
   // Note that this can change the numerical stability of the graph.
+  // Note: this is deprecated.
+  // It is replaced by auto_mixed_precision_onednn_bfloat16
   Toggle auto_mixed_precision_mkl = 25;
+  // Optimize data types for oneDNN (default is OFF).
+  // This will try to use bfloat16 on CPUs, which is faster.
+  // Note that this can change the numerical stability of the graph.
+  // Note: this is equivalent to the deprecated option auto_mixed_precision_mkl
+  Toggle auto_mixed_precision_onednn_bfloat16 = 31;
   // Emulate a model using data type float16 on CPU (default is OFF).
   // This will try to emulate the float16 inputs and outputs of an operator
   // on CPU to have better correlation with float16 on GPU; however the
@@ -121,6 +128,8 @@ message RewriterConfig {
   bool disable_meta_optimizer = 19;
   // Optimizers registered by plugin (default is ON)
   Toggle use_plugin_optimizers = 28;
+  // Conditional code motion (default is ON).
+  Toggle experimental_conditional_code_motion = 30;
 
   // Controls how many times we run the optimizers in meta optimizer (default
   // is once).

tensorboard/compat/proto/rpc_options.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package tensorboard;
+
+option go_package = "github.com/google/tsl/tsl/go/protobuf/for_core_protos_go_proto";
+
+// RPC options for distributed runtime.
+message RPCOptions {
+  // If true, always use RPC to contact the session target.
+  //
+  // If false (the default option), TensorFlow may use an optimized
+  // transport for client-master communication that avoids the RPC
+  // stack. This option is primarily for used testing the RPC stack.
+  bool use_rpc_for_inprocess_master = 1;
+
+  // The compression algorithm to be used. One of "deflate", "gzip".
+  string compression_algorithm = 2;
+
+  // If compression_algorithm is set, the compression level to be used.
+  // From 0 (no compression), up to 3.
+  int32 compression_level = 3;
+
+  // Setting cache_rpc_response to true will enable sender side caching of
+  // response for RecvTensorAsync and RecvBufAsync to allow receiver to retry
+  // requests . This is only necessary when the network fabric is experiencing a
+  // significant error rate.  Without it we'll fail a step on an network error,
+  // while with it we'll be able to complete long steps (like complex
+  // initializations) in the face of some network errors during RecvTensor.
+  bool cache_rpc_response = 4;
+
+  // Disables TCP connection sharing when opening a new RPC channel.
+  bool disable_session_connection_sharing = 5;
+
+  // Setting num_channels_per_target > 0 allows uses of multiple channels to
+  // communicate to the same target. This can be used to improve the aggregate
+  // throughput on high speed links (e.g 100G) where single connection is not
+  // sufficient to maximize link utilization. Note that a single RPC only goes
+  // on a single channel, this only helps in situations where there are multiple
+  // transfers to the same target overlapping in time.
+  int32 num_channels_per_target = 6;
+}

tensorboard/compat/proto/summary.proto

@@ -2,6 +2,8 @@ syntax = "proto3";
 
 package tensorboard;
 
+import public "tensorboard/compat/proto/histogram.proto";
+
 import "tensorboard/compat/proto/tensor.proto";
 
 option cc_enable_arenas = true;
@@ -17,24 +19,6 @@ message SummaryDescription {
   string type_hint = 1;
 }
 
-// Serialization format for histogram module in
-// core/lib/histogram/histogram.h
-message HistogramProto {
-  double min = 1;
-  double max = 2;
-  double num = 3;
-  double sum = 4;
-  double sum_squares = 5;
-
-  // Parallel arrays encoding the bucket boundaries and the bucket values.
-  // bucket(i) is the count for the bucket i.  The range for
-  // a bucket is:
-  //   i == 0:  -DBL_MAX .. bucket_limit(0)
-  //   i != 0:  bucket_limit(i-1) .. bucket_limit(i)
-  repeated double bucket_limit = 6 [packed = true];
-  repeated double bucket = 7 [packed = true];
-}
-
 // A SummaryMetadata encapsulates information on which plugins are able to make
 // use of a certain summary value.
 message SummaryMetadata {

tensorboard/compat/proto/types.proto

@@ -75,3 +75,8 @@ enum DataType {
 //    https://www.tensorflow.org/code/tensorboard/compat/proto/types.cc,
 //    https://www.tensorflow.org/code/tensorboard/compat/proto/dtypes.py,
 //    https://www.tensorflow.org/code/tensorboard/compat/proto/function.py)
+
+// Represents a serialized tf.dtypes.Dtype
+message SerializedDType {
+  DataType datatype = 1;
+}

tensorboard/compat/proto/update.sh

@@ -26,6 +26,10 @@ rsync --existing "$1"/tensorflow/core/protobuf/*.proto tensorboard/compat/proto/
 rsync --existing "$1"/tensorflow/core/profiler/*.proto tensorboard/compat/proto/
 rsync --existing "$1"/tensorflow/core/util/*.proto tensorboard/compat/proto/
 rsync --existing "$1"/tensorflow/python/framework/*.proto tensorboard/compat/proto/
+if [ -d $1"/tensorflow/tsl/protobuf" ] # This directory first appears in TF 2.11
+then
+  rsync --existing "$1"/tensorflow/tsl/protobuf/*.proto tensorboard/compat/proto/
+fi
 
 # Rewrite file paths and package names and disable LINT checks.
 find tensorboard/compat/proto/ -type f  -name '*.proto' -exec perl -pi \
@@ -34,6 +38,7 @@ find tensorboard/compat/proto/ -type f  -name '*.proto' -exec perl -pi \
   -e 's|tensorflow/core/profiler|tensorboard/compat/proto|g;' \
   -e 's|tensorflow/core/util|tensorboard/compat/proto|g;' \
   -e 's|tensorflow/python/framework|tensorboard/compat/proto|g;' \
+  -e 's|tensorflow/tsl/protobuf|tensorboard/compat/proto|g;' \
   -e 's|package tensorflow.tfprof;|package tensorboard;|g;' \
   -e 's|package tensorflow;|package tensorboard;|g;' \
   -e 's|tensorflow\.DataType|tensorboard.DataType|g;' \
@@ -60,9 +65,9 @@ find tensorboard/compat/proto/ -type f  -name '*.proto' -exec perl -pi \
       printf 'tb_proto_library(\n'
       printf '    name = "%s",\n' "${f%.proto}"
       printf '    srcs = ["%s"],\n' "$f"
-      if grep -q '^import "tensorboard/' "$f"; then
+      if grep -q '^import\( public\)* "tensorboard/' "$f"; then
         printf '    deps = [\n'
-        grep '^import "tensorboard/' "$f" | sort |
+        grep '^import\( public\)* "tensorboard/' "$f" | sort |
           sed -e 's#.*compat/proto/\([^.]*\).*#        ":\1",#'
         printf '    ],\n'
       fi