mirror of
https://github.com/vale981/ray
synced 2025-03-12 14:16:39 -04:00
72361c9b44
* Automatic whitespace fixes. * Add metadata handling. * Make create take a buffer instead of a string for the metadata. * Small fixes.
194 lines
7.6 KiB
Python
194 lines
7.6 KiB
Python
from __future__ import print_function
|
|
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import unittest
|
|
import random
|
|
import time
|
|
import tempfile
|
|
|
|
import plasma
|
|
|
|
def random_object_id():
|
|
return "".join([chr(random.randint(0, 255)) for _ in range(20)])
|
|
|
|
def generate_metadata(length):
|
|
metadata = length * ["\x00"]
|
|
if length > 0:
|
|
metadata[0] = chr(random.randint(0, 255))
|
|
metadata[-1] = chr(random.randint(0, 255))
|
|
for _ in range(100):
|
|
metadata[random.randint(0, length - 1)] = chr(random.randint(0, 255))
|
|
return buffer("".join(metadata))
|
|
|
|
def write_to_data_buffer(buff, length):
|
|
if length > 0:
|
|
buff[0] = chr(random.randint(0, 255))
|
|
buff[-1] = chr(random.randint(0, 255))
|
|
for _ in range(100):
|
|
buff[random.randint(0, length - 1)] = chr(random.randint(0, 255))
|
|
|
|
def create_object(client, data_size, metadata_size, seal=True):
|
|
object_id = random_object_id()
|
|
metadata = generate_metadata(metadata_size)
|
|
memory_buffer = client.create(object_id, data_size, metadata)
|
|
write_to_data_buffer(memory_buffer, data_size)
|
|
if seal:
|
|
client.seal(object_id)
|
|
return object_id, memory_buffer, metadata
|
|
|
|
class TestPlasmaClient(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
# Start Plasma.
|
|
plasma_store_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_store")
|
|
store_name = "/tmp/store{}".format(random.randint(0, 10000))
|
|
self.p = subprocess.Popen([plasma_store_executable, "-s", store_name])
|
|
# Connect to Plasma.
|
|
self.plasma_client = plasma.PlasmaClient(store_name)
|
|
|
|
def tearDown(self):
|
|
# Kill the plasma store process.
|
|
self.p.kill()
|
|
|
|
def test_create(self):
|
|
# Create an object id string.
|
|
object_id = random_object_id()
|
|
# Create a new buffer and write to it.
|
|
length = 50
|
|
memory_buffer = self.plasma_client.create(object_id, length)
|
|
for i in range(length):
|
|
memory_buffer[i] = chr(i % 256)
|
|
# Seal the object.
|
|
self.plasma_client.seal(object_id)
|
|
# Get the object.
|
|
memory_buffer = self.plasma_client.get(object_id)
|
|
for i in range(length):
|
|
self.assertEqual(memory_buffer[i], chr(i % 256))
|
|
|
|
def test_create_with_metadata(self):
|
|
for length in range(1000):
|
|
# Create an object id string.
|
|
object_id = random_object_id()
|
|
# Create a random metadata string.
|
|
metadata = generate_metadata(length)
|
|
# Create a new buffer and write to it.
|
|
memory_buffer = self.plasma_client.create(object_id, length, metadata)
|
|
for i in range(length):
|
|
memory_buffer[i] = chr(i % 256)
|
|
# Seal the object.
|
|
self.plasma_client.seal(object_id)
|
|
# Get the object.
|
|
memory_buffer = self.plasma_client.get(object_id)
|
|
for i in range(length):
|
|
self.assertEqual(memory_buffer[i], chr(i % 256))
|
|
# Get the metadata.
|
|
metadata_buffer = self.plasma_client.get_metadata(object_id)
|
|
self.assertEqual(len(metadata), len(metadata_buffer))
|
|
for i in range(len(metadata)):
|
|
self.assertEqual(metadata[i], metadata_buffer[i])
|
|
|
|
def test_illegal_functionality(self):
|
|
# Create an object id string.
|
|
object_id = random_object_id()
|
|
# Create a new buffer and write to it.
|
|
length = 1000
|
|
memory_buffer = self.plasma_client.create(object_id, length)
|
|
# Make sure we cannot access memory out of bounds.
|
|
self.assertRaises(Exception, lambda : memory_buffer[length])
|
|
# Seal the object.
|
|
self.plasma_client.seal(object_id)
|
|
# This test is commented out because it currently fails.
|
|
# # Make sure the object is ready only now.
|
|
# def illegal_assignment():
|
|
# memory_buffer[0] = chr(0)
|
|
# self.assertRaises(Exception, illegal_assignment)
|
|
# Get the object.
|
|
memory_buffer = self.plasma_client.get(object_id)
|
|
# Make sure the object is read only.
|
|
def illegal_assignment():
|
|
memory_buffer[0] = chr(0)
|
|
self.assertRaises(Exception, illegal_assignment)
|
|
|
|
class TestPlasmaManager(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
# Start two PlasmaStores.
|
|
plasma_store_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_store")
|
|
store_name1 = "/tmp/store{}".format(random.randint(0, 10000))
|
|
store_name2 = "/tmp/store{}".format(random.randint(0, 10000))
|
|
self.p2 = subprocess.Popen([plasma_store_executable, "-s", store_name1])
|
|
self.p3 = subprocess.Popen([plasma_store_executable, "-s", store_name2])
|
|
# Start two PlasmaManagers.
|
|
self.port1 = random.randint(10000, 50000)
|
|
self.port2 = random.randint(10000, 50000)
|
|
plasma_manager_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_manager")
|
|
self.p4 = subprocess.Popen([plasma_manager_executable, "-s", store_name1, "-m", "127.0.0.1", "-p", str(self.port1)])
|
|
self.p5 = subprocess.Popen([plasma_manager_executable, "-s", store_name2, "-m", "127.0.0.1", "-p", str(self.port2)])
|
|
time.sleep(0.1)
|
|
# Connect two PlasmaClients.
|
|
self.client1 = plasma.PlasmaClient(store_name1, "127.0.0.1", self.port1)
|
|
self.client2 = plasma.PlasmaClient(store_name2, "127.0.0.1", self.port2)
|
|
time.sleep(0.5)
|
|
|
|
def tearDown(self):
|
|
# Kill the PlasmaStore and PlasmaManager processes.
|
|
self.p2.kill()
|
|
self.p3.kill()
|
|
self.p4.kill()
|
|
self.p5.kill()
|
|
|
|
def test_transfer(self):
|
|
for _ in range(100):
|
|
# Create an object.
|
|
object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000)
|
|
# Transfer the buffer to the the other PlasmaStore.
|
|
self.client1.transfer("127.0.0.1", self.port2, object_id1)
|
|
# Compare the two buffers.
|
|
self.assertEqual(memory_buffer1[:], self.client2.get(object_id1)[:])
|
|
self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])
|
|
self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
|
|
self.assertEqual(self.client1.get_metadata(object_id1)[:], self.client2.get_metadata(object_id1)[:])
|
|
# Transfer the buffer again.
|
|
self.client1.transfer("127.0.0.1", self.port2, object_id1)
|
|
self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
|
|
# Compare the two buffers.
|
|
self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])
|
|
|
|
# Create an object.
|
|
object_id2, memory_buffer2, metadata2 = create_object(self.client2, 20000, 20000)
|
|
# Transfer the buffer to the the other PlasmaStore.
|
|
self.client2.transfer("127.0.0.1", self.port1, object_id2)
|
|
# Compare the two buffers.
|
|
self.assertEqual(memory_buffer2[:], self.client2.get(object_id2)[:])
|
|
self.assertEqual(self.client1.get(object_id2)[:], self.client2.get(object_id2)[:])
|
|
self.assertEqual(metadata2[:], self.client2.get_metadata(object_id2)[:])
|
|
self.assertEqual(self.client1.get_metadata(object_id2)[:], self.client2.get_metadata(object_id2)[:])
|
|
|
|
def test_illegal_functionality(self):
|
|
# Create an object id string.
|
|
object_id = random_object_id()
|
|
# Create a new buffer.
|
|
memory_buffer = self.client1.create(object_id, 20000)
|
|
# This test is commented out because it currently fails.
|
|
# # Transferring the buffer before sealing it should fail.
|
|
# self.assertRaises(Exception, lambda : self.manager1.transfer(1, object_id))
|
|
|
|
def test_stresstest(self):
|
|
a = time.time()
|
|
object_ids = []
|
|
for i in range(10000): # TODO(pcm): increase this to 100000
|
|
object_id = random_object_id()
|
|
object_ids.append(object_id)
|
|
self.client1.create(object_id, 1)
|
|
self.client1.seal(object_id)
|
|
for object_id in object_ids:
|
|
self.client1.transfer("127.0.0.1", self.port2, object_id)
|
|
b = time.time() - a
|
|
|
|
print("it took", b, "seconds to put and transfer the objects")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|