From 666fcde8cafd31785d24cab98fc874c5db2ebdb5 Mon Sep 17 00:00:00 2001
From: SangBin Cho <rkooo567@gmail.com>
Date: Wed, 14 Oct 2020 13:56:41 -0700
Subject: [PATCH] [Placement group] Input validation (#11152)

* Add a basic input validation.

* Addressed code review.
---
 python/ray/tests/test_placement_group.py | 16 +++++++++++-----
 src/ray/common/placement_group.h         | 12 ++++++++++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/python/ray/tests/test_placement_group.py b/python/ray/tests/test_placement_group.py
index 1c07b920a..a5b75636a 100644
--- a/python/ray/tests/test_placement_group.py
+++ b/python/ray/tests/test_placement_group.py
@@ -32,11 +32,17 @@ def test_placement_group_pack(ray_start_cluster):
     ray.init(address=cluster.address)
 
     placement_group = ray.util.placement_group(
-        name="name", strategy="PACK", bundles=[{
-            "CPU": 2
-        }, {
-            "CPU": 2
-        }])
+        name="name",
+        strategy="PACK",
+        bundles=[
+            {
+                "CPU": 2,
+                "GPU": 0  # Test 0 resource spec doesn't break tests.
+            },
+            {
+                "CPU": 2
+            }
+        ])
     ray.get(placement_group.ready())
     actor_1 = Actor.options(
         placement_group=placement_group,
diff --git a/src/ray/common/placement_group.h b/src/ray/common/placement_group.h
index 2b5b8550a..d304631f8 100644
--- a/src/ray/common/placement_group.h
+++ b/src/ray/common/placement_group.h
@@ -77,8 +77,16 @@ class PlacementGroupSpecBuilder {
       auto mutable_bundle_id = message_bundle->mutable_bundle_id();
       mutable_bundle_id->set_bundle_index(i);
       mutable_bundle_id->set_placement_group_id(placement_group_id.Binary());
-      message_bundle->mutable_unit_resources()->insert(resources.begin(),
-                                                       resources.end());
+      auto mutable_unit_resources = message_bundle->mutable_unit_resources();
+      for (auto it = resources.begin(); it != resources.end();) {
+        auto current = it++;
+        // Remove a resource with value 0 because they are not allowed.
+        if (current->second == 0) {
+          resources.erase(current);
+        } else {
+          mutable_unit_resources->insert({current->first, current->second});
+        }
+      }
     }
     return *this;
   }