client.go 5.2 KB
Newer Older
Q
Qiao Longfei 已提交
1
package client
2

3
import (
D
dongzhihong 已提交
4
	"errors"
5 6 7 8
	"hash/fnv"
	"sort"
	"time"

9
	"github.com/PaddlePaddle/Paddle/go/connection"
Q
Qiao Longfei 已提交
10
	"github.com/PaddlePaddle/Paddle/go/pserver"
H
Helin Wang 已提交
11
	log "github.com/sirupsen/logrus"
12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
)

// TODO(helin): add RPC call retry logic

// Selector selects if the client should initialize parameter servers.
type Selector interface {
	Select() bool
}

// Server is the identification of a parameter Server.
type Server struct {
	Index int
	Addr  string
}

// Lister lists currently available parameter servers.
type Lister interface {
	List() []Server
}

H
Helin Wang 已提交
32
// Client is the client to parameter servers.
33
type Client struct {
34 35
	sel      Selector
	pservers []*connection.Conn
36 37
}

H
Helin Wang 已提交
38
// NewClient creates a new client.
39 40 41 42 43 44 45 46 47 48 49 50 51
func NewClient(l Lister, pserverNum int, sel Selector) *Client {
	c := &Client{sel: sel}
	c.pservers = make([]*connection.Conn, pserverNum)
	for i := 0; i < pserverNum; i++ {
		c.pservers[i] = connection.New()
	}
	go c.monitorPservers(l, pserverNum)
	return c
}

// monitorPservers monitors pserver addresses, and updates connection
// when the address changes.
func (c *Client) monitorPservers(l Lister, pserverNum int) {
52
	lastServers := make([]Server, pserverNum)
53 54 55 56 57 58 59 60
	ticker := time.NewTicker(10 * time.Second)
	monitor := func() {
		curServers := make([]Server, pserverNum)
		list := l.List()
		for _, l := range list {
			curServers[l.Index] = l
		}

61
		for i := range lastServers {
H
Helin Wang 已提交
62 63 64
			if lastServers[i].Addr == curServers[i].Addr {
				continue
			}
65

H
Helin Wang 已提交
66 67
			if curServers[i].Addr == "" {
				err := c.pservers[i].Close()
68
				if err != nil {
H
Helin Wang 已提交
69
					log.Errorln(err)
70
				}
H
Helin Wang 已提交
71 72

				continue
73
			}
H
Helin Wang 已提交
74 75 76

			err := c.pservers[i].Connect(curServers[i].Addr)
			if err != nil {
H
Helin Wang 已提交
77
				log.Errorln(err)
H
Helin Wang 已提交
78 79 80 81 82 83 84

				// connect to addr failed, set
				// to last known addr in order
				// to retry next time.
				curServers[i].Addr = lastServers[i].Addr
			}

85 86
		}

87
		lastServers = curServers
88 89 90
	}

	monitor()
H
Helin Wang 已提交
91
	for range ticker.C {
92 93
		monitor()
	}
94 95
}

H
Helin Wang 已提交
96 97 98 99 100 101 102 103
// BeginInitParams begins to initialize parameters on parameter
// servers.
//
// BeginInitParams will be called from multiple trainers, only one
// trainer will be selected to initialize the parameters on parameter
// servers. Other trainers will be blocked until the initialization is
// done, and they need to get the initialized parameters from
// parameter servers using GetParams.
104 105
func (c *Client) BeginInitParams() bool {
	return c.sel.Select()
106 107
}

H
Helin Wang 已提交
108
// InitParam initializes the parameter on parameter servers.
Q
Qiao Longfei 已提交
109
func (c *Client) InitParam(paramWithConfigs pserver.ParameterWithConfig) error {
110
	return c.pservers[c.partition(paramWithConfigs.Param.Name)].Call("Service.InitParam", paramWithConfigs, nil)
111 112
}

H
Helin Wang 已提交
113 114
// FinishInitParams tells parameter servers client has sent all
// parameters to parameter servers as initialization.
115
func (c *Client) FinishInitParams() error {
116
	for _, p := range c.pservers {
117
		err := p.Call("Service.FinishInitParams", 0, nil)
118 119 120 121
		if err != nil {
			return err
		}
	}
122 123 124
	return nil
}

H
Helin Wang 已提交
125 126
// SendGrads sends gradients to parameter servers for updating
// parameters.
Q
Qiao Longfei 已提交
127
func (c *Client) SendGrads(grads []pserver.Gradient) error {
128
	if len(grads) == 0 {
D
dongzhihong 已提交
129
		return errors.New("no gradient received")
130
	}
131 132
	errCh := make(chan error, len(grads))
	for _, g := range grads {
Q
Qiao Longfei 已提交
133
		go func(g pserver.Gradient) {
134
			err := c.pservers[c.partition(g.Name)].Call("Service.SendGrad", g, nil)
135 136 137 138 139 140 141 142 143 144 145 146 147 148 149
			errCh <- err
		}(g)
	}

	recv := 0
	for err := range errCh {
		if err != nil {
			return err
		}

		recv++
		if recv == len(grads) {
			break
		}
	}
150 151 152
	return nil
}

153
type result struct {
H
Helin Wang 已提交
154
	idx   int
Q
Qiao Longfei 已提交
155
	param pserver.Parameter
H
Helin Wang 已提交
156
	err   error
157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
}

type results []result

func (r results) Len() int {
	return len(r)
}

func (r results) Less(i int, j int) bool {
	return r[i].idx < r[j].idx
}

func (r results) Swap(i int, j int) {
	r[i], r[j] = r[j], r[i]
}

H
Helin Wang 已提交
173
// GetParams gets parameters from parameter servers.
Q
Qiao Longfei 已提交
174
func (c *Client) GetParams(names []string) ([]pserver.Parameter, error) {
175 176 177 178
	rCh := make(chan result, len(names))

	for idx, name := range names {
		go func(name string, idx int) {
Q
Qiao Longfei 已提交
179
			var parameter pserver.Parameter
180
			err := c.pservers[c.partition(name)].Call("Service.GetParam", name, &parameter)
H
Helin Wang 已提交
181
			rCh <- result{idx: idx, param: parameter, err: err}
182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199
		}(name, idx)
	}

	var rs results
	recv := 0
	for r := range rCh {
		if r.err != nil {
			return nil, r.err
		}
		rs = append(rs, r)

		recv++
		if recv == len(names) {
			break
		}
	}
	sort.Sort(rs)

Q
Qiao Longfei 已提交
200
	ps := make([]pserver.Parameter, len(rs))
201
	for i := range rs {
H
Helin Wang 已提交
202
		ps[i] = rs[i].param
203 204 205
	}

	return ps, nil
206 207
}

208 209 210 211 212
// Save indicates parameters to save the parameter to the given path.
func (c *Client) Save(path string) error {
	errCh := make(chan error, len(c.pservers))

	for _, p := range c.pservers {
213
		err := p.Call("Service.Save", path, nil)
214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230
		errCh <- err
	}

	recv := 0
	for err := range errCh {
		if err != nil {
			return err
		}

		recv++
		if recv == len(c.pservers) {
			break
		}
	}

	// TODO(helin): there will be many files under path, need to
	// merge them into a single file.
231 232 233
	return nil
}

234 235
func strHash(s string) uint32 {
	h := fnv.New32a()
H
Helin Wang 已提交
236
	_, _ = h.Write([]byte(s))
237 238 239 240 241 242 243 244
	return h.Sum32()
}

// TODO(helin): now partition only select which parameter server to
// send the entire parameter. We need to partition a parameter into
// small blocks and send to different parameter servers.
func (c *Client) partition(key string) int {
	return int(strHash(key) % uint32(len(c.pservers)))
245
}