Merge branch 'jedi_ctx_fix' into codeflash-trace-decorator

Author: alvin-r

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -1,4 +1,6 @@
 from concurrent.futures import ThreadPoolExecutor
+
+
 def funcA(number):
     k = 0
     for i in range(number * 100):
@@ -9,6 +11,7 @@ def funcA(number):
     # Use a generator expression directly in join for more efficiency
     return " ".join(str(i) for i in range(number))
 
+
 def test_threadpool() -> None:
     pool = ThreadPoolExecutor(max_workers=3)
     args = list(range(10, 31, 10))
@@ -19,4 +22,4 @@ def test_threadpool() -> None:
 
 
 if __name__ == "__main__":
-    test_threadpool()
+    test_threadpool()

codeflash/context/code_context_extractor.py

@@ -356,6 +356,7 @@ def get_function_to_optimize_as_function_source(
             name.type == "function"
             and name.full_name
             and name.name == function_to_optimize.function_name
+            and name.full_name.startswith(name.module_name)
             and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name
         ):
             function_source = FunctionSource(
@@ -410,6 +411,7 @@ def get_function_sources_from_jedi(
                         and definition.full_name
                         and definition.type == "function"
                         and not belongs_to_function_qualified(definition, qualified_function_name)
+                        and definition.full_name.startswith(definition.module_name)
                         # Avoid nested functions or classes. Only class.function is allowed
                         and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2
                     ):

codeflash/optimization/function_context.py

@@ -31,7 +31,7 @@ def belongs_to_class(name: Name, class_name: str) -> bool:
 def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool:
     """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name."""
     try:
-        if get_qualified_name(name.module_name, name.full_name) == qualified_function_name:
+        if name.full_name.startswith(name.module_name) and get_qualified_name(name.module_name, name.full_name) == qualified_function_name:
             # Handles function definition and recursive function calls
             return False
         if name := name.parent():

codeflash/tracer.py

@@ -247,14 +247,18 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:
             return
         if self.timeout is not None and (time.time() - self.start_time) > self.timeout:
             sys.setprofile(None)
+            threading.setprofile(None)
             console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.")
             return
         code = frame.f_code
+
         file_name = Path(code.co_filename).resolve()
         # TODO : It currently doesn't log the last return call from the first function
 
         if code.co_name in self.ignored_functions:
             return
+        if not file_name.is_relative_to(self.project_root):
+            return
         if not file_name.exists():
             return
         if self.functions and code.co_name not in self.functions:

tests/scripts/end_to_end_test_tracer_replay.py

@@ -10,7 +10,7 @@ def run_test(expected_improvement_pct: int) -> bool:
         min_improvement_x=0.1,
         expected_unit_tests=1,
         coverage_expectations=[
-            CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[3, 4, 5, 7, 10]),
+            CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[5, 6, 7, 9, 12]),
         ],
     )
     cwd = (

tests/scripts/end_to_end_test_utilities.py

@@ -204,8 +204,8 @@ def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_p
         return False
 
     functions_traced = re.search(r"Traced (\d+) function calls successfully and replay test created at - (.*)$", stdout)
-    if not functions_traced or int(functions_traced.group(1)) != 5:
-        logging.error("Expected 5 traced functions")
+    if not functions_traced or int(functions_traced.group(1)) != 4:
+        logging.error("Expected 4 traced functions")
         return False
 
     replay_test_path = pathlib.Path(functions_traced.group(2))