From 18cd85626ab4a8872c713a3172c58fd854cb74e9 Mon Sep 17 00:00:00 2001 From: Xianming Lei Date: Thu, 10 Sep 2026 19:08:38 +0800 Subject: [PATCH] [CORE] Fix remote fetch failure for empty deserialized RDD blocks A deserialized in-memory block's reported memory size is an estimated JVM object size, not the size of its serialized response. Its serialized response may therefore be empty. Track the location that owns the status returned by the master. Accept an empty response only when it comes from that exact location and the status reports a deserialized in-memory block. Continue treating empty responses from the external shuffle service or any other location as fetch failures, so the existing retry and failure detection behavior is preserved. Add regression coverage for an allowed empty executor response, a rejected external-shuffle-service response, and a response whose location does not match the status-owning location. Generated-by: OpenAI Codex --- .../apache/spark/storage/BlockManager.scala | 13 ++- .../storage/BlockManagerMasterEndpoint.scala | 3 +- .../spark/storage/BlockManagerMessages.scala | 2 + .../spark/storage/BlockManagerSuite.scala | 103 +++++++++++++++++- 4 files changed, 117 insertions(+), 4 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala index 7267348031ed5..c30e1f404eaaf 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala @@ -1408,6 +1408,12 @@ private[spark] class BlockManager( var totalFailureCount = 0 val locations = sortLocations(locationsAndStatus.locations) val maxFetchFailures = locations.size + // For deserialized blocks, memSize is the estimated size of the JVM objects rather than + // the size of their serialized representation. An empty iterator can occupy memory while + // serializing to an empty buffer, so this size cannot be used to reject an empty response. + val isDeserializedMemoryBlock = + locationsAndStatus.status.storageLevel.useMemory && + locationsAndStatus.status.storageLevel.deserialized var locationIterator = locations.iterator while (locationIterator.hasNext) { val loc = locationIterator.next() @@ -1415,7 +1421,12 @@ private[spark] class BlockManager( val data = try { val buf = blockTransferService.fetchBlockSync(loc.host, loc.port, loc.executorId, blockId.toString, tempFileManager) - if (blockSize > 0 && buf.size() == 0) { + // A different location can have a different storage level, including the external shuffle + // service which only serves disk files. Accept an empty response only from the location + // whose status reports a deserialized memory block. + val canServeEmptyDeserializedMemoryBlock = + isDeserializedMemoryBlock && loc == locationsAndStatus.statusLocation + if (!canServeEmptyDeserializedMemoryBlock && blockSize > 0 && buf.size() == 0) { throw SparkException.internalError("Empty buffer received for non empty block " + s"when fetching remote block $blockId from $loc", category = "STORAGE") } diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala index beca0d11d1c5e..7719ff450ceda 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala @@ -1056,9 +1056,10 @@ class BlockManagerMasterEndpoint( BlockLocationsAndStatus( allLocations, blockStatus, + bmId, Option(executorIdToLocalDirs.getIfPresent(bmId.executorId))) } else { - BlockLocationsAndStatus(allLocations, blockStatus, None) + BlockLocationsAndStatus(allLocations, blockStatus, bmId, None) } } .orElse(None) diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala index 3c2ee19ffbd74..a7d60a580fb7e 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala @@ -133,6 +133,7 @@ private[spark] object BlockManagerMessages { /** * The response message of `GetLocationsAndStatus` request. * + * @param statusLocation the block manager location which owns `status`. * @param localDirs if it is persisted-to-disk on the same host as the requester executor is * running on then localDirs will be Some and the cached data will be in a file * in one of those dirs, otherwise it is None. @@ -140,6 +141,7 @@ private[spark] object BlockManagerMessages { case class BlockLocationsAndStatus( locations: Seq[BlockManagerId], status: BlockStatus, + statusLocation: BlockManagerId, localDirs: Option[Array[String]]) { assert(locations.nonEmpty) } diff --git a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala index 70501afd9d921..0ab863f3cd4f2 100644 --- a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala @@ -1832,7 +1832,8 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe val blockManagerIds = (0 to maxFailuresBeforeLocationRefresh) .map { i => BlockManagerId(s"id-$i", s"host-$i", i + 1) } when(mockBlockManagerMaster.getLocationsAndStatus(mc.any[BlockId], mc.any[String])).thenReturn( - Option(BlockLocationsAndStatus(blockManagerIds, BlockStatus.empty, None))) + Option(BlockLocationsAndStatus( + blockManagerIds, BlockStatus.empty, blockManagerIds.head, None))) when(mockBlockManagerMaster.getLocations(mc.any[BlockId])).thenReturn( blockManagerIds) @@ -1905,6 +1906,104 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe assert(store2.getRemoteBytes("item").isEmpty) } + test("allow empty executor response for deserialized memory block") { + val mockBlockManagerMaster = mock(classOf[BlockManagerMaster]) + val executor = BlockManagerId("executor1", "localhost", 8000) + val deserializedMemoryStatus = + BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L) + when(mockBlockManagerMaster.getLocationsAndStatus( + mc.any[BlockId], mc.any[String])).thenReturn( + Option(BlockLocationsAndStatus( + Seq(executor), deserializedMemoryStatus, executor, None))) + val blockFetcher = new MockBlockTransferService(0) { + override def fetchBlockSync( + host: String, + port: Int, + execId: String, + blockId: String, + tempFileManager: DownloadFileManager): ManagedBuffer = { + new NioManagedBuffer(ByteBuffer.allocate(0)) + } + } + val store = makeBlockManager( + 8000, + "executor2", + mockBlockManagerMaster, + transferService = Some(blockFetcher)) + + val remoteBytes = store.getRemoteBytes(rdd(0, 0)) + assert(remoteBytes.isDefined) + assert(remoteBytes.get.size === 0) + } + + test("reject empty shuffle-service response for deserialized memory block") { + val mockBlockManagerMaster = mock(classOf[BlockManagerMaster]) + val externalShuffleServicePort = StorageUtils.externalShuffleServicePort(conf) + val executor = BlockManagerId("executor1", "localhost", 8000) + val shuffleService = BlockManagerId("executor2", "localhost", externalShuffleServicePort) + val deserializedMemoryStatus = + BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L) + when(mockBlockManagerMaster.getLocationsAndStatus( + mc.any[BlockId], mc.any[String])).thenReturn( + Option(BlockLocationsAndStatus( + Seq(executor, shuffleService), deserializedMemoryStatus, executor, None))) + val blockFetcher = new MockBlockTransferService(0) { + override def fetchBlockSync( + host: String, + port: Int, + execId: String, + blockId: String, + tempFileManager: DownloadFileManager): ManagedBuffer = { + if (port == externalShuffleServicePort) { + val transConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numUsableCores = 1) + new FileSegmentManagedBuffer(transConf, new File("missing.file"), 0, 0) + } else { + throw new RuntimeException("simulated executor fetch failure") + } + } + } + val store = makeBlockManager( + 8000, + "executor3", + mockBlockManagerMaster, + transferService = Some(blockFetcher)) + + assert(store.getRemoteBytes("item").isEmpty) + } + + test("reject empty executor response when status belongs to another location") { + val mockBlockManagerMaster = mock(classOf[BlockManagerMaster]) + val staleExecutor = BlockManagerId("executor1", "MockBlockTransferServiceHost", 8000) + val statusExecutor = BlockManagerId("executor2", "other-host", 8000) + val deserializedMemoryStatus = + BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L) + when(mockBlockManagerMaster.getLocationsAndStatus( + mc.any[BlockId], mc.any[String])).thenReturn( + Option(BlockLocationsAndStatus( + Seq(staleExecutor, statusExecutor), deserializedMemoryStatus, statusExecutor, None))) + val blockFetcher = new MockBlockTransferService(0) { + override def fetchBlockSync( + host: String, + port: Int, + execId: String, + blockId: String, + tempFileManager: DownloadFileManager): ManagedBuffer = { + if (execId == staleExecutor.executorId) { + new NioManagedBuffer(ByteBuffer.allocate(0)) + } else { + throw new RuntimeException("simulated status executor fetch failure") + } + } + } + val store = makeBlockManager( + 8000, + "executor3", + mockBlockManagerMaster, + transferService = Some(blockFetcher)) + + assert(store.getRemoteBytes("item").isEmpty) + } + test("test sorting of block locations") { val localHost = "localhost" val otherHost = "otherHost" @@ -2022,7 +2121,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe val blockStatus = BlockStatus(StorageLevel.DISK_ONLY, 0L, 2000L) when(mockBlockManagerMaster.getLocationsAndStatus(mc.any[BlockId], mc.any[String])).thenReturn( - Option(BlockLocationsAndStatus(blockLocations, blockStatus, None))) + Option(BlockLocationsAndStatus(blockLocations, blockStatus, blockLocations.head, None))) when(mockBlockManagerMaster.getLocations(mc.any[BlockId])).thenReturn(blockLocations) val store = makeBlockManager(8000, "executor1", mockBlockManagerMaster,