paasta-tools 1.21.3__py3-none-any.whl

paasta_tools/mesos/mesos_file.py

@@ -0,0 +1,169 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+from . import exceptions
+from paasta_tools.async_utils import async_ttl_cache
+
+
+class File:
+
+    chunk_size = 1024
+
+    def __init__(self, host, task=None, path=None):
+        self.host = host
+        self.task = task
+        self.path = path
+
+        if self.task is None:
+            self._host_path = self.path
+        else:
+            self._host_path = None  # Defer until later (_fetch) so we don't make HTTP requests in __init__.
+
+        self._offset = 0
+
+        # Used during fetch, class level so the dict isn't constantly alloc'd
+        self._params = {
+            "path": self._host_path,
+            "offset": -1,
+            "length": self.chunk_size,
+        }
+
+    def __eq__(self, y):
+        return self.key() == y.key()
+
+    def __hash__(self):
+        return hash(self.__str__())
+
+    def __repr__(self):
+        return f"<open file '{self.path}', for '{self._where}'>"
+
+    def __str__(self):
+        return f"{self._where}:{self.path}"
+
+    def key(self):
+        return "{}:{}".format(self.host.key(), self._host_path)
+
+    @property
+    def _where(self):
+        return self.task["id"] if self.task is not None else self.host.key()
+
+    async def _fetch(self):
+        # fill in path if it wasn't set in __init__
+        if self._params["path"] is None:
+            self._params["path"] = os.path.join(await self.task.directory(), self.path)
+
+        resp = await self.host.fetch("/files/read.json", params=self._params)
+        if resp.status == 404:
+            raise exceptions.FileDoesNotExist("No such file or directory.")
+        return await resp.json()
+
+    async def exists(self):
+        try:
+            await self.size()
+            return True
+        except exceptions.FileDoesNotExist:
+            return False
+        except exceptions.SlaveDoesNotExist:
+            return False
+
+    # When reading a file, it is common to first check whether it exists, then
+    # look at the size to determine where to seek. Instead of requiring
+    # multiple requests to the slave, the size is cached for a very short
+    # period of time.
+    @async_ttl_cache(ttl=0.5, cleanup_self=True)
+    async def size(self):
+        return (await self._fetch())["offset"]
+
+    async def seek(self, offset, whence=os.SEEK_SET):
+        if whence == os.SEEK_SET:
+            self._offset = 0 + offset
+        elif whence == os.SEEK_CUR:
+            self._offset += offset
+        elif whence == os.SEEK_END:
+            self._offset = await self.size() + offset
+
+    def tell(self):
+        return self._offset
+
+    def _length(self, start, size):
+        if size and self.tell() - start + self.chunk_size > size:
+            return size - (self.tell() - start)
+        return self.chunk_size
+
+    async def _get_chunk(self, loc, size=None):
+        if size is None:
+            size = self.chunk_size
+
+        await self.seek(loc, os.SEEK_SET)
+        self._params["offset"] = loc
+        self._params["length"] = size
+
+        data = (await self._fetch())["data"]
+        await self.seek(len(data), os.SEEK_CUR)
+        return data
+
+    async def _read(self, size=None):
+        start = self.tell()
+
+        def pre(x):
+            return x == ""
+
+        def post(x):
+            return size and (self.tell() - start) >= size
+
+        blob = None
+        while blob != "" and not (size and (self.tell() - start) >= size):
+            blob = await self._get_chunk(self.tell(), size=self._length(start, size))
+            yield blob
+
+    async def _read_reverse(self, size=None):
+        fsize = await self.size()
+        if not size:
+            size = fsize
+
+        def next_block():
+            current = fsize
+            while (current - self.chunk_size) > (fsize - size):
+                current -= self.chunk_size
+                yield current
+
+        for pos in next_block():
+            yield await self._get_chunk(pos)
+
+        yield await self._get_chunk(fsize - size, size % self.chunk_size)
+
+    async def _readlines(self, size=None):
+        last = ""
+        async for blob in self._read(size):
+
+            # This is not streaming and assumes small chunk sizes
+            blob_lines = (last + blob).split("\n")
+            for line in blob_lines[: len(blob_lines) - 1]:
+                yield line
+
+            last = blob_lines[-1]
+
+    async def _readlines_reverse(self, size=None):
+        buf = ""
+        async for blob in self._read_reverse(size):
+
+            blob_lines = (blob + buf).split("\n")
+            for line in reversed(blob_lines[1:]):
+                yield line
+
+            buf = blob_lines[0]
+        yield buf

paasta_tools/mesos/parallel.py

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import concurrent.futures
+import contextlib
+
+from . import exceptions
+
+
+@contextlib.contextmanager
+def execute(max_workers):
+    try:
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+        yield executor
+    except KeyboardInterrupt:
+        # Threads in the ThreadPoolExecutor are created with
+        # daemon=True. There is, therefore, an atexit function registered
+        # that allows all the currently running threads to stop before
+        # allowing the interpreter to stop. Because we don't care whether
+        # the worker threads exit cleanly or not, we force shutdown to be
+        # immediate.
+        concurrent.futures.thread._threads_queues.clear()
+        raise
+    finally:
+        executor.shutdown(wait=False)
+
+
+def stream(fn, elements, workers):
+    """Yield the results of fn as jobs complete."""
+    jobs = []
+
+    with execute(workers) as executor:
+        for elem in elements:
+            jobs.append(executor.submit(fn, elem))
+
+        for job in concurrent.futures.as_completed(jobs):
+            try:
+                yield job.result()
+            except exceptions.SkipResult:
+                pass

paasta_tools/mesos/slave.py

@@ -0,0 +1,115 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from urllib.parse import urljoin
+
+import aiohttp
+
+from . import exceptions
+from . import mesos_file
+from . import util
+from paasta_tools.async_utils import async_ttl_cache
+from paasta_tools.utils import get_user_agent
+
+
+class MesosSlave:
+    def __init__(self, config, items):
+        self.config = config
+        self.__items = items
+
+    def __getitem__(self, name):
+        return self.__items[name]
+
+    def __str__(self):
+        return self.key()
+
+    def key(self):
+        return self["pid"].split("@")[-1]
+
+    @property
+    def host(self):
+        return "{}://{}:{}".format(
+            self.config["scheme"], self["hostname"], self["pid"].split(":")[-1]
+        )
+
+    async def fetch(self, url, **kwargs) -> aiohttp.ClientResponse:
+        headers = {"User-Agent": get_user_agent()}
+        async with aiohttp.ClientSession(
+            conn_timeout=self.config["response_timeout"],
+            read_timeout=self.config["response_timeout"],
+        ) as session:
+            try:
+                async with session.get(
+                    urljoin(self.host, url), headers=headers, **kwargs
+                ) as response:
+                    await response.text()
+                    return response
+            except aiohttp.ClientConnectionError:
+                raise exceptions.SlaveDoesNotExist(
+                    f"Unable to connect to the slave at {self.host}"
+                )
+
+    @async_ttl_cache(ttl=5, cleanup_self=True)
+    async def state(self):
+        return await (await self.fetch("/slave(1)/state.json")).json()
+
+    async def frameworks(self):
+        return util.merge(await self.state(), "frameworks", "completed_frameworks")
+
+    async def task_executor(self, task_id):
+        for fw in await self.frameworks():
+            for exc in util.merge(fw, "executors", "completed_executors"):
+                if task_id in list(
+                    map(
+                        lambda x: x["id"],
+                        util.merge(exc, "completed_tasks", "tasks", "queued_tasks"),
+                    )
+                ):
+                    return exc
+        raise exceptions.MissingExecutor("No executor has a task by that id")
+
+    async def file_list(self, path):
+        # The sandbox does not exist on the slave.
+        if path == "":
+            return []
+
+        resp = self.fetch("/files/browse.json", params={"path": path})
+        if resp.status_code == 404:
+            return []
+        return await resp.json()
+
+    def file(self, task, path):
+        return mesos_file.File(self, task, path)
+
+    @async_ttl_cache(ttl=30, cleanup_self=True)
+    async def stats(self):
+        return await (await self.fetch("/monitor/statistics.json")).json()
+
+    def executor_stats(self, _id):
+        return list(filter(lambda x: x["executor_id"]))
+
+    async def task_stats(self, _id):
+        stats = list(filter(lambda x: x["executor_id"] == _id, await self.stats()))
+
+        # Tasks that are not yet in a RUNNING state have no stats.
+        if len(stats) == 0:
+            return {}
+        else:
+            return stats[0]["statistics"]
+
+    @property  # type: ignore
+    @util.memoize
+    def log(self):
+        return mesos_file.File(self, path="/slave/log")

paasta_tools/mesos/task.py

@@ -0,0 +1,94 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import re
+
+import a_sync
+
+from . import exceptions
+from . import framework
+from . import mesos_file
+from paasta_tools.async_utils import async_ttl_cache
+
+
+class Task:
+
+    cmd_re = re.compile(r"\(Command: (.+)\)")
+
+    def __init__(self, master, items):
+        self.master = master
+        self.__items = items
+
+    def __str__(self):
+        return "{}:{}".format(a_sync.block(self.slave), self["id"])
+
+    def __getitem__(self, name):
+        return self.__items[name]
+
+    async def executor(self):
+        return await (await self.slave()).task_executor(self["id"])
+
+    async def framework(self):
+        return framework.Framework(await self.master.framework(self["framework_id"]))
+
+    @async_ttl_cache(cleanup_self=True)
+    async def directory(self):
+        try:
+            return (await self.executor())["directory"]
+        except exceptions.MissingExecutor:
+            return ""
+
+    @async_ttl_cache(cleanup_self=True)
+    async def slave(self):
+        return await self.master.slave(self["slave_id"])
+
+    async def file(self, path):
+        return mesos_file.File(await self.slave(), self, path)
+
+    async def file_list(self, path):
+        return await (await self.slave()).file_list(os.path.join(self.directory, path))
+
+    async def stats(self):
+        try:
+            return await (await self.slave()).task_stats(self["id"])
+        except exceptions.MissingExecutor:
+            return {}
+
+    async def cpu_time(self):
+        st = await self.stats()
+        secs = st.get("cpus_user_time_secs", 0) + st.get("cpus_system_time_secs", 0)
+        return secs
+
+    async def cpu_limit(self):
+        return (await self.stats()).get("cpus_limit", 0)
+
+    async def mem_limit(self):
+        return (await self.stats()).get("mem_limit_bytes", 0)
+
+    async def rss(self):
+        return (await self.stats()).get("mem_rss_bytes", 0)
+
+    async def command(self):
+        try:
+            result = self.cmd_re.search((await self.executor())["name"])
+        except exceptions.MissingExecutor:
+            result = None
+        if not result:
+            return "none"
+        return result.group(1)
+
+    async def user(self):
+        return (await self.framework()).user

paasta_tools/mesos/util.py

@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import functools
+import itertools
+import time
+
+
+def merge(obj, *keys):
+    return itertools.chain(*[obj[k] for k in keys])
+
+
+class CachedProperty:
+    def __init__(self, ttl=300):
+        self.ttl = ttl
+
+    def __call__(self, fget, doc=None):
+        self.fget = fget
+        self.__doc__ = doc or fget.__doc__
+        self.__name__ = fget.__name__
+        self.__module__ = fget.__module__
+        return self
+
+    def __get__(self, inst, owner):
+        try:
+            value, last_update = inst._cache[self.__name__]
+            if self.ttl > 0 and time.time() - last_update > self.ttl:
+                raise AttributeError
+        except (KeyError, AttributeError):
+            value = self.fget(inst)
+            try:
+                cache = inst._cache
+            except AttributeError:
+                cache = inst._cache = {}
+            cache[self.__name__] = (value, time.time())
+        return value
+
+
+def memoize(obj):
+    cache = obj.cache = {}
+
+    @functools.wraps(obj)
+    def memoizer(*args, **kwargs):
+        key = str(args) + str(kwargs)
+        if key not in cache:
+            cache[key] = obj(*args, **kwargs)
+        return cache[key]
+
+    return memoizer
+
+
+def humanize_bytes(b):
+    abbrevs = ((1 << 30, "GB"), (1 << 20, "MB"), (1 << 10, "kB"), (1, "B"))
+    for factor, suffix in abbrevs:
+        if b >= factor:
+            break
+    return "%.*f %s" % (2, b / float(factor), suffix)

paasta_tools/mesos/zookeeper.py

@@ -0,0 +1,37 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import contextlib
+
+import kazoo.client
+import kazoo.exceptions
+import kazoo.handlers.threading
+
+
+TIMEOUT = 1
+
+# Helper for testing
+client_class = kazoo.client.KazooClient
+
+
+@contextlib.contextmanager
+def client(*args, **kwargs):
+    zk = client_class(*args, **kwargs)
+    zk.start(timeout=TIMEOUT)
+    try:
+        yield zk
+    finally:
+        zk.stop()
+        zk.close()