implement task handling for master server's service

doc/design/cluster_train/master_server.md

@@ -10,7 +10,7 @@ A dataset is a list of files in *RecordIO* format. A RecordIO file consists of c
 
 ## Task Queue
 
-As mentioned in [distributed training design doc](./README.md), a *task* is a data shard that the master server assigns to the trainer process to train on. A task consists of one or multiple *blocks* from one or multiple files. The master server maintains *task queues* to track the training progress.
+As mentioned in [distributed training design doc](./README.md), a *task* is a data shard that the master server assigns to the trainer process to train on. A task consists of one or multiple *chunks* from one or multiple files. The master server maintains *task queues* to track the training progress.
 
 ### Task Queue Creation
 
@@ -21,23 +21,23 @@ As mentioned in [distributed training design doc](./README.md), a *task* is a da
 	func (m *RPCServer) ReportDataset(Paths []string, dummy *int) error {
 	}
 	```
-1. The master server will scan through each RecordIO file to generate the *block index* and know how many blocks does each file have. A block can be referenced by the file path and the index of the block within the file. The block index is in memory data structure that enables fast access to each block, and the index of the block with the file is an integer start from 0, representing the n-th block within the file.
+1. The master server will scan through each RecordIO file to generate the *chunk index* and know how many chunks does each file have. A chunk can be referenced by the file path and the index of the chunk within the file. The chunk index is in memory data structure that enables fast access to each chunk, and the index of the chunk with the file is an integer start from 0, representing the n-th chunk within the file.
 
-	The definition of the block is:
+	The definition of the chunk is:
 	```go
-	type Block struct {
-		Idx   int // index of the block within the file
+	type Chunk struct {
+		Idx   int // index of the chunk within the file
 		Path  string
-		Index recordio.Index // block index
+		Index recordio.Index // chunk index
 	}
 	```
-1. Blocks are grouped into tasks, and tasks are filled into the todo queue. The pending queue and the done queue are initialized with no element.
+1. Chunks are grouped into tasks, and tasks are filled into the todo queue. The pending queue and the done queue are initialized with no element.
 
 	The definition of the task is:
 	```go
 	type Task struct {
 		Index  int
-		Blocks []Block
+		Chunks []Chunk
 	}
 	```
 

paddle/go/master/service.go

+package master
+
+import (
+	"errors"
+	"log"
+	"sync"
+	"time"
+
+	"github.com/wangkuiyi/recordio"
+)
+
+const (
+	targetTaskCount = 300
+)
+
+// errors
+var (
+	ErrNoMoreTask          = errors.New("no more task for current pass")
+	ErrPendingTaskNotFound = errors.New("pending task not found")
+)
+
+// Service is the master server service.
+type Service struct {
+	timeoutDur time.Duration
+	timeoutMax int
+
+	mu         sync.Mutex
+	taskQueues taskQueues
+}
+
+// Recover recovers service state from etcd.
+func Recover() (*Service, error) {
+	// TODO(helin): recover from snapshot state from etcd.
+	return nil, nil
+}
+
+func partition(chunks []Chunk, targetTaskCount int) []taskEntry {
+	id := 0
+	chunkPerTask := len(chunks) / targetTaskCount
+	if chunkPerTask <= 0 {
+		chunkPerTask = 1
+	}
+
+	var result []taskEntry
+	var cur taskEntry
+	for i, c := range chunks {
+		if i%chunkPerTask == 0 && len(cur.Task.Chunks) > 0 {
+			cur.Task.ID = id
+			id++
+			result = append(result, cur)
+			cur.Task.Chunks = nil
+		}
+
+		cur.Task.Chunks = append(cur.Task.Chunks, c)
+	}
+
+	if len(cur.Task.Chunks) > 0 {
+		cur.Task.ID = id
+		id++
+		result = append(result, cur)
+	}
+
+	return result
+}
+
+// NewService creates a new service.
+func NewService(chunks []Chunk, timeoutDur time.Duration, timeoutMax int) (*Service, error) {
+	s := &Service{}
+	s.timeoutDur = timeoutDur
+	s.timeoutMax = timeoutMax
+	s.taskQueues = taskQueues{}
+	s.taskQueues.Pending = make(map[int]taskEntry)
+	s.taskQueues.Todo = partition(chunks, targetTaskCount)
+	return s, nil
+}
+
+// Chunk is a chunk of data consisted of several data instances.
+type Chunk struct {
+	Idx   int // index of the chunk within the file
+	Path  string
+	Index recordio.Index // block index
+}
+
+// Task is the basic unit of data instances assigned to trainers.
+type Task struct {
+	ID     int
+	Chunks []Chunk
+}
+
+type taskEntry struct {
+	Epoch      int
+	NumTimeout int
+	Task       Task
+}
+
+type taskQueues struct {
+	Todo    []taskEntry
+	Pending map[int]taskEntry // map from task ID to task entry
+	Done    []taskEntry
+	Failed  []Task
+}
+
+// *must* be called with s.mu being held.
+func (s *Service) snapshot() error {
+	// TODO(helin): snapshot state on etcd.
+	return nil
+}
+
+// GetTask gets a new task from the service.
+func (s *Service) GetTask(dummy int, task *Task) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(s.taskQueues.Todo) == 0 {
+		return ErrNoMoreTask
+	}
+
+	t := s.taskQueues.Todo[0]
+	t.Epoch++
+	s.taskQueues.Todo = s.taskQueues.Todo[1:]
+	s.taskQueues.Pending[t.Task.ID] = t
+	err := s.snapshot()
+	if err != nil {
+		return err
+	}
+
+	time.AfterFunc(s.timeoutDur, func(taskID int, epoch int) func() {
+		return func() {
+			s.mu.Lock()
+			defer s.mu.Unlock()
+
+			t, ok := s.taskQueues.Pending[taskID]
+			if !ok {
+				return
+			}
+
+			if t.Epoch != epoch {
+				// new epoch, task launched after the
+				// schedule of this timeout check.
+				return
+			}
+
+			defer func() {
+				err := s.snapshot()
+				if err != nil {
+					log.Println(err)
+				}
+			}()
+
+			delete(s.taskQueues.Pending, t.Task.ID)
+
+			t.NumTimeout++
+			if t.NumTimeout > s.timeoutMax {
+				s.taskQueues.Failed = append(s.taskQueues.Failed, t.Task)
+				return
+			}
+
+			s.taskQueues.Todo = append(s.taskQueues.Todo, t)
+		}
+	}(t.Task.ID, t.Epoch))
+	return nil
+}
+
+// TaskFinished tell the service that a task is finished.
+func (s *Service) TaskFinished(taskID int, dummy *int) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	t, ok := s.taskQueues.Pending[taskID]
+	if !ok {
+		return ErrPendingTaskNotFound
+	}
+
+	// task finished, reset timeout
+	t.NumTimeout = 0
+	s.taskQueues.Done = append(s.taskQueues.Done, t)
+	delete(s.taskQueues.Pending, taskID)
+	return s.snapshot()
+}

paddle/go/master/service_internal_test.go

+package master
+
+import "testing"
+
+func TestPartitionCount(t *testing.T) {
+	cs := make([]Chunk, 100)
+	ts := partition(cs, 20)
+	if len(ts) != 20 {
+		t.Error(len(ts))
+	}
+
+	cs = make([]Chunk, 101)
+	ts = partition(cs, 20)
+	if len(ts) != 21 {
+		t.Error(len(ts))
+	}
+
+	ts = partition(cs, 200)
+	if len(ts) != 101 {
+		t.Error(len(ts))
+	}
+}
+
+func TestPartionIndex(t *testing.T) {
+	cs := make([]Chunk, 100)
+	ts := partition(cs, 20)
+	for i := range ts {
+		if ts[i].Task.ID != i {
+			t.Error(ts[i], i)
+		}
+	}
+}