@@ -295,49 +295,100 @@ PyTypeObject PyObjectIDType = {
295295 PyType_GenericNew, /* tp_new */
296296};
297297
298- /* Define the PyTask class. */
298+ // Define the PyTask class.
299+
300+ int resource_map_from_python_dict (
301+ PyObject *resource_map,
302+ std::unordered_map<std::string, double > &out) {
303+ RAY_CHECK (out.size () == 0 );
304+
305+ PyObject *key, *value;
306+ Py_ssize_t position = 0 ;
307+ if (!PyDict_Check (resource_map)) {
308+ PyErr_SetString (PyExc_TypeError, " resource_map must be a dictionary"
309+ return -1 ;
310+ }
311+
312+ while (PyDict_Next (resource_map, &position, &key, &value)) {
313+ #if PY_MAJOR_VERSION >= 3
314+ if (!PyUnicode_Check (key)) {
315+ PyErr_SetString (PyExc_TypeError,
316+ " the keys in resource_map must be strings"
317+ return -1 ;
318+ }
319+ #else
320+ if (!PyBytes_Check (key)) {
321+ PyErr_SetString (PyExc_TypeError,
322+ " the keys in resource_map must be strings"
323+ return -1 ;
324+ }
325+ #endif
326+
327+ // Check that the resource quantities are numbers.
328+ if (!(PyFloat_Check (value) || PyInt_Check (value) || PyLong_Check (value))) {
329+ PyErr_SetString (PyExc_TypeError,
330+ " the values in resource_map must be floats"
331+ return -1 ;
332+ }
333+ // Handle the case where the key is a bytes object and the case where it
334+ // is a unicode object.
335+ std::string resource_name;
336+ if (PyUnicode_Check (key)) {
337+ PyObject *ascii_key = PyUnicode_AsASCIIString (key);
338+ resource_name =
339+ std::string (PyBytes_AsString (ascii_key), PyBytes_Size (ascii_key));
340+ Py_DECREF (ascii_key);
341+ } else {
342+ resource_name = std::string (PyBytes_AsString (key), PyBytes_Size (key));
343+ }
344+ out[resource_name] = PyFloat_AsDouble (value);
345+ }
346+ return 0 ;
347+ }
299348
300349static int PyTask_init (PyTask *self, PyObject *args, PyObject *kwds) {
301- /* ID of the driver that this task originates from. */
350+ // ID of the driver that this task originates from.
302351 UniqueID driver_id;
303- /* ID of the actor this task should run on. */
352+ // ID of the actor this task should run on.
304353 UniqueID actor_id = ActorID::nil ();
305- /* ID of the actor handle used to submit this task. */
354+ // ID of the actor handle used to submit this task.
306355 UniqueID actor_handle_id = ActorHandleID::nil ();
307- /* How many tasks have been launched on the actor so far? */
356+ // How many tasks have been launched on the actor so far?
308357 int actor_counter = 0 ;
309- /* True if this is an actor checkpoint task and false otherwise. */
358+ // True if this is an actor checkpoint task and false otherwise.
310359 PyObject *is_actor_checkpoint_method_object = nullptr ;
311- /* ID of the function this task executes. */
360+ // ID of the function this task executes.
312361 FunctionID function_id;
313- /* Arguments of the task (can be PyObjectIDs or Python values). */
362+ // Arguments of the task (can be PyObjectIDs or Python values).
314363 PyObject *arguments;
315- /* Number of return values of this task. */
364+ // Number of return values of this task.
316365 int num_returns;
317- /* The ID of the task that called this task. */
366+ // The ID of the task that called this task.
318367 TaskID parent_task_id;
319- /* The number of tasks that the parent task has called prior to this one. */
368+ // The number of tasks that the parent task has called prior to this one.
320369 int parent_counter;
321370 // The actor creation ID.
322371 ActorID actor_creation_id = ActorID::nil ();
323372 // The dummy object for the actor creation task (if this is an actor method).
324373 ObjectID actor_creation_dummy_object_id = ObjectID::nil ();
325- /* Arguments of the task that are execution-dependent. These must be
326- * PyObjectIDs). */
374+ // Arguments of the task that are execution-dependent. These must be
375+ // PyObjectIDs).
327376 PyObject *execution_arguments = nullptr ;
328- /* Dictionary of resource requirements for this task. */
377+ // Dictionary of resource requirements for this task.
329378 PyObject *resource_map = nullptr ;
379+ // Dictionary of required placement resources for this task.
380+ PyObject *placement_resource_map = nullptr ;
330381 // True if we should use the raylet code path and false otherwise.
331382 PyObject *use_raylet_object = nullptr ;
332383 if (!PyArg_ParseTuple (
333- args, " O&O&OiO&i|O&O&O&O&iOOOO "
384+ args, " O&O&OiO&i|O&O&O&O&iOOOOO "
334385 &PyObjectToUniqueID, &function_id, &arguments, &num_returns,
335386 &PyObjectToUniqueID, &parent_task_id, &parent_counter,
336387 &PyObjectToUniqueID, &actor_creation_id, &PyObjectToUniqueID,
337388 &actor_creation_dummy_object_id, &PyObjectToUniqueID, &actor_id,
338389 &PyObjectToUniqueID, &actor_handle_id, &actor_counter,
339390 &is_actor_checkpoint_method_object, &execution_arguments,
340- &resource_map, &use_raylet_object)) {
391+ &resource_map, &placement_resource_map, & use_raylet_object)) {
341392 return -1 ;
342393 }
343394
@@ -349,48 +400,25 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
349400
350401 // Parse the resource map.
351402 std::unordered_map<std::string, double > required_resources;
403+ std::unordered_map<std::string, double > required_placement_resources;
352404
353- bool found_CPU_requirements = false ;
354- PyObject *key, *value;
355- Py_ssize_t position = 0 ;
356405 if (resource_map != nullptr ) {
357- if (!PyDict_Check (resource_map)) {
358- PyErr_SetString (PyExc_TypeError, " resource_map must be a dictionary"
406+ if (resource_map_from_python_dict (resource_map, required_resources) != 0 ) {
359407 return -1 ;
360408 }
361- while (PyDict_Next (resource_map, &position, &key, &value)) {
362- if (!(PyBytes_Check (key) || PyUnicode_Check (key))) {
363- PyErr_SetString (PyExc_TypeError,
364- " the keys in resource_map must be strings"
365- return -1 ;
366- }
367- if (!(PyFloat_Check (value) || PyInt_Check (value) ||
368- PyLong_Check (value))) {
369- PyErr_SetString (PyExc_TypeError,
370- " the values in resource_map must be floats"
371- return -1 ;
372- }
373- // Handle the case where the key is a bytes object and the case where it
374- // is a unicode object.
375- std::string resource_name;
376- if (PyUnicode_Check (key)) {
377- PyObject *ascii_key = PyUnicode_AsASCIIString (key);
378- resource_name =
379- std::string (PyBytes_AsString (ascii_key), PyBytes_Size (ascii_key));
380- Py_DECREF (ascii_key);
381- } else {
382- resource_name = std::string (PyBytes_AsString (key), PyBytes_Size (key));
383- }
384- if (resource_name == std::string (" CPU"
385- found_CPU_requirements = true ;
386- }
387- required_resources[resource_name] = PyFloat_AsDouble (value);
388- }
389409 }
390- if (!found_CPU_requirements) {
410+
411+ if (required_resources.count (" CPU" 0 ) {
391412 required_resources[" CPU" 1.0 ;
392413 }
393414
415+ if (placement_resource_map != nullptr ) {
416+ if (resource_map_from_python_dict (placement_resource_map,
417+ required_placement_resources) != 0 ) {
418+ return -1 ;
419+ }
420+ }
421+
394422 Py_ssize_t num_args = PyList_Size (arguments);
395423
396424 bool use_raylet = false ;
@@ -463,7 +491,7 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
463491 driver_id, parent_task_id, parent_counter, actor_creation_id,
464492 actor_creation_dummy_object_id, actor_id, actor_handle_id,
465493 actor_counter, function_id, args, num_returns, required_resources,
466- Language::PYTHON);
494+ required_placement_resources, Language::PYTHON);
467495 }
468496
469497 /* Set the task's execution dependencies. */
0 commit comments