Skip to content

Commit

Permalink
Add Setup Logging (#2018)
Browse files Browse the repository at this point in the history
* Explicitly log the following setup states:
1. Starting
2. Started
3. Succeeded
4. Failed
* This gives us better visibility into how long
each of these setup pieces takes.
  • Loading branch information
8W9aG authored Oct 28, 2024
1 parent 5c1908f commit 5e2218f
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion python/cog/server/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@
from urllib3.util.retry import Retry

from .. import schema
from ..base_input import BaseInput
from ..files import put_file_to_signed_endpoint
from ..json import upload_files
from ..predictor import BaseInput
from ..types import PYDANTIC_V2
from .errors import FileUploadError, RunnerBusyError, UnknownPredictionError
from .eventtypes import Done, Log, PredictionOutput, PredictionOutputType
Expand Down Expand Up @@ -163,6 +163,7 @@ def result(self) -> T:

class SetupTask(Task[SetupResult]):
def __init__(self, _clock: Optional[Callable[[], datetime]] = None) -> None:
log.info("starting setup")
self._clock = _clock
if self._clock is None:
self._clock = lambda: datetime.now(timezone.utc)
Expand All @@ -175,6 +176,7 @@ def result(self) -> SetupResult:
return self._result

def track(self, fut: "Future[Done]") -> None:
log.info("started setup")
self._fut = fut
self._fut.add_done_callback(self._handle_done)

Expand All @@ -194,11 +196,13 @@ def append_logs(self, message: str) -> None:
self._result.logs.append(message)

def succeeded(self) -> None:
log.info("setup succeeded")
assert self._clock
self._result.completed_at = self._clock()
self._result.status = schema.Status.SUCCEEDED

def failed(self) -> None:
log.info("setup failed")
assert self._clock
self._result.completed_at = self._clock()
self._result.status = schema.Status.FAILED
Expand Down

0 comments on commit 5e2218f

Please sign in to comment.