wool 0.1rc7__tar.gz → 0.1rc8__tar.gz

{wool-0.1rc7 → wool-0.1rc8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wool
-Version: 0.1rc7
+Version: 0.1rc8
 Summary: A Python framework for distributed multiprocessing.
 Author-email: Conrad Bzura <conrad@wool.io>
 Maintainer-email: maintainers@wool.io
@@ -206,19 +206,24 @@ License:                                  Apache License
            See the License for the specific language governing permissions and
            limitations under the License.
 Classifier: Intended Audience :: Developers
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.10
 Requires-Dist: annotated-types
 Requires-Dist: click
 Requires-Dist: debugpy
+Requires-Dist: grpcio
 Requires-Dist: protobuf
 Requires-Dist: shortuuid
 Requires-Dist: tblib
+Requires-Dist: typing-extensions
 Provides-Extra: dev
 Requires-Dist: pytest; extra == 'dev'
 Requires-Dist: pytest-asyncio; extra == 'dev'
+Requires-Dist: pytest-grpc-aio~=0.2.0; extra == 'dev'
 Requires-Dist: ruff; extra == 'dev'
 Provides-Extra: locking
-Requires-Dist: wool-locking==0.1rc7; extra == 'locking'
+Requires-Dist: wool-locking==0.1rc8; extra == 'locking'
 Description-Content-Type: text/markdown
 
 # Wool

{wool-0.1rc7 → wool-0.1rc8}/pyproject.toml

@@ -1,11 +1,25 @@
 [build-system]
+
 build-backend = "hatchling.build"
 requires = ["debugpy", "hatchling", "packaging", "GitPython", "toml"]
 
 [project]
 authors = [{ name = "Conrad Bzura", email = "conrad@wool.io" }]
-classifiers = ["Intended Audience :: Developers"]
-dependencies = ["annotated-types", "click", "debugpy", "protobuf", "shortuuid", "tblib"]
+classifiers = [
+    "Intended Audience :: Developers",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: POSIX :: Linux",
+]
+dependencies = [
+    "annotated-types",
+    "click",
+    "debugpy",
+    "grpcio",
+    "protobuf",
+    "shortuuid",
+    "tblib",
+    "typing-extensions",
+]
 description = "A Python framework for distributed multiprocessing."
 dynamic = ["version"]
 license = { file = "../LICENSE" }
@@ -15,7 +29,7 @@ readme = "README.md"
 requires-python = ">=3.10"
 
 [project.optional-dependencies]
-dev = ["pytest", "pytest-asyncio", "ruff"]
+dev = ["pytest", "pytest-asyncio", "pytest-grpc-aio~=0.2.0", "ruff"]
 locking = ["wool-locking"]
 
 [project.scripts]
@@ -23,7 +37,7 @@ wool = "wool._cli:cli"
 
 [tool.hatch.build.hooks.protobuf]
 dependencies = ["hatch-protobuf"]
-generate_grpc = false
+generate_grpc = true
 generate_pyi = true
 output_path = "src/wool/_protobuf"
 proto_paths = ["protobuf"]
@@ -55,7 +69,7 @@ select = ["E", "F", "I"]
 quote-style = "double"
 docstring-code-format = true
 
-[tool.ruff.lint.isort] 
+[tool.ruff.lint.isort]
 combine-as-imports = false
 force-single-line = true
 known-first-party = ["wool"]

wool-0.1rc8/wool/_mempool/__init__.py

@@ -0,0 +1,4 @@
+from wool._mempool._mempool import MemoryPool
+from wool._mempool._service import MemoryPoolService
+
+__all__ = ["MemoryPool", "MemoryPoolService"]

wool-0.1rc8/wool/_mempool/_mempool.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import mmap
+import os
+import pathlib
+import shutil
+from contextlib import asynccontextmanager
+from typing import BinaryIO
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+import shortuuid
+
+from wool._mempool._metadata import MetadataMessage
+
+
+class SharedObject:
+    _id: str
+    _mempool: MemoryPool
+    _file: BinaryIO
+    _mmap: mmap.mmap
+    _size: int
+    _md5: bytes
+
+    def __init__(self, id: str, *, mempool: MemoryPool):
+        self._id = id
+        self._mempool = mempool
+        self._file = open(self._path / "dump", "r+b")
+        self._mmap = mmap.mmap(self._file.fileno(), 0)
+        self._size = self.metadata.size
+        self._md5 = self.metadata.md5
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def metadata(self) -> SharedObjectMetadata:
+        return SharedObjectMetadata(self.id, mempool=self._mempool)
+
+    @property
+    def mmap(self) -> mmap.mmap:
+        return self._mmap
+
+    @property
+    def _path(self) -> pathlib.Path:
+        return pathlib.Path(self._mempool.path, self.id)
+
+    def close(self):
+        self.metadata.close()
+        self._mmap.close()
+        self._file.close()
+
+    def refresh(self) -> Self:
+        if self._size != self.metadata.size or self._md5 != self.metadata.md5:
+            self._mmap.close()
+            self._file.close()
+            self._file = open(self._path / "dump", "r+b")
+            self._mmap = mmap.mmap(self._file.fileno(), 0)
+            self._size = self.metadata.size
+            self._md5 = self.metadata.md5
+        return self
+
+
+class SharedObjectMetadata:
+    _id: str
+    _mempool: MemoryPool
+    _file: BinaryIO
+    _mmap: mmap.mmap
+    _instances: dict[str, SharedObjectMetadata] = {}
+
+    def __new__(cls, id: str, *, mempool: MemoryPool):
+        if id in cls._instances:
+            return cls._instances[id]
+        return super().__new__(cls)
+
+    def __init__(self, id: str, mempool: MemoryPool):
+        self._id = id
+        self._mempool = mempool
+        self._file = open(self._path / "meta", "r+b")
+        self._mmap = mmap.mmap(self._file.fileno(), 0)
+        self._instances[id] = self
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def mutable(self) -> bool:
+        return self._metadata.mutable
+
+    @property
+    def size(self) -> int:
+        return self._metadata.size
+
+    @property
+    def md5(self) -> bytes:
+        return self._metadata.md5
+
+    @property
+    def mmap(self) -> mmap.mmap:
+        return self._mmap
+
+    @property
+    def _path(self) -> pathlib.Path:
+        return pathlib.Path(self._mempool.path, self.id)
+
+    @property
+    def _metadata(self) -> MetadataMessage:
+        return MetadataMessage.loads(bytes(self._mmap))
+
+    def close(self):
+        self._mmap.close()
+        self._file.close()
+        del self._instances[self.id]
+
+
+class MemoryPool:
+    _objects: dict[str, SharedObject]
+    _path: pathlib.Path
+
+    def __init__(self, path: str | pathlib.Path = pathlib.Path(".mempool")):
+        if isinstance(path, str):
+            self._path = pathlib.Path(path)
+        else:
+            self._path = path
+        self._lockdir = self._path / "locks"
+        os.makedirs(self._lockdir, exist_ok=True)
+        self._acquire(f"pid-{os.getpid()}")
+        self._objects = dict()
+
+    def __contains__(self, ref: str) -> bool:
+        return ref in self._objects
+
+    def __del__(self):
+        self._release(f"pid-{os.getpid()}")
+
+    @property
+    def path(self) -> pathlib.Path:
+        return self._path
+
+    async def map(self, ref: str | None = None):
+        if ref is not None and ref not in self._objects:
+            if self._locked(f"delete-{ref}"):
+                raise RuntimeError(
+                    f"Reference {ref} is currently locked for deletion"
+                )
+            async with self._reference_lock(ref):
+                self._map(ref)
+        else:
+            for entry in os.scandir(self._path):
+                if entry.is_dir() and (ref := entry.name) != "locks":
+                    if not self._locked(f"delete-{ref}"):
+                        async with self._reference_lock(ref):
+                            self._map(ref)
+
+    async def put(
+        self, dump: bytes, *, mutable: bool = False, ref: str | None = None
+    ) -> str:
+        ref = ref or str(shortuuid.uuid())
+        async with self._reference_lock(ref):
+            self._put(ref, dump, mutable=mutable, exist_ok=False)
+            return ref
+
+    async def post(self, ref: str, dump: bytes) -> bool:
+        if self._locked(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        async with self._reference_lock(ref):
+            if ref not in self._objects:
+                self._map(ref)
+            obj = self._objects[ref]
+            if not obj.metadata.mutable:
+                raise ValueError("Cannot modify an immutable reference")
+            if (size := len(dump)) != obj.metadata.size:
+                try:
+                    obj.mmap.resize(size)
+                    self._post(ref, obj, dump)
+                except SystemError:
+                    self._put(ref, dump, mutable=True, exist_ok=True)
+                return True
+            elif hashlib.md5(dump).digest() != obj.metadata.md5:
+                self._post(ref, obj, dump)
+                return True
+            else:
+                return False
+
+    async def get(self, ref: str) -> bytes:
+        if self._locked(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        async with self._reference_lock(ref):
+            if ref not in self._objects:
+                self._map(ref)
+            return bytes(self._objects[ref].refresh().mmap)
+
+    async def delete(self, ref: str):
+        async with self._delete_lock(ref):
+            async with self._reference_lock(ref):
+                if ref not in self._objects:
+                    self._map(ref)
+                self._objects.pop(ref).close()
+                try:
+                    shutil.rmtree(self.path / ref)
+                except FileNotFoundError:
+                    pass
+
+    def _put(
+        self,
+        ref: str,
+        dump: bytes,
+        *,
+        mutable: bool = False,
+        exist_ok: bool = False,
+    ):
+        metadata = MetadataMessage(
+            ref=ref,
+            mutable=mutable,
+            size=len(dump),
+            md5=hashlib.md5(dump).digest(),
+        )
+
+        refpath = pathlib.Path(self._path, f"{ref}")
+        os.makedirs(refpath, exist_ok=exist_ok)
+
+        with open(refpath / "meta", "wb") as metafile:
+            metafile.write(metadata.dumps())
+
+        with open(refpath / "dump", "wb") as dumpfile:
+            dumpfile.write(dump)
+
+        self._map(ref)
+
+    def _post(self, ref: str, obj: SharedObject, dump: bytes):
+        if not obj.metadata.mutable:
+            raise ValueError("Cannot modify an immutable reference")
+        metadata = MetadataMessage(
+            ref=ref,
+            mutable=True,
+            size=len(dump),
+            md5=hashlib.md5(dump).digest(),
+        )
+        obj.metadata.mmap[:] = metadata.dumps()
+        obj.metadata.mmap.flush()
+        obj.mmap.seek(0)
+        obj.mmap.write(dump)
+
+    def _map(self, ref: str):
+        obj = self._objects.pop(ref, None)
+        if obj:
+            obj.close()
+        self._objects[ref] = SharedObject(id=ref, mempool=self)
+
+    def _lockpath(self, key: str) -> pathlib.Path:
+        return pathlib.Path(self._lockdir, f"{key}.lock")
+
+    def _acquire(self, key: str) -> bool:
+        try:
+            os.symlink(f"{key}", self._lockpath(key))
+            return True
+        except FileExistsError:
+            return False
+
+    def _release(self, key: str):
+        try:
+            if os.path.islink(lock_path := self._lockpath(key)):
+                os.unlink(lock_path)
+        except FileNotFoundError:
+            pass
+
+    def _locked(self, key: str) -> bool:
+        return os.path.islink(self._lockpath(key))
+
+    @asynccontextmanager
+    async def _reference_lock(self, ref: str):
+        try:
+            while not self._acquire(ref):
+                await asyncio.sleep(0)
+            yield
+        finally:
+            self._release(ref)
+
+    @asynccontextmanager
+    async def _delete_lock(self, ref: str):
+        key = f"delete-{ref}"
+        if not self._acquire(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        else:
+            yield
+        self._release(key)

wool-0.1rc7/wool/_mempool/_metadata/__init__.py → wool-0.1rc8/wool/_mempool/_metadata.py

@@ -4,9 +4,7 @@ import logging
 from dataclasses import dataclass
 
 try:
-    from wool._protobuf._mempool._metadata._metadata_pb2 import (
-        _MetadataMessage,
-    )
+    from wool._protobuf.mempool.metadata.metadata_pb2 import _MetadataMessage
 except ImportError:
     logging.error(
         "Failed to import _MetadataMessage. "

wool-0.1rc8/wool/_mempool/_service.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncGenerator
+from typing import Final
+from weakref import WeakSet
+from weakref import WeakValueDictionary
+
+import shortuuid
+from grpc.aio import ServicerContext
+
+from wool._mempool import MemoryPool
+from wool._protobuf.mempool import mempool_pb2 as proto
+from wool._protobuf.mempool import mempool_pb2_grpc as rpc
+
+
+class Session:
+    """
+    A session represents a client connection to the memory pool service and
+    serves as the scope for any shared references acquired over its duration.
+    """
+
+    id: Final[str]
+    queue: Final[asyncio.Queue[proto.SessionResponse]]
+    references: Final[set[Reference]]
+    sessions: Final[WeakValueDictionary[str, Session]] = WeakValueDictionary()
+
+    @classmethod
+    def get(cls, id: str) -> Session | None:
+        return cls.sessions.get(id)
+
+    def __init__(self):
+        self.id = shortuuid.uuid()
+        self.queue = asyncio.Queue()
+        self.references = set()
+        self.sessions[self.id] = self
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, Session):
+            return self.id == other.id
+        return False
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class Reference:
+    id: Final[str]
+    mempool: Final[MemoryPool]
+    sessions: Final[WeakSet[Session]]
+
+    _references: Final[WeakValueDictionary[str, Reference]] = (
+        WeakValueDictionary()
+    )
+    _to_delete: Final[set[str]] = set()
+    _initialized: bool = False
+
+    @classmethod
+    def get(cls, id: str) -> Reference | None:
+        return cls._references.get(id)
+
+    @classmethod
+    def new(cls, id: str, *, mempool: MemoryPool) -> Reference:
+        if id in cls._references:
+            raise ValueError(f"Reference {id} already exists")
+        return cls(id, mempool=mempool)
+
+    def __new__(cls, id: str, *, mempool: MemoryPool):
+        if id in cls._references:
+            if id in cls._to_delete:
+                cls._to_delete.remove(id)
+            return cls._references[id]
+        return super().__new__(cls)
+
+    def __init__(self, id: str, *, mempool: MemoryPool):
+        if not self._initialized:
+            self.id = id
+            self.mempool = mempool
+            self.sessions = WeakSet()
+            self._references[id] = self
+            self._initialized = True
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, Reference):
+            return self.id == other.id
+        return False
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+    def __del__(self):
+        self._to_delete.add(self.id)
+
+        id = self.id
+        to_delete = self._to_delete
+        mempool = self.mempool
+
+        async def _delete():
+            if id in to_delete:
+                try:
+                    to_delete.remove(id)
+                    await mempool.delete(id)
+                except FileNotFoundError:
+                    pass
+
+        try:
+            asyncio.get_running_loop().create_task(_delete())
+        except RuntimeError:
+            asyncio.new_event_loop().run_until_complete(_delete())
+
+
+class MemoryPoolService(rpc.MemoryPoolServicer):
+    def __init__(self, mempool: MemoryPool | None = None):
+        self._mempool = mempool or MemoryPool()
+        self._shutdown = asyncio.Event()
+
+    async def session(
+        self, request: proto.SessionRequest, context: ServicerContext
+    ) -> AsyncGenerator[proto.SessionResponse]:
+        session = Session()
+        yield proto.SessionResponse(session=proto.Session(id=session.id))
+        while True:
+            yield await session.queue.get()
+
+    async def acquire(
+        self, request: proto.AcquireRequest, context: ServicerContext
+    ) -> proto.AcquireResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        if not (reference := Reference.get(request.reference.id)):
+            raise ValueError(f"Reference {request.reference.id} not found")
+        session.references.add(reference)
+        reference.sessions.add(session)
+        return proto.AcquireResponse()
+
+    async def map(
+        self, request: proto.AcquireRequest, context: ServicerContext
+    ) -> proto.AcquireResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        await self._mempool.map(request.reference.id)
+        reference = Reference(request.reference.id, mempool=self._mempool)
+        await self.acquire(
+            proto.AcquireRequest(
+                session=proto.Session(id=session.id),
+                reference=proto.Reference(id=reference.id),
+            ),
+            context,
+        )
+        return proto.AcquireResponse()
+
+    async def put(
+        self, request: proto.PutRequest, context: ServicerContext
+    ) -> proto.PutResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        reference = Reference(
+            id=await self._mempool.put(request.dump, mutable=request.mutable),
+            mempool=self._mempool,
+        )
+        await self.acquire(
+            proto.AcquireRequest(
+                session=proto.Session(id=session.id),
+                reference=proto.Reference(id=reference.id),
+            ),
+            context,
+        )
+        return proto.PutResponse(reference=proto.Reference(id=reference.id))
+
+    async def get(
+        self, request: proto.GetRequest, context: ServicerContext
+    ) -> proto.GetResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        if not (reference := Reference.get(request.reference.id)):
+            raise ValueError(f"Reference {request.reference.id} not found")
+        if reference not in session.references:
+            await self.acquire(
+                proto.AcquireRequest(
+                    session=proto.Session(id=session.id),
+                    reference=proto.Reference(id=reference.id),
+                ),
+                context,
+            )
+        dump = await self._mempool.get(reference.id)
+        return proto.GetResponse(dump=dump)
+
+    async def post(
+        self, request: proto.PostRequest, context: ServicerContext
+    ) -> proto.PostResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        if not (reference := Reference.get(request.reference.id)):
+            raise ValueError(f"Reference {request.reference.id} not found")
+        if reference not in session.references:
+            await self.acquire(
+                proto.AcquireRequest(
+                    session=proto.Session(id=session.id),
+                    reference=proto.Reference(id=reference.id),
+                ),
+                context,
+            )
+        updated = await self._mempool.post(request.reference.id, request.dump)
+        if updated:
+            for session in Reference(
+                id=request.reference.id, mempool=self._mempool
+            ).sessions:
+                if session.id is not request.session.id:
+                    event = proto.Event(
+                        reference=request.reference,
+                        event_type="post",
+                    )
+                    await session.queue.put(proto.SessionResponse(event=event))
+        return proto.PostResponse(updated=updated)
+
+    async def release(
+        self, request: proto.ReleaseRequest, context: ServicerContext
+    ) -> proto.ReleaseResponse:
+        if not (session := Session.get(request.session.id)):
+            raise ValueError(f"Session {request.session.id} not found")
+        if not (reference := Reference.get(request.reference.id)):
+            raise ValueError(f"Reference {request.reference.id} not found")
+        session.references.remove(reference)
+        reference.sessions.remove(session)
+        return proto.ReleaseResponse()

wool-0.1rc8/wool/_protobuf/__init__.py

@@ -0,0 +1,4 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))