diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 97265ad2..17906c3c 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -138,6 +138,7 @@ class QlibDockerConf(DockerConf): default_entry: str = "qrun conf.yaml" extra_volumes: dict = {Path("~/.qlib/").expanduser().resolve(): "/root/.qlib/"} shm_size: str | None = "16g" + enable_gpu: bool = True class DockerEnv(Env[DockerConf]): @@ -161,6 +162,22 @@ class DockerEnv(Env[DockerConf]): except docker.errors.APIError as e: raise RuntimeError(f"Error while pulling the image: {e}") + def _gpu_kwargs(self, client): + """get gpu kwargs based on its availability""" + if not self.conf.enable_gpu: + return {} + gpu_kwargs = { + "device_requests": [ + docker.types.DeviceRequest(count=-1, capabilities=[['gpu']]) + ] if self.conf.enable_gpu else None, + } + try: + client.containers.run(self.conf.image, "nvidia-smi", **gpu_kwargs) + logger.info("GPU Devices are available.") + except docker.errors.APIError as e: + return {} + return gpu_kwargs + def run(self, entry: str | None = None, local_path: str | None = None, env: dict | None = None): if env is None: env = {} @@ -177,6 +194,7 @@ class DockerEnv(Env[DockerConf]): volumns[lp] = {"bind": rp, "mode": "rw"} log_output = "" + try: container: docker.models.containers.Container = client.containers.run( image=self.conf.image, @@ -188,6 +206,7 @@ class DockerEnv(Env[DockerConf]): # auto_remove=True, # remove too fast might cause the logs not to be get network=self.conf.network, shm_size=self.conf.shm_size, + **self._gpu_kwargs(client) ) logs = container.logs(stream=True) for log in logs: diff --git a/test/utils/test_env.py b/test/utils/test_env.py index 79c7a901..3572a5ee 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -43,6 +43,7 @@ class EnvUtils(unittest.TestCase): """ qtde = QTDockerEnv() qtde.prepare() # you can prepare for multiple times. It is expected to handle it correctly + # qtde.run("nvidia-smi") # NOTE: you can check your GPU with this command # the stdout are returned as result result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry="qrun conf.yaml")