refactor: ModelService start logic (alibaba#182)

Author: BCeZn

rock/sdk/sandbox/agent/base.py

@@ -1,11 +1,13 @@
 from abc import ABC, abstractmethod
 
 from rock.actions.sandbox.base import AbstractSandbox
+from rock.sdk.sandbox.model_service.base import ModelService
 
 
 class Agent(ABC):
     def __init__(self, sandbox: AbstractSandbox):
         self._sandbox = sandbox
+        self.model_service: ModelService | None = None
 
     @abstractmethod
     async def init(self):

rock/sdk/sandbox/agent/iflow_cli.py

@@ -469,11 +469,6 @@ async def run(
         logger.debug(f"[{sandbox_id}] Project path: {project_path}, Problem statement: {problem_statement[:100]}...")
 
         try:
-            # Start ModelService if configured
-            if self.model_service and not self.model_service.is_started:
-                logger.info(f"[{sandbox_id}] Starting ModelService")
-                await self.model_service.start()
-
             # Step 1: Change to project directory
             logger.info(f"[{sandbox_id}] Changing working directory to: {project_path}")
             result = await self._sandbox.arun(
@@ -558,13 +553,6 @@ async def run(
                 exc_info=True,
             )
             raise
-        finally:
-            # Clean up ModelService if started
-            if self.model_service and self.model_service.is_started:
-                try:
-                    await self.model_service.stop()
-                except Exception as e:
-                    logger.warning(f"[{sandbox_id}] Failed to stop ModelService: {str(e)}")
 
     async def _agent_run(
         self,
@@ -644,3 +632,9 @@ async def _agent_run(
         except Exception as e:
             error_msg = f"Failed to execute nohup command '{cmd}': {str(e)}"
             return Observation(output=error_msg, exit_code=1, failure_reason=error_msg)
+
+    async def start_model_service(self):
+        if not self.model_service:
+            raise RuntimeError(f"ModelService is not initialized in {self.config.agent_type}!")
+
+        await self.model_service.start()

rock/sdk/sandbox/agent/swe_agent.py

@@ -464,11 +464,6 @@ async def run(
         logger.info(f"[{sandbox_id}] SWE-agent execution started")
 
         try:
-            # Start ModelService if configured
-            if self.model_service:
-                logger.info(f"[{sandbox_id}] Starting ModelService")
-                await self.model_service.start()
-
             with self._config_template_context(problem_statement, project_path, instance_id) as generated_config_path:
                 config_filename = Path(generated_config_path).name
 
@@ -534,14 +529,6 @@ async def run(
                 exc_info=True,
             )
             raise
-        finally:
-            # Clean up ModelService if started
-            if self.model_service and self.model_service.is_started:
-                try:
-                    logger.info(f"[{sandbox_id}] Stopping ModelService")
-                    await self.model_service.stop()
-                except Exception as e:
-                    logger.warning(f"[{sandbox_id}] Failed to stop ModelService: {str(e)}")
 
     async def _agent_run(
         self,
@@ -621,3 +608,9 @@ async def _agent_run(
         except Exception as e:
             error_msg = f"Failed to execute nohup command '{cmd}': {str(e)}"
             return Observation(output=error_msg, exit_code=1, failure_reason=error_msg)
+
+    async def start_model_service(self):
+        if not self.model_service:
+            raise RuntimeError(f"ModelService is not initialized in {self.config.agent_type}!")
+
+        await self.model_service.start()