Add image logging to k6 load script

G-Research · Sep 4, 2024 · 1413d99 · 1413d99
1 parent 62bc122
commit 1413d99
Showing 1 changed file with 46 additions and 19 deletions.
diff --git a/docs/example/k6_load.js b/docs/example/k6_load.js
@@ -1,19 +1,19 @@
 import http from 'k6/http';
 
-const MAX_METRICS_PER_BATCH = 200
+const MAX_METRICS_PER_BATCH = 200;
 
 export default function () {
-  const namespace = 'default'
-  const numberOfExperiments = 1
-  const runsPerExperiment = 2
-  const paramsPerRun = 1
-  const metricsPerRun = 2000
-  const stepsPerMetric = 4
+  const namespace = 'default';
+  const numberOfExperiments = 1;
+  const runsPerExperiment = 2;
+  const paramsPerRun = 1;
+  const metricsPerRun = 2000;
+  const stepsPerMetric = 4;
 
   for (let i = 0; i < numberOfExperiments; i++) {
-    const experimentId = createExperiment(namespace)
+    const experimentId = createExperiment(namespace);
     for (let j = 0; j < runsPerExperiment; j++) {
-      createRun(namespace, experimentId, paramsPerRun, metricsPerRun, stepsPerMetric)
+      createRun(namespace, experimentId, paramsPerRun, metricsPerRun, stepsPerMetric);
     }
   }
 }
@@ -35,7 +35,6 @@ function createExperiment(namespace) {
   return exp_response.json().experiment_id;
 }
 
-
 function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
   const base_url = `http://localhost:5000/ns/${namespace}/api/2.0/mlflow/`;
 
@@ -59,12 +58,12 @@ function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
   );
   const run_id = run_response.json().run.info.run_id;
 
-  let params = []
+  let params = [];
   for (let id = 1; id <= numParams; id++) {
     params.push({
       key: `param${id}`,
       value: `${id * Math.random()}`,
-    })
+    });
   }
   http.post(
     base_url + 'runs/log-batch',
@@ -82,13 +81,12 @@ function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
   let metrics = [];
   for (let step = 1; step <= numSteps; step++) {
     for (let id = 1; id <= numMetrics; id++) {
-      let ctx = {}
-      let rnd = Math.random()
+      let ctx = {};
+      let rnd = Math.random();
       if (rnd < 0.3) {
-        ctx = { type: 'training' }
-      }
-      else if (rnd > 0.6) {
-        ctx = { type: 'testing' }
+        ctx = { type: 'training' };
+      } else if (rnd > 0.6) {
+        ctx = { type: 'testing' };
       }
 
       metrics.push({
@@ -97,7 +95,7 @@ function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
         timestamp: Date.now(),
         step: step,
         context: ctx,
-      })
+      });
 
       if (metrics.length >= MAX_METRICS_PER_BATCH) {
         http.post(
@@ -131,6 +129,8 @@ function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
       );
     }
 
+    logImageArtifact(namespace, run_id);
+
     http.post(
       base_url + 'runs/update',
       JSON.stringify({
@@ -146,3 +146,30 @@ function createRun(namespace, experimentId, numParams, numMetrics, numSteps) {
     );
   }
 }
+
+function logImageArtifact(namespace, runId) {
+  const base_url = `http://localhost:5000/ns/${namespace}/api/2.0/mlflow/`;
+
+  const imageLogRequest = {
+    name: 'example image name',
+    iter: 1,
+    step: 1,
+    caption: 'example image caption',
+    index: 0,
+    width: 1024,
+    run_id: runId,
+    height: 768,
+    format: 'png',
+    blob_uri: 'example_image.png'
+  };
+
+  http.post(
+    base_url + 'runs/log-artifact',
+    JSON.stringify(imageLogRequest),
+    {
+      headers: {
+        'Content-Type': 'application/json'
+      },
+    }
+  );
+}