Fixed perf daemon artifacts generation (#238)

Author: filipecosta90

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redisbench-admin"
-version = "0.5.14"
+version = "0.5.15"
 description = "Redis benchmark run helper. A wrapper around Redis and Redis Modules benchmark tools ( ftsb_redisearch, memtier_benchmark, redis-benchmark, aibench, etc... )."
 authors = ["filipecosta90 <[email protected]>","Redis Performance Group <[email protected]>"]
 readme = "README.md"

redisbench_admin/profilers/daemon.py

@@ -54,6 +54,24 @@ def set_app_loggers(self, app):
         app.logger.addHandler(handler)
         self.perf.set_logger(app.logger)
 
+    def update_ec2_vars_from_request(self, request):
+        aws_access_key_id = None
+        aws_secret_access_key = None
+        aws_session_token = None
+        region_name = "us-east-2"
+        if request.is_json:
+            data = request.get_json()
+            if "aws_access_key_id" in data:
+                aws_access_key_id = data["aws_access_key_id"]
+            if "aws_secret_access_key" in data:
+                aws_secret_access_key = data["aws_secret_access_key"]
+            if "aws_session_token" in data:
+                aws_session_token = data["aws_session_token"]
+            if "region_name" in data:
+                region_name = data["region_name"]
+
+        return aws_access_key_id, aws_secret_access_key, aws_session_token, region_name
+
     def update_vars_from_request(self, request, app):
         app.logger.info("Updating vars from request")
         self.dso = ""
@@ -158,9 +176,15 @@ def profile_status(profiler_name, pid):
         def profile_stop(profiler_name, pid):
             profile_res = self.perf.stop_profile()
             profilers_artifacts_matrix = []
-            bucket_location = "us-east-2"
+
             primary_id = 1
             total_primaries = 1
+            (
+                aws_access_key_id,
+                aws_secret_access_key,
+                aws_session_token,
+                region_name,
+            ) = self.update_ec2_vars_from_request(request)
             if profile_res is True:
                 # Generate:
                 #  - artifact with Flame Graph SVG
@@ -209,7 +233,10 @@ def profile_stop(profiler_name, pid):
                         [profile_artifact],
                         S3_BUCKET_NAME,
                         s3_bucket_path,
-                        bucket_location,
+                        aws_access_key_id,
+                        aws_secret_access_key,
+                        aws_session_token,
+                        region_name,
                     )
                     s3_link = list(url_map.values())[0]
                 profilers_artifacts_matrix.append(

redisbench_admin/profilers/perf_daemon_caller.py

@@ -30,6 +30,11 @@ def __init__(self, remote_endpoint, **kwargs):
         self.pid = None
         self.test_name = ""
         self.setup_name = ""
+        self.aws_access_key_id = None
+        self.aws_secret_access_key = None
+        self.aws_session_token = None
+        self.region_name = None
+
         if "dso" in kwargs:
             self.dso = kwargs.get("dso")
         if "test_name" in kwargs:
@@ -46,6 +51,14 @@ def __init__(self, remote_endpoint, **kwargs):
             self.github_org_name = kwargs.get("github_org_name")
         if "github_sha" in kwargs:
             self.github_sha = kwargs.get("github_sha")
+        if "aws_access_key_id" in kwargs:
+            self.aws_access_key_id = kwargs["aws_access_key_id"]
+        if "aws_secret_access_key" in kwargs:
+            self.aws_secret_access_key = kwargs["aws_secret_access_key"]
+        if "aws_session_token" in kwargs:
+            self.aws_session_token = kwargs["aws_session_token"]
+        if "region_name" in kwargs:
+            self.region_name = kwargs["region_name"]
 
     def set_logger(self, logger_app):
         self.logger = logger_app
@@ -87,8 +100,17 @@ def stop_profile(self, **kwargs):
             url = "http://{}/profiler/perf/stop/{}".format(
                 self.remote_endpoint, self.pid
             )
+            data = {}
+            if self.aws_access_key_id is not None:
+                data["aws_access_key_id"] = self.aws_access_key_id
+            if self.aws_secret_access_key is not None:
+                data["aws_secret_access_key"] = self.aws_secret_access_key
+            if self.aws_session_token is not None:
+                data["aws_session_token"] = self.aws_session_token
+            if self.region_name is not None:
+                data["region_name"] = self.region_name
 
-            response = requests.post(url)
+            response = requests.post(url, data=None, json=data)
             if response.status_code == 200:
                 result = True
                 status_dict = response.json()

redisbench_admin/run_remote/run_remote.py

@@ -50,6 +50,9 @@
 from redisbench_admin.utils.utils import (
     EC2_PRIVATE_PEM,
     upload_artifacts_to_s3,
+    EC2_ACCESS_KEY,
+    EC2_SECRET_KEY,
+    EC2_REGION,
 )
 
 
@@ -381,6 +384,9 @@ def run_remote_command_logic(args, project_name, project_version):
                                             github_repo_name=tf_github_repo,
                                             github_org_name=tf_github_org,
                                             github_sha=tf_github_sha,
+                                            aws_access_key_id=EC2_ACCESS_KEY,
+                                            aws_secret_access_key=EC2_SECRET_KEY,
+                                            region_name=EC2_REGION,
                                         )
                                         primary_one_pid = redis_conns[0].info()[
                                             "process_id"

redisbench_admin/utils/redisgraph_benchmark_go.py

@@ -15,7 +15,7 @@ def setup_remote_benchmark_agent(
     commands = [
         "apt install python3-pip -y",
         "pip3 install redisbench-admin>=0.5.12",
-        "perf-daemon start",
+        "/usr/local/bin/perf-daemon start",
     ]
     execute_remote_commands(
         client_public_ip, username, private_key, commands, client_ssh_port

redisbench_admin/utils/utils.py

@@ -26,15 +26,28 @@
 
 
 def upload_artifacts_to_s3(
-    artifacts, s3_bucket_name, s3_bucket_path, bucket_location=None
+    artifacts,
+    s3_bucket_name,
+    s3_bucket_path,
+    aws_access_key_id=None,
+    aws_secret_access_key=None,
+    aws_session_token=None,
+    region_name=None,
 ):
     artifacts_map = {}
+    if region_name is None:
+        region_name = EC2_REGION
     logging.info("-- uploading results to s3 -- ")
-    s3 = boto3.resource("s3")
+    if aws_access_key_id is not None and aws_secret_access_key is not None:
+        session = boto3.Session(
+            aws_access_key_id, aws_secret_access_key, aws_session_token, region_name
+        )
+        s3 = session.resource("s3")
+    else:
+        s3 = boto3.resource("s3")
     bucket = s3.Bucket(s3_bucket_name)
     progress = tqdm(unit="files", total=len(artifacts))
-    if bucket_location is None:
-        bucket_location = EC2_REGION
+
     for full_artifact_path in artifacts:
         artifact = os.path.basename(full_artifact_path)
         object_key = "{bucket_path}{filename}".format(
@@ -46,7 +59,7 @@ def upload_artifacts_to_s3(
         object_acl.put(ACL="public-read")
         progress.update()
         url = "https://s3.{0}.amazonaws.com/{1}/{2}{3}".format(
-            bucket_location, s3_bucket_name, s3_bucket_path, quote_plus(artifact)
+            region_name, s3_bucket_name, s3_bucket_path, quote_plus(artifact)
         )
         artifacts_map[artifact] = url
     progress.close()