.github/workflows/CD-publish_to_pypi.yml

@@ -13,10 +13,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.11.0
+      - name: Set up Python 3.11.10
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11.0
+          python-version: 3.11.10
 
       - name: Install pypa/build
         run: >-

.github/workflows/CD-test_publish_to_pypi.yml

@@ -14,10 +14,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.11.0
+      - name: Set up Python 3.11.10
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11.0
+          python-version: 3.11.10
 
       - name: Install pypa/build
         run: >-

.github/workflows/CI-pytests.yml

@@ -15,7 +15,7 @@ jobs:
   run_tests:
     strategy:
       matrix:
-        python-version: [3.8, 3.9, 3.10.12, 3.11.0]
+        python-version: [3.8, 3.9, 3.10.15, 3.11.10]
     runs-on: ubuntu-latest
 
     steps:

pyproject.toml

@@ -54,7 +54,7 @@ runpod = "runpod.cli.entry:runpod_cli"
 test = [
     "asynctest",
     "nest_asyncio",
-    "pylint==3.2.5",
+    "faker",
     "pytest-asyncio",
     "pytest-cov",
     "pytest-timeout",

requirements.txt

@@ -5,7 +5,7 @@ backoff >= 2.2.1
 boto3 >= 1.26.165
 click >= 8.1.7
 colorama >= 0.2.5, < 0.4.7
-cryptography < 44.0.0
+cryptography < 45.0.0
 fastapi[all] >= 0.94.0
 paramiko >= 3.3.1
 prettytable >= 3.9.0

runpod/api/ctl_commands.py

@@ -107,6 +107,8 @@ def create_pod(
     template_id: Optional[str] = None,
     network_volume_id: Optional[str] = None,
     allowed_cuda_versions: Optional[list] = None,
+    min_download = None,
+    min_upload = None,
 ) -> dict:
     """
     Create a pod
@@ -125,7 +127,8 @@ def create_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
     :param template_id: the id of the template to use for the pod
-
+    :param min_download: minimum download speed in Mbps
+    :param min_upload: minimum upload speed in Mbps
     :example:
 
     >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070")
@@ -167,6 +170,8 @@ def create_pod(
             template_id,
             network_volume_id,
             allowed_cuda_versions,
+            min_download,
+            min_upload,
         )
     )
 
@@ -297,24 +302,13 @@ def create_endpoint(
     workers_min: int = 0,
     workers_max: int = 3,
     flashboot=False,
+    allowed_cuda_versions: str = "12.1,12.2,12.3,12.4,12.5",
+    gpu_count: int = 1,
 ):
     """
     Create an endpoint
 
-    :param name: the name of the endpoint
-    :param template_id: the id of the template to use for the endpoint
-    :param gpu_ids: the ids of the GPUs to use for the endpoint
-    :param network_volume_id: the id of the network volume to use for the endpoint
-    :param locations: the locations to use for the endpoint
-    :param idle_timeout: the idle timeout for the endpoint
-    :param scaler_type: the scaler type for the endpoint
-    :param scaler_value: the scaler value for the endpoint
-    :param workers_min: the minimum number of workers for the endpoint
-    :param workers_max: the maximum number of workers for the endpoint
-
-    :example:
-
-    >>> endpoint_id = runpod.create_endpoint("test", "template_id")
+    :param allowed_cuda_versions: Comma-separated string of allowed CUDA versions (e.g., "12.4,12.5").
     """
     raw_response = run_graphql_query(
         endpoint_mutations.generate_endpoint_mutation(
@@ -329,6 +323,8 @@ def create_endpoint(
             workers_min,
             workers_max,
             flashboot,
+            allowed_cuda_versions,
+            gpu_count
         )
     )
 

runpod/api/graphql.py

@@ -21,11 +21,12 @@ def run_graphql_query(query: str) -> Dict[str, Any]:
     from runpod import api_key  # pylint: disable=import-outside-toplevel, cyclic-import
 
     api_url_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io")
-    url = f"{api_url_base}/graphql?api_key={api_key}"
+    url = f"{api_url_base}/graphql"
 
     headers = {
         "Content-Type": "application/json",
         "User-Agent": USER_AGENT,
+        "Authorization": f"Bearer {api_key}",
     }
 
     data = json.dumps({"query": query})

runpod/api/mutations/endpoints.py

@@ -15,6 +15,8 @@ def generate_endpoint_mutation(
     workers_min: int = 0,
     workers_max: int = 3,
     flashboot=False,
+    allowed_cuda_versions: str = "12.1,12.2,12.3,12.4,12.5",
+    gpu_count: int = None,
 ):
     """Generate a string for a GraphQL mutation to create a new endpoint."""
     input_fields = []
@@ -44,6 +46,12 @@ def generate_endpoint_mutation(
     input_fields.append(f"workersMin: {workers_min}")
     input_fields.append(f"workersMax: {workers_max}")
 
+    if allowed_cuda_versions is not None:
+        input_fields.append(f'allowedCudaVersions: "{allowed_cuda_versions}"')
+
+    if gpu_count is not None:
+        input_fields.append(f"gpuCount: {gpu_count}")
+
     # Format the input fields into a string
     input_fields_string = ", ".join(input_fields)
 
@@ -65,11 +73,14 @@ def generate_endpoint_mutation(
             scalerValue
             workersMin
             workersMax
+            allowedCudaVersions
+            gpuCount
         }}
     }}
     """
 
 
+
 def update_endpoint_template_mutation(endpoint_id: str, template_id: str):
     """Generate a string for a GraphQL mutation to update an existing endpoint's template."""
     input_fields = []

runpod/api/mutations/pods.py

@@ -28,6 +28,8 @@ def generate_pod_deployment_mutation(
     template_id=None,
     network_volume_id=None,
     allowed_cuda_versions: Optional[List[str]] = None,
+    min_download=None,
+    min_upload=None,
 ):
     """
     Generates a mutation to deploy a pod on demand.
@@ -89,9 +91,14 @@ def generate_pod_deployment_mutation(
         )
         input_fields.append(f"allowedCudaVersions: [{allowed_cuda_versions_string}]")
 
+    if min_download is not None:
+        input_fields.append(f'minDownload: {min_download}')
+
+    if min_upload is not None:
+        input_fields.append(f'minUpload: {min_upload}')
+
     # Format input fields
     input_string = ", ".join(input_fields)
-
     return f"""
     mutation {{
       podFindAndDeployOnDemand(

runpod/endpoint/asyncio/asyncio_runner.py

@@ -1,4 +1,4 @@
-""" Module for running endpoints asynchronously. """
+"""Module for running endpoints asynchronously."""
 
 # pylint: disable=too-few-public-methods,R0801
 
@@ -89,9 +89,14 @@ async def stream(self) -> Any:
         while True:
             await asyncio.sleep(1)
             stream_partial = await self._fetch_job(source="stream")
-            if stream_partial["status"] not in FINAL_STATES:
+            if (
+                stream_partial["status"] not in FINAL_STATES
+                or len(stream_partial.get("stream", [])) > 0
+            ):
                 for chunk in stream_partial.get("stream", []):
                     yield chunk["output"]
+            elif stream_partial["status"] in FINAL_STATES:
+                break
 
     async def cancel(self) -> dict:
         """Cancels current job

runpod/http_client.py

@@ -5,13 +5,16 @@
 import os
 
 import requests
-from aiohttp import ClientSession, ClientTimeout, TCPConnector
+from aiohttp import ClientSession, ClientTimeout, TCPConnector, ClientResponseError
 
 from .cli.groups.config.functions import get_credentials
-from .tracer import create_aiohttp_tracer, create_request_tracer
 from .user_agent import USER_AGENT
 
 
+class TooManyRequests(ClientResponseError):
+    pass
+
+
 def get_auth_header():
     """
     Produce a header dict with the `Authorization` key derived from
@@ -33,13 +36,11 @@ def AsyncClientSession(*args, **kwargs):  # pylint: disable=invalid-name
     """
     Deprecation from aiohttp.ClientSession forbids inheritance.
     This is now a factory method
-    TODO: use httpx
     """
     return ClientSession(
         connector=TCPConnector(limit=0),
         headers=get_auth_header(),
         timeout=ClientTimeout(600, ceil_threshold=400),
-        trace_configs=[create_aiohttp_tracer()],
         *args,
         **kwargs,
     )
@@ -48,43 +49,5 @@ def AsyncClientSession(*args, **kwargs):  # pylint: disable=invalid-name
 class SyncClientSession(requests.Session):
     """
     Inherits requests.Session to override `request()` method for tracing
-    TODO: use httpx
     """
-
-    def request(self, method, url, **kwargs):  # pylint: disable=arguments-differ
-        """
-        Override for tracing. Not using super().request()
-        to capture metrics for connection and transfer times
-        """
-        with create_request_tracer() as tracer:
-            # Separate out the kwargs that are not applicable to `requests.Request`
-            request_kwargs = {
-                k: v
-                for k, v in kwargs.items()
-                # contains the names of the arguments
-                if k in requests.Request.__init__.__code__.co_varnames
-            }
-
-            # Separate out the kwargs that are applicable to `requests.Request`
-            send_kwargs = {k: v for k, v in kwargs.items() if k not in request_kwargs}
-
-            # Create a PreparedRequest object to hold the request details
-            req = requests.Request(method, url, **request_kwargs)
-            prepped = self.prepare_request(req)
-            tracer.request = prepped  # Assign the request to the tracer
-
-            # Merge environment settings
-            settings = self.merge_environment_settings(
-                prepped.url,
-                send_kwargs.get("proxies"),
-                send_kwargs.get("stream"),
-                send_kwargs.get("verify"),
-                send_kwargs.get("cert"),
-            )
-            send_kwargs.update(settings)
-
-            # Send the request
-            response = self.send(prepped, **send_kwargs)
-            tracer.response = response  # Assign the response to the tracer
-
-            return response
+    pass

runpod/serverless/__init__.py

@@ -10,7 +10,6 @@
 import signal
 import sys
 import time
-import typing
 from typing import Any, Dict
 
 from runpod.serverless import core
@@ -23,6 +22,7 @@
 
 log = RunPodLogger()
 
+
 # ---------------------------------------------------------------------------- #
 #                              Run Time Arguments                              #
 # ---------------------------------------------------------------------------- #