client.cpp 8.37 KB
Newer Older
Matthieu Dorier's avatar
Matthieu Dorier committed
1 2 3 4 5
/*
 * (C) 2018 The University of Chicago
 * 
 * See COPYRIGHT in top-level directory.
 */
Matthieu Dorier's avatar
Matthieu Dorier committed
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
#define BOOST_NO_AUTO_PTR
#include <boost/python.hpp>
#include <boost/python/return_opaque_pointer.hpp>
#include <boost/python/handle.hpp>
#include <boost/python/enum.hpp>
#include <boost/python/def.hpp>
#include <boost/python/module.hpp>
#include <boost/python/return_value_policy.hpp>
#include <string>
#include <vector>
#include <cstring>
#include <iostream>
#include <margo.h>
#include <bake.h>
#include <bake-client.h>
Matthieu Dorier's avatar
Matthieu Dorier committed
21 22 23 24
#if HAS_NUMPY
#include <boost/python/numpy.hpp>
namespace np = boost::python::numpy;
#endif
Matthieu Dorier's avatar
Matthieu Dorier committed
25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41

BOOST_PYTHON_OPAQUE_SPECIALIZED_TYPE_ID(margo_instance)
BOOST_PYTHON_OPAQUE_SPECIALIZED_TYPE_ID(bake_provider_handle)
BOOST_PYTHON_OPAQUE_SPECIALIZED_TYPE_ID(bake_client)
BOOST_PYTHON_OPAQUE_SPECIALIZED_TYPE_ID(hg_addr)

namespace bpl = boost::python;

static bake_client_t pybake_client_init(margo_instance_id mid) {
    bake_client_t result = BAKE_CLIENT_NULL;
    bake_client_init(mid, &result);
    return result;
}

static bake_provider_handle_t pybake_provider_handle_create(
        bake_client_t client,
        hg_addr_t addr,
42
        uint8_t provider_id) {
Matthieu Dorier's avatar
Matthieu Dorier committed
43 44

    bake_provider_handle_t providerHandle = BAKE_PROVIDER_HANDLE_NULL;
45
    bake_provider_handle_create(client, addr, provider_id, &providerHandle);
Matthieu Dorier's avatar
Matthieu Dorier committed
46 47 48
    return providerHandle;
}

49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
static uint64_t pybake_get_eager_limit(
        bake_provider_handle_t ph)
{
    uint64_t limit;
    int ret = bake_provider_handle_get_eager_limit(ph, &limit);
    if(ret != 0) return 0;
    return limit;
}

static bpl::object pybake_probe(
        bake_provider_handle_t ph,
        uint64_t max_targets)
{
    bpl::list result;
    std::vector<bake_target_id_t> targets(max_targets);
    uint64_t num_targets;
Matthieu Dorier's avatar
Matthieu Dorier committed
65 66 67 68
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_probe(ph, max_targets, targets.data(), &num_targets);
    Py_END_ALLOW_THREADS
69 70 71 72 73 74 75 76 77 78 79 80 81 82
    if(ret != 0) return bpl::object();
    for(uint64_t i=0; i < num_targets; i++) {
        result.append(bpl::object(targets[i]));
    }
    return result;
}

static bpl::object pybake_create(
        bake_provider_handle_t ph,
        bake_target_id_t bti,
        size_t region_size)
{
    bake_region_id_t rid;
    std::memset(&rid, 0, sizeof(rid));
Matthieu Dorier's avatar
Matthieu Dorier committed
83 84 85 86
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_create(ph, bti, region_size, &rid);
    Py_END_ALLOW_THREADS
87 88 89 90 91 92 93 94 95 96
    if(ret != 0) return bpl::object();
    else return bpl::object(rid);
}

static bpl::object pybake_write(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid,
        uint64_t offset,
        const std::string& data)
{
Matthieu Dorier's avatar
Matthieu Dorier committed
97 98 99 100
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_write(ph, rid, offset, (const void*)data.data(), data.size());
    Py_END_ALLOW_THREADS
101 102 103 104
    if(ret == 0) return bpl::object(true);
    else return bpl::object(false);
}

Matthieu Dorier's avatar
Matthieu Dorier committed
105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120
#if HAS_NUMPY
static bpl::object pybake_write_numpy(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid,
        uint64_t offset,
        const np::ndarray& data)
{
    if(!(data.get_flags() & np::ndarray::bitflag::V_CONTIGUOUS)) {
        std::cerr << "[pyBAKE error]: non-contiguous numpy arrays not yet supported" << std::endl;
        return bpl::object(false);
    }
    size_t size = data.get_dtype().get_itemsize();
    for(int i = 0; i < data.get_nd(); i++) {
        size *= data.shape(i);
    }
    void* buffer = data.get_data();
Matthieu Dorier's avatar
Matthieu Dorier committed
121 122 123 124
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_write(ph, rid, offset, buffer, size);
    Py_END_ALLOW_THREADS
Matthieu Dorier's avatar
Matthieu Dorier committed
125 126 127 128 129
    if(ret != 0) return bpl::object(false);
    else return bpl::object(true);
}
#endif

130 131 132 133
static bpl::object pybake_persist(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid)
{
Matthieu Dorier's avatar
Matthieu Dorier committed
134 135 136 137
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_persist(ph, rid);
    Py_END_ALLOW_THREADS
138 139 140 141 142 143 144 145 146 147
    if(ret == 0) return bpl::object(true);
    else return bpl::object(false);
}

static bpl::object pybake_create_write_persist(
        bake_provider_handle_t ph,
        bake_target_id_t tid,
        const std::string& data)
{
    bake_region_id_t rid;
Matthieu Dorier's avatar
Matthieu Dorier committed
148 149 150
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_create_write_persist(ph, tid, 
151
            data.data(), data.size(), &rid);
Matthieu Dorier's avatar
Matthieu Dorier committed
152
    Py_END_ALLOW_THREADS
153 154 155 156
    if(ret == 0) return bpl::object(rid);
    else return bpl::object();
}

Matthieu Dorier's avatar
Matthieu Dorier committed
157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
#if HAS_NUMPY
static bpl::object pybake_create_write_persist_numpy(
        bake_provider_handle_t ph,
        bake_target_id_t tid,
        const np::ndarray& data)
{
    bake_region_id_t rid;
    if(!(data.get_flags() & np::ndarray::bitflag::V_CONTIGUOUS)) {
        std::cerr << "[pyBAKE error]: non-contiguous numpy arrays not yet supported" << std::endl;
        return bpl::object();
    }
    size_t size = data.get_dtype().get_itemsize();
    for(int i = 0; i < data.get_nd(); i++) {
        size *= data.shape(i);
    }
    void* buffer = data.get_data();
Matthieu Dorier's avatar
Matthieu Dorier committed
173 174 175
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_create_write_persist(ph, tid, 
Matthieu Dorier's avatar
Matthieu Dorier committed
176
            buffer, size, &rid);
Matthieu Dorier's avatar
Matthieu Dorier committed
177
    Py_END_ALLOW_THREADS
Matthieu Dorier's avatar
Matthieu Dorier committed
178 179 180 181 182
    if(ret == 0) return bpl::object(rid);
    else return bpl::object();
}
#endif

183 184 185 186 187
static bpl::object pybake_get_size(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid)
{
    uint64_t size;
Matthieu Dorier's avatar
Matthieu Dorier committed
188 189 190 191
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_get_size(ph, rid, &size);
    Py_END_ALLOW_THREADS
192 193 194 195 196 197 198 199 200 201 202 203
    if(ret == 0) return bpl::object(size);
    else return bpl::object();
}

static bpl::object pybake_read(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid,
        uint64_t offset,
        size_t size) 
{
    std::string result(size, '\0');
    uint64_t bytes_read;
Matthieu Dorier's avatar
Matthieu Dorier committed
204 205 206 207
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_read(ph, rid, offset, (void*)result.data(), size, &bytes_read);
    Py_END_ALLOW_THREADS
208 209 210 211 212
    if(ret != 0) return bpl::object();
    result.resize(bytes_read);
    return bpl::object(result);
}

213 214 215 216 217 218 219 220
static bpl::object pybake_migrate(
        bake_provider_handle_t source_ph,
        const bake_region_id_t& source_rid,
        bool remove_source,
        const std::string& dest_addr,
        uint16_t dest_provider_id,
        bake_target_id_t dest_target_id) {
    bake_region_id_t dest_rid;
Matthieu Dorier's avatar
Matthieu Dorier committed
221 222 223
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_migrate(source_ph, source_rid,
224 225
            remove_source, dest_addr.c_str(), dest_provider_id,
            dest_target_id, &dest_rid);
Matthieu Dorier's avatar
Matthieu Dorier committed
226
    Py_END_ALLOW_THREADS
227 228 229 230
    if(ret != BAKE_SUCCESS) return bpl::object();
    return bpl::object(dest_rid);
}

Matthieu Dorier's avatar
Matthieu Dorier committed
231 232 233 234 235 236 237 238 239 240 241 242 243
#if HAS_NUMPY
static bpl::object pybake_read_numpy(
        bake_provider_handle_t ph,
        const bake_region_id_t& rid,
        uint64_t offset,
        const bpl::tuple& shape,
        const np::dtype& dtype)
{
    np::ndarray result = np::empty(shape, dtype);
    size_t size = dtype.get_itemsize();
    for(int i=0; i < result.get_nd(); i++) 
        size *= result.shape(i);
    uint64_t bytes_read;
Matthieu Dorier's avatar
Matthieu Dorier committed
244 245 246 247
    int ret;
    Py_BEGIN_ALLOW_THREADS
    ret = bake_read(ph, rid, offset, (void*)result.get_data(), size, &bytes_read);
    Py_END_ALLOW_THREADS
Matthieu Dorier's avatar
Matthieu Dorier committed
248 249 250 251 252 253
    if(ret != 0) return bpl::object();
    if(bytes_read != size) return bpl::object();
    else return result;
}
#endif

Matthieu Dorier's avatar
Matthieu Dorier committed
254 255 256
BOOST_PYTHON_MODULE(_pybakeclient)
{
#define ret_policy_opaque bpl::return_value_policy<bpl::return_opaque_pointer>()
Matthieu Dorier's avatar
Matthieu Dorier committed
257 258 259
#if HAS_NUMPY
    np::initialize();
#endif
260
    bpl::import("_pybaketarget");
Matthieu Dorier's avatar
Matthieu Dorier committed
261 262 263 264 265 266 267
    bpl::opaque<bake_client>();
    bpl::opaque<bake_provider_handle>();
    bpl::def("client_init", &pybake_client_init, ret_policy_opaque);
    bpl::def("client_finalize", &bake_client_finalize);
    bpl::def("provider_handle_create", &pybake_provider_handle_create, ret_policy_opaque);
    bpl::def("provider_handle_ref_incr", &bake_provider_handle_ref_incr);
    bpl::def("provider_handle_release", &bake_provider_handle_release);
268 269 270 271 272 273 274 275 276 277
    bpl::def("get_eager_limit", &pybake_get_eager_limit);
    bpl::def("set_eager_limit", &bake_provider_handle_set_eager_limit);
    bpl::def("probe", &pybake_probe);
    bpl::def("create", &pybake_create);
    bpl::def("write", &pybake_write);
    bpl::def("persist", &pybake_persist);
    bpl::def("create_write_persist", &pybake_create_write_persist);
    bpl::def("get_size", &pybake_get_size);
    bpl::def("read", &pybake_read);
    bpl::def("remove", &bake_remove);
278
    bpl::def("migrate", &pybake_migrate);
Matthieu Dorier's avatar
Matthieu Dorier committed
279
    bpl::def("shutdown_service", &bake_shutdown_service);
Matthieu Dorier's avatar
Matthieu Dorier committed
280 281 282 283 284
#if HAS_NUMPY
    bpl::def("write_numpy", &pybake_write_numpy);
    bpl::def("create_write_persist_numpy", &pybake_create_write_persist_numpy);
    bpl::def("read_numpy", &pybake_read_numpy);
#endif
Matthieu Dorier's avatar
Matthieu Dorier committed
285 286 287

#undef ret_policy_opaque
}