io.ray.serve.router.ReplicaSet Maven / Gradle / Ivy
package io.ray.serve.router;
import com.google.common.collect.Sets;
import io.ray.api.ActorHandle;
import io.ray.api.BaseActorHandle;
import io.ray.api.ObjectRef;
import io.ray.api.PyActorHandle;
import io.ray.api.Ray;
import io.ray.api.call.PyActorTaskCaller;
import io.ray.api.function.PyActorMethod;
import io.ray.serve.common.Constants;
import io.ray.serve.exception.RayServeException;
import io.ray.serve.generated.ActorNameList;
import io.ray.serve.replica.RayServeWrappedReplica;
import io.ray.serve.util.CollectionUtil;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;
import org.apache.commons.lang3.RandomUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Data structure representing a set of replica actor handles. */
public class ReplicaSet { // TODO ReplicaScheduler
private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaSet.class);
// The key is the name of the actor, and the value is a set of all flight queries objectrefs of
// the actor.
private final Map>> inFlightQueries;
// Map the actor name to the handle of the actor.
private final Map allActorHandles;
private boolean hasPullReplica = false;
public ReplicaSet(String deploymentName) {
this.inFlightQueries = new ConcurrentHashMap<>();
this.allActorHandles = new ConcurrentHashMap<>();
}
public synchronized void updateWorkerReplicas(Object actorSet) {
if (null != actorSet) {
Set actorNameSet = new HashSet<>(((ActorNameList) actorSet).getNamesList());
Set added = new HashSet<>(Sets.difference(actorNameSet, inFlightQueries.keySet()));
Set removed = new HashSet<>(Sets.difference(inFlightQueries.keySet(), actorNameSet));
added.forEach(
name -> {
Optional handleOptional =
Ray.getActor(name, Constants.SERVE_NAMESPACE);
if (handleOptional.isPresent()) {
allActorHandles.put(name, handleOptional.get());
inFlightQueries.put(name, Sets.newConcurrentHashSet());
} else {
LOGGER.warn("Can not get actor handle. actor name is {}", name);
}
});
removed.forEach(inFlightQueries::remove);
removed.forEach(allActorHandles::remove);
if (added.size() > 0 || removed.size() > 0) {
LOGGER.info("ReplicaSet: +{}, -{} replicas.", added.size(), removed.size());
}
}
hasPullReplica = true;
}
/**
* Given a query, submit it to a replica and return the object ref. This method will keep track of
* the in flight queries for each replicas and only send a query to available replicas (determined
* by the max_concurrent_quries value.)
*
* @param query the incoming query.
* @return ray.ObjectRef
*/
public ObjectRef