diff --git a/README.md b/README.md index 7a99b693..6e6bf143 100644 --- a/README.md +++ b/README.md @@ -142,9 +142,15 @@ pod = runpod.get_pod(pod.id) # Create a pod with GPU pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") +# Create a pod with GPU and encrypted volume +pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070", encrypt_volume=True) + # Create a pod with CPU pod = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4") +# Create a pod with CPU and encrypted volume +pod = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4", encrypt_volume=True) + # Stop the pod runpod.stop_pod(pod.id) diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index adfe7ff8..64bd277d 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -110,6 +110,7 @@ def create_pod( min_download = None, min_upload = None, instance_id: Optional[str] = None, + encrypt_volume: bool = False, ) -> dict: """ Create a pod @@ -131,6 +132,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 encrypt_volume: whether to encrypt the volume :example: >>> # Create GPU pod @@ -185,6 +187,7 @@ def create_pod( min_download, min_upload, instance_id, + encrypt_volume, ) ) diff --git a/runpod/api/mutations/pods.py b/runpod/api/mutations/pods.py index f3be747b..eee1db54 100644 --- a/runpod/api/mutations/pods.py +++ b/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, + encrypt_volume: bool = False, ) -> str: """ Generates a mutation to deploy a pod on demand. @@ -59,6 +60,7 @@ def generate_pod_deployment_mutation( min_download: Minimum download speed in Mbps min_upload: Minimum upload speed in Mbps instance_id: Instance ID for CPU pods + encrypt_volume: Whether to encrypt the volume Returns: str: GraphQL mutation string @@ -125,6 +127,8 @@ def generate_pod_deployment_mutation( input_fields.append(f'minDownload: {min_download}') if min_upload is not None: input_fields.append(f'minUpload: {min_upload}') + if encrypt_volume: + input_fields.append("encryptVolume: true") mutation_type = "podFindAndDeployOnDemand" if gpu_type_id else "deployCpuPod" input_string = ", ".join(input_fields) diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index 63552830..1faba8ea 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -38,8 +38,11 @@ def list_pods(): @click.option( "--support-public-ip", default=True, help="Whether or not to support a public IP." ) +@click.option( + "--encrypt-volume/--no-encrypt-volume", default=False, help="Whether or not to encrypt the volume." +) def create_new_pod( - name, image, gpu_type, gpu_count, support_public_ip + name, image, gpu_type, gpu_count, support_public_ip, encrypt_volume ): # pylint: disable=too-many-arguments """ Creates a pod. @@ -47,6 +50,7 @@ def create_new_pod( kwargs = { "gpu_count": gpu_count, "support_public_ip": support_public_ip, + "encrypt_volume": encrypt_volume, } if not name: diff --git a/runpod/cli/groups/project/functions.py b/runpod/cli/groups/project/functions.py index 4fe01157..d89c7c7b 100644 --- a/runpod/cli/groups/project/functions.py +++ b/runpod/cli/groups/project/functions.py @@ -49,7 +49,8 @@ def _launch_dev_pod(): selected_gpu_types.append(config["project"]["gpu"]) # Attempt to launch a pod with the given configuration - new_pod = attempt_pod_launch(config, environment_variables) + encrypt_volume = config["project"].get("encrypt_volume", False) + new_pod = attempt_pod_launch(config, environment_variables, encrypt_volume) if new_pod is None: print( "Selected GPU types unavailable, try again later or use a different type." @@ -144,6 +145,7 @@ def create_new_project( project_table.add("volume_mount_path", "/runpod-volume") project_table.add("ports", "8080/http, 22/tcp") project_table.add("container_disk_size_gb", 10) + project_table.add("encrypt_volume", False) project_table.add("env_vars", ENV_VARS) toml_config.add("project", project_table) diff --git a/runpod/cli/groups/project/helpers.py b/runpod/cli/groups/project/helpers.py index be7c7d98..791d3a86 100644 --- a/runpod/cli/groups/project/helpers.py +++ b/runpod/cli/groups/project/helpers.py @@ -57,7 +57,7 @@ def copy_template_files(template_dir, destination): shutil.copy2(source_item, destination_item) -def attempt_pod_launch(config, environment_variables): +def attempt_pod_launch(config, environment_variables, encrypt_volume=False): """Attempt to launch a pod with the given configuration.""" for gpu_type in config["project"].get("gpu_types", []): print(f"Trying to get a pod with {gpu_type}... ", end="") @@ -73,6 +73,7 @@ def attempt_pod_launch(config, environment_variables): volume_mount_path=f'{config["project"]["volume_mount_path"]}', container_disk_in_gb=int(config["project"]["container_disk_size_gb"]), env=environment_variables, + encrypt_volume=encrypt_volume, ) print("Success!") return created_pod diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index 2b2e4301..fc3af918 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -136,8 +136,85 @@ def test_create_pod(self): "cloud_type must be one of ALL, COMMUNITY or SECURE", ) + def test_create_pod_with_encrypt_volume(self): + """ + Tests create_pod with encrypt_volume parameter + """ + with patch("runpod.api.graphql.requests.post") as patch_request, patch( + "runpod.api.ctl_commands.get_gpu" + ) as patch_get_gpu, patch("runpod.api.ctl_commands.get_user") as patch_get_user: + patch_request.return_value.json.return_value = { + "data": {"podFindAndDeployOnDemand": {"id": "POD_ID_ENCRYPTED"}} + } + + patch_get_gpu.return_value = None + patch_get_user.return_value = { + "networkVolumes": [ + {"id": "NETWORK_VOLUME_ID", "dataCenterId": "us-east-1"} + ] + } + + # Test with encryption enabled + pod = ctl_commands.create_pod( + name="POD_NAME_ENCRYPTED", + image_name="IMAGE_NAME", + gpu_type_id="NVIDIA A100 80GB PCIe", + network_volume_id="NETWORK_VOLUME_ID", + encrypt_volume=True + ) + + self.assertEqual(pod["id"], "POD_ID_ENCRYPTED") + + # Verify that the GraphQL mutation was called with encryptVolume: true + called_mutation = patch_request.call_args[1]['data'] + self.assertIn("encryptVolume: true", called_mutation) + + # Test with encryption explicitly disabled + patch_request.return_value.json.return_value = { + "data": {"podFindAndDeployOnDemand": {"id": "POD_ID_NO_ENCRYPTION"}} + } + + pod = ctl_commands.create_pod( + name="POD_NAME_NO_ENCRYPTION", + image_name="IMAGE_NAME", + gpu_type_id="NVIDIA A100 80GB PCIe", + network_volume_id="NETWORK_VOLUME_ID", + encrypt_volume=False + ) + + self.assertEqual(pod["id"], "POD_ID_NO_ENCRYPTION") + + # Verify that the GraphQL mutation was not called with encryptVolume + called_mutation = patch_request.call_args[1]['data'] + self.assertNotIn("encryptVolume", called_mutation) + + # Test with default value (should not include encryptVolume) + patch_request.return_value.json.return_value = { + "data": {"podFindAndDeployOnDemand": {"id": "POD_ID_DEFAULT"}} + } + + pod = ctl_commands.create_pod( + name="POD_NAME_DEFAULT", + image_name="IMAGE_NAME", + gpu_type_id="NVIDIA A100 80GB PCIe", + network_volume_id="NETWORK_VOLUME_ID" + ) + + self.assertEqual(pod["id"], "POD_ID_DEFAULT") + + # Verify that the GraphQL mutation was not called with encryptVolume (default) + called_mutation = patch_request.call_args[1]['data'] + self.assertNotIn("encryptVolume", called_mutation) + + def test_create_pod_missing_image_and_template(self): + """ + Tests create_pod error if neither image_name nor template_id are provided + """ + with patch("runpod.api.graphql.requests.post"), patch( + "runpod.api.ctl_commands.get_gpu" + ), patch("runpod.api.ctl_commands.get_user"): with self.assertRaises(ValueError) as context: - pod = ctl_commands.create_pod( + ctl_commands.create_pod( name="POD_NAME", gpu_type_id="NVIDIA A100 80GB PCIe", network_volume_id="NETWORK_VOLUME_ID", @@ -147,7 +224,8 @@ def test_create_pod(self): str(context.exception), "Either image_name or template_id must be provided", ) - + + def test_stop_pod(self): """ Test stop_pod diff --git a/tests/test_api/test_mutations_pods.py b/tests/test_api/test_mutations_pods.py index 77f28986..3da3239c 100644 --- a/tests/test_api/test_mutations_pods.py +++ b/tests/test_api/test_mutations_pods.py @@ -60,6 +60,50 @@ def test_generate_pod_deployment_mutation(self): self.assertIn("mutation", cpu_result) self.assertIn("deployCpuPod", cpu_result) + def test_generate_pod_deployment_mutation_with_encrypt_volume(self): + """ + Test generate_pod_deployment_mutation with encrypt_volume parameter + """ + # Test GPU pod deployment with encryption enabled + gpu_result_encrypted = pods.generate_pod_deployment_mutation( + name="test-encrypted", + image_name="test_image", + gpu_type_id="1", + cloud_type="cloud", + encrypt_volume=True + ) + + # Test CPU pod deployment with encryption enabled + cpu_result_encrypted = pods.generate_pod_deployment_mutation( + name="test-cpu-encrypted", + image_name="test_image", + cloud_type="cloud", + instance_id="cpu3c-2-4", + encrypt_volume=True + ) + + # Test GPU pod deployment with encryption explicitly disabled + gpu_result_no_encryption = pods.generate_pod_deployment_mutation( + name="test-no-encryption", + image_name="test_image", + gpu_type_id="1", + cloud_type="cloud", + encrypt_volume=False + ) + + # Check that encrypted mutations contain encryptVolume: true + self.assertIn("encryptVolume: true", gpu_result_encrypted) + self.assertIn("encryptVolume: true", cpu_result_encrypted) + + # Check that non-encrypted mutations do not contain encryptVolume + self.assertNotIn("encryptVolume", gpu_result_no_encryption) + + # Check basic mutation structure is preserved + self.assertIn("mutation", gpu_result_encrypted) + self.assertIn("podFindAndDeployOnDemand", gpu_result_encrypted) + self.assertIn("mutation", cpu_result_encrypted) + self.assertIn("deployCpuPod", cpu_result_encrypted) + def test_generate_pod_stop_mutation(self): """ Test generate_pod_stop_mutation diff --git a/tests/test_cli/test_cli_groups/test_pod_commands.py b/tests/test_cli/test_cli_groups/test_pod_commands.py index ae594847..a7304f1d 100644 --- a/tests/test_cli/test_cli_groups/test_pod_commands.py +++ b/tests/test_cli/test_cli_groups/test_pod_commands.py @@ -75,9 +75,106 @@ def test_create_new_pod( gpu_count=1, support_public_ip=True, ports="22/tcp", + encrypt_volume=False, # Accept the new default argument ) mock_echo.assert_called_with("Pod sample_id has been created.") + @patch("runpod.cli.groups.pod.commands.click.prompt") + @patch("runpod.cli.groups.pod.commands.click.confirm") + @patch("runpod.cli.groups.pod.commands.click.echo") + @patch("runpod.cli.groups.pod.commands.create_pod") + def test_create_new_pod_with_encrypt_volume( + self, mock_create_pod, mock_echo, mock_confirm, mock_prompt + ): # pylint: disable=too-many-arguments,line-too-long + """ + Test create_new_pod with --encrypt-volume option + """ + # Mock values + mock_confirm.return_value = True # for the quick_launch option + mock_prompt.return_value = "RunPod-CLI-Pod-Encrypted" + mock_create_pod.return_value = {"id": "encrypted_pod_id"} + + runner = CliRunner() + result = runner.invoke(runpod_cli, ["pod", "create", "--encrypt-volume"]) + + # Assertions + assert result.exit_code == 0, result.exception + mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod") + mock_echo.assert_called_with("Pod encrypted_pod_id has been created.") + mock_create_pod.assert_called_with( + "RunPod-CLI-Pod-Encrypted", + "runpod/base:0.0.0", + "NVIDIA GeForce RTX 3090", + gpu_count=1, + support_public_ip=True, + ports="22/tcp", + encrypt_volume=True, # This is the key assertion + ) + + @patch("runpod.cli.groups.pod.commands.click.prompt") + @patch("runpod.cli.groups.pod.commands.click.confirm") + @patch("runpod.cli.groups.pod.commands.click.echo") + @patch("runpod.cli.groups.pod.commands.create_pod") + def test_create_new_pod_with_no_encrypt_volume( + self, mock_create_pod, mock_echo, mock_confirm, mock_prompt + ): # pylint: disable=too-many-arguments,line-too-long + """ + Test create_new_pod with --no-encrypt-volume option + """ + # Mock values + mock_confirm.return_value = True # for the quick_launch option + mock_prompt.return_value = "RunPod-CLI-Pod-No-Encryption" + mock_create_pod.return_value = {"id": "no_encryption_pod_id"} + + runner = CliRunner() + result = runner.invoke(runpod_cli, ["pod", "create", "--no-encrypt-volume"]) + + # Assertions + assert result.exit_code == 0, result.exception + mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod") + mock_echo.assert_called_with("Pod no_encryption_pod_id has been created.") + mock_create_pod.assert_called_with( + "RunPod-CLI-Pod-No-Encryption", + "runpod/base:0.0.0", + "NVIDIA GeForce RTX 3090", + gpu_count=1, + support_public_ip=True, + ports="22/tcp", + encrypt_volume=False, # This is the key assertion + ) + + @patch("runpod.cli.groups.pod.commands.click.prompt") + @patch("runpod.cli.groups.pod.commands.click.confirm") + @patch("runpod.cli.groups.pod.commands.click.echo") + @patch("runpod.cli.groups.pod.commands.create_pod") + def test_create_new_pod_default_encrypt_volume( + self, mock_create_pod, mock_echo, mock_confirm, mock_prompt + ): # pylint: disable=too-many-arguments,line-too-long + """ + Test create_new_pod with default encrypt_volume (should be False) + """ + # Mock values + mock_confirm.return_value = True # for the quick_launch option + mock_prompt.return_value = "RunPod-CLI-Pod-Default" + mock_create_pod.return_value = {"id": "default_pod_id"} + + runner = CliRunner() + result = runner.invoke(runpod_cli, ["pod", "create"]) + + # Assertions + assert result.exit_code == 0, result.exception + mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod") + mock_echo.assert_called_with("Pod default_pod_id has been created.") + mock_create_pod.assert_called_with( + "RunPod-CLI-Pod-Default", + "runpod/base:0.0.0", + "NVIDIA GeForce RTX 3090", + gpu_count=1, + support_public_ip=True, + ports="22/tcp", + encrypt_volume=False, # Default should be False + ) + @patch("runpod.cli.groups.pod.commands.click.echo") @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection") def test_connect_to_pod(self, mock_ssh_connection, mock_echo):