spark/resource/JavaResourceProfileSuite.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.resource;
0019
0020 import java.util.Map;
0021
0022 import static org.junit.Assert.*;
0023 import org.junit.Test;
0024
0025 // Test the ResourceProfile and Request api's from Java
0026 public class JavaResourceProfileSuite {
0027
0028   String GpuResource = "resource.gpu";
0029   String FPGAResource = "resource.fpga";
0030
0031   @Test
0032   public void testResourceProfileAccessFromJava() throws Exception {
0033     ExecutorResourceRequests execReqGpu =
0034       new ExecutorResourceRequests().resource(GpuResource, 2,"myscript", "");
0035     ExecutorResourceRequests execReqFpga =
0036       new ExecutorResourceRequests().resource(FPGAResource, 3, "myfpgascript", "nvidia");
0037
0038     ResourceProfileBuilder rprof = new ResourceProfileBuilder();
0039     rprof.require(execReqGpu);
0040     rprof.require(execReqFpga);
0041     TaskResourceRequests taskReq1 = new TaskResourceRequests().resource(GpuResource, 1);
0042     rprof.require(taskReq1);
0043
0044     assertEquals(rprof.executorResources().size(), 2);
0045     Map<String, ExecutorResourceRequest> eresources = rprof.executorResourcesJMap();
0046     assert(eresources.containsKey(GpuResource));
0047     ExecutorResourceRequest gpuReq = eresources.get(GpuResource);
0048     assertEquals(gpuReq.amount(), 2);
0049     assertEquals(gpuReq.discoveryScript(), "myscript");
0050     assertEquals(gpuReq.vendor(), "");
0051
0052     assert(eresources.containsKey(FPGAResource));
0053     ExecutorResourceRequest fpgaReq = eresources.get(FPGAResource);
0054     assertEquals(fpgaReq.amount(), 3);
0055     assertEquals(fpgaReq.discoveryScript(), "myfpgascript");
0056     assertEquals(fpgaReq.vendor(), "nvidia");
0057
0058     assertEquals(rprof.taskResources().size(), 1);
0059     Map<String, TaskResourceRequest> tresources = rprof.taskResourcesJMap();
0060     assert(tresources.containsKey(GpuResource));
0061     TaskResourceRequest taskReq = tresources.get(GpuResource);
0062     assertEquals(taskReq.amount(), 1.0, 0);
0063     assertEquals(taskReq.resourceName(), GpuResource);
0064   }
0065 }
0066