[SG-3403] autoscale shared runner

Dockerfile

@@ -13,4 +13,4 @@ COPY aws_service.py ${LAMBDA_TASK_ROOT}/
 COPY lambda.py ${LAMBDA_TASK_ROOT}/
 
 # Set the CMD to your handler
-CMD ["lambda.lambda_handler"]
+CMD ["lambda.lambda_handler"]

Dockerfile.pyinstaller

@@ -1,6 +1,9 @@
 # Multi-stage build for AWS Lambda with PyInstaller bundling
 # Stage 1: Build stage with PyInstaller - use same base as Lambda runtime
-FROM public.ecr.aws/lambda/python:3.13 as builder
+FROM public.ecr.aws/lambda/python:3.13 AS builder
+
+ARG type="shared-external"
+ENV SG_RUNNER_TYPE="${type}"
 
 # Install system dependencies for PyInstaller
 RUN microdnf update -y && microdnf install -y \

Makefile

@@ -8,13 +8,14 @@ PROD_ACCOUNT_ID =
 
 # Registry and image configuration
 REGISTRY = $(ACCOUNT_ID).dkr.ecr.eu-central-1.amazonaws.com
-IMAGE_NAME = private-runner/autoscaler
+IMAGE_NAME = autoscaler/runner
 FULL_IMAGE = $(REGISTRY)/$(IMAGE_NAME)
 
 # Docker build configuration
 PLATFORM = linux/arm64
 DOCKERFILE = Dockerfile.pyinstaller
 DOCKER_BUILD_ARGS = --push --pull --platform $(PLATFORM) --provenance=false
+SG_RUNNER_TYPE = shared-external
 
 # AWS profiles
 DASH_PROFILE = default
@@ -44,22 +45,32 @@ build: login ## Build and push Docker image using VERSION variable
 	@echo "Building and pushing image: $(FULL_IMAGE):$(VERSION)"
 	@docker buildx build \
 		-f $(DOCKERFILE) $(DOCKER_BUILD_ARGS) \
+		--build-arg type="$(SG_RUNNER_TYPE)" \
+		-t $(FULL_IMAGE):latest \
 		-t $(FULL_IMAGE):$(VERSION) .
 	@echo "Build completed: $(FULL_IMAGE):$(VERSION)"
 
 dash: PROFILE=$(DASH_PROFILE)
 dash: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
+dash: IMAGE_NAME="autoscaler/external-runner"
+dash: SG_RUNNER_TYPE="external"
 dash: build ## Build and push Docker image for Dash environment with compiled and tagged code
 
-dash-raw: PROFILE=$(DASH_PROFILE)
-dash-raw: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
-dash-raw: DOCKERFILE=Dockerfile
-dash-raw: build ## Build and push Docker image for Dash environment with source Python code
-
-dash-lts: PROFILE=$(DASH_PROFILE)
-dash-lts: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
-dash-lts: VERSION="latest"
-dash-lts: build ## Build and push Docker image for Dash environment with compiled and latest tag
+dash-shared: PROFILE=$(DASH_PROFILE)
+dash-shared: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
+dash-shared: IMAGE_NAME="autoscaler/shared-runner"
+dash-shared: SG_RUNNER_TYPE="shared-external"
+dash-shared: build ## Build and push Docker image for Dash environment with shared runner
+
+# dash-raw: PROFILE=$(DASH_PROFILE)
+# dash-raw: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
+# dash-raw: DOCKERFILE=Dockerfile
+# dash-raw: build ## Build and push Docker image for Dash environment with source Python code
+
+# dash-lts: PROFILE=$(DASH_PROFILE)
+# dash-lts: ACCOUNT_ID=$(DASH_ACCOUNT_ID)
+# dash-lts: VERSION="latest"
+# dash-lts: build ## Build and push Docker image for Dash environment with compiled and latest tag
 
 prod: PROFILE=$(PROD_PROFILE)
 prod: ACCOUNT_ID=$(PROD_ACCOUNT_ID)

lambda.py

@@ -0,0 +1,53 @@
+import json
+import sys
+import logging
+from stackguardian_autoscaler import StackGuardianAutoscaler
+from aws_service import AwsService
+
+# Setup logger
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+def lambda_handler(event, context):
+    """
+    AWS Lambda function handler.
+
+    Args:
+        event (dict): The event data passed to the Lambda function.
+        context (LambdaContext): The context object provided by AWS Lambda.
+
+    Returns:
+        dict: A response object containing the status code and message.
+    """
+    # Log the incoming event for debugging
+    # print("Received event:", event)
+    # logger.info(f"Received event: {event}")
+
+    # Process the event (this is a placeholder for your actual logic)
+    autoscaler = StackGuardianAutoscaler(cloud_service=AwsService())
+    try:
+        autoscaler.start()
+
+        # Create a response object
+        response = {"statusCode": 200, "body": "success"}
+    except Exception as e:
+        response = {"statusCode": 500, "body": str(e)}
+
+    return response
+
+
+if __name__ == "__main__":
+    # Read event and context from stdin when executed as binary
+    try:
+        input_data = json.loads(sys.stdin.read())
+        event = input_data.get("event", {})
+        context = input_data.get("context", "")
+
+        result = lambda_handler(event, context)
+        sys.exit(0 if result["statusCode"] == 200 else 1)
+    except Exception as e:
+        print(f"Exception in main: {e}", file=sys.stderr)
+        error_response = {"statusCode": 500, "body": str(e)}
+        print(json.dumps(error_response))
+        sys.exit(1)

stackguardian_autoscaler.py

@@ -95,6 +95,8 @@ def __init__(self, cloud_service: CloudService):
         self.SG_ORG = os.getenv("SG_ORG")
         self.SG_RUNNER_GROUP = os.getenv("SG_RUNNER_GROUP")
 
+        self.SG_RUNNER_TYPE = os.getenv("SG_RUNNER_TYPE")
+
         self.cloud_service = cloud_service
 
         self.scale_in_cooldown_duration = timedelta(
@@ -278,6 +280,7 @@ def _refresh_sg_runner_group(self):
         res.raise_for_status()
 
         self.sg_runner_group = res.json()
+
         sg_runners = []
         for runner in self.sg_runner_group.get("msg").get(
             "ContainerInstances"
@@ -287,8 +290,15 @@ def _refresh_sg_runner_group(self):
         self.sg_runners = sg_runners
 
     def _refresh_queued_jobs(self) -> int:
+        if self.SG_RUNNER_TYPE == "shared-external":
+            queued_workflows_count_key = "EcsTaskFailAgentCount"
+        elif self.SG_RUNNER_TYPE == "external":
+            queued_workflows_count_key = "QueuedWorkflowsCount"
+        else:
+            raise Exception("Invalid runner type")
+
         queued_jobs = self.sg_runner_group.get("msg").get(
-            "QueuedWorkflowsCount"
+            queued_workflows_count_key
         )
 
         if queued_jobs is None: