Merge pull request #107 from PrimeIntellect-ai/improvement/shm-size-setting

Improvement/shm size setting

Author: JannikSt

miner/src/cli/command.rs

@@ -225,9 +225,15 @@ pub async fn execute_command(
                 None => None,
             };
 
+            let system_memory = node_config
+                .compute_specs
+                .as_ref()
+                .map(|specs| specs.ram_mb.unwrap_or(0));
+
             let docker_service = Arc::new(DockerService::new(
                 cancellation_token.clone(),
                 has_gpu,
+                system_memory,
                 task_bridge.socket_path.clone(),
                 docker_storage_path,
             ));

miner/src/docker/docker_manager.rs

@@ -95,6 +95,7 @@ impl DockerManager {
         Ok(())
     }
 
+    #[allow(clippy::too_many_arguments)]
     /// Start a new container with the given image and configuration
     pub async fn start_container(
         &self,
@@ -105,6 +106,7 @@ impl DockerManager {
         gpu_enabled: bool,
         // Simple Vec of (host_path, container_path, read_only)
         volumes: Option<Vec<(String, String, bool)>>,
+        shm_size: Option<u64>,
     ) -> Result<String, DockerError> {
         println!("Starting to pull image: {}", image);
 
@@ -178,6 +180,7 @@ impl DockerManager {
                     options: Some(HashMap::new()),
                 }]),
                 binds: volume_binds,
+                shm_size: shm_size.map(|s| s as i64),
                 ..Default::default()
             })
         } else {

miner/src/docker/service.rs

@@ -18,6 +18,7 @@ pub struct DockerService {
     cancellation_token: CancellationToken,
     pub state: Arc<DockerState>,
     has_gpu: bool,
+    system_memory_mb: Option<u32>,
     task_bridge_socket_path: String,
 }
 
@@ -27,6 +28,7 @@ impl DockerService {
     pub fn new(
         cancellation_token: CancellationToken,
         has_gpu: bool,
+        system_memory_mb: Option<u32>,
         task_bridge_socket_path: String,
         storage_path: Option<String>,
     ) -> Self {
@@ -36,6 +38,7 @@ impl DockerService {
             cancellation_token,
             state: Arc::new(DockerState::new()),
             has_gpu,
+            system_memory_mb,
             task_bridge_socket_path,
         }
     }
@@ -142,6 +145,7 @@ impl DockerService {
                                     let manager_clone = manager_clone.clone();
                                     let state_clone = state.clone();
                                     let has_gpu = self.has_gpu;
+                                    let system_memory_mb = self.system_memory_mb;
                                     let task_bridge_socket_path = self.task_bridge_socket_path.clone();
                                     let handle = tokio::spawn(async move {
                                         let payload = task_clone.unwrap();
@@ -170,8 +174,14 @@ impl DockerService {
                                                 false,
                                             )
                                         ];
-
-                                        match manager_clone.start_container(&payload.image, &container_task_id, Some(env_vars), Some(cmd), has_gpu, Some(volumes)).await {
+                                        let shm_size = match system_memory_mb {
+                                            Some(mem_mb) => (mem_mb as u64) * 1024 * 1024 / 2, // Convert MB to bytes and divide by 2
+                                            None => {
+                                                Console::warning("System memory not available, using default shm size");
+                                                67108864 // Default to 64MB in bytes
+                                            }
+                                        };
+                                        match manager_clone.start_container(&payload.image, &container_task_id, Some(env_vars), Some(cmd), has_gpu, Some(volumes), Some(shm_size)).await {
                                             Ok(container_id) => {
                                                 Console::info("DockerService", &format!("Container started with id: {}", container_id));
                                             },
@@ -276,6 +286,7 @@ mod tests {
         let docker_service = DockerService::new(
             cancellation_token.clone(),
             false,
+            Some(1024),
             "/tmp/com.prime.miner/metrics.sock".to_string(),
             None,
         );
@@ -319,6 +330,7 @@ mod tests {
         let docker_service = DockerService::new(
             cancellation_token.clone(),
             false,
+            Some(1024),
             "/tmp/com.prime.miner/metrics.sock".to_string(),
             None,
         );