# go-gpuallocator

**Repository Path**: mirrors_NVIDIA/go-gpuallocator

## Basic Information

- **Project Name**: go-gpuallocator
- **Description**: Go Abstraction for Allocating NVIDIA GPUs with Custom Policies
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2020-08-18
- **Last Updated**: 2026-03-21

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

The `gpuallocator` package
-----------------------
The `gpuallocator` package provides a generic abstraction for performing GPU
allocations independent of the larger system the `gpuallocator` is integrated
with.

The abstractions provided by this package are not meant to do actual
"allocation" of GPUs to any specific entity, but rather run the algorithm
responsible for deciding which GPUs should be chosen for allocation based on
the set of GPUs available in the system and the number of GPUs being requested.

Different policies can be hooked in to run different allocation algorithms
depending on the specific needs of the system.

For example, a system like Kubernetes would use this package to help decide
which subset of GPUs to hand out to a container once it has decided how many
GPUs it should be granted. The policy it would choose would be based on the
topological ordering of GPUs to ensure optimal affinity for groups of GPUs
allocated to the same container.

The primary object provided by this package is the `Allocator` object, and the
primary interface used to decide how allocation should actually occur is called
`Policy`. 

More details on each of these can be found below.

The `Allocator` Object
----------------------
The primary object provided by the `gpuallocator` package is that of an
`Allocator`.

A new `Allocator` can be instantiated as follows:

```
func NewAllocator(policy Policy) (*Allocator, error)
```

Once instantiated, an `Allocator` relies on NVML to do GPU discovery and
maintains an internal list of all GPUs available on a node.

Using this list it then allocates GPUs to callers of the `Allocate()` or
`AllocateSpecific()` function and frees GPUs back to this list from callers of
the `Free()` function, as seen below:

```
func (a *Allocator) Allocate(num int) []*Device
func (a *Allocator) AllocateSpecific(devices... *Device) []*Device
func (a *Allocator) Free(devices... *Device)
```

The `Policy` Interface
----------------------
```
type Policy interface {
	Allocate(devices []*Device, num int) []*Device
}
```

The `Policy` interface contains a single function `Allocate()`, which takes a
slice of devices and size `num` as arguments and returns a subset of that
slice of length `num`.

Implementers of this interface take on the heavy lifting
of implementing the actual allocation logic used by the `Allocator`.

The following default policies are implemented as part of this package:
```
func NewSimplePolicy() Policy
func NewBestEffortPolicy() Policy
func NewStaticDGX1Policy(gpuType GPUType) Policy
func NewStaticDGX2Policy() Policy
```

With the following convenience wrappers for simple and best effort allocators:
```
func NewSimpleAllocator() (*Allocator, error)
func NewBestEffortAllocator() (*Allocator, error)
```

`Simple` takes a slice of GPU devices and simply allocates `num` GPUs from the
front of it.

`BestEffort` attempts to allocate GPUs in topological order, considering both
NVLINKs between GPUs and their placement in the PCIe hierarchy.
The choice of GPUs to allocate is optimized to assume that all future
allocations will be of size 'num' as well.

Sample Usage
------------
```
package main

import (
	"fmt"
	"os"

	"github.com/NVIDIA/go-gpuallocator/gpuallocator"
)

func main() {
	allocator, err := gpuallocator.NewSimpleAllocator()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
	}

	for _, gpu := range allocator.GPUs {
		fmt.Printf("%v\n", gpu.Details())
	}

	fmt.Printf("\n")

	for _, i := range []int{1, 2, 4, 8} {
		gpus := allocator.Allocate(i)
		fmt.Printf("Simple allocation of %v GPUs: %v\n", i, gpus)
		allocator.Free(gpus...)
	}
}
```