Merge pull request #2 from srobo/comp-matches

Add support for running competition matches

assets/user_readme.md

@@ -141,7 +141,7 @@ The API for the simulator is the same as the API for the physical robot, so you
 
 
 As well as the logs being displayed in the console, they are also saved to a file.
-This file is saved in the `zone_0` folder and has a name in the format `log-<date>.log`.
+This file is saved in the `zone_0` folder and has a name in the format `log-zone-<zone>-<date>.log`.
 The date is when that simulation was run.
 
 ### Simulation of Time

scripts/generate_release.py

@@ -78,8 +78,14 @@
     logger.info("Copying helper scripts to temp directory")
     shutil.copy(project_root / "scripts/setup.py", temp_dir / "setup.py")
     for script in project_root.glob("scripts/run_*.py"):
+        if "run_comp_" in str(script):
+            continue
         shutil.copy(script, temp_dir)
 
+    script_dir = temp_dir / "scripts"
+    script_dir.mkdir()
+    shutil.copy(project_root / "scripts/run_comp_match.py", script_dir)
+
     logger.info("Copying example code to temp directory")
     shutil.copytree(project_root / "example_robots", temp_dir / "example_robots")
 

scripts/run_comp_match.py

@@ -0,0 +1,298 @@
+#!/usr/bin/env python3
+"""A script to run a competition match."""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from zipfile import ZipFile
+
+if (Path(__file__).parents[1] / 'simulator/VERSION').exists():
+    # Running in release mode, run_simulator will be in folder above
+    sys.path.append(str(Path(__file__).parents[1]))
+
+from run_simulator import get_webots_parameters
+
+NUM_ZONES = 4
+GAME_DURATION_SECONDS = 150
+
+
+class MatchParams(argparse.Namespace):
+    """Parameters for running a competition match."""
+
+    archives_dir: Path
+    match_num: int
+    teams: list[str]
+    duration: int
+    video_enabled: bool
+    video_resolution: tuple[int, int]
+
+
+def load_team_code(
+    usercode_dir: Path,
+    arena_root: Path,
+    match_parameters: MatchParams,
+) -> None:
+    """Load the team code into the arena root."""
+    for zone_id, tla in enumerate(match_parameters.teams):
+        zone_path = arena_root / f"zone_{zone_id}"
+
+        if zone_path.exists():
+            shutil.rmtree(zone_path)
+
+        if tla == '-':
+            # no team in this zone
+            continue
+
+        zone_path.mkdir()
+        with ZipFile(usercode_dir / f'{tla}.zip') as zipfile:
+            zipfile.extractall(zone_path)
+
+
+def generate_match_file(save_path: Path, match_parameters: MatchParams) -> None:
+    """Write the match file to the arena root."""
+    match_file = save_path / 'match.json'
+
+    # Use a format that is compatible with SRComp
+    match_file.write_text(json.dumps(
+        {
+            'match_number': match_parameters.match_num,
+            'arena_id': 'Simulator',
+            'teams': {
+                tla: {'zone': idx}
+                for idx, tla in enumerate(match_parameters.teams)
+                if tla != '-'
+            },
+            'duration': match_parameters.duration,
+            'recording_config': {
+                'enabled': match_parameters.video_enabled,
+                'resolution': list(match_parameters.video_resolution)
+            }
+        },
+        indent=4,
+    ))
+
+
+def set_comp_mode(arena_root: Path) -> None:
+    """Write the mode file to indicate that the competition is running."""
+    (arena_root / 'mode.txt').write_text('comp')
+
+
+def archive_zone_files(
+    team_archives_dir: Path,
+    arena_root: Path,
+    zone: int,
+    match_id: str,
+) -> None:
+    """Zip the files in the zone directory and save them to the team archives directory."""
+    zone_dir = arena_root / f'zone_{zone}'
+
+    shutil.make_archive(str(team_archives_dir / f'{match_id}-zone-{zone}'), 'zip', zone_dir)
+
+
+def archive_zone_folders(
+    archives_dir: Path,
+    arena_root: Path,
+    teams: list[str],
+    match_id: str,
+) -> None:
+    """Zip the zone folders and save them to the archives directory."""
+    for zone_id, tla in enumerate(teams):
+        if tla == '-':
+            # no team in this zone
+            continue
+
+        tla_dir = archives_dir / tla
+        tla_dir.mkdir(exist_ok=True)
+
+        archive_zone_files(tla_dir, arena_root, zone_id, match_id)
+
+
+def archive_match_recordings(archives_dir: Path, arena_root: Path, match_id: str) -> None:
+    """Copy the video, animation, and image files to the archives directory."""
+    recordings_dir = archives_dir / 'recordings'
+    recordings_dir.mkdir(exist_ok=True)
+
+    match_recordings = arena_root / 'recordings'
+
+    # Copy the video file
+    video_file = match_recordings / f'{match_id}.mp4'
+    if video_file.exists():
+        shutil.copy(video_file, recordings_dir)
+
+    # Copy the animation files
+    animation_files = [
+        match_recordings / f'{match_id}.html',
+        match_recordings / f'{match_id}.json',
+        match_recordings / f'{match_id}.x3d',
+        match_recordings / f'{match_id}.css',
+    ]
+    for animation_file in animation_files:
+        shutil.copy(animation_file, recordings_dir)
+
+    # Copy the animation textures
+    # Every match will have the same textures, so we only need one copy of them
+    textures_dir = match_recordings / 'textures'
+    shutil.copytree(textures_dir, recordings_dir / 'textures', dirs_exist_ok=True)
+
+    # Copy the image file
+    image_file = match_recordings / f'{match_id}.jpg'
+    shutil.copy(image_file, recordings_dir)
+
+
+def archive_match_file(archives_dir: Path, match_file: Path, match_number: int) -> None:
+    """
+    Copy the match file (which may contain scoring data) to the archives directory.
+
+    This also renames the file to be compatible with SRComp.
+    """
+    matches_dir = archives_dir / 'matches'
+    matches_dir.mkdir(exist_ok=True)
+
+    # SRComp expects YAML files. JSON is a subset of YAML, so we can just rename the file.
+    completed_match_file = matches_dir / f'{match_number:0>3}.yaml'
+
+    shutil.copy(match_file, completed_match_file)
+
+
+def archive_supervisor_log(archives_dir: Path, arena_root: Path, match_id: str) -> None:
+    """Archive the supervisor log file."""
+    log_archive_dir = archives_dir / 'supervisor_logs'
+    log_archive_dir.mkdir(exist_ok=True)
+
+    log_file = arena_root / f'supervisor-log-{match_id}.txt'
+
+    shutil.copy(log_file, log_archive_dir)
+
+
+def execute_match(arena_root: Path) -> None:
+    """Run Webots with the right world."""
+    # Webots is only on the PATH on Linux so we have a helper function to find it
+    try:
+        webots, world_file = get_webots_parameters()
+    except RuntimeError:
+        raise FileNotFoundError("Webots executable not found.")
+
+    sim_env = os.environ.copy()
+    sim_env['ARENA_ROOT'] = str(arena_root)
+    try:
+        subprocess.check_call(
+            [
+                str(webots),
+                '--batch',
+                '--stdout',
+                '--stderr',
+                '--mode=realtime',
+                str(world_file),
+            ],
+            env=sim_env,
+        )
+    except subprocess.CalledProcessError as e:
+        # TODO review log output here
+        raise RuntimeError(f"Webots failed with return code {e.returncode}") from e
+
+
+def run_match(match_parameters: MatchParams) -> None:
+    """Run the match in a temporary directory and archive the results."""
+    with TemporaryDirectory(suffix=f'match-{match_parameters.match_num}') as temp_folder:
+        arena_root = Path(temp_folder)
+        match_num = match_parameters.match_num
+        match_id = f'match-{match_num}'
+        archives_dir = match_parameters.archives_dir
+
+        # unzip teams code into zone_N folders under this folder
+        load_team_code(archives_dir, arena_root, match_parameters)
+        # Create info file to tell the comp supervisor what match this is
+        # and how to handle recordings
+        generate_match_file(arena_root, match_parameters)
+        # Set mode file to comp
+        set_comp_mode(arena_root)
+
+        try:
+            # Run webots with the right world
+            execute_match(arena_root)
+        except (FileNotFoundError, RuntimeError) as e:
+            print(f"Failed to run match: {e}")
+            # Save the supervisor log as it may contain useful information
+            archive_supervisor_log(archives_dir, arena_root, match_id)
+            raise
+
+        # Archive the supervisor log first in case any collation fails
+        archive_supervisor_log(archives_dir, arena_root, match_id)
+        # Zip up and collect all files for each zone
+        archive_zone_folders(archives_dir, arena_root, match_parameters.teams, match_id)
+        # Collect video, animation & image
+        archive_match_recordings(archives_dir, arena_root, match_id)
+        # Collect ancillary files
+        archive_match_file(archives_dir, arena_root / 'match.json', match_num)
+
+
+def parse_args() -> MatchParams:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Run a competition match.")
+
+    parser.add_argument(
+        'archives_dir',
+        help=(
+            "The directory containing the teams' robot code, as Zip archives "
+            "named for the teams' TLAs. This directory will also be used as the "
+            "root for storing the resulting logs and recordings."
+        ),
+        type=Path,
+    )
+    parser.add_argument(
+        'match_num',
+        type=int,
+        help="The number of the match to run.",
+    )
+    parser.add_argument(
+        'teams',
+        nargs=NUM_ZONES,
+        help=(
+            "TLA of the team in each zone, in order from zone 0 to "
+            f"{NUM_ZONES - 1}. Use dash (-) for an empty zone. "
+            "Must specify all zones."
+        ),
+        metavar='tla',
+    )
+    parser.add_argument(
+        '--duration',
+        help="The duration of the match (in seconds).",
+        type=int,
+        default=GAME_DURATION_SECONDS,
+    )
+    parser.add_argument(
+        '--no-record',
+        help=(
+            "Inhibit creation of the MPEG video, the animation is unaffected. "
+            "This can greatly increase the execution speed on GPU limited systems "
+            "when the video is not required."
+        ),
+        action='store_false',
+        dest='video_enabled',
+    )
+    parser.add_argument(
+        '--resolution',
+        help="Set the resolution of the produced video.",
+        type=int,
+        nargs=2,
+        default=[1920, 1080],
+        metavar=('width', 'height'),
+        dest='video_resolution',
+    )
+    return parser.parse_args(namespace=MatchParams())
+
+
+def main() -> None:
+    """Run a competition match entrypoint."""
+    match_parameters = parse_args()
+    run_match(match_parameters)
+
+
+if __name__ == '__main__':
+    main()