Add Multi-version Capability

Adds support for having multiple evaluation python versions installed in
the docker container. A utility to automatically generate correct
dockerfile instructions and nsjail mounts based on the available
versions is also included.

Signed-off-by: Hassan Abouelela <[email protected]>

8 files changed, 245 insertions, 8 deletions

diff --git a/.gitignore b/.gitignore
index ac702d4..2fedf7c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -106,3 +106,6 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+# Project
+config/versions-dev.json
diff --git a/Dockerfile b/Dockerfile
index 6355dac..015e9d1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-# syntax=docker/dockerfile:1
+# THIS FILE IS AUTOGENERATED, DO NOT MODIFY! #
 FROM python:3.11-slim-buster as builder
 
 WORKDIR /nsjail
@@ -20,8 +20,16 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \
 RUN make
 
 # ------------------------------------------------------------------------------
+FROM python:3.11-slim-buster as base-first
+
+FROM python:3.10-slim-buster as base-3-10
+COPY --from=base-first / /
+
+# ------------------------------------------------------------------------------
 FROM python:3.11-slim-buster as base
 
+COPY --from=base-3-10 / /
+
 # Everything will be a user install to allow snekbox's dependencies to be kept
 # separate from the packages exposed during eval.
 ENV PATH=/root/.local/bin:$PATH \
diff --git a/Makefile b/Makefile
index a385a8e..97461c3 100644
--- a/Makefile
+++ b/Makefile
@@ -34,8 +34,12 @@ test:
 report: setup
 	coverage report
 
+.PHONY: prepare-dockerfile
+prepare-versions:
+	python scripts/set_versions.py
+
 .PHONY: build
-build:
+build: prepare-versions
 	docker build -t ghcr.io/python-discord/snekbox:latest .
 
 .PHONY: devsh
diff --git a/config/snekbox.cfg b/config/snekbox.cfg
index 5dd63da..f10fa5d 100644
--- a/config/snekbox.cfg
+++ b/config/snekbox.cfg
@@ -14,7 +14,6 @@ envar: "OPENBLAS_NUM_THREADS=5"
 envar: "MKL_NUM_THREADS=5"
 envar: "VECLIB_MAXIMUM_THREADS=5"
 envar: "NUMEXPR_NUM_THREADS=5"
-envar: "PYTHONPATH=/snekbox/user_base/lib/python3.11/site-packages"
 envar: "PYTHONIOENCODING=utf-8:strict"
 envar: "HOME=home"
 
@@ -99,12 +98,23 @@ mount {
     rw: false
 }
 
+# THE FOLLOWING SECTION IS AUTOGENERATED, DO NOT MODIFY BY HAND
+# mount-section-key
+mount {
+    src: "/usr/local/bin/python3.10"
+    dst: "/usr/local/bin/python3.10"
+    is_bind: true
+    rw: false
+}
+
 mount {
     src: "/usr/local/bin/python3.11"
     dst: "/usr/local/bin/python3.11"
     is_bind: true
     rw: false
 }
+# mount-section-key-end
+# END AUTOGENERATED SECTION
 
 cgroup_mem_max: 52428800
 cgroup_mem_swap_max: 0
@@ -114,8 +124,3 @@ cgroup_pids_max: 6
 cgroup_pids_mount: "/sys/fs/cgroup/pids"
 
 iface_no_lo: true
-
-exec_bin {
-    path: "/usr/local/bin/python"
-    arg: "-BSqu"
-}
diff --git a/config/versions.json b/config/versions.json
new file mode 100644
index 0000000..fd3b997
--- /dev/null
+++ b/config/versions.json
@@ -0,0 +1,14 @@
+[
+  {
+    "image_tag": "3.10-slim-buster",
+    "version_name": "3.10",
+    "display_name": "Python 3.10",
+    "is_main": false
+  },
+  {
+    "image_tag": "3.11-slim-buster",
+    "version_name": "3.11",
+    "display_name": "Python 3.11",
+    "is_main": true
+  }
+]
diff --git a/scripts/in.Dockerfile b/scripts/in.Dockerfile
new file mode 100644
index 0000000..55bc263
--- /dev/null
+++ b/scripts/in.Dockerfile
@@ -0,0 +1,81 @@
+FROM python:{main_version_tag} as builder
+
+WORKDIR /nsjail
+
+RUN apt-get -y update \
+    && apt-get install -y \
+        bison=2:3.3.* \
+        flex=2.6.* \
+        g++=4:8.3.* \
+        gcc=4:8.3.* \
+        git=1:2.20.* \
+        libprotobuf-dev=3.6.* \
+        libnl-route-3-dev=3.4.* \
+        make=4.2.* \
+        pkg-config=0.29-6 \
+        protobuf-compiler=3.6.*
+RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \
+    && git checkout dccf911fd2659e7b08ce9507c25b2b38ec2c5800
+RUN make
+
+# ------------------------------------------------------------------------------
+{python_install_commands}
+# ------------------------------------------------------------------------------
+FROM python:{main_version_tag} as base
+
+COPY --from=base-{final_base} / /
+
+# Everything will be a user install to allow snekbox's dependencies to be kept
+# separate from the packages exposed during eval.
+ENV PATH=/root/.local/bin:$PATH \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_NO_CACHE_DIR=false \
+    PIP_USER=1
+
+RUN apt-get -y update \
+    && apt-get install -y \
+        gcc=4:8.3.* \
+        git=1:2.20.* \
+        libnl-route-3-200=3.4.* \
+        libprotobuf17=3.6.* \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY --from=builder /nsjail/nsjail /usr/sbin/
+RUN chmod +x /usr/sbin/nsjail
+
+# ------------------------------------------------------------------------------
+FROM base as venv
+
+COPY requirements/ /snekbox/requirements/
+WORKDIR /snekbox
+
+# pip installs to the default user site since PIP_USER is set.
+RUN pip install -U -r requirements/requirements.pip
+
+# This must come after the first pip command! From the docs:
+# All RUN instructions following an ARG instruction use the ARG variable
+# implicitly (as an environment variable), thus can cause a cache miss.
+ARG DEV
+
+# Install numpy when in dev mode; one of the unit tests needs it.
+RUN if [ -n "${DEV}" ]; \
+    then \
+        pip install -U -r requirements/coverage.pip \
+        && PYTHONUSERBASE=/snekbox/user_base pip install numpy~=1.19; \
+    fi
+
+# At the end to avoid re-installing dependencies when only a config changes.
+COPY config/ /snekbox/config/
+
+ENTRYPOINT ["gunicorn"]
+CMD ["-c", "config/gunicorn.conf.py"]
+
+# ------------------------------------------------------------------------------
+FROM venv
+
+# Use a separate directory to avoid importing the source over the installed pkg.
+# The venv already installed dependencies, so nothing besides snekbox itself
+# will be installed. Note requirements.pip cannot be used as a constraint file
+# because it contains extras, which pip disallows.
+RUN --mount=source=.,target=/snekbox_src,rw \
+    pip install /snekbox_src[gunicorn,sentry]
diff --git a/scripts/python_version.py b/scripts/python_version.py
new file mode 100644
index 0000000..6c8f25c
--- /dev/null
+++ b/scripts/python_version.py
@@ -0,0 +1,54 @@
+"""
+Parse and return python version information from the versions file.
+
+The version file is read from the environment variable VERSIONS_CONFIG,
+and defaults to config/versions.json otherwise.
+"""
+
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+VERSIONS_FILE = Path(os.getenv("VERSIONS_CONFIG", "config/versions.json"))
+
+
+@dataclass(frozen=True)
+class Version:
+    """A python image available for eval."""
+
+    image_tag: str
+    version_name: str
+    display_name: str
+    is_main: bool
+
+
+_ALL_VERSIONS = None
+_MAIN_VERSION = None
+
+
+def get_all_versions() -> tuple[list[Version], Version]:
+    """
+    Get a list of all available versions for this evaluation.
+
+    Returns a tuple of all versions, and the main version.
+    """
+    # Return a cached result
+    global _ALL_VERSIONS, _MAIN_VERSION
+    if _ALL_VERSIONS is not None:
+        return _ALL_VERSIONS, _MAIN_VERSION
+
+    versions = []
+    main_version: Version | None = None
+
+    for version_json in json.loads(VERSIONS_FILE.read_text("utf-8")):
+        version = Version(**version_json)
+        if version.is_main:
+            main_version = version
+        versions.append(version)
+
+    if main_version is None:
+        raise Exception("Exactly one version must be configured as the main version.")
+
+    _ALL_VERSIONS, _MAIN_VERSION = versions, main_version
+    return versions, main_version
diff --git a/scripts/set_versions.py b/scripts/set_versions.py
new file mode 100644
index 0000000..f98c52c
--- /dev/null
+++ b/scripts/set_versions.py
@@ -0,0 +1,68 @@
+"""Generate a Dockerfile from in.Dockerfile and a version JSON file, and write version info."""
+
+import re
+from pathlib import Path
+from textwrap import dedent
+
+from scripts import python_version
+
+DOCKERFILE_TEMPLATE = Path("scripts/in.Dockerfile").read_text("utf-8")
+DOCKERFILE = Path("Dockerfile")
+JAIL_CONFIG = Path("config/snekbox.cfg")
+
+versions, main_version = python_version.get_all_versions()
+
+# Download and copy multiple python images into one layer
+python_build = ""
+jail_mounts = ""
+previous_layer = "first"
+
+for version in versions:
+    # Configure NSJail mounts
+    jail_mounts += dedent(
+        f"""
+        mount {{
+            src: "/usr/local/bin/python{version.version_name}"
+            dst: "/usr/local/bin/python{version.version_name}"
+            is_bind: true
+            rw: false
+        }}
+    """
+    )
+
+    if version.is_main:
+        # Main is handled separately later
+        continue
+
+    # Add the current version to the Dockerfile
+    layer_name = version.version_name.replace(".", "-")  # Dots aren't valid in layer names
+    python_build += dedent(
+        f"""
+        FROM python:{version.image_tag} as base-{layer_name}
+        COPY --from=base-{previous_layer} / /
+    """
+    )
+    previous_layer = layer_name
+
+# Main version is installed twice, once at the very beginning to make sure
+# its files aren't overwritten, and once at the end which actually makes use of the version
+python_build = f"FROM python:{main_version.image_tag} as base-first\n" + python_build
+
+# Update mounts for python binaries in the NSJail config
+new_config = re.sub(
+    r"(?<=# mount-section-key)[\s\S]+(?=# mount-section-key-end)",
+    jail_mounts,
+    JAIL_CONFIG.read_text("utf-8"),
+)
+JAIL_CONFIG.write_text(new_config, "utf-8")
+
+# Write new dockerfile
+DOCKERFILE.write_text(
+    "# THIS FILE IS AUTOGENERATED, DO NOT MODIFY! #\n"
+    + DOCKERFILE_TEMPLATE.replace("{python_install_commands}", python_build)
+    .replace("{final_base}", previous_layer)
+    .replace("{main_version_tag}", main_version.image_tag),
+    "utf-8",
+)
+
+print("Finished!")