feat: add container_registry_auth_id to create_pod

Author: golithe

runpod/api/ctl_commands.py

@@ -134,6 +134,7 @@ def create_pod(
     min_download = None,
     min_upload = None,
     instance_id: Optional[str] = None,
+    container_registry_auth_id: Optional[str] = None,
 ) -> dict:
     """
     Create a pod
@@ -155,6 +156,7 @@ def create_pod(
     :param min_download: minimum download speed in Mbps
     :param min_upload: minimum upload speed in Mbps
     :param instance_id: the id of a specific instance to deploy to (for CPU pods)
+    :param container_registry_auth_id: the id of the container registry auth to use for pulling private images (GPU pods only)
     :example:
 
     >>> # Create GPU pod
@@ -209,6 +211,7 @@ def create_pod(
             min_download,
             min_upload,
             instance_id,
+            container_registry_auth_id,
         )
     )
 

runpod/api/mutations/pods.py

@@ -31,6 +31,7 @@ def generate_pod_deployment_mutation(
     min_download: Optional[int] = None,
     min_upload: Optional[int] = None,
     instance_id: Optional[str] = None,
+    container_registry_auth_id: Optional[str] = None,
 ) -> str:
     """
     Generates a mutation to deploy a pod on demand.
@@ -93,6 +94,9 @@ def generate_pod_deployment_mutation(
             cuda_versions = ", ".join(f'"{v}"' for v in allowed_cuda_versions)
             input_fields.append(f"allowedCudaVersions: [{cuda_versions}]")
 
+        if container_registry_auth_id is not None:
+            input_fields.append(f'containerRegistryAuthId: "{container_registry_auth_id}"')
+
     # CPU Pod Fields
     else:
         if instance_id is not None:

tests/test_api/test_mutations_pods.py

@@ -15,7 +15,7 @@ def test_generate_pod_deployment_mutation(self):
         # Test GPU pod deployment
         gpu_result = pods.generate_pod_deployment_mutation(
             name="test",
-            image_name="test_image", 
+            image_name="test_image",
             gpu_type_id="1",
             cloud_type="cloud",
             data_center_id="1",
@@ -34,6 +34,15 @@ def test_generate_pod_deployment_mutation(self):
             allowed_cuda_versions=["11.8", "12.0"],
         )
 
+        # Test GPU pod deployment with container registry auth
+        gpu_result_with_auth = pods.generate_pod_deployment_mutation(
+            name="test",
+            image_name="test_image",
+            gpu_type_id="1",
+            cloud_type="ALL",
+            container_registry_auth_id="auth123",
+        )
+
         # Test CPU pod deployment
         cpu_result = pods.generate_pod_deployment_mutation(
             name="test-cpu",
@@ -55,10 +64,16 @@ def test_generate_pod_deployment_mutation(self):
         # Check GPU pod mutation structure
         self.assertIn("mutation", gpu_result)
         self.assertIn("podFindAndDeployOnDemand", gpu_result)
-        
-        # Check CPU pod mutation structure  
+        self.assertNotIn("containerRegistryAuthId", gpu_result)
+
+        # Check containerRegistryAuthId is included when provided, absent when not
+        self.assertIn("containerRegistryAuthId", gpu_result_with_auth)
+        self.assertIn('"auth123"', gpu_result_with_auth)
+
+        # Check CPU pod mutation structure
         self.assertIn("mutation", cpu_result)
         self.assertIn("deployCpuPod", cpu_result)
+        self.assertNotIn("containerRegistryAuthId", cpu_result)
 
     def test_generate_pod_stop_mutation(self):
         """