service.go 3.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
package pserver

import (
	"errors"
	"fmt"
	"sync"
)

// ElementType is the type of elements of a Parameter.
type ElementType int

12 13 14 15
const (
	AlreadyInitialized = "pserver already initialized"
	Uninitialized      = "pserver not fully initialized"
)
16 17 18 19 20 21 22 23 24 25 26 27 28 29 30

// Supported element types
const (
	Int32 ElementType = iota
	UInt32
	Int64
	UInt64
	Float32
	Float64
)

// Parameter is a piece of data to sync with the parameter server.
type Parameter struct {
	Name        string
	ElementType ElementType
D
dzhwinter 已提交
31
	Content     []byte
32 33 34 35 36 37 38 39 40 41 42
}

// ParameterWithConfig contains the parameter and the configuration.
type ParameterWithConfig struct {
	Param  Parameter
	Config []byte // parameter configuration in Proto Buffer format
}

// Gradient is the gradient of the parameter.
type Gradient Parameter

H
Helin Wang 已提交
43
// Service is the RPC service for pserver.
44 45 46
type Service struct {
	initialized chan struct{}

D
dzhwinter 已提交
47 48
	mu sync.Mutex
	// injection from parameter to optimizer
D
dzhwinter 已提交
49
	optMap map[string]*optimizer
50 51
}

H
Helin Wang 已提交
52
// NewService creates a new service.
53
func NewService() *Service {
D
dzhwinter 已提交
54 55
	s := &Service{}
	s.optMap = make(map[string]*optimizer)
56 57 58 59
	s.initialized = make(chan struct{})
	return s
}

H
Helin Wang 已提交
60
// InitParam initializes a parameter.
61 62 63
func (s *Service) InitParam(paramWithConfigs ParameterWithConfig, dummy *int) error {
	select {
	case <-s.initialized:
64
		return errors.New(AlreadyInitialized)
65 66 67 68 69 70 71 72 73 74 75
	default:
	}

	// TODO(helin): parse parameter config

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

	// TODO(helin): check if paramWithConfigs.Param.Content is
	// properly memory aligned, if not, make copy to a memory
	// aligned region.
D
dzhwinter 已提交
76
	s.optMap[paramWithConfigs.Param.Name] = newOptimizer(paramWithConfigs)
77 78 79
	return nil
}

H
Helin Wang 已提交
80 81
// FinishInitParams tells the parameter server that the parameter
// initialization has finished.
82 83 84
func (s *Service) FinishInitParams(dummy0 int, dummy1 *int) error {
	select {
	case <-s.initialized:
85
		return errors.New(AlreadyInitialized)
86 87 88 89 90 91 92
	default:
	}

	close(s.initialized)
	return nil
}

93
// SendGrad sends gradient to parameter servers for parameter
H
Helin Wang 已提交
94
// optimization.
95
func (s *Service) SendGrad(g Gradient, dummy *int) error {
96 97 98
	select {
	case <-s.initialized:
	default:
99
		return errors.New(Uninitialized)
100
	}
101 102

	s.mu.Lock()
H
Helin Wang 已提交
103
	defer s.mu.Unlock()
104

D
dzhwinter 已提交
105 106
	o, ok := s.optMap[g.Name]
	if !ok {
D
dzhwinter 已提交
107
		return fmt.Errorf("parameter: %s does not exist", g.Name)
D
dzhwinter 已提交
108
	}
109

D
dzhwinter 已提交
110
	return o.UpdateParameter(p, g)
111 112
}

113 114
// GetParam gets parameters from the parameter server.
func (s *Service) GetParam(name string, parameter *Parameter) error {
115 116
	<-s.initialized
	s.mu.Lock()
H
Helin Wang 已提交
117
	defer s.mu.Unlock()
118

119 120 121
	p, ok := s.paramMap[name]
	if !ok {
		return fmt.Errorf("parameter: %s does not exist", name)
122 123
	}

124 125 126 127 128 129 130 131
	// The parameter content (a byte slice) may change
	// during RPC serialization due to write from other
	// goroutine, we allow it since mini-batch based deep
	// learning optimization methods are stochastic in
	// nature. This race condition is allowed deliberately
	// to save the program from making a copy of the
	// paramter content.
	*parameter = p
132 133 134
	return nil
}

H
Helin Wang 已提交
135 136
// Save tells the parameter server to save parameters.
func (s *Service) Save(path string, dummy *int) error {
137 138 139 140 141
	<-s.initialized

	// TODO
	return nil
}