Merge pull request #96 from runpod/api-improvements

API Improvements With Error Handling

Author: justinmerrell

docs/api/handling_errors.md

@@ -0,0 +1,10 @@
+# Handling Errors
+
+```Python
+import runpod
+
+try:
+    # Use runpod to make a request
+except runpod.error.AuthenticationError as err:
+    # Authentication with the API failed
+```

runpod/api_wrapper/ctl_commands.py

@@ -4,6 +4,7 @@
 # pylint: disable=too-many-arguments,too-many-locals
 
 from typing import Optional
+
 from .queries import gpus
 from .queries import pods as pod_queries
 from .graphql import run_graphql_query
@@ -14,19 +15,19 @@ def get_gpus() -> dict:
     '''
     Get all GPU types
     '''
-    raw_return = run_graphql_query(gpus.QUERY_GPU_TYPES)
-    cleaned_return = raw_return["data"]["gpuTypes"]
+    raw_response = run_graphql_query(gpus.QUERY_GPU_TYPES)
+    cleaned_return = raw_response["data"]["gpuTypes"]
     return cleaned_return
 
 
 def get_gpu(gpu_id : str):
     '''
     Get a specific GPU type
-    
+
     :param gpu_id: the id of the gpu
     '''
-    raw_return = run_graphql_query(gpus.generate_gpu_query(gpu_id))
-    cleaned_return = raw_return["data"]["gpuTypes"][0]
+    raw_response = run_graphql_query(gpus.generate_gpu_query(gpu_id))
+    cleaned_return = raw_response["data"]["gpuTypes"][0]
     return cleaned_return
 
 def get_pods() -> dict:
@@ -56,7 +57,7 @@ def create_pod(name : str, image_name : str, gpu_type_id : str, cloud_type : str
     :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?
-    :param env: the environment variables to inject into the pod, 
+    :param env: the environment variables to inject into the pod,
                 for example {EXAMPLE_VAR:"example_value", EXAMPLE_VAR2:"example_value 2"}, will
                 inject EXAMPLE_VAR and EXAMPLE_VAR2 into the pod with the mentioned values
 

runpod/api_wrapper/graphql.py

@@ -7,6 +7,9 @@
 
 import requests
 
+from runpod import error
+
+HTTP_STATUS_UNAUTHORIZED = 401
 
 def run_graphql_query(query: str) -> Dict[str, Any]:
     '''
@@ -19,4 +22,13 @@ def run_graphql_query(query: str) -> Dict[str, Any]:
     }
     data = json.dumps({"query": query})
     response = requests.post(url, headers=headers, data=data, timeout=30)
+
+    print(response.json())
+
+    if response.status_code == HTTP_STATUS_UNAUTHORIZED:
+        raise error.AuthenticationError("Unauthorized request, please check your API key.")
+
+    if "errors" in response.json():
+        raise error.QueryError(response.json()["errors"][0]["message"])
+
     return response.json()

runpod/error.py

@@ -0,0 +1,30 @@
+'''
+runpd | error.py
+
+This file contains the error classes for the runpod package.
+'''
+
+from typing import Optional
+
+
+class RunPodError(Exception):
+    '''
+    Base class for all runpod errors
+    '''
+    def __init__(self, message: Optional[str] = None):
+        super().__init__(message)
+
+        self.message = message
+
+
+
+class AuthenticationError(RunPodError):
+    '''
+    Raised when authentication fails
+    '''
+
+
+class QueryError(RunPodError):
+    '''
+    Raised when a GraphQL query fails
+    '''

tests/test_api_wrapper/test_ctl_commands.py

@@ -123,6 +123,35 @@ def test_terminate_pod(self):
 
             self.assertIsNone(ctl_commands.terminate_pod(pod_id="POD_ID"))
 
+    def test_raised_error(self):
+        '''
+        Test raised_error
+        '''
+        with patch("runpod.api_wrapper.graphql.requests.post") as patch_request:
+            patch_request.return_value.json.return_value = {
+                "errors": [
+                    {
+                        "message": "Error Message"
+                    }
+                ]
+            }
+
+            with self.assertRaises(Exception) as context:
+                ctl_commands.get_gpus()
+
+            self.assertEqual(str(context.exception), "Error Message")
+
+
+        # Test Unauthorized with status code 401
+        with patch("runpod.api_wrapper.graphql.requests.post") as patch_request:
+            patch_request.return_value.status_code = 401
+
+            with self.assertRaises(Exception) as context:
+                ctl_commands.get_gpus()
+
+            self.assertEqual(
+                str(context.exception), "Unauthorized request, please check your API key.")
+
     def test_get_pods(self):
         '''
         Tests get_pods