Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

New Progress Bar, Backoff, Batching #165

Open
wants to merge 16 commits into
base: main
Choose a base branch
from
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -177,7 +177,7 @@ features = ["pyo3/extension-module"]
where = ["src"]

[tool.setuptools.package-data]
dolma = ["py.typed", "data/*"]
dolma = ["py.typed", "data/*", "*.pyi"]

[tool.black]
line-length = 115
19 changes: 19 additions & 0 deletions python/dolma/core/mp_tools.pyi
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from collections.abc import Callable, Iterable
from multiprocessing.managers import SyncManager
from multiprocessing.pool import ApplyResult, Pool
from typing import Any

class ResultWithDebug(ApplyResult): ... # noqa: E701,E302
class ManagerWithDebug(SyncManager): ... # noqa: E701

class PoolWithDebug(Pool): # noqa: E302
def __init__( # noqa: E704
self,
processes: int | None = None,
initializer: Callable[..., Any] | None = None,
initargs: Iterable[Any] = (),
maxtasksperchild: int | None = None,
debug: bool = False,
): ...

def get_manager(pool: Pool) -> SyncManager: ... # noqa: E701, E704, E302