diff --git a/src/benchmarks/ann-benchmarks.py b/src/benchmarks/ann-benchmarks.py index 18504a7d0..869e7a4e7 100644 --- a/src/benchmarks/ann-benchmarks.py +++ b/src/benchmarks/ann-benchmarks.py @@ -18,6 +18,27 @@ import boto3 import paramiko + +# Obtain session token programmatically +def get_temporary_session_token(): + sts_client = boto3.client("sts") + response = sts_client.get_session_token(DurationSeconds=3600) + return response["Credentials"] + + +# Retrieve temporary credentials +temp_credentials = get_temporary_session_token() + +# Initialize a session using the temporary credentials +session = boto3.Session( + aws_access_key_id=temp_credentials["AccessKeyId"], + aws_secret_access_key=temp_credentials["SecretAccessKey"], + aws_session_token=temp_credentials["SessionToken"], +) + +# Use the session to create a client +ec2 = session.client("ec2") + installations = ["tiledb"] algorithms = ["tiledb-ivf-flat", "tiledb-ivf-pq", "tiledb-flat", "tiledb-vamana"] @@ -27,16 +48,21 @@ installations += [ "flann", "faiss", + "faiss_hnsw", "hnswlib", # "weaviate" # "milvus", "pgvector", + "annoy", + "qdrant", + "elasticsearch", ] algorithms += [ "flann", "faiss-ivf", # "faiss-lsh", "faiss-ivfpqfs", + "hnsw(faiss)", "hnswlib", # "weaviate", # "milvus-flat", @@ -45,6 +71,9 @@ # "milvus-scann", # "milvus-hnsw", "pgvector", + "annoy", + "qdrant", + "elasticsearch", ] # If you want to skip running benchmarks and keep the instance running after the script finishes, @@ -65,6 +94,7 @@ ) # You do not need to change these. +# Note this security group is in eu-central-1 (Frankfurt). security_group_ids = ["sg-04258b401ce76d246"] # 64 vCPU, 512 GiB, EBS-Only. instance_type = "r6i.16xlarge"