From 4fb7fdb240b389775534d2d7a1b7092fe6366e79 Mon Sep 17 00:00:00 2001 From: Xin Yin Date: Sat, 4 Feb 2017 13:20:39 -0500 Subject: [PATCH] [jvm-packages] Fixed java.nio.BufferUnderFlow issue in Scala Rabit tracker. (#1993) * [jvm-packages] Scala implementation of the Rabit tracker. A Scala implementation of RabitTracker that is interface-interchangable with the Java implementation, ported from `tracker.py` in the [dmlc-core project](https://github.com/dmlc/dmlc-core). * [jvm-packages] Updated Akka dependency in pom.xml. * Refactored the RabitTracker directory structure. * Fixed premature stopping of connection handler. Added a new finite state "AwaitingPortNumber" to explicitly wait for the worker to send the port, and close the connection. Stopping the actor prematurely sends a TCP RST to the worker, causing the worker to crash on AssertionError. * Added interface IRabitTracker so that user can switch implementations. * Default timeout duration changes. * Dependency for Akka tests. * Removed the main function of RabitTracker. * A skeleton for testing Akka-based Rabit tracker. * waitFor() in RabitTracker no longer throws exceptions. * Completed unit test for the 'start' command of Rabit tracker. * Preliminary support for Rabit Allreduce via JNI (no prepare function support yet.) * Fixed the default timeout duration. * Use Java container to avoid serialization issues due to intermediate wrappers. * Added tests for Allreduce/model training using Scala Rabit tracker. * Added spill-over unit test for the Scala Rabit tracker. * Fixed a typo. * Overhaul of RabitTracker interface per code review. - Removed methods start() waitFor() (no arguments) from IRabitTracker. - The timeout in start(timeout) is now worker connection timeout, as tcp socket binding timeout is less intuitive. - Dropped time unit from start(...) and waitFor(...) methods; the default time unit is millisecond. - Moved random port number generation into the RabitTrackerHandler. - Moved all Rabit-related classes to package ml.dmlc.xgboost4j.scala.rabit. * More code refactoring and comments. * Unified timeout constants. Readable tracker status code. * Add comments to indicate that allReduce is for tests only. Removed all other variants. * Removed unused imports. * Simplified signatures of training methods. - Moved TrackerConf into parameter map. - Changed GeneralParams so that TrackerConf becomes a standalone parameter. - Updated test cases accordingly. * Changed monitoring strategies. * Reverted monitoring changes. * Update test case for Rabit AllReduce. * Mix in UncaughtExceptionHandler into IRabitTracker to prevent tracker from hanging due to exceptions thrown by workers. * More comprehensive test cases for exception handling and worker connection timeout. * Handle executor loss due to unknown cause: the newly spawned executor will attempt to connect to the tracker. Interrupt tracker in such case. * Per code-review, removed training timeout from TrackerConf. Timeout logic must be implemented explicitly and externally in the driver code. * Reverted scalastyle-config changes. * Visibility scope change. Interface tweaks. * Use match pattern to handle tracker_conf parameter. * Minor clarification in JNI code. * Clearer intent in match pattern to suppress warnings. * Removed Future from constructor. Block in start() and waitFor() instead. * Revert inadvertent comment changes. * Removed debugging information. * Updated test cases that are a bit finicky. * Added comments on the reasoning behind the unit tests for testing Rabit tracker robustness. * Fixed BufferUnderFlow bug in decoding tracker 'print' command. * Merge conflicts resolution. --- .../rabit/handler/RabitWorkerHandler.scala | 29 +++++++++++++------ .../RabitTrackerConnectionHandlerTest.scala | 26 +++++++++++++++++ 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/rabit/handler/RabitWorkerHandler.scala b/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/rabit/handler/RabitWorkerHandler.scala index 31acfc1ce..963799884 100644 --- a/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/rabit/handler/RabitWorkerHandler.scala +++ b/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/rabit/handler/RabitWorkerHandler.scala @@ -84,12 +84,15 @@ private[scala] class RabitWorkerHandler(host: String, worldSize: Int, tracker: A def getNeighboringWorkers: Set[Int] = neighboringWorkers def decodeCommand(buffer: ByteBuffer): TrackerCommand = { - val rank = buffer.getInt() - val worldSize = buffer.getInt() - val jobId = buffer.getString + val readBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) + readBuffer.flip() - val command = buffer.getString - command match { + val rank = readBuffer.getInt() + val worldSize = readBuffer.getInt() + val jobId = readBuffer.getString + + val command = readBuffer.getString + val trackerCommand = command match { case "start" => WorkerStart(rank, worldSize, jobId) case "shutdown" => transient = true @@ -99,8 +102,11 @@ private[scala] class RabitWorkerHandler(host: String, worldSize: Int, tracker: A WorkerRecover(rank, worldSize, jobId) case "print" => transient = true - WorkerTrackerPrint(rank, worldSize, jobId, buffer.getString) + WorkerTrackerPrint(rank, worldSize, jobId, readBuffer.getString) } + + stashSpillOver(readBuffer) + trackerCommand } startWith(AwaitingHandshake, DataStruct()) @@ -120,9 +126,14 @@ private[scala] class RabitWorkerHandler(host: String, worldSize: Int, tracker: A case Event(Tcp.Received(bytes), validator) => bytes.asByteBuffers.foreach { buf => readBuffer.put(buf) } if (validator.verify(readBuffer)) { - readBuffer.flip() - tracker ! decodeCommand(readBuffer) - stashSpillOver(readBuffer) + Try(decodeCommand(readBuffer)) match { + case scala.util.Success(decodedCommand) => + tracker ! decodedCommand + case scala.util.Failure(th: java.nio.BufferOverflowException) => + // BufferOverflowException would occur if the message to print has not arrived yet. + // Do nothing, wait for next Tcp.Received event + case scala.util.Failure(th: Throwable) => throw th + } } stay diff --git a/jvm-packages/xgboost4j/src/test/scala/ml/dmlc/xgboost4j/scala/rabit/RabitTrackerConnectionHandlerTest.scala b/jvm-packages/xgboost4j/src/test/scala/ml/dmlc/xgboost4j/scala/rabit/RabitTrackerConnectionHandlerTest.scala index ee4febe39..42994baca 100644 --- a/jvm-packages/xgboost4j/src/test/scala/ml/dmlc/xgboost4j/scala/rabit/RabitTrackerConnectionHandlerTest.scala +++ b/jvm-packages/xgboost4j/src/test/scala/ml/dmlc/xgboost4j/scala/rabit/RabitTrackerConnectionHandlerTest.scala @@ -172,6 +172,32 @@ class RabitTrackerConnectionHandlerTest trackerProbe.expectMsg(printCmd) } + it should "handle fragmented print command without throwing exception" in { + val trackerProbe = TestProbe() + val connProbe = TestProbe() + + val fsm = TestFSMRef(new RabitWorkerHandler("localhost", 4, + trackerProbe.ref, connProbe.ref)) + fsm.stateName shouldEqual RabitWorkerHandler.AwaitingHandshake + + fsm ! Tcp.Received(magic) + connProbe.expectMsg(Tcp.Write(magic)) + + fsm.stateName shouldEqual RabitWorkerHandler.AwaitingCommand + fsm.stateData shouldEqual RabitWorkerHandler.StructTrackerCommand + // ResumeReading should be seen once state transitions + connProbe.expectMsg(Tcp.ResumeReading) + + val printCmd = WorkerTrackerPrint(0, 4, "print", "hello world!") + // 4 + 4 + 4 + 5 = 17 + val (partialMessage, remainder) = printCmd.encode.splitAt(17) + + fsm ! Tcp.Received(partialMessage) + fsm ! Tcp.Received(remainder) + + trackerProbe.expectMsg(printCmd) + } + it should "handle spill-over Tcp data correctly between state transition" in { val trackerProbe = TestProbe() val connProbe = TestProbe()