Trainer library discover master by etcd

go/.gitignore

@@ -0,0 +1,4 @@
+CMakeCache.txt
+cmake_install.cmake
+CMakeFiles
+go

go/master/c/client.go

@@ -54,6 +54,13 @@ func (a addresser) Address() string {
 	return string(a)
 }
 
+//export paddle_new_etcd_master_client
+func paddle_new_etcd_master_client(etcdEndpoints *C.char, bufSize int) C.paddle_master_client {
+	p := C.GoString(etcdEndpoints)
+	c := master.NewEtcdClient(addresser(p), bufSize)
+	return add(c)
+}
+
 //export paddle_new_master_client
 func paddle_new_master_client(addr *C.char, bufSize int) C.paddle_master_client {
 	a := C.GoString(addr)

go/master/client.go

@@ -1,14 +1,19 @@
 package master
 
 import (
+	"context"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/PaddlePaddle/Paddle/go/connection"
 	"github.com/PaddlePaddle/recordio"
+	"github.com/coreos/etcd/clientv3"
 	log "github.com/sirupsen/logrus"
 )
 
+const masterAddrPath = "/master"
+
 // Addresser provide the address of the master server.
 type Addresser interface {
 	Address() string
@@ -20,6 +25,51 @@ type Client struct {
 	ch   chan []byte
 }
 
+// MasterAddresser provide master address
+type masterAddresser struct {
+	client    *clientv3.Client
+	timeout   time.Duration
+	endpoints []string
+}
+
+// Address return the address
+func (m masterAddresser) Address() string {
+	for {
+		ctx, cancel := context.WithTimeout(context.Background(), m.timeout)
+		resp, err := m.client.Get(ctx, masterAddrPath)
+		cancel()
+		if err != nil {
+			log.Errorf("Fetch master addr failed, reconnecting to etcd, %v", err)
+			err := m.client.Close()
+			if err != nil {
+				log.Errorln(err)
+				time.Sleep(m.timeout)
+				continue
+			}
+			// reconnect to etcd server
+			m.client, err = clientv3.New(clientv3.Config{
+				Endpoints:   m.endpoints,
+				DialTimeout: m.timeout,
+			})
+			if err != nil {
+				log.Errorf("Reconnecting etcd failed, sleep for %d seconds ...\n%v", m.timeout, err)
+				time.Sleep(m.timeout)
+				continue
+			}
+			continue
+		}
+		kvs := resp.Kvs
+		if len(kvs) == 0 {
+			log.Infoln("Waiting for master be ready ...")
+			time.Sleep(m.timeout)
+			continue
+		}
+		mAddr := kvs[0].Value
+		log.Debugf("Fetched master address: %s\n", mAddr)
+		return string(mAddr)
+	}
+}
+
 // NewClient creates a new Client.
 //
 // bufSize is the record buffer size. NextRecord will read from this
@@ -33,6 +83,32 @@ func NewClient(addr Addresser, bufSize int) *Client {
 	return c
 }
 
+// NewEtcdClient create a new master client by etcd
+//
+// endpoints is the endpoints for etcd and separated by ",", such as
+// "172.0.1.0:2379,172.0.1.1:2379"
+// timeout is the timeout for etcd calls
+// bufSize is the record buffer size. NextRecord will read from this buffer.
+func NewEtcdClient(endpoints string, timeout int, bufSize int) *Client {
+	t := time.Second * time.Duration(timeout)
+	ep := strings.Split(endpoints, ",")
+	cli, err := clientv3.New(clientv3.Config{
+		Endpoints:   ep,
+		DialTimeout: t,
+	})
+	if err != nil {
+		log.Errorf("Init etcd connection failed: %v", err)
+		panic(err)
+	}
+	log.Debugf("Connected to etcd: %s\n", endpoints)
+	mAddresser := masterAddresser{
+		client:    cli,
+		timeout:   t,
+		endpoints: ep,
+	}
+	return NewClient(mAddresser, bufSize)
+}
+
 func (c *Client) getRecords() {
 	for {
 		t, err := c.getTask()

go/master/client_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/PaddlePaddle/Paddle/go/master"
 	"github.com/PaddlePaddle/recordio"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestNextRecord(t *testing.T) {
@@ -77,3 +78,8 @@ func TestNextRecord(t *testing.T) {
 		}
 	}
 }
+
+func TestNewEtcdClientFailed(t *testing.T) {
+	assert.Panics(t, func() { master.NewEtcdClient("localhost:1235", 3, 1) },
+		"Invalid etcd address should be panic.")
+}