[gRPC] Migrate gcs data structures to protobuf (#5024)

This commit is contained in:
Hao Chen
2019-06-26 05:31:19 +08:00
committed by Philipp Moritz
parent bd8aceb896
commit 0131353d42
52 changed files with 1465 additions and 1642 deletions
@@ -1,7 +1,7 @@
package org.ray.runtime.gcs;
import com.google.common.base.Preconditions;
import java.nio.ByteBuffer;
import com.google.protobuf.InvalidProtocolBufferException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@@ -13,10 +13,10 @@ import org.ray.api.id.BaseId;
import org.ray.api.id.TaskId;
import org.ray.api.id.UniqueId;
import org.ray.api.runtimecontext.NodeInfo;
import org.ray.runtime.generated.ActorCheckpointIdData;
import org.ray.runtime.generated.ClientTableData;
import org.ray.runtime.generated.EntryType;
import org.ray.runtime.generated.TablePrefix;
import org.ray.runtime.generated.Gcs.ActorCheckpointIdData;
import org.ray.runtime.generated.Gcs.ClientTableData;
import org.ray.runtime.generated.Gcs.ClientTableData.EntryType;
import org.ray.runtime.generated.Gcs.TablePrefix;
import org.ray.runtime.util.IdUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -51,7 +51,7 @@ public class GcsClient {
}
public List<NodeInfo> getAllNodeInfo() {
final String prefix = TablePrefix.name(TablePrefix.CLIENT);
final String prefix = TablePrefix.CLIENT.toString();
final byte[] key = ArrayUtils.addAll(prefix.getBytes(), UniqueId.NIL.getBytes());
List<byte[]> results = primary.lrange(key, 0, -1);
@@ -63,36 +63,42 @@ public class GcsClient {
Map<UniqueId, NodeInfo> clients = new HashMap<>();
for (byte[] result : results) {
Preconditions.checkNotNull(result);
ClientTableData data = ClientTableData.getRootAsClientTableData(ByteBuffer.wrap(result));
final UniqueId clientId = UniqueId.fromByteBuffer(data.clientIdAsByteBuffer());
ClientTableData data = null;
try {
data = ClientTableData.parseFrom(result);
} catch (InvalidProtocolBufferException e) {
throw new RuntimeException("Received invalid protobuf data from GCS.");
}
final UniqueId clientId = UniqueId
.fromByteBuffer(data.getClientId().asReadOnlyByteBuffer());
if (data.entryType() == EntryType.INSERTION) {
if (data.getEntryType() == EntryType.INSERTION) {
//Code path of node insertion.
Map<String, Double> resources = new HashMap<>();
// Compute resources.
Preconditions.checkState(
data.resourcesTotalLabelLength() == data.resourcesTotalCapacityLength());
for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
resources.put(data.resourcesTotalLabel(i), data.resourcesTotalCapacity(i));
data.getResourcesTotalLabelCount() == data.getResourcesTotalCapacityCount());
for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
resources.put(data.getResourcesTotalLabel(i), data.getResourcesTotalCapacity(i));
}
NodeInfo nodeInfo = new NodeInfo(
clientId, data.nodeManagerAddress(), true, resources);
clientId, data.getNodeManagerAddress(), true, resources);
clients.put(clientId, nodeInfo);
} else if (data.entryType() == EntryType.RES_CREATEUPDATE) {
} else if (data.getEntryType() == EntryType.RES_CREATEUPDATE) {
Preconditions.checkState(clients.containsKey(clientId));
NodeInfo nodeInfo = clients.get(clientId);
for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
nodeInfo.resources.put(data.resourcesTotalLabel(i), data.resourcesTotalCapacity(i));
for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
nodeInfo.resources.put(data.getResourcesTotalLabel(i), data.getResourcesTotalCapacity(i));
}
} else if (data.entryType() == EntryType.RES_DELETE) {
} else if (data.getEntryType() == EntryType.RES_DELETE) {
Preconditions.checkState(clients.containsKey(clientId));
NodeInfo nodeInfo = clients.get(clientId);
for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
nodeInfo.resources.remove(data.resourcesTotalLabel(i));
for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
nodeInfo.resources.remove(data.getResourcesTotalLabel(i));
}
} else {
// Code path of node deletion.
Preconditions.checkState(data.entryType() == EntryType.DELETION);
Preconditions.checkState(data.getEntryType() == EntryType.DELETION);
NodeInfo nodeInfo = new NodeInfo(clientId, clients.get(clientId).nodeAddress,
false, clients.get(clientId).resources);
clients.put(clientId, nodeInfo);
@@ -107,7 +113,7 @@ public class GcsClient {
*/
public boolean actorExists(UniqueId actorId) {
byte[] key = ArrayUtils.addAll(
TablePrefix.name(TablePrefix.ACTOR).getBytes(), actorId.getBytes());
TablePrefix.ACTOR.toString().getBytes(), actorId.getBytes());
return primary.exists(key);
}
@@ -115,7 +121,7 @@ public class GcsClient {
* Query whether the raylet task exists in Gcs.
*/
public boolean rayletTaskExistsInGcs(TaskId taskId) {
byte[] key = ArrayUtils.addAll(TablePrefix.name(TablePrefix.RAYLET_TASK).getBytes(),
byte[] key = ArrayUtils.addAll(TablePrefix.RAYLET_TASK.toString().getBytes(),
taskId.getBytes());
RedisClient client = getShardClient(taskId);
return client.exists(key);
@@ -126,19 +132,26 @@ public class GcsClient {
*/
public List<Checkpoint> getCheckpointsForActor(UniqueId actorId) {
List<Checkpoint> checkpoints = new ArrayList<>();
final String prefix = TablePrefix.name(TablePrefix.ACTOR_CHECKPOINT_ID);
final String prefix = TablePrefix.ACTOR_CHECKPOINT_ID.toString();
final byte[] key = ArrayUtils.addAll(prefix.getBytes(), actorId.getBytes());
RedisClient client = getShardClient(actorId);
byte[] result = client.get(key);
if (result != null) {
ActorCheckpointIdData data =
ActorCheckpointIdData.getRootAsActorCheckpointIdData(ByteBuffer.wrap(result));
UniqueId[] checkpointIds = IdUtil.getUniqueIdsFromByteBuffer(
data.checkpointIdsAsByteBuffer());
ActorCheckpointIdData data = null;
try {
data = ActorCheckpointIdData.parseFrom(result);
} catch (InvalidProtocolBufferException e) {
throw new RuntimeException("Received invalid protobuf data from GCS.");
}
UniqueId[] checkpointIds = new UniqueId[data.getCheckpointIdsCount()];
for (int i = 0; i < checkpointIds.length; i++) {
checkpointIds[i] = UniqueId
.fromByteBuffer(data.getCheckpointIds(i).asReadOnlyByteBuffer());
}
for (int i = 0; i < checkpointIds.length; i++) {
checkpoints.add(new Checkpoint(checkpointIds[i], data.timestamps(i)));
checkpoints.add(new Checkpoint(checkpointIds[i], data.getTimestamps(i)));
}
}
checkpoints.sort((x, y) -> Long.compare(y.timestamp, x.timestamp));
@@ -16,7 +16,7 @@ import org.ray.api.id.ObjectId;
import org.ray.runtime.AbstractRayRuntime;
import org.ray.runtime.RayDevRuntime;
import org.ray.runtime.config.RunMode;
import org.ray.runtime.generated.ErrorType;
import org.ray.runtime.generated.Gcs.ErrorType;
import org.ray.runtime.util.IdUtil;
import org.ray.runtime.util.Serializer;
import org.slf4j.Logger;
@@ -29,12 +29,12 @@ public class ObjectStoreProxy {
private static final Logger LOGGER = LoggerFactory.getLogger(ObjectStoreProxy.class);
private static final byte[] WORKER_EXCEPTION_META = String.valueOf(ErrorType.WORKER_DIED)
.getBytes();
private static final byte[] ACTOR_EXCEPTION_META = String.valueOf(ErrorType.ACTOR_DIED)
.getBytes();
private static final byte[] WORKER_EXCEPTION_META = String
.valueOf(ErrorType.WORKER_DIED.getNumber()).getBytes();
private static final byte[] ACTOR_EXCEPTION_META = String
.valueOf(ErrorType.ACTOR_DIED.getNumber()).getBytes();
private static final byte[] UNRECONSTRUCTABLE_EXCEPTION_META = String
.valueOf(ErrorType.OBJECT_UNRECONSTRUCTABLE).getBytes();
.valueOf(ErrorType.OBJECT_UNRECONSTRUCTABLE.getNumber()).getBytes();
private static final byte[] RAW_TYPE_META = "RAW".getBytes();