Package org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.gpu

Class GpuResourceAllocator

java.lang.Object
org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.gpu.GpuResourceAllocator
public class GpuResourceAllocator extends Object
Allocate GPU resources according to requirements.

  • Constructor Details

    • GpuResourceAllocator

      public GpuResourceAllocator(Context ctx)

  • Method Details

    • addGpu

      public void addGpu(GpuDevice gpuDevice)
      Add GPU to the allowed list of GPUs.
      Parameters:
      gpuDevice - gpu device

    • getAvailableGpus

      @VisibleForTesting public int getAvailableGpus()

    • recoverAssignedGpus

      public void recoverAssignedGpus(org.apache.hadoop.yarn.api.records.ContainerId containerId) throws ResourceHandlerException
      Throws:
      ResourceHandlerException

    • getRequestedGpus

      public static int getRequestedGpus(org.apache.hadoop.yarn.api.records.Resource requestedResource)
      Get number of requested GPUs from resource.
      Parameters:
      requestedResource - requested resource
      Returns:
      #gpus.

    • assignGpus

      public org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.gpu.GpuResourceAllocator.GpuAllocation assignGpus(Container container) throws ResourceHandlerException
      Assign GPU to the specified container.
      Parameters:
      container - container to allocate
      Returns:
      allocation results.
      Throws:
      ResourceHandlerException - When failed to assign GPUs.

    • unassignGpus

      public void unassignGpus(org.apache.hadoop.yarn.api.records.ContainerId containerId)
      Clean up all GPUs assigned to containerId.
      Parameters:
      containerId - containerId

    • getDeviceAllocationMapping

      @VisibleForTesting public Map<GpuDevice,org.apache.hadoop.yarn.api.records.ContainerId> getDeviceAllocationMapping()

    • getAllowedGpus

      public List<GpuDevice> getAllowedGpus()

    • getAssignedGpus

      public List<AssignedGpuDevice> getAssignedGpus()

    • toString

      public String toString()
      Overrides:
      toString in class Object