Setup Bench (#1048)

.github/actions/poetry_setup/action.yml

@@ -0,0 +1,88 @@
+# An action for setting up poetry install with caching.
+# Using a custom action since the default action does not
+# take poetry install groups into account.
+# Action code from:
+# https://github.com/actions/setup-python/issues/505#issuecomment-1273013236
+name: poetry-install-with-caching
+description: Poetry install with support for caching of dependency groups.
+
+inputs:
+  python-version:
+    description: Python version, supporting MAJOR.MINOR only
+    required: true
+
+  poetry-version:
+    description: Poetry version
+    required: true
+
+  cache-key:
+    description: Cache key to use for manual handling of caching
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/setup-python@v5
+      name: Setup python ${{ inputs.python-version }}
+      id: setup-python
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - uses: actions/cache@v3
+      id: cache-bin-poetry
+      name: Cache Poetry binary - Python ${{ inputs.python-version }}
+      env:
+        SEGMENT_DOWNLOAD_TIMEOUT_MIN: "1"
+      with:
+        path: |
+          /opt/pipx/venvs/poetry
+        # This step caches the poetry installation, so make sure it's keyed on the poetry version as well.
+        key: bin-poetry-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-${{ inputs.poetry-version }}
+
+    - name: Refresh shell hashtable and fixup softlinks
+      if: steps.cache-bin-poetry.outputs.cache-hit == 'true'
+      shell: bash
+      env:
+        POETRY_VERSION: ${{ inputs.poetry-version }}
+        PYTHON_VERSION: ${{ inputs.python-version }}
+      run: |
+        set -eux
+
+        # Refresh the shell hashtable, to ensure correct `which` output.
+        hash -r
+
+        # `actions/cache@v3` doesn't always seem able to correctly unpack softlinks.
+        # Delete and recreate the softlinks pipx expects to have.
+        rm /opt/pipx/venvs/poetry/bin/python
+        cd /opt/pipx/venvs/poetry/bin
+        ln -s "$(which "python$PYTHON_VERSION")" python
+        chmod +x python
+        cd /opt/pipx_bin/
+        ln -s /opt/pipx/venvs/poetry/bin/poetry poetry
+        chmod +x poetry
+
+        # Ensure everything got set up correctly.
+        /opt/pipx/venvs/poetry/bin/python --version
+        /opt/pipx_bin/poetry --version
+
+    - name: Install poetry
+      if: steps.cache-bin-poetry.outputs.cache-hit != 'true'
+      shell: bash
+      env:
+        POETRY_VERSION: ${{ inputs.poetry-version }}
+        PYTHON_VERSION: ${{ inputs.python-version }}
+      # Install poetry using the python version installed by setup-python step.
+      run: pipx install "poetry==$POETRY_VERSION" --python '${{ steps.setup-python.outputs.python-path }}' --verbose
+
+    - name: Restore pip and poetry cached dependencies
+      uses: actions/cache@v3
+      env:
+        SEGMENT_DOWNLOAD_TIMEOUT_MIN: "4"
+      with:
+        path: |
+          ~/.cache/pip
+          ~/.cache/pypoetry/virtualenvs
+          ~/.cache/pypoetry/cache
+          ~/.cache/pypoetry/artifacts
+          ./.venv
+        key: py-deps-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-poetry-${{ inputs.poetry-version }}-${{ inputs.cache-key }}-${{ hashFiles('./poetry.lock') }}

.github/workflows/py-baseline.yml

@@ -0,0 +1,37 @@
+name: py-baseline
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [main]
+    paths:
+      - "python/langsmith/**"
+
+env:
+  POETRY_VERSION: "1.7.1"
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: python
+    steps:
+      - uses: actions/checkout@v4
+      - run: SHA=$(git rev-parse HEAD) && echo "SHA=$SHA" >> $GITHUB_ENV
+      - name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
+        with:
+          python-version: "3.11"
+          poetry-version: ${{ env.POETRY_VERSION }}
+          cache-key: py-benchi
+      - name: Install dependencies
+        run: poetry install --with dev
+      - name: Run benchmarks
+        run: OUTPUT=out/benchmark-baseline.json make -s benchmark
+      - name: Save outputs
+        uses: actions/cache/save@v4
+        with:
+          key: ${{ runner.os }}-benchmark-baseline-${{ env.SHA }}
+          path: |
+            python/out/benchmark-baseline.json

.github/workflows/py-bench.yml

@@ -0,0 +1,71 @@
+name: py-bench
+
+on:
+  pull_request:
+    paths:
+      - "python/langsmith/**"
+
+env:
+  POETRY_VERSION: "1.7.1"
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: python
+    steps:
+      - uses: actions/checkout@v4
+      - id: files
+        name: Get changed files
+        uses: Ana06/[email protected]
+        with:
+          format: json
+      - name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
+        with:
+          python-version: "3.11"
+          poetry-version: ${{ env.POETRY_VERSION }}
+          cache-key: py-bench
+      - name: Install dependencies
+        run: poetry install --with dev
+      - name: Download baseline
+        uses: actions/cache/restore@v4
+        with:
+          key: ${{ runner.os }}-benchmark-baseline
+          restore-keys: |
+            ${{ runner.os }}-benchmark-baseline-
+          fail-on-cache-miss: true
+          path: |
+            python/out/benchmark-baseline.json
+      - name: Run benchmarks
+        id: benchmark
+        run: |
+          {
+            echo 'OUTPUT<<EOF'
+            make -s benchmark
+            echo EOF
+          } >> "$GITHUB_OUTPUT"
+      - name: Compare benchmarks
+        id: compare
+        run: |
+          {
+            echo 'OUTPUT<<EOF'
+            mv out/benchmark-baseline.json out/main.json
+            mv out/benchmark.json out/changes.json
+            poetry run pyperf compare_to out/main.json out/changes.json --table --group-by-speed
+            echo EOF
+          } >> "$GITHUB_OUTPUT"
+      - name: Annotation
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const file = JSON.parse(`${{ steps.files.outputs.added_modified_renamed }}`)[0]
+            core.notice(`${{ steps.benchmark.outputs.OUTPUT }}`, {
+              title: 'Benchmark results',
+              file,
+            })
+            core.notice(`${{ steps.compare.outputs.OUTPUT }}`, {
+              title: 'Comparison against main',
+              file,
+            })

.github/workflows/python_test.yml

@@ -29,18 +29,12 @@ jobs:
         working-directory: python
     steps:
       - uses: actions/checkout@v3
-      - uses: actions/setup-python@v5
-        name: Setup python ${{ matrix.python-version }}
-        id: setup-python
+      - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
         with:
           python-version: ${{ matrix.python-version }}
-          cache: "pip"
-      - name: Install poetry
-        shell: bash
-        env:
-          PYTHON_VERSION: ${{ matrix.python-version }}
-        # Install poetry using the python version installed by setup-python step.
-        run: pipx install "poetry==$POETRY_VERSION" --python '${{ steps.setup-python.outputs.python-path }}' --verbose
+          poetry-version: ${{ env.POETRY_VERSION }}
+          cache-key: build-and-test
       - name: Install dependencies
         run: |
           poetry install --with dev,lint
@@ -51,4 +45,4 @@ jobs:
         run: make lint
       - name: Run Unit tests ${{ matrix.python-version }}
         run: make tests
-        shell: bash
+        shell: bash

python/Makefile

@@ -1,4 +1,17 @@
-.PHONY: tests lint format build publish doctest integration_tests integration_tests_fast evals
+.PHONY: tests lint format build publish doctest integration_tests integration_tests_fast evals benchmark benchmark-fast
+
+
+OUTPUT ?= out/benchmark.json
+
+benchmark:
+	mkdir -p out
+	rm -f $(OUTPUT)
+	poetry run python -m bench -o $(OUTPUT) --rigorous
+
+benchmark-fast:
+	mkdir -p out
+	rm -f $(OUTPUT)
+	poetry run python -m bench -o $(OUTPUT) --fast
 
 tests:
 	 PYTHONDEVMODE=1 PYTHONASYNCIODEBUG=1 poetry run python -m pytest --disable-socket --allow-unix-socket -n auto --durations=10 tests/unit_tests

python/bench/__main__.py

@@ -0,0 +1,76 @@
+from pyperf._runner import Runner
+
+from bench.create_run_tree import create_run_trees
+from bench.dumps_json import (
+    DeeplyNestedModel,
+    DeeplyNestedModelV1,
+    create_nested_instance,
+)
+from langsmith.client import _dumps_json
+
+
+class MyClass:
+    def __init__(self):
+        self.vals = {}
+
+
+benchmarks = (
+    (
+        "create_5_000_run_trees",
+        create_run_trees,
+        5_000,
+    ),
+    (
+        "create_10_000_run_trees",
+        create_run_trees,
+        10_000,
+    ),
+    (
+        "create_20_000_run_trees",
+        create_run_trees,
+        10_000,
+    ),
+    (
+        "dumps_class_nested_py_branch_and_leaf_200x150",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(
+            200, 150, branch_constructor=MyClass, leaf_constructor=MyClass
+        ),
+    ),
+    (
+        "dumps_class_nested_py_leaf_50x100",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(50, 100, leaf_constructor=MyClass),
+    ),
+    (
+        "dumps_class_nested_py_leaf_200x400",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(200, 400, leaf_constructor=MyClass),
+    ),
+    (
+        "dumps_dataclass_nested_200x150",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(200, 150),
+    ),
+    (
+        "dumps_pydantic_nested_200x400",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(200, 400, branch_constructor=DeeplyNestedModel),
+    ),
+    (
+        "dumps_pydantic_nested_50x100",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(50, 100, branch_constructor=DeeplyNestedModel),
+    ),
+    (
+        "dumps_pydanticv1_nested_200x150",
+        lambda x: _dumps_json({"input": x}),
+        create_nested_instance(200, 150, branch_constructor=DeeplyNestedModelV1),
+    ),
+)
+
+
+r = Runner()
+
+for name, fn, input_ in benchmarks:
+    r.bench_func(name, fn, input_)

python/bench/create_run_tree.py

@@ -0,0 +1,9 @@
+from unittest.mock import patch
+
+from langsmith import RunTree
+
+
+def create_run_trees(N: int):
+    with patch("langsmith.client.requests.Session", autospec=True):
+        for i in range(N):
+            RunTree(name=str(i)).post()