Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240"""Stdlib regressions for build, publish and preview. Run: python3 test/site_tools.py."""import contextlibimport hashlibimport jsonimport osfrom pathlib import Pathimport reimport signalimport socketimport subprocessimport tempfileimport timeimport unittestfrom urllib.error import HTTPError, URLErrorfrom urllib.request import ProxyHandler, build_opener
HTTP = build_opener(ProxyHandler({}))REPO = Path(__file__).resolve().parents[1]SCRIPT_ENV = """RELEASE BASE MODELS_URL RUNTIME_URL SOURCE CF_ANALYTICS_TOKEN MODEL_ID SIZESTATIC_ROOT FILE_MODE DIR_MODE MODEL RUNTIME OLDER PORT NEWS FEATURES TRAINEDISOLATE ISOLATE_PREFIX ONNX_BYTES_PER_SECOND BASH_ENV ENV""".split()
class SiteTools(unittest.TestCase): def setUp(self): temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.root = Path(temporary.name) self.env = os.environ.copy() for key in SCRIPT_ENV: self.env.pop(key, None)
def model(self, name, content=b"fixture model\n"): path = self.root / name (path / "onnx").mkdir(parents=True) (path / "config.json").write_text('{"model_type": "fixture"}\n') (path / "onnx/model_q4.onnx").write_bytes(content) (path / "SHA256SUMS").write_text("".join( f"{hashlib.sha256((path / file).read_bytes()).hexdigest()} {file}\n" for file in ("config.json", "onnx/model_q4.onnx"))) return path
def runtime(self, name, label): path = self.root / name path.mkdir(parents=True) (path / "transformers.min.js").write_text(f"export const fixture = '{label}';\n") (path / "ort.wasm").write_bytes(label.encode()) return path
def tree(self, path): result = {} for entry in path.rglob("*"): name = entry.relative_to(path).as_posix() result[name] = (("link", os.readlink(entry)) if entry.is_symlink() else ("dir",) if entry.is_dir() else ("file", entry.read_bytes())) return result
def run_script(self, name, *args, success=True, **env): result = subprocess.run([str(REPO / name), *map(str, args)], cwd=REPO, env={**self.env, **env}, capture_output=True, text=True, timeout=20) message = result.stdout + result.stderr if success: self.assertEqual(result.returncode, 0, message) else: self.assertNotEqual(result.returncode, 0, message) return result
def config(self, page): match = re.search(r'id="pet-config">(.*?)</script>', page, re.S) self.assertIsNotNone(match) return json.loads(match.group(1))
def assert_manifest(self, path): for line in (path / "SHA256SUMS").read_text().splitlines(): digest, name = line.split(" ", 1) self.assertEqual(hashlib.sha256((path / name).read_bytes()).hexdigest(), digest)
@staticmethod def fetch(url): with HTTP.open(url, timeout=1) as response: return response.read(), response.headers
@contextlib.contextmanager def preview(self, output, model, runtime): with socket.socket() as listener: listener.bind(("127.0.0.1", 0)) port = listener.getsockname()[1] url = f"http://127.0.0.1:{port}" with tempfile.TemporaryFile(mode="w+") as log: process = subprocess.Popen([str(REPO / "test/dev-site.sh"), str(output)], cwd=REPO, env={**self.env, "MODEL": str(model), "RUNTIME": str(runtime), "BASE": "pet", "MODEL_ID": "demo", "PORT": str(port)}, stdout=log, stderr=subprocess.STDOUT, start_new_session=True) try: deadline = time.monotonic() + 15 while time.monotonic() < deadline: if process.poll() is not None: log.seek(0) self.fail(f"preview exited with {process.returncode}: {log.read()}") try: self.fetch(url + "/pet/") break except (URLError, TimeoutError): time.sleep(0.05) else: log.seek(0) self.fail(f"preview did not start: {log.read()}") yield url finally: if process.poll() is None: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGTERM) try: process.wait(timeout=5) except subprocess.TimeoutExpired: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) process.wait(timeout=5)
def test_failed_build_preserves_rendered_site_and_published_directories(self): output = self.root / "site" self.model("site/models/demo") self.runtime("site/runtime/fixture-runtime", "preserved") published = {name: self.tree(output / name) for name in ("models", "runtime")} release = {"modelId": "demo", "size": 14, "runtime": "fixture-runtime", "trained": "fixture build", "news": "fixture release", "features": [], "older": []} release_file = self.root / "release.json" release_file.write_text(json.dumps(release)) options = {"RELEASE": str(release_file), "BASE": "/pet/"} self.run_script("site/build.sh", output, **options) for name in published: self.assertEqual(self.tree(output / name), published[name]) config = self.config((output / "index.html").read_text()) self.assertEqual(config["modelId"], "demo") (output / "old-entry.txt").write_text("old output must survive failures\n") before = self.tree(output) siblings = set(self.root.iterdir()) failures = [ ("invalid analytics", json.dumps(release), {"CF_ANALYTICS_TOKEN": "invalid"}), ("malformed release", "{", {}), ("unfilled render", json.dumps({**release, "trained": "@UNFILLED@"}), {}), ("invalid rendered config", json.dumps(release), {"MODEL_ID": 'bad"id'}), ] for name, content, overrides in failures: with self.subTest(name=name): release_file.write_text(content) self.run_script("site/build.sh", output, success=False, **options, **overrides) self.assertEqual(self.tree(output), before) self.assertEqual(self.config((output / "index.html").read_text()), config) self.assertEqual(set(self.root.iterdir()), siblings) release_file.write_text(json.dumps({**release, "news": "rebuilt release"})) self.run_script("site/build.sh", output, **options) self.assertFalse((output / "old-entry.txt").exists()) self.assertEqual(self.config((output / "index.html").read_text())["news"], "rebuilt release") for name in published: self.assertEqual(self.tree(output / name), published[name])
def test_publish_runtime_failure_allows_same_id_retry(self): model = self.model("model") runtime = self.runtime("runtime", "A") root = self.root / "published" options = {"STATIC_ROOT": str(root), "BASE": "pet", "FILE_MODE": "0640", "DIR_MODE": "0750"} # Exercise both a missing runtime and a copy that fails after copying files. broken = self.runtime("broken-runtime", "broken") (broken / "z-directory").mkdir() for source in (self.root / "missing-runtime", broken): with self.subTest(runtime=source.name): self.run_script("train/publish.sh", model, "demo", source, "fixture-runtime", success=False, **options) self.assertEqual(list((root / "pet/models").iterdir()), []) self.assertEqual(list((root / "pet/runtime").iterdir()), []) self.run_script("train/publish.sh", model, "demo", runtime, "fixture-runtime", **options) published_model = root / "pet/models/demo" published_runtime = root / "pet/runtime/fixture-runtime" self.assertEqual(self.tree(published_model), self.tree(model)) payload = self.tree(published_runtime) payload.pop("SHA256SUMS") self.assertEqual(payload, self.tree(runtime)) for path in (published_model, published_runtime): self.assert_manifest(path) for entry in (path, *path.rglob("*")): self.assertEqual(entry.stat().st_mode & 0o777, 0o750 if entry.is_dir() else 0o640)
def test_published_model_and_existing_runtime_are_immutable(self): model_a = self.model("model-a", b"model A\n") model_b = self.model("model-b", b"model B\n") runtime_a = self.runtime("runtime-a", "A") runtime_b = self.runtime("runtime-b", "B") root = self.root / "published" options = {"STATIC_ROOT": str(root), "BASE": "pet"} self.run_script("train/publish.sh", model_a, "demo", runtime_a, "fixture-runtime", **options) before = self.tree(root) self.run_script("train/publish.sh", model_b, "demo", runtime_b, "other-runtime", success=False, **options) self.assertEqual(self.tree(root), before) self.run_script("train/publish.sh", model_b, "next", runtime_b, "fixture-runtime", **options) self.assertEqual(self.tree(root / "pet/models/demo"), self.tree(model_a)) self.assertEqual(self.tree(root / "pet/models/next"), self.tree(model_b)) original_runtime = {key.removeprefix("pet/runtime/fixture-runtime/"): value for key, value in before.items() if key.startswith("pet/runtime/fixture-runtime/")} self.assertEqual(self.tree(root / "pet/runtime/fixture-runtime"), original_runtime) self.assert_manifest(root / "pet/runtime/fixture-runtime")
def test_full_preview_refreshes_runtime_but_rebuild_does_not(self): model = self.model("model") runtime_a = self.runtime("runtime-a", "A") (runtime_a / "obsolete.wasm").write_bytes(b"obsolete") runtime_b = self.runtime("runtime-b", "B") (runtime_b / ".fixture").write_text("hidden runtime metadata\n") (runtime_b / "nested").mkdir() (runtime_b / "nested/fixture.txt").write_text("nested runtime metadata\n") output = self.root / "preview" selected = output / "pet/runtime/runtime" with self.preview(output, model, runtime_a) as url: body, headers = self.fetch(url + "/pet/runtime/runtime/transformers.min.js") self.assertEqual(body, (runtime_a / "transformers.min.js").read_bytes()) self.assertEqual(headers["Cross-Origin-Opener-Policy"], "same-origin") self.assertEqual(headers["Cross-Origin-Embedder-Policy"], "require-corp") self.run_script("test/dev-site.sh", output, "--rebuild", MODEL=str(model), RUNTIME=str(runtime_b), BASE="pet", MODEL_ID="demo") self.assertEqual(self.tree(selected), self.tree(runtime_a)) with self.preview(output, model, runtime_b) as url: body, _ = self.fetch(url + "/pet/runtime/runtime/transformers.min.js") self.assertEqual(body, (runtime_b / "transformers.min.js").read_bytes()) self.assertEqual(self.tree(selected), self.tree(runtime_b)) with self.assertRaises(HTTPError) as error: self.fetch(url + "/pet/runtime/runtime/obsolete.wasm") with error.exception: self.assertEqual(error.exception.code, 404) previous = self.tree(selected) siblings = set(selected.parent.iterdir()) self.run_script("test/dev-site.sh", output, success=False, MODEL=str(model), RUNTIME=str(self.root / "missing-runtime"), BASE="pet", MODEL_ID="demo") self.assertEqual(self.tree(selected), previous) self.assertEqual(set(selected.parent.iterdir()), siblings)
if __name__ == "__main__": unittest.main()