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)