def test_local_folder(self, _, compute_job_raises):
channel_name = 'dummy_channel'
compute_plan_tag = 'cp1'
compute_plan_key = str(uuid.uuid4())
file = 'model.txt'
initial_value = 'initial value'
updated_value = 'updated value'
subtuple = {
'key': str(uuid.uuid4()),
'compute_plan_key': compute_plan_key,
'rank': 1,
'algo': {
'key': 'some key'
}
}
# Write an initial value into the compute plan local folder
cp_local_folder = get_cp_local_folder(compute_plan_key)
os.makedirs(cp_local_folder, exist_ok=True)
with open(os.path.join(cp_local_folder, file), 'w') as x:
x.write(initial_value)
# Call `do_task`, which will:
# 1. write a new value to the subtuple local folder
# 2. and then:
# - complete successfully (compute_job_raises == False)
# - or raise an exception (compute_job_raises == True)
with mock.patch('substrapp.ledger.api.query_ledger') as mquery_ledger,\
mock.patch('substrapp.tasks.tasks.generate_command'),\
mock.patch('substrapp.tasks.tasks.save_models'),\
mock.patch('substrapp.tasks.tasks.compute_job') as mcompute_job:
def compute_job(*args, **kwargs):
for vol in kwargs['volumes']:
if vol.endswith('/local'):
with open(os.path.join(vol, file), 'w') as x:
x.write(updated_value)
if compute_job_raises:
raise Exception('Boom!')
mquery_ledger.return_value = {'tag': compute_plan_tag}
mcompute_job.side_effect = compute_job
try:
build_subtuple_folders(subtuple)
do_task(channel_name, subtuple, TRAINTUPLE_TYPE)
except Exception:
if compute_job_raises:
# exception expected
pass
# Check the compute plan local folder value is correct:
# - If do_task did raise an exception then the local value should be unchanged
# - If do_task did not raise an exception then the local value should be updated
with open(os.path.join(cp_local_folder, file), 'r') as x:
content = x.read()
self.assertEqual(
content, initial_value if compute_job_raises else updated_value)
def remove_local_folders(compute_plan_key):
if not settings.ENABLE_REMOVE_LOCAL_CP_FOLDERS:
logger.info(
f'Skipping deletion of local volume for compute plan {compute_plan_key}'
)
return
try:
local_folder = get_cp_local_folder(compute_plan_key)
logger.info(f'Deleting local folder {local_folder}')
shutil.rmtree(local_folder)
except FileNotFoundError:
logger.info(f'No local folder with path {local_folder}')
pass
except Exception:
logger.error(f'Cannot delete volume {local_folder}', exc_info=True)
if settings.TASK['CHAINKEYS_ENABLED']:
chainkeys_directory = get_chainkeys_directory(compute_plan_key)
try:
shutil.rmtree(chainkeys_directory)
except Exception:
logger.error(f'Cannot delete volume {chainkeys_directory}',
exc_info=True)
def prepare_volumes(subtuple_directory, tuple_type, compute_plan_key,
compute_plan_tag):
model_path = path.join(subtuple_directory, 'model')
output_model_path = path.join(subtuple_directory, 'output_model')
pred_path = path.join(subtuple_directory, 'pred')
export_path = path.join(subtuple_directory, 'export')
perf_path = path.join(subtuple_directory, 'perf')
opener_path = path.join(subtuple_directory, 'opener')
symlinks_volume = {}
data_path = path.join(subtuple_directory, 'data')
for subfolder in os.listdir(data_path):
real_path = os.path.realpath(os.path.join(data_path, subfolder))
symlinks_volume[real_path] = {'bind': f'{real_path}', 'mode': 'ro'}
for subtuple_folder in ['opener', 'model', 'metrics']:
for subitem in os.listdir(
path.join(subtuple_directory, subtuple_folder)):
real_path = os.path.realpath(
os.path.join(subtuple_directory, subtuple_folder, subitem))
if real_path != os.path.join(subtuple_directory, subtuple_folder,
subitem):
symlinks_volume[real_path] = {
'bind': f'{real_path}',
'mode': 'ro'
}
volumes = {
data_path: {
'bind': DATA_FOLDER,
'mode': 'ro'
},
opener_path: {
'bind': '/sandbox/opener',
'mode': 'ro'
}
}
if tuple_type == TESTTUPLE_TYPE:
volumes[pred_path] = {'bind': '/sandbox/pred', 'mode': 'rw'}
volumes[export_path] = {'bind': '/sandbox/export', 'mode': 'rw'}
volumes[perf_path] = {'bind': '/sandbox/perf', 'mode': 'rw'}
model_volume = {
model_path: {
'bind': MODEL_FOLDER,
'mode': 'ro'
},
output_model_path: {
'bind': OUTPUT_MODEL_FOLDER,
'mode': 'rw'
}
}
# local volume content shared by subtuples in compute plan
if compute_plan_key is not None:
subtuple_local_folder = path.join(subtuple_directory, 'local')
cp_local_folder = get_cp_local_folder(compute_plan_key)
# Copy content of compute plan local folder to subtuple local folder.
# It allows the compute plan subtuple to be retried without compromising the
# local data.
if os.path.exists(cp_local_folder):
distutils.dir_util.copy_tree(cp_local_folder,
subtuple_local_folder)
logger.info(
f'Content from compute plan local folder ({cp_local_folder}) '
f'has been copied to subtuple local folder ({subtuple_local_folder})'
)
mode = 'ro' if tuple_type == TESTTUPLE_TYPE else 'rw'
model_volume[subtuple_local_folder] = {
'bind': LOCAL_FOLDER,
'mode': mode
}
chainkeys_volume = {}
if compute_plan_key is not None and settings.TASK['CHAINKEYS_ENABLED']:
chainkeys_volume = prepare_chainkeys(compute_plan_key,
compute_plan_tag,
subtuple_directory)
return {**volumes, **symlinks_volume}, {**model_volume, **chainkeys_volume}
def do_task(channel_name, subtuple, tuple_type):
subtuple_directory = get_subtuple_directory(subtuple['key'])
# compute plan / federated learning variables
compute_plan_key = None
rank = None
compute_plan_tag = None
if 'compute_plan_key' in subtuple and subtuple['compute_plan_key']:
compute_plan_key = subtuple['compute_plan_key']
rank = int(subtuple['rank'])
compute_plan = get_object_from_ledger(channel_name, compute_plan_key,
'queryComputePlan')
compute_plan_tag = compute_plan['tag']
common_volumes, compute_volumes = prepare_volumes(subtuple_directory,
tuple_type,
compute_plan_key,
compute_plan_tag)
# Add node index to environment variable for the compute
node_index = os.getenv('NODE_INDEX')
if node_index:
environment = {'NODE_INDEX': node_index}
else:
environment = {}
# Use tag to tranfer or not performances and models
tag = subtuple.get("tag")
if tuple_type == TESTTUPLE_TYPE:
if tag and TAG_VALUE_FOR_TRANSFER_BUCKET in tag:
environment['TESTTUPLE_TAG'] = TAG_VALUE_FOR_TRANSFER_BUCKET
job_name = f'{tuple_type.replace("_", "-")}-{subtuple["key"][0:8]}-{TUPLE_COMMANDS[tuple_type]}'.lower(
)
command = generate_command(tuple_type, subtuple, rank)
# train or predict
compute_job(subtuple_key=subtuple["key"],
compute_plan_key=compute_plan_key,
dockerfile_path=subtuple_directory,
image_name=get_algo_image_name(subtuple['algo']['key']),
job_name=job_name,
volumes={
**common_volumes,
**compute_volumes
},
command=command,
remove_image=compute_plan_key is None
and not settings.TASK['CACHE_DOCKER_IMAGES'],
remove_container=settings.TASK['CLEAN_EXECUTION_ENVIRONMENT'],
capture_logs=settings.TASK['CAPTURE_LOGS'],
environment=environment)
# Handle model and result from tuple
models = save_models(subtuple_directory, tuple_type,
subtuple['key']) # Can be empty if testtuple
result = extract_result_from_models(tuple_type,
models) # Can be empty if testtuple
# local volume content shared by subtuples in compute plan
if compute_plan_key is not None:
subtuple_local_folder = path.join(subtuple_directory, 'local')
cp_local_folder = get_cp_local_folder(compute_plan_key)
# Copy content of subtuple local folder to compute plan local folder in order
# to fetch modifications made in the subtuple execution
distutils.dir_util.copy_tree(subtuple_local_folder, cp_local_folder)
logger.info(
f'Content from subtuple local folder ({subtuple_local_folder}) '
f'has been copied to compute plan local folder ({cp_local_folder})'
)
# Evaluation
if tuple_type == TESTTUPLE_TYPE:
# We set pred folder to ro during evalutation
pred_path = path.join(subtuple_directory, 'pred')
common_volumes[pred_path]['mode'] = 'ro'
# eval
compute_job(
subtuple_key=subtuple["key"],
compute_plan_key=compute_plan_key,
dockerfile_path=f'{subtuple_directory}/metrics',
image_name=f'substra/metrics_{subtuple["objective"]["key"][0:8]}'.
lower(),
job_name=
f'{tuple_type.replace("_", "-")}-{subtuple["key"][0:8]}-eval'.
lower(),
volumes=common_volumes,
command=f'--output-perf-path {OUTPUT_PERF_PATH}',
remove_image=compute_plan_key is None
and not settings.TASK['CACHE_DOCKER_IMAGES'],
remove_container=settings.TASK['CLEAN_EXECUTION_ENVIRONMENT'],
capture_logs=settings.TASK['CAPTURE_LOGS'],
environment=environment)
pred_path = path.join(subtuple_directory, 'pred')
export_path = path.join(subtuple_directory, 'export')
perf_path = path.join(subtuple_directory, 'perf')
# load performance
with open(path.join(perf_path, 'perf.json'), 'r') as perf_file:
perf = json.load(perf_file)
result['global_perf'] = perf['all']
if tag and TAG_VALUE_FOR_TRANSFER_BUCKET in tag:
transfer_to_bucket(subtuple['key'],
[pred_path, perf_path, export_path])
return result