84 خطوط
4.0 KiB
Python
84 خطوط
4.0 KiB
Python
"""Run `python -m app.worker`; add replicas to scale parsing independently of API."""
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from sqlalchemy import update
|
|
from .config import settings
|
|
from .cleanup import cleanup
|
|
from .db import Job, claim, initialize, sessions
|
|
from .storage import storage
|
|
from .job_types import job_kind, comparison_sources
|
|
|
|
log = logging.getLogger(__name__)
|
|
stopping = False
|
|
|
|
def process(job):
|
|
config = settings()
|
|
prefix = f"{job['id']}/attempt-{job['attempt']}"
|
|
try:
|
|
with tempfile.TemporaryDirectory(prefix='schdoc-') as directory:
|
|
root = Path(directory)
|
|
kind = job_kind(job['input_key'])
|
|
if kind is None: raise ValueError('Unknown job input kind')
|
|
pcb = kind.startswith('pcb_')
|
|
source = root / ('input.PcbDoc' if pcb else 'input.SchDoc')
|
|
comparison = kind.endswith('_comparison')
|
|
if comparison:
|
|
manifest = json.loads(storage().get(job['input_key']))
|
|
source = root / 'comparison.json'
|
|
source.write_text(json.dumps(manifest), encoding='utf-8')
|
|
for key in comparison_sources(job['input_key']):
|
|
(root / key.rsplit('/', 1)[-1]).write_bytes(storage().get(key))
|
|
else:
|
|
source.write_bytes(storage().get(job['input_key']))
|
|
env = {k: v for k, v in os.environ.items() if k in ('PATH', 'LANG', 'LC_ALL', 'PYTHONPATH')}
|
|
# Drain no untrusted output into logs. Kill/wait is handled by subprocess.run.
|
|
task = 'app.pcb_task' if pcb else 'app.compare_task' if comparison else 'app.parse_task'
|
|
args = [kind] if pcb else []
|
|
subprocess.run([sys.executable, '-m', task, str(source), str(root), job['filename'], job['id'], *args],
|
|
timeout=config.parse_timeout, check=True, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
result = json.loads((root / 'result.json').read_text())
|
|
if result['ok']:
|
|
storage().put(prefix + '/data.json', (root / 'data.json').read_bytes(), 'application/json')
|
|
for name in (('old.svg','new.svg') if comparison else ('drawing.svg',)):
|
|
storage().put(prefix + '/' + name, (root / name).read_bytes(), 'image/svg+xml')
|
|
values = dict(status='completed', output_prefix=prefix, summary=json.dumps(result['summary']))
|
|
else:
|
|
values = dict(status='failed', error_code=result['code'], error_message=result['message'])
|
|
except subprocess.TimeoutExpired:
|
|
values = dict(status='failed', error_code='parse_timeout', error_message='Schematic exceeded the parsing time limit')
|
|
except subprocess.CalledProcessError:
|
|
values = dict(status='failed', error_code='parse_resource_limit', error_message='Parser exited before completion (resource limit or invalid input)')
|
|
except Exception:
|
|
log.exception('Worker infrastructure failure for %s', job['id'])
|
|
values = dict(status='failed', error_code='processing_error', error_message='Storage or processing service unavailable')
|
|
with sessions().begin() as db:
|
|
# An expired worker may finish late. Its outputs must never replace a newer attempt.
|
|
db.execute(update(Job).where(Job.id == job['id'], Job.attempt == job['attempt'], Job.status == 'processing')
|
|
.values(**values, lease_until=None, updated_at=time.time()))
|
|
|
|
def main():
|
|
global stopping
|
|
logging.basicConfig(level=logging.INFO)
|
|
def stop(*_):
|
|
global stopping
|
|
stopping = True
|
|
signal.signal(signal.SIGTERM, stop)
|
|
signal.signal(signal.SIGINT, stop)
|
|
initialize()
|
|
last_cleanup=0
|
|
while not stopping:
|
|
if time.monotonic()-last_cleanup>60:
|
|
cleanup(); last_cleanup=time.monotonic()
|
|
job = claim()
|
|
if job: process(job)
|
|
else: time.sleep(0.5)
|
|
|
|
if __name__ == '__main__': main()
|