syntax = "proto3";

// Rename with caution. The package name is part of gRPC service name.
// Existing clients won't find the service if the package name changes.
package executor_api_pb;

// ===== DataPayload =====
enum DataPayloadEncoding {
    DATA_PAYLOAD_ENCODING_UNKNOWN = 0;
    DATA_PAYLOAD_ENCODING_UTF8_JSON = 1;
    DATA_PAYLOAD_ENCODING_UTF8_TEXT = 2;
    DATA_PAYLOAD_ENCODING_BINARY_PICKLE = 3;
    DATA_PAYLOAD_ENCODING_BINARY_ZIP = 4;
    DATA_PAYLOAD_ENCODING_RAW = 5;
}

message DataPayload {
    // URI of the BLOB where the data is stored.
    // S3 URI if the data is stored in S3.
    // Starts with "file://"" prefix if the data is stored on a local file system.
    optional string uri = 1;
    optional DataPayloadEncoding encoding = 2;
    // Not set and ignored by Server right now.
    optional uint64 encoding_version = 3;

    optional string content_type = 4;

    // DataPayload metadata is stored at the head of the binary 
    // blob. This metadata is used optionally to store additional 
    // metadata for the type of the object that is stored 
    // so that we can deserialize the object correctly.
    optional uint64 metadata_size = 5;

    optional uint64 offset = 6;
    // Includes data and metadata.
    optional uint64 size = 7;
    optional string sha256_hash = 8;

    // ID of the function call that returned this data payload as it's return value.
    // Payloads which are not a function call return value, don't have this field set.
    optional string source_function_call_id = 9;
    
    // Opaque ID. Generated by Server if not set by Executor.
    optional string id = 10;

}

enum GPUModel {
    GPU_MODEL_UNKNOWN = 0;
    GPU_MODEL_NVIDIA_A100_40GB = 1;
    GPU_MODEL_NVIDIA_A100_80GB = 2;
    GPU_MODEL_NVIDIA_H100_80GB = 3;
    GPU_MODEL_NVIDIA_TESLA_T4 = 4;
    GPU_MODEL_NVIDIA_A6000 = 5;
    GPU_MODEL_NVIDIA_A10 = 6;
}

// Free GPUs available at the Executor.
message GPUResources {
    optional uint32 count = 1;
    optional GPUModel model = 2;
}

// Resources that we're currently tracking and limiting on Executor.
message HostResources {
    optional uint32 cpu_count = 1;
    optional uint64 memory_bytes = 2;
    optional uint64 disk_bytes = 3;
    // Not set if no GPUs are available.
    optional GPUResources gpu = 4; 
}

// Specification of a single function that is allowed to be run on the Executor.
message AllowedFunction {
    optional string namespace = 1;
    optional string application_name = 2;
    optional string function_name = 3;
    // If none then any version of the application is allowed to run on the Executor.
    optional string application_version = 4;
}

enum FunctionExecutorStatus {
    FUNCTION_EXECUTOR_STATUS_UNKNOWN = 0;
    // FE is being created.
    FUNCTION_EXECUTOR_STATUS_PENDING = 1;
    // FE is running and ready to accept allocations.
    FUNCTION_EXECUTOR_STATUS_RUNNING = 2;
    // FE is terminated, all resources are freed.
    FUNCTION_EXECUTOR_STATUS_TERMINATED = 3;
}

// The reasons why an Executor decided to terminate a Function Executor.
enum FunctionExecutorTerminationReason {
    FUNCTION_EXECUTOR_TERMINATION_REASON_UNKNOWN = 0;
    // Internal error aka platform error on FE startup.
    FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_INTERNAL_ERROR = 1;
    // A clear function constructor code error on FE startup.
    // Typically an exception raised from the constructor.
    FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_FUNCTION_ERROR = 2;
    // Timeout on FE startup while running the function constructor.
    FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_FUNCTION_TIMEOUT = 3;

    // FE was terminated because it failed a health check.
    FUNCTION_EXECUTOR_TERMINATION_REASON_UNHEALTHY = 12;
    // FE was terminated due to an unrecoverable internal error on Executor.
    FUNCTION_EXECUTOR_TERMINATION_REASON_INTERNAL_ERROR = 13;
    // FE was terminated because Function code exceeded its configured timeout.
    // FE termination is the only way to reliably stop the function code execution.
    FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_TIMEOUT = 14;
    // FE was terminated because function allocation currently running on it was removed
    // from Executor desired state.
    // FE termination is the only way to reliably stop the function code execution.
    FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_CANCELLED = 15;
}

// Immutable information that identifies and describes a Function Executor.
message FunctionExecutorResources {
    // 1000 CPU ms per sec is one full CPU core.
    // 2000 CPU ms per sec is two full CPU cores.
    optional uint32 cpu_ms_per_sec = 1;
    optional uint64 memory_bytes = 2;
    optional uint64 disk_bytes = 3;
    optional GPUResources gpu = 4;
}

message FunctionRef {
    optional string namespace = 1;
    optional string application_name = 2;
    optional string function_name = 3;
    optional string application_version = 4;
}

message FunctionExecutorDescription {
    optional string id = 1;
    optional FunctionRef function = 2;
    repeated string secret_names = 3;
    // Timeout for customer code duration during FE creation.
    optional uint32 initialization_timeout_ms = 4;
    optional DataPayload application = 5;
    optional FunctionExecutorResources resources = 6;
    // URI prefix for the startup output payloads.
    // S3 URI if the data is stored in S3.
    // Starts with "file://"" prefix followed by an absolute directory path if the data is stored on a local file system.
    // Deprecated: most probably going to be removed once external FE logs ingestion pipeline gets implemented.
    optional string output_payload_uri_prefix = 7;
    optional uint32 max_concurrency = 8;
    // Timeout for running allocations on FE.
    optional uint32 allocation_timeout_ms = 9;
}

message FunctionExecutorState {
    optional FunctionExecutorDescription description = 1;
    optional FunctionExecutorStatus status = 2;
    optional FunctionExecutorTerminationReason termination_reason = 3;
    repeated string allocation_ids_caused_termination = 4;
}

enum ExecutorStatus {
    EXECUTOR_STATUS_UNKNOWN = 0;
    // Executor is starting up, not ready to accept allocations and FEs.
    EXECUTOR_STATUS_STARTING_UP = 1;
    // Executor is ready to accept allocations and FEs.
    EXECUTOR_STATUS_RUNNING = 2;
    // No new allocations or FEs should be places on the Executor.
    EXECUTOR_STATUS_DRAINED = 3;
    // Executor is stopped, all resources are freed, all FEs are terminated.
    EXECUTOR_STATUS_STOPPED = 4;
}

message ExecutorState {
    optional string executor_id = 1;
    optional string hostname = 3;
    optional string version = 5;
    optional ExecutorStatus status = 6;
    // Total resources at the Executor.
    optional HostResources total_resources = 13;
    // Total resources usable by Function Executors.
    optional HostResources total_function_executor_resources = 7;
    // Empty allowed_functions list means that any function can run on the Executor.
    repeated AllowedFunction allowed_functions = 8;
    repeated FunctionExecutorState function_executor_states = 9;
    map<string, string> labels = 10;
    optional string state_hash = 11;
    // Server supplied clock value of the latest desired executor state that was
    // reconciled by Executor. Not included into state_hash.
    // Initial value on Executor startup is 0.
    optional uint64 server_clock = 12;
}

// Updates that Executor wants to report to Server. If report_executor_state RPC is successful
// then the updates from it won't be included in the next RPC.
message ExecutorUpdate {
    optional string executor_id = 1;
    repeated AllocationResult allocation_results = 2;
}

message ReportExecutorStateRequest {
    optional ExecutorState executor_state = 1;
    optional ExecutorUpdate executor_update = 2;
}

// A message sent by Server to Executor to acknowledge the receipt of ReportExecutorStateRequest.
message ReportExecutorStateResponse {
}

message Allocation {
    // Version is not set in function ref.
    optional FunctionRef function = 1;
    optional string allocation_id = 2;
    optional string function_call_id = 3;
    optional string request_id = 4;
    repeated DataPayload args = 5;
    // URI prefix for the output payloads.
    // S3 URI if the data is stored in S3.
    // Starts with "file://"" prefix followed by an absolute directory path if the data is stored on a local file system.
    optional string output_payload_uri_prefix = 6;
    // BLOB URI prefix for the request error payloads.
    optional string request_error_payload_uri_prefix = 7;
    optional string function_executor_id = 8;

    // Always set. An empty bytes array if no metadata, e.g. for API function call.
    optional bytes function_call_metadata = 9;
}

// A message sent by Executor to Server to open the stream of desired Executor States for the Executor.
message GetDesiredExecutorStatesRequest {
    optional string executor_id = 1;
}

// A message sent from Server to Executor that describes the desired state of the Executor at the moment.
// Executor compares this state with its current state and make necessary changes to match the desired state.
message DesiredExecutorState {
    repeated FunctionExecutorDescription function_executors = 1;
    repeated Allocation allocations = 2;
    // Server supplied clock value used to deduplicate messages. Executor records max clock value
    // it observed and ignores all the messages with clock value <= the max observed value.
    optional uint64 clock = 3;
}

enum AllocationOutcomeCode {
    ALLOCATION_OUTCOME_CODE_UNKNOWN = 0;
    ALLOCATION_OUTCOME_CODE_SUCCESS = 1;
    ALLOCATION_OUTCOME_CODE_FAILURE = 2;
}

enum AllocationFailureReason {
    ALLOCATION_FAILURE_REASON_UNKNOWN = 0;
    // Internal error on Executor aka platform error.
    ALLOCATION_FAILURE_REASON_INTERNAL_ERROR = 1;
    // Clear function code failure typically by raising an exception from the function code.
    // Also a grey failure where we can't determine the exact cause. We attribute these to
    // functions to prevent service abuse but not billed intenionally failing functions.
    ALLOCATION_FAILURE_REASON_FUNCTION_ERROR = 2;
    // Function code run time exceeded its configured timeout.
    ALLOCATION_FAILURE_REASON_FUNCTION_TIMEOUT = 3;
    // Function code raised RequestException to mark the request as permanently failed.
    ALLOCATION_FAILURE_REASON_REQUEST_ERROR = 4;
    // Server removed the allocation from Executor desired state. The allocation didn't finish before the removal.
    ALLOCATION_FAILURE_REASON_ALLOCATION_CANCELLED = 5;
    // Function Executor terminated - can't run the allocation on it anymore.
    ALLOCATION_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED = 6;
}

message ExecutionPlanUpdate {
  oneof op {
    FunctionCall function_call = 10;
    ReduceOp     reduce        = 12;
  }
}

message FunctionCall {
    // The ID of the function call. 
    // When all the dependencies of the function call are resolved, the 
    // resulting DataPayload will have the same ID. 
    // This allows resolving datapayloads of resolved functions calls when it's upstream is called again.
    // The ID is unique within the scope of the request.
    optional string id = 1;
    optional FunctionRef target = 2;
    repeated FunctionArg args = 3;

    // This required metadata allows SDK to restore original function call from the arguments.
    optional bytes call_metadata = 5;
}

message FunctionArg {
  oneof source {
    string function_call_id = 1;
    DataPayload inline_data = 2;
  }
}

message ReduceOp {
    // The DataPayload of the last function call in the reducer will have the function call id set
    // as the id of the ReduceOp.
  optional string id = 1;
  // Contains at least two items.
  repeated FunctionArg collection = 2; 
  optional FunctionRef reducer = 3;
  // This required metadata allows SDK to restore original function call from the arguments.
  optional bytes call_metadata = 4;
}

message ExecutionPlanUpdates {
  repeated ExecutionPlanUpdate updates = 1;
  optional string root_function_call_id = 2;
}

message AllocationResult {
    optional FunctionRef function = 1;
    optional string allocation_id = 2;
    optional string function_call_id = 3;
    optional string request_id = 4;
    optional AllocationOutcomeCode outcome_code = 5;
    optional AllocationFailureReason failure_reason = 6;

    oneof return_value {
      DataPayload value = 7;
      ExecutionPlanUpdates updates = 8;
    }
    // User payload for request error if allocation failed with request error.
    optional DataPayload request_error = 9;
    optional uint64 execution_duration_ms = 10;
}

message FunctionCallRequest {
    // ID of the parent request.
    // Should be set by the Executor for safety.
    optional string parent_request_id = 1;
    // Reference of the target function.
    optional FunctionRef function = 2;

    // Inputs to the function.
    repeated DataPayload inputs = 3;

    // Metadata for the function call.
    optional bytes call_metadata = 4;

    // ID of the source function call.
    // Should be set by the Executor for safety.
    optional string source_function_call_id = 5;
}

message FunctionCallResult {
    oneof result {
        DataPayload output = 1;
        string exception = 2;
    }
}

message FunctionCallResponse {
    oneof event {
        string update = 1;
        FunctionCallResult result = 2;
    }
}

// Internal API for scheduling and running allocations on Executors. Executors are acting as clients of this API.
// Server is responsible for scheduling allocations on Executors and Executors are responsible for running the allocations.
//
// Rename with caution. Existing clients won't find the service if the service name changes. A HTTP2 ingress proxy
// might use the service name in it HTTP2 path based routing rules. See how gRPC uses service names in its HTTP2 paths
// at https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md.
service ExecutorAPI {
    // Called by Executor every 5 seconds to report that it's still alive and provide its current state.
    //
    // Missing 3 reports will result in the Executor being deregistered by Server.
    rpc report_executor_state(ReportExecutorStateRequest) returns (ReportExecutorStateResponse) {}

    // Called by Executor to open a stream of its desired states. When Server wants Executor to change something
    // it puts a message on the stream with the new desired state of the Executor.
    rpc get_desired_executor_states(GetDesiredExecutorStatesRequest) returns (stream DesiredExecutorState) {}

    // Called by the user code to invoke a blocking function call.
    // returns a stream until the function call is completed.
    rpc invoke_function(FunctionCallRequest) returns (stream FunctionCallResponse) {}
}
