diff --git a/great_ai/__init__.py b/great_ai/__init__.py index 211373a..40c27e9 100644 --- a/great_ai/__init__.py +++ b/great_ai/__init__.py @@ -3,7 +3,7 @@ __version__ = "0.0.12" from .context import configure -from .deploy import GreatAI +from .deploy import GreatAI, RouteConfig from .exceptions import ( ArgumentValidationError, MissingArgumentError, diff --git a/great_ai/__main__.py b/great_ai/__main__.py index 0e573cd..dfc30d5 100644 --- a/great_ai/__main__.py +++ b/great_ai/__main__.py @@ -8,7 +8,6 @@ from threading import Event from typing import Optional import uvicorn -from parse_arguments import parse_arguments from uvicorn._subprocess import get_subprocess from uvicorn.config import LOGGING_CONFIG, Config from uvicorn.supervisors.basereload import BaseReload @@ -21,6 +20,8 @@ from great_ai.deploy import GreatAI from great_ai.exceptions import ArgumentValidationError, MissingArgumentError from great_ai.utilities import get_logger +from .parse_arguments import parse_arguments + logger = get_logger(SERVER_NAME) @@ -28,11 +29,11 @@ GREAT_AI_LOGGING_CONFIG = { **LOGGING_CONFIG, "formatters": { "default": { - "()": "great_ai.logger.CustomFormatter", + "()": "great_ai.utilities.logger.CustomFormatter", "fmt": "%(asctime)s | %(levelname)8s | %(message)s", }, "access": { - "()": "great_ai.logger.CustomFormatter", + "()": "great_ai.utilities.logger.CustomFormatter", "fmt": "%(asctime)s | %(levelname)8s | %(message)s", # noqa: E501 }, }, diff --git a/great_ai/context.py b/great_ai/context.py index 3f076d6..19010e1 100644 --- a/great_ai/context.py +++ b/great_ai/context.py @@ -56,12 +56,12 @@ def configure( *, log_level: int = DEBUG, seed: int = 42, - tracing_database: Optional[Type[TracingDatabaseDriver]] = None, + tracing_database_factory: Optional[Type[TracingDatabaseDriver]] = None, large_file_implementation: Optional[Type[LargeFileBase]] = None, should_log_exception_stack: Optional[bool] = None, prediction_cache_size: int = 512, disable_se4ml_banner: bool = False, - dashboard_table_size: int = 50, + dashboard_table_size: int = 20, ) -> None: global _context logger = get_logger("great_ai", level=log_level) @@ -76,17 +76,17 @@ def configure( _set_seed(seed) - tracing_database = _initialize_tracing_database(tracing_database, logger=logger)() + tracing_database_factory = _initialize_tracing_database( + tracing_database_factory, logger=logger + ) + tracing_database = tracing_database_factory() if not tracing_database.is_production_ready: + message = f"The selected tracing database ({tracing_database_factory.__name__}) is not recommended for production" if is_production: - logger.error( - f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production" - ) + logger.error(message) else: - logger.warning( - f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production" - ) + logger.warning(message) _context = Context( tracing_database=tracing_database, diff --git a/great_ai/deploy/__init__.py b/great_ai/deploy/__init__.py index 49f9bc1..63e52be 100644 --- a/great_ai/deploy/__init__.py +++ b/great_ai/deploy/__init__.py @@ -1 +1,2 @@ from .great_ai import GreatAI +from .routes import RouteConfig diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index 813bd48..3024545 100644 --- a/great_ai/deploy/great_ai.py +++ b/great_ai/deploy/great_ai.py @@ -1,266 +1,216 @@ -import inspect from functools import lru_cache, partial, wraps +from textwrap import dedent from typing import ( Any, + Awaitable, Callable, Generic, - Iterable, List, Optional, - Type, + Sequence, TypeVar, + Union, cast, overload, ) -from fastapi import APIRouter, FastAPI, status -from pydantic import BaseModel, create_model +from async_lru import alru_cache +from fastapi import FastAPI from ..constants import DASHBOARD_PATH from ..context import get_context -from ..helper import ( - freeze_arguments, - get_function_metadata_store, - snake_case_to_text, - use_http_exceptions, -) +from ..helper import freeze_arguments, get_function_metadata_store, snake_case_to_text from ..models import model_versions from ..parameters import automatically_decorate_parameters from ..tracing.tracing_context import TracingContext from ..utilities import parallel_map -from ..views import ApiMetadata, CacheStatistics, HealthCheckResponse, Trace -from .routes import ( - bootstrap_docs_endpoints, - bootstrap_feedback_endpoints, - bootstrap_trace_endpoints, -) +from ..views import ApiMetadata, Trace from .routes.bootstrap_dashboard import bootstrap_dashboard +from .routes.bootstrap_docs_endpoints import bootstrap_docs_endpoints +from .routes.bootstrap_feedback_endpoints import bootstrap_feedback_endpoints +from .routes.bootstrap_meta_endpoints import bootstrap_meta_endpoints +from .routes.bootstrap_prediction_endpoint import bootstrap_prediction_endpoint +from .routes.bootstrap_trace_endpoints import bootstrap_trace_endpoints +from .routes.route_config import RouteConfig -T = TypeVar("T") +T = TypeVar("T", bound=Union[Trace, Awaitable[Trace]]) +V = TypeVar("V") -class GreatAI(Generic[T]): - def __init__(self, func: Callable[..., Any], version: str, return_raw_result: bool): - is_asynchronous = inspect.iscoroutinefunction(func) +class GreatAI(Generic[T, V]): + __name__: str + __doc__: str + + def __init__( + self, + func: Callable[..., Union[V, Awaitable[V]]], + version: str, + route_config: RouteConfig, + ): func = automatically_decorate_parameters(func) get_function_metadata_store(func).is_finalised = True - self._func = func - - def func_in_tracing_context_sync( - *args: Any, do_not_persist_traces: bool = False, **kwargs: Any - ) -> Trace[T]: - with TracingContext[T]( - func.__name__, do_not_persist_traces=do_not_persist_traces - ) as t: - result = func(*args, **kwargs) - output = t.finalise(output=result) - return result if return_raw_result else output - - async def func_in_tracing_context_async( - *args: Any, do_not_persist_traces: bool = False, **kwargs: Any - ) -> Trace[T]: - with TracingContext[T]( - func.__name__, do_not_persist_traces=do_not_persist_traces - ) as t: - result = await func(*args, **kwargs) - output = t.finalise(output=result) - return result if return_raw_result else output - - func_in_tracing_context = ( - func_in_tracing_context_async - if is_asynchronous - else func_in_tracing_context_sync - ) - - self._cached_func = lru_cache(get_context().prediction_cache_size)( - func_in_tracing_context - ) # cannot put decorator on method, because it require the context to be setup + self._cached_func = self._get_cached_traced_function(func) + self._wrapped_func = wraps(func)(freeze_arguments(self._cached_func)) wraps(func)(self) + self.__doc__ = f"GreatAI wrapper for interacting with the `{self.__name__}` function.\n\n{dedent(self.__doc__ or '')}" - self._version = version + self.version = version + flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions) + if flat_model_versions: + self.version += f"+{flat_model_versions}" self.app = FastAPI( - title=self.name, + title=snake_case_to_text(self.__name__), version=self.version, - description=self.documentation + description=self.__doc__ + f"\n\nFind out more in the [dashboard]({DASHBOARD_PATH}).", docs_url=None, redoc_url=None, ) + self._bootstrap_rest_api(route_config) + @overload @staticmethod - def create( - func: Optional[Callable[..., T]] = None, - ) -> "GreatAI[T]": + def create( # type: ignore + # "Overloaded function signatures 1 and 2 overlap with incompatible return types" + # https://github.com/python/mypy/issues/12759 + func: Callable[..., Awaitable[V]], + ) -> "GreatAI[Awaitable[Trace[V]], V]": ... @overload @staticmethod def create( - version: str, - return_raw_result: bool, - disable_rest_api: bool, - disable_docs: bool, - disable_dashboard: bool, - ) -> Callable[[Callable[..., T]], "GreatAI[T]"]: + func: Callable[..., V], + ) -> "GreatAI[Trace[V], V]": + ... + + @overload + @staticmethod + def create( + func: Optional[Union[Callable[..., V], Callable[..., Awaitable[V]]]] = ..., + *, + version: str = ..., + route_config: RouteConfig = ..., + ) -> Callable: ... @staticmethod def create( - func: Optional[Callable[..., T]] = None, + func: Optional[Callable] = None, *, version: str = "0.0.1", - return_raw_result: bool = False, - disable_rest_api: bool = False, - disable_docs: bool = False, - disable_dashboard: bool = False, - ): + route_config: RouteConfig = RouteConfig(), + ) -> Union[Callable, "GreatAI"]: if func is None: - return cast( - Callable[[Callable[..., T]], GreatAI[T]], - partial( - GreatAI.create, - version=version, - return_raw_result=return_raw_result, - disable_rest_api=disable_rest_api, - disable_docs=disable_docs, - disable_dashboard=disable_dashboard, - ), - ) - instance = GreatAI[T]( - func, version=version, return_raw_result=return_raw_result + @overload + def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]: + ... + + @overload + def inner(func: Callable[..., V]) -> GreatAI[Trace[V], V]: + ... + + def inner(func): # type: ignore + return GreatAI.create( + func, + version=version, + route_config=route_config, + ) + + return inner + + return GreatAI[T, V]( + func, + version=version, + route_config=route_config, ) - if not disable_rest_api: - instance._bootstrap_rest_api( - disable_docs=disable_docs, disable_dashboard=disable_dashboard - ) - - return instance - - @freeze_arguments - def __call__(self, *args: Any, **kwargs: Any) -> Trace[T]: - return self._cached_func(*args, **kwargs) + def __call__(self, *args: Any, **kwargs: Any) -> T: + return self._wrapped_func(*args, **kwargs) def process_batch( self, - batch: Iterable[Any], + batch: Sequence, concurrency: Optional[int] = None, - do_not_persist_traces: bool = False, - ) -> List[Trace[T]]: + do_not_persist_traces: Optional[bool] = False, + ) -> List[Trace[V]]: return list( parallel_map( - freeze_arguments( - partial( - self._cached_func, do_not_persist_traces=do_not_persist_traces - ) + partial( + self._wrapped_func, do_not_persist_traces=do_not_persist_traces ), batch, concurrency=concurrency, ) ) - @property - def name(self) -> str: - return snake_case_to_text(self._func.__name__) + @staticmethod + def _get_cached_traced_function( + func: Callable[..., Union[V, Awaitable[V]]] + ) -> Callable[..., T]: + @lru_cache(maxsize=get_context().prediction_cache_size) + def func_in_tracing_context_sync( + *args: Any, + do_not_persist_traces: bool = False, + **kwargs: Any, + ) -> Trace[V]: + with TracingContext[V]( + func.__name__, do_not_persist_traces=do_not_persist_traces + ) as t: + result = func(*args, **kwargs) + return t.finalise(output=result) - @property - def version(self) -> str: - flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions) - if flat_model_versions: - flat_model_versions = f"+{flat_model_versions}" + @alru_cache(maxsize=get_context().prediction_cache_size) + async def func_in_tracing_context_async( + *args: Any, + do_not_persist_traces: bool = False, + **kwargs: Any, + ) -> Trace[V]: + with TracingContext[V]( + func.__name__, do_not_persist_traces=do_not_persist_traces + ) as t: + result = await cast(Callable[..., Awaitable], func)(*args, **kwargs) + return t.finalise(output=result) - return f"{self._version}{flat_model_versions}" - - @property - def documentation(self) -> str: return ( - f"GreatAI wrapper for interacting with the `{self._func.__name__}` function.\n\n" - + ( - "\n".join( - line.strip() - for line in (self._func.__doc__ or "").split("\n") - if line.strip() - ) - ) + func_in_tracing_context_async + if get_function_metadata_store(func).is_asynchronous + else func_in_tracing_context_sync ) - def _bootstrap_rest_api(self, disable_docs: bool, disable_dashboard: bool) -> None: - self._bootstrap_prediction_endpoint() + def _bootstrap_rest_api(self, route_config: RouteConfig) -> None: + if route_config.prediction_endpoint_enabled: + bootstrap_prediction_endpoint(self.app, self._wrapped_func) - if not disable_docs: + if route_config.docs_endpoints_enabled: bootstrap_docs_endpoints(self.app) - if not disable_dashboard: + if route_config.dashboard_enabled: bootstrap_dashboard( self.app, - function_name=self._func.__name__, - documentation=self.documentation, + function_name=self.__name__, + documentation=self.__doc__, ) + + if route_config.trace_endpoints_enabled: bootstrap_trace_endpoints(self.app) - bootstrap_feedback_endpoints(self.app) - self._bootstrap_meta_endpoints() + if route_config.feedback_endpoints_enabled: + bootstrap_feedback_endpoints(self.app) - def _bootstrap_prediction_endpoint(self) -> None: - router = APIRouter( - tags=["predictions"], - ) - - schema = self._get_schema() - - @router.post( - "/predict", status_code=status.HTTP_200_OK, response_model=Trace[T] - ) - @use_http_exceptions - def predict(input_value: schema) -> Trace[T]: # type: ignore - return self(**cast(BaseModel, input_value).dict()) - - self.app.include_router(router) - - def _get_schema(self) -> Type[BaseModel]: - signature = inspect.signature(self._func) - parameters = { - p.name: ( - p.annotation if p.annotation != inspect._empty else Any, - p.default if p.default != inspect._empty else ..., + if route_config.meta_endpoints_enabled: + bootstrap_meta_endpoints( + self.app, + self._cached_func, + ApiMetadata( + name=self.__name__, + version=self.version, + documentation=self.__doc__, + configuration=get_context().to_flat_dict(), + ), ) - for p in signature.parameters.values() - if p.name in get_function_metadata_store(self._func).input_parameter_names - } - - schema: Type[BaseModel] = create_model("InputModel", **parameters) # type: ignore - return schema - - def _bootstrap_meta_endpoints(self) -> None: - router = APIRouter( - tags=["meta"], - ) - - @router.get("/health", status_code=status.HTTP_200_OK) - def check_health() -> HealthCheckResponse: - hits, misses, maxsize, cache_size = self._cached_func.cache_info() - cache_statistics = CacheStatistics( - hits=hits, misses=misses, size=cache_size, max_size=maxsize - ) - - return HealthCheckResponse( - is_healthy=True, cache_statistics=cache_statistics - ) - - @router.get( - "/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK - ) - def get_version() -> ApiMetadata: - return ApiMetadata( - name=self.name, - version=self.version, - documentation=self.documentation, - configuration=get_context().to_flat_dict(), - ) - - self.app.include_router(router) diff --git a/great_ai/deploy/routes/__init__.py b/great_ai/deploy/routes/__init__.py index b5c2966..5ac4e86 100644 --- a/great_ai/deploy/routes/__init__.py +++ b/great_ai/deploy/routes/__init__.py @@ -1,4 +1,7 @@ from .bootstrap_dashboard import bootstrap_dashboard from .bootstrap_docs_endpoints import bootstrap_docs_endpoints from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints +from .bootstrap_meta_endpoints import bootstrap_meta_endpoints +from .bootstrap_prediction_endpoint import bootstrap_prediction_endpoint from .bootstrap_trace_endpoints import bootstrap_trace_endpoints +from .route_config import RouteConfig diff --git a/great_ai/deploy/routes/bootstrap_dashboard.py b/great_ai/deploy/routes/bootstrap_dashboard.py index 660db1e..2a0069a 100644 --- a/great_ai/deploy/routes/bootstrap_dashboard.py +++ b/great_ai/deploy/routes/bootstrap_dashboard.py @@ -6,7 +6,7 @@ from fastapi.responses import RedirectResponse from fastapi.staticfiles import StaticFiles from ...constants import DASHBOARD_PATH -from .dashboard import create_dash_app +from .dashboard.create_dash_app import create_dash_app PATH = Path(__file__).parent.resolve() diff --git a/great_ai/deploy/routes/bootstrap_meta_endpoints.py b/great_ai/deploy/routes/bootstrap_meta_endpoints.py new file mode 100644 index 0000000..d419c0b --- /dev/null +++ b/great_ai/deploy/routes/bootstrap_meta_endpoints.py @@ -0,0 +1,26 @@ +from typing import Any + +from fastapi import APIRouter, FastAPI, status + +from ...views import ApiMetadata, CacheStatistics, HealthCheckResponse + + +def bootstrap_meta_endpoints(app: FastAPI, func: Any, metadata: ApiMetadata) -> None: + router = APIRouter( + tags=["meta"], + ) + + @router.get("/health", status_code=status.HTTP_200_OK) + def check_health() -> HealthCheckResponse: + hits, misses, maxsize, cache_size = func.cache_info() + cache_statistics = CacheStatistics( + hits=hits, misses=misses, size=cache_size, max_size=maxsize + ) + + return HealthCheckResponse(is_healthy=True, cache_statistics=cache_statistics) + + @router.get("/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK) + def get_version() -> ApiMetadata: + return metadata + + app.include_router(router) diff --git a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py new file mode 100644 index 0000000..0ed6f5b --- /dev/null +++ b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py @@ -0,0 +1,51 @@ +import inspect +from typing import Any, Awaitable, Callable, Type, Union, cast + +from fastapi import APIRouter, FastAPI, HTTPException, status +from pydantic import BaseModel, create_model + +from ...helper import get_function_metadata_store +from ...views import Trace + + +def bootstrap_prediction_endpoint( + app: FastAPI, func: Callable[..., Union[Trace, Awaitable[Trace]]] +) -> None: + router = APIRouter( + tags=["predictions"], + ) + + schema = _get_schema(func) + + @router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace) + async def predict(input_value: schema) -> Trace: # type: ignore + try: + if inspect.iscoroutinefunction(func): + return await cast(Callable[..., Awaitable[Trace]], func)( + **cast(BaseModel, input_value).dict() + ) + return cast(Callable[..., Trace], func)( + **cast(BaseModel, input_value).dict() + ) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"The following exception has occurred: {type(e).__name__}: {e}", + ) + + app.include_router(router) + + +def _get_schema(func: Callable) -> Type[BaseModel]: + signature = inspect.signature(func) + parameters = { + p.name: ( + p.annotation if p.annotation != inspect._empty else Any, + p.default if p.default != inspect._empty else ..., + ) + for p in signature.parameters.values() + if p.name in get_function_metadata_store(func).input_parameter_names + } + + schema: Type[BaseModel] = create_model("InputModel", **parameters) # type: ignore + return schema diff --git a/great_ai/deploy/routes/dashboard/assets/index.css b/great_ai/deploy/routes/dashboard/assets/index.css index e287b4f..2241c2b 100644 --- a/great_ai/deploy/routes/dashboard/assets/index.css +++ b/great_ai/deploy/routes/dashboard/assets/index.css @@ -129,13 +129,17 @@ main > header > div > h1 { .version-tag { border-radius: var(--border-radius); - background: #ddd; display: inline-block; font-size: 1rem; - padding: 3px 6px; + padding: 3px 8px; margin-left: var(--small-padding) } +main > header .version-tag { + background: var(--background-color); + vertical-align: 4px; +} + main > header > *:nth-child(2) { min-width: 250px; max-width: 550px; diff --git a/great_ai/deploy/routes/dashboard/create_dash_app.py b/great_ai/deploy/routes/dashboard/create_dash_app.py index 54803dc..51d2d48 100644 --- a/great_ai/deploy/routes/dashboard/create_dash_app.py +++ b/great_ai/deploy/routes/dashboard/create_dash_app.py @@ -21,12 +21,13 @@ from .get_traces_table import get_traces_table def create_dash_app(function_name: str, version: str, function_docs: str) -> Flask: accent_color = text_to_hex_color(function_name) + function_name = snake_case_to_text(function_name) app = Dash( function_name, requests_pathname_prefix=DASHBOARD_PATH + "/", server=Flask(__name__), - title=snake_case_to_text(function_name), + title=function_name, update_title=None, external_stylesheets=[ "/assets/index.css", diff --git a/great_ai/deploy/routes/dashboard/get_footer.py b/great_ai/deploy/routes/dashboard/get_footer.py index e6c2fe4..0e09328 100644 --- a/great_ai/deploy/routes/dashboard/get_footer.py +++ b/great_ai/deploy/routes/dashboard/get_footer.py @@ -1,5 +1,7 @@ from dash import html +from great_ai import __version__ + from ....constants import GITHUB_LINK @@ -8,7 +10,9 @@ def get_footer() -> html.Footer: [ html.Div( [ - html.H6("GreatAI"), + html.H6( + ["GreatAI", html.Span(__version__, className="version-tag")] + ), html.P( "A human-friendly framework for robust end-to-end AI deployments." ), diff --git a/great_ai/deploy/routes/route_config.py b/great_ai/deploy/routes/route_config.py new file mode 100644 index 0000000..33c8ebd --- /dev/null +++ b/great_ai/deploy/routes/route_config.py @@ -0,0 +1,10 @@ +from pydantic import BaseModel + + +class RouteConfig(BaseModel): + prediction_endpoint_enabled: bool = True + docs_endpoints_enabled: bool = True + dashboard_enabled: bool = True + feedback_endpoints_enabled: bool = True + trace_endpoints_enabled: bool = True + meta_endpoints_enabled: bool = True diff --git a/great_ai/helper/__init__.py b/great_ai/helper/__init__.py index c956174..ea285ab 100644 --- a/great_ai/helper/__init__.py +++ b/great_ai/helper/__init__.py @@ -5,4 +5,3 @@ from .hashable_base_model import HashableBaseModel from .snake_case_to_text import snake_case_to_text from .strip_lines import strip_lines from .text_to_hex_color import text_to_hex_color -from .use_http_exceptions import use_http_exceptions diff --git a/great_ai/helper/assert_function_is_not_finalised.py b/great_ai/helper/assert_function_is_not_finalised.py index 55169ab..e445811 100644 --- a/great_ai/helper/assert_function_is_not_finalised.py +++ b/great_ai/helper/assert_function_is_not_finalised.py @@ -1,10 +1,10 @@ -from typing import Any, Callable +from typing import Callable from ..exceptions import WrongDecoratorOrderError from .get_function_metadata_store import get_function_metadata_store -def assert_function_is_not_finalised(func: Callable[..., Any]) -> None: +def assert_function_is_not_finalised(func: Callable) -> None: error_message = ( "The outer-most (first) decorator has to be `@GreatAI.deploy`. " + f"In the case of `{func.__name__}`, it is not: fix this by moving `@GreatAI.deploy` to the top." diff --git a/great_ai/helper/freeze_arguments.py b/great_ai/helper/freeze_arguments.py index 427899b..c50b13c 100644 --- a/great_ai/helper/freeze_arguments.py +++ b/great_ai/helper/freeze_arguments.py @@ -1,41 +1,50 @@ +import inspect from functools import wraps -from typing import Any, Callable, Dict, List, Set, Union +from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast from pydantic import BaseModel +F = TypeVar("F", bound=Callable) + class FrozenDict(dict): - def __hash__(self) -> int: + def __hash__(self) -> int: # type: ignore return hash(frozenset((k, freeze(v)) for k, v in self.items())) class FrozenList(list): - def __hash__(self) -> int: + def __hash__(self) -> int: # type: ignore return hash(tuple(freeze(i) for i in self)) class FrozenSet(set): - def __hash__(self) -> int: + def __hash__(self) -> int: # type: ignore return hash(frozenset(freeze(i) for i in self)) -def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]: - """Transform mutable dictionary - Into immutable - Useful to be compatible with cache - source: https://stackoverflow.com/questions/6358481/using-functools-lru-cache-with-dictionary-arguments - """ - +def freeze_arguments(func: F) -> F: @wraps(func) - def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any: + def wrapper(*args: Any, **kwargs: Any) -> Any: args = tuple(freeze(arg) for arg in args) kwargs = {k: freeze(v) for k, v in kwargs.items()} return func(*args, **kwargs) - return wrapper + @wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + return await wrapper(*args, **kwargs) + + return cast(F, async_wrapper if inspect.iscoroutinefunction(func) else wrapper) -def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any: +def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel]) -> Any: + """ + >>> class MyClass(BaseModel): + ... a: int + >>> my_object = MyClass(a=3) + >>> my_other_object = MyClass(a=4) + >>> freeze(my_object) == freeze(my_other_object), freeze(my_object) == freeze(my_object) + (False, True) + """ if isinstance(value, dict): return FrozenDict(value) @@ -47,7 +56,7 @@ def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any: if isinstance(value, BaseModel): - class HashableValue(type(value)): + class HashableValue(type(value)): # type: ignore def __hash__(self) -> int: return hash(frozenset((k, freeze(v)) for k, v in self.dict().items())) diff --git a/great_ai/helper/get_arguments.py b/great_ai/helper/get_arguments.py index 91fb44c..be05f2c 100644 --- a/great_ai/helper/get_arguments.py +++ b/great_ai/helper/get_arguments.py @@ -3,7 +3,7 @@ from typing import Any, Callable, Dict, Mapping, Sequence def get_arguments( - func: Callable[..., Any], args: Sequence[Any], kwargs: Mapping[str, Any] + func: Callable, args: Sequence[Any], kwargs: Mapping[str, Any] ) -> Dict[str, Any]: """Return mapping from parameter names to actual argument values""" diff --git a/great_ai/helper/get_function_metadata_store.py b/great_ai/helper/get_function_metadata_store.py index 8694636..5f35be3 100644 --- a/great_ai/helper/get_function_metadata_store.py +++ b/great_ai/helper/get_function_metadata_store.py @@ -1,12 +1,14 @@ +import inspect from typing import Any, Callable, cast from ..views.function_metadata import FunctionMetadata -def get_function_metadata_store(func: Callable[..., Any]) -> FunctionMetadata: +def get_function_metadata_store(func: Callable) -> FunctionMetadata: any_func = cast(Any, func) if not hasattr(any_func, "_great_ai_metadata"): - any_func._great_ai_metadata = FunctionMetadata() + is_asynchronous = inspect.iscoroutinefunction(func) + any_func._great_ai_metadata = FunctionMetadata(is_asynchronous=is_asynchronous) return any_func._great_ai_metadata diff --git a/great_ai/helper/text_to_hex_color.py b/great_ai/helper/text_to_hex_color.py index e906a0c..62acb98 100644 --- a/great_ai/helper/text_to_hex_color.py +++ b/great_ai/helper/text_to_hex_color.py @@ -4,10 +4,14 @@ from hashlib import md5 def text_to_hex_color(text: str) -> str: ascii_bytes = text.encode("ascii") + digest = md5( ascii_bytes ).hexdigest() # the built-in hash function is salted differently in each process + integer = int(digest, 16) hue = integer % 6311 / 6311.0 - rgb = colorsys.hsv_to_rgb(hue, 0.75, 0.6) + + rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.6) + return "#" + "".join("%02X" % round(i * 255) for i in rgb) diff --git a/great_ai/helper/use_http_exceptions.py b/great_ai/helper/use_http_exceptions.py deleted file mode 100644 index 8eae30a..0000000 --- a/great_ai/helper/use_http_exceptions.py +++ /dev/null @@ -1,20 +0,0 @@ -from functools import wraps -from typing import Any, Callable, Dict, List, TypeVar, cast - -from fastapi import HTTPException, status - -F = TypeVar("F", bound=Callable[..., Any]) - - -def use_http_exceptions(func: F) -> F: - @wraps(func) - def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any: - try: - return func(*args, **kwargs) - except Exception as e: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"The following exception has occurred: {type(e).__name__}: {e}", - ) - - return cast(F, wrapper) diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py index 621eef4..2f9ff2b 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -89,9 +89,7 @@ class LargeFileBase(ABC): cls.configure_credentials(**ConfigFile(secrets_path)) @classmethod - def configure_credentials( - cls, - ) -> None: + def configure_credentials(cls, **kwargs: str) -> None: cls.initialized = True def __enter__(self) -> IO: diff --git a/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/large_file/large_file/large_file_mongo.py index 60a1601..673ac5c 100644 --- a/great_ai/large_file/large_file/large_file_mongo.py +++ b/great_ai/large_file/large_file/large_file_mongo.py @@ -1,7 +1,7 @@ import re from functools import cached_property from pathlib import Path -from typing import Any, List, Mapping +from typing import Any, List from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket from pymongo import MongoClient @@ -27,7 +27,7 @@ class LargeFileMongo(LargeFileBase): *, mongo_connection_string: str, mongo_database: str, - **_: Mapping[str, Any], + **_: Any, ) -> None: cls.mongo_connection_string = mongo_connection_string cls.mongo_database = mongo_database diff --git a/great_ai/large_file/large_file/large_file_s3.py b/great_ai/large_file/large_file/large_file_s3.py index a8c8005..0104eba 100644 --- a/great_ai/large_file/large_file/large_file_s3.py +++ b/great_ai/large_file/large_file/large_file_s3.py @@ -1,6 +1,6 @@ from functools import cached_property from pathlib import Path -from typing import Any, List, Mapping, Optional +from typing import Any, List, Optional import boto3 @@ -58,7 +58,7 @@ class LargeFileS3(LargeFileBase): aws_secret_access_key: str, large_files_bucket_name: str, aws_endpoint_url: Optional[str] = None, - **_: Mapping[str, Any], + **_: Any, ) -> None: cls.region_name = aws_region_name cls.access_key_id = aws_access_key_id diff --git a/great_ai/models/use_model.py b/great_ai/models/use_model.py index 3e830ab..f2d1c14 100644 --- a/great_ai/models/use_model.py +++ b/great_ai/models/use_model.py @@ -21,7 +21,7 @@ from ..helper.assert_function_is_not_finalised import assert_function_is_not_fin from ..tracing.tracing_context import TracingContext from ..views import Model -F = TypeVar("F", bound=Callable[..., Any]) +F = TypeVar("F", bound=Callable) def use_model( @@ -70,4 +70,6 @@ def _load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]: return path, file.version with file as f: - return load(f), file.version + loaded = load(f) + + return loaded, file.version diff --git a/great_ai/parameters/automatically_decorate_parameters.py b/great_ai/parameters/automatically_decorate_parameters.py index 8a3b017..9d4e9f2 100644 --- a/great_ai/parameters/automatically_decorate_parameters.py +++ b/great_ai/parameters/automatically_decorate_parameters.py @@ -1,10 +1,10 @@ import inspect -from typing import Any, Callable, TypeVar +from typing import Callable, TypeVar from ..helper.get_function_metadata_store import get_function_metadata_store from .parameter import parameter -F = TypeVar("F", bound=Callable[..., Any]) +F = TypeVar("F", bound=Callable) def automatically_decorate_parameters(func: F) -> F: diff --git a/great_ai/parameters/parameter.py b/great_ai/parameters/parameter.py index 1d8e4b5..cfcfce6 100644 --- a/great_ai/parameters/parameter.py +++ b/great_ai/parameters/parameter.py @@ -8,14 +8,13 @@ from ..helper import get_arguments, get_function_metadata_store from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised from ..tracing.tracing_context import TracingContext -T = TypeVar("T") -F = TypeVar("F", bound=Callable[..., Any]) +F = TypeVar("F", bound=Callable) def parameter( parameter_name: str, *, - validator: Callable[[T], bool] = lambda _: True, + validator: Callable[[Any], bool] = lambda _: True, disable_logging: bool = False, ) -> Callable[[F], F]: def decorator(func: F) -> F: diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py index 1cc6d26..7c2600d 100644 --- a/great_ai/persistence/mongodb_driver.py +++ b/great_ai/persistence/mongodb_driver.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, List, Mapping, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple from pymongo import MongoClient @@ -20,6 +20,9 @@ operator_mapping = { class MongodbDriver(TracingDatabaseDriver): is_production_ready = True + mongo_connection_string: str + mongo_database: str + def __init__(self) -> None: super().__init__() if self.mongo_connection_string is None or self.mongo_database is None: @@ -33,7 +36,7 @@ class MongodbDriver(TracingDatabaseDriver): *, mongo_connection_string: str, mongo_database: str, - **_: Mapping[str, Any], + **_: Any, ) -> None: cls.mongo_connection_string = mongo_connection_string cls.mongo_database = mongo_database @@ -43,21 +46,23 @@ class MongodbDriver(TracingDatabaseDriver): serialized = trace.to_flat_dict() serialized["_id"] = trace.trace_id - with MongoClient(self.mongo_connection_string) as client: - return client[self.mongo_database].traces.insert_one(serialized) + with MongoClient[Any](self.mongo_connection_string) as client: + return client[self.mongo_database].traces.insert_one(serialized).inserted_id def save_batch(self, documents: List[Trace]) -> List[str]: serialized = [d.to_flat_dict() for d in documents] for s in serialized: s["_id"] = s["trace_id"] - with MongoClient(self.mongo_connection_string) as client: - return client[self.mongo_database].traces.insert_many( - serialized, ordered=False + with MongoClient[Any](self.mongo_connection_string) as client: + return ( + client[self.mongo_database] + .traces.insert_many(serialized, ordered=False) + .inserted_ids ) def get(self, id: str) -> Optional[Trace]: - with MongoClient(self.mongo_connection_string) as client: + with MongoClient[Any](self.mongo_connection_string) as client: value = client[self.mongo_database].traces.find_one(id) if value: @@ -83,15 +88,8 @@ class MongodbDriver(TracingDatabaseDriver): sort_by: Sequence[SortBy] = [], ) -> Tuple[List[Trace], int]: - query = { - "filter": { - "$and": [{"tags": tag} for tag in conjunctive_tags] - + [ - {f.property: {self._get_operator(f): f.value}} - for f in conjunctive_filters - ] - + [{}] - }, + query: Dict[str, Any] = { + "filter": {}, "sort": [ (col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by ], @@ -103,35 +101,43 @@ class MongodbDriver(TracingDatabaseDriver): if take: query["limit"] = take + and_query: List[Dict[str, Any]] = [{}] + and_query.extend({"tags": tag} for tag in conjunctive_tags) + and_query.extend( + {f.property: {self._get_operator(f): f.value}} for f in conjunctive_filters + ) + if since: - query["filter"]["$and"].append({"created": {"$gte": since}}) + and_query.append({"created": {"$gte": since}}) if until: - query["filter"]["$and"].append({"created": {"$lte": until}}) + and_query.append({"created": {"$lte": until}}) if has_feedback is not None: - query["filter"]["$and"].append( + and_query.append( {"feedback": {"$ne": None}} if has_feedback else {"feedback": None} ) + query["filter"]["$and"] = and_query - with MongoClient(self.mongo_connection_string) as client: + with MongoClient[Any](self.mongo_connection_string) as client: values = client[self.mongo_database].traces.find(**query) - documents = [Trace.parse_obj(t) for t in values] + documents = [Trace[Any].parse_obj(t) for t in values] return documents, len(documents) def update(self, id: str, new_version: Trace) -> None: serialized = new_version.to_flat_dict() serialized["_id"] = new_version.trace_id - with MongoClient(self.mongo_connection_string) as client: - client[self.mongo_database].traces.update_one(id, new_version) + + with MongoClient[Any](self.mongo_connection_string) as client: + client[self.mongo_database].traces.update_one({"_id": id}, serialized) def delete(self, id: str) -> None: - with MongoClient(self.mongo_connection_string) as client: - client[self.mongo_database].traces.delete_one(id) + with MongoClient[Any](self.mongo_connection_string) as client: + client[self.mongo_database].traces.delete_one({"_id": id}) - def delete_batch(self, ids: List[str]) -> List[str]: + def delete_batch(self, ids: List[str]) -> None: delete_filter = {"_id": {"$in": ids}} - with MongoClient(self.mongo_connection_string) as client: - return client[self.mongo_database].traces.delete_many(delete_filter) + with MongoClient[Any](self.mongo_connection_string) as client: + client[self.mongo_database].traces.delete_many(delete_filter) diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py index 61b0761..28d364b 100644 --- a/great_ai/persistence/parallel_tinydb_driver.py +++ b/great_ai/persistence/parallel_tinydb_driver.py @@ -1,7 +1,7 @@ from datetime import datetime from multiprocessing import Lock from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import pandas as pd from tinydb import TinyDB @@ -48,14 +48,8 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): def does_match(d: Dict[str, Any]) -> bool: return ( not set(conjunctive_tags) - set(d["tags"]) - and ( - since is None - or cast(datetime, datetime.fromisoformat(d["created"])) >= since - ) - and ( - until is None - or cast(datetime, datetime.fromisoformat(d["created"])) <= until - ) + and (since is None or datetime.fromisoformat(d["created"]) >= since) + and (until is None or datetime.fromisoformat(d["created"]) <= until) and ( has_feedback is None or has_feedback == (d["feedback"] is not None) ) @@ -102,7 +96,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): def delete(self, id: str) -> None: self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id)) - def delete_batch(self, ids: List[str]) -> List[str]: + def delete_batch(self, ids: List[str]) -> None: for i in ids: self.delete(i) diff --git a/great_ai/tracing/add_ground_truth.py b/great_ai/tracing/add_ground_truth.py index c98d2d1..78cc7cc 100644 --- a/great_ai/tracing/add_ground_truth.py +++ b/great_ai/tracing/add_ground_truth.py @@ -51,7 +51,7 @@ def add_ground_truth( created = datetime.utcnow().isoformat() traces = [ - Trace( + Trace[T]( created=created, original_execution_time_ms=0, logged_values=X if isinstance(X, dict) else {"input": X}, diff --git a/great_ai/tracing/delete_ground_truth.py b/great_ai/tracing/delete_ground_truth.py index ef1620e..e5010bc 100644 --- a/great_ai/tracing/delete_ground_truth.py +++ b/great_ai/tracing/delete_ground_truth.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import List, Optional, Union +from typing import List, Optional, Union, cast from ..context import get_context @@ -19,4 +19,4 @@ def delete_ground_truth( conjunctive_tags=tags, until=until, since=since, has_feedback=True ) - db.delete_batch([i.trace_id for i in items]) + db.delete_batch([cast(str, i.trace_id) for i in items]) diff --git a/great_ai/utilities/chunk.py b/great_ai/utilities/chunk.py index 694407f..13db74d 100644 --- a/great_ai/utilities/chunk.py +++ b/great_ai/utilities/chunk.py @@ -3,7 +3,7 @@ from typing import Iterable, List, TypeVar T = TypeVar("T") -def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]: +def chunk(values: Iterable[T], chunk_size: int) -> Iterable[List[T]]: assert chunk_size >= 1 result: List[T] = [] diff --git a/great_ai/utilities/config_file/config_file.py b/great_ai/utilities/config_file/config_file.py index 079069b..92a4951 100644 --- a/great_ai/utilities/config_file/config_file.py +++ b/great_ai/utilities/config_file/config_file.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import Dict, Iterable, Tuple, Union +from typing import Dict, ItemsView, Iterator, KeysView, Mapping, Union, ValuesView from ..logger import get_logger from .parse_error import ParseError @@ -11,7 +11,7 @@ ENVIRONMENT_VARIABLE_KEY_PREFIX = "ENV" logger = get_logger("ConfigFile") -class ConfigFile: +class ConfigFile(Mapping[str, str]): def __init__(self, path: Union[Path, str]) -> None: if not isinstance(path, Path): path = Path(path) @@ -24,7 +24,7 @@ class ConfigFile: self._parse() - def _parse(self): + def _parse(self) -> None: with open(self._path, encoding="utf-8") as f: lines: str = f.read() @@ -72,20 +72,20 @@ class ConfigFile: __getitem__ = __getattr__ - def __iter__(self) -> Iterable[Tuple[str, str]]: + def __iter__(self) -> Iterator[str]: return iter(self._key_values) def __len__(self) -> int: return len(self._key_values) - def keys(self): + def keys(self) -> KeysView[str]: return self._key_values.keys() - def values(self): + def values(self) -> ValuesView[str]: return self._key_values.values() - def items(self): + def items(self) -> ItemsView[str, str]: return self._key_values.items() - def __repr__(self): + def __repr__(self) -> str: return f"{type(self).__name__}(path={self._path}) {self._key_values}" diff --git a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py index 78d9a80..f5e4c51 100644 --- a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py +++ b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, TypeVar import matplotlib import matplotlib.pyplot as plt @@ -9,9 +9,11 @@ from sklearn.metrics import average_precision_score, precision_recall_curve from ..unique import unique from .draw_f1_iso_lines import draw_f1_iso_lines +T = TypeVar("T", str, float) + def evaluate_ranking( - expected: List[Union[str, float]], + expected: List[T], actual_scores: List[float], target_recall: float, title: Optional[str] = "", @@ -20,7 +22,7 @@ def evaluate_ranking( output_svg: Optional[Path] = None, reverse_order: bool = False, plot: bool = True, -) -> Dict[Union[str, float], float]: +) -> Dict[T, float]: assert 0 <= target_recall <= 1 if plot and axes is None: @@ -39,7 +41,7 @@ def evaluate_ranking( draw_f1_iso_lines(axes=ax) - results: Dict[Union[str, float], float] = {} + results: Dict[T, float] = {} for i in range(len(classes) - 1): binarized_expected = [ (v < classes[i]) if reverse_order else (v > classes[i]) diff --git a/great_ai/utilities/parallel_map/get_config.py b/great_ai/utilities/parallel_map/get_config.py index 4f96cdc..98ffbe4 100644 --- a/great_ai/utilities/parallel_map/get_config.py +++ b/great_ai/utilities/parallel_map/get_config.py @@ -17,14 +17,15 @@ def get_config( ) -> ParallelMapConfiguration: is_input_sequence = hasattr(input_values, "__len__") + input_length = len(input_values) if is_input_sequence else None # type: ignore if concurrency is None: concurrency = os.cpu_count() or 1 assert concurrency >= 1, "At least one mapper process has to be created" if chunk_size is None: - if is_input_sequence: - chunk_size = max(1, ceil(len(input_values) / concurrency / 10)) + if input_length is not None: + chunk_size = max(1, ceil(input_length / concurrency / 10)) else: raise ValueError( "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`" @@ -32,19 +33,14 @@ def get_config( assert chunk_size >= 1, "Chunks have to contain at least one element" chunk_count: Optional[int] = None - if is_input_sequence: - chunk_count = ceil(len(input_values) / chunk_size) + if input_length is not None: + chunk_count = ceil(input_length / chunk_size) if chunk_count < concurrency: logger.warning( f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks" ) concurrency = chunk_count - if concurrency == 1: - logger.warning("Running in series, there is no reason for parallelism") - - input_length = len(input_values) if is_input_sequence else None - config = ParallelMapConfiguration( concurrency=concurrency, chunk_count=chunk_count, diff --git a/great_ai/utilities/parallel_map/manage_communication.py b/great_ai/utilities/parallel_map/manage_communication.py index 48f4db0..1fc9630 100644 --- a/great_ai/utilities/parallel_map/manage_communication.py +++ b/great_ai/utilities/parallel_map/manage_communication.py @@ -1,10 +1,11 @@ import multiprocessing as mp import queue import traceback -from typing import Dict, Iterable, TypeVar, Union +from typing import Dict, Iterable, List, TypeVar, Union from ..chunk import chunk from ..logger import get_logger +from .map_result import MapResult from .worker_exception import WorkerException logger = get_logger("parallel_map") @@ -27,6 +28,7 @@ def manage_communication( next_output_index = 0 read_input_length = 0 is_iteration_over = False + while not is_iteration_over or next_output_index < read_input_length: if not is_iteration_over: try: @@ -36,30 +38,30 @@ def manage_communication( except StopIteration: is_iteration_over = True except Exception as e: - if not ignore_exceptions: - raise - else: + if ignore_exceptions: logger.error( f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" ) + else: + raise try: - result_chunk = output_queue.get_nowait() - - for index, value, exception in result_chunk: - if exception is not None: - e, tb = exception - if not ignore_exceptions: - raise WorkerException from e - else: + result_chunk: List[MapResult] = output_queue.get_nowait() + for r in result_chunk: + if r.exception is not None: + if ignore_exceptions: logger.error( - f"Exception {e} encountered in worker, traceback:\n{tb}" + f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}" ) + else: + raise WorkerException from r.exception + if unordered: - yield value + + yield r.value next_output_index += 1 else: - indexed_results[index] = value + indexed_results[r.order] = r.value if not unordered: while next_output_index in indexed_results: diff --git a/great_ai/utilities/parallel_map/manage_serial.py b/great_ai/utilities/parallel_map/manage_serial.py deleted file mode 100644 index ae3f8c5..0000000 --- a/great_ai/utilities/parallel_map/manage_serial.py +++ /dev/null @@ -1,36 +0,0 @@ -import traceback -from typing import Callable, Iterable, TypeVar - -from ..logger import get_logger -from .worker_exception import WorkerException - -logger = get_logger("parallel_map") - -T = TypeVar("T") -V = TypeVar("V") - - -def manage_serial( - *, - function: Callable[[T], V], - input_values: Iterable[T], - ignore_exceptions: bool, -) -> Iterable[V]: - try: - for v in input_values: - try: - yield function(v) - except Exception as e: - if not ignore_exceptions: - raise WorkerException from e - else: - logger.error( - f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}" - ) - except Exception as e: - if not ignore_exceptions: - raise - else: - logger.error( - f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" - ) diff --git a/great_ai/utilities/parallel_map/map_result.py b/great_ai/utilities/parallel_map/map_result.py new file mode 100644 index 0000000..81d3dbb --- /dev/null +++ b/great_ai/utilities/parallel_map/map_result.py @@ -0,0 +1,8 @@ +from typing import Any, NamedTuple, Optional + + +class MapResult(NamedTuple): + order: int + value: Any + exception: Optional[Exception] = None + worker_traceback: Optional[str] = None diff --git a/great_ai/utilities/parallel_map/mapper_function.py b/great_ai/utilities/parallel_map/mapper_function.py index c294a0a..921065a 100644 --- a/great_ai/utilities/parallel_map/mapper_function.py +++ b/great_ai/utilities/parallel_map/mapper_function.py @@ -1,42 +1,71 @@ +import asyncio +import inspect import multiprocessing as mp import queue import threading import traceback -from typing import Callable, Union +from multiprocessing.synchronize import Event +from typing import Any, Awaitable, Callable, List, TypeVar, Union, cast import dill +from .map_result import MapResult + +T = TypeVar("T") +V = TypeVar("V") + def mapper_function( input_queue: Union[mp.Queue, queue.Queue], output_queue: Union[mp.Queue, queue.Queue], - should_stop: Union[mp.Event, threading.Event], - map_function: Union[bytes, Callable], + should_stop: Union[Event, threading.Event], + func: Union[bytes, Callable[[T], V], Callable[[T], Awaitable[V]]], ) -> None: try: - if isinstance(map_function, bytes): - map_function = dill.loads(map_function) + if isinstance(func, bytes): + func = cast(Callable[[T], V], dill.loads(func)) - last_chunk = None + is_asynchronous = inspect.iscoroutinefunction(func) + + last_chunk: List[MapResult] = [] while not should_stop.wait(0.1): - if last_chunk is None: + if not last_chunk: try: input_chunk = input_queue.get_nowait() - last_chunk = [] - for i, v in input_chunk: - result, exception = None, None - try: - result = map_function(v) - except Exception as e: - exception = e, traceback.format_exc() - last_chunk.append((i, result, exception)) + if is_asynchronous: + + async def safe(i: int, value: T) -> Any: + try: + return MapResult( + i, + await cast(Callable[[T], Awaitable[V]], func)( + value + ), + ) + except Exception as e: + return MapResult(i, None, e, traceback.format_exc()) + + async def main() -> List[MapResult]: + return await asyncio.gather( + *[safe(i, v) for i, v in input_chunk] + ) + + last_chunk = asyncio.run(main()) + else: + for i, value in input_chunk: + try: + last_chunk.append(MapResult(i, func(value))) + except Exception as e: + last_chunk.append( + MapResult(i, None, e, traceback.format_exc()) + ) except queue.Empty: pass - if last_chunk is not None: + if last_chunk: try: output_queue.put_nowait(last_chunk) - last_chunk = None + last_chunk = [] except queue.Full: pass except (KeyboardInterrupt, BrokenPipeError): diff --git a/great_ai/utilities/parallel_map/parallel_map.py b/great_ai/utilities/parallel_map/parallel_map.py index e95604d..c4c3bac 100644 --- a/great_ai/utilities/parallel_map/parallel_map.py +++ b/great_ai/utilities/parallel_map/parallel_map.py @@ -1,11 +1,20 @@ import multiprocessing as mp -from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload +from typing import ( + Awaitable, + Callable, + Iterable, + Literal, + Optional, + Sequence, + TypeVar, + Union, + overload, +) import dill from .get_config import get_config from .manage_communication import manage_communication -from .manage_serial import manage_serial from .mapper_function import mapper_function from .worker_exception import WorkerException @@ -15,79 +24,72 @@ V = TypeVar("V") @overload def parallel_map( - function: Callable[[T], V], + func: Callable[[T], Union[V, Awaitable[V]]], input_values: Sequence[T], *, - chunk_size: Optional[int], - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: Optional[Literal[False]], -) -> Iterable[V]: - ... - - -@overload -def parallel_map( - function: Callable[[T], V], - input_values: Iterable[T], - *, - chunk_size: int, - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: Optional[Literal[False]], -) -> Iterable[V]: - ... - - -@overload -def parallel_map( - function: Callable[[T], V], - input_values: Sequence[T], - *, - chunk_size: Optional[int], - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: True, + ignore_exceptions: Literal[True], + chunk_size: Optional[int] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., ) -> Iterable[Optional[V]]: ... @overload def parallel_map( - function: Callable[[T], V], - input_values: Iterable[T], + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], *, chunk_size: int, - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: True, + ignore_exceptions: Literal[True], + concurrency: Optional[int] = ..., + unordered: bool = ..., ) -> Iterable[Optional[V]]: ... +@overload def parallel_map( - function, - input_values, + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Sequence[T], *, - chunk_size=None, - concurrency=None, - unordered=False, - ignore_exceptions=False, -): + chunk_size: Optional[int] = ..., + ignore_exceptions: Literal[False] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., +) -> Iterable[V]: + ... + + +@overload +def parallel_map( + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], + *, + chunk_size: int, + ignore_exceptions: Literal[False] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., +) -> Iterable[V]: + ... + + +def parallel_map( + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], + *, + chunk_size: Optional[int] = None, + ignore_exceptions: bool = False, + concurrency: Optional[int] = None, + unordered: bool = False, +) -> Iterable[Optional[V]]: config = get_config( - function=function, + function=func, input_values=input_values, chunk_size=chunk_size, concurrency=concurrency, ) - if config.concurrency == 1: - yield from manage_serial( - function=function, - input_values=input_values, - ignore_exceptions=ignore_exceptions, - ) - ctx = ( mp.get_context("fork") if "fork" in mp.get_all_start_methods() @@ -99,7 +101,7 @@ def parallel_map( output_queue = manager.Queue(config.concurrency * 2) should_stop = ctx.Event() - serialized_map_function = dill.dumps(function, byref=True, recurse=False) + serialized_map_function = dill.dumps(func, byref=True, recurse=False) processes = [ ctx.Process( @@ -110,7 +112,7 @@ def parallel_map( input_queue=input_queue, output_queue=output_queue, should_stop=should_stop, - map_function=serialized_map_function, + func=serialized_map_function, ), ) for i in range(config.concurrency) diff --git a/great_ai/utilities/parallel_map/threaded_parallel_map.py b/great_ai/utilities/parallel_map/threaded_parallel_map.py index 122b692..9e153b1 100644 --- a/great_ai/utilities/parallel_map/threaded_parallel_map.py +++ b/great_ai/utilities/parallel_map/threaded_parallel_map.py @@ -1,10 +1,19 @@ import queue import threading -from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload +from typing import ( + Awaitable, + Callable, + Iterable, + Literal, + Optional, + Sequence, + TypeVar, + Union, + overload, +) from .get_config import get_config from .manage_communication import manage_communication -from .manage_serial import manage_serial from .mapper_function import mapper_function T = TypeVar("T") @@ -13,81 +22,74 @@ V = TypeVar("V") @overload def threaded_parallel_map( - function: Callable[[T], V], + func: Callable[[T], Union[V, Awaitable[V]]], input_values: Sequence[T], *, - chunk_size: Optional[int], - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: Optional[Literal[False]], -) -> Iterable[V]: - ... - - -@overload -def threaded_parallel_map( - function: Callable[[T], V], - input_values: Iterable[T], - *, - chunk_size: int, - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: Optional[Literal[False]], -) -> Iterable[V]: - ... - - -@overload -def threaded_parallel_map( - function: Callable[[T], V], - input_values: Sequence[T], - *, - chunk_size: Optional[int], - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: True, + ignore_exceptions: Literal[True], + chunk_size: Optional[int] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., ) -> Iterable[Optional[V]]: ... @overload def threaded_parallel_map( - function: Callable[[T], V], - input_values: Iterable[T], + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], *, chunk_size: int, - concurrency: Optional[int], - unordered: Optional[bool], - ignore_exceptions: True, + ignore_exceptions: Literal[True], + concurrency: Optional[int] = ..., + unordered: bool = ..., ) -> Iterable[Optional[V]]: ... +@overload def threaded_parallel_map( - function, - input_values, + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Sequence[T], *, - chunk_size=None, - concurrency=None, - unordered=False, - ignore_exceptions=False, -): + chunk_size: Optional[int] = ..., + ignore_exceptions: Literal[False] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., +) -> Iterable[V]: + ... + + +@overload +def threaded_parallel_map( + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], + *, + chunk_size: int, + ignore_exceptions: Literal[False] = ..., + concurrency: Optional[int] = ..., + unordered: bool = ..., +) -> Iterable[V]: + ... + + +def threaded_parallel_map( + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Union[Iterable[T], Sequence[T]], + *, + chunk_size: Optional[int] = None, + ignore_exceptions: bool = False, + concurrency: Optional[int] = None, + unordered: bool = False, +) -> Iterable[Optional[V]]: config = get_config( - function=function, + function=func, input_values=input_values, chunk_size=chunk_size, concurrency=concurrency, ) - if config.concurrency == 1: - yield from manage_serial( - function=function, - input_values=input_values, - ignore_exceptions=ignore_exceptions, - ) - - input_queue = queue.Queue(config.concurrency * 2) - output_queue = queue.Queue(config.concurrency * 2) + input_queue: queue.Queue = queue.Queue(config.concurrency * 2) + output_queue: queue.Queue = queue.Queue(config.concurrency * 2) should_stop = threading.Event() threads = [ @@ -99,7 +101,7 @@ def threaded_parallel_map( input_queue=input_queue, output_queue=output_queue, should_stop=should_stop, - map_function=function, + func=func, ), ) for i in range(config.concurrency) diff --git a/great_ai/views/function_metadata.py b/great_ai/views/function_metadata.py index 55369c0..e59d1ff 100644 --- a/great_ai/views/function_metadata.py +++ b/great_ai/views/function_metadata.py @@ -4,6 +4,7 @@ from pydantic import BaseModel class FunctionMetadata(BaseModel): + is_asynchronous: bool input_parameter_names: List[str] = [] model_parameter_names: List[str] = [] is_finalised: bool = False diff --git a/great_ai/views/trace.py b/great_ai/views/trace.py index 287c42a..9413d85 100644 --- a/great_ai/views/trace.py +++ b/great_ai/views/trace.py @@ -25,10 +25,10 @@ class Trace(Generic[T], HashableBaseModel): extra = Extra.ignore @validator("trace_id", always=True) - def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]: - if not v: - return str(uuid4()) - return v + def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> str: + if v: + return v + return str(uuid4()) @property def input(self) -> Any: diff --git a/pyproject.toml b/pyproject.toml index ba4faa5..a9a6939 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "typeguard >= 2.10.0", "pymongo >= 3.0.0", "dill >= 0.3.5.0", + "async_lru >= 1.0.0", "aiohttp[speedups] >= 3.8.0", ] @@ -45,6 +46,7 @@ dev = [ "mkdocs", "mkdocstrings[python]", "mkdocs-material", + "mkdocs-jupyter", "autoflake", "isort", "black[jupyter]", @@ -53,7 +55,7 @@ dev = [ "pytest", "pytest-cov", "pytest-subtests", - "pytest-asyncio" + "pytest-asyncio", ] [project.urls] diff --git a/scripts/format-python.sh b/scripts/format-python.sh index 8cbf26d..bed0322 100755 --- a/scripts/format-python.sh +++ b/scripts/format-python.sh @@ -26,5 +26,7 @@ for dir in "$@"; do python3 -m mypy --namespace-packages --ignore-missing-imports --install-types --non-interactive --disallow-untyped-defs --disallow-incomplete-defs --follow-imports=silent --exclude=external/ --exclude=/build/ --pretty . || true fi + python3 -m flake8 . --count --show-source --statistics --exclude=__init__.py,.env,external --ignore=E501,E402,F821,W503,E722,E203 + cd - done diff --git a/great_ai/deploy/routes/dashboard/assets/__init__.py b/tests/__init__.py similarity index 100% rename from great_ai/deploy/routes/dashboard/assets/__init__.py rename to tests/__init__.py diff --git a/tests/great_ai/test_great_ai.py b/tests/great_ai/test_great_ai.py deleted file mode 100644 index 7cc7104..0000000 --- a/tests/great_ai/test_great_ai.py +++ /dev/null @@ -1,9 +0,0 @@ -from great_ai import GreatAI - - -def test_process_batch() -> None: - @GreatAI.create(return_raw_result=True) - def f(x): - return x + 2 - - assert f.process_batch([3, 9, 34]) == [5, 11, 36] diff --git a/tests/great_ai/test_async_starters.py b/tests/test_async_starters.py similarity index 54% rename from tests/great_ai/test_async_starters.py rename to tests/test_async_starters.py index 5236258..9a3c5fc 100644 --- a/tests/great_ai/test_async_starters.py +++ b/tests/test_async_starters.py @@ -1,13 +1,7 @@ from asyncio import sleep import pytest - -from great_ai import ( - ArgumentValidationError, - GreatAI, - WrongDecoratorOrderError, - parameter, -) +from great_ai import GreatAI, WrongDecoratorOrderError, parameter @pytest.mark.asyncio @@ -17,28 +11,28 @@ async def test_create_trivial_cases() -> None: await sleep(0.5) return f"Hello {name}!" - assert (await hello_world_1("andras").output) == "Hello andras!" + assert (await hello_world_1("andras")).output == "Hello andras!" @GreatAI.create async def hello_world_2(name: str) -> str: await sleep(0.5) return f"Hello {name}!" - assert (await hello_world_2("andras").output) == "Hello andras!" + assert (await hello_world_2("andras")).output == "Hello andras!" @GreatAI.create() async def hello_world_3(name: str) -> str: await sleep(0.5) return f"Hello {name}!" - assert (await hello_world_3("andras").output) == "Hello andras!" + assert (await hello_world_3("andras")).output == "Hello andras!" @GreatAI.create() - async def hello_world_4(name): + async def hello_world_4(name): # type: ignore await sleep(0.5) return f"Hello {name}!" - assert (await hello_world_4("andras").output) == "Hello andras!" + assert (await hello_world_4("andras")).output == "Hello andras!" @pytest.mark.asyncio @@ -49,14 +43,23 @@ async def test_with_parameter() -> None: await sleep(0.5) return f"Hello {name}!" - assert (await hello_world("andras").output) == "Hello andras!" - with pytest.raises(ArgumentValidationError): - await hello_world("short") +@pytest.mark.asyncio +async def test_with_parameters() -> None: + @GreatAI.create + @parameter("name", validator=lambda v: len(v) > 5) + @parameter("unused", disable_logging=True) + async def hello_world(name: str, unused) -> str: # type: ignore + await sleep(0.5) + return f"Hello {name}!" + assert (await hello_world("andras", "fr")).output == "Hello andras!" + + +def test_wrong_order() -> None: with pytest.raises(WrongDecoratorOrderError): @parameter("name", validator=lambda v: len(v) > 5) @GreatAI.create - def hello_world(name: str) -> str: + async def hello_world(name: str) -> str: return f"Hello {name}!" diff --git a/tests/great_ai/test_basic_starters.py b/tests/test_basic_starters.py similarity index 81% rename from tests/great_ai/test_basic_starters.py rename to tests/test_basic_starters.py index d03684c..47f6617 100644 --- a/tests/great_ai/test_basic_starters.py +++ b/tests/test_basic_starters.py @@ -1,7 +1,6 @@ from functools import lru_cache import pytest - from great_ai import ( ArgumentValidationError, GreatAI, @@ -18,7 +17,7 @@ def test_create_trivial_cases() -> None: assert hello_world_1("andras").output == "Hello andras!" @GreatAI.create - def hello_world_2(name): + def hello_world_2(name): # type: ignore return f"Hello {name}!" assert hello_world_2("andras").output == "Hello andras!" @@ -30,7 +29,7 @@ def test_create_trivial_cases() -> None: assert hello_world_3("andras").output == "Hello andras!" @GreatAI.create() - def hello_world_4(name): + def hello_world_4(name): # type: ignore return f"Hello {name}!" assert hello_world_4("andras").output == "Hello andras!" @@ -39,17 +38,17 @@ def test_create_trivial_cases() -> None: def test_create_with_other_decorator() -> None: @GreatAI.create @lru_cache - def hello_world(name: str) -> str: + def hello_world_1(name: str) -> str: return f"Hello {name}!" - assert hello_world("andras").output == "Hello andras!" + assert hello_world_1("andras").output == "Hello andras!" @lru_cache @GreatAI.create() - def hello_world(name: str) -> str: + def hello_world_2(name: str) -> str: return f"Hello {name}!" - assert hello_world("andras").output == "Hello andras!" + assert hello_world_2("andras").output == "Hello andras!" def test_with_parameter() -> None: @@ -63,6 +62,8 @@ def test_with_parameter() -> None: with pytest.raises(ArgumentValidationError): hello_world("short") + +def test_wrong_order() -> None: with pytest.raises(WrongDecoratorOrderError): @parameter("name", validator=lambda v: len(v) > 5) diff --git a/tests/test_great_ai.py b/tests/test_great_ai.py new file mode 100644 index 0000000..17e5e65 --- /dev/null +++ b/tests/test_great_ai.py @@ -0,0 +1,22 @@ +from asyncio import sleep + +import pytest +from great_ai import GreatAI + + +def test_process_batch() -> None: + @GreatAI.create + def f(x: int) -> int: + return x + 2 + + assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36] + + +@pytest.mark.asyncio +async def test_process_batch_async() -> None: + @GreatAI.create + async def f(x: int) -> int: + await sleep(0.2) + return x + 2 + + assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36] diff --git a/tests/utilities/test_chunk.py b/tests/utilities/test_chunk.py index 5e706c6..7976fd0 100644 --- a/tests/utilities/test_chunk.py +++ b/tests/utilities/test_chunk.py @@ -1,5 +1,4 @@ import pytest - from great_ai.utilities import chunk @@ -23,7 +22,7 @@ def test_bad_argument() -> None: def test_generator() -> None: - def my_generator(): + def my_generator(): # type: ignore for i in range(1, 11): yield i diff --git a/tests/utilities/test_config_file.py b/tests/utilities/test_config_file.py index d26d0c5..eac0aa2 100644 --- a/tests/utilities/test_config_file.py +++ b/tests/utilities/test_config_file.py @@ -2,7 +2,6 @@ import os from pathlib import Path import pytest - from great_ai.utilities import ConfigFile DATA_PATH = Path(__file__).parent.resolve() / "data" diff --git a/tests/utilities/test_language.py b/tests/utilities/test_language.py index 5b08f2b..60df338 100644 --- a/tests/utilities/test_language.py +++ b/tests/utilities/test_language.py @@ -20,11 +20,11 @@ def test_is_english() -> None: assert not is_english(None) -def english_name_of_language() -> None: +def test_english_name_of_language() -> None: assert english_name_of_language("en") == "English" assert english_name_of_language("hu") == "Hungarian" assert english_name_of_language("zh") == "Chinese" - assert english_name_of_language("zh-TW") == "Chinese" + assert english_name_of_language("zh-TW") == "Chinese (Taiwan)" assert english_name_of_language("und") == "Unknown language" assert english_name_of_language("") == "Unknown language" assert english_name_of_language(None) == "Unknown language" diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index 5d4bd27..9fe2b8f 100644 --- a/tests/utilities/test_parallel_map.py +++ b/tests/utilities/test_parallel_map.py @@ -1,11 +1,13 @@ -import pytest +from typing import Any, Iterable +import pytest from great_ai.utilities import WorkerException, parallel_map +from typing_extensions import Never COUNT = int(1e5) + 3 -def test_simple_case_with_progress_bar() -> None: +def test_simple_case() -> None: assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=4)) == [ v**2 for v in range(COUNT) ] @@ -14,7 +16,7 @@ def test_simple_case_with_progress_bar() -> None: def test_with_iterable() -> None: from time import sleep - def my_generator(): + def my_generator() -> Iterable[int]: for i in range(10): yield i sleep(0.1) @@ -26,12 +28,6 @@ def test_with_iterable() -> None: ) -def test_simple_case_without_progress_bar() -> None: - assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=2)) == [ - v**2 for v in range(COUNT) - ] - - def test_simple_case_invalid_values() -> None: with pytest.raises(AssertionError): list(parallel_map(lambda v: v**2, range(COUNT), concurrency=0)) @@ -41,7 +37,7 @@ def test_simple_case_invalid_values() -> None: def test_this_process_exception() -> None: - def my_generator(): + def my_generator() -> Iterable[int]: yield 1 yield 2 yield 3 @@ -59,7 +55,7 @@ def test_this_process_exception() -> None: def test_ignore_this_process_exception() -> None: - def my_generator(): + def my_generator() -> Iterable[float]: yield 1 yield 2 yield 3 @@ -74,19 +70,10 @@ def test_ignore_this_process_exception() -> None: ignore_exceptions=True, ) ) == [1, 4] - assert list( - parallel_map( - lambda v: v**2, - my_generator(), - concurrency=1, - chunk_size=2, - ignore_exceptions=True, - ) - ) == [1, 4, 9] def test_worker_process_exception() -> None: - def oh_no(_): + def oh_no(_: Any) -> Never: raise ValueError("hi") with pytest.raises(WorkerException): @@ -97,7 +84,7 @@ def test_worker_process_exception() -> None: def test_ignore_worker_process_exception() -> None: - def oh_no(_): + def oh_no(_: Any) -> Never: raise ValueError("hi") assert ( diff --git a/tests/utilities/test_threaded_parallel_map.py b/tests/utilities/test_threaded_parallel_map.py index be6f33a..e2efeb7 100644 --- a/tests/utilities/test_threaded_parallel_map.py +++ b/tests/utilities/test_threaded_parallel_map.py @@ -1,11 +1,13 @@ -import pytest +from typing import Any, Iterable +import pytest from great_ai.utilities import WorkerException, threaded_parallel_map +from typing_extensions import Never COUNT = int(1e5) + 3 -def test_simple_case_with_progress_bar() -> None: +def test_simple_case() -> None: assert list( threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4) ) == [v**2 for v in range(COUNT)] @@ -14,7 +16,7 @@ def test_simple_case_with_progress_bar() -> None: def test_with_iterable() -> None: from time import sleep - def my_generator(): + def my_generator() -> Iterable[int]: for i in range(10): yield i sleep(0.1) @@ -27,12 +29,6 @@ def test_with_iterable() -> None: ) -def test_simple_case_without_progress_bar() -> None: - assert list( - threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=2) - ) == [v**2 for v in range(COUNT)] - - def test_simple_case_invalid_values() -> None: with pytest.raises(AssertionError): list(threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=0)) @@ -42,7 +38,7 @@ def test_simple_case_invalid_values() -> None: def test_this_worker_exception() -> None: - def my_generator(): + def my_generator() -> Iterable[int]: yield 1 yield 2 yield 3 @@ -64,7 +60,7 @@ def test_this_worker_exception() -> None: def test_ignore_this_worker_exception() -> None: - def my_generator(): + def my_generator() -> Iterable[float]: yield 1 yield 2 yield 3 @@ -78,20 +74,14 @@ def test_ignore_this_worker_exception() -> None: chunk_size=2, ignore_exceptions=True, ) - ) == [1, 4] - assert list( - threaded_parallel_map( - lambda v: v**2, - my_generator(), - concurrency=1, - chunk_size=2, - ignore_exceptions=True, - ) - ) == [1, 4, 9] + ) == [ + 1, + 4, + ] # the second chunk is ruined because of the error def test_worker_worker_exception() -> None: - def oh_no(_): + def oh_no(_: Any) -> Never: raise ValueError("hi") with pytest.raises(WorkerException): @@ -102,7 +92,7 @@ def test_worker_worker_exception() -> None: def test_ignore_worker_worker_exception() -> None: - def oh_no(_): + def oh_no(_: Any) -> Never: raise ValueError("hi") assert (