Codeflash trace decorator

.github/workflows/codeflash-optimize.yaml

@@ -68,4 +68,4 @@ jobs:
         id: optimize_code
         run: |
           source .venv/bin/activate
-          poetry run codeflash
+          poetry run codeflash --benchmark

.github/workflows/end-to-end-test-benchmark-bubblesort.yaml

@@ -0,0 +1,41 @@
+name: end-to-end-test
+
+on:
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  benchmark-bubble-sort-optimization:
+    runs-on: ubuntu-latest
+    env:
+      CODEFLASH_AIS_SERVER: prod
+      POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }}
+      CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
+      COLUMNS: 110
+      MAX_RETRIES: 3
+      RETRY_DELAY: 5
+      EXPECTED_IMPROVEMENT_PCT: 5
+      CODEFLASH_END_TO_END: 1
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set up Python 3.11 for CLI
+        uses: astral-sh/setup-uv@v5
+        with:
+          python-version: 3.11.6
+
+      - name: Install dependencies (CLI)
+        run: |
+          uv tool install poetry
+          uv venv
+          source .venv/bin/activate
+          poetry install --with dev
+
+      - name: Run Codeflash to optimize code
+        id: optimize_code_with_benchmarks
+        run: |
+          source .venv/bin/activate
+          poetry run python tests/scripts/end_to_end_test_benchmark_sort.py

.github/workflows/unit-tests.yaml

@@ -32,7 +32,7 @@ jobs:
         run: uvx poetry install --with dev
 
       - name: Unit tests
-        run: uvx poetry run pytest tests/ --cov --cov-report=xml
+        run: uvx poetry run pytest tests/ --cov --cov-report=xml --benchmark-skip -m "not ci_skip"
 
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v5

code_to_optimize/bubble_sort.py

@@ -7,4 +7,4 @@ def sorter(arr):
                 arr[j] = arr[j + 1]
                 arr[j + 1] = temp
     print(f"result: {arr}")
-    return arr
+    return arr

code_to_optimize/bubble_sort_codeflash_trace.py

@@ -0,0 +1,64 @@
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+@codeflash_trace
+def sorter(arr):
+    for i in range(len(arr)):
+        for j in range(len(arr) - 1):
+            if arr[j] > arr[j + 1]:
+                temp = arr[j]
+                arr[j] = arr[j + 1]
+                arr[j + 1] = temp
+    return arr
+
+@codeflash_trace
+def recursive_bubble_sort(arr, n=None):
+    # Initialize n if not provided
+    if n is None:
+        n = len(arr)
+
+    # Base case: if n is 1, the array is already sorted
+    if n == 1:
+        return arr
+
+    # One pass of bubble sort - move the largest element to the end
+    for i in range(n - 1):
+        if arr[i] > arr[i + 1]:
+            arr[i], arr[i + 1] = arr[i + 1], arr[i]
+
+    # Recursively sort the remaining n-1 elements
+    return recursive_bubble_sort(arr, n - 1)
+
+class Sorter:
+    @codeflash_trace
+    def __init__(self, arr):
+        self.arr = arr
+    @codeflash_trace
+    def sorter(self, multiplier):
+        for i in range(len(self.arr)):
+            for j in range(len(self.arr) - 1):
+                if self.arr[j] > self.arr[j + 1]:
+                    temp = self.arr[j]
+                    self.arr[j] = self.arr[j + 1]
+                    self.arr[j + 1] = temp
+        return self.arr * multiplier
+
+    @staticmethod
+    @codeflash_trace
+    def sort_static(arr):
+        for i in range(len(arr)):
+            for j in range(len(arr) - 1):
+                if arr[j] > arr[j + 1]:
+                    temp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = temp
+        return arr
+
+    @classmethod
+    @codeflash_trace
+    def sort_class(cls, arr):
+        for i in range(len(arr)):
+            for j in range(len(arr) - 1):
+                if arr[j] > arr[j + 1]:
+                    temp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = temp
+        return arr

code_to_optimize/bubble_sort_multithread.py

@@ -0,0 +1,23 @@
+# from code_to_optimize.bubble_sort_codeflash_trace import sorter
+from code_to_optimize.bubble_sort_codeflash_trace import sorter
+import concurrent.futures
+
+
+def multithreaded_sorter(unsorted_lists: list[list[int]]) -> list[list[int]]:
+    # Create a list to store results in the correct order
+    sorted_lists = [None] * len(unsorted_lists)
+
+    # Use ThreadPoolExecutor to manage threads
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        # Submit all sorting tasks and map them to their original indices
+        future_to_index = {
+            executor.submit(sorter, unsorted_list): i
+            for i, unsorted_list in enumerate(unsorted_lists)
+        }
+
+        # Collect results as they complete
+        for future in concurrent.futures.as_completed(future_to_index):
+            index = future_to_index[future]
+            sorted_lists[index] = future.result()
+
+    return sorted_lists

code_to_optimize/bubble_sort_picklepatch_test_unused_socket.py

@@ -0,0 +1,18 @@
+
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+
+
+@codeflash_trace
+def bubble_sort_with_unused_socket(data_container):
+    # Extract the list to sort, leaving the socket untouched
+    numbers = data_container.get('numbers', []).copy()
+
+    return sorted(numbers)
+
+@codeflash_trace
+def bubble_sort_with_used_socket(data_container):
+    # Extract the list to sort, leaving the socket untouched
+    numbers = data_container.get('numbers', []).copy()
+    socket = data_container.get('socket')
+    socket.send("Hello from the optimized function!")
+    return sorted(numbers)

code_to_optimize/bubble_sort_picklepatch_test_used_socket.py

@@ -0,0 +1,46 @@
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+
+@codeflash_trace
+def bubble_sort_with_used_socket(data_container):
+    """
+    Performs a bubble sort on a list within the data_container. The data container has the following schema:
+    - 'numbers' (list): The list to be sorted.
+    - 'socket' (socket): A socket
+
+    Args:
+        data_container: A dictionary with at least 'numbers' (list) and 'socket' keys
+
+    Returns:
+        list: The sorted list of numbers
+    """
+    # Extract the list to sort and socket
+    numbers = data_container.get('numbers', []).copy()
+    socket = data_container.get('socket')
+
+    # Track swap count
+    swap_count = 0
+
+    # Classic bubble sort implementation
+    n = len(numbers)
+    for i in range(n):
+        # Flag to optimize by detecting if no swaps occurred
+        swapped = False
+
+        # Last i elements are already in place
+        for j in range(0, n - i - 1):
+            # Swap if the element is greater than the next element
+            if numbers[j] > numbers[j + 1]:
+                # Perform the swap
+                numbers[j], numbers[j + 1] = numbers[j + 1], numbers[j]
+                swapped = True
+                swap_count += 1
+
+        # If no swapping occurred in this pass, the list is sorted
+        if not swapped:
+            break
+
+    # Send final summary
+    summary = f"Bubble sort completed with {swap_count} swaps"
+    socket.send(summary.encode())
+
+    return numbers

code_to_optimize/process_and_bubble_sort.py

@@ -0,0 +1,28 @@
+from code_to_optimize.bubble_sort import sorter
+
+
+def calculate_pairwise_products(arr):
+    """
+    Calculate the average of all pairwise products in the array.
+    """
+    sum_of_products = 0
+    count = 0
+
+    for i in range(len(arr)):
+        for j in range(len(arr)):
+            if i != j:
+                sum_of_products += arr[i] * arr[j]
+                count += 1
+
+    # The average of all pairwise products
+    return sum_of_products / count if count > 0 else 0
+
+
+def compute_and_sort(arr):
+    # Compute pairwise sums average
+    pairwise_average = calculate_pairwise_products(arr)
+
+    # Call sorter function
+    sorter(arr.copy())
+
+    return pairwise_average

code_to_optimize/process_and_bubble_sort_codeflash_trace.py

@@ -0,0 +1,28 @@
+from code_to_optimize.bubble_sort import sorter
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+
+def calculate_pairwise_products(arr):
+    """
+    Calculate the average of all pairwise products in the array.
+    """
+    sum_of_products = 0
+    count = 0
+
+    for i in range(len(arr)):
+        for j in range(len(arr)):
+            if i != j:
+                sum_of_products += arr[i] * arr[j]
+                count += 1
+
+    # The average of all pairwise products
+    return sum_of_products / count if count > 0 else 0
+
+@codeflash_trace
+def compute_and_sort(arr):
+    # Compute pairwise sums average
+    pairwise_average = calculate_pairwise_products(arr)
+
+    # Call sorter function
+    sorter(arr.copy())
+
+    return pairwise_average

code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py

@@ -0,0 +1,13 @@
+import pytest
+
+from code_to_optimize.bubble_sort import sorter
+
+
+def test_sort(benchmark):
+    result = benchmark(sorter, list(reversed(range(500))))
+    assert result == list(range(500))
+
+# This should not be picked up as a benchmark test
+def test_sort2():
+    result = sorter(list(reversed(range(500))))
+    assert result == list(range(500))

code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py

@@ -0,0 +1,8 @@
+from code_to_optimize.process_and_bubble_sort import compute_and_sort
+from code_to_optimize.bubble_sort import sorter
+def test_compute_and_sort(benchmark):
+    result = benchmark(compute_and_sort, list(reversed(range(500))))
+    assert result == 62208.5
+
+def test_no_func(benchmark):
+    benchmark(sorter, list(reversed(range(500))))

code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py

@@ -0,0 +1,4 @@
+from code_to_optimize.bubble_sort_multithread import multithreaded_sorter
+
+def test_benchmark_sort(benchmark):
+    benchmark(multithreaded_sorter, [list(range(1000)) for i in range (10)])

code_to_optimize/tests/pytest/benchmarks_socket_test/test_socket.py

@@ -0,0 +1,20 @@
+import socket
+
+from code_to_optimize.bubble_sort_picklepatch_test_unused_socket import bubble_sort_with_unused_socket
+from code_to_optimize.bubble_sort_picklepatch_test_used_socket import bubble_sort_with_used_socket
+
+def test_socket_picklepatch(benchmark):
+    s1, s2 = socket.socketpair()
+    data = {
+        "numbers": list(reversed(range(500))),
+        "socket": s1
+    }
+    benchmark(bubble_sort_with_unused_socket, data)
+
+def test_used_socket_picklepatch(benchmark):
+    s1, s2 = socket.socketpair()
+    data = {
+        "numbers": list(reversed(range(500))),
+        "socket": s1
+    }
+    benchmark(bubble_sort_with_used_socket, data)

code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort_example.py

@@ -0,0 +1,26 @@
+import pytest
+
+from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter
+
+
+def test_sort(benchmark):
+    result = benchmark(sorter, list(reversed(range(500))))
+    assert result == list(range(500))
+
+# This should not be picked up as a benchmark test
+def test_sort2():
+    result = sorter(list(reversed(range(500))))
+    assert result == list(range(500))
+
+def test_class_sort(benchmark):
+    obj = Sorter(list(reversed(range(100))))
+    result1 = benchmark(obj.sorter, 2)
+
+def test_class_sort2(benchmark):
+    result2 = benchmark(Sorter.sort_class, list(reversed(range(100))))
+
+def test_class_sort3(benchmark):
+    result3 = benchmark(Sorter.sort_static, list(reversed(range(100))))
+
+def test_class_sort4(benchmark):
+    result4 = benchmark(Sorter, [1,2,3])

code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort_example.py

@@ -0,0 +1,8 @@
+from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort
+from code_to_optimize.bubble_sort_codeflash_trace import sorter
+def test_compute_and_sort(benchmark):
+    result = benchmark(compute_and_sort, list(reversed(range(500))))
+    assert result == 62208.5
+
+def test_no_func(benchmark):
+    benchmark(sorter, list(reversed(range(500))))

code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py

@@ -0,0 +1,6 @@
+from code_to_optimize.bubble_sort_codeflash_trace import recursive_bubble_sort
+
+
+def test_recursive_sort(benchmark):
+    result = benchmark(recursive_bubble_sort, list(reversed(range(500))))
+    assert result == list(range(500))

code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py

@@ -0,0 +1,11 @@
+import pytest
+from code_to_optimize.bubble_sort_codeflash_trace import sorter
+
+def test_benchmark_sort(benchmark):
+    @benchmark
+    def do_sort():
+        sorter(list(reversed(range(500))))
+
+@pytest.mark.benchmark(group="benchmark_decorator")
+def test_pytest_mark(benchmark):
+    benchmark(sorter, list(reversed(range(500))))