Merge pull request #281 from runpod/flash-boot

feat: add flash boot enable to endpoints

Author: justinmerrell

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,
+        flash_boot=True
     )
 
     print(new_endpoint)

runpod/api/ctl_commands.py

@@ -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, flash_boot=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, flash_boot
         )
     )
 

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, flash_boot=False
 ):
     """ Generate a string for a GraphQL mutation to create a new endpoint. """
     input_fields = []
 
     # ------------------------------ Required Fields ----------------------------- #
+    if flash_boot:
+        name = name + "-fb"
+
     input_fields.append(f'name: "{name}"')
     input_fields.append(f'templateId: "{template_id}"')
     input_fields.append(f'gpuIds: "{gpu_ids}"')

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)