Fix imports in cache clean up scripts (#54)

Author: agrawalreetika

benchmarks/scripts/cache_cleaning_coordinator.py

@@ -1,14 +1,22 @@
-from presto_utils import execute_cluster_call
+from presto_utils import execute_presto_query
 import argparse
 import sys
 
 def clean_directory_list_cache(hostname, username, password, catalog_name):
     query = "CALL " + catalog_name + ".system.invalidate_directory_list_cache()"
-    return execute_cluster_call(hostname, username, password, catalog_name, query)
+    return execute_presto_query(hostname, username, password, catalog_name, query)
 
 def clean_metastore_cache(hostname, username, password, catalog_name):
     query = "CALL " + catalog_name + ".system.invalidate_metastore_cache()"
-    return execute_cluster_call(hostname, username, password, catalog_name, query)
+    return execute_presto_query(hostname, username, password, catalog_name, query)
+
+def clean_statistics_file_cache(hostname, username, password, catalog_name):
+    query = "CALL " + catalog_name + ".system.invalidate_statistics_file_cache()"
+    return execute_presto_query(hostname, username, password, catalog_name, query)
+
+def clean_manifest_file_cache(hostname, username, password, catalog_name):
+    query = "CALL " + catalog_name + ".system.invalidate_manifest_file_cache()"
+    return execute_presto_query(hostname, username, password, catalog_name, query)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Connect to PrestoDB')
@@ -18,13 +26,18 @@ def clean_metastore_cache(hostname, username, password, catalog_name):
 
     args = parser.parse_args()
 
-    catalog_list = ["hive"]
-    is_list_cache_cleanup_enabled = True
-    is_metadata_cache_cleanup_enabled = True
+    hive_catalog_list = ["hive"]
+    iceberg_catalog_list = ["iceberg"]
+
+    # coordinator cache
+    is_list_cache_cleanup_enabled = True  #for hive
+    is_metadata_cache_cleanup_enabled = True #for hive
+    is_statistics_file_cache_cleanup_enabled = True #for iceberg
+    is_manifest_file_cache_cleanup_enabled = True #for iceberg
 
     # Directory list cache clean up
     if is_list_cache_cleanup_enabled:
-        for catalog_name in catalog_list:
+        for catalog_name in hive_catalog_list:
             print("Cleaning up directory list cache for ", catalog_name)
             rows = clean_directory_list_cache(args.host, args.username, args.password, catalog_name)
             print("directory_list_cache_cleanup_query Query Result: ", rows)
@@ -36,7 +49,7 @@ def clean_metastore_cache(hostname, username, password, catalog_name):
 
     # Metadata cache clean up
     if is_metadata_cache_cleanup_enabled:
-        for catalog_name in catalog_list:
+        for catalog_name in hive_catalog_list:
             print("Cleaning up metadata cache for ", catalog_name)
             rows = clean_metastore_cache(args.host, args.username, args.password, catalog_name)
             print("metastore_cache_cleanup_query Query Result: ", rows)
@@ -45,3 +58,27 @@ def clean_metastore_cache(hostname, username, password, catalog_name):
             else:
                 print("Metastore cache clean up is failed for ", catalog_name)
                 sys.exit(1)
+
+    # Statistics file cache clean up
+    if is_statistics_file_cache_cleanup_enabled:
+        for catalog_name in iceberg_catalog_list:
+            print("Cleaning up statistics file cache for ", catalog_name)
+            rows = clean_statistics_file_cache(args.host, args.username, args.password, catalog_name)
+            print("statistics_file_cache_cleanup_query Query Result: ", rows)
+            if rows[0][0] == True:
+                print("Statistics file cache clean up is successful for ", catalog_name)
+            else:
+                print("Statistics file cache clean up is failed for ", catalog_name)
+                sys.exit(1)
+
+    # Manifest file cache clean up
+    if is_manifest_file_cache_cleanup_enabled:
+        for catalog_name in iceberg_catalog_list:
+            print("Cleaning up manifest file cache for ", catalog_name)
+            rows = clean_manifest_file_cache(args.host, args.username, args.password, catalog_name)
+            print("manifest_file_cache_cleanup_query Query Result: ", rows)
+            if rows[0][0] == True:
+                print("Manifest file cache clean up is successful for ", catalog_name)
+            else:
+                print("Manifest file cache clean up is failed for ", catalog_name)
+                sys.exit(1)

benchmarks/scripts/cache_cleaning_workers.py

@@ -14,21 +14,23 @@ def get_workers_public_ips(data):
     return worker_ips
 
 def cleanup_worker_disk_cache(worker_public_ips, login_user, ssh_key_path, worker_http_port):
-    ssd_cache_clean_commands = ["curl", "-X", "GET", f"http://localhost:{worker_http_port}/v1/operation/server/clearCache?type=ssd"]
+    ssd_cache_clean_command = f"docker exec $(docker ps -q --filter 'name=^presto_workers') curl -sS -X GET http://localhost:{worker_http_port}/v1/operation/server/clearCache?type=ssd"
     for worker_ip in worker_public_ips:
-        execute_ssh_command(worker_ip, login_user, ssh_key_path, ssd_cache_clean_commands)
+        execute_ssh_command(worker_ip, login_user, ssh_key_path, ssd_cache_clean_command)
+    print("cleanup_worker_disk_cache is successful!")
 
 def cleanup_worker_os_cache(worker_public_ips, login_user, ssh_key_path):
     for worker_ip in worker_public_ips:
         os_cache_clean_commands = ["sudo sync && echo 3 | sudo tee /proc/sys/vm/drop_caches", "sudo swapoff -a; sudo swapon -a"]
         for command in os_cache_clean_commands:
             execute_ssh_command(worker_ip, login_user, ssh_key_path, command)
+    print("cleanup_worker_os_cache is successful!")
 
 def cleanup_worker_memory_cache(worker_public_ips, login_user, ssh_key_path, worker_http_port):
+    memory_cache_clean_command = f"docker exec $(docker ps -q --filter 'name=^presto_workers') curl -sS -X GET http://localhost:{worker_http_port}/v1/operation/server/clearCache?type=memory"
     for worker_ip in worker_public_ips:
-        memory_cache_clean_commands = ["curl", "-X", "GET", f"http://localhost:{worker_http_port}/v1/operation/server/clearCache?type=memory"]
-        for command in memory_cache_clean_commands:
-            execute_ssh_command(worker_ip, login_user, ssh_key_path, command)
+        execute_ssh_command(worker_ip, login_user, ssh_key_path, memory_cache_clean_command)
+    print("cleanup_worker_memory_cache is successful!")
 
 # Main function to connect and run queries
 if __name__ == "__main__":