python-microscopy 25.11.20__cp313-cp313-win_amd64.whl

PYME/IO/cluster_directory.py

@@ -0,0 +1,419 @@
+"""
+Work in progress ... refactor some of the listing stuff out of clusterIO to
+
+a) reduce the size of clusterIO
+b) reduce duplication between, e.g. locate and listdirectory
+c) ultimately allow pluggable directory management / caching, e.g. to enable a central directory server (unclear if this
+   would solve current performance issues, but potentially worth a try).
+
+"""
+import time
+import socket
+import threading
+import requests
+import posixpath
+import sys
+import numpy as np
+
+from PYME import config
+
+from multiprocessing.pool import ThreadPool
+
+from PYME.misc.computerName import GetComputerName
+local_computer_name = GetComputerName()
+
+import logging
+logger = logging.getLogger(__name__)
+
+SERVICE_CACHE_LIFETIME = 1 #seconds
+DIR_CACHE_TIME = 1 #seconds
+
+import PYME.misc.pyme_zeroconf as pzc
+from PYME.misc import hybrid_ns
+
+_ns = None
+_ns_lock = threading.Lock()
+
+if config.get('clusterIO-hybridns', True):
+    def get_ns():
+        global _ns
+        with _ns_lock:
+            if _ns is None:
+                #stagger query times
+                time.sleep(3 * np.random.rand())
+                #_ns = pzc.getNS('_pyme-http')
+                _ns = hybrid_ns.getNS('_pyme-http')
+                #wait for replies
+                time.sleep(5)
+        
+        return _ns
+else:
+    def get_ns():
+        global _ns
+        with _ns_lock:
+            if _ns is None:
+                #stagger query times
+                time.sleep(3 * np.random.rand())
+                #_ns = pzc.getNS('_pyme-http')
+                _ns = pzc.getNS('_pyme-http')
+                #wait for replies
+                time.sleep(5)
+        
+        return _ns
+
+if not 'sphinx' in sys.modules.keys():
+    # do not start zeroconf if running under sphinx
+    # otherwise, spin up a thread to go and find the ns and get on with the rest of our initialization
+    
+    threading.Thread(target=get_ns).start()
+
+from collections import OrderedDict
+
+class _LimitedSizeDict(OrderedDict):
+    def __init__(self, *args, **kwds):
+        self.size_limit = kwds.pop("size_limit", None)
+        OrderedDict.__init__(self, *args, **kwds)
+
+        self._lock = threading.Lock()
+        self._check_size_limit()
+
+    def __setitem__(self, key, value):
+        with self._lock:
+            OrderedDict.__setitem__(self, key, value)
+            self._check_size_limit()
+
+    def _check_size_limit(self):
+        if self.size_limit is not None:
+            while len(self) > self.size_limit:
+                self.popitem(last=False)
+
+class DirectoryInfoManager(object):
+    def __init__(self, ns=None, serverfilter=''):
+        self._dirCache = _LimitedSizeDict(size_limit=100)
+        self._locateCache = _LimitedSizeDict(size_limit=500)
+        
+        if not ns:
+            self._ns = get_ns()
+        else:
+            self._ns = ns
+            
+        self.serverfilter = serverfilter
+        
+        self._service_expiry_time = 0
+        self._cached_servers = None
+
+        #use one session for each server (to allow http keep-alives)
+        self._sessions = {}
+        self._pool = ThreadPool(10)
+        
+        self._list_dir_lock = threading.Lock()
+
+    def _getSession(self, url):
+        if not isinstance(url, bytes):
+            url = url.encode()
+
+        servinfo = url.split(b'/')[2]
+        try:
+            session = self._sessions[servinfo]
+        except KeyError:
+            session = requests.Session()
+            session.trust_env = False
+            self._sessions[servinfo] = session
+
+        return session
+        
+        
+    @property
+    def dataservers(self):
+        """
+        Find all the data servers belonging to the cluster, caching the results
+
+        """
+        t = time.time()
+        if (t > self._service_expiry_time):
+            self._cached_servers = []
+            services = self._ns.get_advertised_services()
+
+            for name, info in services:
+                if self.serverfilter in name:
+                    # only look at servers associated with our cluster
+                    if info is None or info.address is None or info.port is None:
+                        # handle the case where zeroconf gives us bad name info. This  is a result of a race condition within
+                        # zeroconf, which should probably be fixed instead, but hopefully this workaround is enough.
+                        # FIXME - fix zeroconf module race condition on info update.
+                        logger.error('''Zeroconf gave us NULL info, ignoring and hoping for the best ...
+                        Node with bogus info was: %s
+                        Total number of nodes: %d
+                        ''' % (name, len(services)))
+                    else:
+                        serverurl = 'http://%s:%d/' % (socket.inet_ntoa(info.address), info.port)
+            
+                        if local_computer_name in name:
+                            self._cached_servers.insert(0, serverurl)
+                        else:
+                            self._cached_servers.append(serverurl)
+            
+            self._service_expiry_time = t + SERVICE_CACHE_LIFETIME
+        
+        if len(self._cached_servers) < 1:
+            # check that there is a cluster running
+            raise IOError('No data servers found, is the cluster running?')
+        
+        return self._cached_servers
+            
+            
+    def list_single_node_dir(self, dirurl, nRetries = 1, timeout = 10, strict_caching=False):
+        """
+        List the directory on a single node
+        
+        Parameters
+        ----------
+        dirurl
+        nRetries
+        timeout
+
+        Returns
+        -------
+
+        """
+        t = time.time()
+    
+        try:
+            dirL, rt, dt = self._dirCache[dirurl]
+            time_in_cache = t - rt
+            if time_in_cache > DIR_CACHE_TIME:
+                # use cached value if the last request took longer than the length of time we've been expired as it
+                # doesn't make sense to hit the server again
+                if strict_caching or (time_in_cache > (DIR_CACHE_TIME + dt)):
+                    raise RuntimeError('key is expired')
+                else:
+                    logger.warning('Using expired entry from directory cache on %s as previous request took too long' % dirurl)
+        except (KeyError, RuntimeError):
+            #logger.debug('dir cache miss')
+            # t = time.time()
+            url = dirurl.encode()
+            haveResult = False
+            nTries = 0
+            while nTries < nRetries and not haveResult:
+                try:
+                    nTries += 1
+                    s = self._getSession(url)
+                    r = s.get(url, timeout=timeout)
+                    haveResult = True
+                except (requests.Timeout, requests.ConnectionError) as e:
+                    # s.get sometimes raises ConnectionError instead of ReadTimeoutError
+                    # see https://github.com/requests/requests/issues/2392
+                    logger.exception('Timeout on listing directory')
+                    logger.info('%d retries left' % (nRetries - nTries))
+                    if nTries == nRetries:
+                        raise
+        
+            dt = time.time() - t
+            if not r.status_code == 200:
+                logger.debug('Request for %s failed with error: %d' % (url, r.status_code))
+            
+                #make sure we read a reply so that the far end doesn't hold the connection open
+                dump = r.content
+                return {}, dt
+            try:
+                dirL = r.json()
+            except ValueError:
+                # directory doesn't exist
+                # TODO - cache empty directories???
+                return {}, dt
+        
+            self._dirCache[dirurl] = (dirL, t, dt)
+    
+        return dirL, dt
+
+    def listdirectory(self, dirname, timeout=5):
+        """Lists the contents of a directory on the cluster.
+
+        Returns a dictionary mapping filenames to clusterListing.FileInfo named tuples.
+        """
+        from . import clusterListing as cl
+        dirlist = dict()
+        dirname = dirname.lstrip('/')
+
+        if not dirname.endswith('/'):
+            dirname = dirname + '/'
+  
+        urls = [s + dirname for s in self.dataservers]
+        
+        with self._list_dir_lock:
+            # TODO - is lock needed?
+            listings = self._pool.map(self.list_single_node_dir, urls)
+
+        for dirL, dt in listings:
+            cl.aggregate_dirlisting(dirlist, dirL)
+
+        return dirlist
+
+    def listdir(self, dirname):
+        """Lists the contents of a directory on the cluster. Similar to os.listdir,
+        but directories are indicated by a trailing slash
+        """
+    
+        return sorted(self.listdirectory(dirname).keys())
+
+    def locate_file(self, filename, return_first_hit=False):
+        """
+        Searches the cluster to find which server(s) a given file is stored on
+
+        Parameters
+        ----------
+        filename : str
+            The file name
+        return_first_hit : bool
+            Whether to try and find all locations, or return when we find the first copy
+
+        Returns
+        -------
+
+        """
+    
+        try:
+            locs, t = self._locateCache[filename]
+            return locs
+        except KeyError:
+            locs = []
+            
+            dirname, fn = posixpath.split(filename)
+            if not dirname.endswith('/'):
+                dirname += '/'
+            
+            for s in self.dataservers:
+                #note: servers are ordered so we try those on the local machine first
+                dirurl = s + dirname
+                try:
+                    dirL, rt, dt = self._dirCache[dirurl]
+                    if dirL[fn]: # this will raise a KeyError if filename is not present
+                        locs.append((dirurl + fn, dt))
+                except KeyError:
+                    dirL, dt = self.list_single_node_dir(dirurl)
+                
+                    if dirL.get(fn, None) is not None:
+                        locs.append((dirurl + fn, dt))
+            
+                if return_first_hit and len(locs) > 0:
+                    return locs
+        
+            if len(locs) > 0:
+                #cache if we found something (this is safe due to write-once nature of fs)
+                self._locateCache[filename] = (locs, time.time())
+        
+            return locs
+
+    def isdir(self, name):
+        """
+        Tests if a given path on the cluster is a directory. Analogous to os.path.isdir
+        """
+        import os
+        
+        name = name.rstrip('/')
+        if name in ['/', '']:
+            #special case for root dir
+            return True
+    
+        pn, n = posixpath.split(name)
+        try:
+            d = self.listdirectory(pn)[n + '/']
+        except KeyError:
+            return False
+    
+        return d.type > 0
+
+    def _cglob(self, url, timeout=2, nRetries=1):
+        url = url.encode()
+        haveResult = False
+        nTries = 0
+        while nTries < nRetries and not haveResult:
+            try:
+                nTries += 1
+                s = self._getSession(url)
+                r = s.get(url, timeout=timeout)
+                haveResult = True
+            except (requests.Timeout, requests.ConnectionError) as e:
+                # s.get sometimes raises ConnectionError instead of ReadTimeoutError
+                # see https://github.com/requests/requests/issues/2392
+                logger.exception('Timeout on listing directory')
+                logger.info('%d retries left' % (nRetries - nTries))
+                if nTries == nRetries:
+                    raise
+    
+        if not r.status_code == 200:
+            logger.debug('Request failed with error: %d' % r.status_code)
+        
+            #make sure we read a reply so that the far end doesn't hold the connection open
+            dump = r.content
+            return []
+    
+        try:
+            matches = r.json()
+        except ValueError:
+            # directory doesn't exist
+            return []
+    
+        return list(matches)
+
+    def cglob(self, pattern):
+        """
+        Find files matching a given glob on the cluster. Analogous to the python glob.glob function.
+
+        Parameters
+        ----------
+        pattern : string glob
+        serverfilter : cluster name (optional)
+
+        Returns
+        -------
+
+        a list of files matching the glob
+
+        """
+        urls = [s + '__glob?pattern=%s' % pattern for s in self.dataservers]
+        
+        with self._list_dir_lock:
+            matches = self._pool.map(self._cglob, urls)
+    
+        #concatenate lists
+        matches = sum(matches, [])
+        return list(set(matches))
+
+    def exists(self, name):
+        """
+        Test whether a file exists on the cluster. Analogue to os.path.exists for local files.
+
+        Parameters
+        ----------
+        name : string, file path
+        serverfilter : name of the cluster (optional)
+
+        Returns
+        -------
+
+        True if file exists, else False
+
+        """
+        return (len(self.locate_file(name, True)) > 0) or self.isdir(name)
+    
+    def register_file(self, filename, url, size):
+        """
+        Call after uploading a new file so we can update our caches
+
+        """
+        from . import clusterListing as cl
+        t = time.time()
+        self._locateCache[filename] = ([(url, .1), ], t)
+        try:
+            dirurl, fn = posixpath.split(url)
+            dirurl = dirurl + '/'
+            dirL, rt, dt = self._dirCache[dirurl]
+            if (t - rt) > DIR_CACHE_TIME:
+                pass #cache entry is expired
+            else:
+                dirL[fn] = cl.FileInfo(cl.FILETYPE_NORMAL, size)
+                self._dirCache[dirurl] = (dirL, rt, dt)
+    
+        except KeyError:
+            pass

PYME/IO/cluster_streaming.py

@@ -0,0 +1,216 @@
+from . import clusterIO
+import socket
+import queue
+import threading
+import logging
+
+logger = logging.getLogger(__name__)
+QUEUE_MAX_SIZE = 200 # ~10k frames
+
+class Stream(object):
+    """
+    Class to handle spooling files in an asynchronous / non blocking manner to a single server
+    
+    TODO - use this in the spoolers as well??? It's a little cleaner than the existing code.
+    """
+    def __init__(self, server_address, server_port, dir_manager=None, filter=None):
+        if not isinstance(server_address, str):
+            server_address = socket.inet_ntoa(server_address)
+            
+        self._server_address = server_address
+        self._server_port = server_port
+        
+        self._dir_manager = dir_manager
+        self._filter = filter
+        
+        self._url = 'http://%s:%d/' % (server_address, server_port)
+
+        self._put_queue = queue.Queue(QUEUE_MAX_SIZE)
+        self._rc_queue = queue.Queue()
+        #self._last_flush_time = time.time()
+        self._socket = None
+        self._alive = True
+        
+        self._connect()
+        
+        self._t_send = threading.Thread(target=self._send_loop)
+        self._t_send.daemon = True
+        self._t_send.start()
+
+        self._t_recv = threading.Thread(target=self._recv_loop)
+        self._t_recv.daemon = True
+        self._t_recv.start()
+        
+    def _connect(self):
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        self._socket.settimeout(30)
+        
+        self._socket.connect((self._server_address, self._server_port))
+    
+    def finalize(self, filename=None):
+        """
+        Stop our threads and tidy up
+
+        """
+        if filename:
+            #if we provide a sentinal file, put that
+            self.put(filename, None)
+            self._t_send.join(2)
+        
+        
+        #self._alive = False
+        
+    def close(self):
+        #TODO - wait on the put queue to empty?
+        self._alive = False
+        self._t_send.join()
+        self._t_recv.join()
+        
+        self._socket.close()
+        
+    def put(self, filename, data):
+        """Add a file to the queue to be put
+        
+        calling put with data == None is a sentinel that no more data will follow, and that keep-alive
+        should be switched off. This should normaly be accompanied by a filename = '__part_pyramid_finish/<pyramid_dir>'
+        
+        """
+        self._put_queue.put((filename, data))
+    
+    def _send_loop(self):
+        """
+        Loop which runs in it's own thread, pushing chunks as they come in. Because we now run sending and recieving
+        in separate threads, and keep the socket open for the entire duration, we don't need to batch puts.
+
+        """
+        connection = b'keep-alive'
+        
+        while (self._alive or not self._put_queue.empty()):
+            try:
+                filename, data = self._put_queue.get(timeout=.1)
+
+                if data is None:
+                    connection = b'close'
+                    data = b''
+                    
+                if self._filter is not None:
+                    # allow us to do, e.g. compression in the spooling thread
+                    data = self._filter(data)
+
+                dl = len(data)
+
+                header = b'PUT /%s HTTP/1.1\r\nConnection: %s\r\nContent-Length: %d\r\n\r\n' % (
+                    filename.encode(), connection, dl)
+                
+                logger.debug(header)
+                
+                self._socket.sendall(header)
+                if dl > 0:
+                    self._socket.sendall(data)
+                    
+                self._put_queue.task_done()
+                self._rc_queue.put((filename, dl))
+                #datalen += dl
+                #nChunksSpooled += 1
+                #nChunksRemaining -= 1
+                
+            except queue.Empty:
+                pass
+                
+            
+    def _recv_loop(self):
+        fp = self._socket.makefile('rb', 65536)
+        
+        try:
+            while self._alive:
+                try:
+                    filename, dl = self._rc_queue.get(timeout=.1)
+                    status, reason, msg = clusterIO._parse_response(fp)
+                    if not status == 200:
+                        logger.error(('Error spooling %s: Response %d - status %d' % (filename, status, reason)) + ' msg: ' + str(msg))
+                    else:
+                        if self._dir_manager:
+                            # register file in cluster directory
+                            fn = filename.split('?')[0]
+                            url = self.url_ + fn
+                            self._dir_manager.register_file(fn, url, dl)
+                            
+                except queue.Empty:
+                    pass
+                
+                #TODO - more detailed error handling - e.g. socket timeouts
+                
+        finally:
+            fp.close()
+        
+        
+def distribution_function_round_robin(n_servers, i=None):  
+    if i is None:
+        # distribute at random
+        import random
+        return random.randrange(n_servers)
+
+    return i % n_servers
+
+class Streamer(object):
+    """ Create a spooler instance which keeps one persistent connection to each node on the
+    cluster and allows non-blocking streaming of data to these nodes
+
+    """
+    def __init__(self, serverfilter=clusterIO.local_serverfilter, servers=None, distribution_fcn=distribution_function_round_robin, filter=None):
+        """
+        Parameters
+        ----------
+
+        serverfilter : string
+                The cluster identifier (when multiple clusters on one network)
+        servers : list
+                A manual list of servers. Usage is not recommended
+        distribution_fcn : callable
+                a function which assigns files to frames. Takes (at least) n_servers as an argument (provided by the spooler), but optionally any addition keyword
+                arguments you pass to `put()`. In simple cases, this will be a frame number - e.g. i - see round_robin fcn above. distribution_fcn
+                *must* provide defaults for all parameters other than n_servers so that sensible behaviour is achived when no value is provided.
+        filter: callable
+                a function which performs some operation on the data before it's saved. Typically format conversion and/or compression.
+        """
+        self._directory = clusterIO.get_dir_manager(serverfilter)
+
+        if servers is None:
+            # FIXME - avoid private _ns usage
+            self.servers = [(socket.inet_ntoa(v.address), v.port) for k, v in self._directory._ns.get_advertised_services()]
+        else:
+            self.servers = servers
+
+        assert len(self.servers) > 0, "No servers found for distribution. Make sure that cluster servers are running and can be reached from this device."
+        
+        self._n_servers = len(self.servers)
+        self._streams =  [Stream(address, port, filter=filter) for address, port in self.servers]
+
+        self._distribution_fcn = distribution_fcn
+
+    #@classmethod
+   
+    
+    def put(self, filename, data, **kwargs):
+        """ Put, choosing a stream using the distribution function
+        
+        kwargs are used as arguments for the server distribution function. 
+
+        """
+        idx = self._distribution_fcn(n_servers = self._n_servers, **kwargs)
+
+        self.put_stream(idx, filename, data)
+
+    def put_stream(self, idx, filename, data):
+        """ Put to a specific stream """
+        self._streams[idx].put(filename, data)
+
+    def close(self):
+        for s in self._streams:
+            s.close()
+
+
+def n_cluster_servers(serverfilter):
+    """ Convenience function for hard-coding distribution functions in, e.g. the clusterh5  case"""
+    return len(clusterIO.get_dir_manager(serverfilter).dataservers)

PYME/IO/compatibility.py

@@ -0,0 +1,12 @@
+import numpy
+import pickle
+
+# this should serve as a drop in replacement for numpy.load calls that need
+# some backwards compatibility
+def np_load_legacy(f):
+    try:
+        ret = numpy.load(f, allow_pickle=True)
+    except (OSError, pickle.UnpicklingError):
+        # this one based on tips from https://rebeccabilbro.github.io/convert-py2-pickles-to-py3/
+        ret = numpy.load(f, allow_pickle=True, encoding="latin1")
+    return ret

PYME/IO/countdir/__init__.py

@@ -0,0 +1,5 @@
+from PYME import config
+if config.get('cluster-listing-no-countdir', False):
+    raise ImportError('Not using countdir')
+
+from .countdir import *