From 1061d9ddcdb53497de7c641fc0ab7e9539b2fc81 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 3 Dec 2025 17:21:42 +0100 Subject: [PATCH 1/2] locking --- exca/map.py | 106 +++++++++++++++++++++++++++++++++++-------- exca/test_locking.py | 93 +++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + 3 files changed, 180 insertions(+), 20 deletions(-) create mode 100644 exca/test_locking.py diff --git a/exca/map.py b/exca/map.py index 2b786086..52308872 100644 --- a/exca/map.py +++ b/exca/map.py @@ -17,6 +17,7 @@ from concurrent import futures from pathlib import Path +import filelock import numpy as np import pydantic import submitit @@ -188,6 +189,7 @@ class MapInfra(base.BaseInfra, slurm.SubmititMixin): # - force: cache is ignored, and result is (re)computed (and cached) # - read-only: never compute anything mode: Mode = "cached" + lock_timeout: int = 3600 # internals _recomputed: tp.Set[str] = set() # for mode="force" @@ -487,6 +489,60 @@ def _method_override_futures(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp logger.debug(msg, len(uid_items), self._factory(), self.cache_dict) return (cache_dict[k] for k, _ in uid_items) + def _acquire_locks( + self, items: tp.Sequence[tp.Any], item_uid: tp.Callable[[tp.Any], str] + ) -> tuple[list[tp.Any], list[filelock.FileLock]]: + """Acquire locks for items, return (items_to_process, locks)""" + locks: list[filelock.FileLock] = [] + if self.cache_dict.folder is None: + return list(items), locks + + lock_dir = self.cache_dict.folder / "locks" + lock_dir.mkdir(exist_ok=True) + + work_items = [] + + # First pass: optimistic locking + delayed_items = [] + for item in items: + uid = item_uid(item) + lock = filelock.FileLock(lock_dir / f"{uid}.lock") + try: + lock.acquire(timeout=0.01) + # We have the lock + # Skip cache check for performance (rely on initial filter) + locks.append(lock) + work_items.append(item) + except filelock.Timeout: + # Locked by someone else + delayed_items.append(item) + + # Second pass: wait for delayed items + # If they were locked, someone is working on them. + # We wait for them to finish. If they release the lock and + # it's still not done (e.g. crash), we pick it up. + for item in delayed_items: + uid = item_uid(item) + if uid in self.cache_dict: + continue # Done + + lock = filelock.FileLock(lock_dir / f"{uid}.lock") + try: + # Wait for lock + lock.acquire(timeout=self.lock_timeout) + if uid in self.cache_dict: + lock.release() + continue # Done + locks.append(lock) + work_items.append(item) + except filelock.Timeout: + logger.warning( + f"Could not acquire lock for {uid} after {self.lock_timeout}s" + ) + # We skip it, but it will likely fail downstream if we expected it + + return work_items, locks + def _call_and_store( self, items: tp.Sequence[tp.Any], use_cache_dict: bool = True ) -> dict[str, tp.Any]: @@ -495,26 +551,36 @@ def _call_and_store( if imethod is None: raise RuntimeError(f"Infra was not applied: {self!r}") item_uid = imethod.item_uid - if items: # make sure some overlapping job did not already run stuff - keys = set(d) # update cache dict - items = [item for item in items if item_uid(item) not in keys] - if isinstance(self, slurm.SubmititMixin): # dependence to mixin - if self.workdir is not None and self.cluster is not None and items: - logger.info("Running from working directory: '%s'", os.getcwd()) - outputs = self._run_method(items) - sentinel = base.Sentinel() - with contextlib.ExitStack() as estack: - writer = d - if isinstance(d, CacheDict): - writer = estack.enter_context(d.writer()) # type: ignore - in_out = itertools.zip_longest(_set_tqdm(items), outputs, fillvalue=sentinel) - for item, output in in_out: - if item is sentinel or output is sentinel: - msg = f"Cached function did not yield exactly once per item: {item=!r}, {output=!r}" - raise RuntimeError(msg) - writer[item_uid(item)] = output - # don't return the whole cache dict if data is cached - return {} if use_cache_dict else d + locks: list[filelock.FileLock] = [] + try: + if items: # make sure some overlapping job did not already run stuff + keys = set(d) # update cache dict + items = [item for item in items if item_uid(item) not in keys] + if use_cache_dict: + items, locks = self._acquire_locks(items, item_uid) + + if isinstance(self, slurm.SubmititMixin): # dependence to mixin + if self.workdir is not None and self.cluster is not None and items: + logger.info("Running from working directory: '%s'", os.getcwd()) + outputs = self._run_method(items) + sentinel = base.Sentinel() + with contextlib.ExitStack() as estack: + writer = d + if isinstance(d, CacheDict): + writer = estack.enter_context(d.writer()) # type: ignore + in_out = itertools.zip_longest( + _set_tqdm(items), outputs, fillvalue=sentinel + ) + for item, output in in_out: + if item is sentinel or output is sentinel: + msg = f"Cached function did not yield exactly once per item: {item=!r}, {output=!r}" + raise RuntimeError(msg) + writer[item_uid(item)] = output + # don't return the whole cache dict if data is cached + return {} if use_cache_dict else d + finally: + for lock in locks: + lock.release() @dataclasses.dataclass diff --git a/exca/test_locking.py b/exca/test_locking.py new file mode 100644 index 00000000..94d44229 --- /dev/null +++ b/exca/test_locking.py @@ -0,0 +1,93 @@ +import threading +import time +import typing as tp +from pathlib import Path + +import filelock +import pydantic + +from .map import MapInfra + + +class Worker(pydantic.BaseModel): + infra: MapInfra = MapInfra() + computed: list[int] = [] + + @infra.apply(item_uid=str) + def process(self, items: tp.Sequence[int]) -> tp.Iterator[int]: + for item in items: + self.computed.append(item) + yield item * 2 + + +def test_locking_wait(tmp_path: Path) -> None: + """Test that a worker waits if item is locked""" + worker = Worker(infra={"folder": tmp_path}) + + # Get the actual cache folder (triggers creation) + cache_folder = worker.infra.cache_dict.folder + assert cache_folder is not None + (cache_folder / "locks").mkdir(parents=True, exist_ok=True) + + # Lock item "2" manually + lock_path = cache_folder / "locks" / "2.lock" + lock = filelock.FileLock(lock_path) + lock.acquire() + + # Function to run worker in a thread + def run_worker(): + # This should block on item 2 until we release it + list(worker.process([2])) + + t = threading.Thread(target=run_worker) + t.start() + + # Give it a moment to start and block + time.sleep(0.5) + assert t.is_alive() # Should be still waiting + + # Release lock + lock.release() + + # Should finish now + t.join(timeout=2.0) + assert not t.is_alive() + assert 2 in worker.computed + + +def test_locking_completed(tmp_path: Path) -> None: + """Test that if locked item is completed, second worker skips it""" + worker1 = Worker(infra={"folder": tmp_path}) + # Use same config to get same UID folder + worker2 = Worker(infra={"folder": tmp_path}) + + cache_folder = worker1.infra.cache_dict.folder + assert cache_folder is not None + (cache_folder / "locks").mkdir(parents=True, exist_ok=True) + + # Lock item "2" + lock_path = cache_folder / "locks" / "2.lock" + lock = filelock.FileLock(lock_path) + lock.acquire() + + def run_worker2(): + # Should wait then skip + list(worker2.process([2])) + + t = threading.Thread(target=run_worker2) + t.start() + + time.sleep(0.5) + assert t.is_alive() + + # Pretend we computed it and wrote to cache + with worker1.infra.cache_dict.writer() as w: + w["2"] = 4 + + # Release lock + lock.release() + + t.join(timeout=2.0) + assert not t.is_alive() + # Worker 2 should NOT have computed it + assert 2 not in worker2.computed diff --git a/pyproject.toml b/pyproject.toml index b6ce96e1..1ccf64d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "pydantic>=2.5.0", "submitit>=1.5.1", "orjson", # significantly faster than json + "filelock", ] [project.urls] From 891c3277151e5f5cc0be6586553af189013d159c Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 3 Dec 2025 18:06:46 +0100 Subject: [PATCH 2/2] Use filelocks for queues in MapInfra --- exca/map.py | 194 +++++++++++++------------------------------ exca/test_locking.py | 85 +++++++++++++++++++ exca/utils.py | 148 +++++++++++++++++++++++++++++++++ 3 files changed, 290 insertions(+), 137 deletions(-) diff --git a/exca/map.py b/exca/map.py index 52308872..886b8239 100644 --- a/exca/map.py +++ b/exca/map.py @@ -11,20 +11,16 @@ import itertools import logging import os -import pickle import typing as tp -import uuid from concurrent import futures from pathlib import Path -import filelock import numpy as np import pydantic -import submitit -from submitit.core import utils from . import base, slurm from .cachedict import CacheDict +from .utils import LockManager MapFunc = tp.Callable[[tp.Sequence[tp.Any]], tp.Iterator[tp.Any]] X = tp.TypeVar("X") @@ -66,45 +62,6 @@ def __call__(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp.Any]: return self.infra._method_override(items) -class JobChecker: - """Keeps a record of running jobs in a folder - and enables waiting for them to complete. - """ - - def __init__(self, folder: Path | str) -> None: - basefolder = utils.JobPaths.get_first_id_independent_folder(folder) - self.folder = basefolder / "running-jobs" - - def add(self, jobs: tp.Iterable[tp.Any]) -> None: - """Add jobs to the list of running jobs""" - self.folder.mkdir(exist_ok=True, parents=True) - for job in jobs: - if not job.done(): - job_path = self.folder / (uuid.uuid4().hex[:8] + ".pkl") - with job_path.open("wb") as f: - pickle.dump(job, f) - - def wait(self) -> bool: - """Wait for completion of running jobs""" - waited = False - for fp in self.folder.glob("*.pkl"): - try: # avoid concurrency issues with deleted items - with fp.open("rb") as f: - job: tp.Any = pickle.load(f) - except Exception: # pylint: disable=broad-except - continue - if not job.done(): - msg = "Waiting for completion of pre-existing map job: %s\nin '%s'" - logger.info(msg, job, self.folder) - job.wait() - waited = True - # delete the file as it is not needed anymore - fp.unlink(missing_ok=True) - if waited: - logger.info("Waiting is over") - return waited - - def to_chunks( items: tp.List[X], *, max_chunks: int | None, min_items_per_chunk: int = 1 ) -> tp.Iterator[tp.List[X]]: @@ -335,13 +292,6 @@ def _find_missing(self, items: tp.Dict[str, tp.Any]) -> tp.Dict[str, tp.Any]: if missing: if self.mode == "read-only": raise RuntimeError(f"{self.mode=} but found {len(missing)} missing items") - executor: submitit.Executor | None = self.executor() - if executor is not None: # wait for items being computed - jcheck = JobChecker(folder=executor.folder) - jcheck.wait() - # update cache dict and recheck as actual checking for keys updates the dict - keys = set(self.cache_dict) # update cache dict - missing = {k: item for k, item in missing.items() if k not in keys} if len(items) == len(missing) == 1 and self.forbid_single_item_computation: key, item = next(iter(missing.items())) raise RuntimeError( @@ -400,8 +350,6 @@ def _method_override(self, *args: tp.Any, **kwargs: tp.Any) -> tp.Iterator[tp.An # select a batch/chunk of samples_per_job items to send to a job j = executor.submit(self._call_and_store, chunk, use_cache_dict=True) jobs.append(j) - jcheck = JobChecker(folder=executor.folder) - jcheck.add(jobs) # pylint: disable=expression-not-assigned uid = self.uid() msg = "Sent %s samples for %s into %s jobs on cluster '%s' (eg: %s)" @@ -489,60 +437,6 @@ def _method_override_futures(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp logger.debug(msg, len(uid_items), self._factory(), self.cache_dict) return (cache_dict[k] for k, _ in uid_items) - def _acquire_locks( - self, items: tp.Sequence[tp.Any], item_uid: tp.Callable[[tp.Any], str] - ) -> tuple[list[tp.Any], list[filelock.FileLock]]: - """Acquire locks for items, return (items_to_process, locks)""" - locks: list[filelock.FileLock] = [] - if self.cache_dict.folder is None: - return list(items), locks - - lock_dir = self.cache_dict.folder / "locks" - lock_dir.mkdir(exist_ok=True) - - work_items = [] - - # First pass: optimistic locking - delayed_items = [] - for item in items: - uid = item_uid(item) - lock = filelock.FileLock(lock_dir / f"{uid}.lock") - try: - lock.acquire(timeout=0.01) - # We have the lock - # Skip cache check for performance (rely on initial filter) - locks.append(lock) - work_items.append(item) - except filelock.Timeout: - # Locked by someone else - delayed_items.append(item) - - # Second pass: wait for delayed items - # If they were locked, someone is working on them. - # We wait for them to finish. If they release the lock and - # it's still not done (e.g. crash), we pick it up. - for item in delayed_items: - uid = item_uid(item) - if uid in self.cache_dict: - continue # Done - - lock = filelock.FileLock(lock_dir / f"{uid}.lock") - try: - # Wait for lock - lock.acquire(timeout=self.lock_timeout) - if uid in self.cache_dict: - lock.release() - continue # Done - locks.append(lock) - work_items.append(item) - except filelock.Timeout: - logger.warning( - f"Could not acquire lock for {uid} after {self.lock_timeout}s" - ) - # We skip it, but it will likely fail downstream if we expected it - - return work_items, locks - def _call_and_store( self, items: tp.Sequence[tp.Any], use_cache_dict: bool = True ) -> dict[str, tp.Any]: @@ -551,36 +445,62 @@ def _call_and_store( if imethod is None: raise RuntimeError(f"Infra was not applied: {self!r}") item_uid = imethod.item_uid - locks: list[filelock.FileLock] = [] - try: - if items: # make sure some overlapping job did not already run stuff - keys = set(d) # update cache dict - items = [item for item in items if item_uid(item) not in keys] - if use_cache_dict: - items, locks = self._acquire_locks(items, item_uid) - - if isinstance(self, slurm.SubmititMixin): # dependence to mixin - if self.workdir is not None and self.cluster is not None and items: - logger.info("Running from working directory: '%s'", os.getcwd()) - outputs = self._run_method(items) - sentinel = base.Sentinel() - with contextlib.ExitStack() as estack: - writer = d - if isinstance(d, CacheDict): - writer = estack.enter_context(d.writer()) # type: ignore - in_out = itertools.zip_longest( - _set_tqdm(items), outputs, fillvalue=sentinel - ) - for item, output in in_out: - if item is sentinel or output is sentinel: - msg = f"Cached function did not yield exactly once per item: {item=!r}, {output=!r}" - raise RuntimeError(msg) - writer[item_uid(item)] = output - # don't return the whole cache dict if data is cached - return {} if use_cache_dict else d - finally: - for lock in locks: - lock.release() + + # Filter out already cached items + if items: + keys = set(d) + items = [item for item in items if item_uid(item) not in keys] + + # No locking needed + if not items or not use_cache_dict or self.cache_dict.folder is None: + return self._process_items(items, d, item_uid) + + # With locking: delegate to LockManager + lock_manager = LockManager( + lock_dir=self.cache_dict.folder / "locks", + item_uid=item_uid, + cache_contains=lambda uid: uid in self.cache_dict, + lock_timeout=self.lock_timeout, + ) + result = lock_manager.process_with_locks( + items, lambda itms: self._process_items(itms, d, item_uid) + ) + return {} if isinstance(d, CacheDict) else result + + def _process_items( + self, + items: tp.Sequence[tp.Any], + d: dict[str, tp.Any], + item_uid: tp.Callable[[tp.Any], str], + ) -> dict[str, tp.Any]: + """Process items and store results""" + if not items: + return {} + + if isinstance(self, slurm.SubmititMixin): # dependence to mixin + if self.workdir is not None and self.cluster is not None: + logger.info("Running from working directory: '%s'", os.getcwd()) + + outputs = self._run_method(items) + sentinel = base.Sentinel() + result = {} + + with contextlib.ExitStack() as estack: + writer = d + if isinstance(d, CacheDict): + writer = estack.enter_context(d.writer()) # type: ignore + in_out = itertools.zip_longest(_set_tqdm(items), outputs, fillvalue=sentinel) + for item, output in in_out: + if item is sentinel or output is sentinel: + msg = f"Cached function did not yield exactly once per item: {item=!r}, {output=!r}" + raise RuntimeError(msg) + uid = item_uid(item) + writer[uid] = output + if not isinstance(d, CacheDict): + # don't return the whole output dict if data is cached + result[uid] = output + + return result @dataclasses.dataclass diff --git a/exca/test_locking.py b/exca/test_locking.py index 94d44229..92f95696 100644 --- a/exca/test_locking.py +++ b/exca/test_locking.py @@ -1,3 +1,4 @@ +import os import threading import time import typing as tp @@ -17,6 +18,7 @@ class Worker(pydantic.BaseModel): def process(self, items: tp.Sequence[int]) -> tp.Iterator[int]: for item in items: self.computed.append(item) + time.sleep(0.1) # Simulate work yield item * 2 @@ -91,3 +93,86 @@ def run_worker2(): assert not t.is_alive() # Worker 2 should NOT have computed it assert 2 not in worker2.computed + + +def test_incremental_processing(tmp_path: Path) -> None: + """Test that workers can process non-overlapping items in parallel""" + worker1 = Worker(infra={"folder": tmp_path}) + worker2 = Worker(infra={"folder": tmp_path}) + + start_time = time.time() + + def run_worker1(): + # Process items [1, 2, 3, 4] + list(worker1.process([1, 2, 3, 4])) + + def run_worker2(): + # Process items [3, 4, 5, 6] - overlaps with worker1 on [3, 4] + time.sleep(0.05) # Start slightly after worker1 + list(worker2.process([3, 4, 5, 6])) + + t1 = threading.Thread(target=run_worker1) + t2 = threading.Thread(target=run_worker2) + + t1.start() + t2.start() + + t1.join() + t2.join() + + elapsed = time.time() - start_time + + # Worker1 processes [1, 2, 3, 4] + # Worker2 should process [5, 6] while worker1 works on [1, 2] + # Then worker2 gets [3, 4] from cache (worker1 already computed them) + # So worker2 only computes [5, 6] + + # Check that worker2 only computed non-overlapping items + assert 5 in worker2.computed + assert 6 in worker2.computed + # Worker2 should NOT have computed 3 or 4 (worker1 did) + assert 3 not in worker2.computed + assert 4 not in worker2.computed + + # Verify all items were processed by someone + all_results = set(worker1.computed) | set(worker2.computed) + assert all_results >= {1, 2, 3, 4, 5, 6} + + # Parallel execution should be faster than sequential + # Each worker processes 0.1s per item + # Worker1: 4 items = 0.4s + # Worker2: 2 items (5, 6) = 0.2s, runs in parallel with worker1 + # Total should be ~0.5s if incremental, ~0.6s if sequential + assert elapsed < 1.0, f"Took {elapsed}s, should be < 1.0s if processing incrementally" + + +def test_stale_lock_detection(tmp_path: Path) -> None: + """Test that stale locks from crashed workers are detected and removed""" + worker = Worker( + infra={"folder": tmp_path, "lock_timeout": 2} + ) # Short timeout for test + + cache_folder = worker.infra.cache_dict.folder + assert cache_folder is not None + lock_dir = cache_folder / "locks" + lock_dir.mkdir(parents=True, exist_ok=True) + + # Create a "stale" lock by creating the lock file and making it old + stale_lock_path = lock_dir / "5.lock" + stale_lock_path.touch() + + # Make the lock file old (older than lock_timeout) + old_time = time.time() - 5 # 5 seconds ago + os.utime(stale_lock_path, (old_time, old_time)) + + # Worker should detect stale lock, remove it, and process the item + start = time.time() + result = list(worker.process([5])) + elapsed = time.time() - start + + # Should complete quickly (not wait for full timeout) + assert elapsed < 1.0, f"Took {elapsed}s, should not wait for stale lock" + + # Item should have been processed + assert 5 in worker.computed + assert result == [10] # 5 * 2 diff --git a/exca/utils.py b/exca/utils.py index 0ad8ae03..95367c63 100644 --- a/exca/utils.py +++ b/exca/utils.py @@ -11,10 +11,12 @@ import os import shutil import sys +import time import typing as tp import uuid from pathlib import Path +import filelock import numpy as np import pydantic @@ -29,6 +31,152 @@ T = tp.TypeVar("T", bound=pydantic.BaseModel) +class LockManager: + """Manages distributed file-based locking for item processing. + + Handles lock acquisition, stale lock detection, and incremental processing + across multiple workers that may have overlapping item assignments. + + Parameters + ---------- + lock_dir: Path + directory where lock files will be created + item_uid: callable + function to get unique ID from an item + cache_contains: callable + function to check if item UID is in cache (expensive operation) + lock_timeout: int + timeout in seconds for lock acquisition and stale lock threshold + """ + + def __init__( + self, + lock_dir: Path, + item_uid: tp.Callable[[tp.Any], str], + cache_contains: tp.Callable[[str], bool], + lock_timeout: int = 3600, + ): + self.lock_dir = Path(lock_dir) + self.lock_dir.mkdir(exist_ok=True, parents=True) + self.item_uid = item_uid + self.cache_contains = cache_contains + self.lock_timeout = lock_timeout + + def process_with_locks( + self, + items: tp.Sequence[tp.Any], + process_fn: tp.Callable[[list[tp.Any]], dict[str, tp.Any]], + ) -> dict[str, tp.Any]: + """Acquire locks and process items, handling contention automatically. + + Strategy: + 1. Try to acquire locks optimistically (fast timeout) + 2. Process acquired items immediately + 3. For remaining items, retry with full timeout + """ + result: dict[str, tp.Any] = {} + + # Filter out already cached items + remaining = [ + item for item in items if not self.cache_contains(self.item_uid(item)) + ] + + if not remaining: + return result + + # Pass 1: Optimistic acquisition (don't wait) + acquired, remaining = self._acquire_batch(remaining, timeout=0.01) + if acquired: + result.update(self._process_batch(acquired, process_fn)) + + # Pass 2: Wait for remaining items + if remaining: + acquired, remaining = self._acquire_batch( + remaining, timeout=self.lock_timeout + ) + if acquired: + result.update(self._process_batch(acquired, process_fn)) + + # Log items we couldn't acquire + if remaining: + for item in remaining: + logger.warning( + f"Skipping {self.item_uid(item)} - could not acquire lock " + f"after {self.lock_timeout}s" + ) + + return result + + def _acquire_batch( + self, items: list[tp.Any], timeout: float + ) -> tuple[list[tuple[tp.Any, filelock.FileLock]], list[tp.Any]]: + """Try to acquire locks for items. + + Returns (acquired_with_locks, failed_items) + """ + acquired = [] + failed = [] + + for item in items: + uid = self.item_uid(item) + + # Check cache before locking (might have been completed) + if self.cache_contains(uid): + continue + + lock = filelock.FileLock(self.lock_dir / f"{uid}.lock") + if self._try_acquire_lock(lock, uid, timeout): + # Double-check cache after acquiring (someone else may have finished) + if self.cache_contains(uid): + lock.release() + continue + acquired.append((item, lock)) + else: + failed.append(item) + + return acquired, failed + + def _try_acquire_lock( + self, lock: filelock.FileLock, uid: str, timeout: float + ) -> bool: + """Try to acquire lock with stale detection.""" + try: + lock.acquire(timeout=timeout) + return True + except filelock.Timeout: + # Check for stale lock + lock_file = Path(lock.lock_file) + if lock_file.exists(): + try: + age = time.time() - lock_file.stat().st_mtime + if age > self.lock_timeout: + logger.warning(f"Removing stale lock for {uid} (age={age:.0f}s)") + lock_file.unlink(missing_ok=True) + try: + lock.acquire(timeout=0.1) + return True + except filelock.Timeout: + pass + except (FileNotFoundError, OSError): + pass + return False + + def _process_batch( + self, + items_with_locks: list[tuple[tp.Any, filelock.FileLock]], + process_fn: tp.Callable[[list[tp.Any]], dict[str, tp.Any]], + ) -> dict[str, tp.Any]: + """Process items and release locks.""" + items = [item for item, _ in items_with_locks] + locks = [lock for _, lock in items_with_locks] + + try: + return process_fn(items) + finally: + for lock in locks: + lock.release() + + def _get_uid_info( model: pydantic.BaseModel, ignore_discriminator: bool = False ) -> tp.Dict[str, tp.Set[str]]: