client.go 3.1 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
var nullPtr = unsafe.Pointer(uintptr(0))
H
Helin Wang 已提交
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 53
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
}

54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
//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 已提交
73 74
}

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

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

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

	return C.PADDLE_MASTER_OK
}

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

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

127 128 129 130 131 132 133 134 135 136 137 138
	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 已提交
139 140 141
}

func main() {}