Removing request level initializatino

Author: justinwlin

runpod/endpoint/asyncio/asyncio_runner.py

@@ -12,38 +12,25 @@
 class Job:
     """Class representing a job for an asynchronous endpoint"""
 
-    def __init__(self, endpoint_id: str, job_id: str, session: ClientSession,
-                 api_key: Optional[str] = None):
+    def __init__(self, endpoint_id: str, job_id: str, session: ClientSession, headers: dict):
         """
-        Initialize a Job instance with optional API key.
+        Initialize a Job instance.
         
         Args:
             endpoint_id: The identifier for the endpoint.
             job_id: The identifier for the job.
             session: The aiohttp ClientSession.
-            api_key: Optional API key for this specific job.
+            headers: Headers to use for requests.
         """
         from runpod import (  # pylint: disable=import-outside-toplevel,cyclic-import
-            api_key as global_api_key,
             endpoint_url_base,
         )
 
         self.endpoint_id = endpoint_id
         self.job_id = job_id
         self.session = session
         self.endpoint_url_base = endpoint_url_base
-        
-        # Use provided API key or fall back to global
-        effective_api_key = api_key or global_api_key
-        
-        if effective_api_key is None:
-            raise RuntimeError("API key must be provided or set globally")
-        
-        self.headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {effective_api_key}",
-            "X-Request-ID": job_id,
-        }
+        self.headers = headers
 
         self.job_status = None
         self.job_output = None
@@ -155,82 +142,51 @@ def __init__(self, endpoint_id: str, session: ClientSession,
         if self.api_key is None:
             raise RuntimeError("API key must be provided or set globally")
 
-        # Keep headers attribute for backward compatibility
         self.headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {self.api_key}",
         }
-    
-    def _get_headers(self, api_key: Optional[str] = None) -> dict:
-        """
-        Get headers with the appropriate API key.
-        
-        Raises:
-            ValueError: If request API key conflicts with instance API key.
-        """
-        # Check for conflicting API keys
-        if api_key and self.api_key and api_key != self.api_key:
-            raise ValueError(
-                "Conflicting API keys: Request API key differs from instance API key. "
-                "Use only one API key source to avoid security issues."
-            )
-        
-        effective_api_key = api_key or self.api_key
-        return {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {effective_api_key}",
-        }
 
-    async def run(self, endpoint_input: dict, api_key: Optional[str] = None) -> Job:
+    async def run(self, endpoint_input: dict) -> Job:
         """
         Runs endpoint with specified input.
 
         Args:
             endpoint_input: any dictionary with input
-            api_key: Optional API key to use for this specific request.
 
         Returns:
             Newly created job
         """
-        headers = self._get_headers(api_key)
-        
         async with self.session.post(
-            self.endpoint_url, headers=headers, json={"input": endpoint_input}
+            self.endpoint_url, headers=self.headers, json={"input": endpoint_input}
         ) as resp:
             json_resp = await resp.json()
 
-        # Pass the API key to the Job instance
-        return Job(self.endpoint_id, json_resp["id"], self.session, 
-                   api_key=api_key or self.api_key)
+        # Create job with endpoint's headers
+        job_headers = self.headers.copy()
+        job_headers["X-Request-ID"] = json_resp["id"]
+        return Job(self.endpoint_id, json_resp["id"], self.session, job_headers)
 
-    async def health(self, api_key: Optional[str] = None) -> dict:
+    async def health(self) -> dict:
         """
         Checks health of endpoint
 
-        Args:
-            api_key: Optional API key to use for this specific request.
-
         Returns:
             Health of endpoint
         """
-        headers = self._get_headers(api_key)
         health_url = f"{self.endpoint_url_base}/{self.endpoint_id}/health"
 
-        async with self.session.get(health_url, headers=headers) as resp:
+        async with self.session.get(health_url, headers=self.headers) as resp:
             return await resp.json()
 
-    async def purge_queue(self, api_key: Optional[str] = None) -> dict:
+    async def purge_queue(self) -> dict:
         """
         Purges queue of endpoint
 
-        Args:
-            api_key: Optional API key to use for this specific request.
-
         Returns:
             Purge status
         """
-        headers = self._get_headers(api_key)
         purge_url = f"{self.endpoint_url_base}/{self.endpoint_id}/purge-queue"
 
-        async with self.session.post(purge_url, headers=headers) as resp:
+        async with self.session.post(purge_url, headers=self.headers) as resp:
             return await resp.json()

runpod/endpoint/runner.py

@@ -56,7 +56,7 @@ def __init__(self, api_key: Optional[str] = None):
 
     def _request(
         self, method: str, endpoint: str, data: Optional[dict] = None, 
-        timeout: int = 10, api_key: Optional[str] = None
+        timeout: int = 10
     ):
         """
         Make a request to the specified endpoint using the given HTTP method.
@@ -66,33 +66,17 @@ def _request(
             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.
-            api_key: Optional API key to use for this specific request.
 
         Returns:
             The JSON response from the server.
 
         Raises:
-            ValueError: If request API key conflicts with instance API key.
             RuntimeError: If the response returns a 401 Unauthorized status.
             requests.HTTPError: If the response contains an unsuccessful status code.
         """
-        # Check for conflicting API keys
-        if api_key and self.api_key and api_key != self.api_key:
-            raise ValueError(
-                "Conflicting API keys: Request API key differs from instance API key. "
-                "Use only one API key source to avoid security issues."
-            )
-        
-        # Use request-specific API key if provided, otherwise use instance API key
-        effective_api_key = api_key or self.api_key
-        headers = self.headers if not api_key else {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {effective_api_key}",
-        }
-        
         url = f"{self.endpoint_url_base}/{endpoint}"
         response = self.rp_session.request(
-            method, url, headers=headers, json=data, timeout=timeout
+            method, url, headers=self.headers, json=data, timeout=timeout
         )
 
         if response.status_code == 401:
@@ -101,13 +85,13 @@ def _request(
         response.raise_for_status()
         return response.json()
 
-    def post(self, endpoint: str, data: dict, timeout: int = 10, api_key: Optional[str] = None):
-        """Post to the endpoint with optional API key override."""
-        return self._request("POST", endpoint, data, timeout, api_key=api_key)
+    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, api_key: Optional[str] = None):
-        """Get from the endpoint with optional API key override."""
-        return self._request("GET", endpoint, timeout=timeout, api_key=api_key)
+    def get(self, endpoint: str, timeout: int = 10):
+        """Get from the endpoint."""
+        return self._request("GET", endpoint, timeout=timeout)
 
 
 # ---------------------------------------------------------------------------- #
@@ -116,30 +100,26 @@ def get(self, endpoint: str, timeout: int = 10, api_key: Optional[str] = None):
 class Job:
     """Represents a job to be run on the Runpod service."""
 
-    def __init__(self, endpoint_id: str, job_id: str, client: RunPodClient, 
-                 api_key: Optional[str] = None):
+    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.
-            api_key: Optional API key for this specific job.
         """
         self.endpoint_id = endpoint_id
         self.job_id = job_id
         self.rp_client = client
-        self.api_key = api_key  # Store job-specific API key
 
         self.job_status = None
         self.job_output = None
 
     def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
         """Returns the raw json of the status, raises an exception if invalid"""
         status_url = f"{self.endpoint_id}/{source}/{self.job_id}"
-        # Pass the job-specific API key if available
-        job_state = self.rp_client.get(endpoint=status_url, api_key=self.api_key)
+        job_state = self.rp_client.get(endpoint=status_url)
 
         if is_completed(job_state["status"]):
             self.job_status = job_state["status"]
@@ -197,8 +177,7 @@ def cancel(self, timeout: int = 3) -> Any:
         return self.rp_client.post(
             f"{self.endpoint_id}/cancel/{self.job_id}", 
             data=None, 
-            timeout=timeout,
-            api_key=self.api_key
+            timeout=timeout
         )
 
 
@@ -225,13 +204,12 @@ def __init__(self, endpoint_id: str, api_key: Optional[str] = None):
         self.endpoint_id = endpoint_id
         self.rp_client = RunPodClient(api_key=api_key)
 
-    def run(self, request_input: Dict[str, Any], api_key: Optional[str] = None) -> Job:
+    def run(self, request_input: Dict[str, Any]) -> Job:
         """
         Run the endpoint with the given input.
 
         Args:
             request_input: The input to pass into the endpoint.
-            api_key: Optional API key to use for this specific request.
 
         Returns:
             A Job instance for the run request.
@@ -241,65 +219,57 @@ def run(self, request_input: Dict[str, Any], api_key: Optional[str] = None) -> J
 
         job_request = self.rp_client.post(
             f"{self.endpoint_id}/run", 
-            request_input,
-            api_key=api_key
+            request_input
         )
-        return Job(self.endpoint_id, job_request["id"], self.rp_client, api_key=api_key)
+        return Job(self.endpoint_id, job_request["id"], self.rp_client)
 
     def run_sync(
-        self, request_input: Dict[str, Any], timeout: int = 86400, 
-        api_key: Optional[str] = None
+        self, request_input: Dict[str, Any], timeout: int = 86400
     ) -> Dict[str, Any]:
         """
         Run the endpoint with the given input synchronously.
 
         Args:
             request_input: The input to pass into the endpoint.
             timeout: Maximum time to wait for the job to complete.
-            api_key: Optional API key to use for this specific request.
         """
         if not request_input.get("input"):
             request_input = {"input": request_input}
 
         job_request = self.rp_client.post(
             f"{self.endpoint_id}/runsync", 
             request_input, 
-            timeout=timeout,
-            api_key=api_key
+            timeout=timeout
         )
 
         if job_request["status"] in FINAL_STATES:
             return job_request.get("output", None)
 
         return Job(
-            self.endpoint_id, job_request["id"], self.rp_client, api_key=api_key
+            self.endpoint_id, job_request["id"], self.rp_client
         ).output(timeout=timeout)
 
-    def health(self, timeout: int = 3, api_key: Optional[str] = None) -> Dict[str, Any]:
+    def health(self, timeout: int = 3) -> Dict[str, Any]:
         """
         Check the health of the endpoint (number/state of workers, number/state of requests).
 
         Args:
             timeout: The number of seconds to wait for the server to respond before giving up.
-            api_key: Optional API key to use for this specific request.
         """
         return self.rp_client.get(
             f"{self.endpoint_id}/health", 
-            timeout=timeout, 
-            api_key=api_key
+            timeout=timeout
         )
 
-    def purge_queue(self, timeout: int = 3, api_key: Optional[str] = None) -> Dict[str, Any]:
+    def purge_queue(self, timeout: int = 3) -> Dict[str, Any]:
         """
         Purges the endpoint's job queue and returns the result of the purge request.
 
         Args:
             timeout: The number of seconds to wait for the server to respond before giving up.
-            api_key: Optional API key to use for this specific request.
         """
         return self.rp_client.post(
             f"{self.endpoint_id}/purge-queue", 
             data=None, 
-            timeout=timeout,
-            api_key=api_key
+            timeout=timeout
         )