autoscaler fix

Signed-off-by: Jonathan Nitisastro <[email protected]>

Author: jonathan-anyscale

python/ray/autoscaler/_private/fake_multi_node/node_provider.py

@@ -316,7 +316,7 @@ def create_node_with_resources_and_labels(
                 num_cpus=resources.pop("CPU", 0),
                 num_gpus=resources.pop("GPU", 0),
                 object_store_memory=resources.pop("object_store_memory", None),
-                _gpu_memory=resources.pop("gpu_memory", 0),
+                _gpu_memory=resources.pop("gpu_memory", 0),  # gpu memory = 600,
                 resources=resources,
                 labels=labels,
                 redis_address="{}:6379".format(

python/ray/autoscaler/_private/resource_demand_scheduler.py

@@ -545,7 +545,12 @@ def add_node(node_type, available_resources=None):
             if TAG_RAY_USER_NODE_TYPE in tags:
                 node_type = tags[TAG_RAY_USER_NODE_TYPE]
                 ip = self.provider.internal_ip(node_id)
-                available_resources = unused_resources_by_ip.get(ip)
+                available_resources = copy.deepcopy(unused_resources_by_ip.get(ip))
+                available = self.node_types[node_type]["resources"]
+                if available_resources and "node:gpu_memory_per_gpu" in available:
+                    available_resources["node:gpu_memory_per_gpu"] = available[
+                        "node:gpu_memory_per_gpu"
+                    ]
                 add_node(node_type, available_resources)
 
         for node_type, count in pending_nodes.items():

python/ray/tests/test_autoscaler_fake_multinode.py

@@ -25,6 +25,7 @@ def test_fake_autoscaler_basic_e2e(shutdown_only):
                     "CPU": 2,
                     "GPU": 1,
                     "object_store_memory": 1024 * 1024 * 1024,
+                    "gpu_memory": 1000,
                 },
                 "node_config": {},
                 "min_workers": 0,
@@ -52,6 +53,11 @@ def test_fake_autoscaler_basic_e2e(shutdown_only):
         def f():
             print("gpu ok")
 
+        # Triggers the addition of a GPU memory node.
+        @ray.remote(_gpu_memory=1000)
+        def f2():
+            print("gpu memory ok")
+
         # Triggers the addition of a CPU node.
         @ray.remote(num_cpus=3)
         def g():

src/ray/common/scheduling/cluster_resource_data.cc

@@ -20,6 +20,27 @@
 namespace ray {
 using namespace ::ray::scheduling;
 
+NodeResources::NodeResources(const NodeResourceSet &resources) {
+  NodeResourceSet resources_adjusted = resources;
+  absl::flat_hash_map<std::string, std::string> node_labels;
+  if (resources.Has(ResourceID::GPU_Memory())) {
+    // if gpu_memory is set, default GPU value is 1
+    if (!resources.Has(ResourceID::GPU())) {
+      resources_adjusted.Set(ResourceID::GPU(), 1);
+    }
+    node_labels["_gpu_memory_per_gpu"] =
+        std::to_string(resources.Get(ResourceID::GPU_Memory()).Double() /
+                       resources_adjusted.Get(ResourceID::GPU()).Double());
+    resources_adjusted.Set(ResourceID::GPU_Memory(), 0);
+  } else {
+    node_labels["_gpu_memory_per_gpu"] = "0";
+  }
+
+  this->total = resources_adjusted;
+  this->available = resources_adjusted;
+  this->labels = node_labels;
+}
+
 /// Convert a map of resources to a ResourceRequest data structure.
 ResourceRequest ResourceMapToResourceRequest(
     const absl::flat_hash_map<std::string, double> &resource_map,
@@ -59,8 +80,12 @@ NodeResources ResourceMapToNodeResources(
   auto node_labels_copy = node_labels;
 
   if (resource_map_total.find("gpu_memory") != resource_map_total.end()) {
+    // if gpu_memory is set, default GPU value is 1
+    if (resource_map_total.find("GPU") == resource_map_total.end()) {
+      resource_map_total_copy["GPU"] = 1;
+    }
     node_labels_copy["_gpu_memory_per_gpu"] = std::to_string(
-        resource_map_total.at("gpu_memory") / resource_map_total.at("GPU"));
+        resource_map_total.at("gpu_memory") / resource_map_total_copy.at("GPU"));
     resource_map_total_copy.erase("gpu_memory");
     resource_map_available_copy.erase("gpu_memory");
   } else {
@@ -108,6 +133,7 @@ bool NodeResources::IsAvailable(const ResourceRequest &resource_request,
   }
   const ResourceSet resource_request_adjusted =
       this->ConvertRelativeResources(resource_request.GetResourceSet());
+
   if (!this->normal_task_resources.IsEmpty()) {
     auto available_resources = this->available;
     available_resources -= this->normal_task_resources;
@@ -159,6 +185,8 @@ const ResourceSet NodeResources::ConvertRelativeResources(
       num_gpus_request =
           (resource.Get(ResourceID::GPU_Memory()).Double() / total_gpu_memory_per_gpu) +
           1 / static_cast<double>(2 * kResourceUnitScaling);
+    } else {
+      return resource;
     }
     adjusted_resource.Set(ResourceID::GPU(), num_gpus_request);
     adjusted_resource.Set(ResourceID::GPU_Memory(), 0);
@@ -208,6 +236,8 @@ const ResourceSet NodeResourceInstances::ConvertRelativeResources(
       num_gpus_request =
           (resource.Get(ResourceID::GPU_Memory()).Double() / total_gpu_memory_per_gpu) +
           1 / static_cast<double>(2 * kResourceUnitScaling);
+    } else {
+      return resource;
     }
     adjusted_resource.Set(ResourceID::GPU(), num_gpus_request);
     adjusted_resource.Set(ResourceID::GPU_Memory(), 0);

src/ray/common/scheduling/cluster_resource_data.h

@@ -289,8 +289,8 @@ class TaskResourceInstances {
 class NodeResources {
  public:
   NodeResources() {}
-  NodeResources(const NodeResourceSet &resources)
-      : total(resources), available(resources) {}
+  NodeResources(const NodeResourceSet &resources);
+
   NodeResourceSet total;
   NodeResourceSet available;
   /// Only used by light resource report.

src/ray/raylet/scheduling/cluster_resource_manager.cc

@@ -88,6 +88,7 @@ bool ClusterResourceManager::UpdateNode(
 
   local_view.total = node_resources.total;
   local_view.available = node_resources.available;
+  // might need label transfer here
   local_view.object_pulls_queued = resource_view_sync_message.object_pulls_queued();
 
   // Update the idle duration for the node in terms of resources usage.
@@ -187,10 +188,10 @@ bool ClusterResourceManager::SubtractNodeAvailableResources(
   }
 
   NodeResources *resources = it->second.GetMutableLocalView();
-  const ResourceSet adjusted_resource_request =
+  const ResourceSet resource_request_adjusted =
       resources->ConvertRelativeResources(resource_request.GetResourceSet());
 
-  resources->available -= adjusted_resource_request;
+  resources->available -= resource_request_adjusted;
   resources->available.RemoveNegative();
 
   // TODO(swang): We should also subtract object store memory if the task has
@@ -216,10 +217,10 @@ bool ClusterResourceManager::HasSufficientResource(
     return false;
   }
 
-  const ResourceSet adjusted_resource_request =
+  const ResourceSet resource_request_adjusted =
       resources.ConvertRelativeResources(resource_request.GetResourceSet());
 
-  return resources.available >= adjusted_resource_request;
+  return resources.available >= resource_request_adjusted;
 }
 
 bool ClusterResourceManager::AddNodeAvailableResources(scheduling::NodeID node_id,
@@ -228,13 +229,15 @@ bool ClusterResourceManager::AddNodeAvailableResources(scheduling::NodeID node_i
   if (it == nodes_.end()) {
     return false;
   }
-
   auto node_resources = it->second.GetMutableLocalView();
-  for (auto &resource_id : resource_set.ResourceIds()) {
+  const ResourceSet adjusted_resources =
+      node_resources->ConvertRelativeResources(resource_set);
+
+  for (auto &resource_id : adjusted_resources.ResourceIds()) {
     if (node_resources->total.Has(resource_id)) {
       auto available = node_resources->available.Get(resource_id);
       auto total = node_resources->total.Get(resource_id);
-      auto new_available = available + resource_set.Get(resource_id);
+      auto new_available = available + adjusted_resources.Get(resource_id);
       if (new_available > total) {
         new_available = total;
       }

src/ray/raylet/scheduling/local_resource_manager.cc

@@ -80,17 +80,18 @@ bool LocalResourceManager::AllocateTaskResourceInstances(
     const ResourceRequest &resource_request,
     std::shared_ptr<TaskResourceInstances> task_allocation) {
   RAY_CHECK(task_allocation != nullptr);
-  const ResourceSet adjusted_resource_request =
+  const ResourceSet resource_request_adjusted =
       local_resources_.ConvertRelativeResources(resource_request.GetResourceSet());
+
   if (resource_request.GetResourceSet().Has(ResourceID::GPU_Memory()) &&
-      adjusted_resource_request.Get(ResourceID::GPU()) > 1) {
+      resource_request_adjusted.Get(ResourceID::GPU()) > 1) {
     return false;
   }
   // add adjust_gpu_memory here, added to NodeInstanceResourceSet
-  auto allocation = local_resources_.available.TryAllocate(adjusted_resource_request);
+  auto allocation = local_resources_.available.TryAllocate(resource_request_adjusted);
   if (allocation) {
     *task_allocation = TaskResourceInstances(*allocation);
-    for (const auto &resource_id : adjusted_resource_request.ResourceIds()) {
+    for (const auto &resource_id : resource_request_adjusted.ResourceIds()) {
       SetResourceNonIdle(resource_id);
     }
     return true;

src/ray/raylet/scheduling/policy/scorer.cc

@@ -35,8 +35,11 @@ double LeastResourceScorer::Score(const ResourceRequest &required_resources,
   }
 
   double node_score = 0.;
-  for (auto &resource_id : required_resources.ResourceIds()) {
-    const auto &request_resource = required_resources.Get(resource_id);
+  const ResourceSet resource_request_adjusted =
+      node_resources_ptr->ConvertRelativeResources(required_resources.GetResourceSet());
+
+  for (auto &resource_id : resource_request_adjusted.ResourceIds()) {
+    const auto &request_resource = resource_request_adjusted.Get(resource_id);
     const auto &node_available_resource = node_resources_ptr->available.Get(resource_id);
     auto score = Calculate(request_resource, node_available_resource);
     if (score < 0.) {

src/ray/raylet/scheduling/scheduling_policy.cc

@@ -24,7 +24,8 @@ namespace raylet_scheduling_policy {
 namespace {
 
 bool IsGPURequest(const ResourceRequest &resource_request) {
-  return resource_request.Has(ResourceID::GPU());
+  return resource_request.Has(ResourceID::GPU()) ||
+         resource_request.Has(ResourceID::GPU_Memory());
 }
 
 bool DoesNodeHaveGPUs(const NodeResources &resources) {