giraph-examples/src/test/java/org/apache/giraph/TestCheckpointing.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;

import org.apache.giraph.aggregators.LongSumAggregator;
import org.apache.giraph.bsp.BspService;
import org.apache.giraph.conf.GiraphConfiguration;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.edge.Edge;
import org.apache.giraph.edge.EdgeFactory;
import org.apache.giraph.examples.SimpleSuperstepComputation;
import org.apache.giraph.graph.BasicComputation;
import org.apache.giraph.graph.Vertex;
import org.apache.giraph.job.GiraphJob;
import org.apache.giraph.master.DefaultMasterCompute;
import org.apache.giraph.worker.DefaultWorkerContext;
import org.apache.giraph.zk.ZooKeeperExt;
import org.apache.giraph.zk.ZooKeeperManager;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.FloatWritable;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Writable;
import org.apache.log4j.Logger;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.junit.Assert;
import org.junit.Test;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.List;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/**
 * Tests that worker context and master computation
 * are properly saved and loaded back at checkpoint.
 */
public class TestCheckpointing extends BspCase {

  /** Class logger */
  private static final Logger LOG =
      Logger.getLogger(TestCheckpointing.class);
  /** ID to be used with test job */
  public static final String TEST_JOB_ID = "test_job";

  private static SuperstepCallback SUPERSTEP_CALLBACK;

  /**
   * Create the test case
   */
  public TestCheckpointing() {
    super(TestCheckpointing.class.getName());
  }

  @Test
  public void testBspCheckpoint() throws InterruptedException, IOException, ClassNotFoundException {
    testBspCheckpoint(false);
  }

  @Test
  public void testAsyncMessageStoreCheckpoint() throws InterruptedException, IOException, ClassNotFoundException {
    testBspCheckpoint(true);
  }

  public void testBspCheckpoint(boolean useAsyncMessageStore)
      throws IOException, InterruptedException, ClassNotFoundException {
    Path checkpointsDir = getTempPath("checkpointing");
    GiraphConfiguration conf = new GiraphConfiguration();
    if (useAsyncMessageStore) {
      GiraphConstants.ASYNC_MESSAGE_STORE_THREADS_COUNT.set(conf, 2);
    }

    SUPERSTEP_CALLBACK = null;

    GiraphConstants.CLEANUP_CHECKPOINTS_AFTER_SUCCESS.set(conf, false);
    conf.setCheckpointFrequency(2);

    long idSum = runOriginalJob(checkpointsDir, conf);
    assertEquals(10, idSum);

    SUPERSTEP_CALLBACK = new SuperstepCallback() {
      @Override
      public void superstep(long superstep,
                            ImmutableClassesGiraphConfiguration<LongWritable, IntWritable, FloatWritable> conf) {
        if (superstep < 2) {
          Assert.fail("Restarted JOB should not be executed on superstep " + superstep);
        }
      }
    };

    runRestartedJob(checkpointsDir, conf, idSum, 2);


  }

  private void runRestartedJob(Path checkpointsDir, GiraphConfiguration conf, long idSum, long restartFrom) throws IOException, InterruptedException, ClassNotFoundException {
    Path outputPath;
    LOG.info("testBspCheckpoint: Restarting from the latest superstep " +
        "with checkpoint path = " + checkpointsDir);
    outputPath = getTempPath("checkpointing_restarted");

    GiraphConstants.RESTART_JOB_ID.set(conf, TEST_JOB_ID);
    conf.set("mapred.job.id", "restarted_test_job");
    if (restartFrom >= 0) {
      conf.set(GiraphConstants.RESTART_SUPERSTEP, Long.toString(restartFrom));
    }

    GiraphJob restartedJob = prepareJob(getCallingMethodName() + "Restarted",
        conf, outputPath);

    GiraphConstants.CHECKPOINT_DIRECTORY.set(restartedJob.getConfiguration(),
        checkpointsDir.toString());

    assertTrue(restartedJob.run(true));


    if (!runningInDistributedMode()) {
      long idSumRestarted =
          CheckpointVertexWorkerContext
              .getFinalSum();
      LOG.info("testBspCheckpoint: idSumRestarted = " +
          idSumRestarted);
      assertEquals(idSum, idSumRestarted);
    }
  }

  private long runOriginalJob(Path checkpointsDir,  GiraphConfiguration conf) throws IOException, InterruptedException, ClassNotFoundException {
    Path outputPath = getTempPath("checkpointing_original");
    conf.setComputationClass(
        CheckpointComputation.class);
    conf.setWorkerContextClass(
        CheckpointVertexWorkerContext.class);
    conf.setMasterComputeClass(
        CheckpointVertexMasterCompute.class);
    conf.setVertexInputFormatClass(SimpleSuperstepComputation.SimpleSuperstepVertexInputFormat.class);
    conf.setVertexOutputFormatClass(SimpleSuperstepComputation.SimpleSuperstepVertexOutputFormat.class);
    conf.set("mapred.job.id", TEST_JOB_ID);
    GiraphJob job = prepareJob(getCallingMethodName(), conf, outputPath);

    GiraphConfiguration configuration = job.getConfiguration();
    GiraphConstants.CHECKPOINT_DIRECTORY.set(configuration, checkpointsDir.toString());

    assertTrue(job.run(true));

    long idSum = 0;
    if (!runningInDistributedMode()) {
      FileStatus fileStatus = getSinglePartFileStatus(job.getConfiguration(),
          outputPath);
      idSum = CheckpointVertexWorkerContext
          .getFinalSum();
      LOG.info("testBspCheckpoint: idSum = " + idSum +
          " fileLen = " + fileStatus.getLen());
    }
    return idSum;
  }


  /**
   * Actual computation.
   */
  public static class CheckpointComputation extends
      BasicComputation<LongWritable, IntWritable, FloatWritable,
          FloatWritable> {
    @Override
    public void compute(
        Vertex<LongWritable, IntWritable, FloatWritable> vertex,
        Iterable<FloatWritable> messages) throws IOException {
      CheckpointVertexWorkerContext workerContext = getWorkerContext();
      assertEquals(getSuperstep() + 1, workerContext.testValue);

      if (getSuperstep() > 4) {
        vertex.voteToHalt();
        return;
      }

      aggregate(LongSumAggregator.class.getName(),
          new LongWritable(vertex.getId().get()));

      float msgValue = 0.0f;
      for (FloatWritable message : messages) {
        float curMsgValue = message.get();
        msgValue += curMsgValue;
      }

      int vertexValue = vertex.getValue().get();
      vertex.setValue(new IntWritable(vertexValue + (int) msgValue));
      for (Edge<LongWritable, FloatWritable> edge : vertex.getEdges()) {
        FloatWritable newEdgeValue = new FloatWritable(edge.getValue().get() +
            (float) vertexValue);
        Edge<LongWritable, FloatWritable> newEdge =
            EdgeFactory.create(edge.getTargetVertexId(), newEdgeValue);
        vertex.addEdge(newEdge);
        sendMessage(edge.getTargetVertexId(), newEdgeValue);

      }
    }
  }

  @Test
  public void testManualCheckpointAtTheBeginning()
      throws InterruptedException, IOException, ClassNotFoundException {
    testManualCheckpoint(0);
  }

  @Test
  public void testManualCheckpoint()
      throws InterruptedException, IOException, ClassNotFoundException {
    testManualCheckpoint(2);
  }


  private void testManualCheckpoint(final int checkpointSuperstep)
      throws IOException, InterruptedException, ClassNotFoundException {
    Path checkpointsDir = getTempPath("checkpointing");
    GiraphConfiguration conf = new GiraphConfiguration();

    SUPERSTEP_CALLBACK = new SuperstepCallback() {

      @Override
      public void superstep(long superstep, ImmutableClassesGiraphConfiguration<LongWritable, IntWritable, FloatWritable> conf) {
        if (superstep == checkpointSuperstep) {
          try {
            ZooKeeperExt zooKeeperExt = new ZooKeeperExt(conf.getZookeeperList(),
                conf.getZooKeeperSessionTimeout(),
                conf.getZookeeperOpsMaxAttempts(),
                conf.getZookeeperOpsRetryWaitMsecs(),
                TestCheckpointing.this);
            String basePath = ZooKeeperManager.getBasePath(conf) + BspService.BASE_DIR + "/" + conf.get("mapred.job.id");
            zooKeeperExt.createExt(
                basePath + BspService.FORCE_CHECKPOINT_USER_FLAG,
                null,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT,
                true);
          } catch (IOException | InterruptedException | KeeperException e) {
            throw new RuntimeException(e);
          }
        } else if (superstep > checkpointSuperstep) {
          Assert.fail("Job should be stopped by now " + superstep);
        }
      }
    };

    try {
      runOriginalJob(checkpointsDir, conf);
      fail("Original job should fail after checkpointing");
    } catch (Exception e) {
      LOG.info("Original job failed, that's OK " + e);
    }

    SUPERSTEP_CALLBACK = new SuperstepCallback() {
      @Override
      public void superstep(long superstep,
                            ImmutableClassesGiraphConfiguration<LongWritable, IntWritable, FloatWritable> conf) {
        if (superstep < checkpointSuperstep) {
          Assert.fail("Restarted JOB should not be executed on superstep " + superstep);
        }
      }
    };

    runRestartedJob(checkpointsDir, conf, 10, -1);
  }

  /**
   * Worker context associated.
   */
  public static class CheckpointVertexWorkerContext
      extends DefaultWorkerContext {
    /** User can access this after the application finishes if local */
    private static long FINAL_SUM;

    private int testValue;

    public static long getFinalSum() {
      return FINAL_SUM;
    }

    @Override
    public void postSuperstep() {
      super.postSuperstep();
      sendMessageToMyself(new LongWritable(getSuperstep()));
    }

    /**
     * Send message to all workers (except this worker)
     *
     * @param message Message to send
     */
    private void sendMessageToMyself(Writable message) {
      sendMessageToWorker(message, getMyWorkerIndex());
    }

    @Override
    public void postApplication() {
      setFinalSum(this.<LongWritable>getAggregatedValue(
          LongSumAggregator.class.getName()).get());
      LOG.info("FINAL_SUM=" + FINAL_SUM);
    }

    /**
     * Set the final sum
     *
     * @param value sum
     */
    private static void setFinalSum(long value) {
      FINAL_SUM = value;
    }

    @Override
    public void preSuperstep() {
      assertEquals(getSuperstep(), testValue++);
      if (getSuperstep() > 0) {
        List<Writable> messages = getAndClearMessagesFromOtherWorkers();
        assertEquals(1, messages.size());
        assertEquals(getSuperstep() - 1, ((LongWritable)(messages.get(0))).get());
      }
    }

    @Override
    public void readFields(DataInput dataInput) throws IOException {
      super.readFields(dataInput);
      testValue = dataInput.readInt();
    }

    @Override
    public void write(DataOutput dataOutput) throws IOException {
      super.write(dataOutput);
      dataOutput.writeInt(testValue);
    }
  }

  /**
   * Master compute
   */
  public static class CheckpointVertexMasterCompute extends
      DefaultMasterCompute {

    private int testValue = 0;

    @Override
    public void compute() {
      long superstep = getSuperstep();
      if (SUPERSTEP_CALLBACK != null) {
        SUPERSTEP_CALLBACK.superstep(getSuperstep(), getConf());
      }
      assertEquals(superstep, testValue++);
    }

    @Override
    public void initialize() throws InstantiationException,
        IllegalAccessException {
      registerAggregator(LongSumAggregator.class.getName(),
          LongSumAggregator.class);
    }

    @Override
    public void readFields(DataInput in) throws IOException {
      super.readFields(in);
      testValue = in.readInt();
    }

    @Override
    public void write(DataOutput out) throws IOException {
      super.write(out);
      out.writeInt(testValue);
    }
  }

  private static interface SuperstepCallback {

    public void superstep(long superstep,
                          ImmutableClassesGiraphConfiguration<LongWritable,
                              IntWritable, FloatWritable> conf);

  }

}