Skip to content

[SPARK-5484][GraphX] Periodically do checkpoint in Pregel #15125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dding3 wants to merge 20 commits into from
Closed

[SPARK-5484][GraphX] Periodically do checkpoint in Pregel #15125

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
3834981
test
Sep 16, 2016
166fd6d
period do checkpoint in pregel
Sep 16, 2016
b119e4a
remove unused code
Sep 16, 2016
d183a7c
fix mima
Sep 16, 2016
352dcb2
fix mima
Sep 16, 2016
ad82e45
fix scala style fail
Sep 16, 2016
e786838
rebase upstream
Jan 17, 2017
a25d00c
code rebase
Feb 11, 2017
38e6238
Update a couple of test suite package names, and enhance the
Feb 14, 2017
f2efef6
Merge pull request #1 from VideoAmp/spark-5484-periodic_pregel_checkp…
dding3 Feb 16, 2017
194dc27
update checkpoint interval and some comments
Feb 17, 2017
9d7e796
add config value to controll checkpoint interval in pergel
Feb 17, 2017
dae94aa
document spark.graphx.pregel.checkpointInterval into graphx-programmi…
Feb 19, 2017
dd6c366
remove details of checkpointer in pregel implementation sketch
Feb 19, 2017
2639eb1
remove unused variable in pregel sketch implementation
Feb 19, 2017
11bc349
small changes for documentation
Feb 23, 2017
9a6fd1f
mv PeriodicRDDCheckpointer to rdd.util and explain why use cache in P…
Apr 19, 2017
5015b44
fix scalastyle error
Apr 20, 2017
24d4ad6
turn off checkpoint graph and messages in pregel by default
Apr 22, 2017
ec62659
recommend a good value in checkpoint interval in pregel
Apr 23, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions ...ark/mllib/impl/PeriodicCheckpointer.scala → ...che/spark/util/PeriodicCheckpointer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* limitations under the License.
*/

package org.apache.spark.mllib.impl
package org.apache.spark.util

import scala.collection.mutable

Expand Down Expand Up @@ -58,7 +58,7 @@ import org.apache.spark.storage.StorageLevel
* @param sc SparkContext for the Datasets given to this checkpointer
* @tparam T Dataset type, such as RDD[Double]
*/
private[mllib] abstract class PeriodicCheckpointer[T](
private[spark] abstract class PeriodicCheckpointer[T](
val checkpointInterval: Int,
val sc: SparkContext) extends Logging {

Expand Down Expand Up @@ -127,6 +127,16 @@ private[mllib] abstract class PeriodicCheckpointer[T](
/** Get list of checkpoint files for this given Dataset */
protected def getCheckpointFiles(data: T): Iterable[String]

/**
* Call this to unpersist the Dataset.
*/
def unpersistDataSet(): Unit = {
while (persistedQueue.nonEmpty) {

@felixcheung felixcheung Apr 19, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if it's the goal but this isn't thread-safe?

@mallman mallman Apr 19, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would agree with you, that this is not thread safe. Is that a concern?

@felixcheung felixcheung Apr 19, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

with the limited internal only use, it should be ok

val dataToUnpersist = persistedQueue.dequeue()
unpersist(dataToUnpersist)
}
}

/**
* Call this at the end to delete any remaining checkpoint files.
*/
Expand Down
2 changes: 1 addition & 1 deletion .../mllib/impl/PeriodicRDDCheckpointer.scala → .../spark/util/PeriodicRDDCheckpointer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* limitations under the License.
*/

package org.apache.spark.mllib.impl
package org.apache.spark.util

@felixcheung felixcheung Apr 19, 2017 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's PeriodicRDDCheckpointer, shouldn't this be in the org.apache.spark.rdd.util namespace?


import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD
Expand Down
8 changes: 3 additions & 5 deletions ...b/impl/PeriodicRDDCheckpointerSuite.scala → ...k/util/PeriodicRDDCheckpointerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,16 @@
* limitations under the License.
*/

package org.apache.spark.mllib.impl
package org.apache.spark.util

import org.apache.hadoop.fs.Path

import org.apache.spark.{SparkContext, SparkFunSuite}
import org.apache.spark.mllib.util.MLlibTestSparkContext
import org.apache.spark.{SharedSparkContext, SparkContext, SparkFunSuite}
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.Utils


class PeriodicRDDCheckpointerSuite extends SparkFunSuite with MLlibTestSparkContext {
class PeriodicRDDCheckpointerSuite extends SparkFunSuite with SharedSparkContext {

import PeriodicRDDCheckpointerSuite._

Expand Down
50 changes: 40 additions & 10 deletions docs/graphx-programming-guide.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -708,7 +708,9 @@ messages remaining.
> messaging function. These constraints allow additional optimization within GraphX.

The following is the type signature of the [Pregel operator][GraphOps.pregel] as well as a *sketch*
of its implementation (note calls to graph.cache have been removed):
of its implementation (note: to avoid stackOverflowError due to long lineage chains, graph and

@mallman mallman Feb 20, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think @viirya is suggesting we remove all references to checkpointing from this implementation sketch. @viirya, correct me if I'm wrong. I suggest reverting these changes to graphx-programming-guide.md. Just add a brief note that GraphX periodically checkpoints the graph and message lineages, and this checkpoint interval can be configured with the spark.graphx.pregel.checkpointInterval Spark configuration property.

@viirya @felixcheung Should we document this config property in the Spark configuration document as well?

@mallman mallman Feb 20, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, we should document the default value and that checkpointing can be disabled by setting this config property to -1.

@dding3 dding3 Feb 20, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. I am thinking in the original graphx-programming-guide, only cache is called. After we add do checkpointing and give notes we do that in the sketch, do we need to show the information in the sketch. If there is an agreement on remove all references to checkpoint, I will revert the changes.

@felixcheung felixcheung Feb 21, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it would help reviewers to understand if you could explain those changes in the impl sketch is required to use checkpoint

I'd agree we should probably document the config property

@mallman mallman Feb 21, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it would help reviewers to understand if you could explain those changes in the impl sketch is required to use checkpoint

Sorry @felixcheung, I don't understand what you mean here. Can you elaborate?

@dding3 dding3 Feb 21, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

About document this config property in the Spark configuration document. Is it OK if I add a new Graphx section to include the config, or just add the config under existing section, say "execution behavor"?

@felixcheung felixcheung Feb 24, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I mean if there is a reason to have graphCheckpointer in the example, if you could explain it, it would make it easier for us to review those.

@mallman mallman Feb 28, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Personally, I don't think we need to include the checkpointer in the implementation sketch. I think it's more an implementation detail than an essential part of the algorithm. I think it's enough to simply document that GraphX's implementation of pregel includes checkpointing to avoid unbounded RDD lineages.

@dding3 dding3 Feb 28, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. Have removed references to checkpoint in the sketch and documented the config property in the spark configuration document in the latest update.

messages are periodically checkpoint and the checkpoint interval is set by
"spark.graphx.pregel.checkpointInterval"):

{% highlight scala %}
class GraphOps[VD, ED] {
Expand All @@ -720,25 +722,53 @@ class GraphOps[VD, ED] {
sendMsg: EdgeTriplet[VD, ED] => Iterator[(VertexId, A)],
mergeMsg: (A, A) => A)
: Graph[VD, ED] = {
// Receive the initial message at each vertex
var g = mapVertices( (vid, vdata) => vprog(vid, vdata, initialMsg) ).cache()
val checkpointInterval = graph.vertices.sparkContext.getConf

@viirya viirya Feb 20, 2017 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm, as this is just an implementation sketch, I don't think we should include such details of checkpointer.

@dding3 dding3 Feb 20, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. I will change back then.

.getInt("spark.graphx.pregel.checkpointInterval", 10)
var g = graph.mapVertices((vid, vdata) => vprog(vid, vdata, initialMsg))
val graphCheckpointer = new PeriodicGraphCheckpointer[VD, ED](
checkpointInterval, graph.vertices.sparkContext)
graphCheckpointer.update(g)

// compute the messages
var messages = g.mapReduceTriplets(sendMsg, mergeMsg)
var messages = GraphXUtils.mapReduceTriplets(g, sendMsg, mergeMsg)
val messageCheckpointer = new PeriodicRDDCheckpointer[(VertexId, A)](
checkpointInterval, graph.vertices.sparkContext)
messageCheckpointer.update(messages.asInstanceOf[RDD[(VertexId, A)]])
var activeMessages = messages.count()
// Loop until no messages remain or maxIterations is achieved

// Loop
var prevG: Graph[VD, ED] = null
var i = 0
while (activeMessages > 0 && i < maxIterations) {
// Receive the messages and update the vertices.
g = g.joinVertices(messages)(vprog).cache()
prevG = g
g = g.joinVertices(messages)(vprog)
graphCheckpointer.update(g)

val oldMessages = messages
// Send new messages, skipping edges where neither side received a message. We must cache
// messages so it can be materialized on the next line, allowing us to uncache the previous
// iteration.
messages = g.mapReduceTriplets(
sendMsg, mergeMsg, Some((oldMessages, activeDirection))).cache()
// and periodic checkpoint messages so it can be materialized on the next line, and avoid
// to have a long lineage chain.
messages = GraphXUtils.mapReduceTriplets(
g, sendMsg, mergeMsg, Some((oldMessages, activeDirection)))
// The call to count() materializes `messages` and the vertices of `g`. This hides oldMessages
// (depended on by the vertices of g) and the vertices of prevG (depended on by oldMessages
// and the vertices of g).
messageCheckpointer.update(messages.asInstanceOf[RDD[(VertexId, A)]])
activeMessages = messages.count()

logInfo("Pregel finished iteration " + i)

// Unpersist the RDDs hidden by newly-materialized RDDs
oldMessages.unpersist(blocking = false)
prevG.unpersistVertices(blocking = false)
prevG.edges.unpersist(blocking = false)
// count the iteration
i += 1
}
messageCheckpointer.unpersistDataSet()
graphCheckpointer.deleteAllCheckpoints()
messageCheckpointer.deleteAllCheckpoints()
g
}
}
Expand Down
29 changes: 23 additions & 6 deletions graphx/src/main/scala/org/apache/spark/graphx/Pregel.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,10 @@ package org.apache.spark.graphx

import scala.reflect.ClassTag

import org.apache.spark.graphx.util.PeriodicGraphCheckpointer
import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
import org.apache.spark.util.PeriodicRDDCheckpointer

/**
* Implements a Pregel-like bulk-synchronous message-passing API.
Expand Down Expand Up @@ -122,27 +125,39 @@ object Pregel extends Logging {
require(maxIterations > 0, s"Maximum number of iterations must be greater than 0," +
s" but got ${maxIterations}")

var g = graph.mapVertices((vid, vdata) => vprog(vid, vdata, initialMsg)).cache()
val checkpointInterval = graph.vertices.sparkContext.getConf
.getInt("spark.graphx.pregel.checkpointInterval", 10)

@viirya viirya Feb 18, 2017 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Do we need to document this config into GraphX related document?

@dding3 dding3 Feb 20, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think so. Currently I add the document in graphx-programming-guide.md. But I am not sure if it's the right place, please let me know if there is a better place to add the document.

@mallman mallman Feb 21, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would also suggest incorporating this change into the Spark 2.2 release notes under a section for GraphX, but I don't see where these notes are maintained. The release notes for 2.1 are published at http://spark.apache.org/releases/spark-release-2-1-0.html, but I can't find them in the repo. Anybody know how these are generated or how to contribute to them? Is there another repo for this documentation?

@felixcheung felixcheung Apr 19, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mallman it's at https://github.com/apache/spark-website/tree/asf-site/releases/_posts

var g = graph.mapVertices((vid, vdata) => vprog(vid, vdata, initialMsg))
val graphCheckpointer = new PeriodicGraphCheckpointer[VD, ED](
checkpointInterval, graph.vertices.sparkContext)
graphCheckpointer.update(g)

// compute the messages
var messages = GraphXUtils.mapReduceTriplets(g, sendMsg, mergeMsg)
val messageCheckpointer = new PeriodicRDDCheckpointer[(VertexId, A)](
checkpointInterval, graph.vertices.sparkContext)
messageCheckpointer.update(messages.asInstanceOf[RDD[(VertexId, A)]])

@viirya viirya Feb 10, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to update it before the loop? I think it should be enough to do the update in the loop.

@mallman mallman Feb 10, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree. What do you think, @dding3?

@mallman mallman Feb 10, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, for the sake of simplicity and consistency, I'm going to suggest we keep the checkpointer update calls but remove all .cache() calls. The update calls persist the underlying data, making the calls to .cache() unnecessary.

@dding3 dding3 Feb 12, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we need cache graph/messages here so they don't need to be computed again in the loop. I agree with you and I will keep the checkpointer update calls and remove all .cache calls.

var activeMessages = messages.count()

// Loop
var prevG: Graph[VD, ED] = null
var i = 0
while (activeMessages > 0 && i < maxIterations) {
// Receive the messages and update the vertices.
prevG = g
g = g.joinVertices(messages)(vprog).cache()
g = g.joinVertices(messages)(vprog)
graphCheckpointer.update(g)

val oldMessages = messages
// Send new messages, skipping edges where neither side received a message. We must cache
// messages so it can be materialized on the next line, allowing us to uncache the previous
// iteration.
// and periodic checkpoint messages so it can be materialized on the next line, and avoid
// to have a long lineage chain.
messages = GraphXUtils.mapReduceTriplets(
g, sendMsg, mergeMsg, Some((oldMessages, activeDirection))).cache()
g, sendMsg, mergeMsg, Some((oldMessages, activeDirection)))
// The call to count() materializes `messages` and the vertices of `g`. This hides oldMessages

@felixcheung felixcheung Feb 17, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should the comment here be updated?

// (depended on by the vertices of g) and the vertices of prevG (depended on by oldMessages
// and the vertices of g).
messageCheckpointer.update(messages.asInstanceOf[RDD[(VertexId, A)]])
activeMessages = messages.count()

logInfo("Pregel finished iteration " + i)
Expand All @@ -154,7 +169,9 @@ object Pregel extends Logging {
// count the iteration
i += 1
}
messages.unpersist(blocking = false)
messageCheckpointer.unpersistDataSet()

@mallman mallman Feb 20, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I don't understand this change. Why do we replace

messages.unpersist(blocking = false)

with

messageCheckpointer.unpersistDataSet()

Especially because this adds a new public method to PeriodicCheckpointer that no other code has needed before.

@dding3 dding3 Feb 20, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the thing is we use messageCheckpointer.update to do the cache, to make a pair, we can use it to unpersist data. Please correct me if I understand wrong.
I think it's fine to add this new method as there is already a public method to cache data in PersistQueue, we should provide a public method to clean the queue.

graphCheckpointer.deleteAllCheckpoints()
messageCheckpointer.deleteAllCheckpoints()

@felixcheung felixcheung Feb 17, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't this be inside a finally clause to make sure checkpoint data is cleaned up even in error cases?

@dding3 dding3 Feb 17, 2017

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think when there is an exception during training, if we keep the checkpoints, there is a chance for user to recover from it. I checked in RandomForest/GBT in spark, looks like they only delete the checkpoints when the training successful finished.

g
} // end of apply

Expand Down
10 changes: 5 additions & 5 deletions ...llib/impl/PeriodicGraphCheckpointer.scala → ...aphx/util/PeriodicGraphCheckpointer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,12 @@
* limitations under the License.
*/

package org.apache.spark.mllib.impl
package org.apache.spark.graphx.util

import org.apache.spark.SparkContext
import org.apache.spark.graphx.Graph
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.PeriodicCheckpointer


/**
Expand Down Expand Up @@ -74,9 +75,8 @@ import org.apache.spark.storage.StorageLevel
* @tparam VD Vertex descriptor type
* @tparam ED Edge descriptor type
*
* TODO: Move this out of MLlib?
*/
private[mllib] class PeriodicGraphCheckpointer[VD, ED](
private[spark] class PeriodicGraphCheckpointer[VD, ED](
checkpointInterval: Int,
sc: SparkContext)
extends PeriodicCheckpointer[Graph[VD, ED]](checkpointInterval, sc) {
Expand All @@ -87,10 +87,10 @@ private[mllib] class PeriodicGraphCheckpointer[VD, ED](

override protected def persist(data: Graph[VD, ED]): Unit = {
if (data.vertices.getStorageLevel == StorageLevel.NONE) {
data.vertices.persist()
data.vertices.cache()

@felixcheung felixcheung Feb 17, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't persist better? this could potentially support different storage level later

@mallman mallman Feb 17, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to use cache because persist does not honor the default storage level requested when constructing the graph. Only cache does that. It's confusing, but true. To verify this for yourself, change these values to persist and run the PeriodicGraphCheckpointerSuite tests.

@felixcheung felixcheung Apr 19, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mallman @dding3 it will be good to add a comment on that here

}
if (data.edges.getStorageLevel == StorageLevel.NONE) {
data.edges.persist()
data.edges.cache()
}
}

Expand Down
105 changes: 55 additions & 50 deletions ...impl/PeriodicGraphCheckpointerSuite.scala → ...util/PeriodicGraphCheckpointerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,77 +15,81 @@
* limitations under the License.
*/

package org.apache.spark.mllib.impl
package org.apache.spark.graphx.util

import org.apache.hadoop.fs.Path

@mallman mallman Feb 10, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a reason this test suite isn't moved into the GraphX codebase?

import org.apache.spark.{SparkContext, SparkFunSuite}
import org.apache.spark.graphx.{Edge, Graph}
import org.apache.spark.mllib.util.MLlibTestSparkContext
import org.apache.spark.graphx.{Edge, Graph, LocalSparkContext}
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.Utils


class PeriodicGraphCheckpointerSuite extends SparkFunSuite with MLlibTestSparkContext {
class PeriodicGraphCheckpointerSuite extends SparkFunSuite with LocalSparkContext {

import PeriodicGraphCheckpointerSuite._

test("Persisting") {
var graphsToCheck = Seq.empty[GraphToCheck]

val graph1 = createGraph(sc)
val checkpointer =
new PeriodicGraphCheckpointer[Double, Double](10, graph1.vertices.sparkContext)
checkpointer.update(graph1)
graphsToCheck = graphsToCheck :+ GraphToCheck(graph1, 1)
checkPersistence(graphsToCheck, 1)

var iteration = 2
while (iteration < 9) {
val graph = createGraph(sc)
checkpointer.update(graph)
graphsToCheck = graphsToCheck :+ GraphToCheck(graph, iteration)
checkPersistence(graphsToCheck, iteration)
iteration += 1
withSpark { sc =>
val graph1 = createGraph(sc)
val checkpointer =
new PeriodicGraphCheckpointer[Double, Double](10, graph1.vertices.sparkContext)
checkpointer.update(graph1)
graphsToCheck = graphsToCheck :+ GraphToCheck(graph1, 1)
checkPersistence(graphsToCheck, 1)

var iteration = 2
while (iteration < 9) {
val graph = createGraph(sc)
checkpointer.update(graph)
graphsToCheck = graphsToCheck :+ GraphToCheck(graph, iteration)
checkPersistence(graphsToCheck, iteration)
iteration += 1
}
}
}

test("Checkpointing") {
val tempDir = Utils.createTempDir()
val path = tempDir.toURI.toString
val checkpointInterval = 2
var graphsToCheck = Seq.empty[GraphToCheck]
sc.setCheckpointDir(path)
val graph1 = createGraph(sc)
val checkpointer = new PeriodicGraphCheckpointer[Double, Double](
checkpointInterval, graph1.vertices.sparkContext)
checkpointer.update(graph1)
graph1.edges.count()
graph1.vertices.count()
graphsToCheck = graphsToCheck :+ GraphToCheck(graph1, 1)
checkCheckpoint(graphsToCheck, 1, checkpointInterval)

var iteration = 2
while (iteration < 9) {
val graph = createGraph(sc)
checkpointer.update(graph)
graph.vertices.count()
graph.edges.count()
graphsToCheck = graphsToCheck :+ GraphToCheck(graph, iteration)
checkCheckpoint(graphsToCheck, iteration, checkpointInterval)
iteration += 1
}
withSpark { sc =>
val tempDir = Utils.createTempDir()
val path = tempDir.toURI.toString
val checkpointInterval = 2
var graphsToCheck = Seq.empty[GraphToCheck]
sc.setCheckpointDir(path)
val graph1 = createGraph(sc)
val checkpointer = new PeriodicGraphCheckpointer[Double, Double](
checkpointInterval, graph1.vertices.sparkContext)
checkpointer.update(graph1)
graph1.edges.count()
graph1.vertices.count()
graphsToCheck = graphsToCheck :+ GraphToCheck(graph1, 1)
checkCheckpoint(graphsToCheck, 1, checkpointInterval)

var iteration = 2
while (iteration < 9) {
val graph = createGraph(sc)
checkpointer.update(graph)
graph.vertices.count()
graph.edges.count()
graphsToCheck = graphsToCheck :+ GraphToCheck(graph, iteration)
checkCheckpoint(graphsToCheck, iteration, checkpointInterval)
iteration += 1
}

checkpointer.deleteAllCheckpoints()
graphsToCheck.foreach { graph =>
confirmCheckpointRemoved(graph.graph)
}
checkpointer.deleteAllCheckpoints()
graphsToCheck.foreach { graph =>
confirmCheckpointRemoved(graph.graph)
}

Utils.deleteRecursively(tempDir)
Utils.deleteRecursively(tempDir)
}
}
}

private object PeriodicGraphCheckpointerSuite {
private val defaultStorageLevel = StorageLevel.MEMORY_ONLY_SER

case class GraphToCheck(graph: Graph[Double, Double], gIndex: Int)

Expand All @@ -96,7 +100,8 @@ private object PeriodicGraphCheckpointerSuite {
Edge[Double](3, 4, 0))

def createGraph(sc: SparkContext): Graph[Double, Double] = {
Graph.fromEdges[Double, Double](sc.parallelize(edges), 0)
Graph.fromEdges[Double, Double](
sc.parallelize(edges), 0, defaultStorageLevel, defaultStorageLevel)
}

def checkPersistence(graphs: Seq[GraphToCheck], iteration: Int): Unit = {
Expand All @@ -116,8 +121,8 @@ private object PeriodicGraphCheckpointerSuite {
assert(graph.vertices.getStorageLevel == StorageLevel.NONE)
assert(graph.edges.getStorageLevel == StorageLevel.NONE)
} else {
assert(graph.vertices.getStorageLevel != StorageLevel.NONE)
assert(graph.edges.getStorageLevel != StorageLevel.NONE)
assert(graph.vertices.getStorageLevel == defaultStorageLevel)
assert(graph.edges.getStorageLevel == defaultStorageLevel)
}
} catch {
case _: AssertionError =>
Expand Down
Loading