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