Merge pull request #288 from runpod/main

justinmerrell · web-flow · commit 89c1a6594d52 · 2024-01-29T12:50:40.000-05:00
update branch
diff --git a/.github/workflows/CD-publish_to_pypi.yml b/.github/workflows/CD-publish_to_pypi.yml
@@ -80,7 +80,7 @@ jobs:
         shell: bash
 
       - name: Repository Dispatch
-        uses: peter-evans/repository-dispatch@v2
+        uses: peter-evans/repository-dispatch@v3
         with:
           token: ${{ secrets.RUNPOD_WORKERS_PAT }}
           repository: ${{ matrix.repo }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,9 +1,11 @@
 # Change Log
 
-## Release 1.5.3 (TBD)
+## Release 1.5.3 (1/25/24)
 
 ### Added
 
+- Expose cuda selection when creating a pod.
+- Expose `flashboot` when creating a new endpoint.
 - Expose bucket name for rp_upload.
 - Exposed `containerRegistryAuthId` for template creation.
 
diff --git a/examples/api/create_endpoint.py b/examples/api/create_endpoint.py
@@ -9,7 +9,7 @@
 
     new_template = runpod.create_template(
         name="test",
-        image_name="runpod/base:0.1.0",
+        image_name="runpod/base:0.4.4",
         is_serverless=True
     )
 
@@ -20,7 +20,8 @@
         template_id=new_template["id"],
         gpu_ids="AMPERE_16",
         workers_min=0,
-        workers_max=1
+        workers_max=1,
+        flashboot=True
     )
 
     print(new_endpoint)
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-aiohttp[speedups] == 3.9.1
+aiohttp[speedups] == 3.9.2
 aiohttp-retry >= 2.8.3
 
 backoff >= 2.2.1
diff --git a/runpod/api/ctl_commands.py b/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"]
@@ -249,7 +249,7 @@ def create_endpoint(
         name: str, template_id: str, gpu_ids: str = "AMPERE_16",
         network_volume_id: str = None, locations: str = None,
         idle_timeout: int = 5, scaler_type: str = "QUEUE_DELAY", scaler_value: int = 4,
-        workers_min: int = 0, workers_max: int = 3
+        workers_min: int = 0, workers_max: int = 3, flashboot=False
 ):
     '''
     Create an endpoint
@@ -274,7 +274,7 @@ def create_endpoint(
             name, template_id, gpu_ids,
             network_volume_id, locations,
             idle_timeout, scaler_type, scaler_value,
-            workers_min, workers_max
+            workers_min, workers_max, flashboot
         )
     )
 
diff --git a/runpod/api/mutations/endpoints.py b/runpod/api/mutations/endpoints.py
@@ -7,12 +7,15 @@ def generate_endpoint_mutation(
     name: str, template_id: str, gpu_ids: str = "AMPERE_16",
     network_volume_id: str = None, locations: str = None,
     idle_timeout: int = 5, scaler_type: str = "QUEUE_DELAY", scaler_value: int = 4,
-    workers_min: int = 0, workers_max: int = 3
+    workers_min: int = 0, workers_max: int = 3, flashboot=False
 ):
     """ Generate a string for a GraphQL mutation to create a new endpoint. """
     input_fields = []
 
     # ------------------------------ Required Fields ----------------------------- #
+    if flashboot:
+        name = name + "-fb"
+
     input_fields.append(f'name: "{name}"')
     input_fields.append(f'templateId: "{template_id}"')
     input_fields.append(f'gpuIds: "{gpu_ids}"')
diff --git a/runpod/api/mutations/pods.py b/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)
 
diff --git a/tests/test_api/test_mutation_endpoints.py b/tests/test_api/test_mutation_endpoints.py
@@ -4,6 +4,7 @@
 
 from runpod.api.mutations.endpoints import generate_endpoint_mutation
 
+
 class TestGenerateEndpointMutation(unittest.TestCase):
     """Tests for the endpoint mutation generation."""
 
@@ -20,9 +21,9 @@ def test_all_fields(self):
         """Test all the fields."""
         result = generate_endpoint_mutation(
             "test_name", "test_template_id", "AMPERE_20",
-            "test_volume_id", "US_WEST", 10, "WORKER_COUNT", 5, 2, 4
+            "test_volume_id", "US_WEST", 10, "WORKER_COUNT", 5, 2, 4, True
         )
-        self.assertIn('name: "test_name"', result)
+        self.assertIn('name: "test_name-fb"', result)
         self.assertIn('templateId: "test_template_id"', result)
         self.assertIn('gpuIds: "AMPERE_20"', result)
         self.assertIn('networkVolumeId: "test_volume_id"', result)
diff --git a/tests/test_api/test_mutations_pods.py b/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()

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-aiohttp[speedups] == 3.9.1`
	`1`	`+aiohttp[speedups] == 3.9.2`
`2`	`2`	`aiohttp-retry >= 2.8.3`
`3`	`3`
`4`	`4`	`backoff >= 2.2.1`