client.go 4.3 KB
Newer Older
D
dongzhihong 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

// http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

15 16 17
package main

/*
18 19 20 21 22 23
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define PADDLE_MASTER_OK    0
#define PADDLE_MASTER_ERROR -1
24

H
Helin Wang 已提交
25 26 27
#define PADDLE_SAVE_MODEL_OK   1
#define PADDLE_SAVE_MODEL_SKIP 0

28 29 30 31 32
typedef int paddle_master_client;
*/
import "C"

import (
33
	"strings"
34
	"sync"
35
	"time"
36 37 38
	"unsafe"

	"github.com/PaddlePaddle/Paddle/go/master"
H
Helin Wang 已提交
39
	log "github.com/sirupsen/logrus"
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68
)

var mu sync.Mutex
var handleMap = make(map[C.paddle_master_client]*master.Client)
var curHandle C.paddle_master_client

func add(c *master.Client) C.paddle_master_client {
	mu.Lock()
	defer mu.Unlock()
	client := curHandle
	curHandle++
	handleMap[client] = c
	return client
}

func get(client C.paddle_master_client) *master.Client {
	mu.Lock()
	defer mu.Unlock()
	return handleMap[client]
}

func remove(client C.paddle_master_client) *master.Client {
	mu.Lock()
	defer mu.Unlock()
	h := handleMap[client]
	delete(handleMap, client)
	return h
}

69
//export paddle_new_etcd_master_client
70 71
//
// bufSize is the record buffer size.
72 73
func paddle_new_etcd_master_client(etcdEndpoints *C.char, timeout int, bufSize int) C.paddle_master_client {
	p := C.GoString(etcdEndpoints)
74 75 76 77 78
	endpoints := strings.Split(p, ",")
	c, err := master.NewClient(
		master.WithEtcd(endpoints, time.Duration(timeout)*time.Second),
		master.WithBuffer(bufSize),
	)
79 80 81
	if err != nil {
		panic(err)
	}
82

83
	return add(c)
84 85
}

86
//export paddle_new_master_client
87 88
//
// bufSize is the record buffer size.
89
func paddle_new_master_client(addr *C.char, bufSize int) C.paddle_master_client {
90
	a := C.GoString(addr)
91 92 93 94 95
	c, err := master.NewClient(master.WithAddr(a), master.WithBuffer(bufSize))
	if err != nil {
		panic(err)
	}

96 97 98
	return add(c)
}

99 100 101
//export paddle_release_master_client
func paddle_release_master_client(client C.paddle_master_client) {
	remove(client)
102 103 104 105 106 107 108
}

//export paddle_set_dataset
func paddle_set_dataset(client C.paddle_master_client, path **C.char, size C.int) C.int {
	c := get(client)
	var paths []string
	for i := 0; i < int(size); i++ {
109
		ptr := (**C.char)(unsafe.Pointer(uintptr(unsafe.Pointer(path)) + uintptr(i)*unsafe.Sizeof(*path)))
110 111 112 113 114
		str := C.GoString(*ptr)
		paths = append(paths, str)
	}
	err := c.SetDataset(paths)
	if err != nil {
H
Helin Wang 已提交
115
		log.Errorln(err)
116 117 118 119 120 121
		return C.PADDLE_MASTER_ERROR
	}

	return C.PADDLE_MASTER_OK
}

122 123 124 125
// paddle_next_record gets the nexts training record.
//
// returns number of bytes of the records if success, -1 if failed.
//
126 127 128
//export paddle_next_record
func paddle_next_record(client C.paddle_master_client, record **C.uchar) C.int {
	c := get(client)
G
gongweibao 已提交
129 130 131 132
	r, err := c.NextRecord()
	if err != nil {
		// Error
		// TODO: return the type of error?
133
		*record = (*C.uchar)(nil)
G
gongweibao 已提交
134
		return -1
G
gongweibao 已提交
135 136
	}

137
	if len(r) == 0 {
G
gongweibao 已提交
138
		// Empty record
139
		*record = (*C.uchar)(nil)
140
		return 0
141 142
	}

143 144 145 146 147 148
	size := C.size_t(len(r))
	*record = (*C.uchar)(C.malloc(size))
	C.memcpy(unsafe.Pointer(*record), unsafe.Pointer(&r[0]), size)
	return C.int(size)
}

149 150 151
// paddle_request_save_model requests the master server to approve the
// caller to save the model.
//
H
Helin Wang 已提交
152
// returns 1 if the save the model request is approved, 0 if the
153 154 155 156 157 158 159 160 161
// request is rejected because other trainer is saving the model, -1
// if error happened.
//
//export paddle_request_save_model
func paddle_request_save_model(client C.paddle_master_client, trainerID string, blockMS int) C.int {
	c := get(client)
	need, err := c.RequestSaveModel(trainerID, time.Duration(blockMS)*time.Millisecond)
	if err != nil {
		log.Errorln(err)
H
Helin Wang 已提交
162
		return C.PADDLE_MASTER_ERROR
163 164 165
	}

	if need {
H
Helin Wang 已提交
166
		return C.PADDLE_SAVE_MODEL_OK
167 168
	}

H
Helin Wang 已提交
169
	return C.PADDLE_SAVE_MODEL_SKIP
170 171
}

172 173 174 175 176 177
//export mem_free
func mem_free(p unsafe.Pointer) {
	// "free" may be a better name for this function, but doing so
	// will cause calling any function of this library from Python
	// ctypes hanging.
	C.free(p)
178 179 180
}

func main() {}
新手
引导
客服 返回
顶部