def test_resolve_memory_then_disk(mock_get_step_uuid, plasma_store):
orchest.Config.PIPELINE_DEFINITION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1.
mock_get_step_uuid.return_value = "uuid-1______________"
data_1 = generate_data(KILOBYTE)
transfer.output_to_memory(
data_1,
name=None,
disk_fallback=False,
)
# It is very unlikely you will output through memory and disk in quick
# succession. In addition, the resolve order has a precision of
# seconds. Thus we need to ensure that indeed it can be resolved.
time.sleep(1)
data_1_new = generate_data(KILOBYTE)
transfer.output_to_disk(data_1_new, name=None)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
input_data = input_data[orchest.Config._RESERVED_UNNAMED_OUTPUTS_STR][0]
assert (input_data == data_1_new).all()
def test_receive_input_order(mock_get_step_uuid, plasma_store):
"""Test the order of the inputs of the receiving step.
Note that the order in which the data is output does not determine the
"receive order", it is the order in which it is defined in the
pipeline.json (for the "incoming-connections").
"""
orchest.Config.PIPELINE_DEFINITION_PATH = "tests/userdir/pipeline-order.json"
# Do as if we are uuid-3
data_3 = generate_data(KILOBYTE)
mock_get_step_uuid.return_value = "uuid-3______________"
transfer.output_to_memory(data_3, name=None)
# Do as if we are uuid-1
data_1 = generate_data(KILOBYTE)
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(data_1, name=None)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
input_data = input_data[orchest.Config._RESERVED_UNNAMED_OUTPUTS_STR]
assert (input_data[0] == data_1).all()
assert (input_data[1] == data_3).all()
def test_memory_out_of_memory(mock_get_step_uuid, plasma_store):
data_1 = generate_data((PLASMA_KILOBYTES + 1) * KILOBYTE)
data_size = pa.serialize(data_1).total_bytes
assert data_size > PLASMA_STORE_CAPACITY
orchest.Config.PIPELINE_DESCRIPTION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1
mock_get_step_uuid.return_value = "uuid-1______________"
with pytest.raises(MemoryError):
transfer.output_to_memory(
data_1, disk_fallback=False,
)
def test_memory_disk_fallback(mock_get_step_uuid, plasma_store):
orchest.Config.PIPELINE_DESCRIPTION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1
data_1 = generate_data((PLASMA_KILOBYTES + 1) * KILOBYTE)
data_size = pa.serialize(data_1).total_bytes
assert data_size > PLASMA_STORE_CAPACITY
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(
data_1, disk_fallback=True,
)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
assert (input_data[0] == data_1).all()
def test_memory_pickle_fallback_and_disk_fallback(mock_get_step_uuid, plasma_store):
data_1 = [CustomClass(generate_data(KILOBYTE)) for _ in range(PLASMA_KILOBYTES + 1)]
serialized, _ = transfer._serialize(data_1)
assert serialized.size > PLASMA_STORE_CAPACITY
orchest.Config.PIPELINE_DEFINITION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(
data_1,
name=None,
disk_fallback=True,
)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
input_data = input_data[orchest.Config._RESERVED_UNNAMED_OUTPUTS_STR][0]
assert input_data == data_1
def test_receive_multiple_named_inputs(mock_get_step_uuid, plasma_store):
"""Test receiving multiple named inputs."""
orchest.Config.PIPELINE_DEFINITION_PATH = "tests/userdir/pipeline-order.json"
# Do as if we are uuid-3
data_3 = generate_data(KILOBYTE)
mock_get_step_uuid.return_value = "uuid-3______________"
transfer.output_to_memory(data_3, name="output3")
# Do as if we are uuid-1
data_1 = generate_data(KILOBYTE)
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(data_1, name="output1")
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
assert len(input_data) == 3
assert not input_data[orchest.Config._RESERVED_UNNAMED_OUTPUTS_STR]
assert (input_data["output1"] == data_1).all()
assert (input_data["output3"] == data_3).all()
def test_memory_disk_fallback(mock_get_step_uuid, plasma_store):
orchest.Config.PIPELINE_DEFINITION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1
data_1 = generate_data((PLASMA_KILOBYTES + 1) * KILOBYTE)
ser_data, _ = transfer._serialize(data_1)
data_size = ser_data.size
assert data_size > PLASMA_STORE_CAPACITY
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(
data_1,
name=None,
disk_fallback=True,
)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
input_data = input_data[orchest.Config._RESERVED_UNNAMED_OUTPUTS_STR][0]
assert (input_data == data_1).all()
def test_memory_pickle_fallback_and_disk_fallback(mock_get_step_uuid, plasma_store):
data_1 = [
UnserializableByPyarrowObject(generate_data(KILOBYTE))
for _ in range(PLASMA_KILOBYTES + 1)
]
serialized, _ = transfer.serialize(data_1)
assert serialized.total_bytes > PLASMA_STORE_CAPACITY
orchest.Config.PIPELINE_DESCRIPTION_PATH = "tests/userdir/pipeline-basic.json"
# Do as if we are uuid-1
mock_get_step_uuid.return_value = "uuid-1______________"
transfer.output_to_memory(
data_1, disk_fallback=True,
)
# Do as if we are uuid-2
mock_get_step_uuid.return_value = "uuid-2______________"
input_data = transfer.get_inputs()
assert input_data[0] == data_1
