diff --git a/packages/databricks-vscode/resources/python/00-databricks-init.py b/packages/databricks-vscode/resources/python/00-databricks-init.py
index 0883fc598..b5d6fb261 100644
--- a/packages/databricks-vscode/resources/python/00-databricks-init.py
+++ b/packages/databricks-vscode/resources/python/00-databricks-init.py
@@ -3,11 +3,13 @@
 import json
 from typing import Any, Union, List
 import os
+import sys
 import time
 import shlex
 import warnings
 import tempfile
 
+builtinSum = sys.modules['builtins'].sum
 
 def logError(function_name: str, e: Union[str, Exception]):
     import sys
@@ -403,18 +405,24 @@ def init_ui(self):
         def update_ticks(
             self,
             stages,
-            inflight_tasks: int
+            inflight_tasks: int,
+            done: bool
         ) -> None:
-            total_tasks = sum(map(lambda x: x.num_tasks, stages))
-            completed_tasks = sum(map(lambda x: x.num_completed_tasks, stages))
+            total_tasks = builtinSum(map(lambda x: x.num_tasks, stages))
+            completed_tasks = builtinSum(map(lambda x: x.num_completed_tasks, stages))
             if total_tasks > 0:
                 self._ticks = total_tasks
                 self._tick = completed_tasks
-                self._bytes_read = sum(map(lambda x: x.num_bytes_read, stages))
+                self._bytes_read = builtinSum(map(lambda x: x.num_bytes_read, stages))
+
+                if done:
+                    self._tick = self._ticks
+                    self._running = 0
+
                 if self._tick is not None and self._tick >= 0:
                     self.output()            
                 self._running = inflight_tasks
-
+            
         def output(self) -> None:
             if self._tick is not None and self._ticks is not None:
                 percent_complete = (self._tick / self._ticks) * 100
@@ -495,7 +503,7 @@ def __call__(self,
                 self.op_id = operation_id
                 self.reset()
 
-            self.p.update_ticks(stages, inflight_tasks)
+            self.p.update_ticks(stages, inflight_tasks, done)
                     
     spark.clearProgressHandlers()
     spark.registerProgressHandler(ProgressHandler())