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))))