Skip to content

Commit

Permalink
Update TF compat protos.
Browse files Browse the repository at this point in the history
(cherry picked from commit 9531583)
  • Loading branch information
arcra committed Feb 18, 2025
1 parent a018b73 commit 6bc843e
Show file tree
Hide file tree
Showing 5 changed files with 65 additions and 20 deletions.
38 changes: 34 additions & 4 deletions tensorboard/compat/proto/config.proto
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,14 @@ syntax = "proto3";

package tensorboard;

import "tensorboard/compat/proto/coordination_config.proto";
import "tensorboard/compat/proto/cost_graph.proto";
import "tensorboard/compat/proto/graph.proto";
import "tensorboard/compat/proto/step_stats.proto";
import "tensorboard/compat/proto/cluster.proto";
import "tensorboard/compat/proto/debug.proto";
import "tensorboard/compat/proto/rewriter_config.proto";
import "tensorboard/compat/proto/rpc_options.proto";
import "tensorboard/compat/proto/coordination_config.proto";

option cc_enable_arenas = true;
option java_outer_classname = "ConfigProtos";
Expand Down Expand Up @@ -77,6 +77,11 @@ message GPUOptions {
// name "/device:GPU:<id>") are also called "TF GPU id"s. Please
// refer to third_party/tensorflow/core/common_runtime/gpu/gpu_id.h
// for more information.
// 3. The visible_device_list is also used for PluggableDevice. And
// different types of PluggableDevices share this field. In that case,
// the pluggable_device_type is used to distinguish them, making the
// visible_device_list a list of <pluggable_device_type>:<device_index>,
// e.g. "PluggableDeviceA:0,PluggableDeviceA:1,PluggableDeviceB:0".
string visible_device_list = 5;

// In the event polling loop sleep this many microseconds between
Expand Down Expand Up @@ -306,7 +311,7 @@ message GPUOptions {

// Everything inside experimental is subject to change and is not subject
// to API stability guarantees in
// https://www.tensorflow.org/guide/version_compat.
// https://www.tensorflow.org/guide/versions.
Experimental experimental = 9;
}

Expand Down Expand Up @@ -578,7 +583,7 @@ message ConfigProto {

// Everything inside Experimental is subject to change and is not subject
// to API stability guarantees in
// https://www.tensorflow.org/guide/version_compat.
// https://www.tensorflow.org/guide/versions.
message Experimental {
// Task name for group resolution.
string collective_group_leader = 1;
Expand Down Expand Up @@ -750,7 +755,7 @@ message ConfigProto {
bool disable_functional_ops_lowering = 21;

// Provides a hint to XLA auto clustering to prefer forming a single large
// cluster that encompases most of the graph.
// cluster that encompasses most of the graph.
bool xla_prefer_single_graph_cluster = 22;

// Distributed coordination service configurations.
Expand Down Expand Up @@ -983,3 +988,28 @@ message CallableOptions {

// Next: 9
}

message BatchingOptions {
// Number of scheduling threads for processing batches of work. Determines
// the number of batches processed in parallel. This should be roughly in line
// with the number of TPU cores available.
int32 num_batch_threads = 1;

// The maximum allowed batch size. Can be larger than allowed_batch_sizes to
// utilize large batch splitting.
int32 max_batch_size = 2;

// Maximum number of microseconds to wait before outputting an incomplete
// batch.
int32 batch_timeout_micros = 3;

// Optional list of allowed batch sizes. If left empty, does nothing.
// Otherwise, supplies a list of batch sizes, causing the op to pad batches up
// to one of those sizes. The entries must increase monotonically, and the
// final entry must be equal or less than the max_batch_size.
repeated int32 allowed_batch_sizes = 4;

// Maximum number of batches enqueued for processing before requests are
// failed fast.
int32 max_enqueued_batches = 5;
}
4 changes: 4 additions & 0 deletions tensorboard/compat/proto/coordination_config.proto
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@ message CoordinationServiceConfig {
// Maximum wait time for all members in the cluster to be registered.
int64 cluster_register_timeout_in_ms = 4;

// Denotes if we should synchronize the agents' register attempts by blocking
// on a barrier. This is useful for synchronized restarts.
bool cluster_register_with_barrier = 14;

// Heartbeat timeout, if a task does not record heartbeat in this time
// window, it will be considered disconnected.
// Note: This is also used as a grace period to accept any heartbeats after
Expand Down
19 changes: 11 additions & 8 deletions tensorboard/compat/proto/types.proto
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,13 @@ enum DataType {
DT_FLOAT8_E5M2 = 24; // 5 exponent bits, 2 mantissa bits.
DT_FLOAT8_E4M3FN = 25; // 4 exponent bits, 3 mantissa bits, finite-only, with
// 2 NaNs (0bS1111111).
// TODO - b/299182407: Leaving room for remaining float8 types.
// DT_FLOAT8_E4M3FNUZ = 26;
// DT_FLOAT8_E4M3B11FNUZ = 27;
// DT_FLOAT8_E5M2FNUZ = 28;
DT_FLOAT8_E4M3FNUZ = 26; // 4 exponent bits, 3 mantissa bits, finite-only,
// with NaN.
DT_FLOAT8_E4M3B11FNUZ = 27; // 4 exponent bits, 3 mantissa bits, 11 bits
// bias, finite-only, with NaNs.
DT_FLOAT8_E5M2FNUZ = 28; // 5 exponent bits, 2 mantissa bits, finite-only,
// with NaN.

DT_INT4 = 29;
DT_UINT4 = 30;

Expand Down Expand Up @@ -78,10 +81,10 @@ enum DataType {
DT_UINT64_REF = 123;
DT_FLOAT8_E5M2_REF = 124;
DT_FLOAT8_E4M3FN_REF = 125;
// TODO - b/299182407: Leaving room for remaining float8 types.
// DT_FLOAT8_E4M3FNUZ_REF = 126;
// DT_FLOAT8_E4M3B11FNUZ_REF = 127;
// DT_FLOAT8_E5M2FNUZ_REF = 128;

DT_FLOAT8_E4M3FNUZ_REF = 126;
DT_FLOAT8_E4M3B11FNUZ_REF = 127;
DT_FLOAT8_E5M2FNUZ_REF = 128;
DT_INT4_REF = 129;
DT_UINT4_REF = 130;
}
Expand Down
Binary file modified tensorboard/data/server/descriptor.bin
Binary file not shown.
24 changes: 16 additions & 8 deletions tensorboard/data/server/tensorboard.pb.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 6bc843e

Please sign in to comment.