func main() {
var err error
flag.Parse()
defer exit()
log.Println("Loading NVIDIA management library")
assert(nvidia.Init())
defer func() { assert(nvidia.Shutdown()) }()
log.Println("Loading NVIDIA unified memory")
assert(nvidia.LoadUVM())
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() {
var option string
var n int
args := os.Args[1:]
defer exit()
assert(nvidia.Init())
defer func() { assert(nvidia.Shutdown()) }()
command, i, err := docker.ParseArgs(args)
assert(err)
if command != "" {
option, n, err = docker.ParseArgs(args[i+1:], command)
i += n + 1
assert(err)
}
switch command {
case "create":
fallthrough
case "run":
if option != "" {
a := GenerateDockerArgs(option)
args = append(args[:i], append(a, args[i:]...)...)
}
case "volume":
if option == "setup" {
assert(CreateLocalVolumes())
return
}
default:
}
assert(nvidia.LoadUVM())
assert(docker.Docker(args...))
}