Example #1
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")
}
Example #2
0
func main() {
	args := os.Args[1:]
	defer exit()

	assert(LoadEnvironment())

	command, off, err := docker.ParseArgs(args)
	assert(err)

	if command != "create" && command != "run" {
		if command == "version" {
			fmt.Printf("NVIDIA Docker: %s\n\n", Version)
		}
		assert(docker.Docker(args...))
	}

	opt, i, err := docker.ParseArgs(args[off+1:], command)
	assert(err)
	off += i + 1

	if (command == "create" || command == "run") && opt != "" {
		vols, err := VolumesNeeded(opt)
		assert(err)

		if vols != nil {
			var nargs []string
			var err error

			if Host != nil {
				nargs, err = GenerateRemoteArgs(opt, vols)
			} else {
				assert(nvidia.LoadUVM())
				assert(nvidia.Init())
				nargs, err = GenerateLocalArgs(opt, vols)
				nvidia.Shutdown()
			}
			assert(err)
			args = append(args[:off], append(nargs, args[off:]...)...)
		}
	}

	assert(docker.Docker(args...))
}