OpenTelemetry Tracing

.github/workflows/benchmark.yml

@@ -0,0 +1,146 @@
+name: benchmark
+
+on:
+  workflow_dispatch:
+  pull_request:
+    paths:
+      - "adbpyg_adapter/*.py"
+      - "benchmark/*.py"
+  push:
+      branches:
+        - master
+      paths:
+        - "adbpyg_adapter/*.py"
+        - "benchmark/*.py"
+
+jobs:
+  generate_master_traces:
+    if: github.event_name == 'push'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Create new branch
+        run: git checkout -b actions/benchmark
+
+      - name: Set branch upstream
+        run: git push -u origin actions/benchmark
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          cache: 'pip'
+          cache-dependency-path: setup.py
+
+      - name: Start ArangoDB Instance
+        run: docker run -d --name arangodb -p 8529:8529 -e ARANGO_ROOT_PASSWORD= arangodb/arangodb
+
+      - name: Start Jaeger Instance
+        run: docker run -d --name jaeger --rm -e COLLECTOR_ZIPKIN_HOST_PORT=:9411 -p 16686:16686 -p 4317:4317 -p 4318:4318 -p 9411:9411 jaegertracing/all-in-one:latest
+
+      - name: Install packages
+        run: |
+          pip install torch==2.1.0
+          pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-$(python -c 'import torch; print(torch.__version__.split("+")[0])')+cpu.html
+          pip install -e '.[dev, tracing]'
+
+      - name: Run Python Script
+        run: python benchmark/write.py --output_dir master
+
+      - name: Echo PyG to ArangoDB
+        run: cat benchmark/traces/master/pyg_to_arangodb.json | jq .
+
+      - name: Echo ArangoDB to PyG 
+        run: cat benchmark/traces/master/arangodb_to_pyg.json | jq .
+
+      - name: Make commit for auto-generated benchmark files
+        uses: EndBug/add-and-commit@v9
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          add: "./benchmark/traces/master/*.json"
+          new_branch: actions/benchmark
+          message: "generate benchmark files for $GITHUB_SHA"
+
+      - name: Create pull request for the auto generated benchmark
+        run: |
+          echo "PR_URL=$(gh pr create \
+            --title "benchmark: $GITHUB_SHA" \
+            --body "beep boop, i am a robot ($GITHUB_SHA)" \
+            --label documentation)" >> $GITHUB_ENV
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Alert developer of open PR
+        run: echo "Benchmark $PR_URL is ready to be merged by developer."
+
+  compare_traces:
+    if: github.event_name == 'pull_request'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+          cache: 'pip'
+          cache-dependency-path: setup.py
+
+      - name: Start ArangoDB Instance
+        run: docker run -d --name arangodb -p 8529:8529 -e ARANGO_ROOT_PASSWORD= arangodb/arangodb
+
+      - name: Start Jaeger Instance
+        run: docker run -d --name jaeger --rm -e COLLECTOR_ZIPKIN_HOST_PORT=:9411 -p 16686:16686 -p 4317:4317 -p 4318:4318 -p 9411:9411 jaegertracing/all-in-one:latest
+
+      - name: Install packages
+        run: |
+          pip install torch==2.1.0
+          pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-$(python -c 'import torch; print(torch.__version__.split("+")[0])')+cpu.html
+          pip install -e '.[dev, tracing]'
+
+      - name: Write PR traces
+        run: python benchmark/write.py --output_dir branch
+
+      - name: Compare PR & Master Traces
+        run: |
+          echo "ROOT_SPAN_COMPARISON<<EOF" >> $GITHUB_ENV
+          python benchmark/compare.py >> $GITHUB_ENV
+          echo "EOF" >> $GITHUB_ENV
+
+      - name: Prepare PR Comment
+        run: |
+          echo 'MESSAGE<<EOF' >> $GITHUB_ENV
+          echo "Benchmark (${{ github.event.pull_request.head.sha }})" >> $GITHUB_ENV
+          echo '' >> $GITHUB_ENV
+          echo '```json' >> $GITHUB_ENV
+          echo "$ROOT_SPAN_COMPARISON" | jq . >> $GITHUB_ENV
+          echo '```' >> $GITHUB_ENV
+          echo '' >> $GITHUB_ENV
+          echo "See the full diff [here](https://github.com/arangoml/pyg-adapter/actions/runs/$GITHUB_RUN_ID)" >> $GITHUB_ENV
+          echo 'EOF' >> $GITHUB_ENV
+
+      - name: Post PR Comment
+        uses: actions/github-script@v7
+        with:
+          github-token: ${{secrets.GITHUB_TOKEN}}
+          script: |
+            const prNumber = context.payload.pull_request.number;
+            const message = process.env.MESSAGE;
+            github.rest.issues.createComment({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              issue_number: prNumber,
+              body: message
+            });
+
+      - name: Echo Full PyG to ArangoDB Diff
+        run: cat benchmark/traces/diff/pyg_to_arangodb.json | jq .
+
+      - name: Echo Full ArangoDB to PyG Diff
+        run: cat benchmark/traces/diff/arangodb_to_pyg.json | jq .

.github/workflows/build.yml

@@ -10,6 +10,7 @@ env:
 jobs:
   build:
     runs-on: ubuntu-latest
+    continue-on-error: true
     strategy:
       matrix:
         python: ["3.8", "3.9", "3.10", "3.11"] # "3.12"
@@ -34,7 +35,7 @@ jobs:
         run: |
           pip install torch==2.1.0
           pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-$(python -c 'import torch; print(torch.__version__.split("+")[0])')+cpu.html
-          pip install -e .[dev]
+          pip install -e '.[dev, tracing]'
 
       - name: Run black
         run: black --check --verbose --diff --color ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}

.gitignore

@@ -118,4 +118,8 @@ adbpyg_adapter/version.py
 .vscode
 
 # PyG Data
-tests/data/pyg
+tests/data/pyg
+
+# "Current" Benchmark Results
+benchmark/traces/branch/*.json
+benchmark/traces/diff/*.json

README.md

@@ -289,15 +289,16 @@ Prerequisite: `arangorestore`
 2. `cd pyg-adapter`
 3. (create virtual environment of choice)
 4. `pip install torch`
-5. `pip install -e .[dev]`
+5. `pip install -e '.[dev, tracing]'`
 6. (create an ArangoDB instance with method of choice)
 7. `pytest --url <> --dbName <> --username <> --password <>`
 
-**Note**: A `pytest` parameter can be omitted if the endpoint is using its default value:
+**Note**: `pytest` parameters:
 ```python
 def pytest_addoption(parser):
     parser.addoption("--url", action="store", default="http://localhost:8529")
     parser.addoption("--dbName", action="store", default="_system")
     parser.addoption("--username", action="store", default="root")
     parser.addoption("--password", action="store", default="")
+    parser.addoption("--otlp_endpoint", action="append", default=[]) # OpenTelemetry Protocol Exporter endpoint
 ```