Restart client to free memory

Author: gabegma

azimuth/app.py

@@ -349,7 +349,6 @@ def run_validation_module(pipeline_index=None):
         for pipeline_index in range(len(config.pipelines)):
             run_validation_module(pipeline_index)
     task_manager.clear_worker_cache()
-    task_manager.restart()
 
 
 def run_startup_tasks(azimuth_config: AzimuthConfig, cluster: SpecCluster):
@@ -368,6 +367,7 @@ def run_startup_tasks(azimuth_config: AzimuthConfig, cluster: SpecCluster):
         run_validation(DatasetSplitName.train, task_manager, azimuth_config)
     if _dataset_split_managers.get(DatasetSplitName.eval):
         run_validation(DatasetSplitName.eval, task_manager, azimuth_config)
+    task_manager.restart()
 
     azimuth_config.save()  # Save only after the validation modules ran successfully
 

azimuth/startup.py

@@ -251,12 +251,16 @@ def startup_tasks(
 
     mods = start_tasks_for_dms(config, dataset_split_managers, task_manager, start_up_tasks)
 
-    # Start a thread to monitor the status.
-    th = threading.Thread(
-        target=wait_for_startup, args=(mods, task_manager), name=START_UP_THREAD_NAME
-    )
-    th.setDaemon(True)
-    th.start()
+    startup_ready = all(m.done() for m in mods.values())
+    if startup_ready:
+        log.info("Loading the application from cache. It should be accessible now.")
+    else:
+        # Start a thread to monitor the status.
+        th = threading.Thread(
+            target=wait_for_startup, args=(mods, task_manager), name=START_UP_THREAD_NAME
+        )
+        th.setDaemon(True)
+        th.start()
 
     return mods
 
@@ -346,4 +350,3 @@ def wait_for_startup(startup_mods: Dict[str, DaskModule], task_manager: TaskMana
     task_manager.restart()
     # After restarting, it is safe to unlock the task manager.
     task_manager.unlock()
-    log.info("Cluster restarted to free memory.")

azimuth/task_manager.py

@@ -1,6 +1,7 @@
 # Copyright ServiceNow, Inc. 2021 – 2022
 # This source code is licensed under the Apache 2.0 license found in the LICENSE file
 # in the root directory of this source tree.
+import time
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import structlog
@@ -217,6 +218,8 @@ def clear_worker_cache(self):
         self.client.run(ArtifactManager.clear_cache)
 
     def restart(self):
-        # Clear futures to free memory.
+        log.info("Cluster restarted to free memory.")
         for task_name, module in self.current_tasks.items():
             module.future = None
+        self.client.restart()
+        time.sleep(2)  # Without that, the test routers fail because some tasks can't get scheduled.