Skip to content

Commit

Permalink
Vine: Add IO test for function calls (#3981)
Browse files Browse the repository at this point in the history
* add IO test for function calls

* add more checks

* change var name
  • Loading branch information
tphung3 authored Nov 22, 2024
1 parent 46122bd commit 96a03b1
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 2 deletions.
9 changes: 8 additions & 1 deletion taskvine/test/TR_vine_python_serverless.sh
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@ export PYTHONPATH=$(pwd)/../../test_support/python_modules/${CCTOOLS_PYTHON_TEST

STATUS_FILE=vine.status
PORT_FILE=vine.port
PYTHON_SCRIPT=vine_python_serverless.py
TEST_INPUT_FILE=${PYTHON_SCRIPT}.input
TEST_OUTPUT_FILE=${PYTHON_SCRIPT}.output

check_needed()
{
Expand All @@ -33,12 +36,14 @@ prepare()
{
rm -f $STATUS_FILE
rm -f $PORT_FILE
rm -f $TEST_INPUT_FILE
rm -f $TEST_OUTPUT_FILE
return 0
}

run()
{
( ${CCTOOLS_PYTHON_TEST_EXEC} vine_python_serverless.py $PORT_FILE; echo $? > $STATUS_FILE ) &
( ${CCTOOLS_PYTHON_TEST_EXEC} ${PYTHON_SCRIPT} $PORT_FILE; echo $? > $STATUS_FILE ) &

# wait at most 15 seconds for vine to find a port.
wait_for_file_creation $PORT_FILE 15
Expand Down Expand Up @@ -77,6 +82,8 @@ clean()
{
rm -f $STATUS_FILE
rm -f $PORT_FILE
rm -f $TEST_INPUT_FILE
rm -f $TEST_OUTPUT_FILE
rm -rf vine-run-info
rm worker.log
exit 0
Expand Down
44 changes: 43 additions & 1 deletion taskvine/test/vine_python_serverless.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import math
import json
from ndcctools.taskvine.utils import load_variable_from_library
import os

def exp(x, y=3):
return {'base_val': x**y}
Expand Down Expand Up @@ -41,6 +42,12 @@ def double(x, with_library=False):
return base_val + m.prod([x,2])
return m.prod([x, 2])

def func_copy_input_to_output(input_filename, output_filename):
with open(input_filename, 'r') as f:
contents = f.read()
with open(output_filename, 'w') as f:
f.write(contents)
return 0

def main():
parser = argparse.ArgumentParser("Test for taskvine python bindings.")
Expand Down Expand Up @@ -77,6 +84,9 @@ def main():
exec("def dyn_fn(x):\n return x + 2", globals(), globals())
libtask_with_special_fns = q.create_library_from_functions('test-library-with-special-fns', lambda_fn, dyn_fn, add_env=False, exec_mode='fork')

# define a function that copies an input file to an output file.
# this is to test the input/output staging of function calls
libtask_with_io_fn = q.create_library_from_functions('test-library-with-io-fn', func_copy_input_to_output, add_env=False, exec_mode='fork')

# Just take default resources for the library, this will cause it to fill the whole worker.
# And the number of functions slots will match the number of cores available.
Expand All @@ -86,7 +96,13 @@ def main():
q.install_library(libtask_with_context_direct)
q.install_library(libtask_with_context_fork)
q.install_library(libtask_with_special_fns)
lib_task_names = ['test-library-no-context-direct', 'test-library-no-context-fork', 'test-library-with-context-direct', 'test-library-with-context-fork', 'test-library-with-special-fns']
q.install_library(libtask_with_io_fn)
lib_task_names = ['test-library-no-context-direct',
'test-library-no-context-fork',
'test-library-with-context-direct',
'test-library-with-context-fork',
'test-library-with-special-fns',
'test-library-with-io-fn']
print("Submitting function call tasks...")

tasks = 100
Expand All @@ -102,6 +118,20 @@ def main():

s_task = vine.FunctionCall(lib_name, 'dyn_fn', 1)
q.submit(s_task)
elif lib_name == 'test-library-with-io-fn':
input_filename = os.path.basename(__file__) + '.input'
output_filename = os.path.basename(__file__) + '.output'
with open(input_filename, 'w') as f:
print('Test IO with function calls', file=f)
s_task = vine.FunctionCall(lib_name, 'func_copy_input_to_output', input_filename, output_filename)
input_file = q.declare_file(input_filename)
output_file = q.declare_file(output_filename)
s_task.add_input(input_file, input_filename)
s_task.add_output(output_file, output_filename)
q.submit(s_task)

# do this test once only
break
else:
s_task = vine.FunctionCall(lib_name, 'divide', 2, 2**2, with_library=with_library)
q.submit(s_task)
Expand Down Expand Up @@ -136,6 +166,18 @@ def main():
print(f"Total: {total_sum}")
print(f"Expected: {expected}")

# Check that IO test passed
with open(input_filename, 'r') as f:
content_input = f.read()

with open(output_filename, 'r') as f:
content_output = f.read()

print(f"Input: {content_input}", end='')
print(f"Output: {content_output}", end='')

assert content_input == content_output

assert total_sum == expected

if __name__ == '__main__':
Expand Down

0 comments on commit 96a03b1

Please sign in to comment.