fix: Use env AWS_REGION in sqs_client or default to us-west-2 (#563)

scaleapi · Jul 11, 2024 · b5e4daf · b5e4daf
1 parent c0cea60
commit b5e4daf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/model-engine/model_engine_server/core/celery/celery_autoscaler.py b/model-engine/model_engine_server/core/celery/celery_autoscaler.py
@@ -54,6 +54,7 @@ def excluded_namespaces():
 
 autoscaler_broker = os.environ.get("BROKER_NAME", SQS_BROKER)
 aws_profile = os.environ.get("AWS_PROFILE")
+aws_region = os.environ.get("AWS_REGION", "us-west-2")
 
 
 @dataclasses.dataclass
@@ -403,7 +404,7 @@ async def get_broker_metrics(
 class SQSBroker(AutoscalerBroker):
     @staticmethod
     def _get_sqs_queue_size(queue_name: str):
-        sqs_client = session(aws_profile).client("sqs", region_name="us-west-2")
+        sqs_client = session(aws_profile).client("sqs", region_name=aws_region)
         try:
             total_start_time = time.time()
             queue_size_hist = []