Commit 38f193fc authored by Stanislav Bohm's avatar Stanislav Bohm

ENH: Scheduler respects cpu resources of tasks

parent 5e5b0325
from .client import Client, LoomException, TaskFailed # noqa
from .plan import Plan # noqa
from .plan import Plan, cpus, cpu1 # noqa
......@@ -18,7 +18,7 @@ _sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor.FileDescriptor(
name='loomplan.proto',
package='loomplan',
serialized_pb=_b('\n\x0eloomplan.proto\x12\x08loomplan\"\xbc\x01\n\x04Task\x12\x11\n\ttask_type\x18\x01 \x02(\x05\x12\x0e\n\x06\x63onfig\x18\x02 \x02(\x0c\x12\x11\n\tinput_ids\x18\x03 \x03(\x05\x12\x36\n\x06policy\x18\x04 \x01(\x0e\x32\x15.loomplan.Task.Policy:\x0fPOLICY_STANDARD\"F\n\x06Policy\x12\x13\n\x0fPOLICY_STANDARD\x10\x01\x12\x11\n\rPOLICY_SIMPLE\x10\x02\x12\x14\n\x10POLICY_SCHEDULER\x10\x03\"9\n\x04Plan\x12\x1d\n\x05tasks\x18\x02 \x03(\x0b\x32\x0e.loomplan.Task\x12\x12\n\nresult_ids\x18\x03 \x03(\x05\x42\x02H\x03')
serialized_pb=_b('\n\x0eloomplan.proto\x12\x08loomplan\"\xe0\x01\n\x04Task\x12\x11\n\ttask_type\x18\x01 \x02(\x05\x12\x0e\n\x06\x63onfig\x18\x02 \x02(\x0c\x12\x11\n\tinput_ids\x18\x03 \x03(\x05\x12\x36\n\x06policy\x18\x04 \x01(\x0e\x32\x15.loomplan.Task.Policy:\x0fPOLICY_STANDARD\x12\"\n\x16resource_request_index\x18\x05 \x01(\x05:\x02-1\"F\n\x06Policy\x12\x13\n\x0fPOLICY_STANDARD\x10\x01\x12\x11\n\rPOLICY_SIMPLE\x10\x02\x12\x14\n\x10POLICY_SCHEDULER\x10\x03\"0\n\x08Resource\x12\x15\n\rresource_type\x18\x01 \x02(\x05\x12\r\n\x05value\x18\x02 \x02(\x05\"8\n\x0fResourceRequest\x12%\n\tresources\x18\x01 \x03(\x0b\x32\x12.loomplan.Resource\"o\n\x04Plan\x12\x34\n\x11resource_requests\x18\x01 \x03(\x0b\x32\x19.loomplan.ResourceRequest\x12\x1d\n\x05tasks\x18\x02 \x03(\x0b\x32\x0e.loomplan.Task\x12\x12\n\nresult_ids\x18\x03 \x03(\x05\x42\x02H\x03')
)
_sym_db.RegisterFileDescriptor(DESCRIPTOR)
......@@ -45,8 +45,8 @@ _TASK_POLICY = _descriptor.EnumDescriptor(
],
containing_type=None,
options=None,
serialized_start=147,
serialized_end=217,
serialized_start=183,
serialized_end=253,
)
_sym_db.RegisterEnumDescriptor(_TASK_POLICY)
......@@ -86,6 +86,13 @@ _TASK = _descriptor.Descriptor(
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
_descriptor.FieldDescriptor(
name='resource_request_index', full_name='loomplan.Task.resource_request_index', index=4,
number=5, type=5, cpp_type=1, label=1,
has_default_value=True, default_value=-1,
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
],
extensions=[
],
......@@ -99,7 +106,74 @@ _TASK = _descriptor.Descriptor(
oneofs=[
],
serialized_start=29,
serialized_end=217,
serialized_end=253,
)
_RESOURCE = _descriptor.Descriptor(
name='Resource',
full_name='loomplan.Resource',
filename=None,
file=DESCRIPTOR,
containing_type=None,
fields=[
_descriptor.FieldDescriptor(
name='resource_type', full_name='loomplan.Resource.resource_type', index=0,
number=1, type=5, cpp_type=1, label=2,
has_default_value=False, default_value=0,
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
_descriptor.FieldDescriptor(
name='value', full_name='loomplan.Resource.value', index=1,
number=2, type=5, cpp_type=1, label=2,
has_default_value=False, default_value=0,
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
],
extensions=[
],
nested_types=[],
enum_types=[
],
options=None,
is_extendable=False,
extension_ranges=[],
oneofs=[
],
serialized_start=255,
serialized_end=303,
)
_RESOURCEREQUEST = _descriptor.Descriptor(
name='ResourceRequest',
full_name='loomplan.ResourceRequest',
filename=None,
file=DESCRIPTOR,
containing_type=None,
fields=[
_descriptor.FieldDescriptor(
name='resources', full_name='loomplan.ResourceRequest.resources', index=0,
number=1, type=11, cpp_type=10, label=3,
has_default_value=False, default_value=[],
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
],
extensions=[
],
nested_types=[],
enum_types=[
],
options=None,
is_extendable=False,
extension_ranges=[],
oneofs=[
],
serialized_start=305,
serialized_end=361,
)
......@@ -111,14 +185,21 @@ _PLAN = _descriptor.Descriptor(
containing_type=None,
fields=[
_descriptor.FieldDescriptor(
name='tasks', full_name='loomplan.Plan.tasks', index=0,
name='resource_requests', full_name='loomplan.Plan.resource_requests', index=0,
number=1, type=11, cpp_type=10, label=3,
has_default_value=False, default_value=[],
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
_descriptor.FieldDescriptor(
name='tasks', full_name='loomplan.Plan.tasks', index=1,
number=2, type=11, cpp_type=10, label=3,
has_default_value=False, default_value=[],
message_type=None, enum_type=None, containing_type=None,
is_extension=False, extension_scope=None,
options=None),
_descriptor.FieldDescriptor(
name='result_ids', full_name='loomplan.Plan.result_ids', index=1,
name='result_ids', full_name='loomplan.Plan.result_ids', index=2,
number=3, type=5, cpp_type=1, label=3,
has_default_value=False, default_value=[],
message_type=None, enum_type=None, containing_type=None,
......@@ -135,14 +216,18 @@ _PLAN = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
serialized_start=219,
serialized_end=276,
serialized_start=363,
serialized_end=474,
)
_TASK.fields_by_name['policy'].enum_type = _TASK_POLICY
_TASK_POLICY.containing_type = _TASK
_RESOURCEREQUEST.fields_by_name['resources'].message_type = _RESOURCE
_PLAN.fields_by_name['resource_requests'].message_type = _RESOURCEREQUEST
_PLAN.fields_by_name['tasks'].message_type = _TASK
DESCRIPTOR.message_types_by_name['Task'] = _TASK
DESCRIPTOR.message_types_by_name['Resource'] = _RESOURCE
DESCRIPTOR.message_types_by_name['ResourceRequest'] = _RESOURCEREQUEST
DESCRIPTOR.message_types_by_name['Plan'] = _PLAN
Task = _reflection.GeneratedProtocolMessageType('Task', (_message.Message,), dict(
......@@ -152,6 +237,20 @@ Task = _reflection.GeneratedProtocolMessageType('Task', (_message.Message,), dic
))
_sym_db.RegisterMessage(Task)
Resource = _reflection.GeneratedProtocolMessageType('Resource', (_message.Message,), dict(
DESCRIPTOR = _RESOURCE,
__module__ = 'loomplan_pb2'
# @@protoc_insertion_point(class_scope:loomplan.Resource)
))
_sym_db.RegisterMessage(Resource)
ResourceRequest = _reflection.GeneratedProtocolMessageType('ResourceRequest', (_message.Message,), dict(
DESCRIPTOR = _RESOURCEREQUEST,
__module__ = 'loomplan_pb2'
# @@protoc_insertion_point(class_scope:loomplan.ResourceRequest)
))
_sym_db.RegisterMessage(ResourceRequest)
Plan = _reflection.GeneratedProtocolMessageType('Plan', (_message.Message,), dict(
DESCRIPTOR = _PLAN,
__module__ = 'loomplan_pb2'
......
......@@ -17,12 +17,53 @@ class Task(object):
id = None
config = ""
policy = POLICY_STANDARD
resource_request = None
def set_message(self, msg, symbols):
def set_message(self, msg, symbols, requests):
msg.config = self.config
msg.task_type = symbols[self.task_type]
msg.input_ids.extend(t.id for t in self.inputs)
msg.policy = self.policy
if self.resource_request:
msg.resource_request_index = requests.index(self.resource_request)
class ResourceRequest(object):
def __init__(self):
self.resources = {}
def add_resource(self, name, value):
self.resources[name] = value
@property
def names(self):
return self.resources.keys()
def set_message(self, msg, symbols):
for name, value in self.resources.items():
r = msg.resources.add()
r.resource_type = symbols[name]
r.value = value
def __eq__(self, other):
if not isinstance(other, ResourceRequest):
return False
return self.resources == other.resources
def __ne__(self, other):
return not self.__eq__(other)
def __hash__(self):
return hash(tuple(sorted(self.resources.items())))
def cpus(value):
r = ResourceRequest()
r.add_resource("loom/resource/cpus", value)
return r
cpu1 = cpus(1)
class Plan(object):
......@@ -49,13 +90,11 @@ class Plan(object):
def __init__(self):
self.tasks = []
self.task_types = set()
def add(self, task):
assert task.id is None
task.id = len(self.tasks)
self.tasks.append(task)
self.task_types.add(task.task_type)
return task
def task_dslice(self, input):
......@@ -100,7 +139,13 @@ class Plan(object):
task.policy = POLICY_SIMPLE
return self.add(task)
def task_run(self, args, inputs=(), outputs=(None,), stdin=None):
def task_run(self,
args,
inputs=(),
outputs=(None,),
stdin=None,
request=cpu1):
if isinstance(args, str):
args = args.split()
......@@ -119,6 +164,7 @@ class Plan(object):
msg.map_outputs.extend(fname if fname else "+out"
for fname in outputs)
task.config = msg.SerializeToString()
task.resource_request = request
return self.add(task)
def task_array_make(self, inputs):
......@@ -160,9 +206,20 @@ class Plan(object):
def create_message(self, symbols):
msg = loomplan_pb2.Plan()
requests = set()
for task in self.tasks:
if task.resource_request:
requests.add(task.resource_request)
requests = list(requests)
for request in requests:
r = msg.resource_requests.add()
request.set_message(r, symbols)
for task in self.tasks:
t = msg.tasks.add()
task.set_message(t, symbols)
task.set_message(t, symbols, requests)
return msg
def write_dot(self, filename, info=None):
......
......@@ -17,6 +17,8 @@ namespace loomplan {
void protobuf_ShutdownFile_loomplan_2eproto() {
delete Task::default_instance_;
delete Resource::default_instance_;
delete ResourceRequest::default_instance_;
delete Plan::default_instance_;
}
......@@ -33,8 +35,12 @@ void protobuf_AddDesc_loomplan_2eproto() {
#endif
Task::default_instance_ = new Task();
Resource::default_instance_ = new Resource();
ResourceRequest::default_instance_ = new ResourceRequest();
Plan::default_instance_ = new Plan();
Task::default_instance_->InitAsDefaultInstance();
Resource::default_instance_->InitAsDefaultInstance();
ResourceRequest::default_instance_->InitAsDefaultInstance();
Plan::default_instance_->InitAsDefaultInstance();
::google::protobuf::internal::OnShutdown(&protobuf_ShutdownFile_loomplan_2eproto);
}
......@@ -80,6 +86,7 @@ const int Task::kTaskTypeFieldNumber;
const int Task::kConfigFieldNumber;
const int Task::kInputIdsFieldNumber;
const int Task::kPolicyFieldNumber;
const int Task::kResourceRequestIndexFieldNumber;
#endif // !_MSC_VER
Task::Task()
......@@ -104,6 +111,7 @@ void Task::SharedCtor() {
task_type_ = 0;
config_ = const_cast< ::std::string*>(&::google::protobuf::internal::GetEmptyStringAlreadyInited());
policy_ = 1;
resource_request_index_ = -1;
::memset(_has_bits_, 0, sizeof(_has_bits_));
}
......@@ -145,7 +153,7 @@ Task* Task::New() const {
}
void Task::Clear() {
if (_has_bits_[0 / 32] & 11) {
if (_has_bits_[0 / 32] & 27) {
task_type_ = 0;
if (has_config()) {
if (config_ != &::google::protobuf::internal::GetEmptyStringAlreadyInited()) {
......@@ -153,6 +161,7 @@ void Task::Clear() {
}
}
policy_ = 1;
resource_request_index_ = -1;
}
input_ids_.Clear();
::memset(_has_bits_, 0, sizeof(_has_bits_));
......@@ -236,6 +245,21 @@ bool Task::MergePartialFromCodedStream(
} else {
goto handle_unusual;
}
if (input->ExpectTag(40)) goto parse_resource_request_index;
break;
}
// optional int32 resource_request_index = 5 [default = -1];
case 5: {
if (tag == 40) {
parse_resource_request_index:
DO_((::google::protobuf::internal::WireFormatLite::ReadPrimitive<
::google::protobuf::int32, ::google::protobuf::internal::WireFormatLite::TYPE_INT32>(
input, &resource_request_index_)));
set_has_resource_request_index();
} else {
goto handle_unusual;
}
if (input->ExpectAtEnd()) goto success;
break;
}
......@@ -288,6 +312,11 @@ void Task::SerializeWithCachedSizes(
4, this->policy(), output);
}
// optional int32 resource_request_index = 5 [default = -1];
if (has_resource_request_index()) {
::google::protobuf::internal::WireFormatLite::WriteInt32(5, this->resource_request_index(), output);
}
output->WriteRaw(unknown_fields().data(),
unknown_fields().size());
// @@protoc_insertion_point(serialize_end:loomplan.Task)
......@@ -317,6 +346,13 @@ int Task::ByteSize() const {
::google::protobuf::internal::WireFormatLite::EnumSize(this->policy());
}
// optional int32 resource_request_index = 5 [default = -1];
if (has_resource_request_index()) {
total_size += 1 +
::google::protobuf::internal::WireFormatLite::Int32Size(
this->resource_request_index());
}
}
// repeated int32 input_ids = 3;
{
......@@ -354,6 +390,9 @@ void Task::MergeFrom(const Task& from) {
if (from.has_policy()) {
set_policy(from.policy());
}
if (from.has_resource_request_index()) {
set_resource_request_index(from.resource_request_index());
}
}
mutable_unknown_fields()->append(from.unknown_fields());
}
......@@ -376,6 +415,7 @@ void Task::Swap(Task* other) {
std::swap(config_, other->config_);
input_ids_.Swap(&other->input_ids_);
std::swap(policy_, other->policy_);
std::swap(resource_request_index_, other->resource_request_index_);
std::swap(_has_bits_[0], other->_has_bits_[0]);
_unknown_fields_.swap(other->_unknown_fields_);
std::swap(_cached_size_, other->_cached_size_);
......@@ -390,6 +430,433 @@ void Task::Swap(Task* other) {
// ===================================================================
#ifndef _MSC_VER
const int Resource::kResourceTypeFieldNumber;
const int Resource::kValueFieldNumber;
#endif // !_MSC_VER
Resource::Resource()
: ::google::protobuf::MessageLite() {
SharedCtor();
// @@protoc_insertion_point(constructor:loomplan.Resource)
}
void Resource::InitAsDefaultInstance() {
}
Resource::Resource(const Resource& from)
: ::google::protobuf::MessageLite() {
SharedCtor();
MergeFrom(from);
// @@protoc_insertion_point(copy_constructor:loomplan.Resource)
}
void Resource::SharedCtor() {
_cached_size_ = 0;
resource_type_ = 0;
value_ = 0;
::memset(_has_bits_, 0, sizeof(_has_bits_));
}
Resource::~Resource() {
// @@protoc_insertion_point(destructor:loomplan.Resource)
SharedDtor();
}
void Resource::SharedDtor() {
#ifdef GOOGLE_PROTOBUF_NO_STATIC_INITIALIZER
if (this != &default_instance()) {
#else
if (this != default_instance_) {
#endif
}
}
void Resource::SetCachedSize(int size) const {
GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
_cached_size_ = size;
GOOGLE_SAFE_CONCURRENT_WRITES_END();
}
const Resource& Resource::default_instance() {
#ifdef GOOGLE_PROTOBUF_NO_STATIC_INITIALIZER
protobuf_AddDesc_loomplan_2eproto();
#else
if (default_instance_ == NULL) protobuf_AddDesc_loomplan_2eproto();
#endif
return *default_instance_;
}
Resource* Resource::default_instance_ = NULL;
Resource* Resource::New() const {
return new Resource;
}
void Resource::Clear() {
#define OFFSET_OF_FIELD_(f) (reinterpret_cast<char*>( \
&reinterpret_cast<Resource*>(16)->f) - \
reinterpret_cast<char*>(16))
#define ZR_(first, last) do { \
size_t f = OFFSET_OF_FIELD_(first); \
size_t n = OFFSET_OF_FIELD_(last) - f + sizeof(last); \
::memset(&first, 0, n); \
} while (0)
ZR_(resource_type_, value_);
#undef OFFSET_OF_FIELD_
#undef ZR_
::memset(_has_bits_, 0, sizeof(_has_bits_));
mutable_unknown_fields()->clear();
}
bool Resource::MergePartialFromCodedStream(
::google::protobuf::io::CodedInputStream* input) {
#define DO_(EXPRESSION) if (!(EXPRESSION)) goto failure
::google::protobuf::uint32 tag;
::google::protobuf::io::StringOutputStream unknown_fields_string(
mutable_unknown_fields());
::google::protobuf::io::CodedOutputStream unknown_fields_stream(
&unknown_fields_string);
// @@protoc_insertion_point(parse_start:loomplan.Resource)
for (;;) {
::std::pair< ::google::protobuf::uint32, bool> p = input->ReadTagWithCutoff(127);
tag = p.first;
if (!p.second) goto handle_unusual;
switch (::google::protobuf::internal::WireFormatLite::GetTagFieldNumber(tag)) {
// required int32 resource_type = 1;
case 1: {
if (tag == 8) {
DO_((::google::protobuf::internal::WireFormatLite::ReadPrimitive<
::google::protobuf::int32, ::google::protobuf::internal::WireFormatLite::TYPE_INT32>(
input, &resource_type_)));
set_has_resource_type();
} else {
goto handle_unusual;
}
if (input->ExpectTag(16)) goto parse_value;
break;
}
// required int32 value = 2;
case 2: {
if (tag == 16) {
parse_value:
DO_((::google::protobuf::internal::WireFormatLite::ReadPrimitive<
::google::protobuf::int32, ::google::protobuf::internal::WireFormatLite::TYPE_INT32>(
input, &value_)));
set_has_value();
} else {
goto handle_unusual;
}
if (input->ExpectAtEnd()) goto success;
break;
}
default: {
handle_unusual:
if (tag == 0 ||
::google::protobuf::internal::WireFormatLite::GetTagWireType(tag) ==
::google::protobuf::internal::WireFormatLite::WIRETYPE_END_GROUP) {
goto success;
}
DO_(::google::protobuf::internal::WireFormatLite::SkipField(
input, tag, &unknown_fields_stream));
break;
}
}
}
success:
// @@protoc_insertion_point(parse_success:loomplan.Resource)
return true;
failure:
// @@protoc_insertion_point(parse_failure:loomplan.Resource)
return false;
#undef DO_
}
void Resource::SerializeWithCachedSizes(
::google::protobuf::io::CodedOutputStream* output) const {
// @@protoc_insertion_point(serialize_start:loomplan.Resource)
// required int32 resource_type = 1;
if (has_resource_type()) {
::google::protobuf::internal::WireFormatLite::WriteInt32(1, this->resource_type(), output);
}
// required int32 value = 2;
if (has_value()) {
::google::protobuf::internal::WireFormatLite::WriteInt32(2, this->value(), output);
}
output->WriteRaw(unknown_fields().data(),
unknown_fields().size());
// @@protoc_insertion_point(serialize_end:loomplan.Resource)
}
int Resource::ByteSize() const {
int total_size = 0;
if (_has_bits_[0 / 32] & (0xffu << (0 % 32))) {
// required int32 resource_type = 1;
if (has_resource_type()) {
total_size += 1 +
::google::protobuf::internal::WireFormatLite::Int32Size(
this->resource_type());
}
// required int32 value = 2;
if (has_value()) {
total_size += 1 +
::google::protobuf::internal::WireFormatLite::Int32Size(
this->value());
}
}
total_size += unknown_fields().size();
GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
_cached_size_ = total_size;
GOOGLE_SAFE_CONCURRENT_WRITES_END();
return total_size;
}
void Resource::CheckTypeAndMergeFrom(
const ::google::protobuf::MessageLite& from) {
MergeFrom(*::google::protobuf::down_cast<const Resource*>(&from));
}
void Resource::MergeFrom(const Resource& from) {
GOOGLE_CHECK_NE(&from, this);
if (from._has_bits_[0 / 32] & (0xffu << (0 % 32))) {
if (from.has_resource_type()) {
set_resource_type(from.resource_type());
}
if (from.has_value()) {
set_value(from.value());
}
}
mutable_unknown_fields()->append(from.unknown_fields());
}
void Resource::CopyFrom(const Resource& from) {
if (&from == this) return;
Clear();
MergeFrom(from);
}
bool Resource::IsInitialized() const {
if ((_has_bits_[0] & 0x00000003) != 0x00000003) return false;
return true;
}
void Resource::Swap(Resource* other) {
if (other != this) {