diff --git a/plugin/Dockerfile.build b/plugin/Dockerfile.build new file mode 100644 index 00000000..2582dcb2 --- /dev/null +++ b/plugin/Dockerfile.build @@ -0,0 +1,34 @@ +# Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +FROM golang + +ENV NVIDIA_GPGKEY_SUM bd841d59a27a406e513db7d405550894188a4c1cd96bf8aa4f82f1b39e0b5c1c +ENV NVIDIA_GPGKEY_FPR 889bee522da690103c4b085ed88c3d385c37d3be +ENV NVIDIA_GDK_SUM 1e32e58f69fe29ee67b845233e7aa9347f37994463252bccbc8bfc8a7104ab5a + +RUN apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/GPGKEY && \ + apt-key adv --export --no-emit-version -a $NVIDIA_GPGKEY_FPR | tail -n +2 > cudasign.pub && \ + echo "$NVIDIA_GPGKEY_SUM cudasign.pub" | sha256sum -c --strict - && rm cudasign.pub && \ + echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1404/x86_64 /" > /etc/apt/sources.list.d/cuda.list + +RUN apt-get update && apt-get install -y --no-install-recommends --force-yes \ + cuda-cudart-dev-6-5=6.5-19 cuda-misc-headers-6-5=6.5-19 \ + && rm -rf /var/lib/apt/lists/* + +RUN objcopy --redefine-sym memcpy=memcpy@GLIBC_2.2.5 /usr/local/cuda-6.5/lib64/libcudart_static.a + +RUN wget -O gdk.run -q http://developer.download.nvidia.com/compute/cuda/7.5/Prod/local_installers/cuda_352_39_gdk_linux.run && \ + echo "$NVIDIA_GDK_SUM gdk.run" | sha256sum -c --strict - && \ + chmod +x gdk.run && ./gdk.run --silent && rm gdk.run + +COPY src /go/src +VOLUME /go/bin + +ENV CGO_CFLAGS "-I /usr/local/cuda-6.5/include -I /usr/include/nvidia/gdk" +ENV CGO_LDFLAGS "-L /usr/local/cuda-6.5/lib64 -L /usr/src/gdk/nvml/lib -ldl -lrt" + +ARG UID +RUN useradd --uid $UID build +USER build + +CMD go get -v -ldflags="-s" plugin diff --git a/plugin/Makefile b/plugin/Makefile new file mode 100644 index 00000000..e596c163 --- /dev/null +++ b/plugin/Makefile @@ -0,0 +1,26 @@ +# Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +MAKE_DIR := $(shell dirname $(abspath $(lastword $(MAKEFILE_LIST)))) +BIN_DIR := $(MAKE_DIR)/bin +USER_ID := $(shell id -u) + +IMAGE := nvdocker-build +PREFIX := /usr/local/nvidia +TARGET := nvidia-docker-plugin +PLUGIN := $(BIN_DIR)/plugin + +.PHONY: all install clean + +all : $(PLUGIN) + +$(PLUGIN) : + @docker build --build-arg UID=$(USER_ID) -t $(IMAGE) -f Dockerfile.build $(MAKE_DIR) + @mkdir -p $(BIN_DIR) + @docker run --rm -v $(BIN_DIR):/go/bin $(IMAGE) + +install: all + install -D -T -m 755 $(PLUGIN) $(PREFIX)/$(TARGET) + +clean : + -@docker rmi -f $(IMAGE) golang 2> /dev/null + @rm -rf $(BIN_DIR) diff --git a/plugin/src/cuda/cuda.go b/plugin/src/cuda/cuda.go new file mode 100644 index 00000000..c848e6df --- /dev/null +++ b/plugin/src/cuda/cuda.go @@ -0,0 +1,118 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package cuda + +// #cgo LDFLAGS: -lcudart_static +// #include +// #include +import "C" + +import ( + "errors" + "fmt" + "unsafe" +) + +type MemoryInfo struct { + ECC bool + Global uint + Shared uint // includes L1 cache + Constant uint + L2Cache uint + Bandwidth uint +} + +type Device struct { + handle C.int + + Gen string + Arch string + Cores uint + Memory MemoryInfo +} + +func cudaErr(ret C.cudaError_t) error { + if ret == C.cudaSuccess { + return nil + } + err := C.GoString(C.cudaGetErrorString(ret)) + return errors.New(err) +} + +var archToGen = map[string]string{ + "1": "Tesla", + "2": "Fermi", + "3": "Kepler", + "5": "Maxwell", +} + +var archToCoresPerSM = map[string]uint{ + "1.0": 8, // Tesla Generation (SM 1.0) G80 class + "1.1": 8, // Tesla Generation (SM 1.1) G8x G9x class + "1.2": 8, // Tesla Generation (SM 1.2) GT21x class + "1.3": 8, // Tesla Generation (SM 1.3) GT20x class + "2.0": 32, // Fermi Generation (SM 2.0) GF100 GF110 class + "2.1": 48, // Fermi Generation (SM 2.1) GF10x GF11x class + "3.0": 192, // Kepler Generation (SM 3.0) GK10x class + "3.2": 192, // Kepler Generation (SM 3.2) TK1 class + "3.5": 192, // Kepler Generation (SM 3.5) GK11x GK20x class + "3.7": 192, // Kepler Generation (SM 3.7) GK21x class + "5.0": 128, // Maxwell Generation (SM 5.0) GM10x class + "5.2": 128, // Maxwell Generation (SM 5.2) GM20x class + "5.3": 128, // Maxwell Generation (SM 5.3) TX1 class +} + +func GetDriverVersion() (string, error) { + var driver C.int + + err := cudaErr(C.cudaDriverGetVersion(&driver)) + d := fmt.Sprintf("%d.%d", int(driver)/1000, int(driver)%100/10) + return d, err +} + +func NewDevice(busID string) (*Device, error) { + var ( + dev C.int + prop C.struct_cudaDeviceProp + ) + + id := C.CString(busID) + if err := cudaErr(C.cudaDeviceGetByPCIBusId(&dev, id)); err != nil { + return nil, err + } + C.free(unsafe.Pointer(id)) + + if err := cudaErr(C.cudaGetDeviceProperties(&prop, dev)); err != nil { + return nil, err + } + arch := fmt.Sprintf("%d.%d", prop.major, prop.minor) + cores, ok := archToCoresPerSM[arch] + if !ok { + return nil, fmt.Errorf("unsupported CUDA arch: %s", arch) + } + + // Destroy the active CUDA context + cudaErr(C.cudaDeviceReset()) + + return &Device{ + handle: dev, + Gen: archToGen[arch[:1]], + Arch: arch, + Cores: cores * uint(prop.multiProcessorCount), + Memory: MemoryInfo{ + ECC: bool(prop.ECCEnabled != 0), + Global: uint(prop.totalGlobalMem / (1024 * 1024)), + Shared: uint(prop.sharedMemPerMultiprocessor / 1024), + Constant: uint(prop.totalConstMem / 1024), + L2Cache: uint(prop.l2CacheSize / 1024), + Bandwidth: 2 * uint((prop.memoryClockRate/1000)*(prop.memoryBusWidth/8)) / 1000, + }, + }, nil +} + +func CanAccessPeer(dev1, dev2 *Device) (bool, error) { + var ok C.int + + err := cudaErr(C.cudaDeviceCanAccessPeer(&ok, dev1.handle, dev2.handle)) + return (ok != 0), err +} diff --git a/plugin/src/graceful/graceful.go b/plugin/src/graceful/graceful.go new file mode 100644 index 00000000..75c4d1e9 --- /dev/null +++ b/plugin/src/graceful/graceful.go @@ -0,0 +1,100 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package graceful + +import ( + "net" + "net/http" + "sync" + "time" + + middleware "github.com/justinas/alice" + "gopkg.in/tylerb/graceful.v1" +) + +const timeout = 5 * time.Second + +type HTTPServer struct { + sync.Mutex + + network string + router *http.ServeMux + server *graceful.Server + err error +} + +func recovery(handler http.Handler) http.Handler { + f := func(w http.ResponseWriter, r *http.Request) { + defer func() { + if recover() != nil { + w.WriteHeader(http.StatusInternalServerError) + } + }() + handler.ServeHTTP(w, r) + } + return http.HandlerFunc(f) +} + +func NewHTTPServer(net, addr string, mw ...middleware.Constructor) *HTTPServer { + r := http.NewServeMux() + + return &HTTPServer{ + network: net, + router: r, + server: &graceful.Server{ + Timeout: timeout, + Server: &http.Server{ + Addr: addr, + Handler: middleware.New(recovery).Append(mw...).Then(r), + ReadTimeout: timeout, + WriteTimeout: timeout, + }, + }, + } +} + +func (s *HTTPServer) Handle(method, route string, handler http.HandlerFunc) { + f := func(w http.ResponseWriter, r *http.Request) { + if r.Method != method { + http.NotFound(w, r) + return + } + handler.ServeHTTP(w, r) + } + s.router.HandleFunc(route, f) +} + +func (s *HTTPServer) Serve() <-chan struct{} { + l, err := net.Listen(s.network, s.server.Addr) + if err != nil { + s.Lock() + s.err = err + s.Unlock() + c := make(chan struct{}) + close(c) + return c + } + + c := s.server.StopChan() + go func() { + s.Lock() + defer s.Unlock() + + err = s.server.Serve(l) + if e, ok := err.(*net.OpError); !ok || (ok && e.Op != "accept") { + s.err = err + } + }() + return c +} + +func (s *HTTPServer) Stop() { + s.server.Stop(timeout) +} + +func (s *HTTPServer) Error() error { + s.Lock() + defer s.Unlock() + + return s.err +} diff --git a/plugin/src/ldcache/ldcache.go b/plugin/src/ldcache/ldcache.go new file mode 100644 index 00000000..ab669a42 --- /dev/null +++ b/plugin/src/ldcache/ldcache.go @@ -0,0 +1,193 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package ldcache + +import ( + "bytes" + "encoding/binary" + "errors" + "os" + "path/filepath" + "syscall" + "unsafe" +) + +const ldcachePath = "/etc/ld.so.cache" + +const ( + magicString1 = "ld.so-1.7.0" + magicString2 = "glibc-ld.so.cache" + magicVersion = "1.1" +) + +const ( + flagTypeMask = 0x00ff + flagTypeELF = 0x0001 + + flagArchMask = 0xff00 + flagArchI386 = 0x0000 + flagArchX8664 = 0x0300 + flagArchX32 = 0x0800 +) + +var ErrInvalidCache = errors.New("invalid ld.so.cache file") + +type Header1 struct { + Magic [len(magicString1) + 1]byte // include null delimiter + NLibs uint32 +} + +type Entry1 struct { + Flags int32 + Key, Value uint32 +} + +type Header2 struct { + Magic [len(magicString2)]byte + Version [len(magicVersion)]byte + NLibs uint32 + TableSize uint32 + _ [3]uint32 // unused + _ uint64 // force 8 byte alignment +} + +type Entry2 struct { + Flags int32 + Key, Value uint32 + OSVersion uint32 + HWCap uint64 +} + +type LDCache struct { + *bytes.Reader + + data, libs []byte + header Header2 + entries []Entry2 +} + +func Open() (*LDCache, error) { + f, err := os.Open(ldcachePath) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + d, err := syscall.Mmap(int(f.Fd()), 0, int(fi.Size()), + syscall.PROT_READ, syscall.MAP_PRIVATE) + if err != nil { + return nil, err + } + + cache := &LDCache{data: d, Reader: bytes.NewReader(d)} + return cache, cache.parse() +} + +func (c *LDCache) Close() error { + return syscall.Munmap(c.data) +} + +func (c *LDCache) Magic() string { + return string(c.header.Magic[:]) +} + +func (c *LDCache) Version() string { + return string(c.header.Version[:]) +} + +func strn(b []byte, n int) string { + return string(b[:n]) +} + +func (c *LDCache) parse() error { + var header Header1 + + // Check for the old format (< glibc-2.2) + if c.Len() <= int(unsafe.Sizeof(header)) { + return ErrInvalidCache + } + if strn(c.data, len(magicString1)) == magicString1 { + if err := binary.Read(c, binary.LittleEndian, &header); err != nil { + return err + } + n := int64(header.NLibs) * int64(unsafe.Sizeof(Entry1{})) + offset, err := c.Seek(n, 1) // skip old entries + if err != nil { + return err + } + n = (-offset) & int64(unsafe.Alignof(c.header)) + _, err = c.Seek(n, 1) // skip padding + if err != nil { + return err + } + } + + c.libs = c.data[c.Size()-int64(c.Len()):] // kv offsets start here + if err := binary.Read(c, binary.LittleEndian, &c.header); err != nil { + return err + } + if c.Magic() != magicString2 || c.Version() != magicVersion { + return ErrInvalidCache + } + c.entries = make([]Entry2, c.header.NLibs) + if err := binary.Read(c, binary.LittleEndian, &c.entries); err != nil { + return err + } + return nil +} + +func (c *LDCache) Lookup(libs ...string) (paths32, paths64 []string) { + type void struct{} + var paths *[]string + + set := make(map[string]void) + prefix := make([][]byte, len(libs)) + + for i := range libs { + prefix[i] = []byte(libs[i]) + } + for _, e := range c.entries { + if ((e.Flags & flagTypeMask) & flagTypeELF) == 0 { + continue + } + switch e.Flags & flagArchMask { + case flagArchX8664: + paths = &paths64 + case flagArchX32: + fallthrough + case flagArchI386: + paths = &paths32 + default: + continue + } + if e.Key > uint32(len(c.libs)) || e.Value > uint32(len(c.libs)) { + continue + } + lib := c.libs[e.Key:] + value := c.libs[e.Value:] + + for _, p := range prefix { + if bytes.HasPrefix(lib, p) { + n := bytes.IndexByte(value, 0) + if n < 0 { + break + } + path, err := filepath.EvalSymlinks(strn(value, n)) + if err != nil { + break + } + if _, ok := set[path]; ok { + break + } + set[path] = void{} + *paths = append(*paths, path) + break + } + } + } + return +} diff --git a/plugin/src/nvml/cpu_mask.go b/plugin/src/nvml/cpu_mask.go new file mode 100644 index 00000000..7780d8b7 --- /dev/null +++ b/plugin/src/nvml/cpu_mask.go @@ -0,0 +1,82 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package nvml + +import ( + "bytes" + "encoding/binary" + "encoding/hex" + "io/ioutil" + "os" + "path" + "path/filepath" +) + +type cpuMask [4]uint64 + +func (m cpuMask) cpuNode() (uint, error) { + masks, err := getMasks() + if err != nil { + return 0, err + } + for i, mask := range masks { + if cmpMasks(m, mask) { + return uint(i), nil + } + } + return 0, ErrCPUAffinity +} + +var cache []cpuMask + +func getMasks() (masks []cpuMask, err error) { + if cache != nil { + return cache, nil + } + err = filepath.Walk("/sys/devices/system/node", + func(p string, fi os.FileInfo, err error) error { + var mask cpuMask + + if err != nil { + return err + } + if path.Base(p) != "cpumap" { + return nil + } + f, err := ioutil.ReadFile(p) + if err != nil { + return err + } + buf := bytes.Split(f[:len(f)-1], []byte(",")) + if len(buf)/2 != len(mask) { + return ErrCPUAffinity + } + + for i := range mask { + h := hex32ToUint64(buf[i*2]) + l := hex32ToUint64(buf[i*2+1]) + mask[len(mask)-1-i] = (h << 32) | l + } + masks = append(masks, mask) + return nil + }) + cache = masks + return +} + +func cmpMasks(m1, m2 cpuMask) bool { + for i := range m1 { + if m1[i] != m2[i] { + return false + } + } + return true +} + +func hex32ToUint64(b []byte) uint64 { + var n uint32 + + h, _ := hex.DecodeString(string(b)) + binary.Read(bytes.NewBuffer(h), binary.BigEndian, &n) + return uint64(n) +} diff --git a/plugin/src/nvml/nvml.go b/plugin/src/nvml/nvml.go new file mode 100644 index 00000000..03c30480 --- /dev/null +++ b/plugin/src/nvml/nvml.go @@ -0,0 +1,332 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package nvml + +// #cgo LDFLAGS: -lnvidia-ml +// #include "nvml_dl.h" +import "C" + +import ( + "errors" + "fmt" +) + +const ( + szDriver = C.NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE + szName = C.NVML_DEVICE_NAME_BUFFER_SIZE + szUUID = C.NVML_DEVICE_UUID_BUFFER_SIZE + szProcs = 32 + szProcName = 64 +) + +var ErrCPUAffinity = errors.New("failed to retrieve CPU affinity") + +type P2PLinkType uint + +const ( + P2PLinkUnknown P2PLinkType = iota + P2PLinkCrossCPU + P2PLinkSameCPU + P2PLinkHostBridge + P2PLinkMultiSwitch + P2PLinkSingleSwitch + P2PLinkSameBoard +) + +type P2PLink struct { + BusID string + Link P2PLinkType +} + +func (t P2PLinkType) String() string { + switch t { + case P2PLinkCrossCPU: + return "Cross CPU socket" + case P2PLinkSameCPU: + return "Same CPU socket" + case P2PLinkHostBridge: + return "Host PCI bridge" + case P2PLinkMultiSwitch: + return "Multiple PCI switches" + case P2PLinkSingleSwitch: + return "Single PCI switch" + case P2PLinkSameBoard: + return "Same board" + case P2PLinkUnknown: + } + return "???" +} + +type ClockInfo struct { + Graphics uint + SM uint + Memory uint +} + +type PCIInfo struct { + BusID string + BAR1 uint64 + Bandwidth uint +} + +type Device struct { + handle C.nvmlDevice_t + + Name string + UUID string + Path string + Power uint + CPUAffinity uint + PCI PCIInfo + Clocks ClockInfo + Topology []P2PLink +} + +type UtilizationInfo struct { + GPU uint + Encoder uint + Decoder uint +} + +type PCIStatusInfo struct { + BAR1Used uint64 + Throughput []uint +} + +type MemoryInfo struct { + GlobalUsed uint64 + ECCErrors []uint64 +} + +type ProcessInfo struct { + PID uint + Name string +} + +type DeviceStatus struct { + Power uint + Temperature uint + Utilization UtilizationInfo + Memory MemoryInfo + Clocks ClockInfo + PCI PCIStatusInfo + Processes []ProcessInfo +} + +func nvmlErr(ret C.nvmlReturn_t) error { + if ret == C.NVML_SUCCESS { + return nil + } + err := C.GoString(C.nvmlErrorString(ret)) + return errors.New(err) +} + +func assert(ret C.nvmlReturn_t) { + if err := nvmlErr(ret); err != nil { + panic(err) + } +} + +func Init() error { + if err := C.nvmlInit_dl(); err != nil { + return errors.New(C.GoString(err)) + } + return nvmlErr(C.nvmlInit()) +} + +func Shutdown() error { + C.nvmlShutdown_dl() + return nvmlErr(C.nvmlShutdown()) +} + +func GetDeviceCount() (uint, error) { + var n C.uint + + err := nvmlErr(C.nvmlDeviceGetCount(&n)) + return uint(n), err +} + +func GetDriverVersion() (string, error) { + var driver [szDriver]C.char + + err := nvmlErr(C.nvmlSystemGetDriverVersion(&driver[0], szDriver)) + return C.GoString(&driver[0]), err +} + +var pcieGenToBandwidth = map[int]uint{ + 1: 250, // MB/s + 2: 500, + 3: 985, + 4: 1969, +} + +func NewDevice(idx uint) (device *Device, err error) { + var ( + dev C.nvmlDevice_t + name [szName]C.char + uuid [szUUID]C.char + pci C.nvmlPciInfo_t + minor C.uint + bar1 C.nvmlBAR1Memory_t + power C.uint + clock [3]C.uint + pciel [2]C.uint + mask cpuMask + ) + + defer func() { + if r := recover(); r != nil { + err = r.(error) + } + }() + + assert(C.nvmlDeviceGetHandleByIndex(C.uint(idx), &dev)) + assert(C.nvmlDeviceGetName(dev, &name[0], szName)) + assert(C.nvmlDeviceGetUUID(dev, &uuid[0], szUUID)) + assert(C.nvmlDeviceGetPciInfo(dev, &pci)) + assert(C.nvmlDeviceGetMinorNumber(dev, &minor)) + assert(C.nvmlDeviceGetBAR1MemoryInfo(dev, &bar1)) + assert(C.nvmlDeviceGetPowerManagementLimit(dev, &power)) + assert(C.nvmlDeviceGetMaxClockInfo(dev, C.NVML_CLOCK_GRAPHICS, &clock[0])) + assert(C.nvmlDeviceGetMaxClockInfo(dev, C.NVML_CLOCK_SM, &clock[1])) + assert(C.nvmlDeviceGetMaxClockInfo(dev, C.NVML_CLOCK_MEM, &clock[2])) + assert(C.nvmlDeviceGetMaxPcieLinkGeneration(dev, &pciel[0])) + assert(C.nvmlDeviceGetMaxPcieLinkWidth(dev, &pciel[1])) + assert(C.nvmlDeviceGetCpuAffinity(dev, C.uint(len(mask)), (*C.ulong)(&mask[0]))) + cpu, err := mask.cpuNode() + if err != nil { + return nil, err + } + + device = &Device{ + handle: dev, + Name: C.GoString(&name[0]), + UUID: C.GoString(&uuid[0]), + Path: fmt.Sprintf("/dev/nvidia%d", uint(minor)), + Power: uint(power / 1000), + CPUAffinity: cpu, + PCI: PCIInfo{ + BusID: C.GoString(&pci.busId[0]), + BAR1: uint64(bar1.bar1Total / (1024 * 1024)), + Bandwidth: pcieGenToBandwidth[int(pciel[0])] * uint(pciel[1]) / 1000, + }, + Clocks: ClockInfo{ + Graphics: uint(clock[0]), + SM: uint(clock[1]), + Memory: uint(clock[2]), + }, + } + return +} + +func (d *Device) Status() (status *DeviceStatus, err error) { + var ( + power C.uint + temp C.uint + usage C.nvmlUtilization_t + encoder [2]C.uint + decoder [2]C.uint + mem C.nvmlMemory_t + ecc [3]C.ulonglong + clock [3]C.uint + bar1 C.nvmlBAR1Memory_t + throughput [2]C.uint + procname [szProcName]C.char + procs [szProcs]C.nvmlProcessInfo_t + nprocs = C.uint(szProcs) + ) + + defer func() { + if r := recover(); r != nil { + err = r.(error) + } + }() + + assert(C.nvmlDeviceGetPowerUsage(d.handle, &power)) + assert(C.nvmlDeviceGetTemperature(d.handle, C.NVML_TEMPERATURE_GPU, &temp)) + assert(C.nvmlDeviceGetUtilizationRates(d.handle, &usage)) + assert(C.nvmlDeviceGetEncoderUtilization(d.handle, &encoder[0], &encoder[1])) + assert(C.nvmlDeviceGetDecoderUtilization(d.handle, &decoder[0], &decoder[1])) + assert(C.nvmlDeviceGetMemoryInfo(d.handle, &mem)) + assert(C.nvmlDeviceGetClockInfo(d.handle, C.NVML_CLOCK_GRAPHICS, &clock[0])) + assert(C.nvmlDeviceGetClockInfo(d.handle, C.NVML_CLOCK_SM, &clock[1])) + assert(C.nvmlDeviceGetClockInfo(d.handle, C.NVML_CLOCK_MEM, &clock[2])) + assert(C.nvmlDeviceGetBAR1MemoryInfo(d.handle, &bar1)) + assert(C.nvmlDeviceGetComputeRunningProcesses(d.handle, &nprocs, &procs[0])) + + status = &DeviceStatus{ + Power: uint(power / 1000), + Temperature: uint(temp), + Utilization: UtilizationInfo{ + GPU: uint(usage.gpu), + Encoder: uint(encoder[0]), + Decoder: uint(decoder[0]), + }, + Memory: MemoryInfo{ + GlobalUsed: uint64(mem.used / (1024 * 1024)), + }, + Clocks: ClockInfo{ + Graphics: uint(clock[0]), + SM: uint(clock[1]), + Memory: uint(clock[2]), + }, + PCI: PCIStatusInfo{ + BAR1Used: uint64(bar1.bar1Used / (1024 * 1024)), + }, + } + + r := C.nvmlDeviceGetMemoryErrorCounter(d.handle, C.NVML_MEMORY_ERROR_TYPE_UNCORRECTED, C.NVML_VOLATILE_ECC, + C.NVML_MEMORY_LOCATION_L1_CACHE, &ecc[0]) + if r != C.NVML_ERROR_NOT_SUPPORTED { // only supported on Tesla cards + assert(r) + assert(C.nvmlDeviceGetMemoryErrorCounter(d.handle, C.NVML_MEMORY_ERROR_TYPE_UNCORRECTED, C.NVML_VOLATILE_ECC, + C.NVML_MEMORY_LOCATION_L2_CACHE, &ecc[1])) + assert(C.nvmlDeviceGetMemoryErrorCounter(d.handle, C.NVML_MEMORY_ERROR_TYPE_UNCORRECTED, C.NVML_VOLATILE_ECC, + C.NVML_MEMORY_LOCATION_DEVICE_MEMORY, &ecc[2])) + status.Memory.ECCErrors = []uint64{uint64(ecc[0]), uint64(ecc[1]), uint64(ecc[2])} + } + + r = C.nvmlDeviceGetPcieThroughput(d.handle, C.NVML_PCIE_UTIL_RX_BYTES, &throughput[0]) + if r != C.NVML_ERROR_NOT_SUPPORTED { // only supported on Maxwell or newer + assert(r) + assert(C.nvmlDeviceGetPcieThroughput(d.handle, C.NVML_PCIE_UTIL_TX_BYTES, &throughput[1])) + status.PCI.Throughput = []uint{uint(throughput[0]), uint(throughput[1])} + } + + status.Processes = make([]ProcessInfo, nprocs) + for i := range status.Processes { + status.Processes[i].PID = uint(procs[i].pid) + assert(C.nvmlSystemGetProcessName(procs[i].pid, &procname[0], szProcName)) + status.Processes[i].Name = C.GoString(&procname[0]) + } + return +} + +func GetP2PLink(dev1, dev2 *Device) (link P2PLinkType, err error) { + var level C.nvmlGpuTopologyLevel_t + + r := C.nvmlDeviceGetTopologyCommonAncestor_dl(dev1.handle, dev2.handle, &level) + if r == C.NVML_ERROR_FUNCTION_NOT_FOUND { + return P2PLinkUnknown, nil + } + if err = nvmlErr(r); err != nil { + return + } + switch level { + case C.NVML_TOPOLOGY_INTERNAL: + link = P2PLinkSameBoard + case C.NVML_TOPOLOGY_SINGLE: + link = P2PLinkSingleSwitch + case C.NVML_TOPOLOGY_MULTIPLE: + link = P2PLinkMultiSwitch + case C.NVML_TOPOLOGY_HOSTBRIDGE: + link = P2PLinkHostBridge + case C.NVML_TOPOLOGY_CPU: + link = P2PLinkSameCPU + case C.NVML_TOPOLOGY_SYSTEM: + link = P2PLinkCrossCPU + default: + err = errors.New("unsupported P2P link type") + } + return +} diff --git a/plugin/src/nvml/nvml_dl.c b/plugin/src/nvml/nvml_dl.c new file mode 100644 index 00000000..0fd82e2c --- /dev/null +++ b/plugin/src/nvml/nvml_dl.c @@ -0,0 +1,38 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +#include +#include + +#include "nvml_dl.h" + +#define DLSYM(x, sym) \ +do { \ + x = dlsym(handle, #sym); \ + if (dlerror() != NULL) { \ + return NVML_ERROR_FUNCTION_NOT_FOUND; \ + } \ +} while (0) + +typedef nvmlReturn_t (*nvmlSym_t)(); + +static void *handle; + +char *NVML_DL(nvmlInit)(void) +{ + handle = dlopen(NULL, RTLD_NOW); + return (handle ? NULL : dlerror()); +} + +void NVML_DL(nvmlShutdown)(void) +{ + dlclose(handle); +} + +nvmlReturn_t NVML_DL(nvmlDeviceGetTopologyCommonAncestor)( + nvmlDevice_t dev1, nvmlDevice_t dev2, nvmlGpuTopologyLevel_t *info) +{ + nvmlSym_t sym; + + DLSYM(sym, nvmlDeviceGetTopologyCommonAncestor); + return (*sym)(dev1, dev2, info); +} diff --git a/plugin/src/nvml/nvml_dl.h b/plugin/src/nvml/nvml_dl.h new file mode 100644 index 00000000..fe0380b2 --- /dev/null +++ b/plugin/src/nvml/nvml_dl.h @@ -0,0 +1,15 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +#ifndef _NVML_DL_H_ +#define _NVML_DL_H_ + +#include + +#define NVML_DL(x) x##_dl + +extern char *NVML_DL(nvmlInit)(void); +extern void NVML_DL(nvmlShutdown)(void); +extern nvmlReturn_t NVML_DL(nvmlDeviceGetTopologyCommonAncestor)( + nvmlDevice_t, nvmlDevice_t, nvmlGpuTopologyLevel_t *); + +#endif // _NVML_DL_H_ diff --git a/plugin/src/plugin/devices.go b/plugin/src/plugin/devices.go new file mode 100644 index 00000000..7b01635c --- /dev/null +++ b/plugin/src/plugin/devices.go @@ -0,0 +1,85 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "sort" + + "cuda" + "nvml" +) + +type NVMLDev nvml.Device +type CUDADev cuda.Device + +type Device struct { + *NVMLDev + *CUDADev +} + +func (d *Device) Status() (*nvml.DeviceStatus, error) { + return (*nvml.Device)(d.NVMLDev).Status() +} + +type deviceSorter struct { + devs []Device +} + +func (s *deviceSorter) Len() int { + return len(s.devs) +} + +func (s *deviceSorter) Swap(i, j int) { + s.devs[i], s.devs[j] = s.devs[j], s.devs[i] +} + +func (s *deviceSorter) Less(i, j int) bool { + return s.devs[i].PCI.BusID < s.devs[j].PCI.BusID +} + +func GetDevices() (devs []Device, err error) { + var i uint + + n, err := nvml.GetDeviceCount() + if err != nil { + return nil, err + } + devs = make([]Device, 0, n) + + for i = 0; i < n; i++ { + nd, err := nvml.NewDevice(i) + if err != nil { + return nil, err + } + cd, err := cuda.NewDevice(nd.PCI.BusID) + if err != nil { + return nil, err + } + devs = append(devs, Device{(*NVMLDev)(nd), (*CUDADev)(cd)}) + } + + for i = 0; i < n-1; i++ { + for j := i + 1; j < n; j++ { + ok, err := cuda.CanAccessPeer( + (*cuda.Device)(devs[i].CUDADev), + (*cuda.Device)(devs[j].CUDADev), + ) + if err != nil { + return nil, err + } + if ok { + l, err := nvml.GetP2PLink( + (*nvml.Device)(devs[i].NVMLDev), + (*nvml.Device)(devs[j].NVMLDev), + ) + if err != nil { + return nil, err + } + devs[i].Topology = append(devs[i].Topology, nvml.P2PLink{devs[j].PCI.BusID, l}) + devs[j].Topology = append(devs[j].Topology, nvml.P2PLink{devs[i].PCI.BusID, l}) + } + } + } + sort.Sort(&deviceSorter{devs}) + return +} diff --git a/plugin/src/plugin/main.go b/plugin/src/plugin/main.go new file mode 100644 index 00000000..88e6e784 --- /dev/null +++ b/plugin/src/plugin/main.go @@ -0,0 +1,103 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "flag" + "log" + "os" + "os/exec" + + "nvml" +) + +var ( + ListenAddr string + VolumePrefix string + SocketPath string + + Devices []Device + Volumes VolumeMap +) + +func init() { + log.SetPrefix("nvidia-docker-plugin | ") + + flag.StringVar(&ListenAddr, "l", "localhost:3476", "Server listen address") + flag.StringVar(&VolumePrefix, "p", "", "Volumes prefix path (default is to use mktemp)") + flag.StringVar(&SocketPath, "s", "/run/docker/plugins/nvidia.sock", "NVIDIA plugin socket path") +} + +func assert(err error) { + if err != nil { + log.Panicln("Error:", err) + } +} + +func exit() { + code := 0 + if recover() != nil { + code = 1 + } + os.Exit(code) +} + +func modprobeUVM() error { + return exec.Command("nvidia-modprobe", "-u", "-c=0").Run() +} + +func main() { + var err error + + flag.Parse() + defer exit() + + log.Println("Loading NVIDIA management library") + assert(nvml.Init()) + defer func() { assert(nvml.Shutdown()) }() + + log.Println("Loading NVIDIA unified memory module") + assert(modprobeUVM()) + + log.Println("Discovering GPU devices") + Devices, err = GetDevices() + assert(err) + + if VolumePrefix == "" { + log.Println("Creating volumes") + } else { + log.Println("Creating volumes at", VolumePrefix) + } + Volumes, err = GetVolumes(VolumePrefix) + 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") +} diff --git a/plugin/src/plugin/plugin_api.go b/plugin/src/plugin/plugin_api.go new file mode 100644 index 00000000..87409e37 --- /dev/null +++ b/plugin/src/plugin/plugin_api.go @@ -0,0 +1,67 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "encoding/json" + "net/http" + "os" + "path" + + "graceful" +) + +type plugin interface { + implement() string + register(*PluginAPI) +} + +type PluginAPI struct { + *graceful.HTTPServer + + plugins []plugin +} + +func accept(handler http.Handler) http.Handler { + const header = "application/vnd.docker.plugins.v1.1+json" + + f := func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Accept") != header { + w.WriteHeader(http.StatusNotAcceptable) + return + } + handler.ServeHTTP(w, r) + } + return http.HandlerFunc(f) +} + +func NewPluginAPI(addr string) *PluginAPI { + os.MkdirAll(path.Dir(addr), 0700) + + a := &PluginAPI{ + HTTPServer: graceful.NewHTTPServer("unix", addr, accept), + } + a.Handle("POST", "/Plugin.Activate", a.activate) + + a.register( + new(pluginVolume), + ) + return a +} + +func (a *PluginAPI) register(plugins ...plugin) { + for _, p := range plugins { + p.register(a) + a.plugins = append(a.plugins, p) + } +} + +func (a *PluginAPI) activate(resp http.ResponseWriter, req *http.Request) { + r := struct{ Implements []string }{} + + r.Implements = make([]string, len(a.plugins)) + for i, p := range a.plugins { + r.Implements[i] = p.implement() + } + assert(json.NewEncoder(resp).Encode(r)) +} diff --git a/plugin/src/plugin/plugin_volume.go b/plugin/src/plugin/plugin_volume.go new file mode 100644 index 00000000..92332345 --- /dev/null +++ b/plugin/src/plugin/plugin_volume.go @@ -0,0 +1,86 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "encoding/json" + "net/http" +) + +type pluginVolume struct{} + +func (p *pluginVolume) implement() string { return "VolumeDriver" } + +func (p *pluginVolume) register(api *PluginAPI) { + prefix := "/" + p.implement() + + api.Handle("POST", prefix+".Create", p.create) + api.Handle("POST", prefix+".Remove", p.remove) + api.Handle("POST", prefix+".Mount", p.mount) + api.Handle("POST", prefix+".Unmount", p.unmount) + api.Handle("POST", prefix+".Path", p.path) +} + +func sptr(s string) *string { return &s } + +func errVolumeUnknown(vol string) *string { + return sptr("No such volume: " + vol) +} + +func (p *pluginVolume) create(resp http.ResponseWriter, req *http.Request) { + q := struct{ Name string }{} + r := struct{ Err *string }{} + + assert(json.NewDecoder(req.Body).Decode(&q)) + if v, ok := Volumes[q.Name]; ok { + if err := v.Create(); err != nil { + r.Err = sptr(err.Error()) + } + } else { + r.Err = errVolumeUnknown(q.Name) + } + assert(json.NewEncoder(resp).Encode(r)) +} + +func (p *pluginVolume) remove(resp http.ResponseWriter, req *http.Request) { + q := struct{ Name string }{} + r := struct{ Err *string }{} + + assert(json.NewDecoder(req.Body).Decode(&q)) + if v, ok := Volumes[q.Name]; ok { + if err := v.Remove(); err != nil { + r.Err = sptr(err.Error()) + } + } else { + r.Err = errVolumeUnknown(q.Name) + } + assert(json.NewEncoder(resp).Encode(r)) +} + +func (p *pluginVolume) mount(resp http.ResponseWriter, req *http.Request) { + q := struct{ Name string }{} + r := struct{ Mountpoint, Err *string }{} + + assert(json.NewDecoder(req.Body).Decode(&q)) + if v, ok := Volumes[q.Name]; ok { + r.Mountpoint = &v.Path + } else { + r.Err = errVolumeUnknown(q.Name) + } + assert(json.NewEncoder(resp).Encode(r)) +} + +func (p *pluginVolume) unmount(resp http.ResponseWriter, req *http.Request) { + q := struct{ Name string }{} + r := struct{ Err *string }{} + + assert(json.NewDecoder(req.Body).Decode(&q)) + if _, ok := Volumes[q.Name]; !ok { + r.Err = errVolumeUnknown(q.Name) + } + assert(json.NewEncoder(resp).Encode(r)) +} + +func (p *pluginVolume) path(resp http.ResponseWriter, req *http.Request) { + p.mount(resp, req) +} diff --git a/plugin/src/plugin/remote_api.go b/plugin/src/plugin/remote_api.go new file mode 100644 index 00000000..06c499f0 --- /dev/null +++ b/plugin/src/plugin/remote_api.go @@ -0,0 +1,50 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "net/http" + + "graceful" +) + +type restapi interface { + version() string + + gpuInfo(http.ResponseWriter, *http.Request) + gpuStatus(http.ResponseWriter, *http.Request) + cli(http.ResponseWriter, *http.Request) +} + +type RemoteAPI struct { + *graceful.HTTPServer + + apis []restapi +} + +func NewRemoteAPI(addr string) *RemoteAPI { + a := &RemoteAPI{ + HTTPServer: graceful.NewHTTPServer("tcp", addr), + } + a.register( + new(remoteV10), + ) + return a +} + +func (a *RemoteAPI) register(apis ...restapi) { + for i, api := range apis { + prefix := "/" + api.version() + + handlers: + a.Handle("GET", prefix+"/gpu/info", api.gpuInfo) + a.Handle("GET", prefix+"/gpu/status", api.gpuStatus) + a.Handle("GET", prefix+"/docker/cli", api.cli) + + if i == len(apis)-1 && prefix != "" { + prefix = "" + goto handlers + } + a.apis = append(a.apis, api) + } +} diff --git a/plugin/src/plugin/remote_v1.go b/plugin/src/plugin/remote_v1.go new file mode 100644 index 00000000..a9e22e6b --- /dev/null +++ b/plugin/src/plugin/remote_v1.go @@ -0,0 +1,167 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "bytes" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "text/tabwriter" + "text/template" + + "cuda" + "nvml" +) + +type remoteV10 struct{} + +func (r *remoteV10) version() string { return "v1.0" } + +func (r *remoteV10) gpuInfo(resp http.ResponseWriter, req *http.Request) { + const tpl = ` + Driver version: {{nvmlDriverVersion}} + Supported CUDA version: {{cudaDriverVersion}} + {{range $i, $e := .}} + Device #{{$i}} + Name: {{.Name}} + UUID: {{.UUID}} + Path: {{.Path}} + Gen: {{.Gen}} + Arch: {{.Arch}} + Cores: {{.Cores}} + Power: {{.Power}} W + CPU Affinity: NUMA node{{.CPUAffinity}} + PCI + BusID: {{.PCI.BusID}} + BAR1: {{.PCI.BAR1}} MiB + Bandwidth: {{.PCI.Bandwidth}} GB/s + Memory + ECC: {{.Memory.ECC}} + Global: {{.Memory.Global}} MiB + Constant: {{.Memory.Constant}} KiB + L1 / Shared: {{.Memory.Shared}} KiB + L2 Cache: {{.Memory.L2Cache}} KiB + Bandwidth: {{.Memory.Bandwidth}} GB/s + Clocks + SM: {{.Clocks.SM}} MHz + Memory: {{.Clocks.Memory}} MHz + Graphics: {{.Clocks.Graphics}} MHz + P2P Available{{if len .Topology | eq 0}}: N/A{{else}}{{range .Topology}} + {{.BusID}} - {{(.Link.String)}}{{end}}{{end}} + {{end}} + ` + m := template.FuncMap{ + "nvmlDriverVersion": nvml.GetDriverVersion, + "cudaDriverVersion": cuda.GetDriverVersion, + } + t := template.Must(template.New("").Funcs(m).Parse(tpl)) + w := tabwriter.NewWriter(resp, 0, 4, 0, ' ', 0) + + assert(t.Execute(w, Devices)) + assert(w.Flush()) +} + +func (r *remoteV10) gpuStatus(resp http.ResponseWriter, req *http.Request) { + const tpl = `{{range $i, $e := .}}{{$s := (.Status)}} + Device #{{$i}} + Power: {{$s.Power}} / {{.Power}} W + Temperature: {{$s.Temperature}} °C + Utilization + GPU: {{$s.Utilization.GPU}} % + Encoder: {{$s.Utilization.Encoder}} % + Decoder: {{$s.Utilization.Decoder}} % + Memory + Global: {{$s.Memory.GlobalUsed}} / {{.Memory.Global}} MiB + ECC Errors{{if len $s.Memory.ECCErrors | eq 0}}: N/A{{else}} + L1 Cache: {{index $s.Memory.ECCErrors 0}} + L2 Cache: {{index $s.Memory.ECCErrors 1}} + Memory: {{index $s.Memory.ECCErrors 2}}{{end}} + PCI + BAR1: {{$s.PCI.BAR1Used}} / {{.PCI.BAR1}} MiB + Throughput{{if len $s.PCI.Throughput | eq 0}}: N/A{{else}} + RX: {{index $s.PCI.Throughput 0}} KB/s + TX: {{index $s.PCI.Throughput 1}} KB/s{{end}} + Clocks + SM: {{$s.Clocks.SM}} MHz + Memory: {{$s.Clocks.Memory}} MHz + Graphics: {{$s.Clocks.Graphics}} MHz + Processes{{if len $s.Processes | eq 0}}: N/A{{else}}{{range $s.Processes}} + {{.PID}} - {{.Name}}{{end}}{{end}} + {{end}} + ` + + t := template.Must(template.New("").Parse(tpl)) + w := tabwriter.NewWriter(resp, 0, 4, 0, ' ', 0) + + assert(t.Execute(w, Devices)) + assert(w.Flush()) +} + +func (r *remoteV10) cli(resp http.ResponseWriter, req *http.Request) { + var body bytes.Buffer + + ids := strings.Split(req.FormValue("dev"), " ") + if err := cliDevice(&body, ids); err != nil { + http.Error(resp, err.Error(), http.StatusBadRequest) + return + } + + names := strings.Split(req.FormValue("vol"), " ") + if err := cliVolume(&body, names); err != nil { + http.Error(resp, err.Error(), http.StatusBadRequest) + return + } + + resp.Write(body.Bytes()) +} + +func cliDevice(wr io.Writer, ids []string) error { + const tpl = "--device=/dev/nvidiactl --device=/dev/nvidia-uvm {{range .}}--device={{.}} {{end}}" + + devs := make([]string, 0, len(Devices)) + + if len(ids) == 1 && (ids[0] == "*" || ids[0] == "") { + for i := range Devices { + devs = append(devs, Devices[i].Path) + } + } else { + for _, id := range ids { + i, err := strconv.Atoi(id) + if err != nil || i < 0 || i >= len(Devices) { + return fmt.Errorf("invalid device: %s", id) + } + devs = append(devs, Devices[i].Path) + } + } + + t := template.Must(template.New("").Parse(tpl)) + assert(t.Execute(wr, devs)) + return nil +} + +func cliVolume(wr io.Writer, names []string) error { + const tpl = "--volume-driver=nvidia {{range .}}--volume={{.}} {{end}}" + + vols := make([]string, 0, len(Volumes)) + + if len(names) == 1 && (names[0] == "*" || names[0] == "") { + for _, v := range Volumes { + vols = append(vols, fmt.Sprintf("%s:%s", v.Name, v.Mountpoint)) + } + } else { + for _, n := range names { + v, ok := Volumes[n] + if !ok { + return fmt.Errorf("invalid volume: %s", n) + } + vols = append(vols, fmt.Sprintf("%s:%s", v.Name, v.Mountpoint)) + } + } + + t := template.Must(template.New("").Parse(tpl)) + assert(t.Execute(wr, vols)) + return nil +} diff --git a/plugin/src/plugin/volumes.go b/plugin/src/plugin/volumes.go new file mode 100644 index 00000000..8b8bc6ce --- /dev/null +++ b/plugin/src/plugin/volumes.go @@ -0,0 +1,185 @@ +// Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved. + +package main + +import ( + "bufio" + "bytes" + "debug/elf" + "io" + "io/ioutil" + "os" + "os/exec" + "path" + "path/filepath" + "strings" + + "ldcache" +) + +const ( + tempDir = "nvidia-volumes-" + lib32Dir = "lib" + lib64Dir = "lib64" +) + +type volumeDir struct { + name string + files []string +} + +type volumeInfo struct { + Path string + dirs []volumeDir +} + +type Volume struct { + Name string + Mountpoint string + components []string + + *volumeInfo +} + +type VolumeMap map[string]*Volume + +var volumes = []Volume{ + { + "bin", + "/usr/local/nvidia/bin", + []string{ + "nvidia-cuda-mps-control", + "nvidia-cuda-mps-server", + "nvidia-debugdump", + "nvidia-persistenced", + "nvidia-smi", + }, nil, + }, + { + "cuda", + "/usr/local/nvidia", + []string{ + "libcuda.so", + "libnvcuvid.so", + "libnvidia-compiler.so", + "libnvidia-encode.so", + "libnvidia-ml.so", + }, nil, + }, +} + +func (v *Volume) Create() (err error) { + if err = os.MkdirAll(v.Path, 0755); err != nil { + return + } + defer func() { + if err != nil { + v.Remove() + } + }() + + for _, d := range v.dirs { + dir := path.Join(v.Path, d.name) + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + for _, f := range d.files { + obj, err := elf.Open(f) + if err != nil { + return err + } + soname, err := obj.DynString(elf.DT_SONAME) + obj.Close() + if err != nil { + return err + } + + l := path.Join(dir, path.Base(f)) + if err := os.Link(f, l); err != nil { + return err + } + if len(soname) > 0 { + f = path.Join(v.Mountpoint, d.name, path.Base(f)) + l = path.Join(dir, soname[0]) + if err := os.Symlink(f, l); err != nil && + !os.IsExist(err.(*os.LinkError).Err) { + return err + } + } + } + } + return nil +} + +func (v *Volume) Remove() error { + return os.RemoveAll(v.Path) +} + +func which(bins ...string) ([]string, error) { + paths := make([]string, 0, len(bins)) + + out, _ := exec.Command("which", bins...).Output() + r := bufio.NewReader(bytes.NewBuffer(out)) + for { + p, err := r.ReadString('\n') + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if p = strings.TrimSpace(p); !path.IsAbs(p) { + continue + } + path, err := filepath.EvalSymlinks(p) + if err != nil { + return nil, err + } + paths = append(paths, path) + } + return paths, nil +} + +func GetVolumes(prefix string) (vols VolumeMap, err error) { + if prefix == "" { + prefix, err = ioutil.TempDir("", tempDir) + if err != nil { + return + } + } + + cache, err := ldcache.Open() + if err != nil { + return nil, err + } + defer func() { + if e := cache.Close(); err == nil { + err = e + } + }() + + vols = make(VolumeMap, len(volumes)) + + for i := range volumes { + vol := &volumes[i] + vol.volumeInfo = &volumeInfo{ + Path: path.Join(prefix, vol.Name), + } + + if vol.Name == "bin" { + bins, err := which(vol.components...) + if err != nil { + return nil, err + } + vol.dirs = append(vol.dirs, volumeDir{".", bins}) + } else { + libs32, libs64 := cache.Lookup(vol.components...) + vol.dirs = append(vol.dirs, + volumeDir{lib32Dir, libs32}, + volumeDir{lib64Dir, libs64}, + ) + } + vols[vol.Name] = vol + } + return +}