// Local scheduler protocol specification

include "gcs.fbs";


// TODO(swang): We put the flatbuffer types in a separate namespace for now to
// avoid conflicts with legacy Ray types.
namespace ray.protocol;

enum MessageType:int {
  // Task is submitted to the local scheduler. This is sent from a worker to a
  // local scheduler.
  SubmitTask = 1,
  // Notify the local scheduler that a task has finished. This is sent from a
  // worker to a local scheduler.
  TaskDone,
  // Log a message to the event table. This is sent from a worker to a local
  // scheduler.
  EventLogMessage,
  // Send an initial connection message to the local scheduler. This is sent
  // from a worker or driver to a local scheduler.
  RegisterClientRequest,
  // Send a reply confirming the successful registration of a worker or driver.
  // This is sent from the local scheduler to a worker or driver.
  RegisterClientReply,
  // Notify the local scheduler that this client is disconnecting unexpectedly.
  // This is sent from a worker to a local scheduler.
  DisconnectClient,
  // Notify the local scheduler that this client is disconnecting gracefully.
  // This is sent from a worker to a local scheduler.
  IntentionalDisconnectClient,
  // Get a new task from the local scheduler. This is sent from a worker to a
  // local scheduler.
  GetTask,
  // Tell a worker to execute a task. This is sent from a local scheduler to a
  // worker.
  ExecuteTask,
  // Reconstruct or fetch possibly lost objects. This is sent from a worker to
  // a local scheduler.
  FetchOrReconstruct,
  // For a worker that was blocked on some object(s), tell the local scheduler
  // that the worker is now unblocked. This is sent from a worker to a local
  // scheduler.
  NotifyUnblocked,
  // A request to get the task frontier for an actor, called by the actor when
  // saving a checkpoint.
  GetActorFrontierRequest,
  // The ActorFrontier response to a GetActorFrontierRequest. The local
  // scheduler returns the actor's per-handle task counts and execution
  // dependencies, which can later be used as the argument to SetActorFrontier
  // when resuming from the checkpoint.
  GetActorFrontierReply,
  // A request to set the task frontier for an actor, called when resuming from
  // a checkpoint. The local scheduler will update the actor's per-handle task
  // counts and execution dependencies, discard any tasks that already executed
  // before the checkpoint, and make any tasks on the frontier runnable by
  // making their execution dependencies available.
  SetActorFrontier,
  // A node manager request to process a task forwarded from another node manager.
  ForwardTaskRequest,
  // Wait for objects to be ready either from local or remote Plasma stores.
  WaitRequest,
  // The response message to WaitRequest; replies with the objects found and objects
  // remaining.
  WaitReply,
  // Push an error to the relevant driver. This is sent from a worker to the
  // node manager.
  PushErrorRequest,
  // Push some profiling events to the GCS. When sending this message to the
  // node manager, the message itself is serialized as a ProfileTableData object.
  PushProfileEventsRequest,
  // Free the objects in objects store.
  FreeObjectsInObjectStoreRequest,
}

table TaskExecutionSpecification {
  // A list of object IDs representing the dependencies of this task that may
  // change at execution time.
  dependencies: [string];
  // The last time this task was received for scheduling.
  last_timestamp: double;
  // The number of times this task was spilled back by local schedulers.
  num_forwards: int;
}

table Task {
  task_specification: string;
  task_execution_spec: TaskExecutionSpecification;
}

table SubmitTaskRequest {
  execution_dependencies: [string];
  task_spec: string;
}

// This message describes a given resource that is reserved for a worker.
table ResourceIdSetInfo {
  // The name of the resource.
  resource_name: string;
  // The resource IDs reserved for this worker.
  resource_ids: [long];
  // The fraction of each resource ID that is reserved for this worker. Note
  // that the length of this list must be the same as the length of
  // resource_ids.
  resource_fractions: [double];
}

// This message is sent from a worker to the node manager.
table DisconnectClient {
}

// This message is sent from the local scheduler to a worker.
table GetTaskReply {
  // A string of bytes representing the task specification.
  task_spec: string;
  // A list of the resources reserved for this worker.
  fractional_resource_ids: [ResourceIdSetInfo];
}

// This struct is used to register a new worker with the local scheduler.
// It is shipped as part of local_scheduler_connect.
table RegisterClientRequest {
  // True if the client is a worker and false if the client is a driver.
  is_worker: bool;
  // The ID of the worker or driver.
  client_id: string;
  // The process ID of this worker.
  worker_pid: long;
  // The driver ID. This is non-nil if the client is a driver.
  driver_id: string;
  // Language of this worker.
  language: Language;
}

table RegisterClientReply {
  // The IDs of the GPUs that are reserved for this worker.
  gpu_ids: [int];
}

table RegisterNodeManagerRequest {
  // GCS ClientID of the connecting node manager.
  client_id: string;
}

table ForwardTaskRequest {
  // The ID of the task to be forwarded.
  task_id: string;
  // The tasks in the uncommitted lineage of the forwarded task. This
  // should include task_id.
  uncommitted_tasks: [Task];
}

table FetchOrReconstruct {
  // List of object IDs of the objects that we want to reconstruct or fetch.
  object_ids: [string];
  // Do we only want to fetch the objects or also reconstruct them?
  fetch_only: bool;
  // The current task ID. If fetch_only is false, then this task is blocked.
  task_id: string;
}

table NotifyUnblocked {
  // The current task ID. This task is no longer blocked.
  task_id: string;
}

table WaitRequest {
  // List of object ids we'll be waiting on.
  object_ids: [string];
  // Number of objects expected to be returned, if available.
  num_ready_objects: int;
  // timeout
  timeout: long;
  // Whether to wait until objects appear locally.
  wait_local: bool;
  // The current task ID. If there are less than num_ready_objects local, then
  // this task is blocked.
  task_id: string;
}

table WaitReply {
  // List of object ids found.
  found: [string];
  // List of object ids not found.
  remaining: [string];
}

// This struct is the same as ErrorTableData.
table PushErrorRequest {
  // The ID of the job that the error is for.
  job_id: string;
  // The type of the error.
  type: string;
  // The error message.
  error_message: string;
  // The timestamp of the error message.
  timestamp: double;
}

table FreeObjectsRequest {
  // Whether keep this request with local object store
  // or send it to all the object stores.
  local_only: bool;
  // List of object ids we'll delete from object store.
  object_ids: [string];
}
