wehub-resource-sync
141 lines
5.6 KiB
Cython
141 lines
5.6 KiB
Cython
from libcpp.string cimport string as c_string
|
|
from libcpp cimport bool as c_bool
|
|
from libcpp.vector cimport vector as c_vector
|
|
from libcpp.unordered_map cimport unordered_map
|
|
from libcpp.memory cimport unique_ptr
|
|
from libc.stdint cimport (
|
|
int32_t as c_int32_t,
|
|
uint32_t as c_uint32_t,
|
|
int64_t as c_int64_t,
|
|
)
|
|
from ray.includes.unique_ids cimport (
|
|
CActorID,
|
|
CJobID,
|
|
CNodeID,
|
|
CObjectID,
|
|
CWorkerID,
|
|
CPlacementGroupID,
|
|
)
|
|
from ray.includes.common cimport (
|
|
CRayStatus,
|
|
CGcsClientOptions,
|
|
)
|
|
from ray.includes.optional cimport (
|
|
optional
|
|
)
|
|
|
|
cdef extern from "ray/gcs_rpc_client/global_state_accessor.h" nogil:
|
|
cdef cppclass CGlobalStateAccessor "ray::gcs::GlobalStateAccessor":
|
|
CGlobalStateAccessor(const CGcsClientOptions&)
|
|
c_bool Connect()
|
|
void Disconnect()
|
|
c_vector[c_string] GetAllJobInfo(
|
|
c_bool skip_submission_job_info_field, c_bool skip_is_running_tasks_field)
|
|
CJobID GetNextJobID()
|
|
c_vector[c_string] GetAllNodeInfo()
|
|
c_vector[c_string] GetAllAvailableResources()
|
|
c_vector[c_string] GetAllTotalResources()
|
|
unordered_map[CNodeID, c_int64_t] GetDrainingNodes()
|
|
unique_ptr[c_string] GetInternalKV(
|
|
const c_string &namespace, const c_string &key)
|
|
c_vector[c_string] GetAllTaskEvents()
|
|
unique_ptr[c_string] GetObjectInfo(const CObjectID &object_id)
|
|
unique_ptr[c_string] GetAllResourceUsage()
|
|
c_vector[c_string] GetAllActorInfo(
|
|
optional[CActorID], optional[CJobID], optional[c_string])
|
|
unique_ptr[c_string] GetActorInfo(const CActorID &actor_id)
|
|
unique_ptr[c_string] GetWorkerInfo(const CWorkerID &worker_id)
|
|
c_vector[c_string] GetAllWorkerInfo()
|
|
c_bool AddWorkerInfo(const c_string &serialized_string)
|
|
c_bool UpdateWorkerDebuggerPort(const CWorkerID &worker_id,
|
|
const c_uint32_t debuger_port)
|
|
c_bool UpdateWorkerNumPausedThreads(const CWorkerID &worker_id,
|
|
const c_int32_t num_paused_threads_delta)
|
|
c_uint32_t GetWorkerDebuggerPort(const CWorkerID &worker_id)
|
|
unique_ptr[c_string] GetPlacementGroupInfo(
|
|
const CPlacementGroupID &placement_group_id)
|
|
unique_ptr[c_string] GetPlacementGroupByName(
|
|
const c_string &placement_group_name,
|
|
const c_string &ray_namespace,
|
|
)
|
|
c_vector[c_string] GetAllPlacementGroupInfo()
|
|
c_string GetSystemConfig()
|
|
CRayStatus GetNode(
|
|
const c_string &node_id_hex_str,
|
|
c_string *node_info)
|
|
|
|
cdef extern from * namespace "ray::gcs" nogil:
|
|
"""
|
|
#include <thread>
|
|
#include "ray/gcs/store_client_kv.h"
|
|
#include "ray/gcs/store_client/redis_store_client.h"
|
|
#include "ray/util/clock.h"
|
|
#include "ray/util/raii.h"
|
|
namespace ray {
|
|
namespace gcs {
|
|
|
|
bool RedisGetKeySync(const std::string& host,
|
|
int32_t port,
|
|
const std::string& username,
|
|
const std::string& password,
|
|
bool use_ssl,
|
|
const std::string& config,
|
|
const std::string& key,
|
|
std::string* data) {
|
|
// Logging default value see class `RayLog`.
|
|
InitShutdownRAII ray_log_shutdown_raii(ray::RayLog::StartRayLog,
|
|
ray::RayLog::ShutDownRayLog,
|
|
"ray_init",
|
|
ray::RayLogLevel::WARNING,
|
|
/*log_filepath=*/"",
|
|
/*err_log_filepath=*/"",
|
|
/*log_rotation_max_size=*/1ULL << 29,
|
|
/*log_rotation_file_num=*/10);
|
|
|
|
std::string config_list;
|
|
RAY_CHECK(absl::Base64Unescape(config, &config_list));
|
|
RayConfig::instance().initialize(config_list);
|
|
|
|
instrumented_io_context io_service{/*enable_lag_probe=*/false, /*running_on_single_thread=*/true};
|
|
RedisClientOptions options{host, port, username, password, use_ssl};
|
|
Clock clock;
|
|
auto client = std::make_unique<StoreClientInternalKV>(
|
|
std::make_unique<RedisStoreClient>(io_service, options, clock));
|
|
|
|
bool ret_val = false;
|
|
client->Get("session", key, {[&](std::optional<std::string> result) {
|
|
if (result.has_value()) {
|
|
*data = result.value();
|
|
ret_val = true;
|
|
} else {
|
|
RAY_LOG(INFO) << "Failed to retrieve the key " << key
|
|
<< " from persistent storage.";
|
|
ret_val = false;
|
|
}
|
|
}, io_service});
|
|
io_service.run_for(std::chrono::milliseconds(1000));
|
|
|
|
return ret_val;
|
|
}
|
|
|
|
}
|
|
}
|
|
"""
|
|
c_bool RedisGetKeySync(const c_string& host,
|
|
c_int32_t port,
|
|
const c_string& username,
|
|
const c_string& password,
|
|
c_bool use_ssl,
|
|
const c_string& config,
|
|
const c_string& key,
|
|
c_string* data)
|
|
|
|
|
|
cdef extern from "ray/gcs/store_client/redis_store_client.h" namespace "ray::gcs" nogil:
|
|
c_bool RedisDelKeyPrefixSync(const c_string& host,
|
|
c_int32_t port,
|
|
const c_string& username,
|
|
const c_string& password,
|
|
c_bool use_ssl,
|
|
const c_string& key_prefix)
|