working version

Author: misrasaurabh1

codeflash/code_utils/checkpoint.py

@@ -11,8 +11,8 @@
 
 
 class CodeflashRunCheckpoint:
-    def __init__(self, module_path: Path, checkpoint_dir: str = "/tmp") -> None:
-        self.module_path = module_path
+    def __init__(self, module_root: Path, checkpoint_dir: str = "/tmp") -> None:
+        self.module_root = module_root
         self.checkpoint_dir = Path(checkpoint_dir)
         # Create a unique checkpoint file name
         unique_id = str(uuid.uuid4())[:8]
@@ -26,7 +26,7 @@ def _initialize_checkpoint_file(self) -> None:
         """Create a new checkpoint file with metadata."""
         metadata = {
             "type": "metadata",
-            "module_path": str(self.module_path),
+            "module_root": str(self.module_root),
             "created_at": time.time(),
             "last_updated": time.time(),
         }
@@ -82,7 +82,7 @@ def _update_metadata_timestamp(self) -> None:
             f.write(rest_content)
 
 
-def get_all_historical_functions(checkpoint_dir: Path, module_path) -> dict[str, dict[str, str]]:
+def get_all_historical_functions(module_root: Path, checkpoint_dir: Path) -> dict[str, dict[str, str]]:
     """Get information about all processed functions, regardless of status.
 
     Returns:
@@ -97,29 +97,27 @@ def get_all_historical_functions(checkpoint_dir: Path, module_path) -> dict[str,
             # Skip the first line (metadata)
             first_line = next(f)
             metadata = json.loads(first_line)
-            if metadata.get("timestamp"):
-                metadata["timestamp"] = datetime.datetime.fromtimestamp(metadata["timestamp"])
-                if metadata["timestamp"] >= datetime.datetime.now() - datetime.timedelta(days=7):
+            if metadata.get("last_updated"):
+                last_updated = datetime.datetime.fromtimestamp(metadata["last_updated"])
+                if datetime.datetime.now() - last_updated >= datetime.timedelta(days=7):
                     to_delete.append(file)
                     continue
-            else:
-                metadata["timestamp"] = datetime.datetime.now()
-            if metadata.get("module_path") != module_path:
+            if metadata.get("module_root") != str(module_root):
                 continue
 
             for line in f:
                 entry = json.loads(line)
                 if entry.get("type") == "function":
                     processed_functions[entry["function_name"]] = entry
     for file in to_delete:
-        file.unlink()
+        file.unlink(missing_ok=True)
     return processed_functions
 
 
 def ask_should_use_checkpoint_get_functions(args: argparse.Namespace) -> Optional[dict[str, dict[str, str]]]:
     previous_checkpoint_functions = None
     if args.all and (sys.platform == "linux" or sys.platform == "darwin") and Path("/tmp").is_dir():
-        previous_checkpoint_functions = get_all_historical_functions(args.module_path, Path("/tmp"))
+        previous_checkpoint_functions = get_all_historical_functions(args.module_root, Path("/tmp"))
         if previous_checkpoint_functions and click.confirm(
             "Previous Checkpoint detected from an incomplete optimization run, shall I continue the optimization from that point?",
             default=True,

codeflash/discovery/functions_to_optimize.py

@@ -8,7 +8,7 @@
 from collections import defaultdict
 from functools import cache
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import git
 import libcst as cst
@@ -416,7 +416,7 @@ def filter_functions(
     ignore_paths: list[Path],
     project_root: Path,
     module_root: Path,
-    previous_checkpoint_functions: dict[Path, list[FunctionToOptimize]] | None = None,
+    previous_checkpoint_functions: dict[Path, dict[str, Any]] | None = None,
     disable_logs: bool = False,
 ) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
     blocklist_funcs = get_blocklisted_functions()
@@ -480,9 +480,7 @@ def filter_functions(
         if previous_checkpoint_functions:
             functions_tmp = []
             for function in _functions:
-                if function.file_path in previous_checkpoint_functions and function in previous_checkpoint_functions[
-                    function.file_path
-                ]:
+                if function.qualified_name_with_modules_from_root(project_root) in previous_checkpoint_functions:
                     previous_checkpoint_functions_removed_count += 1
                     continue
                 functions_tmp.append(function)
@@ -500,7 +498,7 @@ def filter_functions(
             f"{ignore_paths_removed_count} file{'s' if ignore_paths_removed_count != 1 else ''} from ignored paths": ignore_paths_removed_count,
             f"{submodule_ignored_paths_count} file{'s' if submodule_ignored_paths_count != 1 else ''} from ignored submodules": submodule_ignored_paths_count,
             f"{blocklist_funcs_removed_count} function{'s' if blocklist_funcs_removed_count != 1 else ''} as previously optimized": blocklist_funcs_removed_count,
-            f"{previous_checkpoint_functions_removed_count} function{'s' if previous_checkpoint_functions_removed_count != 1 else ''} as previously optimized from checkpoint": previous_checkpoint_functions_removed_count,
+            f"{previous_checkpoint_functions_removed_count} function{'s' if previous_checkpoint_functions_removed_count != 1 else ''} skipped from checkpoint": previous_checkpoint_functions_removed_count,
         }
         log_string = "\n".join([k for k, v in log_info.items() if v > 0])
         if log_string:

codeflash/optimization/function_optimizer.py

@@ -242,7 +242,9 @@ def optimize_function(self) -> Result[BestOptimization, str]:
         # request for new optimizations but don't block execution, check for completion later
         # adding to control and experiment set but with same traceid
         best_optimization = None
-        for _u, (candidates, exp_type) in enumerate(zip([optimizations_set.control, optimizations_set.experiment],["EXP0","EXP1"])):
+        for _u, (candidates, exp_type) in enumerate(
+            zip([optimizations_set.control, optimizations_set.experiment], ["EXP0", "EXP1"])
+        ):
             if candidates is None:
                 continue
 
@@ -254,7 +256,14 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                 file_path_to_helper_classes=file_path_to_helper_classes,
                 exp_type=exp_type,
             )
-            ph("cli-optimize-function-finished", {"function_trace_id": self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id})
+            ph(
+                "cli-optimize-function-finished",
+                {
+                    "function_trace_id": self.function_trace_id[:-4] + exp_type
+                    if self.experiment_id
+                    else self.function_trace_id
+                },
+            )
 
             generated_tests = remove_functions_from_generated_tests(
                 generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove
@@ -324,7 +333,9 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                         explanation=explanation,
                         existing_tests_source=existing_tests,
                         generated_original_test_source=generated_tests_str,
-                        function_trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
+                        function_trace_id=self.function_trace_id[:-4] + exp_type
+                        if self.experiment_id
+                        else self.function_trace_id,
                         coverage_message=coverage_message,
                         git_remote=self.args.git_remote,
                     )
@@ -379,15 +390,19 @@ def determine_best_candidate(
         # Start a new thread for AI service request, start loop in main thread
         # check if aiservice request is complete, when it is complete, append result to the candidates list
         with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-            ai_service_client = self.aiservice_client if exp_type=="EXP0" else self.local_aiservice_client
+            ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
             future_line_profile_results = executor.submit(
                 ai_service_client.optimize_python_code_line_profiler,
                 source_code=code_context.read_writable_code,
                 dependency_code=code_context.read_only_context_code,
                 trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
                 line_profiler_results=original_code_baseline.line_profile_results["str_out"],
                 num_candidates=10,
-                experiment_metadata=ExperimentMetadata(id=self.experiment_id, group= "control" if exp_type == "EXP0" else "experiment") if self.experiment_id  else None,
+                experiment_metadata=ExperimentMetadata(
+                    id=self.experiment_id, group="control" if exp_type == "EXP0" else "experiment"
+                )
+                if self.experiment_id
+                else None,
             )
             try:
                 candidate_index = 0
@@ -462,7 +477,7 @@ def determine_best_candidate(
                                 f"loop{'s' if candidate_result.max_loop_count > 1 else ''})"
                             )
                             tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%")
-                            tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X")
+                            tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X")
                             replay_perf_gain = {}
                             if self.args.benchmark:
                                 test_results_by_benchmark = (
@@ -528,7 +543,9 @@ def determine_best_candidate(
         )
         return best_optimization
 
-    def log_successful_optimization(self, explanation: Explanation, generated_tests: GeneratedTestsList, exp_type: str) -> None:
+    def log_successful_optimization(
+        self, explanation: Explanation, generated_tests: GeneratedTestsList, exp_type: str
+    ) -> None:
         explanation_panel = Panel(
             f"⚡️ Optimization successful! 📄 {self.function_to_optimize.qualified_name} in {explanation.file_path}\n"
             f"📈 {explanation.perf_improvement_line}\n"
@@ -555,7 +572,9 @@ def log_successful_optimization(self, explanation: Explanation, generated_tests:
         ph(
             "cli-optimize-success",
             {
-                "function_trace_id": self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
+                "function_trace_id": self.function_trace_id[:-4] + exp_type
+                if self.experiment_id
+                else self.function_trace_id,
                 "speedup_x": explanation.speedup_x,
                 "speedup_pct": explanation.speedup_pct,
                 "best_runtime": explanation.best_runtime_ns,

codeflash/optimization/optimizer.py

@@ -163,7 +163,7 @@ def run(self) -> None:
             console.rule()
             ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
             if self.args.all:
-                self.functions_checkpoint = CodeflashRunCheckpoint(self.args.module_path)
+                self.functions_checkpoint = CodeflashRunCheckpoint(self.args.module_root)
 
             for original_module_path in file_to_funcs_to_optimize:
                 logger.info(f"Examining file {original_module_path!s}…")