Example #1
0
func devicesArgs() ([]string, error) {
	var args []string

	cdevs, err := nvidia.GetControlDevicePaths()
	if err != nil {
		return nil, err
	}
	for i := range cdevs {
		args = append(args, fmt.Sprintf("--device=%s", cdevs[i]))
	}

	devs, err := nvidia.LookupDevices(nvidia.LookupMinimal)
	if err != nil {
		return nil, err
	}

	if len(GPU) == 0 {
		for i := range devs {
			args = append(args, fmt.Sprintf("--device=%s", devs[i].Path))
		}
	} else {
		devs, err := nvidia.FilterDevices(devs, GPU)
		if err != nil {
			return nil, err
		}
		for i := range devs {
			args = append(args, fmt.Sprintf("--device=%s", devs[i].Path))
		}
	}
	return args, nil
}
Example #2
0
func main() {
	var err error

	flag.Parse()
	defer exit()

	if PrintVersion {
		fmt.Printf("NVIDIA Docker plugin: %s\n", Version)
		return
	}

	log.Println("Loading NVIDIA unified memory")
	assert(nvidia.LoadUVM())

	log.Println("Loading NVIDIA management library")
	assert(nvidia.Init())
	defer func() { assert(nvidia.Shutdown()) }()

	log.Println("Discovering GPU devices")
	Devices, err = nvidia.LookupDevices()
	assert(err)

	log.Println("Provisioning volumes at", VolumesPath)
	Volumes, err = nvidia.LookupVolumes(VolumesPath)
	assert(err)

	plugin := NewPluginAPI(SocketPath)
	remote := NewRemoteAPI(ListenAddr)

	log.Println("Serving plugin API at", SocketPath)
	log.Println("Serving remote API at", ListenAddr)
	p := plugin.Serve()
	r := remote.Serve()

	join, joined := make(chan int, 2), 0
L:
	for {
		select {
		case <-p:
			remote.Stop()
			p = nil
			join <- 1
		case <-r:
			plugin.Stop()
			r = nil
			join <- 1
		case j := <-join:
			if joined += j; joined == cap(join) {
				break L
			}
		}
	}
	assert(plugin.Error())
	assert(remote.Error())
	log.Println("Successfully terminated")
}