Support LocalEnv

rdagent/utils/env.py

@@ -6,12 +6,13 @@
 
 """
 import os
-from abc import abstractmethod
-from pathlib import Path
-from typing import Generic, TypeVar
-
+import sys
 import docker
+import subprocess
+from abc import abstractmethod
 from pydantic import BaseModel
+from typing import Generic, TypeVar, Optional, Dict
+from pathlib import Path
 
 ASpecificBaseModel = TypeVar("ASpecificBaseModel", bound=BaseModel)
 
@@ -62,15 +63,49 @@ def run(self, entry: str | None, local_path: str | None = None, env: dict | None
 
 
 class LocalConf(BaseModel):
-    py_entry: str  # where you can find your python path
+    py_entry: str
+    default_entry: str
 
 
 class LocalEnv(Env[LocalConf]):
-    """
-    Sometimes local environment may be more convinient for testing
-    """
 
-    conf: LocalConf
+    def prepare(self):
+        qtde = QTDockerEnv()
+        qtde.prepare()
+
+    def run(self,
+            entry: str | None = None,
+            local_path: Optional[str] = None) -> str:
+        # if env is None:
+        #     env = {}
+
+        if entry is None:
+            entry = self.conf.default_entry
+        # 获取当前路径
+        current_path = os.getcwd()
+
+        # 输出当前路径
+        print("当前路径:", current_path)
+        # Use the configured Python interpreter
+        command = str(Path(self.conf.py_entry).joinpath(entry)).split(" ")
+
+        # Set the working directory
+        cwd = None
+        if local_path:
+            cwd = Path(local_path).resolve()
+
+        result = subprocess.run(
+            command,
+            # env={**os.environ, **env},
+            cwd=cwd,
+            capture_output=True,
+            text=True
+        )
+
+        if result.returncode != 0:
+            raise RuntimeError(f"Error while running the command: {result.stderr}")
+
+        return result.stdout
 
 
 ## Docker Environment -----

test/utils/test_env.py

@@ -1,13 +1,11 @@
 import os
-import subprocess
 import sys
 import unittest
 from pathlib import Path
-
 sys.path.append(str(Path(__file__).resolve().parent.parent))
+from rdagent.utils.env import QTDockerEnv, LocalEnv, LocalConf
 import shutil
 
-from rdagent.utils.env import QTDockerEnv
 
 DIRNAME = Path(__file__).absolute().resolve().parent
 
@@ -23,6 +21,19 @@ def tearDown(self):
         #     shutil.rmtree(mlrun_p)
         ...
 
+    # NOTE: Since I don't know the exact environment in which it will be used, here's just an example.
+    # def test_local(self):
+    #     local_conf = LocalConf(
+    #         py_entry="/home/v-linlanglv/miniconda3/envs/RD-Agent-310/bin",
+    #         default_entry="qrun conf.yaml",
+    #     )
+    #     qle = LocalEnv(conf=local_conf)
+    #     qle.prepare()
+    #     conf_path = str(DIRNAME / "env_tpl" / "conf.yaml") 
+    #     qle.run(entry="qrun " + conf_path)
+    #     mlrun_p = DIRNAME / "env_tpl" / "mlruns" 
+    #     self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found")
+
     def test_docker(self):
         """
         We will mount `env_tpl` into the docker image.