Pserver Save state

go/pserver/optimizer.go

@@ -40,17 +40,23 @@ func newOptimizer(paramWithConfigs ParameterWithConfig) *optimizer {
 	o.elementType = paramWithConfigs.Param.ElementType
 	p := paramWithConfigs.Param
 	c := paramWithConfigs.Config
+	s := paramWithConfigs.State
 	log.WithFields(log.Fields{
 		"ElementType": p.ElementType,
 		"ParamSize":   len(p.Content),
 		"ConfigSize":  len(c),
+		"StateSize":   len(s),
 	}).Info("New Optimizer Created with config:")
 	var cbuffer unsafe.Pointer
 	cbuffer = C.malloc(C.size_t(len(p.Content)))
 	C.memcpy(cbuffer, unsafe.Pointer(&p.Content[0]), C.size_t(len(p.Content)))
+	var cstate unsafe.Pointer
+	if len(s) != 0 {
+		cstate = unsafe.Pointer(&s[0])
+	}
+
 	o.opt = C.paddle_create_optimizer((*C.uchar)(&c[0]), C.int(len(c)),
-		C.paddle_element_type(p.ElementType), cbuffer, C.int(len(p.Content)/C.sizeof_float),
-		(*C.char)(nullPtr), 0)
+		C.paddle_element_type(p.ElementType), cbuffer, C.int(len(p.Content)/C.sizeof_float), (*C.char)(cstate), C.int(len(s)))
 	return o
 }
 
@@ -60,6 +66,12 @@ func (o *optimizer) GetWeights() []byte {
 	return cArrayToSlice(buffer, int(buffer_len)*C.sizeof_float)
 }
 
+func (o *optimizer) GetStates() []byte {
+	var cbuffer *C.char
+	cbuffer_len := C.paddle_optimizer_get_state(o.opt, &cbuffer)
+	return cArrayToSlice(unsafe.Pointer(cbuffer), int(cbuffer_len))
+}
+
 func (o *optimizer) UpdateParameter(g Gradient) error {
 	if o.elementType != g.ElementType {
 		return fmt.Errorf("Name: %s, parameter and gradient element type not match, parameter: %v, gradient: %v", g.Name, o.elementType, g.ElementType)

go/pserver/service.go

@@ -1,9 +1,19 @@
 package pserver
 
 import (
+	"bufio"
+	"bytes"
+	"crypto/md5"
+	"encoding/gob"
+	"encoding/hex"
 	"errors"
 	"fmt"
+	"os"
+	"strconv"
 	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
 )
 
 // ElementType is the type of elements of a Parameter.
@@ -14,6 +24,10 @@ const (
 	Uninitialized      = "pserver not fully initialized"
 )
 
+const (
+	checkpoint_path = "./checkpoints/"
+)
+
 // Supported element types
 const (
 	Int32 ElementType = iota
@@ -38,6 +52,7 @@ type Parameter struct {
 type ParameterWithConfig struct {
 	Param  Parameter
 	Config []byte // parameter configuration in Proto Buffer format
+	State  []byte // parameter training state
 }
 
 // Gradient is the gradient of the parameter.
@@ -52,14 +67,34 @@ type Service struct {
 	optMap map[string]*optimizer
 }
 
+type checkpoint struct {
+	Uuid      string
+	Md5sum    string
+	Timestamp string
+}
+
+//serialize ParameterWithConfig to byte stream
+func GetBytes(content ...interface{}) ([]byte, error) {
+
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+	err := encoder.Encode(content)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
 // NewService creates a new service, will bypass etcd registration if no
 // endpoints specified.
 func NewService(idx int) (*Service, error) {
 	s := &Service{
 		idx: idx,
 	}
-  s.optMap = make(map[string]*optimizer)
+	s.optMap = make(map[string]*optimizer)
 	s.initialized = make(chan struct{})
+	gob.Register(ParameterWithConfig{})
+	gob.Register(checkpoint{})
 	return s, nil
 }
 
@@ -142,8 +177,51 @@ func (s *Service) GetParam(name string, parameter *Parameter) error {
 
 // Save tells the parameter server to save parameters.
 func (s *Service) Save(path string, dummy *int) error {
+	//FIXME: checkpoint is only used by pserver
+	// and has a constant path of */checkpoints/{pserver_idx}*
 	<-s.initialized
-
-	// TODO
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	var paramWithConfig ParameterWithConfig
+	for name, opt := range s.optMap {
+		paramWithConfig.Param.Name = name
+		paramWithConfig.Param.ElementType = opt.elementType
+		paramWithConfig.Param.Content = opt.GetWeights()
+		paramWithConfig.State = opt.GetStates()
+		content, err := GetBytes(paramWithConfig)
+		if err != nil {
+			log.Errorln(err)
+		}
+		ck := checkpoint{}
+		h := md5.New()
+		ck.Md5sum = hex.EncodeToString(h.Sum(content))
+		ck.Timestamp = time.Now().String()
+		ck.Uuid = checkpoint_path + strconv.Itoa(s.idx)
+		ckbytes, err := GetBytes(ck)
+		if err != nil {
+			log.Errorln(err)
+		}
+		// TODO: according design doc, need to save Uuid to etcd in json format
+		// {\"Uuid\": [UUID], \"md5\", \"MD5 sum\", \"Timestamp\": xxxx}
+		log.Infof("parameter checkpoint %s", ckbytes)
+
+		if _, err = os.Stat(ck.Uuid); os.IsNotExist(err) {
+			log.Info("checkpoint not exists.")
+		} else {
+			err = os.Remove(ck.Uuid)
+			log.Infof("remove %s", ck.Uuid)
+		}
+		f, err := os.Create(ck.Uuid)
+		defer f.Close()
+		if err != nil {
+			log.Errorln(err)
+		}
+		writer := bufio.NewWriter(f)
+		_, err = writer.Write(content)
+		writer.Flush()
+		if err != nil {
+			log.Errorln(err)
+		}
+	}
 	return nil
 }

go/pserver/service_test.go

@@ -79,6 +79,8 @@ func TestServiceFull(t *testing.T) {
 	if !reflect.DeepEqual(param1, p) {
 		t.FailNow()
 	}
+	var dummy int
+	s.Save("", &dummy)
 }
 
 func TestMultipleInit(t *testing.T) {
@@ -166,3 +168,7 @@ func TestBlockUntilInitialized(t *testing.T) {
 
 	wg.Wait()
 }
+
+func TestCheckpointSpeed(t *testing.T) {
+	//TODO: test speed
+}