// 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.federatedcompute.v1;

import "google/protobuf/duration.proto";
import "fcp/protos/federated_api.proto";
import "fcp/protos/federatedcompute/common.proto";

option java_package = "com.google.internal.federatedcompute.v1";
option java_multiple_files = true;

service TaskAssignments {
  // A request sent by the device to check if it should participate in the
  // current task.
  //
  // Clients should use the same `ForwardingInfo` (incl. the headers) as used in
  // this request to construct the URI to poll the operation service to check
  // for task assignment completion.
  //
  // When the task assignment is complete, the result of the operation will
  // either contain an error or the resulting `StartTaskAssignmentResponse` in
  // the response.
  //
  // If the client no longer needs a task assignment because it is interrupted
  // or drops out or an error occurs during polling the long running operation,
  // the client should make a best effort to call CancelOperation.
  //
  // If the returned operation is not complete, clients should poll for status
  // at the rate specified in the StartTaskAssignmentMetadata.
  rpc StartTaskAssignment(StartTaskAssignmentRequest)
      returns (StartTaskAssignmentResponse) {}

  // A request sent by the device to participate in multiple tasks
  // simultaneously.
  //
  // Unlike StartTaskAssignment, which returns at most one task assignment of
  // the server's choice, this RPC attempts to return assignments for *all*
  // tasks requested by the client.
  rpc PerformMultipleTaskAssignments(PerformMultipleTaskAssignmentsRequest)
      returns (PerformMultipleTaskAssignmentsResponse) {}

  // A request sent by the device to report the result of running the task
  // provided by `StartTaskAssignmentResponse`.
  //
  // Clients should use the same `ForwardingInfo` as used in the
  // `StartTaskAssignment` request to construct the URI for this request.
  //
  // A result with a status code other than Code.OK indicates client session
  // termination. The client may not send any future requests with the given
  // session_id.
  rpc ReportTaskResult(ReportTaskResultRequest)
      returns (ReportTaskResultResponse) {}
}

message StartTaskAssignmentRequest {
  // The name of the population this client belongs to.
  //
  // Note that http clients set this value in the request URL instead of the
  // request body.
  string population_name = 1
      ;

  // The session id returned by the server.
  //
  // Note that http clients set this value in the request URL instead of the
  // request body.
  string session_id = 2
      ;

  reserved 3;

  ClientVersion client_version = 4;

  // The client computes this message using the task returned by a previous
  // `RequestEligibilityEvalTask` call.
  //
  // If this field is set, it describes to the server which tasks the client is
  // (in)eligible. The server must take this information into account when
  // deciding task to serve in response to this request.
  //
  // If this field is unset, it may indicate that the client previously received
  // an `EligibilityEvalTask` without an `EligibilityEvalPayload` message (i.e.
  // the population did not have an eligibility-computing task configured at the
  // time of the request). It may also indicate a client for which the
  // eligibility-computing task feature has been disabled, or an old client that
  // does not support this feature yet.
  //
  // If this field is unset but the population has an eligibility-computing task
  // configured, then the server must reject this client, since the server has
  // no way to determine which tasks the client is (in)eligible for.
  //
  // If this field is unset and the population does not have an
  // eligibility-computing task configured, then the server may serve this
  // client any task.
  //
  google.internal.federatedml.v2.TaskEligibilityInfo task_eligibility_info = 5;

  // The client's capabilities when downloading and processing resources.
  ResourceCapabilities resource_capabilities = 6;
}

message StartTaskAssignmentMetadata {
  // The suggested duration between instances of polling the StartTaskAssignment
  // LRO.
  google.protobuf.Duration polling_interval = 1;
}

message StartTaskAssignmentResponse {
  // One of two outcomes, depending on server's decision on participation of the
  // client.
  oneof result {
    // If the client joined the task with this call, information on how to
    // proceed.
    TaskAssignment task_assignment = 1;

    // If the client was not accepted, information how to proceed.
    RejectionInfo rejection_info = 2;
  }
}

message PerformMultipleTaskAssignmentsRequest {
  // The name of the population this client belongs to.
  //
  // Note that http clients set this value in the request URL instead of the
  // request body.
  string population_name = 1
      ;

  // The session id returned by the server.
  string session_id = 2
      ;

  // The client's version information.
  ClientVersion client_version = 3;

  // The client's capabilities when downloading and processing resources.
  ResourceCapabilities resource_capabilities = 4;

  // The list of tasks for which the client would like TaskAssignments. These
  // names are produced by running the population's Eligibility Eval task (see
  // /EligibilityEvalTasks.RequestEligibilityEvalTask).
  repeated string task_names = 5
      ;
}

message PerformMultipleTaskAssignmentsResponse {
  // The TaskAssignments requested by the client. The server may omit requested
  // TaskAssignments, e.g. for any tasks that no longer exist or are not
  // currently accepting client contributions; these cases should be infrequent.
  repeated TaskAssignment task_assignments = 1;
}

// When client (device) is accepted for the current task, this data structure
// carries information necessary to begin task execution.
message TaskAssignment {
  // Information to construct the URI to use for providing resulting aggregation
  // data after task execution.
  ForwardingInfo aggregation_data_forwarding_info = 1;

  message AggregationInfo {}

  message SecureAggregationInfo {
    // The minimum number of clients' values that must be aggregated together
    // before the server can gain access to the aggregate,
    // even transiently (e.g. in RAM).
    // This isn't needed by Secure Aggregation protocol on the client side but
    // shared by the server with clients for transparency or policy reasons.
    int32 minimum_clients_in_server_visible_aggregate = 1;
  }

  oneof aggregation_type {
    // If set, indicates that the Aggregations service (see
    // `aggregations.proto`) should be used to perform aggregation.
    AggregationInfo aggregation_info = 9;

    // If set, indicates that the SecureAggregations service (see
    // `secure_aggregations.proto`) should be used to perform aggregation.
    SecureAggregationInfo secure_aggregation_info = 10;
  }

  // Unique identifier for the client's protocol session.
  string session_id = 5;

  // The opaque id of the aggregation session the client has joined. This is a
  // string generated by the server and MUST NOT contain any information that
  // could be used to identify a specific device.
  string aggregation_id = 2;

  // Unique identifier for the client's participation in an aggregation session.
  string authorization_token = 6;

  // The name identifying the task that was assigned.
  string task_name = 7;

  // The checkpoint from which to start execution.
  Resource init_checkpoint = 3;

  // The plan to be used for execution.
  Resource plan = 4;

  // Info for how to generate URIs for fetching slices at runtime.
  FederatedSelectUriInfo federated_select_uri_info = 8;
}

// Info for how to generate URIs for fetching slices that the task might request
// to be downloaded at runtime.
//
// When one or more slices are requested by the task, the template specified
// here should be used to form a URI from which the client can download the
// slice data, by replacing the "{served_at_id}" and "{key_base10}" substrings
// with the `google.internal.federated.plan.SlicesSelector.served_at_id` and the
// base-10 representation of the `SlicesSelector.keys` value. The client must
// not perform any URI escaping to the values that the substrings are replaced
// with.
message FederatedSelectUriInfo {
  // The URI template to use for fetching slices.
  //
  // This template must always start with "https://".
  //
  // This template must contain the following substrings: "{served_at_id}" and
  // "{key_base10}", as per the above documentation.
  string uri_template = 1;
}

message ReportTaskResultRequest {
  // The name of the population this client belongs to.
  //
  // Note that http clients set this value in the request URL instead of the
  // request body.
  string population_name = 1
      ;

  // The session id returned by the server.
  //
  // Note that http clients set this value in the request URL instead of the
  // request body.
  string session_id = 2
      ;

  // The opaque id of the aggregation session the client has joined. This is a
  // string generated by the server and MUST NOT contain any information that
  // could be used to identify a specific device.
  string aggregation_id = 5
      ;

  // The name identifying the task that was assigned.
  string task_name = 6;

  // Computation status code reported by client.
  // Code.OK indicates that the client computation completed successfully. Any
  // other code indicates unsuccessful computation execution and termination of
  // the protocol session.
  int32 computation_status_code = 3;

  // Stats produced during the client-side execution of the task.
  ClientStats client_stats = 4;
}

// This message is used to report client stats and execution duration to the
// server.
message ClientStats {
  // The time spent on running the task (includes I/O such as reading examples,
  // but does not include time spent on the network for retrieving the task
  // or uploading results).
  google.protobuf.Duration computation_execution_duration = 1;
}

message ReportTaskResultResponse {}