ai.h2o.automl.WorkAllocations Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of h2o-automl Show documentation
H2O AutoML
The newest version!
package ai.h2o.automl;

import water.Iced;
import water.util.ArrayUtils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Predicate;
import java.util.stream.Stream;

public class WorkAllocations extends Iced {

  public enum JobType {
    Unknown,
    ModelBuild,
    HyperparamSearch,
    Selection,
    Dynamic,
  }

  public static class Work extends Iced {
    String _id;
    IAlgo _algo;
    JobType _type;
    int _priorityGroup;
    int _weight;

    Work(String id, IAlgo algo, JobType type, int priorityGroup, int weight) {
      this._algo = algo;
      this._type = type;
      this._id = id;
      this._priorityGroup = priorityGroup;
      this._weight = weight;
    }

    public int consume() {
      int consumed = _weight;
      _weight = 0;
      return consumed;
    }

    @Override
    public String toString() {
      final StringBuilder sb = new StringBuilder("Work{")
              .append(_id).append(", ")
              .append(_algo.name()).append(", ")
              .append(_type).append(", ")
              .append("group=").append(_priorityGroup).append(", ")
              .append("weight=").append(_weight)
              .append('}');
      return sb.toString();
    }
  }

  private boolean frozen;
  private Work[] allocations = new Work[0];

  WorkAllocations allocate(Work work) {
    if (frozen) throw new IllegalStateException("Can not allocate new work.");
    allocations = ArrayUtils.append(allocations, work);
    return this;
  }

  WorkAllocations freeze() {
    frozen = true;
    return this;
  }

  void remove(IAlgo algo) {
    if (frozen) throw new IllegalStateException("Can not modify allocations.");
    List filtered = new ArrayList<>(allocations.length);
    for (Work alloc : allocations) {
      if (!algo.name().equals(alloc._algo.name())) {
        filtered.add(alloc);
      }
    }
    allocations = filtered.toArray(new Work[0]);
  }

  public Work getAllocation(String id, IAlgo algo) {
    for (Work alloc : allocations) {
      if (alloc._algo.name().equals(algo.name()) && alloc._id.equals(id)) return alloc;
    }
    return null;
  }

  public Work[] getAllocations(Predicate predicate) {
    return Stream.of(allocations)
            .filter(predicate)
            .toArray(Work[]::new);
  }

  private int sum(Work[] workItems) {
    int tot = 0;
    for (Work item : workItems) {
      if (item._weight > 0)
        tot += item._weight;
    }
    return tot;
  }

  int remainingWork() {
    return sum(allocations);
  }

  int remainingWork(Predicate predicate) {
    return sum(getAllocations(predicate));
  }

  float remainingWorkRatio(Work work) {
    return (float) work._weight / remainingWork();
  }

  float remainingWorkRatio(Work work, Predicate predicate) {
    return (float) work._weight / remainingWork(predicate);
  }

  @Override
  public String toString() {
    return Arrays.toString(allocations);
  }
}