giraph-core/src/main/java/org/apache/giraph/worker/WorkerAggregatorHandler.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.giraph.worker;

import java.io.IOException;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.giraph.bsp.CentralizedServiceWorker;
import org.apache.giraph.comm.GlobalCommType;
import org.apache.giraph.comm.aggregators.AggregatorUtils;
import org.apache.giraph.comm.aggregators.AllAggregatorServerData;
import org.apache.giraph.comm.aggregators.GlobalCommValueOutputStream;
import org.apache.giraph.comm.aggregators.OwnerAggregatorServerData;
import org.apache.giraph.comm.aggregators.WorkerAggregatorRequestProcessor;
import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.reducers.ReduceOperation;
import org.apache.giraph.reducers.Reducer;
import org.apache.giraph.utils.UnsafeByteArrayOutputStream;
import org.apache.giraph.utils.UnsafeReusableByteArrayInput;
import org.apache.giraph.utils.WritableUtils;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.util.Progressable;
import org.apache.log4j.Logger;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

/** Handler for reduce/broadcast on the workers */
public class WorkerAggregatorHandler implements WorkerThreadGlobalCommUsage {
  /** Class logger */
  private static final Logger LOG =
      Logger.getLogger(WorkerAggregatorHandler.class);
  /** Map of broadcasted values */
  private final Map<String, Writable> broadcastedMap =
      Maps.newHashMap();
  /** Map of reducers currently being reduced */
  private final Map<String, Reducer<Object, Writable>> reducerMap =
      Maps.newHashMap();

  /** Service worker */
  private final CentralizedServiceWorker<?, ?, ?> serviceWorker;
  /** Progressable for reporting progress */
  private final Progressable progressable;
  /** How big a single aggregator request can be */
  private final int maxBytesPerAggregatorRequest;
  /** Giraph configuration */
  private final ImmutableClassesGiraphConfiguration conf;

  /**
   * Constructor
   *
   * @param serviceWorker Service worker
   * @param conf          Giraph configuration
   * @param progressable  Progressable for reporting progress
   */
  public WorkerAggregatorHandler(
      CentralizedServiceWorker<?, ?, ?> serviceWorker,
      ImmutableClassesGiraphConfiguration conf,
      Progressable progressable) {
    this.serviceWorker = serviceWorker;
    this.progressable = progressable;
    this.conf = conf;
    maxBytesPerAggregatorRequest = conf.getInt(
        AggregatorUtils.MAX_BYTES_PER_AGGREGATOR_REQUEST,
        AggregatorUtils.MAX_BYTES_PER_AGGREGATOR_REQUEST_DEFAULT);
  }

  @Override
  public <B extends Writable> B getBroadcast(String name) {
    B value = (B) broadcastedMap.get(name);
    if (value == null) {
      LOG.warn("getBroadcast: " +
          AggregatorUtils.getUnregisteredBroadcastMessage(name,
              broadcastedMap.size() != 0, conf));
    }
    return value;
  }

  @Override
  public void reduce(String name, Object value) {
    Reducer<Object, Writable> reducer = reducerMap.get(name);
    if (reducer != null) {
      progressable.progress();
      synchronized (reducer) {
        reducer.reduce(value);
      }
    } else {
      throw new IllegalStateException("reduce: " +
          AggregatorUtils.getUnregisteredReducerMessage(name,
              reducerMap.size() != 0, conf));
    }
  }

  /**
   * Combine partially reduced value into currently reduced value.
   * @param name Name of the reducer
   * @param valueToReduce Partial value to reduce
   */
  @Override
  public void reduceMerge(String name, Writable valueToReduce) {
    Reducer<Object, Writable> reducer = reducerMap.get(name);
    if (reducer != null) {
      progressable.progress();
      synchronized (reducer) {
        reducer.reduceMerge(valueToReduce);
      }
    } else {
      throw new IllegalStateException("reduce: " +
          AggregatorUtils.getUnregisteredReducerMessage(name,
              reducerMap.size() != 0, conf));
    }
  }

  /**
   * Prepare aggregators for current superstep
   *
   * @param requestProcessor Request processor for aggregators
   */
  public void prepareSuperstep(
      WorkerAggregatorRequestProcessor requestProcessor) {
    broadcastedMap.clear();
    reducerMap.clear();

    if (LOG.isDebugEnabled()) {
      LOG.debug("prepareSuperstep: Start preparing aggregators");
    }
    AllAggregatorServerData allGlobalCommData =
        serviceWorker.getServerData().getAllAggregatorData();
    // Wait for my aggregators
    Iterable<byte[]> dataToDistribute =
        allGlobalCommData.getDataFromMasterWhenReady(
            serviceWorker.getMasterInfo());
    try {
      // Distribute my aggregators
      requestProcessor.distributeReducedValues(dataToDistribute);
    } catch (IOException e) {
      throw new IllegalStateException("prepareSuperstep: " +
          "IOException occurred while trying to distribute aggregators", e);
    }
    // Wait for all other aggregators and store them
    allGlobalCommData.fillNextSuperstepMapsWhenReady(
        getOtherWorkerIdsSet(), broadcastedMap,
        reducerMap);
    if (LOG.isDebugEnabled()) {
      LOG.debug("prepareSuperstep: Aggregators prepared");
    }
  }

  /**
   * Send aggregators to their owners and in the end to the master
   *
   * @param requestProcessor Request processor for aggregators
   */
  public void finishSuperstep(
      WorkerAggregatorRequestProcessor requestProcessor) {
    if (LOG.isInfoEnabled()) {
      LOG.info("finishSuperstep: Start gathering aggregators, " +
          "workers will send their aggregated values " +
          "once they are done with superstep computation");
    }
    OwnerAggregatorServerData ownerGlobalCommData =
        serviceWorker.getServerData().getOwnerAggregatorData();
    // First send partial aggregated values to their owners and determine
    // which aggregators belong to this worker
    for (Map.Entry<String, Reducer<Object, Writable>> entry :
        reducerMap.entrySet()) {
      try {
        boolean sent = requestProcessor.sendReducedValue(entry.getKey(),
            entry.getValue().getCurrentValue());
        if (!sent) {
          // If it's my aggregator, add it directly
          ownerGlobalCommData.reduce(entry.getKey(),
              entry.getValue().getCurrentValue());
        }
      } catch (IOException e) {
        throw new IllegalStateException("finishSuperstep: " +
            "IOException occurred while sending aggregator " +
            entry.getKey() + " to its owner", e);
      }
      progressable.progress();
    }
    try {
      // Flush
      requestProcessor.flush();
    } catch (IOException e) {
      throw new IllegalStateException("finishSuperstep: " +
          "IOException occurred while sending aggregators to owners", e);
    }

    // Wait to receive partial aggregated values from all other workers
    Iterable<Map.Entry<String, Writable>> myReducedValues =
        ownerGlobalCommData.getMyReducedValuesWhenReady(
            getOtherWorkerIdsSet());

    // Send final aggregated values to master
    GlobalCommValueOutputStream globalOutput =
        new GlobalCommValueOutputStream(false);
    for (Map.Entry<String, Writable> entry : myReducedValues) {
      try {
        int currentSize = globalOutput.addValue(entry.getKey(),
            GlobalCommType.REDUCED_VALUE,
            entry.getValue());
        if (currentSize > maxBytesPerAggregatorRequest) {
          requestProcessor.sendReducedValuesToMaster(
              globalOutput.flush());
        }
        progressable.progress();
      } catch (IOException e) {
        throw new IllegalStateException("finishSuperstep: " +
            "IOException occurred while writing aggregator " +
            entry.getKey(), e);
      }
    }
    try {
      requestProcessor.sendReducedValuesToMaster(globalOutput.flush());
    } catch (IOException e) {
      throw new IllegalStateException("finishSuperstep: " +
          "IOException occured while sending aggregators to master", e);
    }
    // Wait for master to receive aggregated values before proceeding
    serviceWorker.getWorkerClient().waitAllRequests();

    ownerGlobalCommData.reset();
    if (LOG.isDebugEnabled()) {
      LOG.debug("finishSuperstep: Aggregators finished");
    }
  }

  /**
   * Create new aggregator usage which will be used by one of the compute
   * threads.
   *
   * @return New aggregator usage
   */
  public WorkerThreadGlobalCommUsage newThreadAggregatorUsage() {
    if (AggregatorUtils.useThreadLocalAggregators(conf)) {
      return new ThreadLocalWorkerGlobalCommUsage();
    } else {
      return this;
    }
  }

  @Override
  public void finishThreadComputation() {
    // If we don't use thread-local aggregators, all the aggregated values
    // are already in this object
  }

  /**
   * Get set of all worker task ids except the current one
   *
   * @return Set of all other worker task ids
   */
  public Set<Integer> getOtherWorkerIdsSet() {
    Set<Integer> otherWorkers = Sets.newHashSetWithExpectedSize(
        serviceWorker.getWorkerInfoList().size());
    for (WorkerInfo workerInfo : serviceWorker.getWorkerInfoList()) {
      if (workerInfo.getTaskId() != serviceWorker.getWorkerInfo().getTaskId()) {
        otherWorkers.add(workerInfo.getTaskId());
      }
    }
    return otherWorkers;
  }

  /**
  * Not thread-safe implementation of {@link WorkerThreadGlobalCommUsage}.
  * We can use one instance of this object per thread to prevent
  * synchronizing on each aggregate() call. In the end of superstep,
  * values from each of these will be aggregated back to {@link
  * WorkerThreadGlobalCommUsage}
  */
  public class ThreadLocalWorkerGlobalCommUsage
    implements WorkerThreadGlobalCommUsage {
    /** Thread-local reducer map */
    private final Map<String, Reducer<Object, Writable>> threadReducerMap;

    /**
    * Constructor
    *
    * Creates new instances of all reducers from
    * {@link WorkerAggregatorHandler}
    */
    public ThreadLocalWorkerGlobalCommUsage() {
      threadReducerMap = Maps.newHashMapWithExpectedSize(
          WorkerAggregatorHandler.this.reducerMap.size());

      UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream();
      UnsafeReusableByteArrayInput in = new UnsafeReusableByteArrayInput();

      for (Entry<String, Reducer<Object, Writable>> entry :
          reducerMap.entrySet()) {
        ReduceOperation<Object, Writable> globalReduceOp =
            entry.getValue().getReduceOp();

        ReduceOperation<Object, Writable> threadLocalCopy =
            WritableUtils.createCopy(out, in, globalReduceOp, conf);

        threadReducerMap.put(entry.getKey(), new Reducer<>(threadLocalCopy));
      }
    }

    @Override
    public void reduce(String name, Object value) {
      Reducer<Object, Writable> reducer = threadReducerMap.get(name);
      if (reducer != null) {
        progressable.progress();
        reducer.reduce(value);
      } else {
        throw new IllegalStateException("reduce: " +
            AggregatorUtils.getUnregisteredAggregatorMessage(name,
                threadReducerMap.size() != 0, conf));
      }
    }

    @Override
    public void reduceMerge(String name, Writable value) {
      Reducer<Object, Writable> reducer = threadReducerMap.get(name);
      if (reducer != null) {
        progressable.progress();
        reducer.reduceMerge(value);
      } else {
        throw new IllegalStateException("reduceMerge: " +
            AggregatorUtils.getUnregisteredAggregatorMessage(name,
                threadReducerMap.size() != 0, conf));
      }
    }

    @Override
    public <B extends Writable> B getBroadcast(String name) {
      return WorkerAggregatorHandler.this.getBroadcast(name);
    }

    @Override
    public void finishThreadComputation() {
      // Aggregate the values this thread's vertices provided back to
      // WorkerAggregatorHandler
      for (Entry<String, Reducer<Object, Writable>> entry :
          threadReducerMap.entrySet()) {
        WorkerAggregatorHandler.this.reduceMerge(entry.getKey(),
            entry.getValue().getCurrentValue());
      }
    }
  }

}