Skip to content

Commit

Permalink
delete remaining runner thread code :)
Browse files Browse the repository at this point in the history
Signed-off-by: technillogue <technillogue@gmail.com>
  • Loading branch information
technillogue committed Nov 6, 2023
1 parent 2652fe4 commit 5b0b716
Showing 1 changed file with 5 additions and 11 deletions.
16 changes: 5 additions & 11 deletions python/cog/server/runner.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
import asyncio
import io
import threading
import traceback
from asyncio import Task
from datetime import datetime, timezone
from multiprocessing.pool import AsyncResult, ThreadPool
from typing import Any, Callable, Dict, Optional, Tuple

import requests
Expand Down Expand Up @@ -41,17 +39,14 @@ def __init__(
self,
*,
predictor_ref: str,
shutdown_event: Optional[threading.Event],
shutdown_event: Optional[asyncio.Event],
upload_url: Optional[str] = None,
) -> None:
self._thread = None
self._threadpool = ThreadPool(processes=1)

self._response: Optional[schema.PredictionResponse] = None
self._result: Optional[Task] = None

self._worker = Worker(predictor_ref=predictor_ref)
self._should_cancel = threading.Event()
self._should_cancel = asyncio.Event()

self._shutdown_event = shutdown_event
self._upload_url = upload_url
Expand Down Expand Up @@ -135,8 +130,7 @@ def is_busy(self) -> bool:

def shutdown(self) -> None:
self._worker.terminate()
self._threadpool.terminate()
self._threadpool.join()
# TODO: cancel setup or predict task

def cancel(self, prediction_id: Optional[str] = None) -> None:
if not self.is_busy():
Expand Down Expand Up @@ -316,7 +310,7 @@ async def predict(
worker: Worker,
request: schema.PredictionRequest,
event_handler: PredictionEventHandler,
should_cancel: threading.Event,
should_cancel: asyncio.Event,
) -> schema.PredictionResponse:
# Set up logger context within prediction thread.
structlog.contextvars.clear_contextvars()
Expand All @@ -341,7 +335,7 @@ async def _predict(
worker: Worker,
request: schema.PredictionRequest,
event_handler: PredictionEventHandler,
should_cancel: threading.Event,
should_cancel: asyncio.Event,
) -> schema.PredictionResponse:
initial_prediction = request.dict()

Expand Down

0 comments on commit 5b0b716

Please sign in to comment.