From c68d65e04c1235ec42780b6f3e9d492bbe7b3270 Mon Sep 17 00:00:00 2001
From: Guido Petretto <g.petretto@gmail.com>
Date: Mon, 30 Dec 2024 11:25:28 +0100
Subject: [PATCH] replace docker python interface and initial addition of PBS
 container

---
 pyproject.toml                                |   2 +-
 tests/integration/conftest.py                 | 261 +++++++++---------
 tests/integration/dockerfiles/Dockerfile      |  64 +++++
 tests/integration/dockerfiles/docker-bake.hcl |  10 +
 tests/integration/dockerfiles/pbs_startup.sh  |  23 ++
 tests/integration/test_workers.py             |   1 +
 6 files changed, 237 insertions(+), 124 deletions(-)
 create mode 100644 tests/integration/dockerfiles/pbs_startup.sh

diff --git a/pyproject.toml b/pyproject.toml
index 93f1e498..652f5377 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["pre-commit>=3.0.0"]
-tests = ["docker ~= 7.0", "pytest ~= 8.0", "pytest-cov >= 4,< 6", "pytest-durations ~= 1.3", "pytest-mock ~= 3.14"]
+tests = ["docker ~= 7.0", "pytest ~= 8.0", "pytest-cov >= 4,< 6", "pytest-durations ~= 1.3", "pytest-mock ~= 3.14", "python-on-whales"]
 docs = [
     "autodoc_pydantic>=2.0.0",
     "pydata-sphinx-theme",
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index de9edf41..75dfb57e 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -9,10 +9,10 @@
 from functools import partialmethod
 from pathlib import Path
 
-import docker
 import fabric
 import pytest
-from docker.models.containers import Container
+from python_on_whales import DockerClient
+from python_on_whales import docker as docker_pow
 
 
 @pytest.fixture(autouse=True)
@@ -62,6 +62,12 @@ def sge_ssh_port():
     return _get_free_port()
 
 
+@pytest.fixture(scope="session")
+def pbs_ssh_port():
+    """The exposed local port for SSH connections to the queue container."""
+    return _get_free_port()
+
+
 @pytest.fixture(scope="session")
 def db_port():
     """The exposed local port for connections to the MongoDB stores."""
@@ -70,139 +76,134 @@ def db_port():
 
 @pytest.fixture(scope="session")
 def docker_client():
-    return docker.from_env()
-
+    return docker_pow
 
-def build_and_launch_container(
-    docker_client: docker.client.DockerClient,
-    dockerfile: Path | None = None,
-    image_name: str | None = None,
-    ports: dict[str, int] | None = None,
-    buildargs: dict[str, str] | None = None,
-):
-    """Builds and/or launches a container, returning the container object.
 
-    Parameters
-    ----------
-        docker_client: The local docker client.
-        dockerfile: An optional location of a dockerfile to build.
-        image_name: Either the tag to attach to the built image, or an image
-            name to pull from the web (may require authenticated docker client).
-        ports: A port specification to use for the launched container.
+# @pytest.fixture(scope="session", autouse=True)
+def bake_containers():
+    hcl_path = Path(__file__).parent.resolve() / "dockerfiles/docker-bake.hcl"
+    docker_pow.buildx.bake(
+        targets=["slurm", "sge"],
+        # targets=["slurm", "sge", "pbs"],
+        files=hcl_path,
+        set={"*.context": str(Path(__file__).parent.parent.parent.resolve())},
+    )
 
-    Yields
-    ------
-        The launched container object, then stops the container after use.
 
-    """
-    if dockerfile is not None:
-        print(f" * Building {image_name}")
-        _, logs = docker_client.images.build(
-            path=str(Path(__file__).parent.parent.parent.resolve()),
-            dockerfile=dockerfile,
-            buildargs=buildargs,
-            tag=image_name,
-            rm=True,
-            quiet=False,
-        )
+@pytest.fixture(scope="session", autouse=True)
+def compose_containers(slurm_ssh_port, sge_ssh_port, pbs_ssh_port, db_port):
+    compose_yaml = f"""
+name: jobflow_remote_testing
+services:
+  mongo:
+    image: mongo:7
+    ports:
+      - "{db_port}:27017"
+    restart: always
+    container_name: mongo_container
+    healthcheck:
+      test: ["CMD", "mongosh", "--eval", "db.runCommand('ping').ok"]
+      interval: 1s
+      timeout: 1s
+      retries: 28
+      start_period: 2s
+
+  jobflow_remote_testing_slurm:
+    image: jobflow-remote-testing-slurm:latest
+    container_name: jobflow_testing_slurm
+    ports:
+      - "{slurm_ssh_port}:22"
+    stdin_open: true
+    tty: true
+    healthcheck:
+      test: ["CMD", "bash", "-c", "</dev/tcp/localhost/22"]
+      interval: 1s
+      timeout: 1s
+      retries: 30
+      start_period: 2s
+
+  jobflow_remote_testing_sge:
+    image: jobflow-remote-testing-sge:latest
+    container_name: jobflow_testing_sge
+    ports:
+      - "{sge_ssh_port}:22"
+    stdin_open: true
+    tty: true
+    healthcheck:
+      test: ["CMD", "bash", "-c", "</dev/tcp/localhost/22"]
+      interval: 1s
+      timeout: 1s
+      retries: 30
+      start_period: 2s
+"""
+    with tempfile.NamedTemporaryFile("wt", suffix="compose.yaml", delete=False) as f:
+        f.write(compose_yaml)
+        f.flush()
+
+        docker_client = DockerClient(compose_files=[f.name])
+        try:
+            print("\n * Launching compose...")
 
-        for step in logs:
-            if step.get("stream"):
-                print(step["stream"], end="")
-
-    try:
-        print(f"\n * Launching container for {image_name}...")
-        container = docker_client.containers.run(
-            image_name,
-            detach=True,
-            remove=False,
-            auto_remove=False,
-            tty=True,
-            ports=ports,
-        )
-        assert isinstance(container, Container)
-        print(" * Waiting for container to be ready...", end="")
-        max_retries = 30
-        while (retries := 0) < max_retries:
-            if container.status == "running":
-                print(f"\n{container.logs().decode()}\n")
-                print(f"\n * Container {container.id} launched.")
-                break
-            if container.status == "exited":
-                logs = f"\n{container.logs().decode()}\n"
+            compose_version = docker_client.compose.up(
+                detach=True,
+            )
+            print(" * Waiting for container to be ready...", end="")
+            max_retries = 30
+            retries = 0
+            while retries < max_retries:
+                containers = docker_client.compose.ps()
+
+                if all(
+                    c.state.health and c.state.health.status == "healthy"
+                    for c in containers
+                ):
+                    print(f"\n{docker_client.compose.logs()}\n")
+                    print(f"\n * {compose_version} launched.")
+                    break
+
+                exited = [c.id for c in containers if c.state.status == "exited"]
+                if any(exited):
+                    logs = f"\n{docker_client.compose.logs()}\n"
+                    pytest.fail(
+                        f"Containers {', '.join(exited)} exited before being ready.\nFull logs: {logs}"
+                    )
+
+                print(".", end="")
+                time.sleep(1)
+                retries += 1
+            else:
+                not_started = [
+                    c.id
+                    for c in containers
+                    if not c.state.health or c.state.health.status != "healthy"
+                ]
+                logs = f"\n{docker_client.compose.logs()}\n"
                 pytest.fail(
-                    f"Container {container.name!r} ({container.image}) exited before being ready.\nFull logs: {logs}"
+                    f"Containers {', '.join(not_started)} did not start in time. Full logs: {logs}"
                 )
-            print(".", end="")
-            time.sleep(1)
-            retries += 1
-            container.reload()
-        else:
-            logs = f"\n{container.logs().decode()}\n"
-            pytest.fail(
-                f"Container {container.name!r} ({container.image}) did not start in time. Full\nlogs: {logs}"
-            )
 
-        yield container
-    finally:
-        try:
-            print(f"\n * Stopping container {container.id}...")
-            try:
-                container.stop()
-            except (docker.errors.APIError, docker.errors.NotFound):
-                pass
-            try:
-                container.kill()
-            except (docker.errors.APIError, docker.errors.NotFound):
-                pass
+            yield docker_client
+        finally:
             try:
-                container.remove()
-            except (docker.errors.APIError, docker.errors.NotFound):
-                pass
-            print(" * Done!")
-        except Exception as exc:
-            print(f" x Failed to stop container: {exc}")
+                print("\n * Stopping containers...")
+                try:
+                    docker_client.compose.stop()
+                except Exception:
+                    pass
 
+                try:
+                    docker_client.compose.kill()
+                except Exception:
+                    pass
 
-@pytest.fixture(scope="session", autouse=True)
-def slurm_container(docker_client, slurm_ssh_port):
-    """Build and launch a container running Slurm and SSH, exposed on a random available
-    port."""
-    ports = {"22/tcp": slurm_ssh_port}
-    yield from build_and_launch_container(
-        docker_client,
-        Path("./tests/integration/dockerfiles/Dockerfile"),
-        "jobflow-remote-slurm:latest",
-        ports=ports,
-        buildargs={"QUEUE_SYSTEM": "slurm"},
-    )
-
+                try:
+                    docker_client.compose.rm(volumes=True)
+                except Exception:
+                    pass
 
-@pytest.fixture(scope="session", autouse=True)
-def sge_container(docker_client, sge_ssh_port):
-    """Build and launch a container running SGE and SSH, exposed on a random available
-    port."""
-    ports = {"22/tcp": sge_ssh_port}
-    yield from build_and_launch_container(
-        docker_client,
-        Path("./tests/integration/dockerfiles/Dockerfile"),
-        "jobflow-remote-sge:latest",
-        ports=ports,
-        buildargs={"QUEUE_SYSTEM": "sge"},
-    )
-
-
-@pytest.fixture(scope="session", autouse=True)
-def mongo_container(docker_client, db_port):
-    """Build and launch a container running MongoDB, exposed on a random available
-    port."""
-    ports = {"27017/tcp": db_port}
-    yield from build_and_launch_container(
-        docker_client,
-        dockerfile=None,
-        image_name="mongo:7",
-        ports=ports,
-    )
+                print(" * Done!")
+            except Exception as exc:
+                print(f" x Failed to stop container: {exc}")
 
 
 @pytest.fixture(scope="session")
@@ -216,6 +217,7 @@ def write_tmp_settings(
     store_database_name,
     slurm_ssh_port,
     sge_ssh_port,
+    pbs_ssh_port,
     db_port,
 ):
     """Collects the various sub-configs and writes them to a temporary file in a
@@ -305,6 +307,19 @@ def write_tmp_settings(
                 pre_run="source /home/jobflow/.venv/bin/activate",
                 connect_kwargs={"allow_agent": False, "look_for_keys": False},
             ),
+            # "test_remote_pbs_worker": dict(
+            #     type="remote",
+            #     host="localhost",
+            #     port=pbs_ssh_port,
+            #     scheduler_type="pbs",
+            #     work_dir="/home/jobflow/jfr",
+            #     user="jobflow",
+            #     password="jobflow",
+            #     scheduler_username="jobflow",
+            #     pre_run="source /home/jobflow/.venv/bin/activate",
+            #     connect_kwargs={"allow_agent": False, "look_for_keys": False},
+            #     resources={"walltime": "00:05:00", "select": "select=1:1"},
+            # ),
             "test_remote_limited_worker": dict(
                 type="remote",
                 host="localhost",
diff --git a/tests/integration/dockerfiles/Dockerfile b/tests/integration/dockerfiles/Dockerfile
index 1eea6234..664041a2 100644
--- a/tests/integration/dockerfiles/Dockerfile
+++ b/tests/integration/dockerfiles/Dockerfile
@@ -4,6 +4,7 @@ ARG USERNAME=jobflow
 
 # Install jobflow-remote in its own layer with uv
 FROM --platform=linux/amd64 ubuntu:22.04 as jobflow-remote-install
+# FROM ubuntu:22.04 as jobflow-remote-install
 
 ARG USERNAME=jobflow
 ENV USERNAME=${USERNAME}
@@ -111,8 +112,70 @@ RUN echo "export SGE_ROOT=/opt/sge" >> /etc/profile && \
 
 #HEALTHCHECK CMD qstat -f || exit 1
 
+# -== PBS ==-
+FROM ubuntu:22.04 as pbs
+
+ARG USERNAME=jobflow
+ENV USERNAME=${USERNAME}
+USER root
+
+ENV DEBIAN_FRONTEND noninteractive
+
+RUN apt update && \
+    apt install -y \
+        gcc \
+        make \
+        libtool \
+        libhwloc-dev \
+        libx11-dev \
+        libxt-dev \
+        libedit-dev \
+        libical-dev \
+        ncurses-dev \
+        perl \
+        postgresql-server-dev-all \
+        postgresql-contrib \
+        unzip \
+        tcl-dev \
+        tk-dev \
+        swig \
+        libexpat-dev \
+        libssl-dev \
+        libxext-dev \
+        libxft-dev \
+        autoconf \
+        automake \
+        g++ \
+        expat \
+        libedit2 \
+        postgresql \
+        sendmail-bin \
+        libical3 \
+        libcjson-dev \
+        curl \
+        python3-dev && \
+    apt clean && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /tmp
+RUN cd /tmp &&\
+    curl -L https://github.com/openpbs/openpbs/archive/refs/tags/v23.06.06.tar.gz -o src.tgz&&\
+    tar xfz src.tgz &&\
+    cd openpbs-* &&\
+    ./autogen.sh &&\
+    ./configure --prefix=/opt/pbs &&\
+    make &&\
+    make install &&\
+    /opt/pbs/libexec/pbs_postinstall &&\
+    chmod 04755 /opt/pbs/sbin/pbs_iff /opt/pbs/sbin/pbs_rcp &&\
+    chmod 0700 /opt/pbs/bin/pbs_topologyinfo /opt/pbs/sbin/pbs_mom /opt/pbs/sbin/pbs_sched /opt/pbs/sbin/pbs_server
+
+COPY ./tests/integration/dockerfiles/pbs_startup.sh /etc/startup.sh
+
+
 # 2) Add common options from ubuntu server
 FROM --platform=linux/amd64 ${QUEUE_SYSTEM} as jobflow
+# FROM ${QUEUE_SYSTEM} as jobflow
 
 ARG USERNAME=jobflow
 ENV USERNAME=${USERNAME}
@@ -157,4 +220,5 @@ RUN chown -R ${USERNAME} /home/${USERNAME} && \
     chmod -R 755 /home/${USERNAME}
 USER ${USERNAME}
 
+# HEALTHCHECK --interval=1s  \
 CMD sudo /etc/startup.sh ; /bin/bash -l
diff --git a/tests/integration/dockerfiles/docker-bake.hcl b/tests/integration/dockerfiles/docker-bake.hcl
index 4ea10cef..24931bc7 100644
--- a/tests/integration/dockerfiles/docker-bake.hcl
+++ b/tests/integration/dockerfiles/docker-bake.hcl
@@ -24,3 +24,13 @@ target "sge" {
         "jobflow-remote-testing-sge:latest"
     ]
 }
+
+target "pbs" {
+    dockerfile = "./tests/integration/dockerfiles/Dockerfile"
+    args = {
+        QUEUE_SYSTEM = "pbs"
+    }
+    tags = [
+        "jobflow-remote-testing-pbs:latest"
+    ]
+}
diff --git a/tests/integration/dockerfiles/pbs_startup.sh b/tests/integration/dockerfiles/pbs_startup.sh
new file mode 100644
index 00000000..4229a4d3
--- /dev/null
+++ b/tests/integration/dockerfiles/pbs_startup.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+# Startup script for Slurm container, vendored from https://github.com/nathan-hess/docker-slurm/blob/a62133d66d624d9ff0ccefbd41a0b1b2abcb9925/dockerfile_base/startup.sh
+
+# Determine whether script is running as root
+sudo_cmd=""
+if [ "$(id -u)" != "0" ]; then
+    sudo_cmd="sudo"
+    sudo -k
+fi
+
+sed -i -e "s/PBS_SERVER=.*/PBS_SERVER=$(hostname)/" -e "s/PBS_START_MOM=0/PBS_START_MOM=1/" /etc/pbs.conf
+sed -i "s/\$clienthost .*/\$clienthost $(hostname)/" /var/spool/pbs/mom_priv/config
+LANG=C /etc/init.d/pbs start
+
+#enable history
+/opt/pbs/bin/qmgr -c "set server job_history_enable=True"
+
+service ssh start
+
+# Revoke sudo permissions
+if [[ ${sudo_cmd} ]]; then
+    sudo -k
+fi
diff --git a/tests/integration/test_workers.py b/tests/integration/test_workers.py
index 97e7f4be..8e1097a1 100644
--- a/tests/integration/test_workers.py
+++ b/tests/integration/test_workers.py
@@ -7,6 +7,7 @@
     reason="Only run integration tests in CI, unless forced with 'CI' env var",
 )
 
+# WORKERS = ["test_local_worker", "test_remote_slurm_worker", "test_remote_sge_worker", "test_remote_pbs_worker"]
 WORKERS = ["test_local_worker", "test_remote_slurm_worker", "test_remote_sge_worker"]
 
 MAX_TRY_SECONDS = 120