Skip to content

feat: add support for encrypted pod volumes via encrypt_volume param #440

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
3 changes: 3 additions & 0 deletions runpod/api/ctl_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -185,6 +187,7 @@ def create_pod(
min_download,
min_upload,
instance_id,
encrypt_volume,
)
)

Expand Down
4 changes: 4 additions & 0 deletions runpod/api/mutations/pods.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
6 changes: 5 additions & 1 deletion runpod/cli/groups/pod/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,19 @@ 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.
"""
kwargs = {
"gpu_count": gpu_count,
"support_public_ip": support_public_ip,
"encrypt_volume": encrypt_volume,
}

if not name:
Expand Down
4 changes: 3 additions & 1 deletion runpod/cli/groups/project/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand Down Expand Up @@ -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)

Expand Down
3 changes: 2 additions & 1 deletion runpod/cli/groups/project/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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="")
Expand All @@ -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
Expand Down
82 changes: 80 additions & 2 deletions tests/test_api/test_ctl_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand Down
44 changes: 44 additions & 0 deletions tests/test_api/test_mutations_pods.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
97 changes: 97 additions & 0 deletions tests/test_cli/test_cli_groups/test_pod_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading