client.py 9.55 KB
Newer Older
Matthieu Dorier's avatar
Matthieu Dorier committed
1 2
# (C) 2018 The University of Chicago
# See COPYRIGHT in top-level directory.
Matthieu Dorier's avatar
Matthieu Dorier committed
3
import _pybakeclient
4
from pybake.target import *
Matthieu Dorier's avatar
Matthieu Dorier committed
5 6 7
import pymargo

class BakeClient():
8 9 10 11 12
    """
    The BakeClient class wraps a bake_client_t structure at C level.
    It registers the RPCs necessary to interact with a Bake provider.
    It can be used to create provider handles pointing to Bake providers.
    """
Matthieu Dorier's avatar
Matthieu Dorier committed
13

14 15 16
    def __init__(self, mid):
        """
        Constructor. Initializes a new BakeClient with a MargoInstance.
Matthieu Dorier's avatar
Matthieu Dorier committed
17

18 19 20 21
        Args:
            mid (MargoInstance): MargoInstance on which to register RPCs.
        """
        self._client = _pybakeclient.client_init(mid._mid)
Matthieu Dorier's avatar
Matthieu Dorier committed
22

23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42
    def create_provider_handle(self, addr, provider_id):
        """
        Creates a BakeProviderHandle object pointing to the given
        address and provider id.

        Args:
            addr (MargoAddress): Address of the Bake provider.
            provider_id (int): ID of the provider.
        """
        ph = _pybakeclient.provider_handle_create(self._client, addr.get_hg_addr(), provider_id)
        return BakeProviderHandle(ph)

    def shutdown_service(self, addr):
        """
        Shut down a MargoInstance running at a particular address.

        Args:
            addr (MargoAddress): Address of the MargoInstance to shut down.
        """
        _pybakeclient.shutdown_service(self._client, addr.get_hg_addr())
Matthieu Dorier's avatar
Matthieu Dorier committed
43
	
44 45 46 47 48
    def finalize(self):
        """
        Finalizes the underlying bake_client_t structure.
        """
        _pybakeclient.client_finalize(self._client)
Matthieu Dorier's avatar
Matthieu Dorier committed
49 50

class BakeProviderHandle():
51 52 53 54
    """
    The BakeProviderHandle class represents a handle to a remote Bake provider.
    Internally, this class wraps a bake_provider_handle_t C structure.
    """
Matthieu Dorier's avatar
Matthieu Dorier committed
55

56 57 58 59 60 61 62
    def __init__(self, ph):
        """
        Constructor. This is not supposed to be called by users.
        Users should create a BakeProviderHandle from a BakeClient clt
        by calling clt.create_provider_handle.
        """
        self._ph = ph
Matthieu Dorier's avatar
Matthieu Dorier committed
63
	
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
    def __del__(self):
        """
        Explicit destructor to call provider_handle_release on the underlying
        bake_provider_handle_t C structure.
        """
        _pybakeclient.provider_handle_release(self._ph)

    def get_eager_limit(self):
        """
        Get the threshold size bellow which this provider
        handle will embed data within the RPC arguments.
        """
        return _pybakeclient.get_eager_limit(self._ph)

    def set_eager_limit(self, limit):
        """
        Set the threshold size bellow which this provider handle
        will embed data within the RPC arguments.
        """ 
        return _pybakeclient.set_eager_limit(self._ph, limit)

    def probe(self, max_targets=0):
        """
        Get the list of BakeTargetIDs of targets located in
        this provider. If max_targets is not specified, this
        function will return all the target ids.
        """
        if(max_targets != 0):
            tgts = _pybakeclient.probe(self._ph, max_targets)
        else:
            num_targets = 32
            while(True):
                tgts = _pybakeclient.probe(self._ph, num_targets)
                if(len(tgts) == num_targets):
                    num_targets *= 2
                else:
                    break
        result = []
        for tgt in tgts:
            result.append(BakeTargetID(tgt))
        return result

    def create(self, bti, region_size):
        """
        Creates a region in the specified target, with a given size.

        Args:
            bti (BakeTargetID): ID of the bake target in which to create the region.
            region_size (int): size of the region to create.

        Returns:
115
            A BakeRegionID object representing the region. None if an error occured.
116 117
        """
        rid = _pybakeclient.create(self._ph, bti._tid, region_size)
118 119
        if(rid == None):
            return None
120 121 122 123 124 125 126 127 128 129
        return BakeRegionID(rid)

    def write(self, rid, offset, data):
        """
        Writes data in a region, at a specified offset.

        Args:
            rid (BakeRegionID): region in which to write.
            offset (int): offset at which to write.
            data (str): data to write.
130 131
        Returns:
            True if the data was correctly written, False otherwise.
132 133 134
        """
        return _pybakeclient.write(self._ph, rid._rid, offset, data)

Matthieu Dorier's avatar
Matthieu Dorier committed
135 136 137 138 139 140 141 142
    def write_numpy(self, rid, offset, array):
        """
        Writes a numpy array in a region at a specified offset.

        Args:
            rid (BakeRegionID): region in which to write.
            offset (int): offset at which to write.
            data (numpy.ndarray): numpy array to write.
143 144
        Returns:
            True if the data was correctly written, False otherwise.
Matthieu Dorier's avatar
Matthieu Dorier committed
145 146 147
        """
        return _pybakeclient.write_numpy(self._ph, rid._rid, offset, array)

148 149 150 151 152 153
    def persist(self, rid):
        """
        Make the changes to a given region persist.
        
        Args:
            rid (BakeRegionID): region to persist.
154 155
        Returns:
            True if the region was correctly persisted, False otherwise.
156 157 158 159 160 161 162 163 164 165 166 167 168
        """
        return _pybakeclient.persist(self._ph, rid._rid)

    def create_write_persist(self, bti, data):
        """
        Creates a new region, write data to it at a given offset,
        and persist the region.

        Args:
            bti (BakeTargetID): target id in which to create the region.
            size (int): size of the region to create.
            offset (int): offset at which to write data in the region.
            data (str): data to write.
169 170
        Returns:
            True if the region was correctly written and persisted, False otherwise.
171 172
        """
        rid = _pybakeclient.create_write_persist(self._ph, bti._tid, data)
Matthieu Dorier's avatar
Matthieu Dorier committed
173 174 175 176 177 178 179 180 181 182 183 184 185 186
        if(rid is None):
            return None
        return BakeRegionID(rid)

    def create_write_persist_numpy(self, bti, array):
        """
        Creates a new region, write the numpy array to it at a given offset,
        and persist the region.

        Args:
            bti (BakeTargetID): target id in which to create the region.
            size (int): size of the region to create.
            offset (int): offset at which to write data in the region.
            array (numpy.ndarray): numpy array to write.
187 188
        Returns:
            True if the region was correctly written and persisted, False otherwise.
Matthieu Dorier's avatar
Matthieu Dorier committed
189 190 191 192
        """
        rid = _pybakeclient.create_write_persist_numpy(self._ph, bti._tid, array)
        if(rid is None):
            return None
193 194 195 196 197 198 199 200 201 202
        return BakeRegionID(rid)

    def get_size(self, rid):
        """
        Get the size of a given region.

        Args:
            rid (BakeRegionID): region id.

        Returns:
203
            The size (int) of the provided region. None in case of error.
204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
        """
        return _pybakeclient.get_size(self._ph, rid._rid)

    def read(self, rid, offset=0, size=-1):
        """
        Reads the data contained in a given region, at a given offset
        and with a given size. If the size is not provided,
        this function will first send an RPC to get the current region size
        and then read from the offset up to the end of the region. If the
        offset is not provided, the entire region is read.

        Args:
            rid (BakeRegionID): region id.
            offset (int): offset at which to read.
            size (int): size to read.

        Returns:
            The data read, in the form of a string, or None if an
            error occured.
        """
        if(size < 0):
            size = self.get_size(rid) - offset
        return _pybakeclient.read(self._ph, rid._rid, offset, size)
    
Matthieu Dorier's avatar
Matthieu Dorier committed
228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244
    def read_numpy(self, rid, offset, shape, dtype):
        """
        Reads the data contained in a given region, at a given offset,
        and interpret it as a numpy array of a given shape and datatype.
        This function will fail if the full array cannot be loaded
        (e.g. the size of the region from the provided offset is too small
         compared with the size of the numpy that should result from the call)

        Args:
            rid (BakeRegionID): region id.
            offset (int): offset at which to read.
            shape (tuple): shape of the resulting array.
            dtype (numpy.dtype): datatype of the resuling array.

        Returns:
            A numpy array or None if it could not be read.
        """
245
        return _pybakeclient.read_numpy(self._ph, rid._rid, offset, tuple(shape), dtype)
Matthieu Dorier's avatar
Matthieu Dorier committed
246

247 248 249 250 251 252 253
    def remove(self, rid):
        """
        Remove a region from its target.

        Args:
            rid (BakeRegionID): region to remove.
        """
254 255
        ret = _pybakeclient.remove(self._ph, rid._rid)
        return (ret == 0)
256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276

    def migrate(self, source_rid, dest_addr, dest_provider_id, dest_target, remove_source=True):
        """
        Migrates a give region from its source to a destination designated by
        an address, a provider id, and a target id. This function will also remove
        the original region if remove_source is set to True.

        Args:
            source_rid (BakeRegionID): region to remove.
            dest_addr (str): destination address.
            dest_provider_id (int): destination provider id.
            dest_target (BakeTargetID): destinatin target id.
            remove_source (bool): whether to remove the source region.
        Returns:
            The resulting BakeRegionID if successful, None otherwise.
        """
        ret = _pybakeclient.migrate(self._ph, source_rid._rid, remove_source,
                str(dest_addr), int(dest_provider_id), dest_target._tid)
        if(ret == None):
            return ret
        return BakeRegionID(ret)