Cleanup

generator/engraving/worker.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import collections
+import dataclasses
+import hashlib
+import json
+import pathlib
+import shutil
+import tempfile
+import time
+
+import PIL.Image
+
+import generator.engraving.audit
+import generator.engraving.cache
+import generator.engraving.fonts
+import generator.engraving.layout
+import generator.engraving.notation
+import generator.engraving.raster
+import generator.engraving.types
+import generator.engraving.validation
+import generator.performance
+import generator.validation
+
+@dataclasses.dataclass
+class GroupResult:
+  scores: list = dataclasses.field(default_factory=list)
+  fullscore: int = 0
+  bodyhits: int = 0
+  bodies: int = 0
+  renderedpages: int = 0
+  rejected: int = 0
+  timings: dict = dataclasses.field(default_factory=lambda: dict.fromkeys(('layout', 'pagination', 'raster', 'cache'), 0.0))
+
+class BodyBands:
+  def __init__(self, root, persistent, result, limit=16 * 1024 * 1024):
+    self.root, self.persistent, self.result, self.limit = root, persistent, result, limit
+    self.images = collections.OrderedDict()
+    self.size = 0
+    self.known = {}
+    self.root.mkdir(parents=True, exist_ok=True)
+  def get(self, notation, page, band, fonts):
+    key = generator.engraving.cache.band_fingerprint(notation, page, band)
+    if key in self.images:
+      self.images.move_to_end(key)
+      self.result.bodyhits += 1
+      return self.images[key]
+    started = time.monotonic()
+    entry = self.root / (key + '.rgb')
+    image = None
+    if key in self.known:
+      if entry.is_symlink() or entry.stat().st_size != page.width * (band.bottom - band.top) * 3:
+        raise generator.engraving.types.EngravingError('temporary notation band was corrupted')
+      data = entry.read_bytes()
+      if hashlib.sha256(data).hexdigest() != self.known[key]:
+        raise generator.engraving.types.EngravingError('temporary notation band was corrupted')
+      image = PIL.Image.frombytes('RGB', (page.width, band.bottom - band.top), data)
+    elif self.persistent:
+      cached = self.persistent / 'notation' / key
+      image = generator.engraving.cache.restore_band(cached, key, page.width, band.bottom - band.top)
+      self.result.rejected += int(image is None and (cached.exists() or cached.is_symlink()))
+    self.result.timings['cache'] += time.monotonic() - started
+    if image is None:
+      started = time.monotonic()
+      image = generator.engraving.raster.rasterize_band(page, band.primitives, band.top, band.bottom, fonts)
+      self.result.timings['raster'] += time.monotonic() - started
+      self.result.bodies += 1
+      if self.persistent:
+        started = time.monotonic()
+        generator.engraving.cache.store_band(self.persistent / 'notation' / key, key, image)
+        self.result.timings['cache'] += time.monotonic() - started
+    else:
+      self.result.bodyhits += 1
+    if key not in self.known:
+      started = time.monotonic()
+      # Spill finalized pixels, so eviction never repeats rasterization or PNG decoding.
+      data = image.tobytes()
+      entry.write_bytes(data)
+      self.result.timings['cache'] += time.monotonic() - started
+      self.known[key] = hashlib.sha256(data).hexdigest()
+    size = image.width * image.height * 3
+    while self.images and self.size + size > self.limit:
+      _, old = self.images.popitem(last=False)
+      self.size -= old.width * old.height * 3
+      old.close()
+    self.images[key] = image
+    self.size += size
+    return image
+  def close(self):
+    for image in self.images.values():
+      image.close()
+    self.images.clear()
+    self.size = 0
+
+def clear_pages(destination):
+  if destination.is_symlink() or destination.is_file():
+    destination.unlink()
+  elif destination.exists():
+    shutil.rmtree(destination)
+
+def write_bundle(path, document, pages, fingerprint):
+  entries = [{k: v for k, v in {**page, 'path': path.with_suffix('.pages').name + '/' + page['path']}.items() if k != 'playback'} for page in pages]
+  normal = [entry for entry, page in zip(entries, pages) if not page['playback']]
+  playback = [entry for entry, page in zip(entries, pages) if page['playback']]
+  music = document
+  bundle = {'music': music, 'presentation': {'engine': generator.engraving.validation.ENGINE_NAME, 'style_sha256': fingerprint,
+    'width': generator.engraving.validation.PAGE_WIDTH, 'height': generator.engraving.validation.PAGE_HEIGHT, 'pages': normal, 'playback_pages': playback}}
+  validate_cached_views(document, pages)
+  bundle['performance'] = generator.performance.timing(document)
+  path.write_text(json.dumps(bundle, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
+
+def validate_cached_views(document, pages):
+  generator.performance.validate_views(document, [p for p in pages if not p['playback']], [p for p in pages if p['playback']])
+
+def engrave_group(notation, paths, root, font_path, cache_dir, fingerprint, progress=False):
+  family = generator.engraving.fonts.Fonts(font_path)
+  style = generator.engraving.types.Style()
+  systems = None
+  index = None
+  scores = {}
+  audited = set()
+  result = GroupResult()
+  with tempfile.TemporaryDirectory(prefix='vamp_notation_') as temporary:
+    bodies = BodyBands(pathlib.Path(temporary), cache_dir, result)
+    try:
+      for path, score_key in paths:
+        try:
+          document = json.loads(path.read_text(encoding='utf-8'))
+          style = generator.engraving.types.score_style(document, generator.engraving.types.Style())
+          destination = path.with_suffix('.pages')
+          clear_pages(destination)
+          destination.mkdir()
+          if score_key in scores:
+            source, pages = scores[score_key]
+            for page in pages:
+              shutil.copyfile(source / page['path'], destination / page['path'])
+            result.fullscore += 1
+          else:
+            if systems is None:
+              started = time.monotonic()
+              generator.validation.validate_document(document)
+              index = generator.engraving.notation.ScoreIndex(document)
+              systems = generator.engraving.layout.layout_systems(document, style, family, index=index)
+              result.timings['layout'] += time.monotonic() - started
+            started = time.monotonic()
+            pages_layout, playback_layout = generator.engraving.layout.page_views(systems, document, family, style, index=index)
+            placement = tuple((page.width, page.height, tuple((system.start, system.end, system.box) for system in page.systems)) for page in pages_layout)
+            if placement not in audited:
+              layout = generator.engraving.types.LayoutDocument(pages_layout, style, fingerprint)
+              generator.engraving.audit.audit_layout(document, layout, family)
+              generator.engraving.audit.audit_playback(pages_layout, playback_layout, style, family)
+              audited.add(placement)
+            else:
+              for page in pages_layout + playback_layout:
+                decoration = tuple(p for p in page.primitives if p.role in ('title', 'subtitle', 'page_number'))
+                generator.engraving.audit.audit_primitives(decoration, page.width, page.height, family)
+            result.timings['pagination'] += time.monotonic() - started
+            pages = []
+            rendered = {}
+            for is_playback, page in [(False, p) for p in pages_layout] + [(True, p) for p in playback_layout]:
+              if page.primitives in rendered:
+                name, digest = rendered[page.primitives]
+                pages.append(dict(path=name, sha256=digest, playback=is_playback, **generator.performance.page_timing(document, page)))
+                continue
+              bands = generator.engraving.raster.page_bands(page)
+              if bands is None:
+                started = time.monotonic()
+                image = generator.engraving.raster.rasterize_page(page, family)
+              else:
+                body = bodies.get(notation, page, bands[1], family)
+                started = time.monotonic()
+                image = generator.engraving.raster.compose_page(page, family, bands, body)
+              target = destination / (('playback-' if is_playback else '') + f'{page.number:04}.png')
+              with image:
+                image.save(target, format='PNG')
+              result.timings['raster'] += time.monotonic() - started
+              result.renderedpages += 1
+              digest = hashlib.sha256(target.read_bytes()).hexdigest()
+              rendered[page.primitives] = target.name, digest
+              pages.append(dict(path=target.name, sha256=digest, playback=is_playback, **generator.performance.page_timing(document, page)))
+            scores[score_key] = (destination, pages)
+          write_bundle(path, document, pages, fingerprint)
+          result.scores.append((path, score_key, pages))
+          if progress and len(result.scores) % 50 == 0:
+            print(f'Engraving body {notation[:12]}: {len(result.scores)}/{len(paths)} scores, fullscore {result.fullscore}, bodyhits {result.bodyhits}, bodies {result.bodies}, renderedpages {result.renderedpages}', flush=True)
+        except ValueError as error:
+          raise generator.engraving.types.EngravingError(f'{path.relative_to(root)}: {error}') from error
+    finally:
+      bodies.close()
+  return result