// Local scheduler protocol specification
namespace ray.local_scheduler.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 gracefully.
  // This is sent from a worker to a local scheduler.
  DisconnectClient,
  // 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.
  ReconstructObjects,
  // 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,
  // Add a result table entry for an object put.
  PutObject,
  // 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
}

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

// This message is sent from the local scheduler to a worker.
table GetTaskReply {
  // A string of bytes representing the task specification.
  task_spec: string;
  // The IDs of the GPUs that the worker is allowed to use for this task.
  gpu_ids: [int];
}

table EventLogMessage {
  key: string;
  value: string;
  timestamp: double;
}

// 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;
}

table DisconnectClient {
}

table ReconstructObjects {
  // 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;
}

table PutObject {
  // Task ID of the task that performed the put.
  task_id: string;
  // Object ID of the object that is being put.
  object_id: string;
}

// The ActorFrontier is used to represent the current frontier of tasks that
// the local scheduler has marked as runnable for a particular actor. It is
// used to save the point in an actor's lifetime at which a checkpoint was
// taken, so that the same frontier of tasks can be made runnable again if the
// actor is resumed from that checkpoint.
table ActorFrontier {
  // Actor ID of the actor whose frontier is described.
  actor_id: string;
  // A list of handle IDs, representing the callers of the actor that have
  // submitted a runnable task to the local scheduler. A nil ID represents the
  // creator of the actor.
  handle_ids: [string];
  // A list representing the number of tasks executed so far, per handle. Each
  // count in task_counters corresponds to the handle at the same in index in
  // handle_ids.
  task_counters: [long];
  // A list representing the execution dependency for the next runnable task,
  // per handle. Each execution dependency in frontier_dependencies corresponds
  // to the handle at the same in index in handle_ids.
  frontier_dependencies: [string];
}

table GetActorFrontierRequest {
  actor_id: string;
}
