Add excluded location configs to workload testing (#40704)

* Added Excluded Locations

* Added envoy config generator

* Add template

* Updated Setup Document

Author: allenkim0129

sdk/cosmos/azure-cosmos/tests/workloads/dev.md

@@ -9,31 +9,32 @@ file when run. These logs are named in this format `<file name>-<process id>-<da
    - 32 GB RAM
    - Ubuntu
    - Accelerated networking
-2. Fork and clone this repository
-3. Go to azure cosmos folder
+1. Fork and clone this repository
+1. Go to azure cosmos folder
    - `cd azure-sdk-for-python/sdk/cosmos/azure-cosmos`
-4. Install the required packages and create virtual environment
+1. Install the required packages and create virtual environment
    - `sudo apt-get update`
    - `sudo apt-get install python3-pip`
    - `sudo apt-get install python3.12-venv`
    - `python3 -m venv azure-cosmosdb-sdk-environment`
    - `source azure-cosmosdb-sdk-environment/bin/activate`
    - `pip install -r dev_requirements.txt`
-5. Checkout the branch with the changes to test. 
-6. Install azure-cosmos
+1. Checkout the branch with the changes to test. 
+1. Install azure-cosmos
    - `pip install .`
-7. Go to workloads folder
+1. Go to workloads folder
     - `cd tests/workloads`
-8. Fill out relevant configs in `workload_configs.py`: key, host, etc
-9. Install envoy proxy https://www.envoyproxy.io/docs/envoy/latest/start/install
-10. Update envoy_simple_config.yaml to have the correct account info. Replace <> with account name.
-11. Go to envoy folder and start envoy
+1. Fill out relevant configs in `workload_configs.py`: key, host, etc
+1. Install envoy proxy https://www.envoyproxy.io/docs/envoy/latest/start/install
+1. Go to envoy folder and generate envoy configuration file using template. Template files are in `envoy/templates` directory. `<account_name>` is your Cosmos DB account name.
     - `cd envoy`
+    - `./generate_envoy_config.sh <template_file_path> <output_envoy_config_file> <account_name>`
+1. Start envoy using the generated configuration file
     - `mkdir logs`
-    - `envoy -c <envoy_file>.yaml --log-level debug --log-path logs/debug.txt`
-12. Run the setup workload to create the database and containers and insert data
+    - `envoy -c <envoy_config_file>.yaml --log-level debug --log-path logs/debug.txt`
+1. Run the setup workload to create the database and containers and insert data
     - `python3 initial-setup.py`
-13. Run the scale workloads
+1. Run the scale workloads
     - `./run_workloads.sh <number of clients per workload>`
 
 ### Monitor Run

sdk/cosmos/azure-cosmos/tests/workloads/envoy/generate_envoy_config.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# Check if the correct number of arguments are provided
+if [ "$#" -ne 3 ]; then
+    echo "Usage: $0 <template_file_path> <output_envoy_config_file> <account_name>"
+    exit 1
+fi
+
+# Assign arguments to variables
+template_file_path=$1
+output_envoy_config_file=$2
+account_name=$3
+
+# Replace occurrences of "<>" with the account name and write to the new file
+sed "s/<>/$account_name/g" "$template_file_path" > "$output_envoy_config_file"
+
+echo "Replacement complete. The result is saved in $output_envoy_config_file"

sdk/cosmos/azure-cosmos/tests/workloads/envoy/envoy_simple_config.yaml renamed to sdk/cosmos/azure-cosmos/tests/workloads/envoy/templates/envoy_simple_config.yaml

@@ -6,37 +6,32 @@
 import aiohttp
 
 from azure.cosmos import documents
-from workload_utils import create_logger, read_item_concurrently, query_items_concurrently
-from workload_configs import COSMOS_KEY, PREFERRED_LOCATIONS, CONCURRENT_REQUESTS, COSMOS_PROXY_URI, COSMOS_CONTAINER, \
-    COSMOS_DATABASE, CONCURRENT_QUERIES
-from workload_configs import USE_MULTIPLE_WRITABLE_LOCATIONS
+from workload_utils import *
+from workload_configs import *
 
 sys.path.append(r"/")
 
 from azure.cosmos.aio import CosmosClient as AsyncClient
 from azure.core.pipeline.transport import AioHttpTransport
 import asyncio
 
-from datetime import datetime
-
 async def run_workload(client_id, client_logger):
     async with aiohttp.ClientSession(trust_env=True) as proxied_aio_http_session:
         connectionPolicy = documents.ConnectionPolicy()
         connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-
         transport = AioHttpTransport(session=proxied_aio_http_session, session_owner=False)
-        async with AsyncClient(COSMOS_PROXY_URI, COSMOS_KEY, preferred_locations=PREFERRED_LOCATIONS,
+        async with AsyncClient(COSMOS_PROXY_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                               preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                                enable_diagnostics_logging=True, logger=client_logger, transport=transport,
-                               user_agent=str(client_id) + "-" + datetime.now().strftime(
-                                   "%Y%m%d-%H%M%S"), connection_policy=connectionPolicy) as client:
+                               user_agent=get_user_agent(client_id)) as client:
             db = client.get_database_client(COSMOS_DATABASE)
             cont = db.get_container_client(COSMOS_CONTAINER)
             await asyncio.sleep(1)
 
             while True:
                 try:
-                    await read_item_concurrently(cont, CONCURRENT_REQUESTS)
-                    await query_items_concurrently(cont, CONCURRENT_QUERIES)
+                    await read_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                    await query_items_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
                 except Exception as e:
                     client_logger.info("Exception in application layer")
                     client_logger.error(e)

sdk/cosmos/azure-cosmos/tests/workloads/r_proxy_workload.py

@@ -6,9 +6,8 @@
 import aiohttp
 
 from azure.cosmos import documents
-from workload_utils import create_logger, upsert_item_concurrently, read_item_concurrently, query_items_concurrently
-from workload_configs import COSMOS_KEY, PREFERRED_LOCATIONS, USE_MULTIPLE_WRITABLE_LOCATIONS, \
-    CONCURRENT_REQUESTS, COSMOS_PROXY_URI, COSMOS_CONTAINER, COSMOS_DATABASE, CONCURRENT_QUERIES
+from workload_utils import *
+from workload_configs import *
 
 sys.path.append(r"/")
 
@@ -17,27 +16,25 @@
 import asyncio
 
 import time
-from datetime import datetime
 
 async def run_workload(client_id, client_logger):
-    connectionPolicy = documents.ConnectionPolicy()
-    connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
     async with aiohttp.ClientSession(trust_env=True) as proxied_aio_http_session:
-
+        connectionPolicy = documents.ConnectionPolicy()
+        connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
         transport = AioHttpTransport(session=proxied_aio_http_session, session_owner=False)
-        async with AsyncClient(COSMOS_PROXY_URI, COSMOS_KEY, preferred_locations=PREFERRED_LOCATIONS,
+        async with AsyncClient(COSMOS_PROXY_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                               preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                                enable_diagnostics_logging=True, logger=client_logger, transport=transport,
-                               user_agent=str(client_id) + "-" + datetime.now().strftime(
-                                   "%Y%m%d-%H%M%S"), connection_policy=connectionPolicy) as client:
+                               user_agent=get_user_agent(client_id)) as client:
             db = client.get_database_client(COSMOS_DATABASE)
             cont = db.get_container_client(COSMOS_CONTAINER)
-            time.sleep(1)
+            await asyncio.sleep(1)
 
             while True:
                 try:
-                    await upsert_item_concurrently(cont, CONCURRENT_REQUESTS)
-                    await read_item_concurrently(cont, CONCURRENT_REQUESTS)
-                    await query_items_concurrently(cont, CONCURRENT_QUERIES)
+                    await upsert_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                    await read_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                    await query_items_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
                 except Exception as e:
                     client_logger.info("Exception in application layer")
                     client_logger.error(e)

sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_proxy_workload.py

@@ -4,35 +4,29 @@
 import sys
 
 from azure.cosmos import documents
-from workload_utils import upsert_item_concurrently, read_item_concurrently, query_items_concurrently, create_logger
-from workload_configs import (COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS,
-                                              USE_MULTIPLE_WRITABLE_LOCATIONS,
-                                              CONCURRENT_REQUESTS, COSMOS_DATABASE, COSMOS_CONTAINER,
-                                              CONCURRENT_QUERIES)
+from workload_utils import *
+from workload_configs import *
 sys.path.append(r"/")
 
 from azure.cosmos.aio import CosmosClient as AsyncClient
 import asyncio
 
-from datetime import datetime
-
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    client = AsyncClient(COSMOS_URI, COSMOS_KEY,
-                           enable_diagnostics_logging=True, logger=client_logger,
-                           user_agent=str(client_id) + "-" + datetime.now().strftime(
-                               "%Y%m%d-%H%M%S"), preferred_locations=PREFERRED_LOCATIONS,
-                           connection_policy=connectionPolicy)
+    client = AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                         preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
+                         enable_diagnostics_logging=True, logger=client_logger,
+                         user_agent=get_user_agent(client_id))
     db = client.get_database_client(COSMOS_DATABASE)
     cont = db.get_container_client(COSMOS_CONTAINER)
     await asyncio.sleep(1)
 
     while True:
         try:
-            await upsert_item_concurrently(cont, CONCURRENT_REQUESTS)
-            await read_item_concurrently(cont, CONCURRENT_REQUESTS)
-            await query_items_concurrently(cont, CONCURRENT_QUERIES)
+            await upsert_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+            await read_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+            await query_items_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
         except Exception as e:
             client_logger.info("Exception in application layer")
             client_logger.error(e)

sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_with_incorrect_client_workload.py

@@ -4,33 +4,29 @@
 import sys
 
 from azure.cosmos import documents
-from workload_utils import upsert_item_concurrently, read_item_concurrently, query_items_concurrently, create_logger
-from workload_configs import (COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS, USE_MULTIPLE_WRITABLE_LOCATIONS,
-                                              CONCURRENT_REQUESTS, COSMOS_DATABASE, COSMOS_CONTAINER, CONCURRENT_QUERIES)
+from workload_utils import *
+from workload_configs import *
 sys.path.append(r"/")
 
 from azure.cosmos.aio import CosmosClient as AsyncClient
 import asyncio
 
-from datetime import datetime
-
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY,
+    async with AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                           preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=client_logger,
-                           user_agent=str(client_id) + "-" + datetime.now().strftime(
-                               "%Y%m%d-%H%M%S"), preferred_locations=PREFERRED_LOCATIONS,
-                           connection_policy=connectionPolicy) as client:
+                           user_agent=get_user_agent(client_id)) as client:
         db = client.get_database_client(COSMOS_DATABASE)
         cont = db.get_container_client(COSMOS_CONTAINER)
         await asyncio.sleep(1)
 
         while True:
             try:
-                await upsert_item_concurrently(cont, CONCURRENT_REQUESTS)
-                await read_item_concurrently(cont, CONCURRENT_REQUESTS)
-                await query_items_concurrently(cont, CONCURRENT_QUERIES)
+                await upsert_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                await read_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                await query_items_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
             except Exception as e:
                 client_logger.info("Exception in application layer")
                 client_logger.error(e)

sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload.py

@@ -2,61 +2,30 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 import os
 import sys
+import time
 
-from workload_utils import create_logger, get_random_item
-from workload_configs import (COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS, USE_MULTIPLE_WRITABLE_LOCATIONS,
-                                              CONCURRENT_REQUESTS, COSMOS_DATABASE, COSMOS_CONTAINER, CONCURRENT_QUERIES,
-                                              PARTITION_KEY)
-
+from workload_utils import *
+from workload_configs import *
 sys.path.append(r"/")
 
 from azure.cosmos import CosmosClient, documents
 
-import time
-from datetime import datetime
-
-def upsert_item(container, num_upserts):
-    for _ in range(num_upserts):
-        container.upsert_item(get_random_item(), etag=None, match_condition=None)
-
-
-def read_item(container, num_reads):
-    for _ in range(num_reads):
-        item = get_random_item()
-        container.read_item(item["id"], item[PARTITION_KEY], etag=None, match_condition=None)
-
-
-def query_items(container, num_queries):
-    for _ in range(num_queries):
-        perform_query(container)
-
-
-def perform_query(container):
-    random_item = get_random_item()
-    results = container.query_items(query="SELECT * FROM c where c.id=@id and c.pk=@pk",
-                                    parameters=[{"name": "@id", "value": random_item["id"]},
-                                                {"name": "@pk", "value": random_item["pk"]}],
-                                    partition_key=random_item[PARTITION_KEY])
-    items = [item for item in results]
-
-
 def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    with CosmosClient(COSMOS_URI, COSMOS_KEY,
-                           enable_diagnostics_logging=True, logger=client_logger,
-                           user_agent=str(client_id) + "-" + datetime.now().strftime(
-                               "%Y%m%d-%H%M%S"), preferred_locations=PREFERRED_LOCATIONS,
-                      connection_policy=connectionPolicy) as client:
+    with CosmosClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                      preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
+                      enable_diagnostics_logging=True, logger=client_logger,
+                      user_agent=get_user_agent(client_id)) as client:
         db = client.get_database_client(COSMOS_DATABASE)
         cont = db.get_container_client(COSMOS_CONTAINER)
         time.sleep(1)
 
         while True:
             try:
-                upsert_item(cont, CONCURRENT_REQUESTS)
-                read_item(cont, CONCURRENT_REQUESTS)
-                query_items(cont, CONCURRENT_QUERIES)
+                upsert_item(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                read_item(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                query_items(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
             except Exception as e:
                 client_logger.info("Exception in application layer")
                 client_logger.error(e)

sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload_sync.py

@@ -4,34 +4,28 @@
 import sys
 
 from azure.cosmos import documents
-from workload_utils import create_logger, read_item_concurrently, query_items_concurrently
-from workload_configs import (COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS, USE_MULTIPLE_WRITABLE_LOCATIONS,
-                                              CONCURRENT_REQUESTS, COSMOS_CONTAINER, COSMOS_DATABASE, CONCURRENT_QUERIES)
-
+from workload_utils import *
+from workload_configs import *
 sys.path.append(r"/")
 
 from azure.cosmos.aio import CosmosClient as AsyncClient
 import asyncio
 
-from datetime import datetime
-
 async def run_workload(client_id, client_logger):
-
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY,
+    async with AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+                           preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=client_logger,
-                           user_agent=str(client_id) + "-" + datetime.now().strftime(
-                               "%Y%m%d-%H%M%S"), preferred_locations=PREFERRED_LOCATIONS,
-                           connection_policy=connectionPolicy) as client:
+                           user_agent=get_user_agent(client_id)) as client:
         db = client.get_database_client(COSMOS_DATABASE)
         cont = db.get_container_client(COSMOS_CONTAINER)
         await asyncio.sleep(1)
 
         while True:
             try:
-                await read_item_concurrently(cont, CONCURRENT_REQUESTS)
-                await query_items_concurrently(cont, CONCURRENT_QUERIES)
+                await read_item_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_REQUESTS)
+                await query_items_concurrently(cont, REQUEST_EXCLUDED_LOCATIONS, CONCURRENT_QUERIES)
             except Exception as e:
                 client_logger.info("Exception in application layer")
                 client_logger.error(e)