vastai-sdk 0.4.2.dev2__tar.gz → 0.4.2.dev3__tar.gz

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vastai-sdk
-Version: 0.4.2.dev2
+Version: 0.4.2.dev3
 Summary: SDK for Vast.ai GPU Cloud Service
 License-File: LICENSE
 Author: Chris McKenzie

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vastai-sdk"
-version = "0.4.2.dev2"
+version = "0.4.2.dev3"
 description = "SDK for Vast.ai GPU Cloud Service"
 readme = "README.md"
 authors = [

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/client/client.py

@@ -136,7 +136,6 @@ class Serverless:
 
         return self._ssl_context
 
-
     async def get_endpoint(self, name="") -> Endpoint:
         endpoints = await self.get_endpoints()
         for e in endpoints:
@@ -179,17 +178,6 @@ class Serverless:
                 raise RuntimeError(f"get_endpoint_workers failed: HTTP {resp.status} - {text}")
 
             data = await resp.json(content_type=None)
-
-            # If error message from authenticate_endpoint_apikey_by_id occurs, there is a possibility that
-            # the endpoint's worker instances are not ready to be queried. If an error message occurs,
-            # return an empty list and print the error message to the user. The endpoint get_endpoint_workers
-            # should normally return a list of dictionaries containing worker instance information.
-            if isinstance(data,dict):
-                if 'error_msg' in data.keys():
-                    self.logger.warning(f"Received the following error from get_endpoint_workers:{data['error_msg']}.\nEndpoint may not be ready for query. Check credentials or wait a few minutes and try again.")
-                    return []
-
-            
             if not isinstance(data, list):
                 raise RuntimeError(f"Unexpected response type (wanted list): {type(data)}")
 
@@ -296,7 +284,7 @@ class Serverless:
         session: Session = None,
         serverless_request: Optional[ServerlessRequest] = None,
         cost: int = 100,
-        timeout: Optional[float] = None,
+        max_wait_time: Optional[float] = None,
         retry: bool = True,
         max_retries: int = None,
         stream: bool = False
@@ -308,7 +296,6 @@ class Serverless:
         async def task(request: ServerlessRequest):
             request_idx: int = 0
             total_attempts = 0
-            start_time = time.time()
             try:
                 while True:
                     total_attempts += 1
@@ -317,10 +304,6 @@ class Serverless:
                     auth_data = {}
                     session_id = None
 
-                    # Check total elapsed time
-                    if timeout is not None and (time.time() - start_time) >= timeout:
-                        raise asyncio.TimeoutError(f"Timed out after {time.time() - start_time:.1f}s waiting for worker")
-
                     if session is None:
                         if request_idx == 0:
                             self.logger.debug(f"Sending initial route call for request_idx {request_idx}")
@@ -336,21 +319,19 @@ class Serverless:
                             self.logger.error("Did not get request_idx from initial route")
 
                         poll_interval = 1
-                        poll_elapsed = 0
+                        elapsed_time = 0
                         attempt = 0
                         while route.status != "READY":
                             request.status = "Polling"
-
-                            # Check total elapsed time
-                            if timeout is not None and (time.time() - start_time) >= timeout:
-                                raise asyncio.TimeoutError(f"Timed out after {time.time() - start_time:.1f}s waiting for worker to become ready")
+                            if max_wait_time is not None and elapsed_time >= max_wait_time:
+                                raise asyncio.TimeoutError("Timed out waiting for worker to become ready")
 
                             await asyncio.sleep(poll_interval)
-                            poll_elapsed += poll_interval
+                            elapsed_time += poll_interval
 
                             route = await endpoint._route(cost=cost, req_idx=request_idx, timeout=60.0)
                             request_idx = route.request_idx or request_idx
-
+                            
                             attempt += 1
                             poll_interval = random.uniform(0.1, min((2 ** attempt) + random.uniform(0, 1), self.max_poll_interval))
                             self.logger.debug(f"Polling route, attempt {attempt}")
@@ -385,7 +366,7 @@ class Serverless:
                             body=worker_request_body,
                             method="POST",
                             retries=1,  # avoid stacking retries with the outer loop
-                            timeout=600,
+                            timeout=30,
                             stream=stream
                         )
                     except Exception as ex:
@@ -396,10 +377,6 @@ class Serverless:
 
                     if not result.get("ok"):
                         if retry and result.get("retryable") and (max_retries is None or total_attempts < max_retries):
-                            # Check if we have time left before retrying
-                            if timeout is not None and (time.time() - start_time) >= timeout:
-                                raise asyncio.TimeoutError(f"Request timed out after {time.time() - start_time:.1f}s")
-
                             request.status = "Retrying"
                             await asyncio.sleep(min((2 ** total_attempts) + random.uniform(0, 1), self.max_poll_interval))
                             continue
@@ -412,9 +389,7 @@ class Serverless:
 
                         response = {
                             "response": result.get("json") if result.get("json") is not None else {"error": result.get("text", "")},
-                            "ok": result.get("ok"),
-                            "status": result.get("status"),
-                            "text" : result.get("text"),
+                            "result": result,
                             "latency": (request.complete_time - request.start_time) if request.start_time else None,
                             "url": worker_url,
                             "request_idx": request_idx,
@@ -433,9 +408,7 @@ class Serverless:
 
                     response = {
                         "response": worker_response,
-                        "ok" : result.get("ok"),
-                        "status" : result.get("status"),
-                        "text" : result.get("text"),
+                        "result": result,
                         "latency": request.complete_time - request.start_time,
                         "url": worker_url,
                         "request_idx": request_idx,
@@ -450,7 +423,6 @@ class Serverless:
             except Exception as ex:
                 request.status = "Errored"
                 self.logger.error(f"Request errored: {ex}")
-                request.set_exception(ex)
                 return
 
         bg_task = asyncio.create_task(task(serverless_request))

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/client/connection.py

@@ -3,7 +3,7 @@ import aiohttp
 import asyncio
 import random
 import json
-from typing import AsyncIterator, Dict, Optional, Any
+from typing import AsyncIterator, Dict, Optional, Union, Any
 
 _JITTER_CAP_SECONDS = 5.0
 
@@ -198,10 +198,6 @@ async def _make_request(
                 last_result["stream"] = _stream_iter()
                 return last_result
 
-            except asyncio.TimeoutError as ex:
-                if attempt == retries:
-                    raise TimeoutError(f"Request to {full_url} timed out after {timeout}s") from ex
-                await asyncio.sleep(_backoff_delay(attempt))
             except Exception as ex:
                 if attempt == retries:
                     raise ex
@@ -261,10 +257,6 @@ async def _make_request(
 
                 return result
 
-        except asyncio.TimeoutError as ex:
-            if attempt == retries:
-                raise TimeoutError(f"Request to {full_url} timed out after {timeout}s") from ex
-            await asyncio.sleep(_backoff_delay(attempt))
         except Exception as ex:
             if attempt == retries:
                 raise ex

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/client/endpoint.py

@@ -24,7 +24,7 @@ class Endpoint:
         self.id = id
         self.api_key = api_key
 
-    def request(self, route, payload, serverless_request=None, cost: int = 100, retry: bool = True, stream: bool = False, timeout: float = None, session: "Session" = None):
+    def request(self, route, payload, serverless_request=None, cost: int = 100, retry: bool = True, stream: bool = False, session: "Session" = None):
         return self.client.queue_endpoint_request(
             endpoint=self,
             worker_route=route,
@@ -33,12 +33,13 @@ class Endpoint:
             cost=cost,
             retry=retry,
             stream=stream,
-            timeout=timeout,
             session=session
         )
 
     def close_session(self, session: "Session"):
-        return self.client.end_endpoint_session(session=session)
+        return self.client.end_endpoint_session(
+            session=session
+        )
 
     async def session_healthcheck(self, session: "Session"):
         result = await self.client.get_endpoint_session(
@@ -85,7 +86,7 @@ class Endpoint:
                 },
                 method="POST",
                 timeout=10.0,
-                retries=1,
+                retries=5,
                 stream=False,
             )
         except Exception as ex:
@@ -95,7 +96,8 @@ class Endpoint:
             raise RuntimeError(f"Failed to route endpoint: HTTP {result.get('status')} - {result.get('text','')[:512]}")
 
         return RouteResponse(result.get("json") or {})
-    
+
+
 class RouteResponse:
     status: str
     body: dict

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/server/lib/backend.py

@@ -5,7 +5,7 @@ import base64
 import subprocess
 import dataclasses
 import logging
-from asyncio import sleep, gather, Semaphore, create_task
+from asyncio import wait, sleep, gather, Semaphore, FIRST_COMPLETED, create_task
 from typing import Tuple, Awaitable, NoReturn, List, Union, Callable, Optional, Any, Dict
 from functools import cached_property
 from distutils.util import strtobool
@@ -35,8 +35,9 @@ from .data_types import (
     Session
 )
 
-VERSION = "1.1.0"
+VERSION = "1.0.1"
 
+MSG_HISTORY_LEN = 100
 log = logging.getLogger(__file__)
 
 # defines the minimum wait time between sending updates to autoscaler
@@ -65,6 +66,7 @@ class Backend:
     log_actions: List[Tuple[LogAction, str]]
     reqnum = -1
     version = VERSION
+    msg_history = []
     sem: Semaphore = dataclasses.field(default_factory=Semaphore)
     queue: deque = dataclasses.field(default_factory=deque, repr=False)
     unsecured: bool = dataclasses.field(
@@ -184,7 +186,8 @@ class Backend:
             if session is None:
                 return False
 
-            # Cancel all in-flight request handler tasks
+            session.cancel_event.set()
+
             for req in list(session.requests):
                 try:
                     tr = getattr(req, "transport", None)
@@ -193,15 +196,15 @@ class Backend:
                 except Exception:
                     pass
             session.requests.clear()
+
             request_metrics = self.session_metrics.pop(session_id, None)
 
-        # Run the on_close callback
         try:
             await self.__run_session_on_close(session)
         except Exception:
             pass
 
-        # Update metrics outside lock
+        # metrics outside lock
         if request_metrics is not None:
             self.metrics._request_success(request_metrics)
             self.metrics._request_end(request_metrics)
@@ -310,8 +313,6 @@ class Backend:
         self._total_pubkey_fetch_errors = 0
         self._pubkey = self._fetch_pubkey()
         self.__start_healthcheck: bool = False
-        self.__healthcheck_ready: asyncio.Event = asyncio.Event()
-        self.__healthcheck_succeeded: bool = False
 
     @property
     def pubkey(self) -> Optional[RSA.RsaKey]:
@@ -374,6 +375,7 @@ class Backend:
         request_metrics: RequestMetrics = RequestMetrics(request_idx=auth_data.request_idx, reqnum=auth_data.reqnum, workload=workload, status="Created")
 
         event = asyncio.Event()
+        finished = asyncio.Event()
 
         session = None
         if session_id is not None:
@@ -398,6 +400,14 @@ class Backend:
                 except ValueError:
                     pass
 
+        async def cancel_api_call_if_disconnected() -> None:
+            await request.wait_for_disconnection()
+            if not finished.is_set():
+                self.metrics._request_canceled(request_metrics)
+        
+        async def cancel_if_session_closed() -> None:
+            await session.cancel_event.wait()
+
         async def make_request() -> Union[web.Response, web.StreamResponse]:
             try:
                 response = await self.__call_backend(handler=handler, payload=payload)
@@ -419,7 +429,13 @@ class Backend:
         if handler.max_queue_time is not None and self.metrics.model_metrics.wait_time > handler.max_queue_time:
             self.metrics._request_reject(request_metrics)
             return web.Response(status=429)
+        
+        disconnect_task = create_task(cancel_api_call_if_disconnected())
+        session_cancel_task = None
+        if session is not None:
+            session_cancel_task = create_task(cancel_if_session_closed())
 
+        next_request_task = None
         work_task = None
 
         self.metrics._request_start(request_metrics, session)
@@ -427,7 +443,24 @@ class Backend:
         try:
             if handler.allow_parallel_requests:
                 work_task = create_task(make_request())
-                # Handler cancellation will raise CancelledError on client disconnect
+
+                wait_set = [work_task, disconnect_task]
+                if session_cancel_task is not None:
+                    wait_set.append(session_cancel_task)
+
+                done, pending = await wait(
+                    wait_set,
+                    return_when=FIRST_COMPLETED,
+                )
+
+                for t in pending:
+                    t.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)
+
+                if disconnect_task in done or (session_cancel_task is not None and session_cancel_task in done):
+                    return web.Response(status=499)  # request cancelled
+
+                # otherwise work_task completed
                 return await work_task
 
             # FIFO-queue branch
@@ -438,50 +471,80 @@ class Backend:
                 if self.queue and self.queue[0] is event:
                     event.set()
 
-                # Wait for our turn - CancelledError raised if client disconnects
-                await event.wait()
+                # Race between our request being next and request being cancelled
+                next_request_task = create_task(event.wait())
+
+                wait_set = [next_request_task, disconnect_task]
+                if session_cancel_task is not None:
+                    wait_set.append(session_cancel_task)
+
+                first_done, first_pending = await wait(
+                    wait_set,
+                    return_when=FIRST_COMPLETED,
+                )
+                # If the disconnect task wins the race
+                if disconnect_task in first_done or (session_cancel_task is not None and session_cancel_task in first_done):
+                    # Clean up the next_request_task, then exit
+                    for t in first_pending:
+                        t.cancel()
+                    await asyncio.gather(*first_pending, return_exceptions=True)
+                    return web.Response(status=499)
 
                 # We are the next-up request in the queue
                 if session is not None:
                     log.debug(f"Starting work on request {request_metrics.reqnum}")
 
-                # Execute the work task
+                # Race the backend API call with the disconnect task
                 work_task = create_task(make_request())
+
+                wait_set = [work_task, disconnect_task]
+                if session_cancel_task is not None:
+                    wait_set.append(session_cancel_task)
+
+                done, pending = await wait(
+                    wait_set,
+                    return_when=FIRST_COMPLETED,
+                )
+
+                for t in pending:
+                    t.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)
+
+                if disconnect_task in done or (session_cancel_task is not None and session_cancel_task in done):
+                    return web.Response(status=499)
+
+                # otherwise work_task completed
                 return await work_task
 
         except asyncio.CancelledError:
-            # With handler_cancellation enabled, this indicates client disconnect
-            log.debug(f"Request {request_metrics.reqnum} cancelled (client disconnect)")
-            self.metrics._request_canceled(request_metrics)
             return web.Response(status=499)
-
+        
         except Exception as e:
             log.debug(f"Exception in main handler loop {e}")
             return web.Response(status=500)
         
         finally:
-            try:
-                # Remove request from session if present
-                if session is not None and session_id is not None:
-                    async with self._sessions_lock:
-                        s = self.sessions.get(session_id)
-                        if s is not None:
-                            try:
-                                s.requests.remove(request)
-                            except ValueError:
-                                pass
-
-                if not handler.allow_parallel_requests:
-                    advance_queue_after_completion(event)
-
-                self.metrics._request_end(request_metrics)
-
-                # Cleanup work task if still pending
-                if work_task and not work_task.done():
-                    work_task.cancel()
-                    await asyncio.gather(work_task, return_exceptions=True)
-            except Exception as e:
-                log.error(f"Error during request cleanup: {e}")
+            # Set finished flag so we don't cancel after completion
+            finished.set()
+
+            if session is not None and session_id is not None:
+                async with self._sessions_lock:
+                    s = self.sessions.get(session_id)
+                    if s is not None:
+                        try:
+                            s.requests.remove(request)
+                        except ValueError:
+                            pass
+
+            if not handler.allow_parallel_requests:
+                advance_queue_after_completion(event)
+            self.metrics._request_end(request_metrics, session)
+            cleanup_tasks = [t for t in (next_request_task, work_task, disconnect_task, session_cancel_task) if t]
+            for t in cleanup_tasks:
+                if not t.done():
+                    t.cancel()
+            if cleanup_tasks:
+                await asyncio.gather(*cleanup_tasks, return_exceptions=True)
 
     async def __healthcheck(self) -> None:
         """
@@ -513,17 +576,12 @@ class Backend:
 
                     if status == 200:
                         log.debug("Healthcheck successful")
-                        if not self.__healthcheck_succeeded:
-                            self.__healthcheck_succeeded = True
-                            self.__healthcheck_ready.set()
-                            log.debug("First healthcheck succeeded - model is ready")
-                    else:
+                    elif status == 503:
                         msg = f"Healthcheck failed with status: {status}"
                         log.debug(msg)
-                        # Only report error if we've already had a successful healthcheck
-                        # (i.e., model was working but now is broken)
-                        if self.__healthcheck_succeeded:
-                            self.backend_errored(msg)
+                        self.backend_errored(msg)
+                    else:
+                        log.debug(f"Healthcheck endpoint not ready: {status}")
 
             except CancelledError:
                 log.debug("Healthcheck task cancelled; exiting loop")
@@ -531,10 +589,7 @@ class Backend:
 
             except Exception as e:
                 log.debug(f"Healthcheck failed with exception: {e}")
-                # Only report connection errors AFTER the first successful healthcheck
-                # During startup, connection failures are expected
-                if self.__healthcheck_succeeded:
-                    self.backend_errored(str(e))
+                self.backend_errored(str(e))
 
     async def _start_tracking(self) -> None:
         await gather(
@@ -600,11 +655,17 @@ class Backend:
                 return False
 
         message = {
-            "url" : auth_data.url
+            key: value
+            for (key, value) in (dataclasses.asdict(auth_data).items())
+            if key != "signature" and key != "__request_id"
         }
-
-        if verify_signature(json.dumps(message, indent=4, sort_keys=True), auth_data.signature):
+        if auth_data.reqnum < (self.reqnum - MSG_HISTORY_LEN):
+            log.error(f"Signature error: reqnum failure, got {auth_data.reqnum}, current_reqnum: {self.reqnum}")
+            return False
+        elif verify_signature(json.dumps(message, indent=4, sort_keys=True), auth_data.signature):
             self.reqnum = max(auth_data.reqnum, self.reqnum)
+            self.msg_history.append(message)
+            self.msg_history = self.msg_history[-MSG_HISTORY_LEN:]
             return True
         else:
             log.error(f"Signature error: signature verification failed, sig:{auth_data.signature}, message: {message}")
@@ -688,24 +749,12 @@ class Backend:
                         log.debug(
                             f"Got log line indicating model is loaded: {log_line}"
                         )
+                        # some backends need a few seconds after logging successful startup before
+                        # they can begin accepting requests
+                        # await sleep(5)
                         try:
                             max_throughput = await run_benchmark()
                             self.__start_healthcheck = True
-
-                            # Wait for the first successful healthcheck before marking model as loaded
-                            if self.healthcheck_url:
-                                log.debug("Benchmark succeeded, waiting for healthcheck to confirm model is ready...")
-                                try:
-                                    await asyncio.wait_for(self.__healthcheck_ready.wait(), timeout=300.0)
-                                    log.debug("Healthcheck confirmed - marking model as loaded")
-                                except asyncio.TimeoutError:
-                                    raise Exception("Timed out waiting for healthcheck after benchmark (waited 300s)")
-                            else:
-                                # No healthcheck endpoint defined, wait 10 seconds as fallback
-                                log.debug("No healthcheck endpoint defined, waiting 10 seconds before marking model as loaded...")
-                                await asyncio.sleep(10)
-                                log.debug("Wait complete - marking model as loaded")
-
                             self.metrics._model_loaded(
                                 max_throughput=max_throughput,
                             )

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/server/lib/data_types.py

@@ -353,5 +353,6 @@ class Session:
     on_close_payload: dict
     requests: list[web.Request] = field(default_factory=list)
     created_at: float = field(default_factory=time.time)
+    cancel_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
 
 

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/serverless/server/lib/server.py

@@ -37,7 +37,7 @@ async def start_server_async(backend: Backend, routes: List[web.RouteDef], **kwa
         app.router.add_post("/session/get", backend.session_get_handler)
         app.router.add_post("/session/health", backend.session_health_handler)
 
-        runner = web.AppRunner(app, handler_cancellation=True)
+        runner = web.AppRunner(app)
         await runner.setup()
         site = web.TCPSite(
             runner,

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/vast.py

@@ -3960,10 +3960,10 @@ def show__earnings(args):
     :rtype:
     """
 
-    Minutes = 60.0
-    Hours	= 60.0*Minutes
-    Days	= 24.0*Hours
-    Years	= 365.0*Days
+    Minutes = 60.0;
+    Hours	= 60.0*Minutes;
+    Days	= 24.0*Hours;
+    Years	= 365.0*Days;
     cday    = time.time() / Days
     sday = cday - 1.0
     eday = cday - 1.0
@@ -3979,7 +3979,7 @@ def show__earnings(args):
         try:
             end_date = dateutil.parser.parse(str(args.end_date))
             end_date_txt = end_date.isoformat()
-            end_timestamp = end_date.timestamp()
+            end_timestamp = time.mktime(end_date.timetuple())
             eday = end_timestamp / Days
         except ValueError as e:
             print(f"Warning: Invalid end date format! Ignoring end date! \n {str(e)}")
@@ -3988,20 +3988,21 @@ def show__earnings(args):
         try:
             start_date = dateutil.parser.parse(str(args.start_date))
             start_date_txt = start_date.isoformat()
-            start_timestamp = start_date.timestamp()
+            start_timestamp = time.mktime(start_date.timetuple())
             sday = start_timestamp / Days
-        except ValueError as e:
+        except ValueError:
             print(f"Warning: Invalid start date format! Ignoring start date! \n {str(e)}")
 
+
+
     req_url = apiurl(args, "/users/me/machine-earnings", {"owner": "me", "sday": sday, "eday": eday, "machid" :args.machine_id});
     r = http_get(args, req_url)
     r.raise_for_status()
     rows = r.json()
 
-    if args.raw:
-        return rows
     print(json.dumps(rows, indent=1, sort_keys=True))
 
+
 def sum(X, k):
     y = 0
     for x in X:

{vastai_sdk-0.4.2.dev2 → vastai_sdk-0.4.2.dev3}/vastai/vastai_sdk.py

@@ -348,21 +348,18 @@ class VastAI(VastAIBase):
 
         sig = getattr(func, "mysignature", None)
         sig_help = getattr(func, "mysignature_help", None)
-
         if sig:
             wrapper.__signature__, docappend = self.generate_signature_from_argparse(sig)
+            epi = None
 
-            # append epilog if exists
-            if getattr(sig, "epilog", None):
-                wrapper.__doc__ = f"{wrapper.__doc__.rstrip()}\n\n{sig.epilog.strip()}\n"
-            
-            # if no epilog or func docstring, fall back to parser help text
-            elif sig_help and not hasDoc:
-                wrapper.__doc__ += f"\n\n{sig_help}"
-
-            # finally append the arg details
-            wrapper.__doc__ = f"{wrapper.__doc__.rstrip()}\n\n{docappend}"
+            if sig.epilog:
+                epi = re.sub('Example.?:.*', '', sig.epilog, flags=re.DOTALL|re.M).strip()
+                wrapper.__doc__ += epi
 
+            if not (epi or hasDoc) and sig_help:
+                wrapper.__doc__ += sig_help
+            
+            wrapper.__doc__ = '\n\n'.join([ wrapper.__doc__.rstrip(), docappend ])
         return wrapper
 
     def credentials_on_disk(self):