def get_actual_capacity(plasma_client: plasma.PlasmaClient) -> int:
"""
Get actual capacity of plasma store
Parameters
----------
plasma_client: PlasmaClient
Plasma client.
Returns
-------
size: int
Actual storage size in bytes
"""
store_limit = plasma_client.store_capacity()
left_size = store_limit
alloc_fraction = 1
while True:
allocate_size = int(left_size * alloc_fraction / PAGE_SIZE) * PAGE_SIZE
try:
obj_id = plasma.ObjectID.from_random()
buf = [plasma_client.create(obj_id, allocate_size)]
plasma_client.seal(obj_id)
del buf[:]
break
except plasma.PlasmaStoreFull: # pragma: no cover
alloc_fraction *= 0.99
finally:
plasma_client.evict(allocate_size)
return allocate_size
def _output_to_memory(
obj: pa.SerializedPyObject,
client: plasma.PlasmaClient,
obj_id: Optional[plasma.ObjectID] = None,
metadata: Optional[bytes] = None,
memcopy_threads: int = 6,
) -> plasma.ObjectID:
"""Outputs an object to memory.
Args:
obj: Object to output to memory.
client: A PlasmaClient to interface with the in-memory object
store.
obj_id: The ID to assign to the `obj` inside the plasma store.
If ``None`` then one is randomly generated.
metadata: Metadata to add to the `obj` inside the store.
memcopy_threads: The number of threads to use to write the
`obj` into the object store for large objects.
Returns:
The ID of the object inside the store. Either the given `obj_id`
or a randomly generated one.
Raises:
MemoryError: If the `obj` does not fit in memory.
"""
# Check whether the object to be passed in memory actually fits in
# memory. We check explicitely instead of trying to insert it,
# because inserting an already full Plasma store will start evicting
# objects to free up space. However, we want to maintain control
# over what objects get evicted.
total_size = obj.total_bytes + len(metadata)
occupied_size = sum(obj["data_size"] + obj["metadata_size"]
for obj in client.list().values())
# Take a percentage of the maximum capacity such that the message
# for object eviction always fits inside the store.
store_capacity = Config.MAX_RELATIVE_STORE_CAPACITY * client.store_capacity(
)
available_size = store_capacity - occupied_size
if total_size > available_size:
raise MemoryError("Object does not fit in memory")
# In case no `obj_id` is specified, one has to be generated because
# an ID is required for an object to be inserted in the store.
if obj_id is None:
obj_id = plasma.ObjectID.from_random()
# Write the object to the plasma store. If the obj_id already
# exists, then it first has to be deleted. Essentially we are
# overwriting the data (just like we do for disk)
try:
buffer = client.create(obj_id, obj.total_bytes, metadata=metadata)
except plasma.PlasmaObjectExists:
client.delete([obj_id])
buffer = client.create(obj_id, obj.total_bytes, metadata=metadata)
stream = pa.FixedSizeBufferWriter(buffer)
stream.set_memcopy_threads(memcopy_threads)
obj.write_to(stream)
client.seal(obj_id)
return obj_id