Example usage for org.apache.commons.io FileUtils openOutputStream

List of usage examples for org.apache.commons.io FileUtils openOutputStream

  1. HOME
  2. Java
  3. org.apache.commons
  4. org.apache.commons.io.*
  5. FileUtils
  6. openOutputStream

Introduction

In this page you can find the example usage for org.apache.commons.io FileUtils openOutputStream.

Prototype

public static FileOutputStream openOutputStream(File file) throws IOException

Source Link

Document

Opens a FileOutputStream for the specified file, checking and creating the parent directory if it does not exist.

Usage

From source file:org.alternativevision.gpx.GPXParser.java

public void writeGPX(GPX gpx, File gpxFile)
        throws IOException, ParserConfigurationException, TransformerException {
    OutputStream out = FileUtils.openOutputStream(gpxFile);
    writeGPX(gpx, out);/*ww w  .j a  va 2  s . com*/
    out.flush();
    out.close();
}

From source file:org.apache.hadoop.gateway.services.topology.DefaultTopologyServiceTest.java

private File createFile(File parent, String name, String resource, long timestamp) throws IOException {
    File file = new File(parent, name);
    if (!file.exists()) {
        FileUtils.touch(file);//from  w ww.jav a2s .c o  m
    }
    InputStream input = ClassLoader.getSystemResourceAsStream(resource);
    OutputStream output = FileUtils.openOutputStream(file);
    IOUtils.copy(input, output);
    //KNOX-685: output.flush();
    input.close();
    output.close();
    file.setLastModified(timestamp);
    assertTrue("Failed to create test file " + file.getAbsolutePath(), file.exists());
    assertTrue("Failed to populate test file " + file.getAbsolutePath(), file.length() > 0);

    return file;
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandler.java

@Test
public void testcheckAndDeleteCgroup() throws Exception {
    CgroupsLCEResourcesHandler handler = new CgroupsLCEResourcesHandler();
    handler.setConf(new YarnConfiguration());
    handler.initConfig();//from  w ww.  j  a v  a  2 s .  c  o m

    FileUtils.deleteQuietly(cgroupDir);
    // Test 0
    // tasks file not present, should return false
    Assert.assertFalse(handler.checkAndDeleteCgroup(cgroupDir));

    File tfile = new File(cgroupDir.getAbsolutePath(), "tasks");
    FileOutputStream fos = FileUtils.openOutputStream(tfile);
    File fspy = Mockito.spy(cgroupDir);

    // Test 1, tasks file is empty
    // tasks file has no data, should return true
    Mockito.stub(fspy.delete()).toReturn(true);
    Assert.assertTrue(handler.checkAndDeleteCgroup(fspy));

    // Test 2, tasks file has data
    fos.write("1234".getBytes());
    fos.close();
    // tasks has data, would not be able to delete, should return false
    Assert.assertFalse(handler.checkAndDeleteCgroup(fspy));
    FileUtils.deleteQuietly(cgroupDir);

}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandler.java

@Test
public void testDeleteCgroup() throws Exception {
    final ControlledClock clock = new ControlledClock();
    CgroupsLCEResourcesHandler handler = new CgroupsLCEResourcesHandler();
    handler.setConf(new YarnConfiguration());
    handler.initConfig();/*from w  ww. j a v  a 2s. c  o  m*/
    handler.clock = clock;

    FileUtils.deleteQuietly(cgroupDir);

    // Create a non-empty tasks file
    File tfile = new File(cgroupDir.getAbsolutePath(), "tasks");
    FileOutputStream fos = FileUtils.openOutputStream(tfile);
    fos.write("1234".getBytes());
    fos.close();

    final CountDownLatch latch = new CountDownLatch(1);
    new Thread() {
        @Override
        public void run() {
            latch.countDown();
            try {
                Thread.sleep(200);
            } catch (InterruptedException ex) {
                //NOP
            }
            clock.tickMsec(YarnConfiguration.DEFAULT_NM_LINUX_CONTAINER_CGROUPS_DELETE_TIMEOUT);
        }
    }.start();
    latch.await();
    Assert.assertFalse(handler.deleteCgroup(cgroupDir.getAbsolutePath()));
    FileUtils.deleteQuietly(cgroupDir);
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerAMDGPU.java

@Test
public void testcheckAndDeleteCgroup() throws Exception {
    CgroupsLCEResourcesHandlerGPU handler = new CgroupsLCEResourcesHandlerGPU();
    handler.setConf(new YarnConfiguration());
    handler.initConfig();//from ww  w  . j a  v  a  2 s  .  c  o m

    FileUtils.deleteQuietly(cgroupDir);
    // Test 0
    // tasks file not present, should return false
    Assert.assertFalse(handler.checkAndDeleteCgroup(cgroupDir));

    File tfile = new File(cgroupDir.getAbsolutePath(), "tasks");
    FileOutputStream fos = FileUtils.openOutputStream(tfile);
    File fspy = Mockito.spy(cgroupDir);

    // Test 1, tasks file is empty
    // tasks file has no data, should return true
    Mockito.stub(fspy.delete()).toReturn(true);
    Assert.assertTrue(handler.checkAndDeleteCgroup(fspy));

    // Test 2, tasks file has data
    fos.write("1234".getBytes());
    fos.close();
    // tasks has data, would not be able to delete, should return false
    Assert.assertFalse(handler.checkAndDeleteCgroup(fspy));
    FileUtils.deleteQuietly(cgroupDir);

}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerAMDGPU.java

@Test
public void testDeleteCgroup() throws Exception {
    final MockClock clock = new MockClock();
    clock.time = System.currentTimeMillis();
    CgroupsLCEResourcesHandlerGPU handler = new CgroupsLCEResourcesHandlerGPU();
    handler.setConf(new YarnConfiguration());
    handler.initConfig();/*from w  ww.  j a  v a  2  s  .  c o m*/
    handler.clock = clock;

    FileUtils.deleteQuietly(cgroupDir);

    // Create a non-empty tasks file
    File tfile = new File(cgroupDir.getAbsolutePath(), "tasks");
    FileOutputStream fos = FileUtils.openOutputStream(tfile);
    fos.write("1234".getBytes());
    fos.close();

    final CountDownLatch latch = new CountDownLatch(1);
    new Thread() {
        @Override
        public void run() {
            latch.countDown();
            try {
                Thread.sleep(200);
            } catch (InterruptedException ex) {
                //NOP
            }
            clock.time += YarnConfiguration.DEFAULT_NM_LINUX_CONTAINER_CGROUPS_DELETE_TIMEOUT;
        }
    }.start();
    latch.await();
    Assert.assertFalse(handler.deleteCgroup(cgroupDir.getAbsolutePath()));
    FileUtils.deleteQuietly(cgroupDir);
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerAMDGPU.java

@Test
public void testContainerGPUAllocation() throws IOException {
    LinuxContainerExecutor mockLCE = new MockLinuxContainerExecutor();
    CustomCgroupsLCEResourceHandlerGPU handler = new CustomCgroupsLCEResourceHandlerGPU();
    YarnConfiguration conf = new YarnConfiguration();
    conf.setBoolean(YarnConfiguration.NM_GPU_RESOURCE_ENABLED, true);
    conf.set(YarnConfiguration.NM_GPU_MANAGEMENT_IMPL, "io.hops.management.amd.AMDManagementLibrary");
    final int numGPUs = 8;
    ResourceCalculatorPlugin plugin = Mockito.mock(ResourceCalculatorPlugin.class);
    handler.setConf(conf);//from  w  w  w. jav a2 s. c o m
    handler.initConfig();

    // create mock cgroup
    File cgroupMountDirCPU = createMockCgroupMount(cgroupDir, "cpu");
    File cgroupMountDirGPU = createMockCgroupMount(cgroupDir, "devices");
    File whiteList = new File(cgroupMountDirGPU, "devices.list");
    FileOutputStream fos1 = FileUtils.openOutputStream(whiteList);
    fos1.write(("a *:* rwm\n").getBytes());

    // create mock mtab
    File mockMtab = createMockMTab(cgroupDir);

    ContainerId id = ContainerId.fromString("container_1_1_1_1");
    GPUAllocator gpuAllocator = null;

    conf.setInt(YarnConfiguration.NM_GPUS, numGPUs);

    gpuAllocator = new GPUAllocator(new CustomGPUmanagementLibrary(), conf);

    GPU gpu0 = new GPU(new Device(226, 0), new Device(226, 128));
    GPU gpu1 = new GPU(new Device(226, 1), new Device(226, 129));
    GPU gpu2 = new GPU(new Device(226, 2), new Device(226, 130));
    GPU gpu3 = new GPU(new Device(226, 3), new Device(226, 131));
    GPU gpu4 = new GPU(new Device(226, 4), new Device(226, 132));
    GPU gpu5 = new GPU(new Device(226, 5), new Device(226, 133));
    GPU gpu6 = new GPU(new Device(226, 6), new Device(226, 134));
    GPU gpu7 = new GPU(new Device(226, 7), new Device(226, 135));

    // setup our handler and call init()
    handler.setGPUAllocator(gpuAllocator);
    handler.setMtabFile(mockMtab.getAbsolutePath());
    handler.init(mockLCE, plugin);
    File containerDirGPU = new File(cgroupMountDirGPU, id.toString());
    File denyFile = new File(containerDirGPU, "devices.deny");
    Assert.assertFalse(denyFile.exists());

    HashMap<String, HashSet<GPU>> allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(0, allocations.size());

    //FIRST ALLOCATION
    handler.preExecute(id, Resource.newInstance(1024, 1, 2));
    Assert.assertTrue(containerDirGPU.exists());
    Assert.assertTrue(containerDirGPU.isDirectory());
    Assert.assertTrue(denyFile.exists());

    allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(1, allocations.size());

    Assert.assertTrue(gpuAllocator.getConfiguredAvailableGPUs().size() == 6);
    //SECOND ALLOCATION
    HashSet<GPU> deviceAllocation1 = gpuAllocator.allocate("test_container", 2);

    allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(2, allocations.size());
    Assert.assertEquals(2, allocations.get("test_container").size());
    Assert.assertTrue(allocations.get("test_container").contains(gpu2));
    Assert.assertTrue(allocations.get("test_container").contains(gpu3));

    HashSet<GPU> deniedDevices1 = deviceAllocation1;

    Assert.assertTrue(deniedDevices1.contains(gpu0));
    Assert.assertTrue(deniedDevices1.contains(gpu1));
    Assert.assertFalse(deniedDevices1.contains(gpu2));
    Assert.assertFalse(deniedDevices1.contains(gpu3));
    Assert.assertTrue(deniedDevices1.contains(gpu4));
    Assert.assertTrue(deniedDevices1.contains(gpu5));
    Assert.assertTrue(deniedDevices1.contains(gpu6));
    Assert.assertTrue(deniedDevices1.contains(gpu7));

    handler.postExecute(id);

    HashSet<GPU> availableGPUs = gpuAllocator.getConfiguredAvailableGPUs();

    Assert.assertTrue(availableGPUs.contains(gpu0));
    Assert.assertTrue(availableGPUs.contains(gpu1));
    Assert.assertFalse(availableGPUs.contains(gpu2));
    Assert.assertFalse(availableGPUs.contains(gpu3));
    Assert.assertTrue(availableGPUs.contains(gpu4));
    Assert.assertTrue(availableGPUs.contains(gpu5));
    Assert.assertTrue(availableGPUs.contains(gpu6));
    Assert.assertTrue(availableGPUs.contains(gpu7));

    HashSet<GPU> deviceAllocation2 = gpuAllocator.allocate("test_container2", 3);

    allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(2, allocations.size());

    HashSet<GPU> deniedDevices2 = deviceAllocation2;
    Assert.assertFalse(deniedDevices2.contains(gpu0)); //allocated in first call
    Assert.assertFalse(deniedDevices2.contains(gpu1));
    Assert.assertTrue(deniedDevices2.contains(gpu2));
    Assert.assertTrue(deniedDevices2.contains(gpu3));
    Assert.assertFalse(deniedDevices2.contains(gpu4));
    Assert.assertTrue(deniedDevices2.contains(gpu5));
    Assert.assertTrue(deniedDevices2.contains(gpu6));
    Assert.assertTrue(deniedDevices2.contains(gpu7));

    gpuAllocator.release("test_container");

    allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(1, allocations.size());

    gpuAllocator.release("test_container2");
    allocations = gpuAllocator.getAllocations();
    Assert.assertEquals(0, allocations.size());
    availableGPUs = gpuAllocator.getConfiguredAvailableGPUs();
    Assert.assertTrue(availableGPUs.contains(gpu0)); //allocated in first call
    Assert.assertTrue(availableGPUs.contains(gpu1));
    Assert.assertTrue(availableGPUs.contains(gpu2));
    Assert.assertTrue(availableGPUs.contains(gpu3));
    Assert.assertTrue(availableGPUs.contains(gpu4));
    Assert.assertTrue(availableGPUs.contains(gpu5));
    Assert.assertTrue(availableGPUs.contains(gpu6));
    Assert.assertTrue(availableGPUs.contains(gpu7));

    FileUtils.deleteQuietly(cgroupDir);
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerAMDGPU.java

@Test
public void testSubsetGPUsSchedulable() throws IOException {
    LinuxContainerExecutor mockLCE = new MockLinuxContainerExecutor();
    CustomCgroupsLCEResourceHandlerGPU handler = new CustomCgroupsLCEResourceHandlerGPU();
    YarnConfiguration conf = new YarnConfiguration();
    conf.setBoolean(YarnConfiguration.NM_GPU_RESOURCE_ENABLED, true);
    conf.set(YarnConfiguration.NM_GPU_MANAGEMENT_IMPL, "io.hops.management.amd.AMDManagementLibrary");
    final int numGPUs = 6;
    ResourceCalculatorPlugin plugin = Mockito.mock(ResourceCalculatorPlugin.class);
    handler.setConf(conf);//from   w  w w .  j  av  a  2  s  .  co m
    handler.initConfig();

    // create mock cgroup
    File cgroupMountDirCPU = createMockCgroupMount(cgroupDir, "cpu");
    File cgroupMountDirGPU = createMockCgroupMount(cgroupDir, "devices");
    File whiteList = new File(cgroupMountDirGPU, "devices.list");
    FileOutputStream fos1 = FileUtils.openOutputStream(whiteList);
    fos1.write(("a *:* rwm\n").getBytes());

    // create mock mtab
    File mockMtab = createMockMTab(cgroupDir);

    ContainerId id = ContainerId.fromString("container_1_1_1_1");
    GPUAllocator gpuAllocator = null;

    conf.setInt(YarnConfiguration.NM_GPUS, numGPUs);

    gpuAllocator = new GPUAllocator(new CustomGPUmanagementLibrarySubset(), conf);

    GPU gpu0 = new GPU(new Device(226, 0), new Device(226, 128));
    GPU gpu1 = new GPU(new Device(226, 1), new Device(226, 129));
    GPU gpu2 = new GPU(new Device(226, 2), new Device(226, 130));
    GPU gpu3 = new GPU(new Device(226, 3), new Device(226, 131));
    GPU gpu4 = new GPU(new Device(226, 4), new Device(226, 132));
    GPU gpu5 = new GPU(new Device(226, 5), new Device(226, 133));

    // setup our handler and call init()
    handler.setGPUAllocator(gpuAllocator);
    handler.setMtabFile(mockMtab.getAbsolutePath());
    handler.init(mockLCE, plugin);
    File containerDirGPU = new File(cgroupMountDirGPU, id.toString());
    File denyFile = new File(containerDirGPU, "devices.deny");
    Assert.assertFalse(denyFile.exists());
    //FIRST ALLOCATION
    handler.preExecute(id, Resource.newInstance(1024, 1, 2));
    Assert.assertTrue(containerDirGPU.exists());
    Assert.assertTrue(containerDirGPU.isDirectory());
    Assert.assertTrue(denyFile.exists());

    Assert.assertTrue(gpuAllocator.getConfiguredAvailableGPUs().size() == 4);
    //SECOND ALLOCATION
    HashSet<GPU> deviceAllocation1 = gpuAllocator.allocate("test_container", 2);

    HashSet<GPU> deniedDevices1 = deviceAllocation1;
    //gpu0 and gpu1 already allocated
    Assert.assertTrue(deniedDevices1.contains(gpu0));
    Assert.assertTrue(deniedDevices1.contains(gpu1));
    Assert.assertFalse(deniedDevices1.contains(gpu2));
    Assert.assertFalse(deniedDevices1.contains(gpu3));
    Assert.assertTrue(deniedDevices1.contains(gpu4));
    Assert.assertTrue(deniedDevices1.contains(gpu5));

    handler.postExecute(id);

    HashSet<GPU> deviceAllocation2 = gpuAllocator.allocate("test_container2", 2);

    HashSet<GPU> deniedDevices2 = deviceAllocation2;
    Assert.assertFalse(deniedDevices2.contains(gpu0)); //allocated in first call
    Assert.assertFalse(deniedDevices2.contains(gpu1));
    Assert.assertTrue(deniedDevices2.contains(gpu2));
    Assert.assertTrue(deniedDevices2.contains(gpu3));
    Assert.assertTrue(deniedDevices2.contains(gpu4));
    Assert.assertTrue(deniedDevices2.contains(gpu5));

    FileUtils.deleteQuietly(cgroupDir);
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerAMDGPU.java

@Test
public void testContainerRecovery() throws IOException {
    LinuxContainerExecutor mockLCE = new MockLinuxContainerExecutor();
    CustomCgroupsLCEResourceHandlerGPU handler = new CustomCgroupsLCEResourceHandlerGPU();
    YarnConfiguration conf = new YarnConfiguration();
    final int numGPUs = 8;
    ResourceCalculatorPlugin plugin = Mockito.mock(ResourceCalculatorPlugin.class);
    handler.setConf(conf);/*w  ww  .jav a2s .c  o  m*/
    handler.initConfig();
    // create mock cgroup
    File cgroupMountDirCPU = createMockCgroupMount(cgroupDir, "cpu");
    File cgroupMountDirGPU = createMockCgroupMount(cgroupDir, "devices");
    File whiteList = new File(cgroupMountDirGPU, "devices.list");
    FileOutputStream outputStream = FileUtils.openOutputStream(whiteList);
    outputStream.write(("a *:* rwm\n").getBytes());

    // create mock mtab
    File mockMtab = createMockMTab(cgroupDir);
    handler.setMtabFile(mockMtab.getAbsolutePath());

    conf.setInt(YarnConfiguration.NM_GPUS, numGPUs);
    conf.set(YarnConfiguration.NM_GPU_MANAGEMENT_IMPL, "io.hops.management.amd.AMDManagementLibrary");

    GPUAllocator gpuAllocator = new GPUAllocator(new CustomGPUmanagementLibrary(), conf);

    handler.setGPUAllocator(gpuAllocator);
    handler.init(mockLCE, plugin);

    ContainerId id1 = ContainerId.fromString("container_1_1_1_1");
    handler.preExecute(id1, Resource.newInstance(1024, 1, 2));
    File containerDir1 = new File(cgroupMountDirGPU, id1.toString());
    File listFile1 = new File(containerDir1, "devices.list");
    FileOutputStream fos1 = FileUtils.openOutputStream(listFile1);
    fos1.write(("c 226:0 rwm\n" + "c 226:1 rwm\n").getBytes());
    fos1.write(("c 226:128 rwm\n" + "c 226:129 rwm\n").getBytes());

    File tasksFile1 = new File(containerDir1, "tasks");
    FileOutputStream tasks1 = FileUtils.openOutputStream(tasksFile1);
    tasks1.write(("12934").getBytes());

    ContainerId id2 = ContainerId.fromString("container_1_1_1_2");
    handler.preExecute(id2, Resource.newInstance(1024, 1, 2));
    File containerDir2 = new File(cgroupMountDirGPU, id2.toString());
    File listFile2 = new File(containerDir2, "devices.list");
    FileOutputStream fos2 = FileUtils.openOutputStream(listFile2);
    fos2.write(("c 226:2 rwm\n" + "c 226:3 rwm\n").getBytes());
    fos1.write(("c 226:130 rwm\n" + "c 195:131 rwm\n").getBytes());

    File tasksFile2 = new File(containerDir2, "tasks");
    FileOutputStream tasks2 = FileUtils.openOutputStream(tasksFile2);
    tasks2.write(("12934").getBytes());

    ContainerId id3 = ContainerId.fromString("container_1_1_1_3");
    handler.preExecute(id3, Resource.newInstance(1024, 1, 2));
    File containerDir3 = new File(cgroupMountDirGPU, id3.toString());
    File listFile3 = new File(containerDir3, "devices.list");
    FileOutputStream fos3 = FileUtils.openOutputStream(listFile3);
    fos3.write(("c 226:4 rwm\n" + "c 226:5 rwm\n").getBytes());
    fos1.write(("c 226:132 rwm\n" + "c 226:133 rwm\n").getBytes());

    File tasksFile3 = new File(containerDir3, "tasks");
    FileOutputStream tasks3 = FileUtils.openOutputStream(tasksFile3);
    tasks3.write(("12934").getBytes());

    conf.setInt(YarnConfiguration.NM_GPUS, 8);
    gpuAllocator = new GPUAllocator(new CustomGPUmanagementLibrary(), conf);
    gpuAllocator.initialize(conf);
    handler = new CustomCgroupsLCEResourceHandlerGPU();
    conf = new YarnConfiguration();
    conf.setBoolean(YarnConfiguration.NM_GPU_RESOURCE_ENABLED, true);
    conf.set(YarnConfiguration.NM_GPU_MANAGEMENT_IMPL, "io.hops.management.amd.AMDManagementLibrary");
    plugin = Mockito.mock(ResourceCalculatorPlugin.class);
    //Mockito.doReturn(numGPUs).when(plugin).getNumGPUs();
    handler.setConf(conf);
    handler.initConfig();

    handler.setMtabFile(mockMtab.getAbsolutePath());
    handler.setGPUAllocator(gpuAllocator);
    handler.init(mockLCE, plugin);

    //No recovery yet
    Assert.assertEquals(8, gpuAllocator.getConfiguredAvailableGPUs().size());
    Assert.assertTrue(listFile1.exists());
    handler.recoverDeviceControlSystem(id1);
    Assert.assertEquals(6, gpuAllocator.getConfiguredAvailableGPUs().size());
    Assert.assertTrue(listFile2.exists());
    handler.recoverDeviceControlSystem(id2);
    Assert.assertEquals(4, gpuAllocator.getConfiguredAvailableGPUs().size());
    Assert.assertTrue(listFile3.exists());
    handler.recoverDeviceControlSystem(id3);
    Assert.assertEquals(2, gpuAllocator.getConfiguredAvailableGPUs().size());
    HashSet<GPU> availableGPUs = new HashSet<>(gpuAllocator.getConfiguredAvailableGPUs());
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 0), new Device(226, 128))));
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 1), new Device(226, 129))));
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 2), new Device(226, 130))));
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 3), new Device(226, 131))));
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 4), new Device(226, 132))));
    Assert.assertFalse(availableGPUs.contains(new GPU(new Device(226, 5), new Device(226, 133))));
    Assert.assertTrue(availableGPUs.contains(new GPU(new Device(226, 6), new Device(226, 134))));
    Assert.assertTrue(availableGPUs.contains(new GPU(new Device(226, 7), new Device(226, 135))));

    FileUtils.deleteQuietly(cgroupDir);
}

From source file:org.apache.hadoop.yarn.server.nodemanager.util.TestCgroupsLCEResourcesHandlerGPU.java

@Test
public void testContainerGPUAllocation() throws IOException {
    LinuxContainerExecutor mockLCE = new MockLinuxContainerExecutor();
    CustomCgroupsLCEResourceHandlerGPU handler = new CustomCgroupsLCEResourceHandlerGPU();
    YarnConfiguration conf = new YarnConfiguration();
    conf.setBoolean(YarnConfiguration.NM_GPU_RESOURCE_ENABLED, true);
    final int numGPUs = 8;
    ResourceCalculatorPlugin plugin = Mockito.mock(ResourceCalculatorPlugin.class);
    Mockito.doReturn(numGPUs).when(plugin).getNumGPUs();
    handler.setConf(conf);//w ww  . j  a  va  2 s .  c  o  m
    handler.initConfig();

    // create mock cgroup
    File cgroupMountDirCPU = createMockCgroupMount(cgroupDir, "cpu");
    File cgroupMountDirGPU = createMockCgroupMount(cgroupDir, "devices");
    File whiteList = new File(cgroupMountDirGPU, "devices.list");
    FileOutputStream fos1 = FileUtils.openOutputStream(whiteList);
    fos1.write(("a *:* rwm\n").getBytes());

    // create mock mtab
    File mockMtab = createMockMTab(cgroupDir);

    ContainerId id = ContainerId.fromString("container_1_1_1_1");
    GPUAllocator gpuAllocator = null;

    conf.setInt(YarnConfiguration.NM_GPUS, 8);

    gpuAllocator = new GPUAllocator(new CustomGPUmanagementLibrary(), conf);

    Device gpu0 = new Device(195, 0);
    Device gpu1 = new Device(195, 1);
    Device gpu2 = new Device(195, 2);
    Device gpu3 = new Device(195, 3);
    Device gpu4 = new Device(195, 4);
    Device gpu5 = new Device(195, 5);
    Device gpu6 = new Device(195, 6);
    Device gpu7 = new Device(195, 7);

    // setup our handler and call init()
    handler.setGPUAllocator(gpuAllocator);
    handler.setMtabFile(mockMtab.getAbsolutePath());
    handler.init(mockLCE, plugin);
    File containerDirGPU = new File(cgroupMountDirGPU, id.toString());
    File denyFile = new File(containerDirGPU, "devices.deny");
    Assert.assertFalse(denyFile.exists());
    //FIRST ALLOCATION
    handler.preExecute(id, Resource.newInstance(1024, 1, 2));
    Assert.assertTrue(containerDirGPU.exists());
    Assert.assertTrue(containerDirGPU.isDirectory());
    Assert.assertTrue(denyFile.exists());

    Assert.assertTrue(gpuAllocator.getConfiguredAvailableGPUs().size() == 6);
    //SECOND ALLOCATION
    HashSet<Device> deviceAllocation1 = gpuAllocator.allocate("test_container", 2);

    HashSet<Device> deniedDevices1 = deviceAllocation1;
    //gpu0 and gpu1 already allocated
    Assert.assertTrue(deniedDevices1.contains(gpu0));
    Assert.assertTrue(deniedDevices1.contains(gpu1));
    Assert.assertFalse(deniedDevices1.contains(gpu2));
    Assert.assertFalse(deniedDevices1.contains(gpu3));
    Assert.assertTrue(deniedDevices1.contains(gpu4));
    Assert.assertTrue(deniedDevices1.contains(gpu5));
    Assert.assertTrue(deniedDevices1.contains(gpu6));
    Assert.assertTrue(deniedDevices1.contains(gpu7));

    handler.postExecute(id);

    HashSet<Device> deviceAllocation2 = gpuAllocator.allocate("test_container2", 3);

    HashSet<Device> deniedDevices2 = deviceAllocation2;
    Assert.assertFalse(deniedDevices2.contains(gpu0)); //allocated in first call
    Assert.assertFalse(deniedDevices2.contains(gpu1));
    Assert.assertTrue(deniedDevices2.contains(gpu2));
    Assert.assertTrue(deniedDevices2.contains(gpu3));
    Assert.assertFalse(deniedDevices2.contains(gpu4));
    Assert.assertTrue(deniedDevices2.contains(gpu5));
    Assert.assertTrue(deniedDevices2.contains(gpu6));
    Assert.assertTrue(deniedDevices2.contains(gpu7));

    FileUtils.deleteQuietly(cgroupDir);
}