This repository has been archived by the owner on Jan 22, 2024. It is now read-only.
/
remote.go
105 lines (92 loc) · 2.18 KB
/
remote.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
// Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
package main
import (
"encoding/json"
"fmt"
"io/ioutil"
"log"
"net"
"net/http"
"net/url"
"os"
"strings"
"syscall"
"time"
"golang.org/x/crypto/ssh"
"golang.org/x/crypto/ssh/agent"
"golang.org/x/crypto/ssh/terminal"
)
const timeout = 10 * time.Second
const (
endpointInfo = "http://plugin/gpu/info/json"
endpointCLI = "http://plugin/docker/cli"
)
func GenerateRemoteArgs(image string, vols []string) ([]string, error) {
var info struct {
Version struct{ CUDA string }
}
c := httpClient(Host)
r, err := c.Get(endpointInfo)
if err != nil {
return nil, err
}
defer r.Body.Close()
if err := json.NewDecoder(r.Body).Decode(&info); err != nil {
return nil, err
}
if err := cudaSupported(image, info.Version.CUDA); err != nil {
return nil, err
}
uri := fmt.Sprintf("%s?vol=%s&dev=%s", endpointCLI,
strings.Join(vols, "+"),
strings.Join(GPU, "+"),
)
r2, err := c.Get(uri)
if err != nil {
return nil, err
}
defer r2.Body.Close()
b, err := ioutil.ReadAll(r2.Body)
if err != nil {
return nil, err
}
return strings.Split(string(b), " "), nil
}
func httpClient(addr *url.URL) *http.Client {
dial := func(string, string) (net.Conn, error) {
if addr.Scheme == "ssh" {
c, err := ssh.Dial("tcp", addr.Host, &ssh.ClientConfig{
User: addr.User.Username(),
Auth: sshAuths(addr),
})
if err != nil {
return nil, err
}
return c.Dial("tcp", addr.Opaque)
}
return net.Dial("tcp", addr.Host)
}
return &http.Client{
Timeout: timeout,
Transport: &http.Transport{Dial: dial},
}
}
func sshAuths(addr *url.URL) (methods []ssh.AuthMethod) {
if sock := os.Getenv("SSH_AUTH_SOCK"); sock != "" {
c, err := net.Dial("unix", sock)
if err != nil {
log.Println("Warning: failed to contact the local SSH agent")
} else {
auth := ssh.PublicKeysCallback(agent.NewClient(c).Signers)
methods = append(methods, auth)
}
}
auth := ssh.PasswordCallback(func() (string, error) {
fmt.Printf("%s@%s password: ", addr.User.Username(), addr.Host)
b, err := terminal.ReadPassword(int(syscall.Stdin))
fmt.Print("\n")
return string(b), err
})
methods = append(methods, auth)
return
}