add stress-testing script (#709)

* add stress-testing script

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update stress-testing script to set `delay_between_calls`

* use eventlet

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update stress testing script

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add eventlet

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update how stress testing works

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add healtcheck to docker-compose setup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update Dockerfile to use python native image

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Dockerfile

@@ -1,18 +1,14 @@
-FROM continuumio/miniconda3
+FROM python:3.12
+SHELL ["/bin/bash", "--login", "-c"]
+
 WORKDIR /usr/src/app
-COPY ./ ./
-# required for h5py, chemfiles
+
+# required for h5py
 RUN apt update && apt install -y gcc pkg-config libhdf5-dev build-essential
-RUN conda install python=3.11 nodejs
-# RUN conda install conda-forge::packmol
-RUN npm install -g bun
+RUN curl -fsSL https://bun.sh/install | bash
 
-# Make RUN commands use the new environment:
-SHELL ["conda", "run", "--no-capture-output", "-n", "base", "/bin/bash", "-c"]
+COPY ./ ./
 RUN cd app && bun install && bun vite build && cd ..
-RUN pip install -e .[all]
-
+RUN pip install -e .
 
-EXPOSE 5003
-# # The code to run when container is started:
-ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "base", "zndraw", "--port", "5003", "--no-browser"]
+ENTRYPOINT ["zndraw", "--port", "5003", "--no-browser"]

docker-compose.yaml

@@ -0,0 +1,40 @@
+services:
+  zndraw:
+    build: .
+    healthcheck:
+      test: ["CMD", "zndraw", "--healthcheck", "--url", "http://zndraw:5003"]
+      interval: 30s
+      timeout: 10s
+      retries: 5
+    command: --no-standalone
+    restart: unless-stopped
+    ports:
+      - 5003:5003
+    depends_on:
+      - redis
+      - worker
+    environment:
+      - FLASK_STORAGE=redis://redis:6379/0
+      - FLASK_AUTH_TOKEN=super-secret-token
+
+  worker:
+    build: .
+    healthcheck:
+      test: ["CMD", "zndraw", "--healthcheck", "--url", "http://zndraw:5003"]
+      interval: 30s
+      timeout: 10s
+      retries: 5
+    entrypoint: celery -A zndraw_app.make_celery worker --loglevel=info -P eventlet
+    restart: unless-stopped
+    depends_on:
+      - redis
+    environment:
+      - FLASK_STORAGE=redis://redis:6379/0
+      - FLASK_SERVER_URL="http://zndraw:5003"
+      - FLASK_AUTH_TOKEN=super-secret-token
+
+  redis:
+    image: redis:latest
+    restart: always
+    environment:
+      - REDIS_PORT=6379

examples/stress_testing/README.md

@@ -0,0 +1,3 @@
+# Stress Testing
+The following scripts can be used to perform some stress testing on a given ZnDraw instance.
+Do not use them against public servers which you are not hosting yourself.

examples/stress_testing/multi_connection.py

@@ -0,0 +1,19 @@
+import subprocess
+
+import typer
+
+app = typer.Typer()
+
+
+@app.command()
+def main(file: str, n: int, browser: bool = False):
+    cmd = ["zndraw", file]
+    if not browser:
+        cmd.append("--no-browser")
+    # run cmd n times in parallel
+    for _ in range(n):
+        subprocess.Popen(cmd)
+
+
+if __name__ == "__main__":
+    app()

examples/stress_testing/single_connection.py

@@ -0,0 +1,39 @@
+# Run tests with and without eventlet
+import eventlet
+
+eventlet.monkey_patch()
+
+import datetime  # noqa
+import uuid
+import os
+
+import tqdm
+from rdkit2ase import smiles2conformers
+
+#### Import ZnDraw ####
+from zndraw import ZnDraw
+
+vis = ZnDraw(url=os.environ["ZNDRAW_URL"], token=uuid.uuid4().hex)
+#### ------------- ####
+
+# vis._refresh_client.delay_between_calls = datetime.timedelta(milliseconds=10)
+
+conformers = smiles2conformers("CCCCCCCCCO", numConfs=1000)
+
+# append
+for atoms in tqdm.tqdm(conformers, desc="append", ncols=80):
+    vis.append(atoms)
+
+# read
+for i in tqdm.trange(len(vis), desc="getitem", ncols=80):
+    _ = vis[i]
+
+# delete
+for i in tqdm.tqdm(range(len(vis) - 1, -1, -1), desc="delete", ncols=80):
+    del vis[i]
+
+# extend
+vis.extend(conformers)
+
+# read_all
+print(f"len(vis[:]): {len(vis[:])}")

zndraw/upload.py

@@ -31,7 +31,6 @@ def upload(
 
     if not append:
         del vis[:]
-    typer.echo(f"Reading {fileio.name} ...")
 
     generator = get_generator_from_filename(fileio)
 

zndraw/zndraw.py

@@ -327,6 +327,8 @@ def extend(self, values: list[ase.Atoms]):
             isinstance(x, ase.Atoms) for x in values
         ):
             raise ValueError("Unable to parse provided data object")
+        if len(values) == 0:
+            return
 
         if not self.r.exists(f"room:{self.token}:frames") and self.r.exists(
             "room:default:frames"

zndraw_app/cli.py

@@ -15,6 +15,7 @@
 from zndraw.tasks import read_file, read_plots
 from zndraw.upload import upload
 from zndraw.utils import get_port
+from zndraw_app.healthcheck import run_healthcheck
 
 cli = typer.Typer()
 
@@ -129,11 +130,16 @@ def main(
         help="Convert NaN values to None. This is slow and experimental, but if your file contains NaN/inf values, it is required.",
         envvar="ZNDRAW_CONVERT_NAN",
     ),
+    healthcheck: bool = typer.Option(False, help="Run the healthcheck."),
 ):
     """Start the ZnDraw server.
 
     Visualize Trajectories, Structures, and more in ZnDraw.
     """
+    if healthcheck:
+        if url is None:
+            raise ValueError("You need to provide a URL to use the healthcheck feature.")
+        run_healthcheck(url)
     if plots is None:
         plots = []
     if token is not None and url is None:

zndraw_app/healthcheck.py

@@ -0,0 +1,9 @@
+import sys
+
+
+def run_healthcheck(server: str):
+    """Check the health of a server."""
+    from zndraw import ZnDraw
+
+    _ = ZnDraw(url=server, token="healthcheck")
+    sys.exit(0)