[ft-benchmark] some fixes for benchmark infra

scripts/ft-benchmark-data-sender.py

@@ -4,15 +4,11 @@
 from time import sleep
 import numpy as np
 import subprocess
-import psycopg2
-from psycopg2 import sql
-from os import getenv, chdir
+from os import chdir
 import os
 import json
 import tempfile
-
-# Duration of experiment in hours
-DURATION = 2
+import argparse
 
 # TPS polling interval in seconds
 POLL_INTERVAL = 30
@@ -89,24 +85,45 @@ def get_commit() -> tuple[str, datetime]:
 
 
 def commit_to_db(data: dict) -> None:
+    print(data)
     chdir(os.path.expanduser("~/nearcore/benchmarks/continous/db/tool"))
-    with tempfile.NamedTemporaryFile() as fp:
+    with tempfile.NamedTemporaryFile(mode="w", encoding='utf-8') as fp:
         json.dump(data, fp)
-        fp.close()
+        fp.flush()
+        os.fsync(fp.fileno())
         subprocess.run(f"cargo run -p cli -- insert-ft-transfer {fp.name}",
                        shell=True)
+        fp.close()
 
 
 # TODO: send signal to this process if ft-benchmark.sh decided to switch neard to another commit.
 # add handling of this signal to this script
 if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        description=
+        "Collect data from local prometheus and send to ft-benchmark db.")
+    parser.add_argument('--duration',
+                        type=int,
+                        required=True,
+                        help='Duration of experiment in seconds')
+    parser.add_argument('--users',
+                        type=int,
+                        required=True,
+                        help='Number of users')
+    args = parser.parse_args()
+    DURATION = args.duration / 3600
+
     state_size = (int(
         subprocess.check_output(["du", "-s", "~/.near/localnet/node0/data"],
                                 stderr=subprocess.PIPE,
                                 shell=True).decode("utf-8").split()[0]) * 1024)
     processed_transactions = []
     time_begin = datetime.now()
     while True:
+        print(
+            f"Data sender loop. Time elapsed: {(datetime.now() - time_begin).seconds} seconds"
+        )
         if (datetime.now() - time_begin).seconds / 3600 > DURATION:
             break
         processed_transactions.append(calculate_processed_transactions())
@@ -131,10 +148,9 @@ def commit_to_db(data: dict) -> None:
         "disjoint_workloads":
             False,  # TODO: probably should be filled by terraform
         "num_shards": calculate_shards(),
-        "num_unique_users":
-            1000,  # TODO: probably should be filled by terraform or ft-benchmark.sh
+        "num_unique_users": args.users,
         "size_state_bytes": state_size,
         "tps": int(average_tps),
-        "total_transactions": processed_transactions[-1],
+        "total_transactions": int(processed_transactions[-1]),
     }
     commit_to_db(response)

scripts/run-ft-benchmark.py

@@ -1,9 +1,10 @@
 import argparse
 import os
 import subprocess
+from locust.util.timespan import parse_timespan
 
 LOCK_FILE = "/tmp/run-ft-benchmark.lock"
-REPO_DIR = "~/nearcore"
+REPO_DIR = os.path.expanduser("~/nearcore")
 
 
 def create_lock_file(user: str) -> None:
@@ -25,7 +26,7 @@ def remove_lock_file() -> None:
 def run_benchmark(repo_dir: str, time: str, users: int, shards: int, nodes: int,
                   rump_up: int) -> None:
     benchmark_command = (
-        f"./scripts/start_benchmark.sh {time} {users} {shards} {nodes} {rump_up}"
+        f"./scripts/start-benchmark.sh {time} {users} {shards} {nodes} {rump_up}"
     )
     subprocess.run(benchmark_command, cwd=repo_dir, shell=True, check=True)
 
@@ -49,13 +50,12 @@ def main() -> None:
     parser.add_argument('--nodes', type=int, default=1, help="Number of nodes")
     parser.add_argument('--rump-up', type=int, default=10, help="Rump-up rate")
     parser.add_argument('--user', type=str, default='unknown', help="User name")
-
     args = parser.parse_args()
-
+    time_seconds = parse_timespan(args.time)
     try:
         create_lock_file(args.user)
-        run_benchmark(REPO_DIR, args.time, args.users, args.shards, args.nodes,
-                      args.rump_up)
+        run_benchmark(REPO_DIR, time_seconds, args.users, args.shards,
+                      args.nodes, args.rump_up)
     except RuntimeError as e:
         print(e)
     finally:

scripts/start-benchmark.sh

@@ -26,13 +26,13 @@ export KEY=~/.near/localnet/node0/validator_key.json
 
 # Run benchmark
 cd pytest/tests/loadtest/locust/
-nohup locust -H 127.0.0.1:3030 -f locustfiles/ft.py --funding-key=$KEY -t $TIME -u $USERS -r $RUMP_UP --processes 8 --headless &
+nohup locust -H 127.0.0.1:3030 -f locustfiles/ft.py --funding-key=$KEY -t "${TIME}s" -u $USERS -r $RUMP_UP --processes 8 --headless &
 
-# Give locust 5 minutes to start and rump up
-sleep 300
+# Give locust 0.5 minutes to start and rump up
+sleep 30
 
 # Run data collector
 cd ~/nearcore
-python3 scripts/ft-benchmark-data-sender.py
+python3 scripts/ft-benchmark-data-sender.py --duration $TIME --users $USERS
 
 echo "Benchmark completed."