wool 0.1rc7__tar.gz → 0.1rc9__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wool
-Version: 0.1rc7
+Version: 0.1rc9
 Summary: A Python framework for distributed multiprocessing.
 Author-email: Conrad Bzura <conrad@wool.io>
 Maintainer-email: maintainers@wool.io
@@ -206,19 +206,25 @@ 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: pytest-mock; extra == 'dev'
 Requires-Dist: ruff; extra == 'dev'
 Provides-Extra: locking
-Requires-Dist: wool-locking==0.1rc7; extra == 'locking'
+Requires-Dist: wool-locking==0.1rc9; extra == 'locking'
 Description-Content-Type: text/markdown
 
 # Wool

{wool-0.1rc7 → wool-0.1rc9}/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,13 @@ 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",
+    "pytest-mock",
+    "ruff",
+]
 locking = ["wool-locking"]
 
 [project.scripts]
@@ -23,7 +43,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 +75,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.1rc9/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.1rc9/wool/_mempool/_client.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Final
+from typing import Protocol
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+import grpc
+
+try:
+    from wool._protobuf.mempool import service_pb2 as pb
+    from wool._protobuf.mempool import service_pb2_grpc as rpc
+except ImportError as e:
+    from wool._protobuf import ProtobufImportError
+
+    raise ProtobufImportError(e) from e
+
+
+class EventHandler(Protocol):
+    @staticmethod
+    def __call__(client: MemoryPoolClient, response: pb.Event) -> None: ...
+
+
+class MemoryPoolClient:
+    _event_stream: Final[grpc.aio.UnaryStreamCall]
+    _listener: asyncio.Task
+    _session: pb.Session | None
+
+    def __init__(self, channel: grpc.aio.Channel, event_handler: EventHandler):
+        self._stub = rpc.MemoryPoolStub(channel)
+        self._event_stream = self._stub.session(pb.SessionRequest())
+        self._event_handler = event_handler
+        self._session = None
+
+    async def __aenter__(self) -> Self:
+        if self._session:
+            raise RuntimeError(
+                "Client session may not be re-entered. "
+                "Use 'async with MemoryPoolClient(...)' only once."
+            )
+        response = await self._event_stream.read()
+        assert isinstance(response, pb.SessionResponse)
+        self._session = response.session
+        self._listener = asyncio.create_task(self._listen())
+        return self
+
+    async def __aexit__(self, *_):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        assert self._listener
+        self._listener.cancel()
+        try:
+            await self._listener
+        except asyncio.CancelledError:
+            pass
+
+    def __del__(self):
+        try:
+            self._event_stream.cancel()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        return self._session.id
+
+    async def map(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.AcquireRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.map(request)
+
+    async def get(self, ref: str) -> bytes:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.GetRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        response: pb.GetResponse = await self._stub.get(request)
+        return response.dump
+
+    async def put(self, dump: bytes, *, mutable: bool = False) -> str:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.PutRequest(
+            session=self._session,
+            dump=dump,
+            mutable=mutable,
+        )
+        response: pb.PutResponse = await self._stub.put(request)
+        return response.reference.id
+
+    async def post(self, ref: str, dump: bytes) -> bool:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.PostRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+            dump=dump,
+        )
+        response: pb.PostResponse = await self._stub.post(request)
+        return response.updated
+
+    async def acquire(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.AcquireRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.acquire(request)
+
+    async def release(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.ReleaseRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.release(request)
+
+    async def _listen(self):
+        assert self._event_stream
+        try:
+            while True:
+                if (response := await self._event_stream.read()) is None:
+                    break
+                assert isinstance(response, pb.SessionResponse), (
+                    f"Unexpected event type: {type(response)}"
+                )
+                self._event_handler(self, response.event)
+        except asyncio.CancelledError:
+            self._event_stream.cancel()

wool-0.1rc9/wool/_mempool/_mempool.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import os
+import pathlib
+import shutil
+from contextlib import asynccontextmanager
+from mmap import mmap
+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
+    _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(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:
+        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(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
+    _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(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:
+        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.1rc9/wool/_mempool/_metadata.py

@@ -1,18 +1,14 @@
 from __future__ import annotations
 
-import logging
+from dataclasses import asdict
 from dataclasses import dataclass
 
 try:
-    from wool._protobuf._mempool._metadata._metadata_pb2 import (
-        _MetadataMessage,
-    )
-except ImportError:
-    logging.error(
-        "Failed to import _MetadataMessage. "
-        "Ensure protocol buffers are compiled."
-    )
-    raise
+    from wool._protobuf.mempool import metadata_pb2 as pb
+except ImportError as e:
+    from wool._protobuf import ProtobufImportError
+
+    raise ProtobufImportError(e) from e
 
 
 @dataclass
@@ -24,7 +20,7 @@ class MetadataMessage:
 
     @classmethod
     def loads(cls, data: bytes) -> MetadataMessage:
-        (metadata := _MetadataMessage()).ParseFromString(data)
+        (metadata := pb.MetadataMessage()).ParseFromString(data)
         return cls(
             ref=metadata.ref,
             mutable=metadata.mutable,
@@ -33,9 +29,7 @@ class MetadataMessage:
         )
 
     def dumps(self) -> bytes:
-        return _MetadataMessage(
-            ref=self.ref, mutable=self.mutable, size=self.size, md5=self.md5
-        ).SerializeToString()
+        return pb.MetadataMessage(**asdict(self)).SerializeToString()
 
 
 __all__ = ["MetadataMessage"]