mirror of
https://github.com/wassname/ray.git
synced 2026-07-20 12:40:20 +08:00
Plasma C extensions (#34)
* switch plasma from ctypes to python C API * clang-format * various fixes
This commit is contained in:
Philipp Moritz
committed by
Robert Nishihara
Robert Nishihara
parent
ad6a401740
commit
986ed5c9e8
+24
-107
@@ -1,25 +1,12 @@
|
||||
import ctypes
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
Addr = ctypes.c_ubyte * 4
|
||||
import libplasma
|
||||
|
||||
PLASMA_ID_SIZE = 20
|
||||
ID = ctypes.c_ubyte * PLASMA_ID_SIZE
|
||||
|
||||
class PlasmaID(ctypes.Structure):
|
||||
_fields_ = [("plasma_id", ID)]
|
||||
|
||||
def make_plasma_id(string):
|
||||
if len(string) != PLASMA_ID_SIZE:
|
||||
raise Exception("PlasmaIDs must be {} characters long".format(PLASMA_ID_SIZE))
|
||||
return PlasmaID(plasma_id=ID.from_buffer_copy(string))
|
||||
|
||||
def plasma_id_to_str(plasma_id):
|
||||
return str(bytearray(plasma_id.plasma_id))
|
||||
PLASMA_WAIT_TIMEOUT = 2 ** 36
|
||||
|
||||
class PlasmaBuffer(object):
|
||||
"""This is the type of objects returned by calls to get with a PlasmaClient.
|
||||
@@ -47,7 +34,7 @@ class PlasmaBuffer(object):
|
||||
If the plasma client has been shut down, then don't do anything.
|
||||
"""
|
||||
if self.plasma_client.alive:
|
||||
self.plasma_client.client.plasma_release(self.plasma_client.plasma_conn, self.plasma_id)
|
||||
libplasma.release(self.plasma_client.conn, self.plasma_id)
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Read from the PlasmaBuffer as if it were just a regular buffer."""
|
||||
@@ -80,33 +67,11 @@ class PlasmaClient(object):
|
||||
manager_socket_name (str): Name of the socket the plasma manager is listening at.
|
||||
"""
|
||||
self.alive = True
|
||||
plasma_client_library = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../build/plasma_client.so")
|
||||
self.client = ctypes.cdll.LoadLibrary(plasma_client_library)
|
||||
|
||||
self.client.plasma_connect.restype = ctypes.c_void_p
|
||||
self.client.plasma_create.restype = None
|
||||
self.client.plasma_get.restype = None
|
||||
self.client.plasma_release.restype = None
|
||||
self.client.plasma_contains.restype = None
|
||||
self.client.plasma_seal.restype = None
|
||||
self.client.plasma_delete.restype = None
|
||||
self.client.plasma_subscribe.restype = ctypes.c_int
|
||||
self.client.plasma_wait.restype = ctypes.c_int
|
||||
|
||||
self.buffer_from_memory = ctypes.pythonapi.PyBuffer_FromMemory
|
||||
self.buffer_from_memory.argtypes = [ctypes.c_void_p, ctypes.c_int64]
|
||||
self.buffer_from_memory.restype = ctypes.py_object
|
||||
|
||||
self.buffer_from_read_write_memory = ctypes.pythonapi.PyBuffer_FromReadWriteMemory
|
||||
self.buffer_from_read_write_memory.argtypes = [ctypes.c_void_p, ctypes.c_int64]
|
||||
self.buffer_from_read_write_memory.restype = ctypes.py_object
|
||||
|
||||
if manager_socket_name is not None:
|
||||
self.has_manager_conn = True
|
||||
self.plasma_conn = ctypes.c_void_p(self.client.plasma_connect(store_socket_name, manager_socket_name, release_delay))
|
||||
self.conn = libplasma.connect(store_socket_name, manager_socket_name, release_delay)
|
||||
else:
|
||||
self.has_manager_conn = False
|
||||
self.plasma_conn = ctypes.c_void_p(self.client.plasma_connect(store_socket_name, None, release_delay))
|
||||
self.conn = libplasma.connect(store_socket_name, "", release_delay)
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown the client so that it does not send messages.
|
||||
@@ -115,6 +80,8 @@ class PlasmaClient(object):
|
||||
to, then we can use this method to prevent the client from trying to send
|
||||
messages to the killed processes.
|
||||
"""
|
||||
if self.alive:
|
||||
libplasma.disconnect(self.conn)
|
||||
self.alive = False
|
||||
|
||||
def create(self, object_id, size, metadata=None):
|
||||
@@ -128,13 +95,10 @@ class PlasmaClient(object):
|
||||
metadata (buffer): An optional buffer encoding whatever metadata the user
|
||||
wishes to encode.
|
||||
"""
|
||||
# This is used to hold the address of the buffer.
|
||||
data = ctypes.c_void_p()
|
||||
# Turn the metadata into the right type.
|
||||
metadata = buffer("") if metadata is None else metadata
|
||||
metadata = (ctypes.c_ubyte * len(metadata)).from_buffer_copy(metadata)
|
||||
self.client.plasma_create(self.plasma_conn, make_plasma_id(object_id), size, ctypes.cast(metadata, ctypes.POINTER(ctypes.c_ubyte * len(metadata))), len(metadata), ctypes.byref(data))
|
||||
return PlasmaBuffer(self.buffer_from_read_write_memory(data, size), make_plasma_id(object_id), self)
|
||||
metadata = bytearray("") if metadata is None else metadata
|
||||
buff = libplasma.create(self.conn, object_id, size, metadata)
|
||||
return PlasmaBuffer(buff, object_id, self)
|
||||
|
||||
def get(self, object_id):
|
||||
"""Create a buffer from the PlasmaStore based on object ID.
|
||||
@@ -145,12 +109,8 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
size = ctypes.c_int64()
|
||||
data = ctypes.c_void_p()
|
||||
metadata_size = ctypes.c_int64()
|
||||
metadata = ctypes.c_void_p()
|
||||
self.client.plasma_get(self.plasma_conn, make_plasma_id(object_id), ctypes.byref(size), ctypes.byref(data), ctypes.byref(metadata_size), ctypes.byref(metadata))
|
||||
return PlasmaBuffer(self.buffer_from_memory(data, size), make_plasma_id(object_id), self)
|
||||
buff = libplasma.get(self.conn, object_id)[0]
|
||||
return PlasmaBuffer(buff, object_id, self)
|
||||
|
||||
def get_metadata(self, object_id):
|
||||
"""Create a buffer from the PlasmaStore based on object ID.
|
||||
@@ -161,12 +121,8 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
size = ctypes.c_int64()
|
||||
data = ctypes.c_void_p()
|
||||
metadata_size = ctypes.c_int64()
|
||||
metadata = ctypes.c_void_p()
|
||||
self.client.plasma_get(self.plasma_conn, make_plasma_id(object_id), ctypes.byref(size), ctypes.byref(data), ctypes.byref(metadata_size), ctypes.byref(metadata))
|
||||
return PlasmaBuffer(self.buffer_from_memory(metadata, metadata_size), make_plasma_id(object_id), self)
|
||||
buff = libplasma.get(self.conn, object_id)[1]
|
||||
return PlasmaBuffer(buff, object_id, self)
|
||||
|
||||
def contains(self, object_id):
|
||||
"""Check if the object is present and has been sealed in the PlasmaStore.
|
||||
@@ -174,15 +130,7 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
has_object = ctypes.c_int()
|
||||
self.client.plasma_contains(self.plasma_conn, make_plasma_id(object_id), ctypes.byref(has_object))
|
||||
has_object = has_object.value
|
||||
if has_object == 1:
|
||||
return True
|
||||
elif has_object == 0:
|
||||
return False
|
||||
else:
|
||||
raise Exception("This code should be unreachable.")
|
||||
return libplasma.contains(self.conn, object_id)
|
||||
|
||||
def seal(self, object_id):
|
||||
"""Seal the buffer in the PlasmaStore for a particular object ID.
|
||||
@@ -193,7 +141,7 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
self.client.plasma_seal(self.plasma_conn, make_plasma_id(object_id))
|
||||
libplasma.seal(self.conn, object_id)
|
||||
|
||||
def delete(self, object_id):
|
||||
"""Delete the buffer in the PlasmaStore for a particular object ID.
|
||||
@@ -203,7 +151,7 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
self.client.plasma_delete(self.plasma_conn, make_plasma_id(object_id))
|
||||
libplasma.delete(self.conn, object_id)
|
||||
|
||||
def evict(self, num_bytes):
|
||||
"""Evict some objects until to recover some bytes.
|
||||
@@ -213,8 +161,7 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
num_bytes (int): The number of bytes to attempt to recover.
|
||||
"""
|
||||
num_bytes_evicted = self.client.plasma_evict(self.plasma_conn, num_bytes)
|
||||
return num_bytes_evicted
|
||||
return libplasma.evict(self.conn, num_bytes)
|
||||
|
||||
def transfer(self, addr, port, object_id):
|
||||
"""Transfer local object with id object_id to another plasma instance
|
||||
@@ -224,9 +171,7 @@ class PlasmaClient(object):
|
||||
port (int): Port number of the plasma instance the object is sent to.
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
if not self.has_manager_conn:
|
||||
raise Exception("Not connected to the plasma manager socket")
|
||||
self.client.plasma_transfer(self.plasma_conn, addr, port, make_plasma_id(object_id))
|
||||
return libplasma.transfer(self.conn, object_id, addr, port)
|
||||
|
||||
def fetch(self, object_ids):
|
||||
"""Fetch the object with id object_id from another plasma manager instance.
|
||||
@@ -234,19 +179,9 @@ class PlasmaClient(object):
|
||||
Args:
|
||||
object_id (str): A string used to identify an object.
|
||||
"""
|
||||
object_id_array = (len(object_ids) * PlasmaID)()
|
||||
for i, object_id in enumerate(object_ids):
|
||||
object_id_array[i] = make_plasma_id(object_id)
|
||||
success_array = (len(object_ids) * ctypes.c_int)()
|
||||
if not self.has_manager_conn:
|
||||
raise Exception("Not connected to the plasma manager socket")
|
||||
self.client.plasma_fetch(self.plasma_conn,
|
||||
object_id_array._length_,
|
||||
object_id_array,
|
||||
success_array);
|
||||
return [bool(success) for success in success_array]
|
||||
return libplasma.fetch(self.conn, object_ids)
|
||||
|
||||
def wait(self, object_ids, timeout, num_returns):
|
||||
def wait(self, object_ids, timeout=PLASMA_WAIT_TIMEOUT, num_returns=1):
|
||||
"""Wait until num_returns objects in object_ids are ready.
|
||||
|
||||
Args:
|
||||
@@ -258,30 +193,12 @@ class PlasmaClient(object):
|
||||
ready_ids, waiting_ids (List[str], List[str]): List of object IDs that
|
||||
are ready and list of object IDs we might still wait on respectively.
|
||||
"""
|
||||
if not self.has_manager_conn:
|
||||
raise Exception("Not connected to the plasma manager socket")
|
||||
if num_returns < 0:
|
||||
raise Exception("The argument num_returns cannot be less than one.")
|
||||
if num_returns > len(object_ids):
|
||||
raise Exception("The argument num_returns cannot be greater than len(object_ids): num_returns is {}, len(object_ids) is {}.".format(num_returns, len(object_ids)))
|
||||
if timeout > 2 ** 36:
|
||||
raise Exception("The method wait currently cannot be used with a timeout greater than 2 ** 36.")
|
||||
object_id_array = (len(object_ids) * PlasmaID)()
|
||||
for i, object_id in enumerate(object_ids):
|
||||
object_id_array[i] = make_plasma_id(object_id)
|
||||
return_id_array = (num_returns * PlasmaID)()
|
||||
num_return_objects = self.client.plasma_wait(self.plasma_conn,
|
||||
object_id_array._length_,
|
||||
object_id_array,
|
||||
ctypes.c_int64(timeout),
|
||||
num_returns,
|
||||
return_id_array)
|
||||
ready_ids = map(plasma_id_to_str, return_id_array[num_returns-num_return_objects:])
|
||||
return ready_ids, list(set(object_ids) - set(ready_ids))
|
||||
ready_ids, waiting_ids = libplasma.wait(self.conn, object_ids, timeout, num_returns)
|
||||
return ready_ids, list(waiting_ids)
|
||||
|
||||
def subscribe(self):
|
||||
"""Subscribe to notifications about sealed objects."""
|
||||
fd = self.client.plasma_subscribe(self.plasma_conn)
|
||||
fd = libplasma.subscribe(self.conn)
|
||||
self.notification_sock = socket.fromfd(fd, socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
# Make the socket non-blocking.
|
||||
self.notification_sock.setblocking(0)
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
from setuptools import setup, find_packages
|
||||
import setuptools.command.install as _install
|
||||
|
||||
import subprocess
|
||||
|
||||
class install(_install.install):
|
||||
def run(self):
|
||||
subprocess.check_call(["make"], cwd="../../")
|
||||
subprocess.check_call(["cmake", ".."], cwd="../../build")
|
||||
subprocess.check_call(["make", "install"], cwd="../../build")
|
||||
# Calling _install.install.run(self) does not fetch required packages and
|
||||
# instead performs an old-style install. See command/install.py in
|
||||
# setuptools. So, calling do_egg_install() manually here.
|
||||
self.do_egg_install()
|
||||
|
||||
setup(name="Plasma",
|
||||
version="0.0.1",
|
||||
description="Plasma client for Python",
|
||||
packages=find_packages(),
|
||||
package_data={"plasma": ["libplasma.so"]},
|
||||
cmdclass={"install": install},
|
||||
include_package_data=True,
|
||||
zip_safe=False)
|
||||
Reference in New Issue
Block a user