Merge pull request #282 from runpod/allowed-cuda-pod

feat: add the ability to create pod with selected cuda

Author: justinmerrell

runpod/api/ctl_commands.py

@@ -93,7 +93,7 @@ def create_pod(
     min_vcpu_count: int = 1, min_memory_in_gb: int = 1, docker_args: str = "",
     ports: Optional[str] = None, volume_mount_path: str = "/runpod-volume",
     env: Optional[dict] = None,  template_id: Optional[str] = None,
-    network_volume_id: Optional[str] = None
+    network_volume_id: Optional[str] = None, allowed_cuda_versions: Optional[list] = None
 ) -> dict:
     '''
     Create a pod
@@ -138,7 +138,7 @@ def create_pod(
             cloud_type, support_public_ip, start_ssh,
             data_center_id, country_code, gpu_count,
             volume_in_gb, container_disk_in_gb, min_vcpu_count, min_memory_in_gb, docker_args,
-            ports, volume_mount_path, env, template_id, network_volume_id)
+            ports, volume_mount_path, env, template_id, network_volume_id, allowed_cuda_versions)
     )
 
     cleaned_response = raw_response["data"]["podFindAndDeployOnDemand"]

runpod/api/mutations/pods.py

@@ -3,14 +3,17 @@
 """
 # pylint: disable=too-many-arguments, too-many-locals, too-many-branches
 
+from typing import Optional, List
+
 
 def generate_pod_deployment_mutation(
-        name:str, image_name:str, gpu_type_id:str,
-        cloud_type:str="ALL", support_public_ip:bool=True, start_ssh:bool=True,
+        name: str, image_name: str, gpu_type_id: str,
+        cloud_type: str = "ALL", support_public_ip: bool = True, start_ssh: bool = True,
         data_center_id=None, country_code=None,
         gpu_count=None, volume_in_gb=None, container_disk_in_gb=None, min_vcpu_count=None,
         min_memory_in_gb=None, docker_args=None, ports=None, volume_mount_path=None,
-        env:dict=None, template_id=None, network_volume_id=None):
+        env: dict = None, template_id=None, network_volume_id=None,
+        allowed_cuda_versions: Optional[List[str]] = None):
     '''
     Generates a mutation to deploy a pod on demand.
     '''
@@ -64,6 +67,11 @@ def generate_pod_deployment_mutation(
     if network_volume_id is not None:
         input_fields.append(f'networkVolumeId: "{network_volume_id}"')
 
+    if allowed_cuda_versions is not None:
+        allowed_cuda_versions_string = ", ".join(
+            [f'"{version}"' for version in allowed_cuda_versions])
+        input_fields.append(f'allowedCudaVersions: [{allowed_cuda_versions_string}]')
+
     # Format input fields
     input_string = ", ".join(input_fields)
 

tests/test_api/test_mutations_pods.py

@@ -4,6 +4,7 @@
 
 from runpod.api.mutations import pods
 
+
 class TestPodMutations(unittest.TestCase):
     ''' Test API Wrapper Pod Mutations '''
 
@@ -28,7 +29,9 @@ def test_generate_pod_deployment_mutation(self):
             volume_mount_path="/path",
             env={"ENV": "test"},
             support_public_ip=True,
-            template_id="abcde")
+            template_id="abcde",
+            allowed_cuda_versions=["11.8", "12.0"]
+        )
 
         # Here you should check the correct structure of the result
         self.assertIn("mutation", result)
@@ -57,5 +60,6 @@ def test_generate_pod_terminate_mutation(self):
         # Here you should check the correct structure of the result
         self.assertIn("mutation", result)
 
+
 if __name__ == "__main__":
     unittest.main()