⚡️ Speed up method DefaultExecutor.execute_cell_async by 15%

[codeflash-ai]

📄 15% (0.15x) speedup for DefaultExecutor.execute_cell_async in marimo/_runtime/executor.py

⏱️ Runtime : 431 microseconds → 375 microseconds (best of 44 runs)

📝 Explanation and details

The optimized code achieves a 15% runtime improvement through two key micro-optimizations that reduce overhead in frequently called functions:

1. Cached Module Constant in _is_coroutine:

• Hoisted inspect.CO_COROUTINE to module-level constant _CO_COROUTINE
• Eliminates repeated attribute lookups on the inspect module (from 414.6μs to 326.9μs per hit)
• This function is called 1,313 times in profiling, making the optimization compound significantly

2. Optimized Regex Pattern in _raise_name_error:

• Pre-compiled regex pattern _NAME_ERROR_REGEX at module level instead of recompiling on each call
• Replaced re.findall() with faster re.search().group(1) since only one match is needed
• Reduced regex processing time from 164.4μs to 8.0μs (95% improvement) in the error handling path

The optimizations target hot paths in async cell execution where _is_coroutine is called multiple times per cell (for both body and last_expr evaluation). While the regex optimization affects error cases less frequently, it provides dramatic improvements when NameErrors occur.

Test Case Performance:

• Basic execution tests benefit most from the _is_coroutine optimization
• Concurrent execution tests (50-100 cells) see amplified benefits due to repeated function calls
• Error handling tests benefit from the regex optimization when NameErrors are raised

The throughput remains unchanged at 23,540 ops/sec as these are micro-optimizations that reduce per-operation overhead rather than changing the fundamental execution model.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 732 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import asyncio  # used to run async functions
import types

import pytest  # used for our unit tests
from marimo._runtime.executor import DefaultExecutor


# Minimal stubs for required classes and functions
class MarimoRuntimeException(Exception):
    pass

class MarimoMissingRefError(Exception):
    def __init__(self, missing_name, orig_exc):
        super().__init__(f"Missing ref: {missing_name}")
        self.missing_name = missing_name
        self.orig_exc = orig_exc

class DirectedGraph:
    def __init__(self):
        self.definitions = set()

class CellImpl:
    def __init__(self, body, last_expr):
        self.body = body
        self.last_expr = last_expr

class Executor:
    pass
from marimo._runtime.executor import DefaultExecutor


# Helper to create code objects for exec/eval
def make_code(source, mode="exec"):
    # mode: "exec" for exec, "eval" for eval
    return compile(source, "<cell>", mode)

def make_async_code(source, mode="exec"):
    # Returns a coroutine code object
    # For exec: async def _(): <source>; for eval: async def _(): return <source>
    if mode == "exec":
        src = f"async def _():\n    {source.replace(chr(10), chr(10)+'    ')}"
        ns = {}
        exec(src, ns)
        return ns["_"].__code__
    elif mode == "eval":
        src = f"async def _():\n    return {source}"
        ns = {}
        exec(src, ns)
        return ns["_"].__code__
    else:
        raise ValueError("mode must be 'exec' or 'eval'")

# ========== BASIC TEST CASES ==========

@pytest.mark.asyncio
async def test_execute_cell_async_basic_return_value():
    """Test basic synchronous code execution and return value."""
    cell = CellImpl(
        body=make_code("x = 42", "exec"),
        last_expr=make_code("x", "eval")
    )
    glbls = {}
    result = await DefaultExecutor().execute_cell_async(cell, glbls)

@pytest.mark.asyncio

async def test_execute_cell_async_basic_none_body():
    """Test case where cell.body is None, should return None."""
    cell = CellImpl(
        body=None,
        last_expr=make_code("1", "eval")
    )
    glbls = {}
    result = await DefaultExecutor().execute_cell_async(cell, glbls)

@pytest.mark.asyncio

async def test_execute_cell_async_basic_last_expr_async_after_sync_body():
    """Test sync body, async last_expr."""
    cell = CellImpl(
        body=make_code("w = 13", "exec"),
        last_expr=make_async_code("w", "eval")
    )
    glbls = {}
    result = await DefaultExecutor().execute_cell_async(cell, glbls)

# ========== EDGE TEST CASES ==========

@pytest.mark.asyncio





async def test_execute_cell_async_concurrent_execution():
    """Test concurrent execution of multiple cells."""
    cell1 = CellImpl(
        body=make_code("x = 1", "exec"),
        last_expr=make_code("x", "eval")
    )
    cell2 = CellImpl(
        body=make_code("y = 2", "exec"),
        last_expr=make_code("y", "eval")
    )
    glbls1 = {}
    glbls2 = {}
    results = await asyncio.gather(
        DefaultExecutor().execute_cell_async(cell1, glbls1),
        DefaultExecutor().execute_cell_async(cell2, glbls2)
    )

@pytest.mark.asyncio

async def test_execute_cell_async_large_scale_concurrent():
    """Test concurrent execution of many cells (scalable, bounded <100)."""
    N = 50
    cells = [
        CellImpl(
            body=make_code(f"x{i} = {i}", "exec"),
            last_expr=make_code(f"x{i}", "eval")
        )
        for i in range(N)
    ]
    glbls_list = [{} for _ in range(N)]
    tasks = [
        DefaultExecutor().execute_cell_async(cell, glbls)
        for cell, glbls in zip(cells, glbls_list)
    ]
    results = await asyncio.gather(*tasks)

@pytest.mark.asyncio




#------------------------------------------------
import asyncio  # used to run async functions
from types import CodeType, SimpleNamespace
from typing import Any, Optional

import pytest  # used for our unit tests
from marimo._runtime.executor import DefaultExecutor

# ---- Minimal stubs to support the function's dependencies ----

class MarimoRuntimeException(Exception):
    pass

class MarimoMissingRefError(Exception):
    def __init__(self, missing_name, orig_exc):
        super().__init__(f"Missing ref: {missing_name}")
        self.missing_name = missing_name
        self.orig_exc = orig_exc

class DirectedGraph:
    def __init__(self, definitions=None):
        self.definitions = definitions or set()

class CellImpl:
    def __init__(self, body: Optional[CodeType], last_expr: Optional[CodeType]):
        self.body = body
        self.last_expr = last_expr

# ---- Function under test ----

class Executor:
    pass
from marimo._runtime.executor import DefaultExecutor

# ---- Helper functions for test code generation ----

def make_code_object(source: str, is_expr: bool = False) -> CodeType:
    """Create a code object from source string."""
    mode = "eval" if is_expr else "exec"
    return compile(source, "<cell>", mode)

# ---- Test suite ----

@pytest.mark.asyncio
async def test_execute_cell_async_basic_exec_and_eval():
    # Test basic synchronous exec and eval
    executor = DefaultExecutor()
    # cell.body: exec code, cell.last_expr: eval code
    cell = CellImpl(
        body=make_code_object("x = 5"),
        last_expr=make_code_object("x + 2", is_expr=True)
    )
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_basic_async_body_and_expr():
    # Test async body and async last_expr
    executor = DefaultExecutor()
    # cell.body: async def, cell.last_expr: async def
    # We'll use async code objects for both
    # The body sets y=10, last_expr returns y+5
    body_code = compile(
        "async def _f():\n    globals()['y'] = 10\n    return None\n_f()",
        "<cell>", "exec"
    )
    ns = {}
    exec(body_code, ns)
    body = ns['_f'].__code__

    last_expr_code = compile(
        "async def _g():\n    return y + 5\n_g()",
        "<cell>", "exec"
    )
    exec(last_expr_code, ns)
    last_expr = ns['_g'].__code__

    cell = CellImpl(body=body, last_expr=last_expr)
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_none_body_returns_none():
    # If cell.body is None, should return None
    executor = DefaultExecutor()
    cell = CellImpl(body=None, last_expr=make_code_object("1", is_expr=True))
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_assert_last_expr_not_none():
    # If cell.last_expr is None, should raise AssertionError
    executor = DefaultExecutor()
    cell = CellImpl(body=make_code_object("x = 1"), last_expr=None)
    glbls = {}
    with pytest.raises(AssertionError):
        await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio



async def test_execute_cell_async_concurrent_execution():
    # Test concurrent execution of multiple cells
    executor = DefaultExecutor()
    cell1 = CellImpl(
        body=make_code_object("x = 1"),
        last_expr=make_code_object("x + 1", is_expr=True)
    )
    cell2 = CellImpl(
        body=make_code_object("y = 2"),
        last_expr=make_code_object("y * 2", is_expr=True)
    )
    glbls1 = {}
    glbls2 = {}
    results = await asyncio.gather(
        executor.execute_cell_async(cell1, glbls1),
        executor.execute_cell_async(cell2, glbls2)
    )

@pytest.mark.asyncio
async def test_execute_cell_async_async_last_expr_returns_value():
    # last_expr is async code object, should await and return its value
    executor = DefaultExecutor()
    body_code = make_code_object("x = 10")
    last_expr_code = compile(
        "async def _f():\n    return x * 2\n_f()",
        "<cell>", "exec"
    )
    ns = {"x": 10}
    exec(last_expr_code, ns)
    last_expr = ns["_f"].__code__
    cell = CellImpl(body=body_code, last_expr=last_expr)
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_large_scale_concurrent():
    # Test large scale concurrent execution (up to 100 cells)
    executor = DefaultExecutor()
    N = 100
    cells = [
        CellImpl(
            body=make_code_object(f"x = {i}"),
            last_expr=make_code_object("x + 1", is_expr=True)
        )
        for i in range(N)
    ]
    glbls_list = [{} for _ in range(N)]
    results = await asyncio.gather(
        *[executor.execute_cell_async(cell, glbls)
          for cell, glbls in zip(cells, glbls_list)]
    )
    for i, glbls in enumerate(glbls_list):
        pass

@pytest.mark.asyncio
async def test_execute_cell_async_large_scale_async_expr():
    # Test large scale concurrent execution with async last_expr
    executor = DefaultExecutor()
    N = 50
    cells = []
    glbls_list = []
    for i in range(N):
        body_code = make_code_object(f"x = {i}")
        last_expr_code = compile(
            f"async def _f():\n    return x * 2\n_f()",
            "<cell>", "exec"
        )
        ns = {"x": i}
        exec(last_expr_code, ns)
        last_expr = ns["_f"].__code__
        cells.append(CellImpl(body=body_code, last_expr=last_expr))
        glbls_list.append({})
    results = await asyncio.gather(
        *[executor.execute_cell_async(cell, glbls)
          for cell, glbls in zip(cells, glbls_list)]
    )
    for i, glbls in enumerate(glbls_list):
        pass

@pytest.mark.asyncio



async def test_execute_cell_async_throughput_async_mix():
    # Throughput test: mix of async and sync last_expr
    executor = DefaultExecutor()
    N = 40
    cells = []
    glbls_list = []
    for i in range(N):
        body_code = make_code_object(f"x = {i}")
        if i % 2 == 0:
            last_expr_code = make_code_object("x * 2", is_expr=True)
            last_expr = last_expr_code
        else:
            last_expr_code = compile(
                f"async def _f():\n    return x * 3\n_f()",
                "<cell>", "exec"
            )
            ns = {"x": i}
            exec(last_expr_code, ns)
            last_expr = ns["_f"].__code__
        cells.append(CellImpl(body=body_code, last_expr=last_expr))
        glbls_list.append({})
    results = await asyncio.gather(
        *[executor.execute_cell_async(cell, glbls)
          for cell, glbls in zip(cells, glbls_list)]
    )
    # Even indices: x*2, odd: x*3
    expected = [i*2 if i % 2 == 0 else i*3 for i in range(N)]

@pytest.mark.asyncio
async def test_execute_cell_async_edge_empty_globals():
    # Edge case: empty globals dict, code uses builtins only
    executor = DefaultExecutor()
    cell = CellImpl(
        body=make_code_object(""),
        last_expr=make_code_object("len([1,2,3])", is_expr=True)
    )
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_edge_globals_prepopulated():
    # Edge case: globals dict prepopulated with variables
    executor = DefaultExecutor()
    cell = CellImpl(
        body=make_code_object("x = y + 1"),
        last_expr=make_code_object("x", is_expr=True)
    )
    glbls = {"y": 10}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio
async def test_execute_cell_async_edge_body_is_pass():
    # Edge case: body is 'pass', last_expr is a constant
    executor = DefaultExecutor()
    cell = CellImpl(
        body=make_code_object("pass"),
        last_expr=make_code_object("42", is_expr=True)
    )
    glbls = {}
    result = await executor.execute_cell_async(cell, glbls)

@pytest.mark.asyncio

To edit these changes git checkout codeflash/optimize-DefaultExecutor.execute_cell_async-mh6a03ps and push.