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")
}
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...))
}