Adding basic support for a user-interpretable resource label (#761)

* adding support for the user-interpretable label(UIR)

* more plumbing for num_uirs further upstream; set to infty when specified on cmd line

* pass default num_uirs for actors; update GlobalStateAPI

* support num_uirs in ray.init()

* local scheduler resource accounting: support num_uirs; prep for vectorized resource accounting

* global scheduler test updated

* Fix bug introduced by rebase.

* Rename UIR -> CustomResource and add test.

* Small changes and use constexpr instead of macros.

* Linting and some renaming.

* Reorder some code.

* Remove cpus_in_use and fix bug.

* Add another test and make a small change.

* Rephrase documentation about feature stability.

src/common/format/common.fbs

@@ -9,9 +9,11 @@ enum ResourceIndex:int {
   CPU = 0,
   // A graphics processing unit.
   GPU = 1,
+  // A user-defined custom resource.
+  CustomResource = 2,
   // A dummy entry to make ResourceIndex_MAX equal to the length of
   // a resource vector.
-  DUMMY = 2
+  DUMMY = 3
 }
 
 table Arg {

src/local_scheduler/local_scheduler.cc

@@ -119,8 +119,9 @@ void kill_worker(LocalSchedulerState *state,
   }
 
   /* Release any resources held by the worker. */
-  release_resources(state, worker, worker->cpus_in_use,
-                    worker->gpus_in_use.size());
+  release_resources(state, worker, worker->resources_in_use[ResourceIndex_CPU],
+                    worker->gpus_in_use.size(),
+                    worker->resources_in_use[ResourceIndex_CustomResource]);
 
   /* Clean up the task in progress. */
   if (worker->task_in_progress) {
@@ -412,35 +413,46 @@ LocalSchedulerState *LocalSchedulerState_init(
   return state;
 }
 
+/* TODO(atumanov): vectorize resource counts on input. */
 bool check_dynamic_resources(LocalSchedulerState *state,
                              double num_cpus,
-                             double num_gpus) {
+                             double num_gpus,
+                             double num_custom_resource) {
   if (num_cpus > 0 && state->dynamic_resources[ResourceIndex_CPU] < num_cpus) {
     /* We only use this check when num_cpus is positive so that we can still
      * create actors even when the CPUs are oversubscribed. */
     return false;
   }
+  if (num_custom_resource > 0 &&
+      state->dynamic_resources[ResourceIndex_CustomResource] <
+          num_custom_resource) {
+    return false;
+  }
   if (state->dynamic_resources[ResourceIndex_GPU] < num_gpus) {
     return false;
   }
   return true;
 }
 
+/* TODO(atumanov): just pass the required resource vector of doubles. */
 void acquire_resources(LocalSchedulerState *state,
                        LocalSchedulerClient *worker,
                        double num_cpus,
-                       double num_gpus) {
+                       double num_gpus,
+                       double num_custom_resource) {
   /* Acquire the CPU resources. */
   bool oversubscribed = (state->dynamic_resources[ResourceIndex_CPU] < 0);
   state->dynamic_resources[ResourceIndex_CPU] -= num_cpus;
-  CHECK(worker->cpus_in_use == 0);
-  worker->cpus_in_use += num_cpus;
+  CHECK(worker->resources_in_use[ResourceIndex_CPU] == 0);
+  worker->resources_in_use[ResourceIndex_CPU] += num_cpus;
   /* Log a warning if we are using more resources than we have been allocated,
    * and we weren't already oversubscribed. */
   if (!oversubscribed && state->dynamic_resources[ResourceIndex_CPU] < 0) {
-    LOG_WARN("local_scheduler dynamic resources dropped to %8.4f\t%8.4f\n",
-             state->dynamic_resources[ResourceIndex_CPU],
-             state->dynamic_resources[ResourceIndex_GPU]);
+    LOG_WARN(
+        "local_scheduler dynamic resources dropped to %8.4f\t%8.4f\t%8.4f\n",
+        state->dynamic_resources[ResourceIndex_CPU],
+        state->dynamic_resources[ResourceIndex_GPU],
+        state->dynamic_resources[ResourceIndex_CustomResource]);
   }
 
   /* Acquire the GPU resources. */
@@ -457,16 +469,22 @@ void acquire_resources(LocalSchedulerState *state,
     CHECK(state->dynamic_resources[ResourceIndex_GPU] >= num_gpus);
     state->dynamic_resources[ResourceIndex_GPU] -= num_gpus;
   }
+
+  /* Acquire the custom resources. */
+  state->dynamic_resources[ResourceIndex_CustomResource] -= num_custom_resource;
+  CHECK(worker->resources_in_use[ResourceIndex_CustomResource] == 0);
+  worker->resources_in_use[ResourceIndex_CustomResource] += num_custom_resource;
 }
 
 void release_resources(LocalSchedulerState *state,
                        LocalSchedulerClient *worker,
                        double num_cpus,
-                       double num_gpus) {
+                       double num_gpus,
+                       double num_custom_resource) {
   /* Release the CPU resources. */
-  CHECK(num_cpus == worker->cpus_in_use);
+  CHECK(num_cpus == worker->resources_in_use[ResourceIndex_CPU]);
   state->dynamic_resources[ResourceIndex_CPU] += num_cpus;
-  worker->cpus_in_use = 0;
+  worker->resources_in_use[ResourceIndex_CPU] = 0;
 
   /* Release the GPU resources. */
   if (num_gpus != 0) {
@@ -478,6 +496,12 @@ void release_resources(LocalSchedulerState *state,
     worker->gpus_in_use.clear();
     state->dynamic_resources[ResourceIndex_GPU] += num_gpus;
   }
+
+  /* Release the user-defined custom resource. */
+  CHECK(num_custom_resource ==
+        worker->resources_in_use[ResourceIndex_CustomResource]);
+  state->dynamic_resources[ResourceIndex_CustomResource] += num_custom_resource;
+  worker->resources_in_use[ResourceIndex_CustomResource] = 0;
 }
 
 bool is_driver_alive(LocalSchedulerState *state, WorkerID driver_id) {
@@ -489,9 +513,10 @@ void assign_task_to_worker(LocalSchedulerState *state,
                            int64_t task_spec_size,
                            LocalSchedulerClient *worker) {
   /* Acquire the necessary resources for running this task. */
-  acquire_resources(state, worker,
-                    TaskSpec_get_required_resource(spec, ResourceIndex_CPU),
-                    TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
+  acquire_resources(
+      state, worker, TaskSpec_get_required_resource(spec, ResourceIndex_CPU),
+      TaskSpec_get_required_resource(spec, ResourceIndex_GPU),
+      TaskSpec_get_required_resource(spec, ResourceIndex_CustomResource));
   /* Check that actor tasks don't have GPU requirements. Any necessary GPUs
    * should already have been acquired by the actor worker. */
   if (!ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
@@ -542,16 +567,20 @@ void finish_task(LocalSchedulerState *state, LocalSchedulerClient *worker) {
   if (worker->task_in_progress != NULL) {
     TaskSpec *spec = Task_task_spec(worker->task_in_progress);
     /* Return dynamic resources back for the task in progress. */
-    CHECK(worker->cpus_in_use ==
+    CHECK(worker->resources_in_use[ResourceIndex_CPU] ==
           TaskSpec_get_required_resource(spec, ResourceIndex_CPU));
     if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
       CHECK(worker->gpus_in_use.size() ==
             TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
-      release_resources(state, worker, worker->cpus_in_use,
-                        worker->gpus_in_use.size());
+      release_resources(state, worker,
+                        worker->resources_in_use[ResourceIndex_CPU],
+                        worker->gpus_in_use.size(),
+                        worker->resources_in_use[ResourceIndex_CustomResource]);
     } else {
       CHECK(0 == TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
-      release_resources(state, worker, worker->cpus_in_use, 0);
+      release_resources(state, worker,
+                        worker->resources_in_use[ResourceIndex_CPU], 0,
+                        worker->resources_in_use[ResourceIndex_CustomResource]);
     }
     /* If we're connected to Redis, update tables. */
     if (state->db != NULL) {
@@ -781,8 +810,8 @@ void handle_client_register(LocalSchedulerState *state,
       /* If there are enough GPUs available, allocate them and reply to the
        * actor. */
       double num_gpus_required = (double) message->num_gpus();
-      if (check_dynamic_resources(state, 0, num_gpus_required)) {
-        acquire_resources(state, worker, 0, num_gpus_required);
+      if (check_dynamic_resources(state, 0, num_gpus_required, 0)) {
+        acquire_resources(state, worker, 0, num_gpus_required, 0);
       } else {
         /* TODO(rkn): This means that an actor wants to register but that there
          * aren't enough GPUs for it. We should queue this request, and reply to
@@ -959,7 +988,9 @@ void process_message(event_loop *loop,
       worker->is_blocked = true;
       /* Return the CPU resources that the blocked worker was using, but not
        * GPU resources. */
-      release_resources(state, worker, worker->cpus_in_use, 0);
+      release_resources(state, worker,
+                        worker->resources_in_use[ResourceIndex_CPU], 0,
+                        worker->resources_in_use[ResourceIndex_CustomResource]);
       /* Let the scheduling algorithm process the fact that the worker is
        * blocked. */
       if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
@@ -989,9 +1020,10 @@ void process_message(event_loop *loop,
        * workers explicitly yield and wait to be given back resources before
        * continuing execution. */
       TaskSpec *spec = Task_task_spec(worker->task_in_progress);
-      acquire_resources(state, worker,
-                        TaskSpec_get_required_resource(spec, ResourceIndex_CPU),
-                        0);
+      acquire_resources(
+          state, worker,
+          TaskSpec_get_required_resource(spec, ResourceIndex_CPU), 0,
+          TaskSpec_get_required_resource(spec, ResourceIndex_CustomResource));
       /* Let the scheduling algorithm process the fact that the worker is
        * unblocked. */
       if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
@@ -1039,7 +1071,7 @@ void new_client_connection(event_loop *loop,
   worker->is_worker = true;
   worker->client_id = NIL_WORKER_ID;
   worker->task_in_progress = NULL;
-  worker->cpus_in_use = 0;
+  memset(&worker->resources_in_use[0], 0, sizeof(double) * ResourceIndex_MAX);
   worker->is_blocked = false;
   worker->pid = 0;
   worker->is_child = false;
@@ -1316,9 +1348,14 @@ int main(int argc, char *argv[]) {
   if (!static_resource_list) {
     /* Use defaults for this node's static resource configuration. */
     memset(&static_resource_conf[0], 0, sizeof(static_resource_conf));
-    static_resource_conf[ResourceIndex_CPU] = DEFAULT_NUM_CPUS;
-    static_resource_conf[ResourceIndex_GPU] = DEFAULT_NUM_GPUS;
+    /* TODO(atumanov): Define a default vector and replace individual
+     * constants. */
+    static_resource_conf[ResourceIndex_CPU] = kDefaultNumCPUs;
+    static_resource_conf[ResourceIndex_GPU] = kDefaultNumGPUs;
+    static_resource_conf[ResourceIndex_CustomResource] =
+        kDefaultNumCustomResource;
   } else {
+    /* TODO(atumanov): Switch this tokenizer to reading from ifstream. */
     /* Tokenize the string. */
     const char delim[2] = ",";
     char *token;
@@ -1329,6 +1366,12 @@ int main(int argc, char *argv[]) {
       /* Attempt to get the next token. */
       token = strtok(NULL, delim);
     }
+    if (static_resource_conf[ResourceIndex_CustomResource] < 0) {
+      /* Interpret negative values for the custom resource as deferring to the
+       * default system configuration. */
+      static_resource_conf[ResourceIndex_CustomResource] =
+          kDefaultNumCustomResource;
+    }
   }
   if (!scheduler_socket_name) {
     LOG_FATAL("please specify socket for incoming connections with -s switch");

src/local_scheduler/local_scheduler.h

@@ -1,5 +1,6 @@
 #ifndef LOCAL_SCHEDULER_H
 #define LOCAL_SCHEDULER_H
+#include <math.h>
 
 #include "task.h"
 #include "event_loop.h"
@@ -8,8 +9,9 @@
  * worker SIGKILL. */
 #define KILL_WORKER_TIMEOUT_MILLISECONDS 100
 
-#define DEFAULT_NUM_CPUS INT16_MAX
-#define DEFAULT_NUM_GPUS 0
+constexpr double kDefaultNumCPUs = INT16_MAX;
+constexpr double kDefaultNumGPUs = 0;
+constexpr double kDefaultNumCustomResource = INFINITY;
 
 /**
  * Establish a connection to a new client.
@@ -133,7 +135,8 @@ void start_worker(LocalSchedulerState *state,
  */
 bool check_dynamic_resources(LocalSchedulerState *state,
                              double num_cpus,
-                             double num_gpus);
+                             double num_gpus,
+                             double num_custom_resource);
 
 /**
  * Acquire additional resources (CPUs and GPUs) for a worker.
@@ -147,7 +150,8 @@ bool check_dynamic_resources(LocalSchedulerState *state,
 void acquire_resources(LocalSchedulerState *state,
                        LocalSchedulerClient *worker,
                        double num_cpus,
-                       double num_gpus);
+                       double num_gpus,
+                       double num_custom_resource);
 
 /**
  * Return resources (CPUs and GPUs) being used by a worker to the local
@@ -162,7 +166,8 @@ void acquire_resources(LocalSchedulerState *state,
 void release_resources(LocalSchedulerState *state,
                        LocalSchedulerClient *worker,
                        double num_cpus,
-                       double num_gpus);
+                       double num_gpus,
+                       double num_custom_resource);
 
 /** The following methods are for testing purposes only. */
 #ifdef LOCAL_SCHEDULER_TEST

src/local_scheduler/local_scheduler_algorithm.cc

@@ -314,9 +314,11 @@ bool dispatch_actor_task(LocalSchedulerState *state,
   /* If there are not enough resources available, we cannot assign the task. */
   CHECK(0 ==
         TaskSpec_get_required_resource(first_task.spec, ResourceIndex_GPU));
-  if (!check_dynamic_resources(state, TaskSpec_get_required_resource(
-                                          first_task.spec, ResourceIndex_CPU),
-                               0)) {
+  if (!check_dynamic_resources(
+          state,
+          TaskSpec_get_required_resource(first_task.spec, ResourceIndex_CPU), 0,
+          TaskSpec_get_required_resource(first_task.spec,
+                                         ResourceIndex_CustomResource))) {
     return false;
   }
   /* Assign the first task in the task queue to the worker and mark the worker
@@ -696,7 +698,9 @@ void dispatch_tasks(LocalSchedulerState *state,
     /* Skip to the next task if this task cannot currently be satisfied. */
     if (!check_dynamic_resources(
             state, TaskSpec_get_required_resource(task.spec, ResourceIndex_CPU),
-            TaskSpec_get_required_resource(task.spec, ResourceIndex_GPU))) {
+            TaskSpec_get_required_resource(task.spec, ResourceIndex_GPU),
+            TaskSpec_get_required_resource(task.spec,
+                                           ResourceIndex_CustomResource))) {
       /* This task could not be satisfied -- proceed to the next task. */
       ++it;
       continue;
@@ -924,8 +928,9 @@ bool resource_constraints_satisfied(LocalSchedulerState *state,
   /* 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 < ResourceIndex_MAX; i++) {
-    if (TaskSpec_get_required_resource(spec, i) > state->static_resources[i] ||
-        TaskSpec_get_required_resource(spec, i) > state->dynamic_resources[i]) {
+    double required_resource = TaskSpec_get_required_resource(spec, i);
+    if (required_resource > state->static_resources[i] ||
+        required_resource > state->dynamic_resources[i]) {
       return false;
     }
   }

src/local_scheduler/local_scheduler_shared.h

@@ -97,10 +97,8 @@ struct LocalSchedulerClient {
    *  no task is running on the worker, this will be NULL. This is used to
    *  update the task table. */
   Task *task_in_progress;
-  /** The number of CPUs that the worker is currently using. This will only be
-   *  nonzero when the worker is actively executing a task. If the worker is
-   *  blocked, then this value will be zero. */
-  double cpus_in_use;
+  /** An array of resource counts currently in use by the worker.  */
+  double resources_in_use[ResourceIndex_MAX];
   /** A vector of the IDs of the GPUs that the worker is currently using. If the
    *  worker is an actor, this will be constant throughout the lifetime of the
    *  actor (and will be equal to the number of GPUs requested by the actor). If

src/local_scheduler/test/local_scheduler_tests.cc

@@ -75,8 +75,8 @@ LocalSchedulerMock *LocalSchedulerMock_init(int num_workers,
   const char *node_ip_address = "127.0.0.1";
   const char *redis_addr = node_ip_address;
   int redis_port = 6379;
-  const double static_resource_conf[ResourceIndex_MAX] = {DEFAULT_NUM_CPUS,
-                                                          DEFAULT_NUM_GPUS};
+  const double static_resource_conf[ResourceIndex_MAX] = {kDefaultNumCPUs,
+                                                          kDefaultNumGPUs};
   LocalSchedulerMock *mock =
       (LocalSchedulerMock *) malloc(sizeof(LocalSchedulerMock));
   memset(mock, 0, sizeof(LocalSchedulerMock));