Convert task_spec to flatbuffers (#255)

* convert Ray to C++ * convert task_spec to flatbuffers * fix * it compiles * latest * tests are passing * task2 -> task * fix * fix * fix * fix * fix * linting * fix valgrind * upgrade flatbuffers * use debug mode for valgrind * fix naming and comments * downgrade flatbuffers * fix linting * reintroduce TaskSpec_free * rename TaskSpec -> TaskInfo * refactoring * linting
2026-07-22 13:00:49 +08:00 · 2017-03-05 02:05:02 -08:00
parent 65a8659f3d
commit 0b8d279ef2
36 changed files with 1054 additions and 931 deletions
@@ -16,7 +16,8 @@ void remove_actor(SchedulingAlgorithmState *algorithm_state, ActorID actor_id);

 typedef struct task_queue_entry {
  /** The task that is queued. */
-  task_spec *spec;
+  TaskSpec *spec;
+  int64_t task_spec_size;
  struct task_queue_entry *prev;
  struct task_queue_entry *next;
 } task_queue_entry;
@@ -39,7 +40,9 @@ UT_icd task_queue_entry_icd = {sizeof(task_queue_entry *), NULL, NULL, NULL};

 /** This is used to define the queue of actor task specs for which the
 *  corresponding local scheduler is unknown. */
-UT_icd task_spec_icd = {sizeof(task_spec *), NULL, NULL, NULL};
+UT_icd task_spec_icd = {sizeof(TaskSpec *), NULL, NULL, NULL};
+/** This is used to keep track of task spec sizes in the above queue. */
+UT_icd task_spec_size_icd = {sizeof(int64_t), NULL, NULL, NULL};
 /** This is used to define the queue of available workers. */
 UT_icd worker_icd = {sizeof(LocalSchedulerClient *), NULL, NULL, NULL};

@@ -83,6 +86,8 @@ struct SchedulingAlgorithmState {
   *  about a new local scheduler arrives, we will resubmit all of these tasks
   *  locally. */
  UT_array *cached_submitted_actor_tasks;
+  /** An array of task sizes of cached_submitted_actor_tasks. */
+  UT_array *cached_submitted_actor_task_sizes;
  /** An array of pointers to workers in the worker pool. These are workers
   *  that have registered a PID with us and that are now waiting to be
   *  assigned a task to execute. */
@@ -119,6 +124,9 @@ SchedulingAlgorithmState *SchedulingAlgorithmState_init(void) {
  algorithm_state->dispatch_task_queue = NULL;

  utarray_new(algorithm_state->cached_submitted_actor_tasks, &task_spec_icd);
+  utarray_new(algorithm_state->cached_submitted_actor_task_sizes,
+              &task_spec_size_icd);
+
  algorithm_state->local_actor_infos = NULL;

  utarray_new(algorithm_state->available_workers, &worker_icd);
@@ -132,13 +140,13 @@ void SchedulingAlgorithmState_free(SchedulingAlgorithmState *algorithm_state) {
  task_queue_entry *elt, *tmp1;
  DL_FOREACH_SAFE(algorithm_state->waiting_task_queue, elt, tmp1) {
    DL_DELETE(algorithm_state->waiting_task_queue, elt);
-    free_task_spec(elt->spec);
+    free(elt->spec);
    free(elt);
  }
  /* Free all the tasks in the dispatch queue. */
  DL_FOREACH_SAFE(algorithm_state->dispatch_task_queue, elt, tmp1) {
    DL_DELETE(algorithm_state->dispatch_task_queue, elt);
-    free_task_spec(elt->spec);
+    free(elt->spec);
    free(elt);
  }
  /* Remove all of the remaining actors. */
@@ -152,11 +160,12 @@ void SchedulingAlgorithmState_free(SchedulingAlgorithmState *algorithm_state) {
  /* Free the list of cached actor task specs and the task specs themselves. */
  for (int i = 0;
       i < utarray_len(algorithm_state->cached_submitted_actor_tasks); ++i) {
-    task_spec **spec = (task_spec **) utarray_eltptr(
+    TaskSpec **spec = (TaskSpec **) utarray_eltptr(
        algorithm_state->cached_submitted_actor_tasks, i);
    free(*spec);
  }
  utarray_free(algorithm_state->cached_submitted_actor_tasks);
+  utarray_free(algorithm_state->cached_submitted_actor_task_sizes);
  /* Free the list of available workers. */
  utarray_free(algorithm_state->available_workers);
  utarray_free(algorithm_state->executing_workers);
@@ -195,7 +204,7 @@ void provide_scheduler_info(LocalSchedulerState *state,
      waiting_task_queue_length + dispatch_task_queue_length;
  info->available_workers = utarray_len(algorithm_state->available_workers);
  /* Copy static and dynamic resource information. */
-  for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
+  for (int i = 0; i < ResourceIndex_MAX; i++) {
    info->dynamic_resources[i] = state->dynamic_resources[i];
    info->static_resources[i] = state->static_resources[i];
  }
@@ -259,7 +268,7 @@ void remove_actor(SchedulingAlgorithmState *algorithm_state, ActorID actor_id) {
  task_queue_entry *task_queue_elt, *tmp;
  DL_FOREACH_SAFE(entry->task_queue, task_queue_elt, tmp) {
    DL_DELETE(entry->task_queue, task_queue_elt);
-    free_task_spec(task_queue_elt->spec);
+    free(task_queue_elt->spec);
    free(task_queue_elt);
  }
  /* Remove the entry from the hash table and free it. */
@@ -310,9 +319,10 @@ void handle_actor_worker_disconnect(LocalSchedulerState *state,
 */
 void add_task_to_actor_queue(LocalSchedulerState *state,
                             SchedulingAlgorithmState *algorithm_state,
-                             task_spec *spec,
+                             TaskSpec *spec,
+                             int64_t task_spec_size,
                             bool from_global_scheduler) {
-  ActorID actor_id = task_spec_actor_id(spec);
+  ActorID actor_id = TaskSpec_actor_id(spec);
  char tmp[ID_STRING_SIZE];
  ObjectID_to_string(actor_id, tmp, ID_STRING_SIZE);
  DCHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
@@ -332,7 +342,7 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
    CHECK(entry != NULL);
  }

-  int64_t task_counter = task_spec_actor_counter(spec);
+  int64_t task_counter = TaskSpec_actor_counter(spec);
  /* As a sanity check, the counter of the new task should be greater than the
   * number of tasks that have executed on this actor so far (since we are
   * guaranteeing in-order execution of the tasks on the actor). TODO(rkn): This
@@ -342,8 +352,9 @@ void add_task_to_actor_queue(LocalSchedulerState *state,

  /* Create a new task queue entry. */
  task_queue_entry *elt = (task_queue_entry *) malloc(sizeof(task_queue_entry));
-  elt->spec = (task_spec *) malloc(task_spec_size(spec));
-  memcpy(elt->spec, spec, task_spec_size(spec));
+  elt->spec = (TaskSpec *) malloc(task_spec_size);
+  memcpy(elt->spec, spec, task_spec_size);
+  elt->task_spec_size = task_spec_size;
  /* Add the task spec to the actor's task queue in a manner that preserves the
   * order of the actor task counters. Iterate from the beginning of the queue
   * to find the right place to insert the task queue entry. TODO(pcm): This
@@ -351,15 +362,15 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
   * be optimized. */
  task_queue_entry *current_entry = entry->task_queue;
  while (current_entry != NULL && current_entry->next != NULL &&
-         task_counter > task_spec_actor_counter(current_entry->spec)) {
+         task_counter > TaskSpec_actor_counter(current_entry->spec)) {
    current_entry = current_entry->next;
  }
  DL_APPEND_ELEM(entry->task_queue, current_entry, elt);

  /* Update the task table. */
  if (state->db != NULL) {
-    Task *task =
-        Task_alloc(spec, TASK_STATUS_QUEUED, get_db_client_id(state->db));
+    Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
+                            get_db_client_id(state->db));
    if (from_global_scheduler) {
      /* If the task is from the global scheduler, it's already been added to
       * the task table, so just update the entry. */
@@ -403,7 +414,7 @@ bool dispatch_actor_task(LocalSchedulerState *state,
     * the actor. */
    return false;
  }
-  int64_t next_task_counter = task_spec_actor_counter(entry->task_queue->spec);
+  int64_t next_task_counter = TaskSpec_actor_counter(entry->task_queue->spec);
  if (next_task_counter != entry->task_counter) {
    /* We cannot execute the next task on this actor without violating the
     * in-order execution guarantee for actor tasks. */
@@ -418,12 +429,13 @@ bool dispatch_actor_task(LocalSchedulerState *state,
   * as unavailable. */
  task_queue_entry *first_task = entry->task_queue;
  entry->task_counter += 1;
-  assign_task_to_worker(state, first_task->spec, entry->worker);
+  assign_task_to_worker(state, first_task->spec, first_task->task_spec_size,
+                        entry->worker);
  entry->worker_available = false;
  /* Remove the task from the actor's task queue. */
  DL_DELETE(entry->task_queue, first_task);
  /* Free the task spec and the task queue entry. */
-  free_task_spec(first_task->spec);
+  free(first_task->spec);
  free(first_task);
  return true;
 }
@@ -479,12 +491,12 @@ void fetch_missing_dependency(LocalSchedulerState *state,
 void fetch_missing_dependencies(LocalSchedulerState *state,
                                SchedulingAlgorithmState *algorithm_state,
                                task_queue_entry *task_entry) {
-  task_spec *task = task_entry->spec;
-  int64_t num_args = task_num_args(task);
+  TaskSpec *task = task_entry->spec;
+  int64_t num_args = TaskSpec_num_args(task);
  int num_missing_dependencies = 0;
  for (int i = 0; i < num_args; ++i) {
-    if (task_arg_type(task, i) == ARG_BY_REF) {
-      ObjectID obj_id = task_arg_id(task, i);
+    if (TaskSpec_arg_by_ref(task, i)) {
+      ObjectID obj_id = TaskSpec_arg_id(task, i);
      object_entry *entry;
      HASH_FIND(hh, algorithm_state->local_objects, &obj_id, sizeof(obj_id),
                entry);
@@ -508,11 +520,11 @@ void fetch_missing_dependencies(LocalSchedulerState *state,
 *         task are present in the local object store, otherwise it returns
 *         false.
 */
-bool can_run(SchedulingAlgorithmState *algorithm_state, task_spec *task) {
-  int64_t num_args = task_num_args(task);
+bool can_run(SchedulingAlgorithmState *algorithm_state, TaskSpec *task) {
+  int64_t num_args = TaskSpec_num_args(task);
  for (int i = 0; i < num_args; ++i) {
-    if (task_arg_type(task, i) == ARG_BY_REF) {
-      ObjectID obj_id = task_arg_id(task, i);
+    if (TaskSpec_arg_by_ref(task, i)) {
+      ObjectID obj_id = TaskSpec_arg_id(task, i);
      object_entry *entry;
      HASH_FIND(hh, algorithm_state->local_objects, &obj_id, sizeof(obj_id),
                entry);
@@ -580,7 +592,7 @@ void dispatch_tasks(LocalSchedulerState *state,
    }
    /* Terminate early if there are no more resources available. */
    bool resources_available = false;
-    for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
+    for (int i = 0; i < ResourceIndex_MAX; i++) {
      if (state->dynamic_resources[i] > 0) {
        /* There are still resources left, continue checking tasks. */
        resources_available = true;
@@ -593,8 +605,8 @@ void dispatch_tasks(LocalSchedulerState *state,
    }
    /* Skip to the next task if this task cannot currently be satisfied. */
    bool task_satisfied = true;
-    for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
-      if (task_spec_get_required_resource(elt->spec, i) >
+    for (int i = 0; i < ResourceIndex_MAX; i++) {
+      if (TaskSpec_get_required_resource(elt->spec, i) >
          state->dynamic_resources[i]) {
        /* Insufficient capacity for this task, proceed to the next task. */
        task_satisfied = false;
@@ -612,7 +624,7 @@ void dispatch_tasks(LocalSchedulerState *state,
    LocalSchedulerClient **worker = (LocalSchedulerClient **) utarray_back(
        algorithm_state->available_workers);
    /* Tell the available worker to execute the task. */
-    assign_task_to_worker(state, elt->spec, *worker);
+    assign_task_to_worker(state, elt->spec, elt->task_spec_size, *worker);
    /* Remove the worker from the available queue, and add it to the executing
     * workers. */
    utarray_pop_back(algorithm_state->available_workers);
@@ -620,7 +632,7 @@ void dispatch_tasks(LocalSchedulerState *state,
    /* Dequeue the task and free the struct. */
    print_resource_info(state, elt->spec);
    DL_DELETE(algorithm_state->dispatch_task_queue, elt);
-    free_task_spec(elt->spec);
+    free(elt->spec);
    free(elt);
  } /* End for each task in the dispatch queue. */
 }
@@ -641,20 +653,22 @@ void dispatch_tasks(LocalSchedulerState *state,
 */
 task_queue_entry *queue_task(LocalSchedulerState *state,
                             task_queue_entry **task_queue,
-                             task_spec *spec,
+                             TaskSpec *spec,
+                             int64_t task_spec_size,
                             bool from_global_scheduler) {
  /* Copy the spec and add it to the task queue. The allocated spec will be
   * freed when it is assigned to a worker. */
  task_queue_entry *elt = (task_queue_entry *) malloc(sizeof(task_queue_entry));
-  elt->spec = (task_spec *) malloc(task_spec_size(spec));
-  memcpy(elt->spec, spec, task_spec_size(spec));
+  elt->spec = (TaskSpec *) malloc(task_spec_size);
+  memcpy(elt->spec, spec, task_spec_size);
+  elt->task_spec_size = task_spec_size;
  DL_APPEND((*task_queue), elt);

  /* The task has been added to a local scheduler queue. Write the entry in the
   * task table to notify others that we have queued it. */
  if (state->db != NULL) {
-    Task *task =
-        Task_alloc(spec, TASK_STATUS_QUEUED, get_db_client_id(state->db));
+    Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
+                            get_db_client_id(state->db));
    if (from_global_scheduler) {
      /* If the task is from the global scheduler, it's already been added to
       * the task table, so just update the entry. */
@@ -684,11 +698,13 @@ task_queue_entry *queue_task(LocalSchedulerState *state,
 */
 void queue_waiting_task(LocalSchedulerState *state,
                        SchedulingAlgorithmState *algorithm_state,
-                        task_spec *spec,
+                        TaskSpec *spec,
+                        int64_t task_spec_size,
                        bool from_global_scheduler) {
  LOG_DEBUG("Queueing task in waiting queue");
-  task_queue_entry *task_entry = queue_task(
-      state, &algorithm_state->waiting_task_queue, spec, from_global_scheduler);
+  task_queue_entry *task_entry =
+      queue_task(state, &algorithm_state->waiting_task_queue, spec,
+                 task_spec_size, from_global_scheduler);
  /* If we're queueing this task in the waiting queue, there must be at least
   * one missing dependency, so record it. */
  fetch_missing_dependencies(state, algorithm_state, task_entry);
@@ -707,10 +723,11 @@ void queue_waiting_task(LocalSchedulerState *state,
 */
 void queue_dispatch_task(LocalSchedulerState *state,
                         SchedulingAlgorithmState *algorithm_state,
-                         task_spec *spec,
+                         TaskSpec *spec,
+                         int64_t task_spec_size,
                         bool from_global_scheduler) {
  LOG_DEBUG("Queueing task in dispatch queue");
-  queue_task(state, &algorithm_state->dispatch_task_queue, spec,
+  queue_task(state, &algorithm_state->dispatch_task_queue, spec, task_spec_size,
             from_global_scheduler);
 }

@@ -728,14 +745,17 @@ void queue_dispatch_task(LocalSchedulerState *state,
 */
 void queue_task_locally(LocalSchedulerState *state,
                        SchedulingAlgorithmState *algorithm_state,
-                        task_spec *spec,
+                        TaskSpec *spec,
+                        int64_t task_spec_size,
                        bool from_global_scheduler) {
  if (can_run(algorithm_state, spec)) {
    /* Dependencies are ready, so push the task to the dispatch queue. */
-    queue_dispatch_task(state, algorithm_state, spec, from_global_scheduler);
+    queue_dispatch_task(state, algorithm_state, spec, task_spec_size,
+                        from_global_scheduler);
  } else {
    /* Dependencies are not ready, so push the task to the waiting queue. */
-    queue_waiting_task(state, algorithm_state, spec, from_global_scheduler);
+    queue_waiting_task(state, algorithm_state, spec, task_spec_size,
+                       from_global_scheduler);
  }
 }

@@ -751,7 +771,8 @@ void queue_task_locally(LocalSchedulerState *state,
 */
 void give_task_to_local_scheduler(LocalSchedulerState *state,
                                  SchedulingAlgorithmState *algorithm_state,
-                                  task_spec *spec,
+                                  TaskSpec *spec,
+                                  int64_t task_spec_size,
                                  DBClientID local_scheduler_id) {
  if (DBClientID_equal(local_scheduler_id, get_db_client_id(state->db))) {
    LOG_WARN("Local scheduler is trying to assign a task to itself.");
@@ -759,7 +780,8 @@ void give_task_to_local_scheduler(LocalSchedulerState *state,
  CHECK(state->db != NULL);
  /* Assign the task to the relevant local scheduler. */
  DCHECK(state->config.global_scheduler_exists);
-  Task *task = Task_alloc(spec, TASK_STATUS_SCHEDULED, local_scheduler_id);
+  Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_SCHEDULED,
+                          local_scheduler_id);
  task_table_add_task(state->db, task, NULL, NULL, NULL);
 }

@@ -773,27 +795,27 @@ void give_task_to_local_scheduler(LocalSchedulerState *state,
 */
 void give_task_to_global_scheduler(LocalSchedulerState *state,
                                   SchedulingAlgorithmState *algorithm_state,
-                                   task_spec *spec) {
+                                   TaskSpec *spec,
+                                   int64_t task_spec_size) {
  if (state->db == NULL || !state->config.global_scheduler_exists) {
    /* A global scheduler is not available, so queue the task locally. */
-    queue_task_locally(state, algorithm_state, spec, false);
+    queue_task_locally(state, algorithm_state, spec, task_spec_size, false);
    return;
  }
  /* Pass on the task to the global scheduler. */
  DCHECK(state->config.global_scheduler_exists);
-  Task *task = Task_alloc(spec, TASK_STATUS_WAITING, NIL_ID);
+  Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_WAITING, NIL_ID);
  DCHECK(state->db != NULL);
  task_table_add_task(state->db, task, NULL, NULL, NULL);
 }

 bool resource_constraints_satisfied(LocalSchedulerState *state,
-                                    task_spec *spec) {
+                                    TaskSpec *spec) {
  /* At the local scheduler, if required resource vector exceeds either static
   * or dynamic resource vector, the resource constraint is not satisfied. */
-  for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
-    if (task_spec_get_required_resource(spec, i) > state->static_resources[i] ||
-        task_spec_get_required_resource(spec, i) >
-            state->dynamic_resources[i]) {
+  for (int i = 0; i < ResourceIndex_MAX; i++) {
+    if (TaskSpec_get_required_resource(spec, i) > state->static_resources[i] ||
+        TaskSpec_get_required_resource(spec, i) > state->dynamic_resources[i]) {
      return false;
    }
  }
@@ -802,7 +824,8 @@ bool resource_constraints_satisfied(LocalSchedulerState *state,

 void handle_task_submitted(LocalSchedulerState *state,
                           SchedulingAlgorithmState *algorithm_state,
-                           task_spec *spec) {
+                           TaskSpec *spec,
+                           int64_t task_spec_size) {
  /* TODO(atumanov): if static is satisfied and local objects ready, but dynamic
   * resource is currently unavailable, then consider queueing task locally and
   * recheck dynamic next time. */
@@ -814,10 +837,10 @@ void handle_task_submitted(LocalSchedulerState *state,
  if (resource_constraints_satisfied(state, spec) &&
      (utarray_len(algorithm_state->available_workers) > 0) &&
      can_run(algorithm_state, spec)) {
-    queue_dispatch_task(state, algorithm_state, spec, false);
+    queue_dispatch_task(state, algorithm_state, spec, task_spec_size, false);
  } else {
    /* Give the task to the global scheduler to schedule, if it exists. */
-    give_task_to_global_scheduler(state, algorithm_state, spec);
+    give_task_to_global_scheduler(state, algorithm_state, spec, task_spec_size);
  }

  /* Try to dispatch tasks, since we may have added one to the queue. */
@@ -826,8 +849,9 @@ void handle_task_submitted(LocalSchedulerState *state,

 void handle_actor_task_submitted(LocalSchedulerState *state,
                                 SchedulingAlgorithmState *algorithm_state,
-                                 task_spec *spec) {
-  ActorID actor_id = task_spec_actor_id(spec);
+                                 TaskSpec *spec,
+                                 int64_t task_spec_size) {
+  ActorID actor_id = TaskSpec_actor_id(spec);
  CHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));

  /* Find the local scheduler responsible for this actor. */
@@ -840,6 +864,8 @@ void handle_actor_task_submitted(LocalSchedulerState *state,
     * will be resubmitted (internally by the local scheduler) whenever a new
     * actor notification arrives. */
    utarray_push_back(algorithm_state->cached_submitted_actor_tasks, &spec);
+    utarray_push_back(algorithm_state->cached_submitted_actor_task_sizes,
+                      &task_spec_size);
    return;
  }

@@ -847,13 +873,14 @@ void handle_actor_task_submitted(LocalSchedulerState *state,
                       get_db_client_id(state->db))) {
    /* This local scheduler is responsible for the actor, so handle the task
     * locally. */
-    add_task_to_actor_queue(state, algorithm_state, spec, false);
+    add_task_to_actor_queue(state, algorithm_state, spec, task_spec_size,
+                            false);
    /* Attempt to dispatch tasks to this actor. */
    dispatch_actor_task(state, algorithm_state, actor_id);
  } else {
    /* This local scheduler is not responsible for the task, so assign the task
     * directly to the actor that is responsible. */
-    give_task_to_local_scheduler(state, algorithm_state, spec,
+    give_task_to_local_scheduler(state, algorithm_state, spec, task_spec_size,
                                 entry->local_scheduler_id);
  }
 }
@@ -864,35 +891,43 @@ void handle_actor_creation_notification(
    ActorID actor_id) {
  int num_cached_actor_tasks =
      utarray_len(algorithm_state->cached_submitted_actor_tasks);
+  CHECK(num_cached_actor_tasks ==
+        utarray_len(algorithm_state->cached_submitted_actor_task_sizes));
  for (int i = 0; i < num_cached_actor_tasks; ++i) {
-    task_spec **spec = (task_spec **) utarray_eltptr(
+    TaskSpec **spec = (TaskSpec **) utarray_eltptr(
        algorithm_state->cached_submitted_actor_tasks, i);
+    int64_t *task_spec_size = (int64_t *) utarray_eltptr(
+        algorithm_state->cached_submitted_actor_task_sizes, i);
    /* Note that handle_actor_task_submitted may append the spec to the end of
     * the cached_submitted_actor_tasks array. */
-    handle_actor_task_submitted(state, algorithm_state, *spec);
+    handle_actor_task_submitted(state, algorithm_state, *spec, *task_spec_size);
  }
  /* Remove all the tasks that were resubmitted. This does not erase the tasks
   * that were newly appended to the cached_submitted_actor_tasks array. */
  utarray_erase(algorithm_state->cached_submitted_actor_tasks, 0,
                num_cached_actor_tasks);
+  utarray_erase(algorithm_state->cached_submitted_actor_task_sizes, 0,
+                num_cached_actor_tasks);
 }

 void handle_task_scheduled(LocalSchedulerState *state,
                           SchedulingAlgorithmState *algorithm_state,
-                           task_spec *spec) {
+                           TaskSpec *spec,
+                           int64_t task_spec_size) {
  /* This callback handles tasks that were assigned to this local scheduler by
   * the global scheduler, so we can safely assert that there is a connection to
   * the database. */
  DCHECK(state->db != NULL);
  DCHECK(state->config.global_scheduler_exists);
  /* Push the task to the appropriate queue. */
-  queue_task_locally(state, algorithm_state, spec, true);
+  queue_task_locally(state, algorithm_state, spec, task_spec_size, true);
  dispatch_tasks(state, algorithm_state);
 }

 void handle_actor_task_scheduled(LocalSchedulerState *state,
                                 SchedulingAlgorithmState *algorithm_state,
-                                 task_spec *spec) {
+                                 TaskSpec *spec,
+                                 int64_t task_spec_size) {
  /* This callback handles tasks that were assigned to this local scheduler by
   * the global scheduler or by other workers, so we can safely assert that
   * there is a connection to the database. */
@@ -900,7 +935,7 @@ void handle_actor_task_scheduled(LocalSchedulerState *state,
  DCHECK(state->config.global_scheduler_exists);
  /* Check that the task is meant to run on an actor that this local scheduler
   * is responsible for. */
-  ActorID actor_id = task_spec_actor_id(spec);
+  ActorID actor_id = TaskSpec_actor_id(spec);
  DCHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
  actor_map_entry *entry;
  HASH_FIND(hh, state->actor_mapping, &actor_id, sizeof(actor_id), entry);
@@ -918,7 +953,7 @@ void handle_actor_task_scheduled(LocalSchedulerState *state,
        "corresponding actor_map_entry is not present. This should be rare.");
  }
  /* Push the task to the appropriate queue. */
-  add_task_to_actor_queue(state, algorithm_state, spec, true);
+  add_task_to_actor_queue(state, algorithm_state, spec, task_spec_size, true);
  dispatch_actor_task(state, algorithm_state, actor_id);
 }

@@ -1047,7 +1082,7 @@ void handle_worker_blocked(LocalSchedulerState *state,

      /* Return the resources that the blocked worker was using. */
      CHECK(worker->task_in_progress != NULL);
-      task_spec *spec = Task_task_spec(worker->task_in_progress);
+      TaskSpec *spec = Task_task_spec(worker->task_in_progress);
      update_dynamic_resources(state, spec, true);
      /* Add the worker to the list of blocked workers. */
      worker->is_blocked = true;
@@ -1089,7 +1124,7 @@ void handle_worker_unblocked(LocalSchedulerState *state,
       * fixed by having blocked workers explicitly yield and wait to be given
       * back resources before continuing execution. */
      CHECK(worker->task_in_progress != NULL);
-      task_spec *spec = Task_task_spec(worker->task_in_progress);
+      TaskSpec *spec = Task_task_spec(worker->task_in_progress);
      update_dynamic_resources(state, spec, false);
      /* Add the worker to the list of executing workers. */
      worker->is_blocked = false;
@@ -1171,11 +1206,11 @@ void handle_object_removed(LocalSchedulerState *state,
  task_queue_entry *elt, *tmp;
  /* Track the dependency for tasks that were in the waiting queue. */
  DL_FOREACH(algorithm_state->waiting_task_queue, elt) {
-    task_spec *task = elt->spec;
-    int64_t num_args = task_num_args(task);
+    TaskSpec *task = elt->spec;
+    int64_t num_args = TaskSpec_num_args(task);
    for (int i = 0; i < num_args; ++i) {
-      if (task_arg_type(task, i) == ARG_BY_REF) {
-        ObjectID arg_id = task_arg_id(task, i);
+      if (TaskSpec_arg_by_ref(task, i)) {
+        ObjectID arg_id = TaskSpec_arg_id(task, i);
        if (ObjectID_equal(arg_id, removed_object_id)) {
          fetch_missing_dependency(state, algorithm_state, elt,
                                   removed_object_id);
@@ -1186,11 +1221,11 @@ void handle_object_removed(LocalSchedulerState *state,
  /* Track the dependency for tasks that were in the dispatch queue. Remove
   * these tasks from the dispatch queue and push them to the waiting queue. */
  DL_FOREACH_SAFE(algorithm_state->dispatch_task_queue, elt, tmp) {
-    task_spec *task = elt->spec;
-    int64_t num_args = task_num_args(task);
+    TaskSpec *task = elt->spec;
+    int64_t num_args = TaskSpec_num_args(task);
    for (int i = 0; i < num_args; ++i) {
-      if (task_arg_type(task, i) == ARG_BY_REF) {
-        ObjectID arg_id = task_arg_id(task, i);
+      if (TaskSpec_arg_by_ref(task, i)) {
+        ObjectID arg_id = TaskSpec_arg_id(task, i);
        if (ObjectID_equal(arg_id, removed_object_id)) {
          LOG_DEBUG("Moved task from dispatch queue back to waiting queue");
          DL_DELETE(algorithm_state->dispatch_task_queue, elt);