Repository: qutip/qutip-benchmark
Pull request #57: Setup `python` and `qutip_jax` for buildkite pipeline
Author: login: Langhaarzombie
Created: 2025-03-12T02:41:01Z | Merged: 2025-03-24T08:38:15Z

This adds benchmarks for qutip using the buildkite pipeline.

Things left to do so far:

- [x] Unify output directory
- [x] Fine tune parameters for pytest-benchmark


Diff:
diff --git a/.buildkite/benchmark.yml b/.buildkite/benchmark.yml
index e0e4bfcb0..db262301d 100644
--- a/.buildkite/benchmark.yml
+++ b/.buildkite/benchmark.yml
@@ -1,4 +1,5 @@
 steps:
+  ### Julia ###
   - label: "Run QuantumToolbox.jl benchmarks with CUDA"
     plugins:
       - JuliaCI/julia#v1:
@@ -19,12 +20,25 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  - label: "Run QuTiP benchmarks"
+    command: |
+      apt-get update -y
+      apt-get install -y python3-pip
+
+      python3 --version
+      pip3 --version
+
+      pip3 install -r benchmark/qutip/requirements.txt
+
+      python3 -m benchmark.qutip.runbenchmark -v -s
+    artifact_paths:
+      - "benchmark/results/*"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 60
 
   #- label: "Run dynamiqs benchmarks with CUDA"
   #  ...
   #  env:
-  #    GROUP: "CUDA"
+  #    GROUP: "CUDA"
\ No newline at end of file
diff --git a/benchmark/qutip/requirements.txt b/benchmark/qutip/requirements.txt
new file mode 100644
index 000000000..0c5756b06
--- /dev/null
+++ b/benchmark/qutip/requirements.txt
@@ -0,0 +1,5 @@
+numpy>=1.21
+qutip>=5.1
+jax[cuda12]<0.5
+qutip-jax>=0.1
+pytest-benchmark==5.1.0
\ No newline at end of file
diff --git a/benchmark/qutip/runbenchmark.py b/benchmark/qutip/runbenchmark.py
new file mode 100644
index 000000000..804a7af38
--- /dev/null
+++ b/benchmark/qutip/runbenchmark.py
@@ -0,0 +1,23 @@
+import sys
+import pytest
+
+def run_benchmarks():
+    return pytest.main(
+        [
+            "benchmark/qutip/scripts",
+            "--benchmark-only",
+            "--benchmark-columns=Mean,StdDev,rounds,Iterations",
+            "--benchmark-sort=name",
+            "--benchmark-autosave",
+            "--benchmark-storage=./benchmark/results/"
+        ]
+    )
+
+
+def main():
+    exit_code = run_benchmarks()
+    return exit_code
+
+
+if __name__ == "__main__":
+    sys.exit(main())
\ No newline at end of file
diff --git a/benchmark/qutip/scripts/bench_jax.py b/benchmark/qutip/scripts/bench_jax.py
new file mode 100644
index 000000000..65506367f
--- /dev/null
+++ b/benchmark/qutip/scripts/bench_jax.py
@@ -0,0 +1,23 @@
+import pytest
+import jax
+import jax.numpy as jnp
+import qutip_jax as qj
+from qutip import mesolve, basis, sigmax, sigmaz, CoreOptions
+
+@pytest.mark.jax
+def bench_jax_mesolve(benchmark):
+    benchmark.group = "jax:mesolve"
+    qj.set_as_default()
+
+    with jax.default_device(jax.devices("gpu")[0]):
+        opt = { "method": "diffrax", "normalize_output": False }
+
+        with CoreOptions(default_dtype="jax"):
+            delta = jnp.pi
+            g = 0.2
+            H = delta/2.0 * sigmax()
+            c_ops = [jnp.sqrt(g) * sigmaz()]
+            e_ops = [sigmaz()]
+            psi0 = basis(2, 0)
+            tlist = jnp.linspace(0, 10, 100)
+            benchmark(mesolve, H, psi0, tlist, c_ops, e_ops=e_ops, options=opt)
diff --git a/benchmark/qutip/scripts/bench_mesolve.py b/benchmark/qutip/scripts/bench_mesolve.py
new file mode 100644
index 000000000..f5c4ed592
--- /dev/null
+++ b/benchmark/qutip/scripts/bench_mesolve.py
@@ -0,0 +1,18 @@
+import pytest
+import numpy as np
+from qutip import mesolve, basis, sigmax, sigmaz
+
+@pytest.mark.solvers
+def bench_mesolve(benchmark):
+    benchmark.group = "solvers:mesolve"
+
+    delta = np.pi
+    g = 0.2
+
+    H = delta/2.0 * sigmax()
+    c_ops = [np.sqrt(g) * sigmaz()]
+
+    psi0 = basis(2, 0)
+    tlist = np.linspace(0, 10, 100)
+
+    benchmark(mesolve, H, psi0, tlist, c_ops, e_ops=[sigmaz()])
diff --git a/benchmark/qutip/scripts/pytest.ini b/benchmark/qutip/scripts/pytest.ini
new file mode 100644
index 000000000..edb715515
--- /dev/null
+++ b/benchmark/qutip/scripts/pytest.ini
@@ -0,0 +1,7 @@
+[pytest]
+markers =
+    solvers: benchmarks for qutip solvers
+    jax: benchmarks for qutip with jax
+
+python_files = bench_*.py
+python_functions = bench_*
\ No newline at end of file


Reviews:
user: login: ytdHuang
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-13T02:38:20Z
user: login: ytdHuang
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T07:32:17Z
user: login: Langhaarzombie
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T06:59:25Z
user: login: Langhaarzombie
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:00:23Z
user: login: Langhaarzombie
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T08:14:08Z
user: login: Langhaarzombie
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T08:39:16Z
user: login: ytdHuang
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T13:08:12Z
user: login: ytdHuang
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-19T13:08:46Z

Review comments:
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: I think we can remove the `key`
created_at: 2025-03-13T02:35:49Z
author_association: MEMBER
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: pip install -r benchmark/qutip/requirements.txt
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: maybe try:

```suggestion
    command: "pip install -r benchmark/qutip/requirements.txt"
```
created_at: 2025-03-13T02:36:05Z
author_association: MEMBER
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: pip install -r benchmark/qutip/requirements.txt
+
+  - label: "qutip: run benchmark"
+    depends_on: ["pip"]
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: can remove depends on pip
created_at: 2025-03-13T02:36:27Z
author_association: MEMBER
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: pip install -r benchmark/qutip/requirements.txt
+
+  - label: "qutip: run benchmark"
+    depends_on: ["pip"]
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: python -m benchmark.qutip.benchmark -v -s
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: same here 

```suggestion
    command: "python -m benchmark.qutip.benchmark -v -s"
```
created_at: 2025-03-13T02:36:50Z
author_association: MEMBER
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: pip install -r benchmark/qutip/requirements.txt
+
+  - label: "qutip: run benchmark"
+    depends_on: ["pip"]
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: python -m benchmark.qutip.benchmark -v -s
+    artifact_paths:
+      - ".benchmarks/*"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      GROUP: "CUDA"
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: remember to add a line of `timeout_in_minutes: 60` after this.
created_at: 2025-03-13T02:37:34Z
author_association: MEMBER
diff_hunk: @@ -19,10 +20,27 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "qutip: pip install requirements"
+    key: "pip"
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: pip install -r benchmark/qutip/requirements.txt
+
+  - label: "qutip: run benchmark"
+    depends_on: ["pip"]
+    plugins:
+      - docker#v5.12.0:
+          image: "python:3"
+    command: python -m benchmark.qutip.benchmark -v -s
+    artifact_paths:
+      - ".benchmarks/*"
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: maybe make it same as above first ?

```suggestion
      - "benchmarks/results/*"
```
created_at: 2025-03-13T02:38:12Z
author_association: MEMBER
diff_hunk: @@ -19,12 +20,24 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "Run QuTiP benchmarks"
+    command: |
+      apt-get update -y
+      apt-get install -y python3-pip
+
+      pip3 install -r benchmark/qutip/requirements.txt
+
+      python3 --version
+      pip3 --version
+
+      python3 -m benchmark.qutip.runbenchmark -v -s
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 60
path: .buildkite/benchmark.yml
user: login: ytdHuang
body: @Langhaarzombie 

Can you also try if the following works ?

```suggestion
  ### Python ###
  - label: "Setup Python"
    key: "python"
    command: |
      apt-get update -y
      apt-get install -y python3-pip
      python3 --version
      pip3 --version
    agents:
      queue: "juliagpu"
      cuda: "*"
    timeout_in_minutes: 60

  - label: "Run QuTiP benchmarks with CUDA"
    depends_on: ["python"]
    command: |
      pip3 install -r benchmark/qutip/requirements.txt
      python3 -m benchmark.qutip.runbenchmark -v -s
    artifact_paths:
      - "benchmark/results/*"
    agents:
      queue: "juliagpu"
      cuda: "*"
    env:
      GROUP: "CUDA"
    timeout_in_minutes: 60
```
created_at: 2025-03-18T07:29:17Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,3 @@
+numpy>=1.21
+qutip>=5.1
path: benchmark/qutip/requirements.txt
user: login: ytdHuang
body: Shall we also add `qutip_jax` here?

Just make sure we can install it.
created_at: 2025-03-18T07:31:23Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+import numpy as np
+from qutip import mesolve, basis, sigmax, sigmaz
path: benchmark/qutip/scripts/bench_mesolve.py
user: login: ytdHuang
body: maybe also import `qutip_jax` and run the following codes with `jax` backend ?
created_at: 2025-03-18T07:32:09Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+import numpy as np
+from qutip import mesolve, basis, sigmax, sigmaz
path: benchmark/qutip/scripts/bench_mesolve.py
user: login: Langhaarzombie
body: I added a jax benchmark. Now I am trying to figure out, why the GPU backend is not automatically available...
created_at: 2025-03-19T06:59:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,12 +20,24 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "Run QuTiP benchmarks"
+    command: |
+      apt-get update -y
+      apt-get install -y python3-pip
+
+      pip3 install -r benchmark/qutip/requirements.txt
+
+      python3 --version
+      pip3 --version
+
+      python3 -m benchmark.qutip.runbenchmark -v -s
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 60
path: .buildkite/benchmark.yml
user: login: Langhaarzombie
body: Tried that but somehow pip3 cannot be found in the next step. That's why I put everything into one. 
created_at: 2025-03-19T07:00:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,12 +20,24 @@ steps:
       GROUP: "CUDA"
     timeout_in_minutes: 60
 
-  #- label: "Run qutip benchmarks with CUDA"
-  #  ...
-  #  env:
-  #    GROUP: "CUDA"
+  ### Python ###
+  - label: "Run QuTiP benchmarks"
+    command: |
+      apt-get update -y
+      apt-get install -y python3-pip
+
+      pip3 install -r benchmark/qutip/requirements.txt
+
+      python3 --version
+      pip3 --version
+
+      python3 -m benchmark.qutip.runbenchmark -v -s
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 60
path: .buildkite/benchmark.yml
user: login: Langhaarzombie
body: If I understand correctly, the option "depends_on" does not carry over the installations of the previous step.
created_at: 2025-03-19T08:14:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+import numpy as np
+from qutip import mesolve, basis, sigmax, sigmaz
path: benchmark/qutip/scripts/bench_mesolve.py
user: login: Langhaarzombie
body: Looking good now with the latest commits
created_at: 2025-03-19T08:39:16Z
author_association: CONTRIBUTOR
path: benchmark/qutip/scripts/__init__.py
user: login: ytdHuang
body: @Langhaarzombie 

A small question...

What is this empty file for ?

Should we remove it ?
created_at: 2025-03-19T13:08:12Z
author_association: MEMBER

PR comments:
user: login: hodgestar
created_at: 2025-03-12T08:49:48Z
body: @Langhaarzombie Thanks, this is looking good. Is there some way to run a small benchmark that makes use of CUDA within QuTiP? Or does that happen automatically somehow just with CUDA installed?
author_association: CONTRIBUTOR
user: login: Langhaarzombie
created_at: 2025-03-12T09:05:42Z
body: > @Langhaarzombie Thanks, this is looking good. Is there some way to run a small benchmark that makes use of CUDA within QuTiP? Or does that happen automatically somehow just with CUDA installed?

I am not sure at this point. But I will look into it, once we get the pipeline for python to work. Thanks for the input!
author_association: CONTRIBUTOR
user: login: ytdHuang
created_at: 2025-03-15T06:05:10Z
body: Can follow the similar way I tried in PR #58 to setup `python`.

Although this restricts the `python` version to `3.10` due to Ubuntu's `apt-get`, we can at least make the benchmark work first, and think of installing `miniconda` from scratch in the future.
author_association: MEMBER
user: login: ytdHuang
created_at: 2025-03-19T13:10:35Z
body: @Ericgig 

The buildkite pipeline looks good now.

Do you have any further comments about the setup and related codes for `qutip_jax` ?
author_association: MEMBER