// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

syntax = "proto3";

package google.internal.federated.plan;

import "google/protobuf/any.proto";
import "fcp/protos/confidentialcompute/windowing_schedule.proto";
import "fcp/protos/data_type.proto";

option java_package = "com.google.internal.federated.plan";
option java_multiple_files = true;
option java_outer_classname = "PlanProto";

// Primitives
// ===========

// Container for a metric used by the internal toolkit.
message Metric {
  // Name of an Op to run to read the value.
  string variable_name = 1;

  // A human-readable name for the statistic. Metric names are usually
  // camel case by convention, e.g., 'Loss', 'AbsLoss', or 'Accuracy'.
  // Must be 7-bit ASCII and under 122 characters.
  string stat_name = 2;

  // The human-readable name of another metric by which this metric should be
  // normalized, if any. If empty, this Metric should be aggregated with simple
  // summation. If not empty, the Metric is aggregated according to
  // weighted_metric_sum = sum_i (metric_i * weight_i)
  // weight_sum = sum_i weight_i
  // average_metric_value = weighted_metric_sum / weight_sum
  string weight_name = 3;
}

// Controls the format of output metrics users receive. Represents instructions
// for how metrics are to be output to users, controlling the end format of
// the metric users receive.
message OutputMetric {
  // Metric name.
  string name = 1;

  oneof value_source {
    // A metric representing one stat with aggregation type sum.
    SumOptions sum = 2;

    // A metric representing a ratio between metrics with aggregation
    // type average.
    AverageOptions average = 3;

    // A metric that is not aggregated by the MetricReportAggregator or
    // metrics_loader. This includes metrics like 'num_server_updates' that are
    // aggregated in TensorFlow.
    NoneOptions none = 4;

    // A metric representing one stat with aggregation type only sample.
    // Samples at most 101 clients' values.
    OnlySampleOptions only_sample = 5;
  }
  // Iff True, the metric will be plotted in the default view of the
  // task level Colab automatically.
  oneof visualization_info {
    bool auto_plot = 6 [deprecated = true];
    VisualizationSpec plot_spec = 7;
  }
}

message VisualizationSpec {
  // Different allowable plot types.
  enum VisualizationType {
    NONE = 0;
    DEFAULT_PLOT_FOR_TASK_TYPE = 1;
    LINE_PLOT = 2;
    LINE_PLOT_WITH_PERCENTILES = 3;
    HISTOGRAM = 4;
  }

  // Defines the plot type to provide downstream.
  VisualizationType plot_type = 1;

  // The x-axis which to provide for the given metric. Must be the name of a
  // metric or counter. Recommended x_axis options are source_round, round,
  // or time.
  string x_axis = 2;

  // Iff True, metric will be displayed on a population level dashboard.
  bool plot_on_population_dashboard = 3;
}

// A metric representing one stat with aggregation type sum.
message SumOptions {
  // Name for corresponding Metric stat_name field.
  string stat_name = 1;

  // Iff True, a cumulative sum over rounds will be provided in addition to a
  // sum per round for the value metric.
  bool include_cumulative_sum = 2;

  // Iff True, sample of at most 101 clients' values.
  // Used to calculate quantiles in downstream visualization pipeline.
  bool include_client_samples = 3;
}

// A metric representing a ratio between metrics with aggregation type average.
// Represents: numerator stat / denominator stat.
message AverageOptions {
  // Numerator stat name pointing to corresponding Metric stat_name.
  string numerator_stat_name = 1;

  // Denominator stat name pointing to corresponding Metric stat_name.
  string denominator_stat_name = 2;

  // Name for corresponding Metric stat_name that is the ratio of the
  // numerator stat / denominator stat.
  string average_stat_name = 3;

  // Iff True, sample of at most 101 client's values.
  // Used to calculate quantiles in downstream visualization pipeline.
  bool include_client_samples = 4;
}

// A metric representing one stat with aggregation type none.
message NoneOptions {
  // Name for corresponding Metric stat_name field.
  string stat_name = 1;
}

// A metric representing one stat with aggregation type only sample.
message OnlySampleOptions {
  // Name for corresponding Metric stat_name field.
  string stat_name = 1;
}

// Represents a data set. This is used for testing.
message Dataset {
  // Represents the data set for one client.
  message ClientDataset {
    // A string identifying the client.
    string client_id = 1;

    // A list of serialized tf.Example protos.
    repeated bytes example = 2;

    // Represents a dataset whose examples are selected by an ExampleSelector.
    message SelectedExample {
      ExampleSelector selector = 1;
      repeated bytes example = 2;
    }

    // A list of (selector, dataset) pairs. Used in testing some *TFF-based
    // tasks* that require multiple datasets as client input, e.g., a TFF-based
    // personalization eval task requires each client to provide at least two
    // datasets: one for train, and the other for test.
    repeated SelectedExample selected_example = 3;
  }

  // A list of client data.
  repeated ClientDataset client_data = 1;
}

// Represents predicates over metrics - i.e., expectations. This is used in
// training/eval tests to encode metric names and values expected to be reported
// by a client execution.
message MetricTestPredicates {
  // The value must lie in [lower_bound; upper_bound]. Can also be used for
  // approximate matching (lower == value - epsilon; upper = value + epsilon).
  message Interval {
    double lower_bound = 1;
    double upper_bound = 2;
  }

  // The value must be a real value as long as the value of the weight_name
  // metric is non-zero. If the weight metric is zero, then it is acceptable for
  // the value to be non-real.
  message RealIfNonzeroWeight {
    string weight_name = 1;
  }

  message MetricCriterion {
    // Name of the metric.
    string name = 1;

    // FL training round this metric is expected to appear in.
    int32 training_round_index = 2;

    // If none of the following is set, no matching is performed; but the
    // metric is still expected to be present (with whatever value).
    oneof Criterion {
      // The reported metric must be < lt.
      float lt = 3;
      // The reported metric must be > gt.
      float gt = 4;
      // The reported metric must be <= le.
      float le = 5;
      // The reported metric must be >= ge.
      float ge = 6;
      // The reported metric must be == eq.
      float eq = 7;
      // The reported metric must lie in the interval.
      Interval interval = 8;
      // The reported metric is not NaN or +/- infinity.
      bool real = 9;
      // The reported metric is real (i.e., not NaN or +/- infinity) if the
      // value of an associated weight is not 0.
      RealIfNonzeroWeight real_if_nonzero_weight = 10;
    }
  }

  repeated MetricCriterion metric_criterion = 1;

  reserved 2;
}

// Client Phase
// ============

// A `TensorflowSpec` that is executed on the client in a single `tf.Session`.
// In federated optimization, this will correspond to one `ServerPhase`.
message ClientPhase {
  // If populated, `io_router` must be specified.
  oneof spec {
    // A functional interface for the TensorFlow logic the client should
    // perform.
    TensorflowSpec tensorflow_spec = 4 [lazy = true];
    // Spec for client plans that issue example queries and send the query
    // results directly to an aggregator with no or little additional
    // processing.
    ExampleQuerySpec example_query_spec = 9 [lazy = true];
  }

  // The specification of the inputs coming either from customer apps
  // (Local Compute) or the federated protocol (Federated Compute).
  oneof io_router {
    FederatedComputeIORouter federated_compute = 5 [lazy = true];
    LocalComputeIORouter local_compute = 6 [lazy = true];
    FederatedComputeEligibilityIORouter federated_compute_eligibility = 7
        [lazy = true];
    FederatedExampleQueryIORouter federated_example_query = 8 [lazy = true];
  }

  reserved 1, 2, 3;
}

// TensorflowSpec message describes a single call into TensorFlow, including the
// expected input tensors that must be fed when making that call, which
// output tensors to be fetched, and any operations that have no output but must
// be run. The TensorFlow session will then use the input tensors to do some
// computation, generally reading from one or more datasets, and provide some
// outputs.
//
// Conceptually, client or server code uses this proto along with an IORouter
// to build maps of names to input tensors, vectors of output tensor names,
// and vectors of target nodes:
//
//    CreateTensorflowArguments(
//        TensorflowSpec& spec,
//        IORouter& io_router,
//        const vector<pair<string, Tensor>>* input_tensors,
//        const vector<string>* output_tensor_names,
//        const vector<string>* target_node_names);
//
// Where `input_tensor`, `output_tensor_names` and `target_node_names`
// correspond to the arguments of TensorFlow C++ API for
// `tensorflow::Session:Run()`, and the client executes only a single
// invocation.
//
// Note: the execution engine never sees any concepts related to the federated
// protocol, e.g. input checkpoints or aggregation protocols. This is a "tensors
// in, tensors out" interface. New aggregation methods can be added without
// having to modify the execution engine / TensorflowSpec message, instead they
// should modify the IORouter messages.
//
// Note: both `input_tensor_specs` and `output_tensor_specs` are full
// `tensorflow.TensorSpecProto` messages, though TensorFlow technically
// only requires the names to feed the values into the session. The additional
// dtypes/shape information must always be included in case the runtime
// executing this TensorflowSpec wants to perform additional, optional static
// assertions. The runtimes however are free to ignore the dtype/shapes and only
// rely on the names if so desired.
//
// Assertions:
//   - all names in `input_tensor_specs`, `output_tensor_specs`, and
//     `target_node_names` must appear in the serialized GraphDef where
//     the TF execution will be invoked.
//   - `output_tensor_specs` or `target_node_names` must be non-empty, otherwise
//     there is nothing to execute in the graph.
message TensorflowSpec {
  // The name of a tensor into which a unique token for the current session
  // should be written. The corresponding tensor is a scalar string tensor and
  // is separate from `input_tensors` as there is only one.
  //
  // A session token allows TensorFlow ops such as `ServeSlices` or
  // `ExternalDataset` to refer to callbacks and other session-global objects
  // registered before running the session. In the `ExternalDataset` case, a
  // single dataset_token is valid for multiple `tf.data.Dataset` objects as
  // the token can be thought of as a handle to a dataset factory.
  string dataset_token_tensor_name = 1;

  // tensorflow.TensorSpecProtos for inputs which will be passed to TF.
  //
  // Corresponds to the `feed_dict` parameter of `tf.Session.run()` in
  // TensorFlow's Python API, excluding the dataset_token listed above.
  //
  // Assertions:
  //   - All the tensor names designated as inputs in the corresponding IORouter
  //     must be listed (otherwise the IORouter input work is unused).
  //   - All placeholders in the TF graph must be listed here, with the
  //     exception of the dataset_token which is explicitly set above (otherwise
  //     TensorFlow will fail to execute).
  repeated bytes input_tensor_specs = 2;

  // tensorflow.TensorSpecProtos that should be fetched from TF after execution.
  //
  // Corresponds to the `fetches` parameter of `tf.Session.run()` in
  // TensorFlow's Python API, and the `output_tensor_names` in TensorFlow's C++
  // API.
  //
  // Assertions:
  //   - The set of tensor names here must strictly match the tensor names
  //     designated as outputs in the corresponding IORouter (if any exist).
  repeated bytes output_tensor_specs = 3;

  // Node names in the graph that should be executed, but the output not
  // returned.
  //
  // Corresponds to the `fetches` parameter of `tf.Session.run()` in
  // TensorFlow's Python API, and the `target_node_names` in TensorFlow's C++
  // API.
  //
  // This is intended for use with operations that do not produce tensors, but
  // nonetheless are required to run (e.g. serializing checkpoints).
  repeated string target_node_names = 4;

  // Map of Tensor names to constant inputs (tensorflow.TensorProto).
  // Note: tensors specified via this message should not be included in
  // input_tensor_specs.
  map<string, bytes> constant_inputs = 5;
}

// Configuration describing how clients should generate privacy IDs.
//
// Privacy IDs are user IDs that rotate and are stored in the encrypted part
// of client uploads.
message PrivacyIdConfig {
  // The windowing schedule for the task. Currently only tumbling windows are
  // supported.
  fcp.confidentialcompute.WindowingSchedule windowing_schedule = 1;
}

// ExampleQuerySpec message describes client execution that issues example
// queries and sends the query results directly to an aggregator with no or
// little additional processing.
// This message describes one or more example store queries that perform the
// client side analytics computation in C++. The corresponding output vectors
// will be converted into the expected federated protocol output format.
// This must be used in conjunction with the `FederatedExampleQueryIORouter`.
message ExampleQuerySpec {
  message OutputVectorSpec {
    // The output vector name.
    string vector_name = 1;

    // The data type for each entry in the vector.
    DataType data_type = 2;
  }

  message ExampleQuery {
    // The `ExampleSelector` to issue the query with.
    ExampleSelector example_selector = 1;

    // Indicates that the query returns vector data and must return a single
    // ExampleQueryResult result containing a VectorData entry matching each
    // OutputVectorSpec.vector_name.
    //
    // If the query instead returns no result, then it will be treated as is if
    // an error was returned. In that case, or if the query explicitly returns
    // an error, then the client will abort its session.
    //
    // The keys in the map are the names the vectors should be aggregated under,
    // and must match the keys in FederatedExampleQueryIORouter.aggregations.
    //
    // If this field is set, direct_output_tensor_name must be empty.
    map<string, OutputVectorSpec> output_vector_specs = 2;

    // Indicates that the query should be treated as a raw example store query.
    // The results will be stored as a 1-d string tensor inside a federated
    // compute checkpoint, and this field specifies the name of the tensor.
    // This must match the key in the
    // FederatedExampleQueryIORouter.aggregations.
    //
    // If this field is set, output_vector_specs must be empty.
    string direct_output_tensor_name = 3;

    // The minimum number of output rows required to satisfy the data
    // availability policy. If the number of output rows is less than this
    // value, the client will drop out.
    int32 min_output_row_count = 4;
  }

  // The queries to run.
  repeated ExampleQuery example_queries = 1;

  PrivacyIdConfig privacy_id_config = 2;
}

// The input and output router for Federated Compute plans.
//
// This proto is the glue between the federated protocol and the TensorFlow
// execution engine. This message describes how to prepare data coming from the
// incoming `CheckinResponse` (defined in
// fcp/protos/federated_api.proto) for the `TensorflowSpec`, and what
// to do with outputs from `TensorflowSpec` (e.g. how to aggregate them back on
// the server).
//
// TODO(team) we could replace `input_checkpoint_file_tensor_name` with
// an `input_tensors` field, which would then be a tensor that contains the
// input TensorProtos directly and skipping disk I/O, rather than referring to a
// checkpoint file path.
message FederatedComputeIORouter {
  // ===========================================================================
  //   Inputs
  // ===========================================================================
  // The name of the scalar string tensor that is fed the file path to the
  // initial checkpoint (e.g. as provided via AcceptanceInfo.init_checkpoint).
  //
  // The federated protocol code would copy the `CheckinResponse`'s initial
  // checkpoint to a temporary file and then pass that file path through this
  // tensor.
  //
  // Ops may be added to the client graph that take this tensor as input and
  // reads the path.
  //
  // This field is optional. It may be omitted if the client graph does not use
  // an initial checkpoint.
  string input_filepath_tensor_name = 1;

  // The name of the scalar string tensor that is fed the file path to which
  // client work should serialize the bytes to send back to the server.
  //
  // The federated protocol code generates a temporary file and passes the file
  // path through this tensor.
  //
  // Ops may be be added to the client graph that use this tensor as an argument
  // to write files (e.g. writing checkpoints to disk).
  //
  // This field is optional. It must be omitted if the client graph does not
  // generate any output files (e.g. when all output tensors of `TensorflowSpec`
  // use Secure Aggregation). If this field is not set, then the `ReportRequest`
  // message in the federated protocol will not have the
  // `Report.update_checkpoint` field set. This absence of a value here can be
  // used to validate that the plan only uses Secure Aggregation.
  //
  // Conversely, if this field is set and executing the associated
  // TensorflowSpec does not write to the path is indication of an internal
  // framework error. The runtime should notify the caller that the computation
  // was setup incorrectly.
  string output_filepath_tensor_name = 2;

  // ===========================================================================
  //   Outputs
  // ===========================================================================
  // Describes which output tensors should be aggregated using an aggregation
  // protocol, and the configuration for those protocols.
  //
  // Assertions:
  //   - All keys must exist in the associated `TensorflowSpec` as
  //     `output_tensor_specs.name` values.
  map<string, AggregationConfig> aggregations = 3;
}

// The input and output router for client plans that do not use TensorFlow.
//
// This proto is the glue between the federated protocol and the example query
// execution engine, describing how the query results should ultimately be
// aggregated.
message FederatedExampleQueryIORouter {
  // Describes how each output vector should be aggregated using an aggregation
  // protocol, and the configuration for those protocols.
  // Keys must match the keys in ExampleQuerySpec.output_vector_specs.
  map<string, AggregationConfig> aggregations = 1;
}

// The specification for how to aggregate the associated tensor across clients
// on the server.
message AggregationConfig {
  oneof protocol_config {
    // Indicates that the given output tensor should be processed using Secure
    // Aggregation, using the specified config options.
    SecureAggregationConfig secure_aggregation = 2;

    // Note: in the future we could add a `SimpleAggregationConfig` to add
    // support for simple aggregation without writing to an intermediate
    // checkpoint file first.

    // Indicates that the given output tensor or vector (e.g. as produced by an
    // ExampleQuerySpec) should be placed in an output TF v1 checkpoint or
    // federated compute checkpoint. Currently, which report format is used is
    // determined by the `enable_lightweight_client_report_wire_format` client
    // flag.
    //
    // Currently only ExampleQuerySpec output vectors are supported by this
    // aggregation type (i.e. it cannot be used with TensorflowSpec output
    // tensors). The vectors will be stored in the checkpoint as a 1-D Tensor of
    // its corresponding data type.
    TFV1CheckpointAggregation tf_v1_checkpoint_aggregation = 3;

    // Indicates that the given output tensor or vector should be placed in a
    // federated compute checkpoint.
    FederatedComputeCheckpointAggregation
        federated_compute_checkpoint_aggregation = 4;
  }
}

// Parameters for the SecAgg protocol (go/secagg).
//
// Currently only the server uses the SecAgg parameters, so we only use this
// message to signify usage of SecAgg.
message SecureAggregationConfig {}

// Parameters for the TFV1 Checkpoint Aggregation protocol.
//
// Currently only ExampleQuerySpec output vectors are supported by this
// aggregation type (i.e. it cannot be used with TensorflowSpec output
// tensors). The vectors will be stored in the checkpoint as a 1-D Tensor of
// its corresponding data type.
message TFV1CheckpointAggregation {}

// Parameters for the Federated Compute Checkpoint Aggregation protocol.
//
// Currently only the output from a data upload task is supported by this
// aggregation type. The vectors will be stored in the checkpoint as a 1-D
// String Tensor.
message FederatedComputeCheckpointAggregation {}

// The input and output router for eligibility-computing plans. These plans
// compute which other plans a client is eligible to run, and are returned by
// clients via a `EligibilityEvalCheckinResponse` (defined in
// fcp/protos/federated_api.proto).
message FederatedComputeEligibilityIORouter {
  // The name of the scalar string tensor that is fed the file path to the
  // initial checkpoint (e.g. as provided via
  // `EligibilityEvalPayload.init_checkpoint`).
  //
  // For more detail see the
  // `FederatedComputeIoRouter.input_filepath_tensor_name`, which has the same
  // semantics.
  //
  // This field is optional. It may be omitted if the client graph does not use
  // an initial checkpoint.
  //
  // This tensor name must exist in the associated
  // `TensorflowSpec.input_tensor_specs` list.
  string input_filepath_tensor_name = 1;

  // Name of the output tensor (a string scalar) containing the serialized
  // `google.internal.federatedml.v2.TaskEligibilityInfo` proto output. The
  // client code will parse this proto and place it in the
  // `task_eligibility_info` field of the subsequent `CheckinRequest`.
  //
  // This tensor name must exist in the associated
  // `TensorflowSpec.output_tensor_specs` list.
  string task_eligibility_info_tensor_name = 2;
}

// The input and output router for Local Compute plans.
//
// This proto is the glue between the customers app and the TensorFlow
// execution engine. This message describes how to prepare data coming from the
// customer app (e.g. the input directory the app setup), and the temporary,
// scratch output directory that will be notified to the customer app upon
// completion of `TensorflowSpec`.
message LocalComputeIORouter {
  // ===========================================================================
  //   Inputs
  // ===========================================================================
  // The name of the placeholder tensor representing the input resource path(s).
  // It can be a single input directory or file path (in this case the
  // `input_dir_tensor_name` is populated) or multiple input resources
  // represented as a map from names to input directories or file paths (in this
  // case the `multiple_input_resources` is populated).
  //
  // In the multiple input resources case, the placeholder tensors are
  // represented as a map: the keys are the input resource names defined by the
  // users when constructing the `LocalComputation` Python object, and the
  // values are the corresponding placeholder tensor names created by the local
  // computation plan builder.
  //
  // Apps will have the ability to create contracts between their Android code
  // and `LocalComputation` toolkit code to place files inside the input
  // resource paths with known names (Android code) and create graphs with ops
  // to read from these paths (file names can be specified in toolkit code).
  oneof input_resource {
    string input_dir_tensor_name = 1;
    // Directly using the `map` field is not allowed in `oneof`, so we have to
    // wrap it in a new message.
    MultipleInputResources multiple_input_resources = 3;
  }

  // Scalar string tensor name that will contain the output directory path.
  //
  // The provided directory should be considered temporary scratch that will be
  // deleted, not persisted. It is the responsibility of the calling app to
  // move the desired files to a permanent location once the client returns this
  // directory back to the calling app.
  string output_dir_tensor_name = 2;

  // ===========================================================================
  //   Outputs
  // ===========================================================================
  // NOTE: LocalCompute has no outputs other than what the client graph writes
  // to `output_dir` specified above.
}

// Describes the multiple input resources in `LocalComputeIORouter`.
message MultipleInputResources {
  // The keys are the input resource names (defined by the users when
  // constructing the `LocalComputation` Python object), and the values are the
  // corresponding placeholder tensor names created by the local computation
  // plan builder.
  map<string, string> input_resource_tensor_name_map = 1;
}

// Describes a queue to which input is fed.
message AsyncInputFeed {
  // The op for enqueuing an example input.
  string enqueue_op = 1;

  // The input placeholders for the enqueue op.
  repeated string enqueue_params = 2;

  // The op for closing the input queue.
  string close_op = 3;

  // Whether the work that should be fed asynchronously is the data itself
  // or a description of where that data lives.
  bool feed_values_are_data = 4;
}

message DatasetInput {
  // Initializer of iterator corresponding to tf.data.Dataset object which
  // handles the input data. Stores name of an op in the graph.
  string initializer = 1;

  // Placeholders necessary to initialize the dataset.
  DatasetInputPlaceholders placeholders = 2;

  // Batch size to be used in tf.data.Dataset.
  int32 batch_size = 3;
}

message DatasetInputPlaceholders {
  // Name of placeholder corresponding to filename(s) of SSTable(s) to read data
  // from.
  string filename = 1;

  // Name of placeholder corresponding to key_prefix initializing the
  // SSTableDataset. Note the value fed should be unique user id, not a prefix.
  string key_prefix = 2;

  // Name of placeholder corresponding to number of rounds the local training
  // should be run for.
  string num_epochs = 3;

  // Name of placeholder corresponding to batch size.
  string batch_size = 4;
}

// Specifies an example selection procedure.
message ExampleSelector {
  // Selection criteria following a contract agreed upon between client and
  // model designers.
  google.protobuf.Any criteria = 1;

  // A URI identifying the example collection to read from. Format should adhere
  // to "${COLLECTION}://${APP_NAME}${COLLECTION_NAME}". The URI segments
  // should adhere to the following rules:
  // - The scheme ${COLLECTION} should be one of:
  //   - "app" for app-hosted example
  //   - "simulation" for collections not connected to an app (e.g., if used
  //     purely for simulation)
  // - The authority ${APP_NAME} identifies the owner of the example
  //   collection and should be either the app's package name, or be left empty
  //   (which means "the current app package name").
  // - The path ${COLLECTION_NAME} can be any valid URI path. NB It starts with
  //   a forward slash ("/").
  // - The query and fragment are currently not used, but they may become used
  //   for something in the future. To keep open that possibility they must
  //   currently be left empty.
  //
  // Example: "app://com.google.some.app/someCollection/name"
  // identifies the collection "/someCollection/name" owned and hosted by the
  // app with package name "com.google.some.app".
  //
  // Example: "app:/someCollection/name" or "app:///someCollection/name"
  // both identify the collection "/someCollection/name" owned and hosted by the
  // app associated with the training job in which this URI appears.
  //
  // The path will not be interpreted by the runtime, and will be passed to the
  // example collection implementation for interpretation. Thus, in the case of
  // app-hosted example stores, the path segment's interpretation is a contract
  // between the app's example store developers, and the app's model designers.
  //
  // If an `app://` URI is set, then the `TrainerOptions` collection name must
  // not be set.
  string collection_uri = 2;

  // Resumption token following a contract agreed upon between client and
  // model designers.
  google.protobuf.Any resumption_token = 3;
}

// Selector for slices to fetch as part of a `federated_select` operation.
message SlicesSelector {
  // The string ID under which the slices are served.
  //
  // This value must have been returned by a previous call to the `serve_slices`
  // op run during the `write_client_init` operation.
  string served_at_id = 1;

  // The indices of slices to fetch.
  repeated int32 keys = 2;
}

// Represents slice data to be served as part of a `federated_select` operation.
// This is used for testing.
message SlicesTestDataset {
  // The test data to use. The keys map to the `SlicesSelector.served_at_id`
  // field.  E.g. test slice data for a slice with `served_at_id`="foo" and
  // `keys`=2 would be store in `dataset["foo"].slice_data[2]`.
  map<string, SlicesTestData> dataset = 1;
}
message SlicesTestData {
  // The test slice data to serve. Each entry's index corresponds to the slice
  // key it is the test data for.
  repeated bytes slice_data = 2;
}

// The data that will be persisted on the client. The configuration of where and
// for how long this data will live is up to the client implementation. There
// are no guarantees on how long this information will be kept around.
message ClientPersistedData {
  // The current index of MinimumSeparationPolicy that applied to this
  // computation execution.
  // The MinimumSeparationPolicy now is only used for enabling round-based
  // sampling-without-replacement. The index (i.e., the round number) should be
  // the same no matter what the value of the minimum separation is. So use a
  // single value to record the index.
  optional int64 min_sep_policy_index = 2;

  reserved 1;
}

// Represents a client part of the plan of federated optimization.
// This also used to describe a client-only plan for standalone on-device
// training, known as personalization.
// NEXT_TAG: 7
message ClientOnlyPlan {
  reserved 3;

  // The graph to use for training, in binary form.
  bytes graph = 1;

  // Optional. The flatbuffer used for TFLite training.
  // Whether "graph" or "tflite_graph" is used for training is up to the client
  // code to allow for a flag-controlled a/b rollout.
  bytes tflite_graph = 5;

  // The client phase to execute.
  ClientPhase phase = 2;

  // A TensorFlow ConfigProto.
  google.protobuf.Any tensorflow_config_proto = 4;

  // Optional. The data to be persisted on the client.
  // The configuration of where and for how long this data will live is up to
  // the client implementation. There are no guarantees on how long this
  // information will be kept around.
  optional ClientPersistedData client_persisted_data = 6;
}

message Measurement {
  // Name of a TensorFlow op to run to read/fetch the value of this measurement.
  string read_op_name = 1;

  // A human-readable name for the measurement. Names are usually
  // camel case by convention, e.g., 'Loss', 'AbsLoss', or 'Accuracy'.
  string name = 2;

  reserved 3;

  // A serialized `federated_language.Type` for the measurement.
  bytes tff_type = 4;
}
