File size: 3,061 Bytes
324ecd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """Copy the game files SUPERBEAM needs into envs/superbeam_env/vendor/ so the env dir
(and its Docker image / HF Space) is self-contained. Re-run after changing the game.
python scripts/prepare_superbeam_env.py
The imports of beamboot/superbeam_gym.py are followed automatically (module-level
imports of beamboot.* only). The in-weights computer (chip8_lm, chip8_fast), the scripts
that run and inspect it, and its tests are copied too, so the Hub repo holds the full
source; they need torch + transformers (Beam), which the env server image does not
install (superbeam_gym imports chip8_fast lazily, only for the beam-weights backend).
Run its tests from vendor/: cd vendor && pytest tests/
"""
import ast
import os
import shutil
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEST = os.path.join(ROOT, "envs", "superbeam_env", "vendor")
ROMS = ("superbeam.asm", "superbeam.ch8", "superbeam_inputs.json")
IN_WEIGHTS = ("beamboot/chip8_lm.py", "beamboot/chip8_fast.py",
"scripts/run_chip8_fast.py", "scripts/attention_probe.py", "scripts/deploy_inweights_modal.py",
"tests/test_chip8_lm.py", "tests/test_chip8_fast.py")
def module_level_deps(relpath, seen):
if relpath in seen:
return
seen.add(relpath)
tree = ast.parse(open(os.path.join(ROOT, relpath)).read())
for node in tree.body: # module level only: lazy imports are skipped
names = []
if isinstance(node, ast.ImportFrom) and node.module:
names = [node.module] + [node.module + "." + a.name for a in node.names]
elif isinstance(node, ast.Import):
names = [a.name for a in node.names]
for n in names:
if n.split(".")[0] == "beamboot":
p = n.replace(".", "/") + ".py"
if os.path.exists(os.path.join(ROOT, p)):
module_level_deps(p, seen)
def main():
seen = set()
module_level_deps("beamboot/superbeam_gym.py", seen)
seen.discard("beamboot/__init__.py")
if os.path.isdir(DEST):
shutil.rmtree(DEST)
os.makedirs(os.path.join(DEST, "beamboot"))
os.makedirs(os.path.join(DEST, "roms"))
open(os.path.join(DEST, "beamboot", "__init__.py"), "w").write('"""Bundled subset of beamboot for the SUPERBEAM env."""\n')
for rel in sorted(seen):
shutil.copy2(os.path.join(ROOT, rel), os.path.join(DEST, rel))
for r in ROMS:
shutil.copy2(os.path.join(ROOT, "roms", r), os.path.join(DEST, "roms", r))
for rel in IN_WEIGHTS:
os.makedirs(os.path.dirname(os.path.join(DEST, rel)), exist_ok=True)
shutil.copy2(os.path.join(ROOT, rel), os.path.join(DEST, rel))
# Lets `pytest tests/` run from vendor/ (as the project root), not inside the env package.
open(os.path.join(DEST, "pytest.ini"), "w").write("[pytest]\npythonpath = .\n")
for dirpath, _, files in os.walk(DEST):
for f in files:
print(os.path.relpath(os.path.join(dirpath, f), ROOT))
if __name__ == "__main__":
main()
|