# One image for every role in the cluster environment tests (agent/tests/environments):
# NFS server, Slurm controller, Slurm compute node, plain compute node, squid proxy
# and the relay that puts the test's fake API on the container network.
#
# One image, not five, because the box's disk is the tight resource: torch alone is
# most of the size, and every role but squid/nfs needs it.
FROM python:3.12-slim

ARG TORCH_VERSION=2.8.0

# Package postinst scripts must not try to start services during the build.
RUN printf '#!/bin/sh\nexit 101\n' > /usr/sbin/policy-rc.d && chmod +x /usr/sbin/policy-rc.d \
 && apt-get update \
 && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
      slurm-wlm munge nfs-common nfs-kernel-server rpcbind squid \
      iproute2 iptables procps psmisc curl ca-certificates openssl \
 && rm -rf /var/lib/apt/lists/* /var/cache/apt/*

RUN pip install --no-cache-dir "torch==${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu \
 && pip install --no-cache-dir numpy

# Last, so a new SDK release rebuilds only this layer. The harness normally passes
# an exact version (the newest on PyPI unless PROBE_ENV_SDK_VERSION pins one). With
# PROBE_ENV_SDK_SRC set, lab.py instead passes PROBE_SDK_UNPINNED: the source tree
# it mounts at /sdk-src (PYTHONPATH) shadows whatever lands here, so the installed
# version doesn't matter and pinning would only risk PyPI's simple index refusing
# to serve that exact old version during a maintenance window. Install unpinned
# (pip resolves the newest it can) instead of failing the build over a version
# nothing will actually import.
ARG PROBE_SDK_VERSION
ARG PROBE_SDK_UNPINNED
RUN if [ -n "$PROBE_SDK_UNPINNED" ]; then \
      pip install --no-cache-dir "probe-research"; \
    else \
      test -n "$PROBE_SDK_VERSION" \
      && pip install --no-cache-dir "probe-research==${PROBE_SDK_VERSION}"; \
    fi \
 && python -c "import importlib.metadata as m; print(m.version('probe-research'))" > /etc/probe-sdk-version

# The researcher: the same uid on every node, the way an LDAP account is. Home is
# created at start: on the NFS export for cluster roles, locally otherwise.
RUN groupadd -g 2000 researcher && useradd -M -u 2000 -g 2000 -d /home/researcher -s /bin/bash researcher \
 && mkdir -p /etc/slurm /var/spool/slurmctld /var/spool/slurmd /var/log/slurm /run/munge \
 && (test -s /etc/munge/munge.key || (dd if=/dev/urandom of=/etc/munge/munge.key bs=1024 count=1 2>/dev/null)) \
 && chown -R munge:munge /etc/munge /run/munge && chmod 400 /etc/munge/munge.key

COPY slurm.conf cgroup.conf /etc/slurm/
COPY squid.conf /etc/squid/envtest.conf
COPY entrypoint.sh relay.py /usr/local/bin/
RUN chmod +x /usr/local/bin/entrypoint.sh

ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
