async def judge_consume(self, handler_type): async with self.ws_connect(self.full_url('judge/consume-conn/websocket')) as ws: logger.info('Connected') async for msg in ws: request = json.loads(msg.data) await handler_type(self, request, ws).handle() logger.warning('Connection lost with code %d', ws.close_code)
async def do_submission(self): loop = get_event_loop() logger.info('Submission: %s, %s, %s', self.domain_id, self.pid, self.rid) cases_file_task = loop.create_task(cache_open(self.session, self.domain_id, self.pid)) package = await self.build() with await cases_file_task as cases_file: await self.judge(cases_file, package)
async def judge_consume(self, handler_type): async with self.ws_connect( self.full_url('judge/consume-conn/websocket')) as ws: logger.info('Connected') queue = Queue() async def worker(): try: while True: request = await queue.get() await handler_type(self, request, ws).handle() except CancelledError: raise except Exception as e: logger.exception(e) await ws.close() worker_task = get_event_loop().create_task(worker()) try: while True: queue.put_nowait(await ws.receive_json()) except TypeError: pass logger.warning('Connection lost with code %s', ws.close_code) worker_task.cancel() try: await worker_task except CancelledError: pass
async def do_pretest(self): loop = get_event_loop() logger.info('Pretest: %s, %s, %s', self.domain_id, self.pid, self.rid) cases_data_task = loop.create_task(self.session.record_pretest_data(self.rid)) package = await self.build() with BytesIO(await cases_data_task) as cases_file: await self.judge(cases_file, package)
async def _init(): parallelism = config.get('parallelism', 1) logger.info('Using parallelism: %d', parallelism) for sandbox in await create_sandboxes(parallelism): _sandbox_pool.put_nowait(sandbox) try: with open(_LANGS_FILE) as file: langs_config = yaml.load(file, Loader=yaml.RoundTripLoader) except FileNotFoundError: logger.error('Language file %s not found.', _LANGS_FILE) exit(1) for lang_name, lang_config in langs_config.items(): if lang_config['type'] == 'compiler': compiler = Compiler(lang_config['compiler_file'], shlex.split(lang_config['compiler_args']), lang_config['code_file'], lang_config['execute_file'], shlex.split(lang_config['execute_args'])) _langs[lang_name] = partial( _compiler_build, compiler, time_limit_ns=lang_config.get('time_limit_ms', DEFAULT_TIME_MS) * 1000000, memory_limit_bytes=lang_config.get('memory_limit_kb', DEFAULT_MEM_KB) * 1024, process_limit=lang_config.get('process_limit', PROCESS_LIMIT)) elif lang_config['type'] == 'interpreter': interpreter = Interpreter(lang_config['code_file'], lang_config['execute_file'], shlex.split(lang_config['execute_args'])) _langs[lang_name] = partial(_interpreter_build, interpreter) else: logger.error('Unknown type %s', lang_config['type'])
async def build(self, sandbox, *, output_file=None, cgroup_file=None, config=None): lang_config = config.get('lang') or {} loop = get_event_loop() compiler_file = lang_config.get('compiler_file') or self.compiler_file compiler_args = lang_config.get('compiler_args') or self.compiler_args status = await sandbox.call(SANDBOX_COMPILE, compiler_file, compiler_args, output_file, cgroup_file) if status: return None, status if 'runtime_files' in config: logger.info("Extracting runtime files to sandbox %s", sandbox.out_dir) await loop.run_in_executor(None, config['runtime_files'], sandbox.out_dir, False) package_dir = mkdtemp(prefix='jd4.package.') await loop.run_in_executor(None, copytree, sandbox.out_dir, path.join(package_dir, 'package')) execute_file = lang_config.get('execute_file') or self.execute_file execute_args = lang_config.get('execute_args') or self.execute_args return Package(package_dir, execute_file, execute_args), 0
def try_init_cgroup(): euid = geteuid() if euid == 0: logger.warning('Running as root') cgroups_to_init = list() if not (path.isdir(CPUACCT_CGROUP_ROOT) and access(CPUACCT_CGROUP_ROOT, W_OK)): cgroups_to_init.append(CPUACCT_CGROUP_ROOT) if not (path.isdir(MEMORY_CGROUP_ROOT) and access(MEMORY_CGROUP_ROOT, W_OK)): cgroups_to_init.append(MEMORY_CGROUP_ROOT) if not (path.isdir(PIDS_CGROUP_ROOT) and access(PIDS_CGROUP_ROOT, W_OK)): cgroups_to_init.append(PIDS_CGROUP_ROOT) if cgroups_to_init: if euid == 0: logger.info('Initializing cgroup: %s', ', '.join(cgroups_to_init)) for cgroup_to_init in cgroups_to_init: makedirs(cgroup_to_init, exist_ok=True) elif __stdin__.isatty(): logger.info('Initializing cgroup: %s', ', '.join(cgroups_to_init)) call([ 'sudo', 'sh', '-c', 'mkdir -p "{1}" && chown -R "{0}" "{1}"'.format( euid, '" "'.join(cgroups_to_init)) ]) else: logger.error('Cgroup not initialized')
async def update_problem_data(session): logger.info('Update problem data') result = await session.judge_datalist(config.get('last_update_at', 0)) for pid in result['pids']: await cache_invalidate(pid['domain_id'], str(pid['pid'])) logger.debug('Invalidated %s/%s', pid['domain_id'], str(pid['pid'])) config['last_update_at'] = result['time'] await save_config()
async def pretest(self): domain_id = self.request.pop('domain_id') pid = self.request.pop('pid') rid = self.request.pop('rid') logger.info('Pretest: %s, %s, %s', domain_id, pid, rid) cases_data, package = await gather( self.session.record_pretest_data(rid), self.build()) await self.judge(BytesIO(cases_data), package)
async def do_submission_remote(self): loop = get_event_loop() logger.info('Submission Remote: %s, %s, %s ( %s, %s )', self.domain_id, self.pid, self.rid, self.remote['orig_oj'], self.remote['orig_id']) #cases_file_task = loop.create_task(cache_open(self.session, self.domain_id, self.pid)) #package = await self.build() #with await cases_file_task as cases_file: await self.judge_remote()
async def login_if_needed(self, uname, password): try: await self.judge_noop() logger.info('Session is valid') except VJ4Error as e: if e.name == 'PrivilegeError': await self.login(uname, password) await save_cookies() else: raise
async def login_if_needed(self, uname, password): try: await self.judge_noop() logger.info('Session is valid') except VJ4Error as e: if e.name == 'PrivilegeError': await self.login(uname, password) await get_event_loop().run_in_executor( None, lambda: _COOKIE_JAR.save(_COOKIES_FILE)) else: raise
def _init(): parallelism = config.get('parallelism', 2) if parallelism < 2: logger.warning( 'Parallelism less than 2, custom judge will not be supported.') logger.info('Using parallelism: %d', parallelism) sandboxes_task = create_sandboxes(parallelism) global _lock, _queue _lock = Lock() _queue = LifoQueue() put_sandbox(*get_event_loop().run_until_complete(sandboxes_task))
async def submission(self): domain_id = self.request.pop('domain_id') pid = self.request.pop('pid') rid = self.request.pop('rid') logger.info('Submission: %s, %s, %s', domain_id, pid, rid) cases_file, package = await gather( cache_open(self.session, domain_id, pid), self.build()) try: await self.judge(cases_file, package) finally: cases_file.close()
async def daemon(): async with VJ4Session(config['server_url']) as session: while True: try: await session.login_if_needed(config['uname'], config['password']) await update_problem_data(session) await session.judge_consume(JudgeHandler) except Exception as e: logger.exception(e) logger.info('Retrying after %d seconds', RETRY_DELAY_SEC) await sleep(RETRY_DELAY_SEC)
async def build(self): lang = self.request.pop('lang') self.next(status=STATUS_COMPILING) build_fn = langs.get(lang) if not build_fn: raise SystemError('Unsupported language: {}'.format(lang)) package, message = await build_fn(sandbox, self.request.pop('code').encode()) self.next(compiler_text=message) if not package: logger.info('Compile error: %s', message) raise CompileError(message) return package
async def record_pretest_data(self, rid): logger.info('Getting pretest data: %s', rid) async with self.get(self.full_url('records', rid, 'data'), headers={'accept': 'application/json'}) as response: if response.content_type == 'application/json': response_dict = await response.json() if 'error' in response_dict: error = response_dict['error'] raise VJ4Error(error.get('name', 'unknown'), error.get('message', ''), *error.get('args', [])) raise Exception('unexpected response') if response.status != 200: raise Exception('http error ' + str(response.status)) return await response.read()
def do_lang(self, lang, code): package, message, time_usage_ns, memory_usage_bytes = \ run(build(lang, code)) self.assertIsNotNone(package, 'Compile failed: ' + message) logger.info('Compiled successfully in %d ms time, %d kb memory', time_usage_ns // 1000000, memory_usage_bytes // 1024) if message: logger.warning('Compiler output is not empty: %s', message) for case in self.cases: status, score, time_usage_ns, memory_usage_bytes, stderr = \ run(case.judge(package)) self.assertEqual(status, STATUS_ACCEPTED) self.assertEqual(score, 10) self.assertEqual(stderr, b'') logger.info('Accepted: %d ms time, %d kb memory', time_usage_ns // 1000000, memory_usage_bytes // 1024)
async def daemon(): try_init_cgroup() async with VJ4Session(config['server_url']) as session: while True: try: await session.login_if_needed(config['uname'], config['password']) done, pending = await wait([do_judge(session), do_noop(session)], return_when=FIRST_COMPLETED) for task in pending: task.cancel() await gather(*done) except Exception as e: logger.exception(e) logger.info('Retrying after %d seconds', RETRY_DELAY_SEC) await sleep(RETRY_DELAY_SEC)
async def main(): try_init_cgroup() sandbox = await create_sandbox() gcc = Compiler('/usr/bin/gcc', ['gcc', '-std=c99', '-o', '/out/foo', '/in/foo.c'], 'foo.c', 'foo', ['foo']) await gcc.prepare(sandbox, b"""#include <stdio.h> int main(void) { int a, b; scanf("%d%d", &a, &b); printf("%d\\n", a + b); }""") package, _ = await gcc.build(sandbox) for case in read_legacy_cases('examples/1000.zip'): logger.info(await case.judge(sandbox, package)) for i in range(10): logger.info(await APlusBCase(randint(0, 32767), randint(0, 32767)).judge(sandbox, package))
async def prepare(self, sandbox, code, code_type, config): loop = get_event_loop() await sandbox.reset() if code_type == CODE_TYPE_TEXT: await loop.run_in_executor(None, write_binary_file, path.join(sandbox.in_dir, self.code_file), code) elif code_type == CODE_TYPE_TAR: await loop.run_in_executor(None, extract_tar_file, code, sandbox.in_dir) if 'compile_time_files' in config: logger.info("Extracting compile time files to sandbox %s", sandbox.in_dir) await loop.run_in_executor(None, config['compile_time_files'], sandbox.in_dir)
async def do_record(self): self.tag = self.request.pop('tag') self.type = self.request.pop('type') self.domain_id = self.request.pop('domain_id') self.pid = self.request.pop('pid') self.rid = self.request.pop('rid') self.lang = self.request.pop('lang') self.code_type = self.request.pop('code_type') self.judge_category = self.request.pop('judge_category') self.judge_category = self.judge_category and self.judge_category.split( ',') or [] if self.code_type == CODE_TYPE_TEXT: self.code = self.request.pop('code').encode() else: self.code = path.join(mkdtemp(prefix='jd4.code.')) await self.session.record_code_data(self.rid, path.join(self.code, 'code')) logger.info('code dir: %s', self.code) # TODO(tc-imba) pretest not supported try: await self.prepare() if self.type == 0: await self.do_submission() elif self.type == 1: await self.do_pretest() else: raise Exception('Unsupported type: {}'.format(self.type)) except CompileError: self.end(status=STATUS_COMPILE_ERROR, score=0, time_ms=0, memory_kb=0) except ClientError: raise except Exception as e: logger.exception(e) self.next(judge_text=repr(e)) self.end(status=STATUS_SYSTEM_ERROR, score=0, time_ms=0, memory_kb=0)
async def problem_data(self, domain_id, pid, save_path): logger.info('Getting problem data: %s, %s', domain_id, pid) loop = get_event_loop() async with self.get(self.full_url('d', domain_id, 'p', pid, 'data'), headers={'accept': 'application/json'}) as response: if response.content_type == 'application/json': response_dict = await response.json() if 'error' in response_dict: error = response_dict['error'] raise VJ4Error(error.get('name', 'unknown'), error.get('message', ''), *error.get('args', [])) raise Exception('unexpected response') if response.status != 200: raise Exception('http error ' + str(response.status)) with open(save_path, 'wb') as save_file: while True: buffer = await response.content.read(CHUNK_SIZE) if not buffer: break await loop.run_in_executor(None, save_file.write, buffer)
def do_status(self, expected_status, expected_score, code): package, message, time_usage_ns, memory_usage_bytes = \ run(build('c', code)) self.assertIsNotNone(package, 'Compile failed: ' + message) logger.info('Compiled successfully in %d ms time, %d kb memory', time_usage_ns // 1000000, memory_usage_bytes // 1024) if message: logger.warning('Compiler output is not empty: %s', message) total_status = STATUS_ACCEPTED total_score = 0 for case in self.cases: status, score, time_usage_ns, memory_usage_bytes, stderr = \ run(case.judge(package)) total_status = max(total_status, status) total_score += score self.assertEqual(status, STATUS_ACCEPTED) self.assertEqual(score, 25) self.assertEqual(stderr, b'') logger.info('Accepted: %d ms time, %d kb memory', time_usage_ns // 1000000, memory_usage_bytes // 1024) self.assertEqual(total_status, expected_status) self.assertEqual(total_score, expected_score)
async def judge(self, cases_file, package): self.next(status=STATUS_JUDGING, progress=0) cases = list(read_legacy_cases(cases_file)) total_status = 0 total_score = 0 total_time_usage_ns = 0 total_memory_usage_bytes = 0 for index, case in enumerate(cases): status, score, time_usage_ns, memory_usage_bytes, stderr = await case.judge( sandbox, package) self.next(status=STATUS_JUDGING, case={ 'status': status, 'score': score, 'time_ms': time_usage_ns // 1000000, 'memory_kb': memory_usage_bytes // 1024, 'judge_text': stderr.decode(encoding='utf-8', errors='replace') }, progress=(index + 1) * 100 // len(cases)) logger.debug('Case %d: %d, %g, %g, %g, %s', index, status, score, time_usage_ns / 1000000, memory_usage_bytes / 1024, stderr) total_status = max(total_status, status) total_score += score total_time_usage_ns += time_usage_ns total_memory_usage_bytes = max(total_memory_usage_bytes, memory_usage_bytes) self.end(status=total_status, score=total_score, time_ms=total_time_usage_ns // 1000000, memory_kb=total_memory_usage_bytes // 1024) logger.info('Total: %d, %g, %g, %g', total_status, total_score, total_time_usage_ns / 1000000, total_memory_usage_bytes / 1024)
async def daemon(): # 尝试初始化cgroups,其名称源自控制组群(control groups)的简写,是Linux内核的一个功能,用来限制、控制与分离一个进程组的资源(如CPU、内存、磁盘输入输出等)。 try_init_cgroup() # 创建异步回话连接jv4服务端 # config['server_url'] 获取配置文件数据 async with VJ4Session(config['server_url']) as session: while True: try: # 登陆JV4服务器 await session.login_if_needed(config['uname'], config['password']) # 运行两个协程 done, pending = await wait([do_judge(session), do_noop(session)], return_when=FIRST_COMPLETED) for task in pending: task.cancel() await gather(*done) except Exception as e: logger.exception(e) # 打印等待log logger.info('Retrying after %d seconds', RETRY_DELAY_SEC) # 等待30秒 await sleep(RETRY_DELAY_SEC)
async def main(): sandbox, = await sandboxes_task logger.info('sandbox_dir: %s', sandbox.sandbox_dir) logger.info('return value: %d', await sandbox.backdoor())
async def login(self, uname, password): logger.info('Login') await self.post_json('login', uname=uname, password=password)
async def main(): sandbox = await create_sandbox() logger.info('sandbox_dir: %s', sandbox.sandbox_dir) logger.info('return value: %d', await sandbox.backdoor())
async def do_noop(session): while True: await sleep(3600) logger.info('Updating session') await session.judge_noop()