Propogate InvalidArgumentStatus from LabelSelector

Signed-off-by: Ryan O'Leary <[email protected]>

src/ray/common/scheduling/label_selector.cc

@@ -20,25 +20,28 @@
 namespace ray {
 
 // Constructor to parse LabelSelector data type from proto.
-LabelSelector::LabelSelector(
+StatusOr<LabelSelector> LabelSelector::FromProto(
     const google::protobuf::Map<std::string, std::string> &label_selector) {
+  LabelSelector selector;
   for (const auto &[key, value] : label_selector) {
     if (key.empty()) {
-      // TODO (ryanaoleary@): propagate up an InvalidArgument from here.
-      RAY_LOG(ERROR) << "Empty Label Selector key.";
+      return Status::InvalidArgument("Empty label selector key is not supported.");
     }
-
-    AddConstraint(key, value);
+    RAY_RETURN_NOT_OK(selector.AddConstraint(key, value));
   }
+  return selector;
 }
 
-void LabelSelector::AddConstraint(const std::string &key, const std::string &value) {
-  auto [op, values] = ParseLabelSelectorValue(key, value);
+Status LabelSelector::AddConstraint(const std::string &key, const std::string &value) {
+  RAY_ASSIGN_OR_RETURN(auto parsed, ParseLabelSelectorValue(key, value));
+  auto &[op, values] = parsed;
+
   LabelConstraint constraint(key, op, values);
   AddConstraint(std::move(constraint));
+  return Status::OK();
 }
 
-std::pair<LabelSelectorOperator, absl::flat_hash_set<std::string>>
+StatusOr<std::pair<LabelSelectorOperator, absl::flat_hash_set<std::string>>>
 LabelSelector::ParseLabelSelectorValue(const std::string &key, const std::string &value) {
   bool is_negated = false;
   std::string_view val = value;
@@ -65,8 +68,8 @@ LabelSelector::ParseLabelSelectorValue(const std::string &key, const std::string
     }
 
     if (values.empty()) {
-      // TODO (ryanaoleary@): propagate up an InvalidArgument from here.
-      RAY_LOG(ERROR) << "No values provided for Label Selector key: " << key;
+      return Status::InvalidArgument(
+          "No values provided for Label Selector 'in' operator.");
     }
 
     op = is_negated ? LabelSelectorOperator::LABEL_NOT_IN
@@ -77,7 +80,7 @@ LabelSelector::ParseLabelSelectorValue(const std::string &key, const std::string
                     : LabelSelectorOperator::LABEL_IN;
   }
 
-  return {op, values};
+  return std::make_pair(op, values);
 }
 
 }  // namespace ray

src/ray/common/scheduling/label_selector.h

@@ -19,6 +19,8 @@
 
 #include "absl/container/flat_hash_set.h"
 #include "google/protobuf/map.h"
+#include "ray/common/status.h"
+#include "ray/common/status_or.h"
 
 namespace ray {
 
@@ -57,18 +59,18 @@ class LabelSelector {
  public:
   LabelSelector() = default;
 
-  explicit LabelSelector(
-      const google::protobuf::Map<std::string, std::string> &label_selector);
+  static StatusOr<LabelSelector> FromProto(
+      const google::protobuf::Map<std::string, std::string> &label_selector_dict);
 
-  void AddConstraint(const std::string &key, const std::string &value);
+  Status AddConstraint(const std::string &key, const std::string &value);
 
   void AddConstraint(LabelConstraint constraint) {
     constraints_.push_back(std::move(constraint));
   }
 
   const std::vector<LabelConstraint> &GetConstraints() const { return constraints_; }
 
-  std::pair<LabelSelectorOperator, absl::flat_hash_set<std::string>>
+  StatusOr<std::pair<LabelSelectorOperator, absl::flat_hash_set<std::string>>>
   ParseLabelSelectorValue(const std::string &key, const std::string &value);
 
  private:

src/ray/common/task/task_spec.cc

@@ -88,7 +88,7 @@ bool TaskSpecification::PlacementGroupCaptureChildTasks() const {
   }
 }
 
-void TaskSpecification::ComputeResources() {
+Status TaskSpecification::ComputeResources() {
   auto &required_resources = message_->required_resources();
 
   if (required_resources.empty()) {
@@ -130,8 +130,13 @@ void TaskSpecification::ComputeResources() {
   runtime_env_hash_ = CalculateRuntimeEnvHash(SerializedRuntimeEnv());
 
   // Set LabelSelector required for scheduling if specified. Parses string map
-  // from proto to LabelSelector data type.
-  label_selector_ = std::make_shared<LabelSelector>(message_->label_selector());
+  // from proto to LabelSelector data type. Returns an InvalidArgument Status if
+  // a malformed label selector is passed to the Task.
+  RAY_ASSIGN_OR_RETURN(auto label_selector,
+                       LabelSelector::FromProto(message_->label_selector()));
+  label_selector_ = std::make_shared<LabelSelector>(std::move(label_selector));
+
+  return Status::OK();
 }
 
 // Task specification getter methods.

src/ray/common/task/task_spec.h

@@ -518,7 +518,7 @@ class TaskSpecification : public MessageWrapper<rpc::TaskSpec> {
   TaskAttempt GetTaskAttempt() const;
 
  private:
-  void ComputeResources();
+  Status ComputeResources();
 
   /// Field storing required resources. Initialized in constructor.
   /// TODO(ekl) consider optimizing the representation of ResourceSet for fast copies

src/ray/common/test/label_selector_test.cc

@@ -23,9 +23,11 @@ TEST(LabelSelectorTest, BasicConstruction) {
   label_selector_dict["market-type"] = "spot";
   label_selector_dict["region"] = "us-east";
 
-  LabelSelector selector(label_selector_dict);
-  auto constraints = selector.GetConstraints();
+  auto result = LabelSelector::FromProto(label_selector_dict);
+  ASSERT_TRUE(result.ok()) << result.status().ToString();
+  auto selector = std::move(result.value());
 
+  auto constraints = selector.GetConstraints();
   ASSERT_EQ(constraints.size(), 2);
 
   for (const auto &constraint : constraints) {
@@ -39,7 +41,8 @@ TEST(LabelSelectorTest, BasicConstruction) {
 
 TEST(LabelSelectorTest, InOperatorParsing) {
   LabelSelector selector;
-  selector.AddConstraint("region", "in(us-west,us-east,me-central)");
+  Status status = selector.AddConstraint("region", "in(us-west,us-east,me-central)");
+  ASSERT_TRUE(status.ok());
 
   auto constraints = selector.GetConstraints();
   ASSERT_EQ(constraints.size(), 1);
@@ -55,7 +58,8 @@ TEST(LabelSelectorTest, InOperatorParsing) {
 
 TEST(LabelSelectorTest, NotInOperatorParsing) {
   LabelSelector selector;
-  selector.AddConstraint("tier", "!in(premium,free)");
+  Status status = selector.AddConstraint("tier", "!in(premium,free)");
+  ASSERT_TRUE(status.ok());
 
   auto constraints = selector.GetConstraints();
   ASSERT_EQ(constraints.size(), 1);
@@ -70,7 +74,8 @@ TEST(LabelSelectorTest, NotInOperatorParsing) {
 
 TEST(LabelSelectorTest, SingleValueNotInParsing) {
   LabelSelector selector;
-  selector.AddConstraint("env", "!dev");
+  Status status = selector.AddConstraint("env", "!dev");
+  ASSERT_TRUE(status.ok());
 
   auto constraints = selector.GetConstraints();
   ASSERT_EQ(constraints.size(), 1);
@@ -86,22 +91,18 @@ TEST(LabelSelectorTest, ErrorLogsOnEmptyKey) {
   google::protobuf::Map<std::string, std::string> label_selector_dict;
   label_selector_dict[""] = "value";
 
-  testing::internal::CaptureStderr();
-  LabelSelector selector(label_selector_dict);
-  std::string stderr_output = testing::internal::GetCapturedStderr();
+  auto result = LabelSelector::FromProto(label_selector_dict);
 
-  EXPECT_NE(stderr_output.find("Empty Label Selector key."), std::string::npos);
+  ASSERT_FALSE(result.ok());
+  ASSERT_EQ(result.status().message(), "Empty label selector key is not supported.");
 }
 
 TEST(LabelSelectorTest, ErrorLogsOnEmptyInList) {
   LabelSelector selector;
 
-  testing::internal::CaptureStderr();
-  selector.AddConstraint("key", "in()");
-  std::string stderr_output = testing::internal::GetCapturedStderr();
-
-  EXPECT_NE(stderr_output.find("No values provided for Label Selector key: key"),
-            std::string::npos);
+  Status status = selector.AddConstraint("key", "in()");
+  ASSERT_FALSE(status.ok());
+  ASSERT_EQ(status.message(), "No values provided for Label Selector 'in' operator.");
 }
 
 }  // namespace ray