Skypilot

docs/cli/command_line_interface.md

@@ -0,0 +1,10 @@
+# RunPod CLI
+
+Note: This CLI is not the same as runpodctl and provides a different set of features.
+
+
+## Overview
+
+```bash
+runpod --help
+```

docs/getting_Started.md

@@ -0,0 +1,21 @@
+# Getting Started
+
+The runpod python library is a powerful library providing SDK functions, API access, and CLI commands for interacting with the runpod platform.
+
+## Credentials File
+
+This python library supports a credentials file saved to `~/.runpod/credentials.toml` that contains the following information:
+
+```toml
+[profile]
+api_key = "YOUR_RUNPOD_API_KEY"
+```
+### Profile
+
+By default all credentials are stored under the `default` profile. To switch profiles you can use the `--profile` argument on CLI commands or change the profile within a script.
+
+```python
+import runpod
+
+runpod.profile = "my_profile"
+```

requirements.txt

@@ -2,6 +2,7 @@ aiohttp >= 3.8.4
 
 backoff == 2.2.1
 boto3 >= 1.26.165
+click >= 8.1.7
 pillow >= 9.5.0
 py-cpuinfo == 9.0.0
 python-dotenv >= 1.0.0

runpod/__init__.py

@@ -5,14 +5,21 @@
 from . import serverless
 from .endpoint import Endpoint
 from .endpoint import AsyncioEndpoint, AsyncioJob
-from .api_wrapper.ctl_commands import(
+from .api.ctl_commands import(
     get_gpus, get_gpu,
     get_pods, get_pod,
     create_pod, stop_pod, resume_pod, terminate_pod
 )
 from .cli.config import set_credentials, check_credentials, get_credentials
 
-api_key = None  # pylint: disable=invalid-name
+
+profile = "default"  # pylint: disable=invalid-name
+
+_credentials = get_credentials(profile)
+if _credentials is not None:
+    api_key = _credentials['api_key']  # pylint: disable=invalid-name
+else:
+    api_key = None  # pylint: disable=invalid-name
 
 api_url_base = "https://api.runpod.io"  # pylint: disable=invalid-name
 

runpod/api_wrapper/ctl_commands.py → runpod/api/ctl_commands.py

@@ -20,15 +20,22 @@ def get_gpus() -> dict:
     return cleaned_return
 
 
-def get_gpu(gpu_id : str):
+def get_gpu(gpu_id : str, gpu_quantity : int = 1):
     '''
     Get a specific GPU type
 
     :param gpu_id: the id of the gpu
+    :param gpu_quantity: how many of the gpu should be returned
     '''
-    raw_response = run_graphql_query(gpus.generate_gpu_query(gpu_id))
-    cleaned_return = raw_response["data"]["gpuTypes"][0]
-    return cleaned_return
+    raw_response = run_graphql_query(gpus.generate_gpu_query(gpu_id, gpu_quantity))
+
+    cleaned_return = raw_response["data"]["gpuTypes"]
+
+    if len(cleaned_return) < 1:
+        raise ValueError("No GPU found with the specified ID, "
+                         "run runpod.get_gpus() to get a list of all GPUs")
+
+    return cleaned_return[0]
 
 def get_pods() -> dict:
     '''
@@ -47,12 +54,15 @@ def get_pod(pod_id : str):
     raw_response = run_graphql_query(pod_queries.generate_pod_query(pod_id))
     return raw_response["data"]["pod"]
 
-def create_pod(name : str, image_name : str, gpu_type_id : str, cloud_type : str="ALL",
-               data_center_id : Optional[str]=None, country_code:Optional[str]=None,
-               gpu_count:int=1, volume_in_gb:int=0, container_disk_in_gb:int=5,
-               min_vcpu_count:int=1, min_memory_in_gb:int=1, docker_args:str="",
-               ports:Optional[str]=None, volume_mount_path:str="/workspace",
-               env:Optional[dict]=None):
+def create_pod(
+        name:str, image_name:str, gpu_type_id:str,
+        cloud_type:str="ALL", support_public_ip:bool=False,
+        data_center_id : Optional[str]=None, country_code:Optional[str]=None,
+        gpu_count:int=1, volume_in_gb:int=0, container_disk_in_gb:int=5,
+        min_vcpu_count:int=1, min_memory_in_gb:int=1, docker_args:str="",
+        ports:Optional[str]=None, volume_mount_path:str="/workspace",
+        env:Optional[dict]=None
+    ) -> dict:
     '''
     Create a pod
 
@@ -74,10 +84,16 @@ def create_pod(name : str, image_name : str, gpu_type_id : str, cloud_type : str
 
     >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070")
     '''
+    # Input Validation
+    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")
 
     raw_response = run_graphql_query(
         pod_mutations.generate_pod_deployment_mutation(
-            name, image_name, gpu_type_id, cloud_type, data_center_id, country_code, gpu_count,
+            name, image_name, gpu_type_id,
+            cloud_type, support_public_ip,
+            data_center_id, country_code, gpu_count,
             volume_in_gb, container_disk_in_gb, min_vcpu_count, min_memory_in_gb, docker_args,
             ports, volume_mount_path, env)
     )

runpod/api_wrapper/mutations/pods.py → runpod/api/mutations/pods.py

@@ -5,17 +5,32 @@
 
 
 def generate_pod_deployment_mutation(
-        name, image_name, gpu_type_id, cloud_type=None, data_center_id=None, country_code=None,
+        name:str, image_name:str, gpu_type_id:str,
+        cloud_type:str="ALL", support_public_ip:bool=False,
+        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=None, support_public_ip=None):
+        env=None):
     '''
     Generates a mutation to deploy a pod on demand.
     '''
     input_fields = []
 
-    if cloud_type is not None:
-        input_fields.append(f"cloudType: {cloud_type}")
+    # 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
+    del cloud_type # Temporarily remove cloud_type
+    # input_fields.append(f'cloudType: "{cloud_type}"')
+
+    if support_public_ip:
+        input_fields.append('supportPublicIp: true')
+    else:
+        input_fields.append('supportPublicIp: false')
+
+    # Optional Fields
     if data_center_id is not None:
         input_fields.append(f'dataCenterId: "{data_center_id}"')
     if country_code is not None:
@@ -30,12 +45,6 @@ def generate_pod_deployment_mutation(
         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 gpu_type_id is not None:
-        input_fields.append(f'gpuTypeId: "{gpu_type_id}"')
-    if name is not None:
-        input_fields.append(f'name: "{name}"')
-    if image_name is not None:
-        input_fields.append(f'imageName: "{image_name}"')
     if docker_args is not None:
         input_fields.append(f'dockerArgs: "{docker_args}"')
     if ports is not None:
@@ -46,9 +55,9 @@ def generate_pod_deployment_mutation(
         env_string = ", ".join(
             [f'{{ key: "{key}", value: "{value}" }}' for key, value in env.items()])
         input_fields.append(f"env: [{env_string}]")
-    if support_public_ip is not None:
-        input_fields.append(f"supportPublicIp: {support_public_ip}")
 
+
+    # Format input fields
     input_string = ", ".join(input_fields)
 
     return f"""

runpod/api_wrapper/queries/gpus.py → runpod/api/queries/gpus.py

@@ -1,5 +1,5 @@
 """
-RunPod | API Wrapper | Queries | GPUs
+RunPod | API | Queries | GPUs
 """
 
 QUERY_GPU_TYPES = """

runpod/api_wrapper/queries/pods.py → runpod/api/queries/pods.py

@@ -26,6 +26,15 @@
             vcpuCount
             volumeInGb
             volumeMountPath
+            runtime {
+                ports{
+                    ip
+                    isIpPublic
+                    privatePort
+                    publicPort
+                    type
+                }
+            }
             machine {
                 gpuDisplayName
             }
@@ -62,6 +71,15 @@ def generate_pod_query(pod_id):
             vcpuCount
             volumeInGb
             volumeMountPath
+            runtime {{
+                ports {{
+                    ip
+                    isIpPublic
+                    privatePort
+                    publicPort
+                    type
+                }}
+            }}
             machine {{
                 gpuDisplayName
             }}

runpod/api/queries/user.py

@@ -0,0 +1,5 @@
+'''
+RunPod | API | Queries | User
+
+Query for user information.
+'''

runpod/cli/commands.py

@@ -0,0 +1,44 @@
+'''
+RunPod | CLI Commands
+
+A collection of CLI functions.
+'''
+import sys
+import click
+
+from .config import set_credentials, check_credentials
+
+@click.group()
+def runpod_cli():
+    '''A collection of CLI functions for RunPod.'''
+
+@runpod_cli.command('store_api_key')
+@click.argument('api_key')
+@click.option('--profile', default='default', help='The profile to set the credentials for.')
+def store_api_key(api_key, profile):
+    '''
+    Sets the credentials for a profile.
+    '''
+    try:
+        set_credentials(api_key, profile)
+    except ValueError as err:
+        click.echo(err)
+        sys.exit(1)
+
+    click.echo('Credentials set for profile: ' + profile + ' in ~/.runpod/credentials.toml')
+
+
+@runpod_cli.command('check_creds')
+@click.option('--profile', default='default', help='The profile to check the credentials for.')
+def validate_credentials_file(profile='default'):
+    '''
+    Validates the credentials file.
+    '''
+    click.echo('Validating ~/.runpod/credentials.toml')
+    valid, error = check_credentials(profile)
+
+    if not valid:
+        click.echo(error)
+        sys.exit(1)
+
+    click.echo('Credentials file is valid.')

runpod/cli/config.py

@@ -4,50 +4,60 @@
 A collection of functions to set and validate configurations.
 Configurations are TOML files located under ~/.runpod/
 '''
-
 import os
+from pathlib import Path
 
 import tomli as toml
 
 CREDENTIAL_FILE = os.path.expanduser('~/.runpod/credentials.toml')
 
-def set_credentials(api_key: str) -> None:
+
+def set_credentials(api_key: str, profile:str="default") -> None:
     '''
     Sets the user's credentials in ~/.runpod/credentials.toml
+    If profile already exists user must use `update_credentials` instead.
 
     Args:
         api_key (str): The user's API key.
-
+        profile (str): The profile to set the credentials for.
 
     --- File Structure ---
 
     [default]
     api_key = "RUNPOD_API_KEY"
     '''
+    os.makedirs(os.path.dirname(CREDENTIAL_FILE), exist_ok=True)
+    Path(CREDENTIAL_FILE).touch(exist_ok=True)
+
+    with open(CREDENTIAL_FILE, 'rb') as cred_file:
+        if profile in toml.load(cred_file):
+            raise ValueError('Profile already exists. Use `update_credentials` instead.')
+
     with open(CREDENTIAL_FILE, 'w', encoding="UTF-8") as cred_file:
-        cred_file.write('[default]\n')
+        cred_file.write('[' + profile + ']\n')
         cred_file.write('api_key = "' + api_key + '"\n')
 
 
-def check_credentials():
+def check_credentials(profile:str="default"):
     '''
     Checks if the credentials file exists and is valid.
     '''
     if not os.path.exists(CREDENTIAL_FILE):
-        return False, 'Error: ~/.runpod/credentials.toml does not exist.'
+        return False, '~/.runpod/credentials.toml does not exist.'
 
     # Check for default api_key
     try:
-        config = toml.load(CREDENTIAL_FILE)
+        with open(CREDENTIAL_FILE, 'rb') as cred_file:
+            config = toml.load(cred_file)
 
-        if 'default' not in config:
-            return False, 'Error: ~/.runpod/credentials.toml is missing default section.'
+        if profile not in config:
+            return False, f'~/.runpod/credentials.toml is missing {profile} profile.'
 
-        if 'api_key' not in config['default']:
-            return False, 'Error: ~/.runpod/credentials.toml is missing api_key.'
+        if 'api_key' not in config[profile]:
+            return False, f'~/.runpod/credentials.toml is missing api_key for {profile} profile.'
 
     except (TypeError, ValueError):
-        return False, 'Error: ~/.runpod/credentials.toml is not a valid TOML file.'
+        return False, '~/.runpod/credentials.toml is not a valid TOML file.'
 
     return True, None
 
@@ -56,7 +66,10 @@ def get_credentials(profile='default'):
     '''
     Returns the credentials for the specified profile from ~/.runpod/credentials.toml
     '''
-    with open(CREDENTIAL_FILE, 'r', encoding="UTF-8") as cred_file:
+    if not os.path.exists(CREDENTIAL_FILE):
+        return None
+
+    with open(CREDENTIAL_FILE, 'rb') as cred_file:
         credentials = toml.load(cred_file)
 
     if profile not in credentials:

runpod/serverless/modules/rp_ping.py

@@ -35,7 +35,7 @@ def __init__(self, pool_connections=100, retries=3) -> None:
         retry_strategy = Retry(
             total=retries,
             status_forcelist=[429, 500, 502, 503, 504],
-            method_whitelist=["GET"],
+            allowed_methods=["GET"],
             backoff_factor=1
         )
 

setup.cfg

@@ -29,6 +29,7 @@ install_requires =
     aiohttp >= 3.8.4
     backoff >= 2.2.1
     boto3 >= 1.26.165
+    click >= 8.1.7
     pillow >= 9.5.0
     py-cpuinfo >= 9.0.0
     python-dotenv >= 1.0.0

setup.py

@@ -13,7 +13,7 @@
 
         entry_points={
             'console_scripts': [
-                'runpod = runpod.cli:main'
+                'runpod = runpod.cli.commands:runpod_cli'
             ]
         }
     )

tests/test_api/__init__.py

@@ -0,0 +1 @@
+''' Allows files to be imported from this directory. '''