Github user vanzin commented on a diff in the pull request: https://github.com/apache/spark/pull/20761#discussion_r222053715 --- 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 -- I'm not sure this is testing anything interesting. It's basically testing YARN code, not Spark code, and using reflection for that... The code you're actually adding to `YarnAllocator` is not being tested, which is the code that creates the `resource` field in that class. Instead, you're manually creating a `Resource` in this test. Testing that code would be better than what you have here.
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org