def task_put(self, task):
limits = KolejkaLimits()
limits.cpus = self.config.cpus
limits.memory = self.config.memory
limits.pids = self.config.pids
limits.storage = self.config.storage
limits.time = self.config.time
limits.network = self.config.network
task.limits.update(limits)
if not self.instance_session:
self.login()
for k, f in task.files.items():
if not f.reference or not self.blob_check(
blob_reference=f.reference):
f.reference = None
if f.path:
f.reference = self.blob_put(os.path.join(
task.path, f.path))['key']
else:
raise
info = self.post('/task/task/', data=json.dumps(task.dump()))
if info.status_code == 200:
task = KolejkaTask(None)
task.load(info.json()['task'])
return task
else:
print(info)
print(info.text)
def task_put(self, task):
limits = KolejkaLimits()
limits.cpus = self.config.cpus
limits.memory = self.config.memory
limits.swap = self.config.swap
limits.pids = self.config.pids
limits.storage = self.config.storage
limits.image = self.config.image
limits.workspace = self.config.workspace
limits.time = self.config.time
limits.network = self.config.network
limits.gpus = self.config.gpus
task.limits.update(limits)
if not self.instance_session:
self.login()
for f in task.files.values():
if not f.reference or not self.blob_check(
blob_reference=f.reference):
assert f.path
f.reference = self.blob_put(os.path.join(task.path,
f.path))['key']
response = self.post('/task/task/', data=json.dumps(task.dump()))
task = KolejkaTask(None)
task.load(response.json()['task'])
return task
def dequeue(self, concurency, limits, tags):
if not self.instance_session:
self.login()
response = self.post('/queue/dequeue/',
data=json.dumps({
'concurency': concurency,
'limits': limits.dump(),
'tags': tags
}))
ts = response.json()['tasks']
tasks = list()
for t in ts:
tt = KolejkaTask(None)
tt.load(t)
tasks.append(tt)
return tasks
def execute(args):
kolejka_config(args=args)
client = KolejkaClient()
task = KolejkaTask(args.task)
response = client.task_put(task)
while True:
time.sleep(args.interval)
result = client.result_get(response.id, args.result)
if result is not None:
if args.consume:
shutil.rmtree(args.task)
return result
def execute(args):
kolejka_config(args=args)
client = KolejkaClient()
task = KolejkaTask(args.task)
response = client.task_put(task)
while True:
client.session.close()
time.sleep(args.interval)
try:
result = client.result_get(response.id, args.result)
except KolejkaClientObjectNotFoundError:
continue
if args.consume:
shutil.rmtree(args.task)
break
def task_get(self, task_key, task_path):
if isinstance(task_key, KolejkaTask):
task_key = task_key.id
response = self.get('/task/task/{}/'.format(task_key))
os.makedirs(task_path, exist_ok=True)
task = KolejkaTask(task_path)
task.load(response.json()['task'])
for k, f in task.files.items():
self.blob_get(os.path.join(task.path, k), f.reference)
f.path = k
task.commit()
return task
def task(self, task_path=None):
task = KolejkaTask(task_path)
task.load(self.description)
return task
def kolejka_task(task_dir,
tests,
solution,
judgepy,
exist_ok=False,
debug=False):
kolejka_image = None
kolejka_requires = set()
kolejka_exclusive = False
kolejka_collects = dict()
kolejka_stdout = 'console_stdout.txt'
kolejka_stderr = 'console_stderr.txt'
kolejka_limits = {}
kolejka_system = '--observer'
for test_id, test in tests.items():
kolejka_opts = test.get('kolejka', dict())
if 'image' in kolejka_opts:
if kolejka_image and kolejka_image != kolejka_opts['image']:
raise ValueError
kolejka_image = kolejka_opts['image']
for req in kolejka_opts.get('requires', []):
kolejka_requires.add(req)
kolejka_exclusive = kolejka_exclusive or bool(
kolejka_opts.get('exclusive', False))
kolejka_collects[test_id] = kolejka_opts.get('collect', [])
kolejka_collects[test_id].append(config.SATORI_RESULT + '/**')
test_limits = kolejka_opts.get('limits', {})
if 'time' in test_limits:
kolejka_limits['time'] = kolejka_limits.get(
'time', datetime.timedelta(seconds=1)) + parse_time(
test_limits['time'])
if 'memory' in test_limits:
kolejka_limits['memory'] = max(kolejka_limits.get('memory', 0),
parse_memory(test_limits['memory']))
if 'swap' in test_limits:
kolejka_limits['swap'] = max(kolejka_limits.get('swap', 0),
parse_memory(test_limits['swap']))
if 'cpus' in test_limits:
kolejka_limits['cpus'] = max(kolejka_limits.get('cpus', 1),
int(test_limits['cpus']))
if 'network' in test_limits:
kolejka_limits['network'] = kolejka_limits.get(
'network', False) or bool(test_limits['network'])
if 'pids' in test_limits:
kolejka_limits['pids'] = max(kolejka_limits.get('pids', 16),
int(test_limits['pids']))
if 'storage' in test_limits:
kolejka_limits['storage'] = kolejka_limits.get(
'storage', 0) + parse_memory(test_limits['storage'])
if 'workspace' in test_limits:
kolejka_limits['workspace'] = kolejka_limits.get(
'workspace', 0) + parse_memory(test_limits['workspace'])
kolejka_image = kolejka_image or 'kolejka/satori:judge'
from kolejka.common import KolejkaTask, KolejkaLimits
task_dir = pathlib.Path(task_dir).resolve()
solution = pathlib.Path(solution).resolve()
judgepy = pathlib.Path(judgepy).resolve()
task_dir.mkdir(parents=True, exist_ok=exist_ok)
test_dir = pathlib.PurePath('tests')
(task_dir / test_dir).mkdir()
solution_dir = pathlib.PurePath('solution')
(task_dir / solution_dir).mkdir()
results_dir = pathlib.PurePath('results')
results_yaml = pathlib.PurePath('results.yaml')
kolejka_collect = []
kolejka_collect += [{'glob': str(results_dir / results_yaml)}]
for test_id, collects in kolejka_collects.items():
for collect in collects:
kolejka_collect += [{
'glob':
str(results_dir / str(test_id) / str(collect))
}]
if debug:
kolejka_collect += [
{
'glob': str(test_dir) + '/**'
},
{
'glob': str(solution_dir) + '/**'
},
{
'glob': str(results_dir) + '/**'
},
]
tests_yaml = test_dir / 'tests.yaml'
solution_path = solution_dir / solution.name
(task_dir / solution_path).symlink_to(solution)
judgepy_path = pathlib.PurePath('judge.py')
(task_dir / judgepy_path).symlink_to(judgepy)
lib_path = pathlib.PurePath(config.DISTRIBUTION_PATH)
(task_dir / lib_path).symlink_to(judgepy.parent / lib_path)
if not (judgepy.parent / lib_path).is_file():
logging.warning(
'Kolejka Judge library not present in {}. Try running library update.'
.format(judgepy.parent / lib_path))
task_args = [
'python3',
str(judgepy_path),
]
if debug:
task_args += [
'--debug',
]
task_args += [
'execute',
kolejka_system,
str(tests_yaml),
str(solution_path),
str(results_dir),
'--results',
str(results_yaml),
]
input_map = dict()
class collect:
def __init__(self, input_map):
self.input_count = 0
self.input_map = input_map
def __call__(self, a):
if isinstance(a, InputPath):
a = pathlib.Path(a.path)
if isinstance(a, pathlib.Path):
if a in self.input_map:
return self.input_map[a]
self.input_count += 1
input_path = task_dir / test_dir / (
'%03d' % (self.input_count, )) / a.name
input_path.parent.mkdir(exist_ok=True, parents=True)
input_path.symlink_to(a)
input_path = input_path.relative_to(task_dir)
self.input_map[a] = input_path
return input_path
if isinstance(a, list):
return [self(e) for e in a]
if isinstance(a, dict):
return dict([(self(k), self(v)) for k, v in a.items()])
return a
tests = collect(input_map)(tests)
ctxyaml_dump(tests, task_dir / tests_yaml, work_dir=task_dir)
task = KolejkaTask(
str(task_dir),
image=kolejka_image,
requires=list(kolejka_requires),
exclusive=kolejka_exclusive,
limits=kolejka_limits,
args=task_args,
stdout=kolejka_stdout,
stderr=kolejka_stderr,
files=dict([(str(p), None) for p in [
tests_yaml,
solution_path,
judgepy_path,
lib_path,
] + list(input_map.values())]),
collect=kolejka_collect,
)
task.commit()
def stage0(task_path, result_path, temp_path=None, consume_task_folder=False):
config = worker_config()
cgs = ControlGroupSystem()
task = KolejkaTask(task_path)
if not task.id:
task.id = uuid.uuid4().hex
logging.warning('Assigned id {} to the task'.format(task.id))
if not task.image:
logging.error('Task does not define system image')
sys.exit(1)
if not task.args:
logging.error('Task does not define args')
sys.exit(1)
if not task.files.is_local:
logging.error('Task contains non-local files')
sys.exit(1)
limits = KolejkaLimits()
limits.cpus = config.cpus
limits.memory = config.memory
limits.swap = config.swap
limits.pids = config.pids
limits.storage = config.storage
limits.image = config.image
limits.workspace = config.workspace
limits.time = config.time
limits.network = config.network
limits.gpus = config.gpus
task.limits.update(limits)
docker_task = 'kolejka_worker_{}'.format(task.id)
docker_cleanup = [
['docker', 'kill', docker_task],
['docker', 'rm', docker_task],
]
with tempfile.TemporaryDirectory(dir=temp_path) as jailed_path:
#TODO jailed_path size remains unlimited?
logging.debug('Using {} as temporary directory'.format(jailed_path))
jailed_task_path = os.path.join(jailed_path, 'task')
os.makedirs(jailed_task_path, exist_ok=True)
jailed_result_path = os.path.join(jailed_path, 'result')
os.makedirs(jailed_result_path, exist_ok=True)
jailed = KolejkaTask(os.path.join(jailed_path, 'task'))
jailed.load(task.dump())
jailed.files.clear()
volumes = list()
check_python_volume()
if os.path.exists(OBSERVER_SOCKET):
volumes.append((OBSERVER_SOCKET, OBSERVER_SOCKET, 'rw'))
else:
logging.warning('Observer is not running.')
volumes.append(
(jailed_result_path, os.path.join(WORKER_DIRECTORY,
'result'), 'rw'))
for key, val in task.files.items():
if key != TASK_SPEC:
src_path = os.path.join(task.path, val.path)
dst_path = os.path.join(jailed_path, 'task', key)
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
if consume_task_folder:
shutil.move(src_path, dst_path)
else:
shutil.copy(src_path, dst_path)
jailed.files.add(key)
jailed.files.add(TASK_SPEC)
#jailed.limits = KolejkaLimits() #TODO: Task is limited by docker, no need to limit it again?
jailed.commit()
volumes.append((jailed.path, os.path.join(WORKER_DIRECTORY,
'task'), 'rw'))
if consume_task_folder:
try:
shutil.rmtree(task_path)
except:
logging.warning('Failed to remove {}'.format(task_path))
pass
for spath in [os.path.dirname(__file__)]:
stage1 = os.path.join(spath, 'stage1.sh')
if os.path.isfile(stage1):
volumes.append(
(stage1, os.path.join(WORKER_DIRECTORY,
'stage1.sh'), 'ro'))
break
for spath in [os.path.dirname(__file__)]:
stage2 = os.path.join(spath, 'stage2.py')
if os.path.isfile(stage2):
volumes.append(
(stage2, os.path.join(WORKER_DIRECTORY,
'stage2.py'), 'ro'))
break
docker_call = ['docker', 'run']
docker_call += ['--detach']
docker_call += ['--name', docker_task]
docker_call += [
'--entrypoint',
os.path.join(WORKER_DIRECTORY, 'stage1.sh')
]
for key, val in task.environment.items():
docker_call += ['--env', '{}={}'.format(key, val)]
docker_call += ['--hostname', WORKER_HOSTNAME]
docker_call += ['--init']
if task.limits.cpus is not None:
docker_call += [
'--cpuset-cpus', ','.join([
str(c) for c in cgs.limited_cpuset(cgs.full_cpuset(
), task.limits.cpus, task.limits.cpus_offset)
])
]
if task.limits.gpus is not None and task.limits.gpus > 0:
check_gpu_runtime_availability()
gpus = ','.join(
map(
str,
limited_gpuset(full_gpuset(), task.limits.gpus,
task.limits.gpus_offset)))
docker_call += [
'--runtime=nvidia', '--shm-size=1g', '--gpus',
f'"device={gpus}"'
]
if task.limits.memory is not None:
docker_call += ['--memory', str(task.limits.memory)]
if task.limits.swap is not None:
docker_call += [
'--memory-swap',
str(task.limits.memory + task.limits.swap)
]
if task.limits.storage is not None:
docker_info_run = subprocess.run(
['docker', 'system', 'info', '--format', '{{json .Driver}}'],
stdout=subprocess.PIPE,
check=True)
storage_driver = str(
json.loads(str(docker_info_run.stdout, 'utf-8')))
if storage_driver == 'overlay2':
docker_info_run = subprocess.run([
'docker', 'system', 'info', '--format',
'{{json .DriverStatus}}'
],
stdout=subprocess.PIPE,
check=True)
storage_fs = dict(
json.loads(str(docker_info_run.stdout,
'utf-8')))['Backing Filesystem']
if storage_fs in ['xfs']:
storage_limit = task.limits.storage
docker_call += [
'--storage-opt', 'size=' + str(storage_limit)
]
else:
logging.warning(
"Storage limit on {} ({}) is not supported".format(
storage_driver, storage_fs))
else:
logging.warning("Storage limit on {} is not supported".format(
storage_driver))
if task.limits.network is not None:
if not task.limits.network:
docker_call += ['--network=none']
docker_call += ['--cap-add', 'SYS_NICE']
if task.limits.pids is not None:
docker_call += ['--pids-limit', str(task.limits.pids)]
if task.limits.time is not None:
docker_call += [
'--stop-timeout',
str(int(math.ceil(task.limits.time.total_seconds())))
]
docker_call += [
'--volume',
'{}:{}:{}'.format(WORKER_PYTHON_VOLUME,
os.path.join(WORKER_DIRECTORY, 'python3'), 'ro')
]
for v in volumes:
docker_call += [
'--volume', '{}:{}:{}'.format(os.path.realpath(v[0]), v[1],
v[2])
]
docker_call += ['--workdir', WORKER_DIRECTORY]
docker_image = task.image
docker_call += [docker_image]
docker_call += ['--consume']
if config.debug:
docker_call += ['--debug']
if config.verbose:
docker_call += ['--verbose']
docker_call += [os.path.join(WORKER_DIRECTORY, 'task')]
docker_call += [os.path.join(WORKER_DIRECTORY, 'result')]
logging.debug('Docker call : {}'.format(docker_call))
pull_image = config.pull
if not pull_image:
docker_inspect_run = subprocess.run(
['docker', 'image', 'inspect', docker_image],
stdout=subprocess.DEVNULL,
stderr=subprocess.STDOUT)
if docker_inspect_run.returncode != 0:
pull_image = True
if pull_image:
subprocess.run(['docker', 'pull', docker_image], check=True)
for docker_clean in docker_cleanup:
silent_call(docker_clean)
if os.path.exists(result_path):
shutil.rmtree(result_path)
os.makedirs(result_path, exist_ok=True)
result = KolejkaResult(result_path)
result.id = task.id
result.limits = task.limits
result.stdout = task.stdout
result.stderr = task.stderr
start_time = datetime.datetime.now()
docker_run = subprocess.run(docker_call, stdout=subprocess.PIPE)
cid = str(docker_run.stdout, 'utf-8').strip()
logging.info('Started container {}'.format(cid))
try:
if task.limits.gpus is not None and task.limits.gpus > 0:
result.stats.update(
gpu_stats(gpus=limited_gpuset(full_gpuset(
), task.limits.gpus, task.limits.gpus_offset)))
except:
pass
time.sleep(0.1)
while True:
try:
docker_state_run = subprocess.run(
['docker', 'inspect', '--format', '{{json .State}}', cid],
stdout=subprocess.PIPE)
state = json.loads(str(docker_state_run.stdout, 'utf-8'))
except:
break
try:
result.stats.update(cgs.name_stats(cid))
if task.limits.gpus is not None and task.limits.gpus > 0:
result.stats.update(
gpu_stats(gpus=limited_gpuset(full_gpuset(
), task.limits.gpus, task.limits.gpus_offset)))
except:
pass
time.sleep(0.1)
if not state['Running']:
result.result = state['ExitCode']
try:
result.stats.time = dateutil.parser.parse(
state['FinishedAt']) - dateutil.parser.parse(
state['StartedAt'])
except:
result.stats.time = None
break
if task.limits.time is not None and datetime.datetime.now(
) - start_time > task.limits.time + datetime.timedelta(seconds=2):
docker_kill_run = subprocess.run(
['docker', 'kill', docker_task])
subprocess.run(['docker', 'logs', cid], stdout=subprocess.PIPE)
try:
summary = KolejkaResult(jailed_result_path)
result.stats.update(summary.stats)
except:
pass
stop_time = datetime.datetime.now()
if result.stats.time is None:
result.stats.time = stop_time - start_time
result.stats.pids.usage = None
result.stats.memory.usage = None
result.stats.memory.swap = None
for dirpath, dirnames, filenames in os.walk(jailed_result_path):
for filename in filenames:
abspath = os.path.join(dirpath, filename)
realpath = os.path.realpath(abspath)
if realpath.startswith(
os.path.realpath(jailed_result_path) + '/'):
relpath = abspath[len(jailed_result_path) + 1:]
if relpath != RESULT_SPEC:
destpath = os.path.join(result.path, relpath)
os.makedirs(os.path.dirname(destpath), exist_ok=True)
shutil.move(realpath, destpath)
os.chmod(destpath, 0o640)
result.files.add(relpath)
result.commit()
os.chmod(result.spec_path, 0o640)
for docker_clean in docker_cleanup:
silent_call(docker_clean)
def execute(args):
kolejka_config(args=args)
client = KolejkaClient()
task = KolejkaTask(args.task)
response = client.task_put(task)
print(response.id)
def task(request, key=''):
if request.method == 'POST':
if key != '':
return HttpResponseForbidden()
if not request.user.has_perm('task.add_task'):
return HttpResponseForbidden()
t = KolejkaTask(None)
t.load(request.read())
for image_re, image_sub in settings.IMAGE_NAME_MAPS:
t.image = re.sub(r'^' + image_re + r'$', image_sub, t.image)
accept_image = False
for image_re in settings.LIMIT_IMAGE_NAME:
if re.match(image_re, t.image):
accept_image = True
break
if not accept_image:
return FAILResponse(
message='Image {} is not accepted by the server'.format(
t.image))
local_image = False
for image_re in settings.LOCAL_IMAGE_NAMES:
if re.match(image_re, t.image):
local_image = True
break
t.id = uuid.uuid4().hex
for k, f in t.files.items():
if not f.reference:
return FAILResponse(
message='File {} does not have a reference'.format(k))
f.path = None
refs = list()
for k, f in t.files.items():
try:
ref = Reference.objects.get(key=f.reference)
except Reference.DoesNotExist:
return FAILResponse(
message='Reference for file {} is unknown'.format(k))
if not ref.public:
if not request.user.has_perm(
'blob.view_reference') and request.user != ref.user:
return FAILResponse(
message='Reference for file {} is unknown'.format(k))
refs.append(ref)
limits = KolejkaLimits(
cpus=settings.LIMIT_CPUS,
memory=settings.LIMIT_MEMORY,
swap=settings.LIMIT_SWAP,
pids=settings.LIMIT_PIDS,
storage=settings.LIMIT_STORAGE,
network=settings.LIMIT_NETWORK,
time=settings.LIMIT_TIME,
image=settings.LIMIT_IMAGE,
workspace=settings.LIMIT_WORKSPACE,
)
t.limits.update(limits)
if settings.IMAGE_REGISTRY is not None and settings.IMAGE_REGISTRY_NAME is not None and not local_image:
try:
subprocess.run(['docker', 'pull', t.image], check=True)
docker_inspect_run = subprocess.run([
'docker', 'image', 'inspect', '--format', '{{json .Id}}',
t.image
],
stdout=subprocess.PIPE,
check=True)
image_id = str(
json.loads(str(docker_inspect_run.stdout,
'utf-8'))).split(':')[-1]
logging.info(image_id)
docker_inspect_run = subprocess.run([
'docker', 'image', 'inspect', '--format', '{{json .Size}}',
t.image
],
stdout=subprocess.PIPE,
check=True)
image_size = int(
json.loads(str(docker_inspect_run.stdout, 'utf-8')))
except:
return FAILResponse(
message='Image {} could not be pulled'.format(t.image))
if t.limits.image is not None and image_size > t.limits.image:
return FAILResponse(
message='Image {} exceeds image size limit {}'.format(
t.image, t.limits.image))
image_name = settings.IMAGE_REGISTRY + '/' + settings.IMAGE_REGISTRY_NAME + ':' + image_id
try:
subprocess.run(['docker', 'tag', t.image, image_name],
check=True)
subprocess.run(['docker', 'push', image_name], check=True)
subprocess.run(['docker', 'rmi', image_name], check=True)
except:
return FAILResponse(
message='Image {} could not be pushed to local repository'.
format(t.image))
t.image = image_name
t.limits.image = image_size
task = models.Task(user=request.user,
key=t.id,
description=json.dumps(t.dump()))
task.save()
for ref in refs:
task.files.add(ref)
response = dict()
response['task'] = task.task().dump()
return OKResponse(response)
if not request.user.is_authenticated:
return HttpResponseForbidden()
try:
task = models.Task.objects.get(key=key)
except models.Task.DoesNotExist:
return HttpResponseNotFound()
if not request.user.has_perm(
'task.view_task'
) and request.user != task.user and request.user != task.assignee:
return HttpResponseForbidden()
if request.method == 'PUT':
response = dict()
response['task'] = task.task().dump()
return OKResponse(response)
if request.method == 'DELETE':
if not request.user.has_perm(
'task.delete_task') and request.user != task.user:
return HttpResponseForbidden()
task.delete()
return OKResponse({'deleted': True})
if request.method == 'GET' or request.method == 'HEAD':
response = dict()
response['task'] = task.task().dump()
return OKResponse(response)
return HttpResponseNotAllowed(['HEAD', 'GET', 'POST', 'PUT', 'DELETE'])