def discard_subtasks(
work_dir: dirutils.RequestorTaskDir,
subtask_ids: List[str],
) -> List[str]:
task_manager = DBTaskManager(work_dir)
for subtask_id in subtask_ids:
task_manager.update_subtask_status(subtask_id, SubtaskStatus.ABORTED)
return subtask_ids
async def verify(
work_dir: dirutils.RequestorTaskDir,
subtask_id: str,
) -> Tuple[enums.VerifyResult, Optional[str]]:
with open(work_dir / 'task_params.json', 'r') as f:
task_params = json.load(f)
with open(work_dir / f'subtask{subtask_id}.json', 'r') as f:
params = json.load(f)
subtask_work_dir = work_dir / f'subtask{subtask_id}'
subtask_work_dir.mkdir()
subtask_results_dir = subtask_work_dir / 'results'
subtask_results_dir.mkdir()
subtask_output_dir = subtask_work_dir / 'output'
subtask_output_dir.mkdir()
subtask_outputs_dir = work_dir.subtask_outputs_dir(subtask_id)
zip_file_path = subtask_outputs_dir / f'{subtask_id}.zip'
with zipfile.ZipFile(zip_file_path, 'r') as zip_file:
zip_file.extractall(subtask_results_dir)
task_manager = DBTaskManager(work_dir)
part_num = task_manager.get_part_num(subtask_id)
task_manager.update_subtask_status(subtask_id, SubtaskStatus.VERIFYING)
dir_contents = subtask_results_dir.iterdir()
verdict = await verifier.verify(
[str(entry) for entry in dir_contents if entry.is_file()],
params['borders'],
work_dir.task_inputs_dir / params['scene_file'],
params['resolution'],
params['samples'],
params['frames'],
params['output_format'],
mounted_paths={
'OUTPUT_DIR': str(subtask_output_dir),
'WORK_DIR': str(subtask_work_dir),
}
)
print("Verdict:", verdict)
if not verdict:
task_manager.update_subtask_status(
subtask_id,
SubtaskStatus.FAILURE)
# pylint: disable=fixme
# TODO: provide some extra info why verification failed
return enums.VerifyResult.FAILURE, None
task_manager.update_subtask_status(subtask_id, SubtaskStatus.SUCCESS)
_collect_results(
task_manager,
part_num,
task_params,
params,
work_dir,
subtask_results_dir,
work_dir.task_outputs_dir,
)
return enums.VerifyResult.SUCCESS, None
async def create_task(
work_dir: dirutils.RequestorTaskDir,
max_subtasks_count: int,
params: dict
) -> structs.Task:
frame_count = len(utils.string_to_frames(params['frames']))
if max_subtasks_count <= frame_count:
subtasks_count = max_subtasks_count
else:
subtasks_count = max_subtasks_count // frame_count * frame_count
params['subtasks_count'] = subtasks_count
if not utils.get_scene_file_from_resources(params['resources']):
raise RuntimeError("Scene file not found in resources")
min_memory = await test_task(work_dir, params)
with zipfile.ZipFile(work_dir.subtask_inputs_dir / '0.zip', 'w') as zipf:
for resource in params['resources']:
resource_path = work_dir.task_inputs_dir / resource
zipf.write(resource_path, resource)
with open(work_dir / 'task_params.json', 'w') as f:
json.dump(params, f)
DBTaskManager(work_dir).create_task(subtasks_count)
return envs.create_docker_cpu_task(
image=constants.DOCKER_IMAGE,
tag=constants.VERSION,
inf=Infrastructure(min_memory_mib=min_memory // (1024 * 1024))
)
def task_manager(tmpdir):
manager = DBTaskManager(Path(tmpdir))
try:
yield manager
finally:
database.close()
shutil.rmtree(tmpdir)
def get_next_subtask(
work_dir: dirutils.RequestorTaskDir,
subtask_id: str,
) -> structs.Subtask:
with open(work_dir / 'task_params.json', 'r') as f:
task_params = json.load(f)
task_manager = DBTaskManager(work_dir)
part_num = task_manager.get_next_computable_part_num()
if part_num is None:
raise Exception('No available subtasks at the moment')
print(f'Part number: {part_num}, id: {subtask_id}')
task_manager.start_subtask(part_num, subtask_id)
scene_file = utils.get_scene_file_from_resources(task_params['resources'])
all_frames = utils.string_to_frames(task_params['frames'])
frames, parts = _choose_frames(
all_frames,
part_num,
task_params['subtasks_count'],
)
min_y = (part_num % parts) / parts
max_y = (part_num % parts + 1) / parts
resources = ['0.zip']
borders: List[float] = [0.0, min_y, 1.0, max_y]
subtask_params = {
"scene_file": scene_file,
"resolution": task_params['resolution'],
"use_compositing": False,
"samples": 0,
"frames": frames,
"output_format": task_params['format'],
"borders": borders,
"resources": resources,
}
with open(work_dir / f'subtask{subtask_id}.json', 'w') as f:
json.dump(subtask_params, f)
return structs.Subtask(
params=subtask_params,
resources=resources,
)
def _collect_results(
task_manager: DBTaskManager,
part_num: int,
task_params: dict,
params: dict,
work_dir: Path,
subtask_results_dir: Path,
results_dir: Path) -> None:
frames = utils.string_to_frames(task_params['frames'])
frame_count = len(frames)
out_format = get_expected_extension(params['output_format'])
parts = task_params['subtasks_count'] // frame_count
if parts <= 1:
for frame in params['frames']:
shutil.copy2(
subtask_results_dir / f'result{frame:04d}.{out_format}',
results_dir / f'result{frame:04d}.{out_format}',
)
return
frame_id = part_num // parts
frame = frames[frame_id]
subtasks_nums = list(range(frame_id * parts, (frame_id + 1) * parts))
subtasks_statuses = task_manager.get_subtasks_statuses(subtasks_nums)
all_finished = all([
s[0] == SubtaskStatus.SUCCESS for s in subtasks_statuses.values()
])
if not all_finished:
print('Not all finished, waiting for more results')
return
print('All finished, collecting results')
collector = RenderingTaskCollector(
width=params['resolution'][0],
height=params['resolution'][1],
)
for i in subtasks_nums[::-1]:
result_dir = work_dir / f'subtask{subtasks_statuses[i][1]}' / 'results'
result_img = result_dir / f'result{frame:04d}.{out_format}'
print(f'result_dir: {result_dir}')
for result_file in result_dir.iterdir():
print(f'result_candidate: {result_file}')
print(f"result_img:{result_img.exists()}")
print(f"result_img.size:{result_img.stat()}")
collector.add_img_file(str(result_img))
image = collector.finalize()
if not image:
raise RuntimeError("No accepted image files")
with image as img:
img.save_with_extension(
results_dir / f'result{frame:04d}',
out_format)
def abort_task(work_dir: dirutils.RequestorTaskDir) -> None:
DBTaskManager(work_dir).abort_task()
def has_pending_subtasks(work_dir: dirutils.RequestorTaskDir) -> bool:
return DBTaskManager(work_dir).get_next_computable_part_num() is not None
def abort_subtask(work_dir: dirutils.RequestorTaskDir,
subtask_id: str) -> None:
DBTaskManager(work_dir).update_subtask_status(subtask_id,
SubtaskStatus.ABORTED)