service.go 9.1 KB
Newer Older
1 2 3
package master

import (
4 5 6
	"bytes"
	"compress/gzip"
	"encoding/gob"
7
	"errors"
8 9
	"os"
	"path/filepath"
10 11 12
	"sync"
	"time"

H
Helin Wang 已提交
13 14
	log "github.com/sirupsen/logrus"

15
	"github.com/PaddlePaddle/recordio"
16 17
)

18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33
const (
	dialTimeout = 5 * time.Second
)

// Store is the interface for save and load the master state.
type Store interface {
	Save([]byte) error
	Load() ([]byte, error)
}

// Chunk is a chunk of data consisted of several data instances.
type Chunk struct {
	Path  string
	Index recordio.Index // chunk index
}

G
gongweibao 已提交
34 35 36 37 38 39
// TaskMeta is a struct which stores task's meta info.
type TaskMeta struct {
	ID    int
	Epoch int
}

40 41
// Task is the basic unit of data instances assigned to trainers.
type Task struct {
G
gongweibao 已提交
42
	Meta   TaskMeta
43 44 45 46
	Chunks []Chunk
}

type taskEntry struct {
G
gongweibao 已提交
47 48 49
	Task Task
	// A task fails if it's timeout or trainer reports it exits unnormally.
	NumFailure int
50 51 52 53 54 55
}

type taskQueues struct {
	Todo    []taskEntry
	Pending map[int]taskEntry // map from task ID to task entry
	Done    []taskEntry
G
gongweibao 已提交
56
	Failed  []taskEntry
57 58
}

59 60
// Service is the master server service.
type Service struct {
G
gongweibao 已提交
61 62 63 64 65
	chunksPerTask int
	timeoutDur    time.Duration
	failureMax    int
	ready         chan struct{}
	store         Store
66 67

	mu         sync.Mutex
H
Helin Wang 已提交
68
	initDone   bool
69 70 71
	taskQueues taskQueues
}

H
Helin Wang 已提交
72
func partition(chunks []Chunk, chunksPerTask int) []taskEntry {
73
	id := 0
H
Helin Wang 已提交
74 75
	if chunksPerTask <= 0 {
		chunksPerTask = 1
76 77 78 79 80
	}

	var result []taskEntry
	var cur taskEntry
	for i, c := range chunks {
H
Helin Wang 已提交
81
		if i%chunksPerTask == 0 && len(cur.Task.Chunks) > 0 {
G
gongweibao 已提交
82
			cur.Task.Meta.ID = id
83 84 85 86 87 88 89 90 91
			id++
			result = append(result, cur)
			cur.Task.Chunks = nil
		}

		cur.Task.Chunks = append(cur.Task.Chunks, c)
	}

	if len(cur.Task.Chunks) > 0 {
G
gongweibao 已提交
92
		cur.Task.Meta.ID = id
93 94 95 96 97 98 99
		result = append(result, cur)
	}

	return result
}

// NewService creates a new service.
G
gongweibao 已提交
100
func NewService(store Store, chunksPerTask int, timeoutDur time.Duration, failureMax int) (*Service, error) {
101
	s := &Service{}
102
	s.chunksPerTask = chunksPerTask
103
	s.timeoutDur = timeoutDur
G
gongweibao 已提交
104
	s.failureMax = failureMax
105 106
	s.taskQueues = taskQueues{}
	s.taskQueues.Pending = make(map[int]taskEntry)
107
	s.ready = make(chan struct{})
108 109 110 111 112
	s.store = store
	recovered, err := s.recover()
	if err != nil {
		return nil, err
	}
113

114 115 116 117 118
	if recovered {
		// Recovered. Now the state is already initialized,
		// and the master is ready.
		s.initDone = true
		close(s.ready)
119
		log.Info("Master recovered from saved state.")
120
	}
121

122
	return s, nil
123 124
}

125 126 127 128 129 130
// recover recovers service state from etcd.
func (s *Service) recover() (bool, error) {
	state, err := s.store.Load()
	if err != nil {
		return false, err
	}
131

132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
	if state == nil {
		log.Infoln("No state exists, not recovered.")
		return false, nil
	}

	log.Infof("Loaded snapshot of size: %d bytes.", len(state))
	gr, err := gzip.NewReader(bytes.NewReader(state))
	if err != nil {
		return false, err
	}

	dec := gob.NewDecoder(gr)
	var tqs taskQueues
	err = dec.Decode(&tqs)
	if err != nil {
		return false, err
	}

	err = gr.Close()
	if err != nil {
		// Only close failed, recover actually succeed, so
		// just log error.
		log.Errorln(err)
	}

	s.taskQueues = tqs
	return true, nil
159 160
}

161
// snapshot *must* be called with s.mu being held.
162
func (s *Service) snapshot() error {
163
	// TODO(helin): etcd request has a size limit, so the snapshot
164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
	// size is limited by the max request size. We should either
	// divide the snapshot into smaller chunks and save under
	// different keys, or configure the request size to be big
	// enough:
	// https://github.com/coreos/etcd/blob/2f84f3d8d8ed8f9537ab6ffa44a3a1c7eddfa9b1/embed/config.go#L44
	var buf bytes.Buffer
	gw := gzip.NewWriter(&buf)
	enc := gob.NewEncoder(gw)
	err := enc.Encode(s.taskQueues)
	if err != nil {
		return err
	}
	err = gw.Close()
	if err != nil {
		return err
	}

	state := buf.Bytes()
	log.Infof("Saving snapshot of size: %d bytes.", len(state))
	return s.store.Save(state)
184 185
}

H
Helin Wang 已提交
186
func readChunks(globPaths []string) ([]Chunk, error) {
187 188 189 190 191 192
	var chunks []Chunk
	var paths []string

	for _, s := range globPaths {
		match, err := filepath.Glob(s)
		if err != nil {
H
Helin Wang 已提交
193
			return nil, err
194 195 196 197 198
		}
		paths = append(paths, match...)
	}

	if len(paths) == 0 {
H
Helin Wang 已提交
199
		return nil, errors.New("no valid dataset specified")
200 201 202 203 204
	}

	for _, path := range paths {
		f, err := os.Open(path)
		if err != nil {
H
Helin Wang 已提交
205
			return nil, err
206 207 208 209
		}

		index, err := recordio.LoadIndex(f)
		if err != nil {
H
Helin Wang 已提交
210
			return nil, err
211 212 213
		}
		err = f.Close()
		if err != nil {
H
Helin Wang 已提交
214
			return nil, err
215 216 217
		}

		count := index.NumChunks()
218
		log.Infof("readChunks: file %s has %d chunks", path, count)
219 220 221 222 223 224 225 226 227
		for i := 0; i < count; i++ {
			chunk := Chunk{
				Path:  path,
				Index: *index.ChunkIndex(i),
			}
			chunks = append(chunks, chunk)
		}
	}

H
Helin Wang 已提交
228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248
	return chunks, nil
}

// SetDataset sets dataset to dispatch for the master server.
//
// SetDataset can be call multiple times. But only the first call will
// be honored.
func (s *Service) SetDataset(globPaths []string, dummy *int) error {
	if len(globPaths) == 0 {
		return errors.New("no dataset specified")
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.initDone {
		// Already initialized. All trainer will call
		// SetDataset, but we only handle the first one. Treat
		// other calls as successful but do nothing.
		return nil
	}

H
Helin Wang 已提交
249
	chunks, err := readChunks(globPaths)
H
Helin Wang 已提交
250 251 252 253
	if err != nil {
		return err
	}

254 255
	s.taskQueues.Todo = partition(chunks, s.chunksPerTask)

H
Helin Wang 已提交
256
	err = s.snapshot()
257
	if err != nil {
H
Helin Wang 已提交
258
		log.Errorln(err)
259 260 261 262
		return err
	}

	close(s.ready)
H
Helin Wang 已提交
263
	s.initDone = true
264 265 266
	return nil
}

G
gongweibao 已提交
267 268
func (s *Service) processFailedTask(t taskEntry, epoch int) {
	if t.Task.Meta.Epoch != epoch {
G
gongweibao 已提交
269 270 271 272 273 274 275 276 277 278 279 280
		// new epoch, task launched after the
		// schedule of this timeout check or failed status report.
		return
	}

	defer func() {
		err := s.snapshot()
		if err != nil {
			log.Errorln(err)
		}
	}()

G
gongweibao 已提交
281
	delete(s.taskQueues.Pending, t.Task.Meta.ID)
G
gongweibao 已提交
282

G
gongweibao 已提交
283 284 285
	t.NumFailure++
	if t.NumFailure > s.failureMax {
		log.Warningf("Task %v failed %d times, discard.", t.Task, t.NumFailure)
G
gongweibao 已提交
286 287 288 289
		s.taskQueues.Failed = append(s.taskQueues.Failed, t)
		return
	}

G
gongweibao 已提交
290
	log.Warningf("Task %v failed %d times, discard.", t.Task, t.NumFailure)
G
gongweibao 已提交
291 292 293
	s.taskQueues.Todo = append(s.taskQueues.Todo, t)
}

H
Helin Wang 已提交
294 295 296 297 298 299 300 301 302 303
func (s *Service) checkTimeoutFunc(taskID int, epoch int) func() {
	return func() {
		s.mu.Lock()
		defer s.mu.Unlock()

		t, ok := s.taskQueues.Pending[taskID]
		if !ok {
			return
		}

G
gongweibao 已提交
304
		s.processFailedTask(t, epoch)
H
Helin Wang 已提交
305 306 307
	}
}

H
Helin Wang 已提交
308 309 310 311 312 313 314 315 316 317
// must be called with lock held.
func (s *Service) logFields() log.Fields {
	return log.Fields{
		"todoLen":    len(s.taskQueues.Todo),
		"pendingLen": len(s.taskQueues.Pending),
		"doneLen":    len(s.taskQueues.Done),
		"failedLen":  len(s.taskQueues.Failed),
	}
}

318 319
// GetTask gets a new task from the service.
func (s *Service) GetTask(dummy int, task *Task) error {
320 321 322 323
	select {
	case <-s.ready:
	}

324 325 326 327
	s.mu.Lock()
	defer s.mu.Unlock()

	if len(s.taskQueues.Todo) == 0 {
328 329
		if len(s.taskQueues.Done) == 0 {
			if len(s.taskQueues.Pending) == 0 {
H
Helin Wang 已提交
330
				err := errors.New("all task failed")
H
Helin Wang 已提交
331
				log.WithFields(s.logFields()).Warningln("All tasks failed.")
H
Helin Wang 已提交
332
				return err
333 334 335 336 337
			}

			// TODO(helin): client need to retry in this
			// error case. Gotcha: RPC client can't
			// compare returned error with predefined
H
Helin Wang 已提交
338 339 340 341 342
			// errors like io.EOF, because the error
			// instance deserialized from RPC is a
			// different instance than the error defined
			// in package. So we need to figure out a way
			// for client to check this error correctly.
H
Helin Wang 已提交
343
			err := errors.New("no more available task")
H
Helin Wang 已提交
344
			log.WithFields(s.logFields()).Warningln("No more available task.")
H
Helin Wang 已提交
345
			return err
346 347
		}
		s.taskQueues.Todo = s.taskQueues.Done
H
Helin Wang 已提交
348
		s.taskQueues.Done = nil
H
Helin Wang 已提交
349
		log.WithFields(s.logFields()).Infoln("No more todo task, but trainer is requesting task to do. Move all done task to todo.")
350 351 352
	}

	t := s.taskQueues.Todo[0]
G
gongweibao 已提交
353
	t.Task.Meta.Epoch++
354
	s.taskQueues.Todo = s.taskQueues.Todo[1:]
G
gongweibao 已提交
355
	s.taskQueues.Pending[t.Task.Meta.ID] = t
356 357 358 359 360
	err := s.snapshot()
	if err != nil {
		return err
	}

361
	*task = t.Task
G
gongweibao 已提交
362
	log.WithFields(s.logFields()).Infof("Task #%v dispatched.", t.Task.Meta)
363

G
gongweibao 已提交
364
	time.AfterFunc(s.timeoutDur, s.checkTimeoutFunc(t.Task.Meta.ID, t.Task.Meta.Epoch))
365 366 367 368 369
	return nil
}

// TaskFinished tell the service that a task is finished.
func (s *Service) TaskFinished(taskID int, dummy *int) error {
370 371 372 373
	select {
	case <-s.ready:
	}

374 375 376 377 378
	s.mu.Lock()
	defer s.mu.Unlock()

	t, ok := s.taskQueues.Pending[taskID]
	if !ok {
H
Helin Wang 已提交
379
		log.WithFields(s.logFields()).Warningln("Pending task #%d not found.", taskID)
G
gongweibao 已提交
380
		return nil
381 382 383
	}

	// task finished, reset timeout
G
gongweibao 已提交
384
	t.NumFailure = 0
385 386
	s.taskQueues.Done = append(s.taskQueues.Done, t)
	delete(s.taskQueues.Pending, taskID)
387

H
Helin Wang 已提交
388 389
	log.WithFields(s.logFields()).Infof("Task #%d finished.", taskID)

H
Helin Wang 已提交
390
	if len(s.taskQueues.Pending) == 0 && len(s.taskQueues.Todo) == 0 {
H
Helin Wang 已提交
391
		log.WithFields(s.logFields()).Infoln("No more todo and pending task, start a new pass.")
392
		s.taskQueues.Todo = append(s.taskQueues.Todo, s.taskQueues.Done...)
393 394 395
		s.taskQueues.Done = nil
	}

H
Helin Wang 已提交
396 397 398 399 400
	err := s.snapshot()
	if err != nil {
		log.Errorln(err)
	}
	return err
401
}
G
gongweibao 已提交
402

G
gongweibao 已提交
403
// TaskFailed tells the service that a task is failed.
G
gongweibao 已提交
404
func (s *Service) TaskFailed(meta TaskMeta, dummy *int) error {
G
gongweibao 已提交
405 406 407 408 409 410 411
	select {
	case <-s.ready:
	}

	s.mu.Lock()
	defer s.mu.Unlock()

G
gongweibao 已提交
412
	t, ok := s.taskQueues.Pending[meta.ID]
G
gongweibao 已提交
413
	if !ok {
G
gongweibao 已提交
414
		log.WithFields(s.logFields()).Warningln("TaskFailed:Pending task #%v not found.", t.Task.Meta)
G
gongweibao 已提交
415
		return nil
G
gongweibao 已提交
416 417
	}

G
gongweibao 已提交
418
	s.processFailedTask(t, meta.Epoch)
G
gongweibao 已提交
419 420
	return nil
}