resource leaks

Author: Ismail Mehdi

Makefile

@@ -36,11 +36,13 @@ container:
 	$(CONTAINER_ENGINE) run --rm --net=host --name=evalbench_container \
 		-v ~/.config/gcloud:/root/.config/gcloud \
 		-e GOOGLE_CLOUD_PROJECT=cloud-db-nl2sql \
+		--cap-add=SYS_PTRACE	\
 		-e OPTION=--localhost \
 		-e TYPE=$(TYPE) evalbench:latest /evalbench/run_service.sh 
 
 shell:
 	$(CONTAINER_ENGINE) run -ti --rm --net=host --name=evalbench_container \
+		--cap-add=SYS_PTRACE	\
 		-v ~/.config/gcloud:/root/.config/gcloud \
 		-v ~/.gitconfig:/root/.gitconfig \
 		-v ~/.gitcookies:/root/.gitcookies \

evalbench/databases/postgres.py

@@ -34,6 +34,7 @@
 GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO {DML_USERNAME};
 """
 
+CONNECTOR = Connector()
 
 class PGDB(DB):
     #####################################################
@@ -44,10 +45,9 @@ class PGDB(DB):
 
     def __init__(self, db_config):
         super().__init__(db_config)
-        self.connector = Connector()
 
         def get_conn():
-            conn = self.connector.connect(
+            conn = CONNECTOR.connect(
                 self.db_path,
                 "pg8000",
                 user=self.username,
@@ -74,7 +74,7 @@ def get_engine_args():
 
     def close_connections(self):
         try:
-            self.connector.close()
+            self.engine.dispose()
         except Exception:
             logging.warning(
                 f"Failed to close connections. This may result in idle unused connections."

evalbench/databases/util.py

@@ -8,6 +8,7 @@
 import re
 from dataclasses import dataclass, field
 
+CLIENT=secretmanager_v1.SecretManagerServiceClient()
 
 @dataclass
 class Column:
@@ -47,14 +48,13 @@ def get_db_secret(secret):
             "secret manager path not parsable. Could not recover password for DB."
         )
     secret_path = secret
-    # Create a client
-    client = secretmanager_v1.SecretManagerServiceClient()
     # Initialize request argument(s)
     request = secretmanager_v1.AccessSecretVersionRequest(
         name=secret_path,
     )
     # Make the request
-    response = client.access_secret_version(request=request)
+    response = CLIENT.access_secret_version(request=request)
+
     # Return the secret
     return response.payload.data.decode("utf-8")
 

evalbench/eval_service.py

@@ -28,6 +28,7 @@
     get_dataset_from_request,
 )
 
+import threading
 
 SESSIONMANAGER = get_SessionManager()
 rpc_id_var = contextvars.ContextVar("rpc_id", default="default")
@@ -132,6 +133,7 @@ async def Eval(
             model_config,
             db_configs,
         )
+        logging.info(f"Finished Job ID {job_id} Thread count:{threading.active_count()}")
         return eval_response_pb2.EvalResponse(response=f"{job_id}")
 
 

requirements.txt

@@ -26,4 +26,5 @@ GitPython
 sqlparse
 redis
 anthropic[vertex]
-sqlglot
+sqlglot
+pystack