Github user szyszy commented on a diff in the pull request:
https://github.com/apache/spark/pull/20761#discussion_r223159523
--- Diff:
resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
---
@@ -134,6 +166,42 @@ class YarnAllocatorSuite extends SparkFunSuite with
Matchers with BeforeAndAfter
size should be (0)
}
+ test("custom resource type requested from yarn") {
+ assume(ResourceRequestHelper.isYarnResourceTypesAvailable())
+ TestYarnResourceRequestHelper.initializeResourceTypes(List("gpu"))
+
+ // request a single container and receive it
+ val handler = createAllocatorWithAdditionalConfigs(1, Map(
+ YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "gpu" -> "2G",
+ YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "memory" -> "1G"
+ ))
+ handler.updateResourceRequests()
+ handler.getNumExecutorsRunning should be (0)
+ handler.getPendingAllocate.size should be (1)
+
+ val resource = Resource.newInstance(3072, 6)
+ val resourceTypes = Map("gpu" -> "2G")
+ ResourceRequestHelper.setResourceRequests(resourceTypes, resource)
+
+ val container = createContainerWithResource("host1", resource)
+ handler.handleAllocatedContainers(Array(container))
+
+ // verify custom resource type is part of rmClient.ask set
+ val askField = rmClient.getClass.getDeclaredField("ask")
--- End diff --
Good point.
I found a better way to test what I wanted and it also covers verifying
whether the `resource` field is set correctly.
Please check the new testcase's implementation.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
