Skip to content

Commit 4d58b00

Browse files
cancel tasks with status initializing
1 parent 1361dde commit 4d58b00

File tree

3 files changed

+12
-3
lines changed

3 files changed

+12
-3
lines changed

codeflash/lsp/beta.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,13 @@ def execution_context(**kwargs: str) -> None:
344344
server.execution_context_vars.reset(token)
345345

346346

347+
@server.feature("cleanupCurrentOptimizerSession")
348+
def cleanup_optimizer(_params: any) -> dict[str, str]:
349+
if not server.cleanup_the_optimizer():
350+
return {"status": "error", "message": "Failed to cleanup optimizer"}
351+
return {"status": "success"}
352+
353+
347354
@server.feature("initializeFunctionOptimization")
348355
def initialize_function_optimization(params: FunctionOptimizationInitParams) -> dict[str, str]:
349356
with execution_context(task_id=params.task_id):

codeflash/lsp/server.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,10 +63,10 @@ def show_message_log(self, message: str, message_type: str) -> None:
6363
log_params = LogMessageParams(type=lsp_message_type, message=message)
6464
self.protocol.notify("window/logMessage", log_params)
6565

66-
def cleanup_the_optimizer(self) -> None:
66+
def cleanup_the_optimizer(self) -> bool:
6767
self.current_optimization_init_result = None
6868
if not self.optimizer:
69-
return
69+
return False
7070
try:
7171
self.optimizer.cleanup_temporary_paths()
7272
# restore args and test cfg
@@ -77,6 +77,8 @@ def cleanup_the_optimizer(self) -> None:
7777
self.optimizer.current_function_optimizer = None
7878
except Exception:
7979
self.show_message_log("Failed to cleanup optimizer", "Error")
80+
return False
81+
return True
8082

8183
def shutdown(self) -> None:
8284
"""Gracefully shutdown the server."""

codeflash/optimization/function_optimizer.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1586,7 +1586,7 @@ def establish_original_code_baseline(
15861586
)
15871587

15881588
if test_framework == "pytest":
1589-
with progress_bar("Running line profiling to identify performance bottlenecks..."):
1589+
with progress_bar("Running line profiler to identify performance bottlenecks..."):
15901590
line_profile_results = self.line_profiler_step(
15911591
code_context=code_context, original_helper_code=original_helper_code, candidate_index=0
15921592
)

0 commit comments

Comments
 (0)