mirror of
https://github.com/wassname/ray.git
synced 2026-07-02 03:50:57 +08:00
[gRPC] Migrate gcs data structures to protobuf (#5024)
This commit is contained in:
@@ -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();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user