Github user szyszy commented on a diff in the pull request:
https://github.com/apache/spark/pull/20761#discussion_r223928827
--- Diff:
resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
---
@@ -134,6 +163,29 @@ class YarnAllocatorSuite extends SparkFunSuite with
Matchers with BeforeAndAfter
size should be (0)
}
+ test("custom resource requested from yarn") {
+ assume(ResourceRequestHelper.isYarnResourceTypesAvailable())
+ ResourceRequestTestHelper.initializeResourceTypes(List("gpu"))
+
+ val mockAmClient = mock(classOf[AMRMClient[ContainerRequest]])
+ val handler = createAllocatorWithAdditionalConfigs(1, Map(
+ YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "gpu" -> "2G"), mockAmClient)
+
+ handler.updateResourceRequests()
+ val container = createContainerWithResource("host1", handler.resource)
+ handler.handleAllocatedContainers(Array(container))
+
+ // get amount of memory and vcores from resource, so effectively
skipping their validation
+ val expectedResources =
Resource.newInstance(handler.resource.getMemory(),
+ handler.resource.getVirtualCores)
+ ResourceRequestHelper.setResourceRequests(Map("gpu" -> "2G"),
expectedResources)
+ val captor = ArgumentCaptor.forClass(classOf[ContainerRequest])
+
+ verify(mockAmClient).addContainerRequest(captor.capture())
+ val containerRequest: ContainerRequest = captor.getValue
+ assert(containerRequest.getCapability == expectedResources)
--- End diff --
fixed
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
