client.go 3.0 KB
Newer Older
H
Helin Wang 已提交
1 2 3
package main

/*
4 5 6 7 8 9
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define PADDLE_MASTER_OK    0
#define PADDLE_MASTER_ERROR -1
H
Helin Wang 已提交
10 11 12 13 14 15

typedef int paddle_master_client;
*/
import "C"

import (
16
	"strings"
H
Helin Wang 已提交
17
	"sync"
18
	"time"
H
Helin Wang 已提交
19 20 21
	"unsafe"

	"github.com/PaddlePaddle/Paddle/go/master"
22
	"github.com/coreos/etcd/clientv3"
H
Helin Wang 已提交
23
	log "github.com/sirupsen/logrus"
H
Helin Wang 已提交
24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
)

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
}

53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
//export paddle_new_etcd_master_client
func paddle_new_etcd_master_client(etcdEndpoints *C.char, timeout int, bufSize int) C.paddle_master_client {
	p := C.GoString(etcdEndpoints)
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   strings.Split(p, ","),
		DialTimeout: time.Second * time.Duration(timeout),
	})
	if err != nil {
		panic(err)
	}
	ch := make(chan string, 1)
	a, err := master.GetKey(cli, master.DefaultAddrPath, timeout)
	if err != nil {
		panic(err)
	}
	ch <- a
	go master.WatchKey(cli, master.DefaultAddrPath, ch)
	c := master.NewClient(ch, bufSize)
	return add(c)
H
Helin Wang 已提交
72 73
}

74
//export paddle_new_master_client
75
func paddle_new_master_client(addr *C.char, bufSize int) C.paddle_master_client {
H
Helin Wang 已提交
76
	a := C.GoString(addr)
77 78 79
	ch := make(chan string, 1)
	ch <- a
	c := master.NewClient(ch, bufSize)
H
Helin Wang 已提交
80 81 82
	return add(c)
}

83 84 85
//export paddle_release_master_client
func paddle_release_master_client(client C.paddle_master_client) {
	remove(client)
H
Helin Wang 已提交
86 87 88 89 90 91 92
}

//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++ {
93
		ptr := (**C.char)(unsafe.Pointer(uintptr(unsafe.Pointer(path)) + uintptr(i)*unsafe.Sizeof(*path)))
H
Helin Wang 已提交
94 95 96 97 98
		str := C.GoString(*ptr)
		paths = append(paths, str)
	}
	err := c.SetDataset(paths)
	if err != nil {
H
Helin Wang 已提交
99
		log.Errorln(err)
100 101 102 103 104 105
		return C.PADDLE_MASTER_ERROR
	}

	return C.PADDLE_MASTER_OK
}

G
gongweibao 已提交
106 107
// return value:
//     0:ok
G
gongweibao 已提交
108
//    -1:error
109 110 111
//export paddle_next_record
func paddle_next_record(client C.paddle_master_client, record **C.uchar) C.int {
	c := get(client)
G
gongweibao 已提交
112 113 114 115
	r, err := c.NextRecord()
	if err != nil {
		// Error
		// TODO: return the type of error?
116
		*record = (*C.uchar)(nil)
G
gongweibao 已提交
117
		return -1
G
gongweibao 已提交
118 119
	}

120
	if len(r) == 0 {
G
gongweibao 已提交
121
		// Empty record
122
		*record = (*C.uchar)(nil)
123
		return 0
H
Helin Wang 已提交
124 125
	}

126 127 128 129 130 131 132 133 134 135 136 137
	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)
}

//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)
H
Helin Wang 已提交
138 139 140
}

func main() {}