diff --git a/python/ray/experimental/state.py b/python/ray/experimental/state.py index 482b26a1e..65913faa8 100644 --- a/python/ray/experimental/state.py +++ b/python/ray/experimental/state.py @@ -150,7 +150,9 @@ class GlobalState(object): result = {"ManagerIDs": manager_ids, "TaskID": binary_to_hex(result_table_message.TaskId()), - "IsPut": bool(result_table_message.IsPut())} + "IsPut": bool(result_table_message.IsPut()), + "DataSize": result_table_message.DataSize(), + "Hash": binary_to_hex(result_table_message.Hash())} return result diff --git a/src/common/format/common.fbs b/src/common/format/common.fbs index 78c5ec2c2..2c0b1d079 100644 --- a/src/common/format/common.fbs +++ b/src/common/format/common.fbs @@ -136,6 +136,10 @@ table ResultTableReply { task_id: string; // Whether the task created the object through a ray.put. is_put: bool; + // The size of the object created. + data_size: long; + // The hash of the object created. + hash: string; } root_type ResultTableReply; diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index 9bbd6df78..a8d54207e 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -1,5 +1,4 @@ #include "redismodule.h" - #include #include @@ -832,9 +831,12 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleString *task_id; RedisModuleString *is_put; + RedisModuleString *data_size; + RedisModuleString *hash; RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task", &task_id, "is_put", - &is_put, NULL); + &is_put, "data_size", &data_size, "hash", &hash, NULL); RedisModule_CloseKey(key); + if (task_id == NULL || is_put == NULL) { return RedisModule_ReplyWithNull(ctx); } @@ -851,8 +853,27 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx, /* Make and return the flatbuffer reply. */ flatbuffers::FlatBufferBuilder fbb; - auto message = CreateResultTableReply(fbb, RedisStringToFlatbuf(fbb, task_id), - bool(is_put_integer)); + long long data_size_value; + + if (data_size == NULL) { + data_size_value = -1; + } else { + RedisModule_StringToLongLong(data_size, &data_size_value); + CHECK(RedisModule_StringToLongLong(data_size, &data_size_value) == + REDISMODULE_OK); + } + + flatbuffers::Offset hash_str; + if (hash == NULL) { + hash_str = fbb.CreateString("", strlen("")); + } else { + hash_str = RedisStringToFlatbuf(fbb, hash); + } + + flatbuffers::Offset message = + CreateResultTableReply(fbb, RedisStringToFlatbuf(fbb, task_id), + bool(is_put_integer), data_size_value, hash_str); + fbb.Finish(message); RedisModuleString *reply = RedisModule_CreateString( ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize()); @@ -863,6 +884,14 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModule_FreeString(ctx, is_put); RedisModule_FreeString(ctx, task_id); + if (data_size != NULL) { + RedisModule_FreeString(ctx, data_size); + } + + if (hash != NULL) { + RedisModule_FreeString(ctx, hash); + } + return REDISMODULE_OK; }