From bf08cb88243a19b98521c31ad99e3b3c8a518c72 Mon Sep 17 00:00:00 2001 From: Adonai Vera Date: Sat, 8 Mar 2025 00:43:25 -0500 Subject: [PATCH] Added cpu support for pod creation from python sdk --- README.md | 5 +- runpod/api/ctl_commands.py | 26 +++-- runpod/api/mutations/pods.py | 141 ++++++++++++++++---------- tests/test_api/test_mutations_pods.py | 34 ++++++- 4 files changed, 138 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 5c6e92ea..fcbb8c52 100644 --- a/README.md +++ b/README.md @@ -137,9 +137,12 @@ pods = runpod.get_pods() # Get a specific pod pod = runpod.get_pod(pod.id) -# Create a pod +# Create a pod with GPU pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") +# Create a pod with CPU +pod = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4") + # Stop the pod runpod.stop_pod(pod.id) diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index 57bcecd3..bcf4e69e 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -89,7 +89,7 @@ def get_pod(pod_id: str): def create_pod( name: str, image_name: str, - gpu_type_id: str, + gpu_type_id: Optional[str] = None, cloud_type: str = "ALL", support_public_ip: bool = True, start_ssh: bool = True, @@ -109,17 +109,18 @@ def create_pod( allowed_cuda_versions: Optional[list] = None, min_download = None, min_upload = None, + instance_id: Optional[str] = None, ) -> dict: """ Create a pod :param name: the name of the pod :param image_name: the name of the docker image to be used by the pod - :param gpu_type_id: the gpu type wanted by the pod (retrievable by get_gpus) + :param gpu_type_id: the gpu type wanted by the pod (retrievable by get_gpus). If None, creates a CPU-only pod :param cloud_type: if secure cloud, community cloud or all is wanted :param data_center_id: the id of the data center :param country_code: the code for country to start the pod in - :param gpu_count: how many gpus should be attached to the pod + :param gpu_count: how many gpus should be attached to the pod (ignored for CPU-only pods) :param volume_in_gb: how big should the pod volume be :param ports: the ports to open in the pod, example format - "8888/http,666/tcp" :param volume_mount_path: where to mount the volume? @@ -129,12 +130,19 @@ def create_pod( :param template_id: the id of the template to use for the 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) :example: + >>> # Create GPU pod >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") + >>> # Create CPU pod + >>> pod_id = runpod.create_pod("test", "runpod/stack") + >>> # Create CPU pod on specific instance + >>> pod_id = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4") """ # Input Validation - get_gpu(gpu_type_id) # Check if GPU exists, will raise ValueError if not. + if gpu_type_id is not None: + get_gpu(gpu_type_id) # Check if GPU exists, will raise ValueError if not. if cloud_type not in ["ALL", "COMMUNITY", "SECURE"]: raise ValueError("cloud_type must be one of ALL, COMMUNITY or SECURE") @@ -158,7 +166,7 @@ def create_pod( start_ssh, data_center_id, country_code, - gpu_count, + gpu_count if gpu_type_id is not None else None, volume_in_gb, container_disk_in_gb, min_vcpu_count, @@ -172,10 +180,16 @@ def create_pod( allowed_cuda_versions, min_download, min_upload, + instance_id, ) ) - cleaned_response = raw_response["data"]["podFindAndDeployOnDemand"] + print(f"raw_response: {raw_response}") + + if gpu_type_id is not None: + cleaned_response = raw_response["data"]["podFindAndDeployOnDemand"] + else: + cleaned_response = raw_response["data"]["deployCpuPod"] return cleaned_response diff --git a/runpod/api/mutations/pods.py b/runpod/api/mutations/pods.py index ff21cbf5..07e31b63 100644 --- a/runpod/api/mutations/pods.py +++ b/runpod/api/mutations/pods.py @@ -10,104 +10,133 @@ def generate_pod_deployment_mutation( name: str, image_name: str, - gpu_type_id: str, + gpu_type_id: Optional[str] = None, 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, + data_center_id: Optional[str] = None, + country_code: Optional[str] = None, + gpu_count: Optional[int] = None, + volume_in_gb: Optional[int] = None, + container_disk_in_gb: Optional[int] = None, + min_vcpu_count: Optional[int] = None, + min_memory_in_gb: Optional[int] = None, + docker_args: Optional[str] = None, + ports: Optional[str] = None, + volume_mount_path: Optional[str] = None, + env: Optional[dict] = None, + template_id: Optional[str] = None, + network_volume_id: Optional[str] = None, allowed_cuda_versions: Optional[List[str]] = None, - min_download=None, - min_upload=None, -): + min_download: Optional[int] = None, + min_upload: Optional[int] = None, + instance_id: Optional[str] = None, +) -> str: """ Generates a mutation to deploy a pod on demand. + + Args: + name: Name of the pod + image_name: Docker image name + gpu_type_id: GPU type ID for GPU pods, None for CPU pods + cloud_type: Cloud type (ALL, COMMUNITY, or SECURE) + support_public_ip: Whether to support public IP + start_ssh: Whether to start SSH service + data_center_id: Data center ID + country_code: Country code for pod location + gpu_count: Number of GPUs (for GPU pods) + volume_in_gb: Volume size in GB + container_disk_in_gb: Container disk size in GB + min_vcpu_count: Minimum vCPU count + min_memory_in_gb: Minimum memory in GB + docker_args: Docker arguments + ports: Port mappings (e.g. "8080/tcp,22/tcp") + volume_mount_path: Volume mount path + env: Environment variables dict + template_id: Template ID + network_volume_id: Network volume ID + allowed_cuda_versions: List of allowed CUDA versions + min_download: Minimum download speed in Mbps + min_upload: Minimum upload speed in Mbps + instance_id: Instance ID for CPU pods + + Returns: + str: GraphQL mutation string """ input_fields = [] - # ------------------------------ Required Fields ----------------------------- # - input_fields.append(f'name: "{name}"') - input_fields.append(f'imageName: "{image_name}"') - input_fields.append(f'gpuTypeId: "{gpu_type_id}"') - - # ------------------------------ Default Fields ------------------------------ # - input_fields.append(f"cloudType: {cloud_type}") + # Required Fields + input_fields.extend([ + f'name: "{name}"', + f'imageName: "{image_name}"', + f"cloudType: {cloud_type}" + ]) if start_ssh: input_fields.append("startSsh: true") - if support_public_ip: - input_fields.append("supportPublicIp: true") + # GPU Pod Fields + if gpu_type_id is not None: + input_fields.append(f'gpuTypeId: "{gpu_type_id}"') + input_fields.append(f"supportPublicIp: {str(support_public_ip).lower()}") + + if gpu_count is not None: + input_fields.append(f"gpuCount: {gpu_count}") + if volume_in_gb is not None: + input_fields.append(f"volumeInGb: {volume_in_gb}") + if min_vcpu_count is not None: + input_fields.append(f"minVcpuCount: {min_vcpu_count}") + if min_memory_in_gb is not None: + input_fields.append(f"minMemoryInGb: {min_memory_in_gb}") + if docker_args is not None: + input_fields.append(f'dockerArgs: "{docker_args}"') + if allowed_cuda_versions is not None: + cuda_versions = ", ".join(f'"{v}"' for v in allowed_cuda_versions) + input_fields.append(f"allowedCudaVersions: [{cuda_versions}]") + + # CPU Pod Fields else: - input_fields.append("supportPublicIp: false") + if instance_id is not None: + input_fields.append(f'instanceId: "{instance_id}"') + template_id = template_id or "runpod-ubuntu" - # ------------------------------ Optional Fields ----------------------------- # + # Optional Fields if data_center_id is not None: input_fields.append(f'dataCenterId: "{data_center_id}"') + else: + input_fields.append("dataCenterId: null") + if country_code is not None: input_fields.append(f'countryCode: "{country_code}"') - if gpu_count is not None: - input_fields.append(f"gpuCount: {gpu_count}") - if volume_in_gb is not None: - input_fields.append(f"volumeInGb: {volume_in_gb}") if container_disk_in_gb is not None: input_fields.append(f"containerDiskInGb: {container_disk_in_gb}") - if min_vcpu_count is not None: - input_fields.append(f"minVcpuCount: {min_vcpu_count}") - if min_memory_in_gb is not None: - input_fields.append(f"minMemoryInGb: {min_memory_in_gb}") - if docker_args is not None: - input_fields.append(f'dockerArgs: "{docker_args}"') if ports is not None: - ports = ports.replace(" ", "") - input_fields.append(f'ports: "{ports}"') + input_fields.append(f'ports: "{ports.replace(" ", "")}"') if volume_mount_path is not None: input_fields.append(f'volumeMountPath: "{volume_mount_path}"') if env is not None: - env_string = ", ".join( - [f'{{ key: "{key}", value: "{value}" }}' for key, value in env.items()] - ) - input_fields.append(f"env: [{env_string}]") + env_items = [f'{{ key: "{k}", value: "{v}" }}' for k, v in env.items()] + input_fields.append(f"env: [{', '.join(env_items)}]") if template_id is not None: input_fields.append(f'templateId: "{template_id}"') - 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}]") - if min_download is not None: input_fields.append(f'minDownload: {min_download}') - if min_upload is not None: input_fields.append(f'minUpload: {min_upload}') - # Format input fields + mutation_type = "podFindAndDeployOnDemand" if gpu_type_id else "deployCpuPod" input_string = ", ".join(input_fields) + return f""" mutation {{ - podFindAndDeployOnDemand( + {mutation_type}( input: {{ {input_string} }} ) {{ id - desiredStatus imageName env machineId diff --git a/tests/test_api/test_mutations_pods.py b/tests/test_api/test_mutations_pods.py index 8ebd819e..77f28986 100644 --- a/tests/test_api/test_mutations_pods.py +++ b/tests/test_api/test_mutations_pods.py @@ -10,11 +10,12 @@ class TestPodMutations(unittest.TestCase): def test_generate_pod_deployment_mutation(self): """ - Test generate_pod_deployment_mutation + Test generate_pod_deployment_mutation for both GPU and CPU pods """ - result = pods.generate_pod_deployment_mutation( + # 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", @@ -33,8 +34,31 @@ def test_generate_pod_deployment_mutation(self): allowed_cuda_versions=["11.8", "12.0"], ) - # Here you should check the correct structure of the result - self.assertIn("mutation", result) + # Test CPU pod deployment + cpu_result = pods.generate_pod_deployment_mutation( + name="test-cpu", + image_name="test_image", + cloud_type="cloud", + data_center_id="1", + country_code="US", + volume_in_gb=100, + container_disk_in_gb=10, + min_vcpu_count=2, + min_memory_in_gb=4, + docker_args="args", + ports="8080", + volume_mount_path="/path", + env={"ENV": "test"}, + instance_id="cpu3c-2-4" + ) + + # Check GPU pod mutation structure + self.assertIn("mutation", gpu_result) + self.assertIn("podFindAndDeployOnDemand", gpu_result) + + # Check CPU pod mutation structure + self.assertIn("mutation", cpu_result) + self.assertIn("deployCpuPod", cpu_result) def test_generate_pod_stop_mutation(self): """