Support setting how many GPUs to expose in dockerized environment (#258)

Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>

README.md

@@ -30,6 +30,7 @@
   - [Docker-Only Setup](#docker-only-setup)
     - [1. Create Guest Program](#1-create-guest-program)
     - [2. Create Host](#2-create-host)
+- [Environment Variables](#environment-variables)
 - [Directory Layout](#directory-layout)
 - [Contributing](#contributing)
 - [Disclaimer](#disclaimer)
@@ -380,6 +381,28 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 ```
 
+## Environment Variables
+
+| Variable          | Description                                                                                                     | Default |
+| ----------------- | --------------------------------------------------------------------------------------------------------------- | ------- |
+| `ERE_GPU_DEVICES` | Specifies which GPU devices to use when running Docker containers for GPU-enabled zkVMs. The value is passed to Docker's `--gpus` flag. | `all`   |
+
+Example usage:
+
+```bash
+# Use all GPUs (default)
+ere prove ...
+
+# Use specific GPU devices
+ERE_GPU_DEVICES="device=0" ere prove ...
+
+# Use multiple specific GPUs
+ERE_GPU_DEVICES="device=0,1" ere prove ...
+
+# Can also signal to use any available GPUs
+ERE_GPU_DEVICES="4" ere prove ...
+```
+
 ## Directory Layout
 
 ```

crates/dockerized/src/util/docker.rs

@@ -151,8 +151,9 @@ impl DockerRunCmd {
         self.volume(DOCKER_SOCKET, DOCKER_SOCKET)
     }
 
-    pub fn gpus(self, devices: impl AsRef<str>) -> Self {
-        self.option("gpus", devices)
+    pub fn gpus(self) -> Self {
+        let devices = env::var("ERE_GPU_DEVICES").unwrap_or_else(|_| "all".to_string());
+        self.option("gpus", &devices)
     }
 
     pub fn network(self, name: impl AsRef<str>) -> Self {

crates/dockerized/src/zkvm.rs

@@ -174,14 +174,14 @@ impl ServerContainer {
         // zkVM specific options when using GPU
         if gpu {
             cmd = match zkvm_kind {
-                zkVMKind::Airbender => cmd.gpus("all"),
-                zkVMKind::OpenVM => cmd.gpus("all"),
+                zkVMKind::Airbender => cmd.gpus(),
+                zkVMKind::OpenVM => cmd.gpus(),
                 // SP1 runs docker command to spin up the server to do GPU
                 // proving, to give the client access to the prover service, we
                 // need to use the host networking driver.
                 zkVMKind::SP1 => cmd.mount_docker_socket().network("host"),
-                zkVMKind::Risc0 => cmd.gpus("all").inherit_env("RISC0_DEFAULT_PROVER_NUM_GPUS"),
-                zkVMKind::Zisk => cmd.gpus("all"),
+                zkVMKind::Risc0 => cmd.gpus().inherit_env("RISC0_DEFAULT_PROVER_NUM_GPUS"),
+                zkVMKind::Zisk => cmd.gpus(),
                 _ => cmd,
             }
         }