Merge pull request #216 from runpod/sdk-examples

Sdk examples

Author: justinmerrell

CHANGELOG.md

@@ -1,5 +1,25 @@
 # Change Log
 
+## Release 1.3.3 (11/8/23)
+
+### Added
+
+- Method of creating logs with job id.
+
+### Fixed
+
+- Reduced polling when checking for job completion.
+- Removed print statements for endpoint calls.
+- Serverless progress updates no longer restricted to only strings.
+
+## Changed
+
+- Removed `pillow` dependency.
+- Removed `python-dotenv` dependency.
+- Removed `setuptools_scm` from required dependencies.
+
+---
+
 ## Release 1.3.2 (11/3/23)
 
 ### Changed

examples/call_endpoint.py

@@ -0,0 +1,26 @@
+'''
+Example of calling an endpoint using the RunPod Python Language Library.
+'''
+
+import runpod
+
+# Set your global API key with `runpod config` or uncomment the line below:
+# runpod.api_key = "YOUR_RUNPOD_API_KEY"
+
+endpoint = runpod.Endpoint("sdxl")  # Where "sdxl" is the endpoint ID
+
+run_request = endpoint.run({
+    "input": {
+        "prompt": "a photo of a horse the size of a Boeing 787"
+    }
+})
+
+# Check the status of the run request
+print(run_request.status())
+
+# Get the output of the endpoint run request.
+print(run_request.output())
+
+# Get the output of the endpoint run request.
+# If timeout is greater than 0, this will block until the endpoint run is complete.
+print(run_request.output(timeout=60))

examples/call_endpoint_asyncio.py renamed to examples/endpoints/asyncio_job_request.py

@@ -0,0 +1,25 @@
+'''
+Example of calling an endpoint using the RunPod Python Language Library.
+'''
+
+import runpod
+
+# Set your global API key with `runpod config` or uncomment the line below:
+# runpod.api_key = "YOUR_RUNPOD_API_KEY"
+
+endpoint = runpod.Endpoint("sdxl")  # Where "sdxl" is the endpoint ID
+
+try:
+    # Run the endpoint synchronously, blocking until the endpoint run is complete.
+    run_request = endpoint.run_sync(
+        {
+            "input": {
+                "prompt": "a photo of a horse the size of a Boeing 787"
+            }
+        },
+        timeout=60  # Seconds
+    )
+
+    print(run_request)
+except TimeoutError as err:
+    print("Job timed out.")

examples/endpoints/run.py

@@ -1,168 +1,197 @@
 '''
 RunPod | Python | Endpoint Runner
 '''
-from typing import Any, Union
+from typing import Any, Optional, Dict
 import time
 import requests
 from requests.adapters import HTTPAdapter, Retry
 
+# Exception Messages
+UNAUTHORIZED_MSG = "401 Unauthorized | Make sure Runpod API key is set and valid."
+API_KEY_NOT_SET_MSG = ("Expected `run_pod.api_key` to be initialized. "
+                       "You can solve this by setting `run_pod.api_key = 'your-key'. "
+                       "An API key can be generated at "
+                       "https://runpod.io/console/user/settings")
+
 
 # ---------------------------------------------------------------------------- #
 #                                    Client                                    #
 # ---------------------------------------------------------------------------- #
 class RunPodClient:
-    ''' A client for running endpoint calls. '''
+    """A client for running endpoint calls."""
 
     def __init__(self):
-        '''
-        Initialize the client.
-        '''
+        """
+        Initialize a RunPodClient instance.
+
+        Raises:
+            RuntimeError: If the API key has not been initialized.
+        """
         from runpod import api_key, endpoint_url_base  # pylint: disable=import-outside-toplevel, cyclic-import
+
         if api_key is None:
-            raise RuntimeError(
-                "Expected `run_pod.api_key` to be initialized. "
-                "You can solve this by running `run_pod.api_key = 'your-key'. "
-                "An API key can be generated at "
-                "https://www.runpod.io/console/user/settings"
-            )
+            raise RuntimeError(API_KEY_NOT_SET_MSG)
+
         self.rp_session = requests.Session()
-        retries = Retry(total=5, backoff_factor=0.1, status_forcelist=[429])
+        retries = Retry(total=5, backoff_factor=1, status_forcelist=[429])
         self.rp_session.mount('http://', HTTPAdapter(max_retries=retries))
 
-        self.endpoint_url_base = endpoint_url_base
         self.headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {api_key}"
         }
 
-    def post(self, endpoint : str, data : dict, timeout : int=10):
-        '''
-        Post to the endpoint.
-        '''
-        url = f"{self.endpoint_url_base}/{endpoint}"
-        return self.rp_session.post(url, headers=self.headers, json=data, timeout=timeout)
+        self.endpoint_url_base = endpoint_url_base
+
+    def _request(self,
+                 method: str, endpoint: str, data: Optional[dict] = None, timeout: int = 10):
+        """
+        Make a request to the specified endpoint using the given HTTP method.
+
+        Args:
+            method: The HTTP method to use ('GET' or 'POST').
+            endpoint: The endpoint path to which the request will be made.
+            data: The JSON payload to send with the request.
+            timeout: The number of seconds to wait for the server to send data before giving up.
+
+        Returns:
+            The JSON response from the server.
 
-    def get(self, endpoint : str, timeout : int=10):
-        '''
-        Get from the endpoint.
-        '''
+        Raises:
+            RuntimeError: If the response returns a 401 Unauthorized status.
+            requests.HTTPError: If the response contains an unsuccessful status code.
+        """
         url = f"{self.endpoint_url_base}/{endpoint}"
-        return self.rp_session.get(url, headers=self.headers, timeout=timeout)
+        response = self.rp_session.request(
+            method, url, headers=self.headers, json=data, timeout=timeout)
+
+        if response.status_code == 401:
+            raise RuntimeError(UNAUTHORIZED_MSG)
+
+        response.raise_for_status()
+        return response.json()
+
+    def post(self, endpoint: str, data: dict, timeout: int = 10):
+        """ Post to the endpoint. """
+        return self._request('POST', endpoint, data, timeout)
+
+    def get(self, endpoint: str, timeout: int = 10):
+        """ Get from the endpoint. """
+        return self._request('GET', endpoint, timeout=timeout)
+
 
 # ---------------------------------------------------------------------------- #
 #                                      Job                                     #
 # ---------------------------------------------------------------------------- #
 class Job:
-    ''' Creates a class to run a job. '''
+    """Represents a job to be run on the RunPod service."""
 
-    def __init__(self, endpoint_id : str, job_id : str):
-        ''' Initializes the class. '''
+    def __init__(self, endpoint_id: str, job_id: str, client: RunPodClient):
+        """
+        Initialize a Job instance with the given endpoint ID and job ID.
 
+        Args:
+            endpoint_id: The identifier for the endpoint.
+            job_id: The identifier for the job.
+            client: An instance of the RunPodClient to make requests with.
+        """
         self.endpoint_id = endpoint_id
         self.job_id = job_id
-        self.rp_client = RunPodClient()
+        self.rp_client = client
 
+        self.job_status = None
         self.job_output = None
 
-    def _status_json(self):
-        """
-        Returns the raw json of the status, raises an exception if invalid
-        """
-
+    def _fetch_job(self):
+        """ Returns the raw json of the status, raises an exception if invalid """
         status_url = f"{self.endpoint_id}/status/{self.job_id}"
+        job_state = self.rp_client.get(endpoint=status_url)
 
-        status_request = self.rp_client.get(endpoint=status_url, timeout=10)
-        request_json = status_request.json()
+        if job_state["status"] in ["COMPLETED", "FAILED", "TIMEOUT"]:
+            self.job_status = job_state["status"]
+            self.job_output = job_state.get("output", None)
 
-        if "error" in request_json:
-            raise RuntimeError(f"Error from RunPod Server: '{request_json['error']}'")
-
-        return request_json
+        return job_state
 
     def status(self):
-        '''
-        Returns the status of the job request.
-        '''
-        return self._status_json()["status"]
+        """ Returns the status of the job request. """
+        if self.job_status is not None:
+            return self.job_status
 
-    def output(self, timeout : int=60) -> Union[None, dict]:
-        '''
-        Gets the output of the endpoint run request.
+        return self._fetch_job()["status"]
 
-        :param timeout: after how much time should the request timeout? 
-                        (if it doesn't get a response back)
-        '''
-        while self.status() not in ["COMPLETED", "FAILED", "TIMEOUT"]:
-            time.sleep(.1)
-            timeout -= .1
+    def output(self, timeout: int = 0) -> Any:
+        """
+        Returns the output of the job request.
 
-        if self.job_output is None:
-            status_json = self._status_json()
-            if "output" not in status_json:
-                return None
-            self.job_output = status_json["output"]
+        Args:
+            timeout: The number of seconds to wait for the server to send data before giving up.
+        """
+        if timeout > 0:
+            while self.status() not in ["COMPLETED", "FAILED", "TIMEOUT"]:
+                time.sleep(1)
+                timeout -= 1
+                if timeout <= 0:
+                    raise TimeoutError("Job timed out.")
 
-        return self.job_output
+        if self.job_output is not None:
+            return self.job_output
 
+        return self._fetch_job().get("output", None)
 
 
 # ---------------------------------------------------------------------------- #
 #                                   Endpoint                                   #
 # ---------------------------------------------------------------------------- #
 class Endpoint:
-    '''Creates a class to run an endpoint.'''
-
-    def __init__(self, endpoint_id : str):
-        '''
-        Initializes the class
-
-        :param endpoint_id: the id of the endpoint
-
-        :example:
-
-        >>> endpoint = runpod.Endpoint("ENDPOINT_ID")
-        >>> run_request = endpoint.run(
-                {"your_model_input_key": "your_model_input_value"}
-            )
-        >>> print(run_request.status())
-        
-        >>> print(run_request.output())
-        '''
-        # the endpoint id
-        self.endpoint_id : str = endpoint_id
+    """Manages an endpoint to run jobs on the RunPod service."""
+
+    def __init__(self, endpoint_id: str):
+        """
+        Initialize an Endpoint instance with the given endpoint ID.
+
+        Args:
+            endpoint_id: The identifier for the endpoint.
+
+        Example:
+            >>> endpoint = runpod.Endpoint("ENDPOINT_ID")
+            >>> run_request = endpoint.run({"your_model_input_key": "your_model_input_value"})
+            >>> print(run_request.status())
+            >>> print(run_request.output())
+        """
+        self.endpoint_id = endpoint_id
         self.rp_client = RunPodClient()
 
-        print(f"Initialized endpoint: {self.endpoint_id}")
+    def run(self, request_input: Dict[str, Any]) -> Job:
+        """
+        Run the endpoint with the given input.
 
-    def run(self, endpoint_input : Any) -> Job:
-        '''
-        Runs the endpoint.
+        Args:
+            request_input: The input to pass into the endpoint.
 
-        :param endpoint_input: the input to pass into the endpoint
-        '''
-        job_request = self.rp_client.post(
-            endpoint=f"{self.endpoint_id}/run",
-            data={"input": endpoint_input},
-            timeout=10
-        )
+        Returns:
+            A Job instance for the run request.
+        """
+        if not request_input.get("input"):
+            request_input = {"input": request_input}
 
-        if job_request.status_code == 401:
-            raise RuntimeError("401 Unauthorized | Make sure Runpod API key is set and valid.")
+        job_request = self.rp_client.post(f"{self.endpoint_id}/run", request_input)
+        return Job(self.endpoint_id, job_request["id"], self.rp_client)
 
-        print(f"Started job: {job_request.json()['id']}")
+    def run_sync(self, request_input: Dict[str, Any], timeout: int = 86400) -> Dict[str, Any]:
+        """
+        Run the endpoint with the given input synchronously.
 
-        return Job(self.endpoint_id, job_request.json()["id"])
+        Args:
+            request_input: The input to pass into the endpoint.
+        """
+        if not request_input.get("input"):
+            request_input = {"input": request_input}
 
-    def run_sync(self, endpoint_input : Any) -> dict:
-        '''
-        Blocking run where the job results are returned with the call.
+        job_request = self.rp_client.post(
+            f"{self.endpoint_id}/runsync", request_input, timeout=timeout)
 
-        :param endpoint_input: the input to pass into the endpoint
-        '''
-        job_return = self.rp_client.post(
-            endpoint=f"{self.endpoint_id}/runsync",
-            data={"input": endpoint_input},
-            timeout=60
-        )
+        if job_request["status"] in ["COMPLETED", "FAILED", "TIMEOUT"]:
+            return job_request.get("output", None)
 
-        return job_return.json()
+        return Job(self.endpoint_id, job_request["id"], self.rp_client).output(timeout=timeout)