[runtime env] plugin refactor [7/n]: support runtime env in C++ API (r…

…ay-project#27010)

Signed-off-by: 久龙 <[email protected]>

bazel/BUILD.nlohmann_json

@@ -1,9 +1,9 @@
 cc_library(
     name = "nlohmann_json",
     hdrs = glob([
-        "include/**/*.hpp",
+        "single_include/**/*.hpp",
     ]),
-    includes = ["include"],
+    includes = ["single_include"],
     visibility = ["//visibility:public"],
     alwayslink = 1,
 )

cpp/BUILD.bazel

@@ -66,6 +66,7 @@ cc_library(
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
         "@msgpack",
+        "@nlohmann_json",
     ],
     alwayslink = True,
 )
@@ -102,6 +103,7 @@ cc_binary(
             "@boost//:callable_traits",
             "@boost//:optional",
             "@msgpack",
+            "@nlohmann_json",
         ],
     }),
 )
@@ -123,6 +125,7 @@ genrule(
         mkdir -p "$$PY_CPP_DIR/lib/" &&
         cp -f $(location default_worker) "$$PY_CPP_DIR/" &&
         cp -f -r $$WORK_DIR/external/msgpack/include/* "$$PY_CPP_DIR/include" &&
+        cp -f -r $$WORK_DIR/external/nlohmann_json/single_include/* "$$PY_CPP_DIR/include" &&
         cp -f -r "$$WORK_DIR/external/boost/boost/archive" "$$BOOST_DIR" &&
         cp -f -r "$$WORK_DIR/external/boost/boost/assert" "$$BOOST_DIR" &&
         cp -f -r "$$WORK_DIR/external/boost/boost/bind" "$$BOOST_DIR" &&
@@ -250,6 +253,7 @@ cc_binary(
         "@boost//:callable_traits",
         "@boost//:optional",
         "@msgpack",
+        "@nlohmann_json",
     ],
 )
 
@@ -271,6 +275,7 @@ cc_binary(
         "@boost//:callable_traits",
         "@boost//:optional",
         "@msgpack",
+        "@nlohmann_json",
     ],
 )
 
@@ -308,6 +313,7 @@ cc_binary(
         "@boost//:callable_traits",
         "@boost//:optional",
         "@msgpack",
+        "@nlohmann_json",
     ],
 )
 

cpp/include/ray/api.h

@@ -23,6 +23,7 @@
 #include <ray/api/ray_remote.h>
 #include <ray/api/ray_runtime.h>
 #include <ray/api/ray_runtime_holder.h>
+#include <ray/api/runtime_env.h>
 #include <ray/api/task_caller.h>
 #include <ray/api/wait_result.h>
 

cpp/include/ray/api/actor_creator.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <ray/api/actor_handle.h>
+#include <ray/api/runtime_env.h>
 #include <ray/api/task_options.h>
 
 namespace ray {
@@ -71,6 +72,11 @@ class ActorCreator {
     return *this;
   }
 
+  ActorCreator &SetRuntimeEnv(const ray::RuntimeEnv &runtime_env) {
+    create_options_.serialized_runtime_env_info = runtime_env.SerializeToRuntimeEnvInfo();
+    return *this;
+  }
+
  private:
   RayRuntime *runtime_;
   RemoteFunctionHolder remote_function_holder_;

cpp/include/ray/api/ray_config.h

@@ -14,6 +14,7 @@
 
 #pragma once
 #include <ray/api/ray_exception.h>
+#include <ray/api/runtime_env.h>
 
 #include <memory>
 #include <string>
@@ -52,6 +53,9 @@ class RayConfig {
   // The default actor lifetime type, `DETACHED` or `NON_DETACHED`.
   ActorLifetime default_actor_lifetime = ActorLifetime::NON_DETACHED;
 
+  // The job level runtime environments.
+  boost::optional<RuntimeEnv> runtime_env;
+
   /* The following are unstable parameters and their use is discouraged. */
 
   // Prevents external clients without the password from connecting to Redis if provided.

cpp/include/ray/api/ray_exception.h

@@ -53,5 +53,10 @@ class RayFunctionNotFound : public RayException {
  public:
   RayFunctionNotFound(const std::string &msg) : RayException(msg){};
 };
+
+class RayRuntimeEnvException : public RayException {
+ public:
+  RayRuntimeEnvException(const std::string &msg) : RayException(msg){};
+};
 }  // namespace internal
 }  // namespace ray

cpp/include/ray/api/runtime_env.h

@@ -0,0 +1,102 @@
+// Copyright 2022 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <ray/api/ray_exception.h>
+
+#include <string>
+
+#include "nlohmann/json.hpp"
+
+using json = nlohmann::json;
+
+namespace ray {
+
+/// This class provides interfaces of setting runtime environments for job/actor/task.
+class RuntimeEnv {
+ public:
+  /// Set a runtime env field by name and Object.
+  /// \param[in] name The runtime env plugin name.
+  /// \param[in] value An object with primitive data type or jsonable type of
+  /// nlohmann/json.
+  template <typename T>
+  void Set(const std::string &name, const T &value);
+
+  /// Get the object of a runtime env field.
+  /// \param[in] name The runtime env plugin name.
+  template <typename T>
+  T Get(const std::string &name) const;
+
+  /// Set a runtime env field by name and json string.
+  /// \param[in] name The runtime env plugin name.
+  /// \param[in] json_str A json string represents the runtime env field.
+  void SetJsonStr(const std::string &name, const std::string &json_str);
+
+  /// Get the json string of a runtime env field.
+  /// \param[in] name The runtime env plugin name.
+  std::string GetJsonStr(const std::string &name) const;
+
+  /// Whether a field is contained.
+  /// \param[in] name The runtime env plugin name.
+  bool Contains(const std::string &name) const;
+
+  /// Remove a field by name.
+  /// \param[in] name The runtime env plugin name.
+  /// \return true if remove an existing field, otherwise false.
+  bool Remove(const std::string &name);
+
+  /// Whether the runtime env is empty.
+  bool Empty() const;
+
+  /// Serialize the runtime env to string.
+  std::string Serialize() const;
+
+  /// Serialize the runtime env to RuntimeEnvInfo.
+  std::string SerializeToRuntimeEnvInfo() const;
+
+  /// Deserialize the runtime env from string.
+  /// \return The deserialized RuntimeEnv instance.
+  static RuntimeEnv Deserialize(const std::string &serialized_runtime_env);
+
+ private:
+  json fields_;
+};
+
+// --------- inline implementation ------------
+
+template <typename T>
+inline void RuntimeEnv::Set(const std::string &name, const T &value) {
+  try {
+    json value_j = value;
+    fields_[name] = value_j;
+  } catch (std::exception &e) {
+    throw ray::internal::RayRuntimeEnvException("Failed to set the field " + name + ": " +
+                                                e.what());
+  }
+}
+
+template <typename T>
+inline T RuntimeEnv::Get(const std::string &name) const {
+  if (!Contains(name)) {
+    throw ray::internal::RayRuntimeEnvException("The field " + name + " not found.");
+  }
+  try {
+    return fields_[name].get<T>();
+  } catch (std::exception &e) {
+    throw ray::internal::RayRuntimeEnvException("Failed to get the field " + name + ": " +
+                                                e.what());
+  }
+}
+
+}  // namespace ray

cpp/include/ray/api/task_caller.h

@@ -14,8 +14,10 @@
 
 #pragma once
 
+#include <ray/api/runtime_env.h>
 #include <ray/api/static_check.h>
 #include <ray/api/task_options.h>
+
 namespace ray {
 namespace internal {
 
@@ -50,6 +52,11 @@ class TaskCaller {
     return *this;
   }
 
+  TaskCaller &SetRuntimeEnv(const ray::RuntimeEnv &runtime_env) {
+    task_options_.serialized_runtime_env_info = runtime_env.SerializeToRuntimeEnvInfo();
+    return *this;
+  }
+
  private:
   RayRuntime *runtime_;
   RemoteFunctionHolder remote_function_holder_{};

cpp/include/ray/api/task_options.h

@@ -95,6 +95,7 @@ struct CallOptions {
   std::unordered_map<std::string, double> resources;
   PlacementGroup group;
   int bundle_index;
+  std::string serialized_runtime_env_info;
 };
 
 struct ActorCreationOptions {
@@ -105,6 +106,7 @@ struct ActorCreationOptions {
   int max_concurrency = 1;
   PlacementGroup group;
   int bundle_index;
+  std::string serialized_runtime_env_info;
 };
 }  // namespace internal
 

cpp/src/ray/config_internal.cc

@@ -75,6 +75,13 @@ ABSL_FLAG(std::string,
           "",
           "The default actor lifetime type, `detached` or `non_detached`.");
 
+ABSL_FLAG(std::string, ray_runtime_env, "", "The serialized runtime env.");
+
+ABSL_FLAG(int,
+          ray_runtime_env_hash,
+          -1,
+          "The computed hash of the runtime env for this worker.");
+
 namespace ray {
 namespace internal {
 
@@ -111,6 +118,9 @@ void ConfigInternal::Init(RayConfig &config, int argc, char **argv) {
   if (config.default_actor_lifetime == ActorLifetime::DETACHED) {
     default_actor_lifetime = rpc::JobConfig_ActorLifetime_DETACHED;
   }
+  if (config.runtime_env) {
+    runtime_env = config.runtime_env;
+  }
 
   if (argc != 0 && argv != nullptr) {
     // Parse config from command line.
@@ -158,6 +168,10 @@ void ConfigInternal::Init(RayConfig &config, int argc, char **argv) {
       default_actor_lifetime =
           ParseDefaultActorLifetimeType(FLAGS_ray_default_actor_lifetime.CurrentValue());
     }
+    if (!FLAGS_ray_runtime_env.CurrentValue().empty()) {
+      runtime_env = RuntimeEnv::Deserialize(FLAGS_ray_runtime_env.CurrentValue());
+    }
+    runtime_env_hash = absl::GetFlag<int>(FLAGS_ray_runtime_env_hash);
   }
   worker_type = config.is_worker_ ? WorkerType::WORKER : WorkerType::DRIVER;
   if (worker_type == WorkerType::DRIVER && run_mode == RunMode::CLUSTER) {

cpp/src/ray/config_internal.h

@@ -60,6 +60,10 @@ class ConfigInternal {
 
   std::vector<std::string> head_args = {};
 
+  boost::optional<RuntimeEnv> runtime_env;
+
+  int runtime_env_hash = 0;
+
   // The default actor lifetime type.
   rpc::JobConfig_ActorLifetime default_actor_lifetime =
       rpc::JobConfig_ActorLifetime_NON_DETACHED;

cpp/src/ray/runtime/runtime_env.cc

@@ -0,0 +1,78 @@
+// Copyright 2022 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <google/protobuf/util/json_util.h>
+#include <ray/api/runtime_env.h>
+#include <ray/util/logging.h>
+
+#include "src/ray/protobuf/runtime_env_common.pb.h"
+
+namespace ray {
+
+void RuntimeEnv::SetJsonStr(const std::string &name, const std::string &json_str) {
+  try {
+    json value_j = json::parse(json_str);
+    fields_[name] = value_j;
+  } catch (std::exception &e) {
+    throw ray::internal::RayRuntimeEnvException("Failed to set the field " + name +
+                                                " by json string: " + e.what());
+  }
+}
+
+std::string RuntimeEnv::GetJsonStr(const std::string &name) const {
+  if (!Contains(name)) {
+    throw ray::internal::RayRuntimeEnvException("The field " + name + " not found.");
+  }
+  auto j = fields_[name].get<json>();
+  return j.dump();
+}
+
+bool RuntimeEnv::Contains(const std::string &name) const {
+  return fields_.contains(name);
+}
+
+bool RuntimeEnv::Remove(const std::string &name) {
+  if (Contains(name)) {
+    fields_.erase(name);
+    return true;
+  }
+  return false;
+}
+
+bool RuntimeEnv::Empty() const { return fields_.empty(); }
+
+std::string RuntimeEnv::Serialize() const { return fields_.dump(); }
+
+std::string RuntimeEnv::SerializeToRuntimeEnvInfo() const {
+  rpc::RuntimeEnvInfo runtime_env_info;
+  runtime_env_info.set_serialized_runtime_env(Serialize());
+  std::string serialized_runtime_env_info;
+  RAY_CHECK(google::protobuf::util::MessageToJsonString(runtime_env_info,
+                                                        &serialized_runtime_env_info)
+                .ok());
+  return serialized_runtime_env_info;
+}
+
+RuntimeEnv RuntimeEnv::Deserialize(const std::string &serialized_runtime_env) {
+  RuntimeEnv runtime_env;
+  try {
+    runtime_env.fields_ = json::parse(serialized_runtime_env);
+  } catch (std::exception &e) {
+    throw ray::internal::RayRuntimeEnvException("Failed to deserialize runtime env " +
+                                                serialized_runtime_env + ": " + e.what());
+  }
+  return runtime_env;
+}
+
+}  // namespace ray