From d8ae9de99cd84a6ae1108e80c54367a6892a6bef Mon Sep 17 00:00:00 2001
From: Ion <istoica@cs.berkeley.edu>
Date: Tue, 6 Nov 2018 01:14:09 +0200
Subject: [PATCH] Caching task resource requirements. (#3231)

* caching resource requirements

* small fixes

* avoid copying the resource map
---
 src/ray/raylet/task_spec.cc | 34 +++++++++++++++++-----------------
 src/ray/raylet/task_spec.h  |  5 ++++-
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/src/ray/raylet/task_spec.cc b/src/ray/raylet/task_spec.cc
index 34747c90b..0a914adcb 100644
--- a/src/ray/raylet/task_spec.cc
+++ b/src/ray/raylet/task_spec.cc
@@ -32,6 +32,18 @@ flatbuffers::Offset<Arg> TaskArgumentByValue::ToFlatbuffer(
 
 void TaskSpecification::AssignSpecification(const uint8_t *spec, size_t spec_size) {
   spec_.assign(spec, spec + spec_size);
+  // Initialize required_resources_ and required_placement_resources_
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
+  auto required_resources = map_from_flatbuf(*message->required_resources());
+  auto required_placement_resources =
+      map_from_flatbuf(*message->required_placement_resources());
+  // If the required_placement_resources field is empty, then the placement
+  // resources default to the required resources.
+  if (required_placement_resources.size() == 0) {
+    required_placement_resources = required_resources;
+  }
+  required_resources_ = ResourceSet(required_resources);
+  required_placement_resources_ = ResourceSet(required_placement_resources);
 }
 
 TaskSpecification::TaskSpecification(const flatbuffers::String &string) {
@@ -168,30 +180,18 @@ size_t TaskSpecification::ArgValLength(int64_t arg_index) const {
 }
 
 double TaskSpecification::GetRequiredResource(const std::string &resource_name) const {
-  auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  auto required_resources = map_from_flatbuf(*message->required_resources());
-  auto it = required_resources.find(resource_name);
-  RAY_CHECK(it != required_resources.end());
+  RAY_CHECK(required_resources_.GetResourceMap().empty() == false);
+  auto it = required_resources_.GetResourceMap().find(resource_name);
+  RAY_CHECK(it != required_resources_.GetResourceMap().end());
   return it->second;
 }
 
 const ResourceSet TaskSpecification::GetRequiredResources() const {
-  auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  auto required_resources = map_from_flatbuf(*message->required_resources());
-  return ResourceSet(required_resources);
+  return required_resources_;
 }
 
 const ResourceSet TaskSpecification::GetRequiredPlacementResources() const {
-  auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  auto required_placement_resources =
-      map_from_flatbuf(*message->required_placement_resources());
-  // If the required_placement_resources field is empty, then the placement
-  // resources default to the required resources.
-  if (required_placement_resources.size() == 0) {
-    required_placement_resources = map_from_flatbuf(*message->required_resources());
-  }
-
-  return ResourceSet(required_placement_resources);
+  return required_placement_resources_;
 }
 
 bool TaskSpecification::IsDriverTask() const {
diff --git a/src/ray/raylet/task_spec.h b/src/ray/raylet/task_spec.h
index ae7237576..5a86a443c 100644
--- a/src/ray/raylet/task_spec.h
+++ b/src/ray/raylet/task_spec.h
@@ -204,7 +204,10 @@ class TaskSpecification {
   const uint8_t *data() const;
   /// Get the size in bytes of the task specification.
   size_t size() const;
-
+  /// Field storing required resources. Initalized in constructor.
+  ResourceSet required_resources_;
+  /// Field storing required placement resources. Initalized in constructor.
+  ResourceSet required_placement_resources_;
   /// The task specification data.
   std::vector<uint8_t> spec_;
 };