Skip to content

feat(wrapper): add gpu support in containers #254

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Fastiraz wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

feat(wrapper): add gpu support in containers #254

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions exegol/utils/DockerUtils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,11 @@ def createContainer(self, model: ExegolContainerTemplate, temporary: bool = Fals
"tty": model.config.tty,
"mounts": model.config.getVolumes(),
"working_dir": model.config.getWorkingDir()}
gpu_flag = self.isGPUAvailable()
if gpu_flag:
docker_args["device_requests"] = [
docker.types.DeviceRequest(count=-1 if gpu_flag == "all" else 1,capabilities=[["gpu"]])
]
if temporary:
# Only the 'run' function support the "remove" parameter
docker_create_function = self.__client.containers.run
Expand Down Expand Up @@ -198,6 +203,21 @@ def getContainer(self, tag: str) -> ExegolContainer:
# In this case, ObjectNotFound is raised
raise ObjectNotFound

def isGPUAvailable(self) -> Optional[str]:
"""Check if the GPU is available and return the appropriate device ('cuda', 'mps', or 'cpu').
Return the appropriate value for Docker's --gpus flag."""
try:
import torch, numpy
if os.name == "nt" or os.name == "posix":
if torch.cuda.is_available():
num_gpus = torch.cuda.device_count()
return f'"device=0"' if num_gpus == 1 else '"all"'
if "darwin" in os.uname().sysname.lower() and torch.backends.mps.is_available():
return None
except ImportError:
pass
return None

# # # Volumes Section # # #

def __loadDockerVolume(self, volume_path: str, volume_name: str) -> Volume:
Expand Down
2 changes: 2 additions & 0 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,5 @@ GitPython~=3.1.43
PyYAML>=6.0.2
argcomplete~=3.5.0
tzlocal~=5.2; platform_system != 'Linux'
torch~=2.6.0
numpy~=2.2.3
Comment on lines +8 to +9
Copy link
Copy Markdown

@skyf0l skyf0l Apr 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you check for GPU without a installing a 1Go torch dependency ?

Copy link
Copy Markdown
Author

@Fastiraz Fastiraz Apr 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes but only with some hacks.

Instead of checking properly with a Python module, I can just execute nvidia-smi on a subprocess and check if there is any standard output or standard error.
If there is a standard output, that mean NVIDIA drivers are installed.
Otherwise, if there is a standard error, that mean NVIDIA drivers are not installed.

For macOS, I just have to check to platform and the CPU architecture.
If the platform is macOS and the CPU architecture is ARM, we determine that the "GPU" is MPS.

I've already a version like that and it works well on Linux.
Keep in mind that I cannot test it with AMD GPUs and Linux or Windows operating systems.

If you think a version like this is better, I can make a new commit with the new function.

Here's the new function:

def isGPUAvailable(self) -> Optional[str]:
      import platform, subprocess
      from typing import Optional
      system = platform.system().lower()

      try:
          gpu_names = subprocess.run(
              ['nvidia-smi'],
              stdout=subprocess.PIPE,
              stderr=subprocess.DEVNULL,
              check=True
          ).stdout.decode().strip().split('\n')
          return '"device=0"' if len(gpu_names) == 1 else '"all"'
      except (subprocess.CalledProcessError, FileNotFoundError):
          pass

      if system == "darwin" and platform.machine().startswith("arm"):
          return None

      return None