Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@vadiklyutiy
vadiklyutiy committed Dec 24, 2024
1 parent 8db9f39 commit a7341de
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 94 deletions.
86 changes: 76 additions & 10 deletions apps/compile_server/resources/compilation.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
from typing import Dict, Any, List, Tuple
from typing import Dict, Any, Tuple, Sequence, Union
import time
import re
import sys
import os
import traceback
import threading
import requests
import subprocess
import zipfile
import logging
Expand All @@ -26,7 +25,73 @@
commits_dir = os.path.join(os.getcwd(), 'commits')
results_dir = os.path.join(os.getcwd(), 'results')

compile_script = os.path.join(os.path.dirname(__file__), 'compile_worker.py')

def save_response(response, response_file: str):
with open(response_file, 'wb') as f:
pickle.dump(response, f)


def compile_job(job_id: str):
try:
job_file = os.path.join(jobs_dir, job_id + '.pickle')
if not os.path.exists(job_file):
# job not found
return 1

job_lock = os.path.join(jobs_dir, job_id + '.lock')
with FileLock(job_lock):
response_file = os.path.join(jobs_dir, job_id + '.response')
if os.path.exists(response_file):
# job already compiled
return 0

# unpacking the job
with open(job_file, 'rb') as f:
job: Dict[str, Any] = pickle.load(f)

# import the hidet from the commit
commit_id: str = job['commit_id']
commit_dir = os.path.join(commits_dir, commit_id)
sys.path.insert(0, os.path.join(commit_dir, 'python'))
import hidet # import the hidet from the commit

# load the workload
workload: Dict[str, Any] = pickle.loads(job['workload'])
ir_module: Union[hidet.ir.IRModule, Sequence[hidet.ir.IRModule]] = workload['ir_module']
target: str = workload['target']
output_kind: str = workload['output_kind']

# perform the compilation
module_string = str(ir_module)
key = module_string + target + output_kind + commit_id
hash_digest: str = sha256(key.encode()).hexdigest()
zip_file_path: str = os.path.join(results_dir, hash_digest + '.zip')
if not os.path.exists(zip_file_path):
output_dir: str = os.path.join(results_dir, hash_digest)
with FileLock(os.path.join(results_dir, f'{hash_digest}.lock')):
if not os.path.exists(os.path.join(output_dir, 'lib.so')):
hidet.drivers.build_ir_module(
ir_module,
output_dir=output_dir,
target=target,
output_kind=output_kind
)
with zipfile.ZipFile(zip_file_path, 'w') as f:
for root, dirs, files in os.walk(output_dir):
for file in files:
f.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), output_dir))
except Exception:
response = {
'message': '[Remote] ' + traceback.format_exc()
}, 400
save_response(response, response_file)
return 0

response = {
'download_filename': hash_digest + '.zip',
}, 200
save_response(response, response_file)
return 0


def should_update(repo_timestamp) -> bool:
Expand Down Expand Up @@ -142,11 +207,11 @@ def post(self):
job_path = os.path.join(jobs_dir, job_id + '.pickle')
job_response_path = os.path.join(jobs_dir, job_id + '.response')

print('[{}] Received a job: {}'.format(pid, job_id[:16]))
print('[{}] Received a job: {}'.format(pid, job_id[:16]), flush=True)

# check if the job is already done
if os.path.exists(job_response_path):
print('[{}] Job {} has already done before, respond directly'.format(pid, job_id[:16]))
print('[{}] Job {} has already done before, respond directly'.format(pid, job_id[:16]), flush=True)
with open(job_response_path, 'rb') as f:
return pickle.load(f)

Expand All @@ -158,19 +223,20 @@ def post(self):
pickle.dump(job, f)

with lock: # Only one thread can access the following code at the same time
print('[{}] Start compiling: {}'.format(pid, job_id[:16]))
ret = subprocess.run([sys.executable, compile_script, '--job_id', job_id])
print('[{}] Start compiling: {}'.format(pid, job_id[:16]), flush=True)
from .compile_worker import compile_job
compile_job(job_id)

# respond to the client
response_path = os.path.join(jobs_dir, job_id + '.response')
if not os.path.exists(response_path):
raise RuntimeError('Can not find the response file:\n{}{}'.format(ret.stderr, ret.stdout))
raise RuntimeError('Can not find the response file')
else:
print('[{}] Finish compiling: {}'.format(pid, job_id[:16]))
print('[{}] Finish compiling: {}'.format(pid, job_id[:16]), flush=True)
with open(response_path, 'rb') as f:
response: Tuple[Dict, int] = pickle.load(f)
return response
except Exception as e:
msg = traceback.format_exc()
print('[{}] Failed to compile:\n{}'.format(pid, msg))
print('[{}] Failed to compile:\n{}'.format(pid, msg), flush=True)
return {'message': '[Remote] {}'.format(msg)}, 500
85 changes: 1 addition & 84 deletions apps/compile_server/resources/compile_worker.py
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,10 @@
from typing import Dict, Any, List, Tuple, Sequence, Union
from typing import Dict, Any, Sequence, Union
import os
import traceback
import argparse
import sys
import re
import subprocess
import zipfile
import logging
import pickle
import git
from hashlib import sha256
from filelock import FileLock

Expand All @@ -18,82 +14,3 @@
repos_dir = os.path.join(os.getcwd(), 'repos')
commits_dir = os.path.join(os.getcwd(), 'commits')
results_dir = os.path.join(os.getcwd(), 'results')


def save_response(response, response_file: str):
with open(response_file, 'wb') as f:
pickle.dump(response, f)


def compile_job(job_id: str):
try:
job_file = os.path.join(jobs_dir, job_id + '.pickle')
if not os.path.exists(job_file):
# job not found
return 1

job_lock = os.path.join(jobs_dir, job_id + '.lock')
with FileLock(job_lock):
response_file = os.path.join(jobs_dir, job_id + '.response')
if os.path.exists(response_file):
# job already compiled
return 0

# unpacking the job
with open(job_file, 'rb') as f:
job: Dict[str, Any] = pickle.load(f)

# import the hidet from the commit
commit_id: str = job['commit_id']
commit_dir = os.path.join(commits_dir, commit_id)
sys.path.insert(0, os.path.join(commit_dir, 'python'))
import hidet # import the hidet from the commit

# load the workload
workload: Dict[str, Any] = pickle.loads(job['workload'])
ir_module: Union[hidet.ir.IRModule, Sequence[hidet.ir.IRModule]] = workload['ir_module']
target: str = workload['target']
output_kind: str = workload['output_kind']

# perform the compilation
module_string = str(ir_module)
key = module_string + target + output_kind + commit_id
hash_digest: str = sha256(key.encode()).hexdigest()
zip_file_path: str = os.path.join(results_dir, hash_digest + '.zip')
if not os.path.exists(zip_file_path):
output_dir: str = os.path.join(results_dir, hash_digest)
with FileLock(os.path.join(results_dir, f'{hash_digest}.lock')):
if not os.path.exists(os.path.join(output_dir, 'lib.so')):
hidet.drivers.build_ir_module(
ir_module,
output_dir=output_dir,
target=target,
output_kind=output_kind
)
with zipfile.ZipFile(zip_file_path, 'w') as f:
for root, dirs, files in os.walk(output_dir):
for file in files:
f.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), output_dir))
except Exception:
response = {
'message': '[Remote] ' + traceback.format_exc()
}, 400
save_response(response, response_file)
return 0

response = {
'download_filename': hash_digest + '.zip',
}, 200
save_response(response, response_file)
return 0


def main():
parser = argparse.ArgumentParser()
parser.add_argument('--job_id', type=str, required=True)
args = parser.parse_args()
exit(compile_job(args.job_id))


if __name__ == '__main__':
main()
2 changes: 2 additions & 0 deletions apps/compile_server/run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ def main():
# download the compile server code from official repo
if not os.path.exists('compile_server/hidet'):
run('git clone https://github.com/hidet-org/hidet.git compile_server/hidet')
else:
run('git pull', cwd='compile_server/hidet')

# install the dependencies
run('pip install -r requirements.txt', cwd='compile_server/hidet/apps/compile_server')
Expand Down

0 comments on commit a7341de

Please sign in to comment.