Cached Generation

generator/engraving/__init__.py

@@ -1,16 +1,20 @@
 from __future__ import annotations
 
+from concurrent.futures import ProcessPoolExecutor
 import hashlib
 import json
+import multiprocessing
 from pathlib import Path
 import shutil
 
 from PIL import Image, ImageChops
 
 from generator.engraving.fonts import Fonts
+from generator.engraving.cache import renderer_fingerprint, restore_pages, score_fingerprint, store_pages
 from generator.engraving.layout import layout_score
 from generator.engraving.raster import rasterize_page
 from generator.engraving.types import EngravingError, Style
+from generator.model import validate_document
 
 ENGINE_NAME = 'Vamp Python'
 ENGINE_VERSION = '1.0.0'
@@ -31,31 +35,63 @@ def check_page(path: Path) -> None:
     if bounds[0] < 4 or bounds[1] < 4 or bounds[2] > PAGE_WIDTH - 4 or bounds[3] > PAGE_HEIGHT - 4:
       raise EngravingError(f'{path.name}: notation reaches the page edge; refusing potentially clipped output')
 
-def engrave_tree(root: Path, fonts: Path) -> None:
-  paths = sorted(root.rglob('*.json'))
-  if not paths:
-    return
+def engrave_scores(paths: list[Path], root: Path, fonts: Path, cache_dir: Path | None, progress: bool, worker: int) -> tuple[int, int]:
   family = Fonts(fonts)
   style = Style()
-  for path in paths:
+  fingerprint = family.fingerprint(style)
+  renderer = renderer_fingerprint(family, style) if cache_dir else ''
+  reused = 0
+  for index, path in enumerate(paths, 1):
     document = json.loads(path.read_text(encoding='utf-8'))
     try:
-      layout = layout_score(document, style, family)
+      validate_document(document)
       destination = path.with_suffix('.pages')
       if destination.exists():
         shutil.rmtree(destination)
-      destination.mkdir()
-      entries = []
-      for page in layout.pages:
-        name = f'{page.number:04}.png'
-        target = destination / name
-        rasterize_page(page, family).save(target, format='PNG')
-        check_page(target)
-        entries.append({'path': destination.name + '/' + name, 'sha256': hashlib.sha256(target.read_bytes()).hexdigest()})
+      entry = cache_dir / score_fingerprint(document, renderer) if cache_dir else None
+      pages = restore_pages(entry, destination) if entry else None
+      if pages is None:
+        if destination.exists():
+          shutil.rmtree(destination)
+        destination.mkdir()
+        layout = layout_score(document, style, family)
+        pages = []
+        for page in layout.pages:
+          name = f'{page.number:04}.png'
+          target = destination / name
+          rasterize_page(page, family).save(target, format='PNG')
+          check_page(target)
+          pages.append({'path': name, 'sha256': hashlib.sha256(target.read_bytes()).hexdigest()})
+        if entry:
+          store_pages(entry, destination, pages)
+      else:
+        reused += 1
+      entries = [{**page, 'path': destination.name + '/' + page['path']} for page in pages]
       music = {k: v for k, v in document.items() if k not in ('schema_version', 'playback')}
-      bundle = {'schema_version': 3, 'music': music, 'presentation': {'engine': ENGINE_NAME, 'engine_version': ENGINE_VERSION, 'style_sha256': layout.fingerprint, 'width': PAGE_WIDTH, 'height': PAGE_HEIGHT, 'pages': entries}}
+      bundle = {'schema_version': 3, 'music': music, 'presentation': {'engine': ENGINE_NAME, 'engine_version': ENGINE_VERSION, 'style_sha256': fingerprint, 'width': PAGE_WIDTH, 'height': PAGE_HEIGHT, 'pages': entries}}
       if 'playback' in document:
         bundle['playback'] = document['playback']
       path.write_text(json.dumps(bundle, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
+      if progress and (index % 50 == 0 or index == len(paths)):
+        print(f'Engraving worker {worker}: {index}/{len(paths)} scores, {reused} reused', flush=True)
     except ValueError as error:
       raise EngravingError(f'{path.relative_to(root)}: {error}') from error
+  return len(paths), reused
+
+def engrave_tree(root: Path, fonts: Path, *, cache_dir: Path | None = None, jobs: int = 1, progress: bool = False) -> None:
+  if type(jobs) is not int or not 1 <= jobs <= 16:
+    raise ValueError('engraving jobs must be between 1 and 16')
+  if cache_dir and (cache_dir.resolve().is_relative_to(root.resolve()) or root.resolve().is_relative_to(cache_dir.resolve())):
+    raise ValueError('engraving cache and generated assets must not overlap')
+  paths = sorted(root.rglob('*.json'))
+  if not paths:
+    return
+  jobs = min(jobs, len(paths))
+  if jobs == 1:
+    results = [engrave_scores(paths, root, fonts, cache_dir, progress, 1)]
+  else:
+    with ProcessPoolExecutor(max_workers=jobs, mp_context=multiprocessing.get_context('spawn')) as pool:
+      futures = [pool.submit(engrave_scores, paths[index::jobs], root, fonts, cache_dir, progress, index + 1) for index in range(jobs)]
+      results = [future.result() for future in futures]
+  if progress:
+    print(f'Engraving complete: {sum(total for total, _ in results)} scores, {sum(reused for _, reused in results)} reused', flush=True)

generator/engraving/cache.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import hashlib
+import importlib.metadata
+import json
+from pathlib import Path
+import shutil
+import sys
+import tempfile
+
+from PIL import features
+
+def renderer_fingerprint(fonts, style) -> str:
+  root = Path(__file__).resolve().parents[1]
+  versions = {name: importlib.metadata.version(name) for name in ('Pillow', 'fonttools', 'aggdraw')}
+  versions.update({name: features.version(name) for name in features.get_supported()})
+  digest = hashlib.sha256(json.dumps([fonts.fingerprint(style), sys.version, versions], sort_keys=True).encode())
+  for path in sorted([*root.joinpath('engraving').glob('*.py'), root / 'model.py', root / 'harmony.py', root / 'requirements.txt']):
+    digest.update(path.relative_to(root).as_posix().encode())
+    digest.update(path.read_bytes())
+  return digest.hexdigest()
+
+def score_fingerprint(document: dict, renderer: str) -> str:
+  music = {key: value for key, value in document.items() if key != 'playback'}
+  return hashlib.sha256(json.dumps([renderer, music], sort_keys=True, ensure_ascii=False).encode()).hexdigest()
+
+def restore_pages(entry: Path, destination: Path) -> list[dict] | None:
+  from generator.engraving import check_page
+  try:
+    if entry.is_symlink():
+      return None
+    manifest = entry / 'pages.json'
+    if manifest.is_symlink():
+      return None
+    pages = json.loads(manifest.read_text(encoding='utf-8'))
+    if not isinstance(pages, list) or not 0 < len(pages) <= 10000:
+      return None
+    destination.mkdir(parents=True, exist_ok=True)
+    for index, page in enumerate(pages, 1):
+      if not isinstance(page, dict) or set(page) != {'path', 'sha256'} or page['path'] != f'{index:04}.png':
+        return None
+      path = entry / page['path']
+      if path.is_symlink():
+        return None
+      target = destination / page['path']
+      shutil.copyfile(path, target)
+      if hashlib.sha256(target.read_bytes()).hexdigest() != page['sha256']:
+        return None
+      check_page(target)
+    return pages
+  except (OSError, ValueError, TypeError):
+    return None
+
+def store_pages(entry: Path, destination: Path, pages: list[dict]) -> None:
+  entry.parent.mkdir(parents=True, exist_ok=True)
+  if entry.is_symlink() or entry.is_file():
+    entry.unlink(missing_ok=True)
+  entry.mkdir(exist_ok=True)
+  with tempfile.TemporaryDirectory(prefix='.engraving-', dir=entry.parent) as temporary:
+    staged = Path(temporary)
+    for page in pages:
+      shutil.copyfile(destination / page['path'], staged / page['path'])
+    (staged / 'pages.json').write_text(json.dumps(pages), encoding='utf-8')
+    # Identical keys have identical pages; publish complete files, then the manifest.
+    for name in [page['path'] for page in pages] + ['pages.json']:
+      target = entry / name
+      if target.is_dir() and not target.is_symlink():
+        shutil.rmtree(target, ignore_errors=True)
+      (staged / name).replace(target)