Add D2 diagram and add it to the README

.github/workflows/update-readme.yml

@@ -0,0 +1,56 @@
+name: Update README.md
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  update_readme:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+        with:
+          persist-credentials: false
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v3
+        with:
+          python-version: '3.x'
+
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install markdown-code-runner
+          pip install -e .
+
+      - name: Install d2
+        run: |
+          curl -fsSL https://d2lang.com/install.sh | sh -s --
+
+      - name: Run markdown-code-runner
+        run: markdown-code-runner README.md
+
+      - name: Commit updated README.md
+        id: commit
+        run: |
+          git add -u .
+          git config --local user.email "github-actions[bot]@users.noreply.github.com"
+          git config --local user.name "github-actions[bot]"
+          if git diff --quiet && git diff --staged --quiet; then
+            echo "No changes in README.md, skipping commit."
+            echo "commit_status=skipped" >> $GITHUB_ENV
+          else
+            git commit -m "Update README.md"
+            echo "commit_status=committed" >> $GITHUB_ENV
+          fi
+
+      - name: Push changes
+        if: env.commit_status == 'committed'
+        uses: ad-m/github-push-action@master
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          branch: ${{ github.head_ref }}

README.md

@@ -18,6 +18,7 @@ This is an asynchronous job scheduler for [`Adaptive`](https://github.com/python
 - [:dart: Design Goals](#dart-design-goals)
 - [:test_tube: How does it work?](#test_tube-how-does-it-work)
 - [:mag: But how does it *really* work?](#mag-but-how-does-it-really-work)
+- [:bar_chart: Flow diagram](#bar_chart-flow-diagram)
 - [:notebook: Jupyter Notebook Example](#notebook-jupyter-notebook-example)
 - [:computer: Installation](#computer-installation)
 - [:hammer_and_wrench: Development](#hammer_and_wrench-development)
@@ -147,6 +148,16 @@ And use `scheduler.cancel(job_names)` to cancel the jobs.
 
 You don't actually ever have to leave the Jupyter notebook; take a look at the [`example notebook`](https://github.com/basnijholt/adaptive-scheduler/blob/main/example.ipynb).
 
+## :bar_chart: Flow diagram
+
+This is a flow diagram of the different objects.
+
+<!-- CODE:BASH:START -->
+<!-- d2 docs/diagram.d2 docs/source/_static/diagram.svg -->
+<!-- CODE:END -->
+
+![](./docs/source/_static/diagram.svg)
+
 ## :notebook: Jupyter Notebook Example
 
 See [`example.ipynb`](https://github.com/basnijholt/adaptive-scheduler/blob/main/example.ipynb).

adaptive_scheduler/_scheduler/slurm.py

@@ -225,11 +225,16 @@ def start_job(self, name: str) -> None:
     def queue(self, *, me_only: bool = True) -> dict[str, dict[str, str]]:
         """Get the queue of jobs."""
         python_format = {
-            "jobid": 100,
-            "name": 100,
-            "state": 100,
-            "numnodes": 100,
-            "reasonlist": 4000,
+            "JobID": 100,
+            "Name": 100,
+            "State": 100,
+            "NumNodes": 100,
+            "NumTasks": 100,
+            "ReasonList": 4000,
+            "SubmitTime": 100,
+            "StartTime": 100,
+            "UserName": 100,
+            "Partition": 100,
         }  # (key -> length) mapping
 
         slurm_format = ",".join(f"{k}:{v}" for k, v in python_format.items())
@@ -263,10 +268,10 @@ def line_to_dict(line: str) -> dict[str, str]:
 
         squeue = [line_to_dict(line) for line in output.split("\n")]
         states = ("PENDING", "RUNNING", "CONFIGURING")
-        squeue = [info for info in squeue if info["state"] in states]
-        running = {info.pop("jobid"): info for info in squeue}
+        squeue = [info for info in squeue if info["State"] in states]
+        running = {info.pop("JobID"): info for info in squeue}
         for info in running.values():
-            info["job_name"] = info.pop("name")
+            info["job_name"] = info.pop("Name")
         return running
 
     @cached_property

docs/diagram.d2

@@ -0,0 +1,83 @@
+Adaptive Scheduler: {
+  near: top-center
+  shape: text
+  style.font-size: 45
+}
+
+Cluster: {
+  icon: https://icons.terrastruct.com/tech/022-server.svg
+}
+Cluster.SLURM: {
+  shape: queue
+  icon: https://icons.terrastruct.com/azure%2FManagement%20and%20Governance%20Service%20Color%2FScheduler%20Job%20Collections.svg
+}
+
+Jupyter notebook: {
+  icon: https://icons.terrastruct.com/dev%2Fpython.svg
+  scheduler: BaseScheduler {
+    shape: class
+    queue(): "list[str, dict[str, Any]]"
+    "write_job_script(name: str, launcher_options: dict[str, Any])": None
+  }
+
+  dbmanager: DataBaseManager {
+    shape: class
+    url: str
+    db_fname: str
+    learners: "list[BaseLearner]"
+    fnames: "list[str | Path]"
+    scheduler: BaseScheduler
+    as_dicts(): "list[dict]"
+    is_done(): bool
+  }
+  jobmanager: JobManager {
+    shape: class
+    job_names: "list[str]"
+    interval: int
+    database_manager: DataBaseManager
+    scheduler: BaseScheduler
+  }
+  runmanager: RunManager {
+    shape: class
+    learners: "list[BaseLearner]"
+    fnames: "list[str | Path]"
+    job_manager: JobManager
+    scheduler: BaseScheduler
+    db_manager: DataBaseManager
+    start(): None
+  }
+
+  scheduler -> dbmanager: Talks to the SLURM and tracks the jobs that are running
+  scheduler -> jobmanager: Talks to the SLURM and DataBaseManager and submits new jobs
+
+  runmanager -> scheduler
+  runmanager -> dbmanager
+  runmanager -> jobmanager
+}
+
+Jupyter notebook.scheduler -> Cluster.SLURM: Talks to the SLURM cluster
+
+Worker node: {
+    icon: https://icons.terrastruct.com/azure%2FDatabases%20Service%20Color%2FVirtual%20Clusters.svg
+  launcher: "laucher.py" {
+    icon: https://icons.terrastruct.com/dev%2Fpython.svg
+    "zmq.Socket"
+    "adaptive.BalancingLearner"
+    learners
+    fnames
+    "adaptive.Runner" {
+        shape: class
+        learner: BaseLearner
+        fname: str | Path
+        start_periodic_saving(): asyncio.Task
+    }
+    "zmq.Socket" -> learners
+    "zmq.Socket" -> fnames
+    learners -> "adaptive.BalancingLearner"
+    fnames -> "adaptive.BalancingLearner"
+    "adaptive.BalancingLearner" -> "adaptive.Runner"
+  }
+}
+
+Worker node.launcher."zmq.Socket" -> Jupyter notebook.dbmanager: "gets learners and fnames"
+Worker node.launcher."zmq.Socket" <- Jupyter notebook.dbmanager: "tell_done()"