Merge branch 'dev' of github.com:jhirschibar/option_bot into dev

Author: jhirschibar

.gitignore

@@ -6,4 +6,4 @@
 data/*
 *data/
 .python-version
-.models
+*.code-workspace

LICENSE

@@ -22,9 +22,7 @@
     R,
     T,
     TaskID,
-    TracebackStr,
 )
-from data_pipeline.polygon_utils import HistoricalQuotes
 
 from option_bot.proj_constants import log
 
@@ -123,11 +121,12 @@ def __init__(
             session_base_url=session_base_url,
         )
         self.o_ticker_count_mapping = o_ticker_count_mapping
-        self.o_ticker_queue_progress: Dict[str, set[int]] = {}  # list of tids per o_ticker pulled from queue
+        self.o_ticker_queue_progress: Dict[str, set[int]] = {}  # tids pulled to execute
+        self.o_ticker_skip_tids: Dict[str, set[int]] = {}  # tids pulled to skip
         self.tid_result_progress: set = set()
-        self.o_ticker: str = ""
         self.empty_tids: set = set()
         self.complete_otkrs: set[str] = set()
+        self.completely_processed_otkrs: list[str] = []
 
     async def run(self):
         if self.init_client_session:
@@ -138,6 +137,7 @@ async def run(self):
             ) as client_session:
                 pending: Dict[asyncio.Future, TaskID] = {}
                 completed: int = 0
+                skipped: int = 0
                 running = True
                 while running or pending:
                     # TTL, Tasks To Live, determines how many tasks to execute before dying
@@ -157,22 +157,26 @@ async def run(self):
 
                         tid, func, args, kwargs = task
 
-                        self.o_ticker = args[0]
-                        if self.o_ticker not in self.o_ticker_queue_progress:
-                            self.o_ticker_queue_progress[self.o_ticker] = set([tid])
+                        # tracking progress
+                        o_ticker = args[0]
+                        if o_ticker not in self.o_ticker_queue_progress:
+                            self.o_ticker_queue_progress[o_ticker] = set([tid])
+
+                        # start work on task, add to pending if not a skipped/complete otkr.
+                        # Otherwise add tid to "results" to mark as done
+                        if o_ticker not in self.complete_otkrs:
+                            self.o_ticker_queue_progress[o_ticker].add(tid)
+                            args = [
+                                *args,
+                                client_session,
+                            ]  # NOTE: adds client session to the args list
+                            future = asyncio.ensure_future(func(*args, **kwargs))
+                            pending[future] = tid
                         else:
-                            self.o_ticker_queue_progress[self.o_ticker].add(tid)
+                            self.rx.put_nowait((tid, None, None))
+                            self.o_ticker_skip_tids[o_ticker].add(tid)
+                            skipped += 1
 
-                        # start work on task, add to pending
-
-                        args = [
-                            *args,
-                            client_session,
-                        ]  # NOTE: adds client session to the args list
-                        future = asyncio.ensure_future(func(*args, **kwargs))
-                        pending[future] = tid
-
-                    # NOTE: this won't initially start processing "pending" until pending is as big as concurrency limit
                     if not pending:
                         await asyncio.sleep(0.005)
                         continue
@@ -190,10 +194,11 @@ async def run(self):
                         tb = None
                         try:
                             result = future.result()
-                            if result[0] is False:
-                                if result[1] not in self.complete_otkrs:
-                                    self.empty_tids.add(tid)
-                                result = None
+                            if result:
+                                if result[0] is False:
+                                    if result[1] not in self.complete_otkrs:
+                                        self.empty_tids.add(tid)
+                            result = None
                         except BaseException as e:
                             if self.exception_handler is not None:
                                 self.exception_handler(e)
@@ -203,25 +208,25 @@ async def run(self):
                         self.tid_result_progress.add(tid)
                         completed += 1
 
-                    k = 15  # indicator that we've passed the listing date for the option
-                    if len(self.empty_tids) > k:
-                        seq_start = self.has_consecutive_sequence(k=k)
-                        if seq_start:
-                            empty_otkr = self.check_completed_otkr(seq_start)
-                            await self.clean_up_queue(empty_otkr)
+                    self.eval_list_date()
+
                     self.clean_o_ticker_progress()
-                    log.debug(
-                        f"loop complete, restarting while again, running: {running}, total processed: {completed}"
-                    )
 
-        log.info(f"worker finished: processed {completed} tasks")
+        log.info(f"worker finished: processed {completed} tasks, and skipped {skipped}")
+
+    def eval_list_date(self):
+        k = 15  # indicator that we've passed the listing date for the option
+        if len(self.empty_tids) > k:
+            seq_start = self.has_consecutive_sequence(k=k)
+            if seq_start:
+                self.check_completed_otkr(seq_start)
 
     def has_consecutive_sequence(self, k=15) -> int | bool:
-        """check if there is a sequence of length 16 or longer in which the tids are consecutive"""
+        """check if there is a sequence of length k or longer in which the tids are consecutive"""
         for tid in self.empty_tids:
             if all((tid + i) in self.empty_tids for i in range(k)):
-                # log.debug(f"consecutive sequence found with {len(self.empty_tids)} empty tids")
-                # log.debug(f"empty tids: {self.empty_tids}")
+                log.debug(f"consecutive sequence of {k} found within {len(self.empty_tids)} total empty tids")
+                log.debug(f"empty tids: {self.empty_tids}")
                 return tid
         return False
 
@@ -231,45 +236,41 @@ def check_completed_otkr(self, tid):
         for otkr in self.o_ticker_queue_progress:
             if tid in self.o_ticker_queue_progress[otkr]:
                 now_empty_otkr = otkr
+                break
         if now_empty_otkr:
             self.complete_otkrs.add(now_empty_otkr)
-        return now_empty_otkr
-
-    async def clean_up_queue(self, otkr: str):
-        """identifies the o_ticker that has the tid with the consecutive sequence.
-        It calculates the remaining tasks that have to be pulled from the queue for the oticker and removes them.
-        Everything that has already been pulled and is in `pending` will still be processed.
-        Removes completed
-        """
-        done_tids = self.o_ticker_queue_progress[otkr]
-        remaining_tasks = self.o_ticker_count_mapping[otkr] - len(done_tids)
-        i = 0
-        while i < remaining_tasks:
-            try:
-                self.tx.get_nowait()
-                i += 1
-            except queue.Empty:
-                await asyncio.sleep(0.001)
-
-        self.empty_tids = self.empty_tids - done_tids
-        self.o_ticker_count_mapping[otkr] -= remaining_tasks
+            self.o_ticker_skip_tids[now_empty_otkr] = set()
+            self.empty_tids -= self.o_ticker_queue_progress[now_empty_otkr]
 
     def clean_o_ticker_progress(self):
-        """removes anything"""
-        completely_done_otkrs = []
-        for otkr in self.o_ticker_queue_progress:
-            if otkr in self.complete_otkrs:
-                self.empty_tids = self.empty_tids - self.o_ticker_queue_progress[otkr]
-            if len(self.o_ticker_queue_progress[otkr]) >= self.o_ticker_count_mapping[otkr]:
-                if len(self.o_ticker_queue_progress[otkr] - self.tid_result_progress) == 0:
-                    completely_done_otkrs.append(otkr)
-
-        if completely_done_otkrs:
-            for otkr in completely_done_otkrs:
-                self.tid_result_progress -= self.o_ticker_queue_progress[otkr]
-                self.o_ticker_queue_progress[otkr].clear()
-                self.o_ticker_queue_progress.pop(otkr)
-            log.info(f"\ncompletely done with otkrs: {completely_done_otkrs}!!")
+        """reports how many o_tickers have had all tids pulled from the queue and cleans internal tracking sets"""
+
+        for otkr in self.complete_otkrs:
+            total_tids = self.o_ticker_count_mapping[otkr]
+            self.empty_tids -= self.o_ticker_queue_progress[otkr]
+            if otkr not in self.completely_processed_otkrs:
+                if (
+                    len(self.o_ticker_skip_tids.get(otkr, []) | self.o_ticker_queue_progress.get(otkr, []))
+                    >= total_tids
+                ):
+                    self.completely_processed_otkrs.append(otkr)
+                    log.info(f"all processed for {otkr}! ({total_tids} tasks)")
+
+                elif (
+                    len(
+                        (
+                            self.o_ticker_queue_progress.get(otkr, set())
+                            | self.o_ticker_skip_tids.get(otkr, set())
+                        )
+                        - self.tid_result_progress
+                    )
+                    == 0
+                ):
+                    self.completely_processed_otkrs.append(otkr)
+                    log.info(f"all processed for {otkr}!! \
+({len(self.o_ticker_queue_progress.get(otkr, []))} processed, \
+{total_tids - len(self.o_ticker_queue_progress.get(otkr, []))} will be skipped, \
+{total_tids} expected)")
 
 
 class QuotePool(Pool):
@@ -287,7 +288,6 @@ def __init__(
         init_client_session: bool = False,
         session_base_url: Optional[str] = None,
         o_ticker_count_mapping: Dict[str, int] = None,
-        paginator: HistoricalQuotes = None,
     ) -> None:
         self.o_ticker_count_mapping: dict[str, int] = o_ticker_count_mapping
         scheduler = QuoteScheduler(self.o_ticker_count_mapping)
@@ -338,9 +338,9 @@ def queue_work(
 
         return task_id
 
-    def finish_work(self, task_id: TaskID, value: Any, tb: Optional[TracebackStr]):
-        """overwriting the inherited function. Not using ._results in the pool"""
-        self.scheduler.complete_task(task_id)
+    # def finish_work(self, task_id: TaskID, value: Any, tb: Optional[TracebackStr]):
+    #     """overwriting the inherited function. Not using ._results in the pool"""
+    #     self.scheduler.complete_task(task_id)
 
     def starmap(
         self,
@@ -373,7 +373,6 @@ def create_worker(
 
         :meta private:
         """
-        # paginator = copy.deepcopy(self.paginator)
         tx, rx = self.queues[qid]
         process = QuoteWorker(
             tx,
@@ -387,7 +386,6 @@ def create_worker(
             init_client_session=self.init_client_session,
             session_base_url=self.session_base_url,
             o_ticker_count_mapping=self.o_ticker_count_mapping,
-            # paginator=paginator,
         )
         process.start()
         return process

option_bot.code-workspace

@@ -1,4 +1,3 @@
-from asyncio import Queue
 from datetime import datetime
 
 from aiomultiprocess import Pool
@@ -128,37 +127,29 @@ async def download_options_snapshots(o_tickers: list[OptionTicker]):
     await api_pool_downloader(paginator=op_snapshots, pool_kwargs=pool_kwargs, args_data=o_tickers)
 
 
-async def download_options_quotes(
-    tickers: list[str], o_tickers: list[OptionTicker], queue: Queue, months_hist: int = 24
-):
+async def download_options_quotes(ticker: str, o_tickers: list[OptionTicker], months_hist: int = 24):
     """This function downloads options quotes from polygon and stores it as local json.
 
     Args:
         o_tickers: list of OptionTicker tuples
         month_hist: number of months of history to pull
     """
-    pool_kwargs = {"childconcurrency": 100, "maxtasksperchild": 1000, "processes": 30}
+    pool_kwargs = {"childconcurrency": 40, "maxtasksperchild": 1000, "processes": 30}
     o_ticker_lookup = {x.o_ticker: x.id for x in o_tickers}
     op_quotes = HistoricalQuotes(months_hist=months_hist, o_ticker_lookup=o_ticker_lookup)
-    ticker_counter = 0
     BATCH_SIZE_OTICKERS = 1000
-    for ticker in tickers:
-        ticker_counter += 1
-        batch_o_tickers = [x for x in o_tickers if x.underlying_ticker == ticker]
-        for i in range(0, len(batch_o_tickers), BATCH_SIZE_OTICKERS):
-            small_batch = batch_o_tickers[i : i + BATCH_SIZE_OTICKERS]
-            log.info(
-                f"downloading quotes for {ticker} ({ticker_counter}/{len(tickers)}) \
-                with {i+BATCH_SIZE_OTICKERS}/{len(batch_o_tickers)} o_tickers"
-            )
-            await api_quote_downloader(
-                paginator=op_quotes,
-                pool_kwargs=pool_kwargs,
-                args_data=small_batch,
-            )
-            log.info(f"Completely done downloading {ticker}. \nPassing {ticker} to upload queue")
-            await queue.put(ticker)
-    await queue.put(None)
+    batch_o_tickers = [x for x in o_tickers if x.underlying_ticker == ticker]
+    for i in range(0, len(batch_o_tickers), BATCH_SIZE_OTICKERS):
+        small_batch = batch_o_tickers[i : i + BATCH_SIZE_OTICKERS]
+
+        log.info(f"downloading quotes for {i+BATCH_SIZE_OTICKERS}/{len(batch_o_tickers)} batch of o_tickers")
+
+        await api_quote_downloader(
+            paginator=op_quotes,
+            pool_kwargs=pool_kwargs,
+            args_data=small_batch,
+        )
+    log.info(f"-- Completely done downloading {ticker}")
 
 
 async def api_quote_downloader(
@@ -182,13 +173,11 @@ async def api_quote_downloader(
     url_args, o_ticker_count_mapping = paginator.generate_request_args(args_data)
     if url_args:
         log.info("fetching data from polygon api")
-        log.debug(f"tasks: {len(url_args)}")
+        log.info(f"tasks: {len(url_args)}")
         pool_kwargs = dict(**pool_kwargs, **{"init_client_session": True, "session_base_url": POLYGON_BASE_URL})
         pool_kwargs = pool_kwarg_config(pool_kwargs)
         log.info("creating quote pool")
-        async with QuotePool(
-            **pool_kwargs, o_ticker_count_mapping=o_ticker_count_mapping, paginator=paginator
-        ) as pool:
+        async with QuotePool(**pool_kwargs, o_ticker_count_mapping=o_ticker_count_mapping) as pool:
             log.info("deploying QuoteWorkers in Pool")
             await pool.starmap(paginator.download_data, url_args)