locust-cloud 1.0.17__py3-none-any.whl → 1.0.18__py3-none-any.whl

locust_cloud/cloud.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import math
 import os
 import sys
 import time
@@ -13,19 +14,13 @@ import requests
 from botocore.exceptions import ClientError
 from locust_cloud.constants import (
     DEFAULT_CLUSTER_NAME,
+    DEFAULT_LAMBDA_URL,
     DEFAULT_NAMESPACE,
     DEFAULT_REGION_NAME,
-    LAMBDA_URL,
+    USERS_PER_WORKER,
 )
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
-logging.basicConfig(
-    format="[LOCUST-CLOUD] %(levelname)s: %(message)s",
-    level=logging.INFO,
-)
-logger = logging.getLogger(__name__)
-
-
 LOCUST_ENV_VARIABLE_IGNORE_LIST = ["LOCUST_BUILD_PATH", "LOCUST_SKIP_MONKEY_PATCH"]
 
 
@@ -67,9 +62,9 @@ parser = configargparse.ArgumentParser(
             configargparse.DefaultConfigFileParser,
         ]
     ),
-    description="""A tool for Locust Cloud users to deploy clusters.
+    description="""Launches distributed Locust runs on locust.cloud infrastructure.
 
-Example: locust-cloud -f locust.py --aws-access-key-id 123 --aws-secret-access-key 456""",
+Example: locust-cloud -f my_locustfile.py --aws-region-name us-east-1 --users 1000""",
     epilog="""Any parameters not listed here are forwarded to locust master unmodified, so go ahead and use things like --users, --host, --run-time, ...
 Locust config can also be set using config file (~/.locust.conf, locust.conf, pyproject.toml, ~/.cloud.conf or cloud.conf).
 Parameters specified on command line override env vars, which in turn override config files.""",
@@ -113,61 +108,104 @@ parser.add_argument(
     help="Sets the namespace for scoping the deployed cluster",
     env_var="KUBE_NAMESPACE",
 )
+parser.add_argument(
+    "--lambda-url",
+    type=str,
+    default=DEFAULT_LAMBDA_URL,
+    help="Sets the namespace for scoping the deployed cluster",
+    env_var="LOCUST_CLOUD_LAMBDA",
+)
 parser.add_argument(
     "--aws-access-key-id",
     type=str,
-    help="Authentication for deploying with Locust Cloud",
+    help=configargparse.SUPPRESS,
     env_var="AWS_ACCESS_KEY_ID",
+    default=None,
 )
 parser.add_argument(
     "--aws-secret-access-key",
     type=str,
-    help="Authentication for deploying with Locust Cloud",
+    help=configargparse.SUPPRESS,
     env_var="AWS_SECRET_ACCESS_KEY",
+    default=None,
+)
+parser.add_argument(
+    "--username",
+    type=str,
+    help="Authentication for deploying with Locust Cloud",
+    env_var="LOCUST_CLOUD_USERNAME",
+    default=None,
+)
+parser.add_argument(
+    "--password",
+    type=str,
+    help="Authentication for deploying with Locust Cloud",
+    env_var="LOCUST_CLOUD_PASSWORD",
+    default=None,
+)
+parser.add_argument(
+    "--loglevel",
+    "-L",
+    type=str,
+    help="Log level",
+    env_var="LOCUST_CLOUD_LOGLEVEL",
+    default="INFO",
+)
+parser.add_argument(
+    "--workers",
+    type=str,
+    help="Number of workers to use for the deployment",
+    env_var="LOCUST_CLOUD_WORKERS",
+    default=None,
 )
 
 options, locust_options = parser.parse_known_args()
-
-username = os.environ.get("LOCUST_CLOUD_USERNAME")
-password = os.environ.get("LOCUST_CLOUD_PASSWORD")
+logging.basicConfig(
+    format="[LOCUST-CLOUD] %(levelname)s: %(message)s",
+    level=options.loglevel.upper(),
+)
+logger = logging.getLogger(__name__)
+# Restore log level for other libs. Yes, this can be done more nicely
+logging.getLogger("botocore").setLevel(logging.INFO)
+logging.getLogger("boto3").setLevel(logging.INFO)
+logging.getLogger("s3transfer").setLevel(logging.INFO)
+logging.getLogger("requests").setLevel(logging.INFO)
+logging.getLogger("urllib3").setLevel(logging.INFO)
 
 
 def main() -> None:
     s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
     deployed_pods: list[Any] = []
+    worker_count = options.workers or math.ceil(locust_options.users / USERS_PER_WORKER)
+    worker_count = worker_count if worker_count > 2 else 2
 
     try:
-        if options.aws_access_key_id and options.aws_secret_access_key:
-            credential_manager = CredentialManager(
-                lambda_url=LAMBDA_URL,
-                region_name=options.aws_region_name,
-                access_key=options.aws_access_key_id,
-                secret_key=options.aws_secret_access_key,
-            )
-        elif username and password:
-            credential_manager = CredentialManager(
-                lambda_url=LAMBDA_URL,
-                username=username,
-                password=password,
-                region_name=options.aws_region_name,
-            )
-        else:
+        if not (
+            (options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)
+        ):
             logger.error(
-                "Authentication is required to use Locust Cloud. Provide either AWS credentials or set the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables."
+                "Authentication is required to use Locust Cloud. Please ensure the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
             )
             sys.exit(1)
 
+        logger.info(f"Logging you into Locust Cloud ({options.lambda_url}, {options.aws_region_name})")
+
+        credential_manager = CredentialManager(
+            lambda_url=options.lambda_url,
+            access_key=options.aws_access_key_id,
+            secret_key=options.aws_secret_access_key,
+            username=options.username,
+            password=options.password,
+            region_name=options.aws_region_name,
+        )
+
         credentials = credential_manager.get_current_credentials()
-        cognito_client_id_token: str = credentials["cognito_client_id_token"]
+        cognito_client_id_token = credentials["cognito_client_id_token"]
         aws_access_key_id = credentials.get("access_key")
         aws_secret_access_key = credentials.get("secret_key")
-        aws_session_token = credentials.get("token")
-
-        if not all([aws_access_key_id, aws_secret_access_key]):
-            logger.error("Authentication failed: Missing AWS credentials.")
-            sys.exit(1)
+        aws_session_token = credentials.get("token", "")
 
-        logger.info(f"Uploading {options.locustfile} to S3 bucket {s3_bucket}...")
+        logger.info(f"Uploading {options.locustfile}")
         s3 = credential_manager.session.client("s3")
         try:
             s3.upload_file(options.locustfile, s3_bucket, os.path.basename(options.locustfile))
@@ -176,17 +214,17 @@ def main() -> None:
                 Params={"Bucket": s3_bucket, "Key": os.path.basename(options.locustfile)},
                 ExpiresIn=3600,
             )
-            logger.info(f"Uploaded {options.locustfile} successfully.")
+            logger.debug(f"Uploaded {options.locustfile} successfully")
         except FileNotFoundError:
             logger.error(f"File not found: {options.locustfile}")
             sys.exit(1)
         except ClientError as e:
-            logger.error(f"Failed to upload {options.locustfile} to S3: {e}")
+            logger.error(f"Failed to upload {options.locustfile}: {e}")
             sys.exit(1)
 
         requirements_url = ""
         if options.requirements:
-            logger.info(f"Uploading {options.requirements} to S3 bucket {s3_bucket} as requirements.txt...")
+            logger.info(f"Uploading {options.requirements}")
             try:
                 s3.upload_file(options.requirements, s3_bucket, "requirements.txt")
                 requirements_url = s3.generate_presigned_url(
@@ -194,44 +232,46 @@ def main() -> None:
                     Params={"Bucket": s3_bucket, "Key": "requirements.txt"},
                     ExpiresIn=3600,
                 )
-                logger.info(f"Uploaded {options.requirements} successfully.")
+                logger.debug(f"Uploaded {options.requirements} successfully")
             except FileNotFoundError:
                 logger.error(f"File not found: {options.requirements}")
                 sys.exit(1)
             except ClientError as e:
-                logger.error(f"Failed to upload {options.requirements} to S3: {e}")
+                logger.error(f"Failed to upload {options.requirements}: {e}")
                 sys.exit(1)
 
-        logger.info("Deploying load generators via API Gateway...")
+        logger.info("Deploying load generators")
         locust_env_variables = [
             {"name": env_variable, "value": str(os.environ[env_variable])}
             for env_variable in os.environ
             if env_variable.startswith("LOCUST_") and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{LAMBDA_URL}/{options.kube_cluster_name}"
+        deploy_endpoint = f"{options.lambda_url}/{options.kube_cluster_name}"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_LOCUSTFILE", "value": locustfile_url},
                 {"name": "LOCUST_REQUIREMENTS_URL", "value": requirements_url},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
+                {"name": "LOCUST_WEB_HOST_DISPLAY_NAME", "value": "_____"},
                 *locust_env_variables,
-            ]
+            ],
+            "worker_count": worker_count,
         }
         headers = {
             "Authorization": f"Bearer {cognito_client_id_token}",
             "Content-Type": "application/json",
             "AWS_ACCESS_KEY_ID": aws_access_key_id,
             "AWS_SECRET_ACCESS_KEY": aws_secret_access_key,
-            "AWS_SESSION_TOKEN": aws_session_token if aws_session_token else "",
+            "AWS_SESSION_TOKEN": aws_session_token,
         }
         try:
             response = requests.post(deploy_endpoint, json=payload, headers=headers)
         except requests.exceptions.RequestException as e:
-            logger.error(f"HTTP request failed: {e}")
+            logger.error(f"Failed to deploy the load generators: {e}")
             sys.exit(1)
+
         if response.status_code == 200:
             deployed_pods = response.json().get("pods", [])
-            logger.info("Load generators deployed successfully.")
         else:
             logger.error(
                 f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
@@ -241,13 +281,20 @@ def main() -> None:
         logger.error(f"Credential error: {ce}")
         sys.exit(1)
 
+    log_group_name = f"/eks/{options.kube_cluster_name}-{options.kube_namespace}"
+    master_pod_name = next((pod for pod in deployed_pods if "master" in pod), None)
+
+    if not master_pod_name:
+        logger.error(
+            "Master pod not found among deployed pods. Something went wrong during the load generator deployment, please try again or contact an administrator for assistance"
+        )
+        sys.exit(1)
+
+    logger.debug("Load generators deployed successfully!")
+
     try:
         logger.info("Waiting for pods to be ready...")
-        log_group_name = f"/eks/{options.kube_cluster_name}-{options.kube_namespace}"
-        master_pod_name = next((pod for pod in deployed_pods if "master" in pod), None)
-        if not master_pod_name:
-            logger.error("Master pod not found among deployed pods.")
-            sys.exit(1)
+
         log_stream: str | None = None
         while log_stream is None:
             try:
@@ -264,7 +311,7 @@ def main() -> None:
             except ClientError as e:
                 logger.error(f"Error describing log streams: {e}")
                 time.sleep(5)
-        logger.info("Pods are ready, switching to Locust logs.")
+        logger.debug("Pods are ready, switching to Locust logs")
 
         timestamp = int((datetime.now(UTC) - timedelta(minutes=5)).timestamp() * 1000)
 
@@ -291,10 +338,10 @@ def main() -> None:
             except ClientError as e:
                 error_code = e.response.get("Error", {}).get("Code", "")
                 if error_code == "ExpiredTokenException":
-                    logger.warning("AWS session token expired during log streaming. Refreshing credentials...")
+                    logger.debug("AWS session token expired during log streaming. Refreshing credentials.")
                     time.sleep(5)
     except KeyboardInterrupt:
-        logger.debug("Interrupted by user.")
+        logger.debug("Interrupted by user")
     except Exception as e:
         logger.exception(e)
         sys.exit(1)
@@ -315,28 +362,26 @@ def main() -> None:
                 headers["AWS_SESSION_TOKEN"] = token
 
             response = requests.delete(
-                f"{LAMBDA_URL}/{options.kube_cluster_name}",
+                f"{options.lambda_url}/{options.kube_cluster_name}",
                 headers=headers,
                 params={"namespace": options.kube_namespace} if options.kube_namespace else {},
             )
 
             if response.status_code != 200:
                 logger.error(
-                    f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
+                    f"Could not automatically tear down Locust Cloud: HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
                 )
-            else:
-                logger.info("Cluster teardown initiated successfully.")
         except Exception as e:
             logger.error(f"Could not automatically tear down Locust Cloud: {e}")
 
         try:
-            logger.info(f"Cleaning up S3 bucket: {s3_bucket}")
+            logger.info("Cleaning up locust files")
             s3 = credential_manager.session.resource("s3")
             bucket = s3.Bucket(s3_bucket)
             bucket.objects.all().delete()
-            logger.info(f"S3 bucket {s3_bucket} cleaned up successfully.")
+            logger.info("Done! ✨")
         except ClientError as e:
-            logger.error(f"Failed to clean up S3 bucket {s3_bucket}: {e}")
+            logger.error(f"Failed to clean up locust files: {e}")
             sys.exit(1)
 
 

locust_cloud/constants.py

@@ -1,4 +1,5 @@
 DEFAULT_REGION_NAME = "us-east-1"
 DEFAULT_CLUSTER_NAME = "dmdb"
 DEFAULT_NAMESPACE = "default"
-LAMBDA_URL = "https://api.locust.cloud/1"
+DEFAULT_LAMBDA_URL = "https://api.locust.cloud/1"
+USERS_PER_WORKER = 500

locust_cloud/credential_manager.py

@@ -10,10 +10,6 @@ from botocore.credentials import RefreshableCredentials
 from botocore.session import Session as BotocoreSession
 from locust_cloud.constants import DEFAULT_REGION_NAME
 
-logging.basicConfig(
-    format="[LOCUST-CLOUD] %(levelname)s: %(message)s",
-    level=logging.INFO,
-)
 logger = logging.getLogger(__name__)