From 4cbb401d1298f962a944ca4dd22ec48fa95de0be Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Tue, 22 Apr 2025 11:52:45 -0300
Subject: [PATCH] feat: add optional data removal with adjusted quotas

---
 benchmarks/codex/config.py                     | 2 ++
 config/codex/experiments.k8s.yaml              | 3 +++
 config/codex/experiments.local.yaml            | 1 +
 k8s/charts/codex/templates/_helpers.tpl        | 9 +++++++--
 k8s/charts/codex/templates/testrunner-job.yaml | 2 ++
 k8s/charts/codex/values.yaml                   | 4 ++++
 6 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/benchmarks/codex/config.py b/benchmarks/codex/config.py
index f856659..3f02233 100644
--- a/benchmarks/codex/config.py
+++ b/benchmarks/codex/config.py
@@ -74,6 +74,7 @@ class CodexExperimentConfig(
     )
 
     download_metric_unit_bytes: int = 1
+    remove_data: bool = False
 
     def build(self) -> CodexDisseminationExperiment:
         node_specs = (
@@ -90,6 +91,7 @@ class CodexExperimentConfig(
             CodexNode(
                 codex_api_url=parse_url(f"http://{str(node.address)}:{node.api_port}"),
                 agent=agents[i],
+                remove_data=self.remove_data,
             )
             for i, node in enumerate(node_specs)
         ]
diff --git a/config/codex/experiments.k8s.yaml b/config/codex/experiments.k8s.yaml
index a26a8f9..e4c32d2 100644
--- a/config/codex/experiments.k8s.yaml
+++ b/config/codex/experiments.k8s.yaml
@@ -4,6 +4,9 @@ codex_experiment:
   seeders: ${SEEDERS}
   file_size: ${FILE_SIZE}
   repetitions: ${REPETITIONS}
+  # Should we delete the data at the end of each experiment (slower, uses less space), or we just
+  # leave it there (faster, uses more space)?
+  remove_data: ${REMOVE_DATA}
   # No need for cooldown as Codex takes forever to remove files, so there's plenty of time to log stuff. :-)
   logging_cooldown: 0
 
diff --git a/config/codex/experiments.local.yaml b/config/codex/experiments.local.yaml
index f8d2d27..07eb16d 100644
--- a/config/codex/experiments.local.yaml
+++ b/config/codex/experiments.local.yaml
@@ -2,6 +2,7 @@ codex_experiment:
   seeders: 1
   file_size: 52428800
   repetitions: 3
+  remove_data: true
 
   nodes:
     - name: codex-1
diff --git a/k8s/charts/codex/templates/_helpers.tpl b/k8s/charts/codex/templates/_helpers.tpl
index b1f6446..e5323b9 100644
--- a/k8s/charts/codex/templates/_helpers.tpl
+++ b/k8s/charts/codex/templates/_helpers.tpl
@@ -9,8 +9,13 @@ Expand the name of the chart.
 {{- mul $sizeNum (index $size $sizeUnit) -}}
 {{- end -}}
 
-{{- define "codex.quota" }}
-{{- div (mul (include "filesize.bytes" .) 13) 10 -}}
+{{- define "experiment.count" -}}
+{{- mul .Values.experiment.seederSets .Values.experiment.repetitions -}}
+{{- end -}}
+
+{{- define "codex.quota" -}}
+{{- $mulFactor := .Values.experiment.removeData | ternary 1 (include "experiment.count" .) -}}
+{{- div (mul (mul (include "filesize.bytes" .) $mulFactor) 13) 10 -}}
 {{- end -}}
 
 {{- define "experiment.groupId" -}}
diff --git a/k8s/charts/codex/templates/testrunner-job.yaml b/k8s/charts/codex/templates/testrunner-job.yaml
index 8bc1d1a..2ceb34d 100644
--- a/k8s/charts/codex/templates/testrunner-job.yaml
+++ b/k8s/charts/codex/templates/testrunner-job.yaml
@@ -36,6 +36,8 @@ spec:
               value: {{ .Values.experiment.repetitions | quote }}
             - name: SEEDER_SETS
               value: {{ .Values.experiment.seederSets | quote }}
+            - name: REMOVE_DATA
+              value: {{ .Values.experiment.removeData | quote }}
             - name: FILE_SIZE
               value: {{ include "filesize.bytes" . | quote }}
             - name: CODEX_STATEFULSET
diff --git a/k8s/charts/codex/values.yaml b/k8s/charts/codex/values.yaml
index 19f4a2d..d24a70f 100644
--- a/k8s/charts/codex/values.yaml
+++ b/k8s/charts/codex/values.yaml
@@ -21,6 +21,10 @@ experiment:
   # If set to false, does not deploy a test runner (useful if you just want the network).
   testRunner: true
 
+  # If set to false, does not delete the data at the end of each experiment. This is faster
+  # than deleting it, but requires more space.
+  removeData: true
+
 deployment:
   appName: ""