add request only resource for gpu memory

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

Author: jonathan-anyscale

src/ray/common/ray_config_def.h

@@ -718,6 +718,12 @@ RAY_CONFIG(std::string, predefined_unit_instance_resources, "GPU")
 /// When set it to "neuron_cores,TPU,FPGA", we will also treat FPGA as unit_instance.
 RAY_CONFIG(std::string, custom_unit_instance_resources, "neuron_cores,TPU")
 
+/// The scheduler will treat these resource as resource which can be requested
+/// but not stored as NodeResources. The main reason is the resource is different
+/// representation of other resource stored in NodeResources.
+/// For example: gpu_memory and GPU.
+RAY_CONFIG(std::string, request_only_resources, "gpu_memory")
+
 // Maximum size of the batches when broadcasting resources to raylet.
 RAY_CONFIG(uint64_t, resource_broadcast_batch_size, 512)
 

src/ray/common/scheduling/cluster_resource_data.cc

@@ -59,10 +59,9 @@ NodeResources ResourceMapToNodeResources(
   auto node_labels_copy = node_labels;
   // move gpu_memory to node labels
   if (resource_map_total.find("gpu_memory") != resource_map_total.end()) {
-    RAY_LOG(INFO) << resource_map_total.at("gpu_memory");
-    node_labels_copy["gpu_memory"] =
-        std::to_string(resource_map_total.at("gpu_memory") / resource_map_total.at("GPU"));
-    //RAY_CHECK(std::stod(node_labels_copy.at("gpu_memory")) == 1000);
+    node_labels_copy["gpu_memory"] = std::to_string(resource_map_total.at("gpu_memory") /
+                                                    resource_map_total.at("GPU"));
+    // RAY_CHECK(std::stod(node_labels_copy.at("gpu_memory")) == 1000);
     resource_map_total_copy.erase("gpu_memory");
     resource_map_available_copy.erase("gpu_memory");
   } else {
@@ -152,6 +151,7 @@ const ResourceSet NodeResources::ConvertRelativeResource(
   if (resource.Has(ResourceID::GPU_Memory())) {
     double total_gpu_memory = 0;
     if (this->labels.find("gpu_memory") != this->labels.end()) {
+      // TODO: raise exception if this is not true
       total_gpu_memory = std::stod(this->labels.at("gpu_memory"));
     }
     double num_gpus_request = 0;

src/ray/common/scheduling/scheduling_ids.cc

@@ -113,6 +113,25 @@ absl::flat_hash_set<int64_t> &ResourceID::UnitInstanceResources() {
   return set;
 }
 
+absl::flat_hash_set<int64_t> &ResourceID::RequestOnlyResources() {
+  static absl::flat_hash_set<int64_t> set{[]() {
+    absl::flat_hash_set<int64_t> res;
+
+    std::string request_only_resources = RayConfig::instance().request_only_resources();
+    if (!request_only_resources.empty()) {
+      std::vector<std::string> results;
+      boost::split(results, request_only_resources, boost::is_any_of(","));
+      for (std::string &result : results) {
+        int64_t resource_id = ResourceID(result).ToInt();
+        res.insert(resource_id);
+      }
+    }
+
+    return res;
+  }()};
+  return set;
+}
+
 }  // namespace scheduling
 
 }  // namespace ray

src/ray/common/scheduling/scheduling_ids.h

@@ -178,6 +178,8 @@ class ResourceID : public BaseSchedulingID<SchedulingIDTag::Resource> {
     return !IsPredefinedResource() && absl::StartsWith(Binary(), kImplicitResourcePrefix);
   }
 
+  bool IsRequestOnlyResource() const { return RequestOnlyResources().contains(id_); }
+
   /// Resource ID of CPU.
   static ResourceID CPU() { return ResourceID(PredefinedResourcesEnum::CPU); }
 
@@ -203,6 +205,9 @@ class ResourceID : public BaseSchedulingID<SchedulingIDTag::Resource> {
  private:
   /// Return the IDs of all unit-instance resources.
   static absl::flat_hash_set<int64_t> &UnitInstanceResources();
+
+  /// Return the IDs of all request-only-instance resources.
+  static absl::flat_hash_set<int64_t> &RequestOnlyResources();
 };
 
 }  // namespace scheduling

src/ray/raylet/local_task_manager.cc

@@ -430,9 +430,11 @@ bool LocalTaskManager::PoppedWorkerHandler(
 
   const auto &required_resource =
       task.GetTaskSpecification().GetRequiredResources().GetResourceMap();
-  for (auto &entry : required_resource) { // it fails here
+  for (auto &entry : required_resource) {
+    scheduling::ResourceID resource_id(entry.first);
     if (!cluster_resource_scheduler_->GetLocalResourceManager().ResourcesExist(
-            scheduling::ResourceID(entry.first)) && entry.first != "gpu_memory") {
+            resource_id) &&
+        !resource_id.IsRequestOnlyResource()) {
       RAY_CHECK(task.GetTaskSpecification().PlacementGroupBundleId().first !=
                 PlacementGroupID::Nil());
       RAY_LOG(DEBUG) << "The placement group: "

src/ray/raylet/scheduling/local_resource_manager.cc

@@ -82,7 +82,8 @@ bool LocalResourceManager::AllocateTaskResourceInstances(
   RAY_CHECK(task_allocation != nullptr);
   const ResourceSet adjusted_resource_request =
       local_resources_.ConvertRelativeResource(resource_request.GetResourceSet());
-  if (resource_request.GetResourceSet().Has(ResourceID::GPU_Memory()) && adjusted_resource_request.Get(ResourceID::GPU()) > 1) {
+  if (resource_request.GetResourceSet().Has(ResourceID::GPU_Memory()) &&
+      adjusted_resource_request.Get(ResourceID::GPU()) > 1) {
     return false;
   }
   // add adjust_gpu_memory here, added to NodeInstanceResourceSet