From 9bcc4ff3fda4470e791807268950ad3d76383801 Mon Sep 17 00:00:00 2001
From: zhangddjs <1005155946@qq.com>
Date: Thu, 4 Sep 2025 13:53:27 +0000
Subject: [PATCH 1/5] feat: performance improvement config option for batch
 problem solving

---
 aiopslab/config.yml.example              |  8 ++++++++
 aiopslab/orchestrator/orchestrator.py    |  3 ++-
 aiopslab/service/apps/socialnet.py       | 14 +++++++++++---
 aiopslab/service/telemetry/prometheus.py | 10 ++++++++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/aiopslab/config.yml.example b/aiopslab/config.yml.example
index 3da81468..427ee766 100644
--- a/aiopslab/config.yml.example
+++ b/aiopslab/config.yml.example
@@ -13,3 +13,11 @@ qualitative_eval: false
 
 # Flag to enable/disable printing the session
 print_session: false
+
+# Helm deployment settings
+helm:
+  image_pull_policy: IfNotPresent  # Options: Always, IfNotPresent, Never
+  # IfNotPresent enables caching - images are only pulled if not already on the node
+
+# Infrastructure cleanup settings
+keep_infra: true  # Keep OpenEBS and Prometheus running after problem completion
diff --git a/aiopslab/orchestrator/orchestrator.py b/aiopslab/orchestrator/orchestrator.py
index cdf109f1..6900e2bd 100644
--- a/aiopslab/orchestrator/orchestrator.py
+++ b/aiopslab/orchestrator/orchestrator.py
@@ -11,6 +11,7 @@
 from aiopslab.utils.status import *
 from aiopslab.utils.critical_section import CriticalSection
 from aiopslab.service.telemetry.prometheus import Prometheus
+from aiopslab.paths import config
 import time
 import inspect
 import asyncio
@@ -203,7 +204,7 @@ async def start_problem(self, max_steps: int):
         # if not self.session.problem.sys_status_after_recovery():
         self.session.problem.app.cleanup()
         
-        if self.session.problem.namespace != "docker":
+        if self.session.problem.namespace != "docker" and not config.get("keep_infra", False):
             self.prometheus.teardown()
             print("Uninstalling OpenEBS...")
             self.kubectl.exec_command("kubectl delete sc openebs-hostpath openebs-device --ignore-not-found")
diff --git a/aiopslab/service/apps/socialnet.py b/aiopslab/service/apps/socialnet.py
index a6570344..cd4a69d3 100644
--- a/aiopslab/service/apps/socialnet.py
+++ b/aiopslab/service/apps/socialnet.py
@@ -8,6 +8,7 @@
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import SOCIAL_NETWORK_METADATA
+from aiopslab.paths import config
 
 
 class SocialNetwork(Application):
@@ -47,11 +48,11 @@ def deploy(self):
         node_architectures = self.kubectl.get_node_architectures()
         is_arm = any(arch in ["arm64", "aarch64"] for arch in node_architectures)
 
+        if "extra_args" not in self.helm_configs:
+            self.helm_configs["extra_args"] = []
+
         if is_arm:
             # Use the ARM-compatible image for media-frontend
-            if "extra_args" not in self.helm_configs:
-                self.helm_configs["extra_args"] = []
-
             self.helm_configs["extra_args"].append(
                 "--set media-frontend.container.image=jacksonarthurclark/media-frontend"
             )
@@ -59,6 +60,13 @@ def deploy(self):
                 "--set media-frontend.container.imageVersion=latest"
             )
 
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            self.helm_configs["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/telemetry/prometheus.py b/aiopslab/service/telemetry/prometheus.py
index 721e659a..76d0d49b 100644
--- a/aiopslab/service/telemetry/prometheus.py
+++ b/aiopslab/service/telemetry/prometheus.py
@@ -9,6 +9,7 @@
 from aiopslab.service.helm import Helm
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.paths import PROMETHEUS_METADATA, BASE_DIR
+from aiopslab.paths import config
 
 
 class Prometheus:
@@ -78,6 +79,15 @@ def deploy(self):
             if not self._pvc_exists(pvc_name):
                 self._apply_pvc()
 
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            if "extra_args" not in self.helm_configs:
+                self.helm_configs["extra_args"] = []
+            self.helm_configs["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.namespace)
 

From 2259222fe95d616f65ad2649d60ab4c932a9f177 Mon Sep 17 00:00:00 2001
From: zhangddjs <1005155946@qq.com>
Date: Thu, 4 Sep 2025 14:44:23 +0000
Subject: [PATCH 2/5] feat: helm pull policy config for all services

---
 aiopslab/service/apps/astronomy_shop.py        | 11 +++++++++++
 aiopslab/service/apps/flight_ticket.py         | 11 +++++++++++
 aiopslab/service/apps/tidb_cluster_operator.py | 10 ++++++++++
 aiopslab/service/apps/train_ticket.py          | 11 +++++++++++
 4 files changed, 43 insertions(+)

diff --git a/aiopslab/service/apps/astronomy_shop.py b/aiopslab/service/apps/astronomy_shop.py
index c1caeb0e..8a216fd3 100644
--- a/aiopslab/service/apps/astronomy_shop.py
+++ b/aiopslab/service/apps/astronomy_shop.py
@@ -6,6 +6,7 @@
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import ASTRONOMY_SHOP_METADATA
+from aiopslab.paths import config
 
 
 class AstronomyShop(Application):
@@ -28,6 +29,16 @@ def deploy(self):
             "open-telemetry",
             "https://open-telemetry.github.io/opentelemetry-helm-charts",
         )
+        
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            if "extra_args" not in self.helm_configs:
+                self.helm_configs["extra_args"] = []
+            self.helm_configs["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/flight_ticket.py b/aiopslab/service/apps/flight_ticket.py
index 62e06839..b2402e03 100644
--- a/aiopslab/service/apps/flight_ticket.py
+++ b/aiopslab/service/apps/flight_ticket.py
@@ -6,6 +6,7 @@
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import FLIGHT_TICKET_METADATA
+from aiopslab.paths import config
 
 
 class FlightTicket(Application):
@@ -28,6 +29,16 @@ def deploy(self):
             "flight-ticket",
             "https://xlab-uiuc.github.io/flight-ticket",
         )
+        
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            if "extra_args" not in self.helm_configs:
+                self.helm_configs["extra_args"] = []
+            self.helm_configs["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/tidb_cluster_operator.py b/aiopslab/service/apps/tidb_cluster_operator.py
index e133b51d..30d80489 100644
--- a/aiopslab/service/apps/tidb_cluster_operator.py
+++ b/aiopslab/service/apps/tidb_cluster_operator.py
@@ -5,6 +5,7 @@
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TIDB_METADATA
+from aiopslab.paths import config
 
 
 class TiDBCluster(Application):
@@ -52,6 +53,15 @@ def install_tidb_operator(self):
         operator_namespace = self.helm_operator_config.get("namespace", "tidb-admin")
         self.kubectl.create_namespace_if_not_exist(operator_namespace)
 
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            if "extra_args" not in self.helm_operator_config:
+                self.helm_operator_config["extra_args"] = []
+            self.helm_operator_config["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+
         print("Installing TiDB Operator...")
         Helm.install(**self.helm_operator_config)
         Helm.assert_if_deployed(operator_namespace)
diff --git a/aiopslab/service/apps/train_ticket.py b/aiopslab/service/apps/train_ticket.py
index efc1fcb3..bba09a9e 100644
--- a/aiopslab/service/apps/train_ticket.py
+++ b/aiopslab/service/apps/train_ticket.py
@@ -6,6 +6,7 @@
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import TRAIN_TICKET_METADATA
+from aiopslab.paths import config
 
 
 class TrainTicket(Application):
@@ -24,6 +25,16 @@ def load_app_json(self):
     def deploy(self):
         """Deploy the Helm configurations."""
         self.kubectl.create_namespace_if_not_exist(self.namespace)
+        
+        # Apply image pull policy from config if specified
+        helm_config = config.get("helm", {})
+        if helm_config and "image_pull_policy" in helm_config:
+            if "extra_args" not in self.helm_configs:
+                self.helm_configs["extra_args"] = []
+            self.helm_configs["extra_args"].append(
+                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+            )
+        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 

From 2931928e1a6871f45526c7f17da1e304eeeb6517 Mon Sep 17 00:00:00 2001
From: zhangddjs <1005155946@qq.com>
Date: Fri, 5 Sep 2025 01:23:38 +0000
Subject: [PATCH 3/5] feat: batch mode config option to 3x batch problem
 solving speed

---
 aiopslab/config.yml.example                    | 11 ++++-------
 aiopslab/orchestrator/orchestrator.py          |  3 ++-
 aiopslab/service/apps/astronomy_shop.py        |  7 +++----
 aiopslab/service/apps/flight_ticket.py         |  7 +++----
 aiopslab/service/apps/socialnet.py             |  7 +++----
 aiopslab/service/apps/tidb_cluster_operator.py |  7 +++----
 aiopslab/service/apps/train_ticket.py          |  7 +++----
 aiopslab/service/telemetry/prometheus.py       |  7 +++----
 8 files changed, 24 insertions(+), 32 deletions(-)

diff --git a/aiopslab/config.yml.example b/aiopslab/config.yml.example
index 427ee766..98df84c1 100644
--- a/aiopslab/config.yml.example
+++ b/aiopslab/config.yml.example
@@ -14,10 +14,7 @@ qualitative_eval: false
 # Flag to enable/disable printing the session
 print_session: false
 
-# Helm deployment settings
-helm:
-  image_pull_policy: IfNotPresent  # Options: Always, IfNotPresent, Never
-  # IfNotPresent enables caching - images are only pulled if not already on the node
-
-# Infrastructure cleanup settings
-keep_infra: true  # Keep OpenEBS and Prometheus running after problem completion
+# Batch mode settings
+batch_mode: false  # Optimizes for running multiple problems sequentially:
+# - Caches Docker images locally to avoid re-pulling
+# - Keeps infrastructure (OpenEBS, Prometheus) running between problems
diff --git a/aiopslab/orchestrator/orchestrator.py b/aiopslab/orchestrator/orchestrator.py
index 6900e2bd..8acbc857 100644
--- a/aiopslab/orchestrator/orchestrator.py
+++ b/aiopslab/orchestrator/orchestrator.py
@@ -204,7 +204,8 @@ async def start_problem(self, max_steps: int):
         # if not self.session.problem.sys_status_after_recovery():
         self.session.problem.app.cleanup()
         
-        if self.session.problem.namespace != "docker" and not config.get("keep_infra", False):
+        # In batch mode, keep infrastructure running between problems
+        if self.session.problem.namespace != "docker" and not config.get("batch_mode", False):
             self.prometheus.teardown()
             print("Uninstalling OpenEBS...")
             self.kubectl.exec_command("kubectl delete sc openebs-hostpath openebs-device --ignore-not-found")
diff --git a/aiopslab/service/apps/astronomy_shop.py b/aiopslab/service/apps/astronomy_shop.py
index 8a216fd3..e367bedf 100644
--- a/aiopslab/service/apps/astronomy_shop.py
+++ b/aiopslab/service/apps/astronomy_shop.py
@@ -30,13 +30,12 @@ def deploy(self):
             "https://open-telemetry.github.io/opentelemetry-helm-charts",
         )
         
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             if "extra_args" not in self.helm_configs:
                 self.helm_configs["extra_args"] = []
             self.helm_configs["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
         
         Helm.install(**self.helm_configs)
diff --git a/aiopslab/service/apps/flight_ticket.py b/aiopslab/service/apps/flight_ticket.py
index b2402e03..10e83b35 100644
--- a/aiopslab/service/apps/flight_ticket.py
+++ b/aiopslab/service/apps/flight_ticket.py
@@ -30,13 +30,12 @@ def deploy(self):
             "https://xlab-uiuc.github.io/flight-ticket",
         )
         
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             if "extra_args" not in self.helm_configs:
                 self.helm_configs["extra_args"] = []
             self.helm_configs["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
         
         Helm.install(**self.helm_configs)
diff --git a/aiopslab/service/apps/socialnet.py b/aiopslab/service/apps/socialnet.py
index cd4a69d3..58fba827 100644
--- a/aiopslab/service/apps/socialnet.py
+++ b/aiopslab/service/apps/socialnet.py
@@ -60,11 +60,10 @@ def deploy(self):
                 "--set media-frontend.container.imageVersion=latest"
             )
 
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             self.helm_configs["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
 
         Helm.install(**self.helm_configs)
diff --git a/aiopslab/service/apps/tidb_cluster_operator.py b/aiopslab/service/apps/tidb_cluster_operator.py
index 30d80489..5d4bf818 100644
--- a/aiopslab/service/apps/tidb_cluster_operator.py
+++ b/aiopslab/service/apps/tidb_cluster_operator.py
@@ -53,13 +53,12 @@ def install_tidb_operator(self):
         operator_namespace = self.helm_operator_config.get("namespace", "tidb-admin")
         self.kubectl.create_namespace_if_not_exist(operator_namespace)
 
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             if "extra_args" not in self.helm_operator_config:
                 self.helm_operator_config["extra_args"] = []
             self.helm_operator_config["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
 
         print("Installing TiDB Operator...")
diff --git a/aiopslab/service/apps/train_ticket.py b/aiopslab/service/apps/train_ticket.py
index bba09a9e..c6361767 100644
--- a/aiopslab/service/apps/train_ticket.py
+++ b/aiopslab/service/apps/train_ticket.py
@@ -26,13 +26,12 @@ def deploy(self):
         """Deploy the Helm configurations."""
         self.kubectl.create_namespace_if_not_exist(self.namespace)
         
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             if "extra_args" not in self.helm_configs:
                 self.helm_configs["extra_args"] = []
             self.helm_configs["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
         
         Helm.install(**self.helm_configs)
diff --git a/aiopslab/service/telemetry/prometheus.py b/aiopslab/service/telemetry/prometheus.py
index 76d0d49b..9465aea9 100644
--- a/aiopslab/service/telemetry/prometheus.py
+++ b/aiopslab/service/telemetry/prometheus.py
@@ -79,13 +79,12 @@ def deploy(self):
             if not self._pvc_exists(pvc_name):
                 self._apply_pvc()
 
-        # Apply image pull policy from config if specified
-        helm_config = config.get("helm", {})
-        if helm_config and "image_pull_policy" in helm_config:
+        # Apply batch mode optimizations if enabled
+        if config.get("batch_mode", False):
             if "extra_args" not in self.helm_configs:
                 self.helm_configs["extra_args"] = []
             self.helm_configs["extra_args"].append(
-                f"--set global.imagePullPolicy={helm_config['image_pull_policy']}"
+                "--set global.imagePullPolicy=IfNotPresent"
             )
 
         Helm.install(**self.helm_configs)

From e6bbf9f353b67eec7f448c7a5140e147de354704 Mon Sep 17 00:00:00 2001
From: zhangddjs <1005155946@qq.com>
Date: Fri, 5 Sep 2025 02:05:35 +0000
Subject: [PATCH 4/5] improve: move batch mode to a central place helm.py

---
 aiopslab/service/apps/astronomy_shop.py        | 8 --------
 aiopslab/service/apps/flight_ticket.py         | 8 --------
 aiopslab/service/apps/socialnet.py             | 6 ------
 aiopslab/service/apps/tidb_cluster_operator.py | 8 --------
 aiopslab/service/apps/train_ticket.py          | 8 --------
 aiopslab/service/helm.py                       | 7 ++++++-
 aiopslab/service/telemetry/prometheus.py       | 8 --------
 7 files changed, 6 insertions(+), 47 deletions(-)

diff --git a/aiopslab/service/apps/astronomy_shop.py b/aiopslab/service/apps/astronomy_shop.py
index e367bedf..ed2bfa51 100644
--- a/aiopslab/service/apps/astronomy_shop.py
+++ b/aiopslab/service/apps/astronomy_shop.py
@@ -30,14 +30,6 @@ def deploy(self):
             "https://open-telemetry.github.io/opentelemetry-helm-charts",
         )
         
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            if "extra_args" not in self.helm_configs:
-                self.helm_configs["extra_args"] = []
-            self.helm_configs["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/flight_ticket.py b/aiopslab/service/apps/flight_ticket.py
index 10e83b35..e30fde14 100644
--- a/aiopslab/service/apps/flight_ticket.py
+++ b/aiopslab/service/apps/flight_ticket.py
@@ -30,14 +30,6 @@ def deploy(self):
             "https://xlab-uiuc.github.io/flight-ticket",
         )
         
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            if "extra_args" not in self.helm_configs:
-                self.helm_configs["extra_args"] = []
-            self.helm_configs["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/socialnet.py b/aiopslab/service/apps/socialnet.py
index 58fba827..da702347 100644
--- a/aiopslab/service/apps/socialnet.py
+++ b/aiopslab/service/apps/socialnet.py
@@ -60,12 +60,6 @@ def deploy(self):
                 "--set media-frontend.container.imageVersion=latest"
             )
 
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            self.helm_configs["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/tidb_cluster_operator.py b/aiopslab/service/apps/tidb_cluster_operator.py
index 5d4bf818..2abdba12 100644
--- a/aiopslab/service/apps/tidb_cluster_operator.py
+++ b/aiopslab/service/apps/tidb_cluster_operator.py
@@ -53,14 +53,6 @@ def install_tidb_operator(self):
         operator_namespace = self.helm_operator_config.get("namespace", "tidb-admin")
         self.kubectl.create_namespace_if_not_exist(operator_namespace)
 
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            if "extra_args" not in self.helm_operator_config:
-                self.helm_operator_config["extra_args"] = []
-            self.helm_operator_config["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-
         print("Installing TiDB Operator...")
         Helm.install(**self.helm_operator_config)
         Helm.assert_if_deployed(operator_namespace)
diff --git a/aiopslab/service/apps/train_ticket.py b/aiopslab/service/apps/train_ticket.py
index c6361767..45916673 100644
--- a/aiopslab/service/apps/train_ticket.py
+++ b/aiopslab/service/apps/train_ticket.py
@@ -26,14 +26,6 @@ def deploy(self):
         """Deploy the Helm configurations."""
         self.kubectl.create_namespace_if_not_exist(self.namespace)
         
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            if "extra_args" not in self.helm_configs:
-                self.helm_configs["extra_args"] = []
-            self.helm_configs["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/helm.py b/aiopslab/service/helm.py
index 2cc2c8f6..266838cb 100644
--- a/aiopslab/service/helm.py
+++ b/aiopslab/service/helm.py
@@ -27,8 +27,13 @@ def install(**args):
         chart_path = args.get("chart_path")
         namespace = args.get("namespace")
         version = args.get("version")
-        extra_args = args.get("extra_args")
+        extra_args = args.get("extra_args", [])
         remote_chart = args.get("remote_chart", False)
+        
+        # Apply batch mode if enabled
+        if config.get("batch_mode", False):
+            extra_args = extra_args or []
+            extra_args.append("--set global.imagePullPolicy=IfNotPresent")
 
         if not remote_chart:
             # Install dependencies for chart before installation
diff --git a/aiopslab/service/telemetry/prometheus.py b/aiopslab/service/telemetry/prometheus.py
index 9465aea9..09e0748d 100644
--- a/aiopslab/service/telemetry/prometheus.py
+++ b/aiopslab/service/telemetry/prometheus.py
@@ -79,14 +79,6 @@ def deploy(self):
             if not self._pvc_exists(pvc_name):
                 self._apply_pvc()
 
-        # Apply batch mode optimizations if enabled
-        if config.get("batch_mode", False):
-            if "extra_args" not in self.helm_configs:
-                self.helm_configs["extra_args"] = []
-            self.helm_configs["extra_args"].append(
-                "--set global.imagePullPolicy=IfNotPresent"
-            )
-
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.namespace)
 

From 0704650e10c21827bb363e28b4fb4f6a90b836a9 Mon Sep 17 00:00:00 2001
From: zhangddjs <1005155946@qq.com>
Date: Fri, 5 Sep 2025 02:00:38 +0000
Subject: [PATCH 5/5] cleanup: do some cleaning for the PR

cleanup

cleanup

cleanup

cleanup
---
 aiopslab/orchestrator/orchestrator.py          | 1 -
 aiopslab/service/apps/astronomy_shop.py        | 2 --
 aiopslab/service/apps/flight_ticket.py         | 3 ---
 aiopslab/service/apps/socialnet.py             | 7 +++----
 aiopslab/service/apps/tidb_cluster_operator.py | 1 -
 aiopslab/service/apps/train_ticket.py          | 3 ---
 aiopslab/service/helm.py                       | 2 +-
 aiopslab/service/telemetry/prometheus.py       | 2 --
 8 files changed, 4 insertions(+), 17 deletions(-)

diff --git a/aiopslab/orchestrator/orchestrator.py b/aiopslab/orchestrator/orchestrator.py
index 8acbc857..5c54e1c1 100644
--- a/aiopslab/orchestrator/orchestrator.py
+++ b/aiopslab/orchestrator/orchestrator.py
@@ -3,7 +3,6 @@
 
 """Orchestrator class that interfaces with the agent and the environment."""
 
-from aiopslab.service.helm import Helm
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.session import Session
 from aiopslab.orchestrator.problems.registry import ProblemRegistry
diff --git a/aiopslab/service/apps/astronomy_shop.py b/aiopslab/service/apps/astronomy_shop.py
index ed2bfa51..c1caeb0e 100644
--- a/aiopslab/service/apps/astronomy_shop.py
+++ b/aiopslab/service/apps/astronomy_shop.py
@@ -6,7 +6,6 @@
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import ASTRONOMY_SHOP_METADATA
-from aiopslab.paths import config
 
 
 class AstronomyShop(Application):
@@ -29,7 +28,6 @@ def deploy(self):
             "open-telemetry",
             "https://open-telemetry.github.io/opentelemetry-helm-charts",
         )
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/flight_ticket.py b/aiopslab/service/apps/flight_ticket.py
index e30fde14..3392ad29 100644
--- a/aiopslab/service/apps/flight_ticket.py
+++ b/aiopslab/service/apps/flight_ticket.py
@@ -4,9 +4,7 @@
 from aiopslab.service.helm import Helm
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
-from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import FLIGHT_TICKET_METADATA
-from aiopslab.paths import config
 
 
 class FlightTicket(Application):
@@ -29,7 +27,6 @@ def deploy(self):
             "flight-ticket",
             "https://xlab-uiuc.github.io/flight-ticket",
         )
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/apps/socialnet.py b/aiopslab/service/apps/socialnet.py
index da702347..a6570344 100644
--- a/aiopslab/service/apps/socialnet.py
+++ b/aiopslab/service/apps/socialnet.py
@@ -8,7 +8,6 @@
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import SOCIAL_NETWORK_METADATA
-from aiopslab.paths import config
 
 
 class SocialNetwork(Application):
@@ -48,11 +47,11 @@ def deploy(self):
         node_architectures = self.kubectl.get_node_architectures()
         is_arm = any(arch in ["arm64", "aarch64"] for arch in node_architectures)
 
-        if "extra_args" not in self.helm_configs:
-            self.helm_configs["extra_args"] = []
-
         if is_arm:
             # Use the ARM-compatible image for media-frontend
+            if "extra_args" not in self.helm_configs:
+                self.helm_configs["extra_args"] = []
+
             self.helm_configs["extra_args"].append(
                 "--set media-frontend.container.image=jacksonarthurclark/media-frontend"
             )
diff --git a/aiopslab/service/apps/tidb_cluster_operator.py b/aiopslab/service/apps/tidb_cluster_operator.py
index 2abdba12..e133b51d 100644
--- a/aiopslab/service/apps/tidb_cluster_operator.py
+++ b/aiopslab/service/apps/tidb_cluster_operator.py
@@ -5,7 +5,6 @@
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
 from aiopslab.paths import TIDB_METADATA
-from aiopslab.paths import config
 
 
 class TiDBCluster(Application):
diff --git a/aiopslab/service/apps/train_ticket.py b/aiopslab/service/apps/train_ticket.py
index 45916673..d1247799 100644
--- a/aiopslab/service/apps/train_ticket.py
+++ b/aiopslab/service/apps/train_ticket.py
@@ -4,9 +4,7 @@
 from aiopslab.service.helm import Helm
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.service.apps.base import Application
-from aiopslab.paths import TARGET_MICROSERVICES
 from aiopslab.paths import TRAIN_TICKET_METADATA
-from aiopslab.paths import config
 
 
 class TrainTicket(Application):
@@ -25,7 +23,6 @@ def load_app_json(self):
     def deploy(self):
         """Deploy the Helm configurations."""
         self.kubectl.create_namespace_if_not_exist(self.namespace)
-        
         Helm.install(**self.helm_configs)
         Helm.assert_if_deployed(self.helm_configs["namespace"])
 
diff --git a/aiopslab/service/helm.py b/aiopslab/service/helm.py
index 266838cb..5574e17b 100644
--- a/aiopslab/service/helm.py
+++ b/aiopslab/service/helm.py
@@ -4,9 +4,9 @@
 """Interface for helm operations"""
 
 import subprocess
-import time
 
 from aiopslab.service.kubectl import KubeCtl
+from aiopslab.paths import config
 
 
 class Helm:
diff --git a/aiopslab/service/telemetry/prometheus.py b/aiopslab/service/telemetry/prometheus.py
index 09e0748d..d8955116 100644
--- a/aiopslab/service/telemetry/prometheus.py
+++ b/aiopslab/service/telemetry/prometheus.py
@@ -4,12 +4,10 @@
 import os
 import json
 import yaml
-import time
 from subprocess import CalledProcessError
 from aiopslab.service.helm import Helm
 from aiopslab.service.kubectl import KubeCtl
 from aiopslab.paths import PROMETHEUS_METADATA, BASE_DIR
-from aiopslab.paths import config
 
 
 class Prometheus: