From 41cd34ba4de367d591f09146cfd49cfc12203a4f Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Wed, 18 Sep 2024 14:55:12 -0400
Subject: [PATCH] Allow specifying job names for benchmarks

---
 api-daemon/pvcapid/flaskapi.py     | 14 +++++++++++++-
 client-cli/pvc/cli/cli.py          | 13 +++++++++++--
 client-cli/pvc/lib/storage.py      |  6 ++++--
 daemon-common/benchmark.py         | 11 +++++++----
 worker-daemon/pvcworkerd/Daemon.py |  8 ++++----
 5 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/api-daemon/pvcapid/flaskapi.py b/api-daemon/pvcapid/flaskapi.py
index 9f27b213..b7ae0bf3 100755
--- a/api-daemon/pvcapid/flaskapi.py
+++ b/api-daemon/pvcapid/flaskapi.py
@@ -5069,6 +5069,10 @@ class API_Storage_Ceph_Benchmark(Resource):
                 "required": True,
                 "helptext": "A valid pool must be specified.",
             },
+            {
+                "name": "name",
+                "required": False,
+            },
         ]
     )
     @Authenticator
@@ -5084,6 +5088,11 @@ class API_Storage_Ceph_Benchmark(Resource):
             type: string
             required: true
             description: The PVC storage pool to benchmark
+          - in: query
+            name: name
+            type: string
+            required: false
+            description: An optional override name for the job
         responses:
           200:
             description: OK
@@ -5101,7 +5110,10 @@ class API_Storage_Ceph_Benchmark(Resource):
             }, 400
 
         task = run_celery_task(
-            "storage.benchmark", pool=reqargs.get("pool", None), run_on="primary"
+            "storage.benchmark",
+            pool=reqargs.get("pool", None),
+            name=reqargs.get("name", None),
+            run_on="primary",
         )
         return (
             {
diff --git a/client-cli/pvc/cli/cli.py b/client-cli/pvc/cli/cli.py
index 22b72935..da860817 100644
--- a/client-cli/pvc/cli/cli.py
+++ b/client-cli/pvc/cli/cli.py
@@ -3755,6 +3755,13 @@ def cli_storage_benchmark():
 @click.command(name="run", short_help="Run a storage benchmark.")
 @connection_req
 @click.argument("pool")
+@click.option(
+    "--name",
+    "name",
+    default=None,
+    show_default=False,
+    help="Use a custom name for the job",
+)
 @click.option(
     "--wait/--no-wait",
     "wait_flag",
@@ -3766,12 +3773,14 @@ def cli_storage_benchmark():
 @confirm_opt(
     "Storage benchmarks take approximately 10 minutes to run and generate significant load on the cluster; they should be run sparingly. Continue"
 )
-def cli_storage_benchmark_run(pool, wait_flag):
+def cli_storage_benchmark_run(pool, name, wait_flag):
     """
     Run a storage benchmark on POOL in the background.
     """
 
-    retcode, retmsg = pvc.lib.storage.ceph_benchmark_run(CLI_CONFIG, pool, wait_flag)
+    retcode, retmsg = pvc.lib.storage.ceph_benchmark_run(
+        CLI_CONFIG, pool, name, wait_flag
+    )
 
     if retcode and wait_flag:
         retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
diff --git a/client-cli/pvc/lib/storage.py b/client-cli/pvc/lib/storage.py
index 1dbca65c..2d3b0be4 100644
--- a/client-cli/pvc/lib/storage.py
+++ b/client-cli/pvc/lib/storage.py
@@ -1725,15 +1725,17 @@ def format_list_snapshot(config, snapshot_list):
 #
 # Benchmark functions
 #
-def ceph_benchmark_run(config, pool, wait_flag):
+def ceph_benchmark_run(config, pool, name, wait_flag):
     """
     Run a storage benchmark against {pool}
 
     API endpoint: POST /api/v1/storage/ceph/benchmark
-    API arguments: pool={pool}
+    API arguments: pool={pool}, name={name}
     API schema: {message}
     """
     params = {"pool": pool}
+    if name:
+        params["name"] = name
     response = call_api(config, "post", "/storage/ceph/benchmark", params=params)
 
     return get_wait_retdata(response, wait_flag)
diff --git a/daemon-common/benchmark.py b/daemon-common/benchmark.py
index 1a359334..404f702f 100644
--- a/daemon-common/benchmark.py
+++ b/daemon-common/benchmark.py
@@ -410,11 +410,14 @@ def run_benchmark_job(
     return resource_data, jstdout
 
 
-def worker_run_benchmark(zkhandler, celery, config, pool):
+def worker_run_benchmark(zkhandler, celery, config, pool, name):
     # Phase 0 - connect to databases
-    cur_time = datetime.now().isoformat(timespec="seconds")
-    cur_primary = zkhandler.read("base.config.primary_node")
-    job_name = f"{cur_time}_{cur_primary}"
+    if not name:
+        cur_time = datetime.now().isoformat(timespec="seconds")
+        cur_primary = zkhandler.read("base.config.primary_node")
+        job_name = f"{cur_time}_{cur_primary}"
+    else:
+        job_name = name
 
     current_stage = 0
     total_stages = 13
diff --git a/worker-daemon/pvcworkerd/Daemon.py b/worker-daemon/pvcworkerd/Daemon.py
index 173752f9..00e2a409 100755
--- a/worker-daemon/pvcworkerd/Daemon.py
+++ b/worker-daemon/pvcworkerd/Daemon.py
@@ -96,12 +96,12 @@ def create_vm(
 
 
 @celery.task(name="storage.benchmark", bind=True, routing_key="run_on")
-def storage_benchmark(self, pool=None, run_on="primary"):
+def storage_benchmark(self, pool=None, name=None, run_on="primary"):
     @ZKConnection(config)
-    def run_storage_benchmark(zkhandler, self, pool):
-        return worker_run_benchmark(zkhandler, self, config, pool)
+    def run_storage_benchmark(zkhandler, self, pool, name):
+        return worker_run_benchmark(zkhandler, self, config, pool, name)
 
-    return run_storage_benchmark(self, pool)
+    return run_storage_benchmark(self, pool, name)
 
 
 @celery.task(name="cluster.autobackup", bind=True, routing_key="run_on")