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()
