Allow scheduling with arbitrary user-defined resource labels. (#1236)

* Enable scheduling with custom resource labels.

* Fix.

* Minor fixes and ref counting fix.

* Linting

* Use .data() instead of .c_str().

* Fix linting.

* Fix ResourcesTest.testGPUIDs test by waiting for workers to start up.

* Sleep in test so that all tasks are submitted before any completes.
This commit is contained in:
Robert Nishihara
2017-12-01 11:41:40 -08:00
committed by Philipp Moritz
parent ac64631043
commit c21e189371
42 changed files with 1072 additions and 805 deletions
+30 -3
View File
@@ -6,10 +6,10 @@ flatbuffers::Offset<flatbuffers::String> to_flatbuf(
return fbb.CreateString((char *) &object_id.id[0], sizeof(object_id.id));
}
ObjectID from_flatbuf(const flatbuffers::String *string) {
ObjectID from_flatbuf(const flatbuffers::String &string) {
ObjectID object_id;
CHECK(string->size() == sizeof(object_id.id));
memcpy(&object_id.id[0], string->data(), sizeof(object_id.id));
CHECK(string.size() == sizeof(object_id.id));
memcpy(&object_id.id[0], string.data(), sizeof(object_id.id));
return object_id;
}
@@ -24,3 +24,30 @@ to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
}
return fbb.CreateVector(results);
}
std::string string_from_flatbuf(const flatbuffers::String &string) {
return std::string(string.data(), string.size());
}
const std::unordered_map<std::string, double> map_from_flatbuf(
const flatbuffers::Vector<flatbuffers::Offset<ResourcePair>>
&resource_vector) {
std::unordered_map<std::string, double> required_resources;
for (int64_t i = 0; i < resource_vector.size(); i++) {
const ResourcePair *resource_pair = resource_vector.Get(i);
required_resources[string_from_flatbuf(*resource_pair->key())] =
resource_pair->value();
}
return required_resources;
}
flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<ResourcePair>>>
map_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
const std::unordered_map<std::string, double> &resource_map) {
std::vector<flatbuffers::Offset<ResourcePair>> resource_vector;
for (auto const &resource_pair : resource_map) {
resource_vector.push_back(CreateResourcePair(
fbb, fbb.CreateString(resource_pair.first), resource_pair.second));
}
return fbb.CreateVector(resource_vector);
}