io.ray.serve.dag.DAGNode Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ray-serve Show documentation
java for ray serve
The newest version!
package io.ray.serve.dag;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;

public abstract class DAGNode implements DAGNodeBase {

  private final Object[] boundArgs;

  private final Map boundOptions;

  private final Map boundOtherArgsToResolve;

  private String stableUuid = UUID.randomUUID().toString().replace("-", "");

  public DAGNode(
      Object[] args, Map options, Map otherArgsToResolve) {
    this.boundArgs = args != null ? args : new Object[0];
    this.boundOptions = options != null ? options : new HashMap<>();
    this.boundOtherArgsToResolve =
        otherArgsToResolve != null ? otherArgsToResolve : new HashMap<>();
  }

  @Override
  public  T applyRecursive(Function fn) {
    if (!(fn instanceof CachingFn)) {
      Function newFun = new CachingFn<>(fn);
      return newFun.apply(applyAndReplaceAllChildNodes(node -> node.applyRecursive(newFun)));
    } else {
      return fn.apply(applyAndReplaceAllChildNodes(node -> node.applyRecursive(fn)));
    }
  }

  @Override
  public  DAGNodeBase applyAndReplaceAllChildNodes(Function fn) {
    Object[] newArgs = new Object[boundArgs.length];
    for (int i = 0; i < boundArgs.length; i++) {
      newArgs[i] =
          boundArgs[i] instanceof DAGNodeBase ? fn.apply((DAGNodeBase) boundArgs[i]) : boundArgs[i];
    }

    return copy(newArgs, boundOptions, boundOtherArgsToResolve);
  }

  @Override
  public DAGNodeBase copy(
      Object[] newArgs, Map newOptions, Map newOtherArgsToResolve) {
    DAGNode instance = (DAGNode) copyImpl(newArgs, newOptions, newOtherArgsToResolve);
    instance.stableUuid = stableUuid;
    return instance;
  }

  public Object[] getBoundArgs() {
    return boundArgs;
  }

  public Map getBoundOptions() {
    return boundOptions;
  }

  public Map getBoundOtherArgsToResolve() {
    return boundOtherArgsToResolve;
  }

  @Override
  public String getStableUuid() {
    return stableUuid;
  }

  private static class CachingFn implements Function {

    private final Map cache = new HashMap<>();

    private final Function function;

    public CachingFn(Function function) {
      this.function = function;
    }

    @Override
    public T apply(DAGNodeBase t) {
      return cache.computeIfAbsent(t.getStableUuid(), key -> function.apply(t));
    }
  }
}