Initial commit

hf3fs/__init__.py

@@ -0,0 +1,310 @@
+from hf3fs_py_usrbio import Client, iovec
+
+#import attrs
+from contextlib import contextmanager, AbstractContextManager
+from dataclasses import dataclass
+import functools
+import os
+from pathlib import PurePosixPath
+import threading as th
+
+from pkg_resources import get_distribution
+
+try:
+    __version__ = get_distribution('hf3fs').version
+except:
+    __version__ = "debug"
+
+DEFAULT_CLIENT = th.local()
+DEFAULT_CLIENT.client = None
+DEFAULT_CLIENT.clients = {}
+
+#@attrs.define
+@dataclass
+class MountInfo:
+    token: str
+    as_super: bool
+
+MOUNT_INFO = {}
+
+def _getDefaultClient(kwargs):
+    if 'client' in kwargs and kwargs['client'] is not None:
+        client = kwargs['client']
+        del kwargs['client']
+    elif 'mount_name' in kwargs and kwargs['mount_name'] is not None:
+        mount_name = kwargs['mount_name']
+        if mount_name not in DEFAULT_CLIENT.clients:
+            setupDefaultClient(mount_name)
+        client = DEFAULT_CLIENT.clients[mount_name]
+        del kwargs['mount_name']
+    elif DEFAULT_CLIENT.client is None:
+        raise RuntimeError("default client not setup")
+    else:
+        client = DEFAULT_CLIENT.client
+
+    return client, kwargs
+
+def _setupH3Method(name):
+    @functools.wraps(getattr(Client, name))
+    def wrapper(*args, **kwargs):
+        nonlocal name
+        client, kwargs = _getDefaultClient(kwargs)
+        return getattr(client, name)(*args, **kwargs)
+
+    globals()[name] = wrapper
+
+for _name in ['stat', 'fstat', 'mkdir', 'rmdir', 'unlink', 'remove', 'realpath', 'readlink', 'opendir', 'readdir',
+              'creat', 'symlink', 'link', 'open', 'close', 'chmod', 'chown', 'chdir', 'ftruncate',
+              'iovalloc', 'iovfree', 'preadv', 'pwritev']:
+    _setupH3Method(_name)
+
+def setMountInfo(mount_name, token, as_super=False):
+    global MOUNT_INFO
+    MOUNT_INFO[mount_name] = MountInfo(token, as_super)
+
+def setupDefaultClient(mount_name):
+    global MOUNT_INFO
+    if mount_name not in MOUNT_INFO:
+        raise ValueError(f"unknown mount name '{mount_name}'")
+
+    mount = MOUNT_INFO[mount_name]
+    
+    global DEFAULT_CLIENT
+    client = DEFAULT_CLIENT.clients[mount_name] = Client(mount_name, mount.token, as_super=mount.as_super)
+    return client
+
+@contextmanager
+def defaultClient(mount_name, token, as_super=False):
+    global DEFAULT_CLIENT
+    lastClient = DEFAULT_CLIENT.client
+    DEFAULT_CLIENT.client = Client(mount_name, token, as_super=as_super)
+
+    try:
+        yield DEFAULT_CLIENT.client
+    finally:
+        DEFAULT_CLIENT.client = lastClient
+
+def withClient(f):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        nonlocal f
+        client, kwargs = _getDefaultClient(kwargs)
+        return f(client=client, *args, **kwargs)
+
+    return wrapper
+
+@withClient
+def listdir(path='.', client=None):
+    if path is None:
+        path = '.'
+
+    dirp = client.opendir(path)
+#    print(dirp, flush=True)
+    fileList = []
+    while True:
+        dent = client.readdir(dirp)
+#        print('dent', dent, flush=True)
+        if dent is None:
+            break
+        fileList.append(dent.d_name)
+
+    return fileList
+
+class DirEntry(object):
+    @withClient
+    def __init__(self, parentPath, name, etype, parentFd, client=None):
+        self._parentPath = PurePosixPath(parentPath)
+        self._name = name
+        self._etype = etype
+        self._parentFd = parentFd
+        self._client = client
+        self._st = None
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def path(self):
+        return str(self._parentPath / self._name)
+
+    _S_IFMT = 0o170000
+    _S_IFLNK = 0o120000
+    _S_IFREG = 0o100000
+    _S_IFDIR = 0o040000
+
+    _DT_DIR = 4
+    _DT_REG = 8
+    _DT_LNK = 10
+
+    def _checkWFollow(self, against, follow_symlinks):
+        if self._etype == against[0]:
+            return True
+        elif follow_symlinks and self.is_symlink():
+            st = self.stat(True)
+            return (st.st_mode & self._S_IFMT) == against[1]
+        else:
+            return False
+
+    def is_dir(self, follow_symlinks=True):
+        return self._checkWFollow((self._DT_DIR, self._S_IFDIR), follow_symlinks)
+
+    def is_file(self, follow_symlinks=True):
+        return self._checkWFollow((self._DT_REG, self._S_IFREG), follow_symlinks)
+
+    def is_symlink(self):
+        return self._etype == self._DT_LNK
+
+    def stat(self, follow_symlinks=True):
+        if self._st is None:
+            self._st = self._client.stat(self._name, dir_fd=self._parentFd, follow_symlinks=follow_symlinks)
+        return self._st
+
+@withClient
+def scandir(path='.', dir_fd=None, client=None):
+    class DirEntryIter(AbstractContextManager):
+        def __init__(self, path, client, dir_fd):
+            self._path = path
+            self._client = client
+            self._dir_fd = dir_fd
+            self._fd = None
+            self._dirp = None
+
+        @property
+        def dir_fd(self):
+            return self._fd
+
+        def close(self):
+            if self._fd is not None:
+                self._client.close(self._fd)
+                self._fd = None
+                
+            self._dirp = None
+
+        def __iter__(self):
+            self._fd = self._client.open(self._path, os.O_DIRECTORY | os.O_PATH, dir_fd=self._dir_fd)
+#            print('dirfd to scan', self._fd, flush=True)
+            self._dirp = self._client.opendir(self._path, dir_fd=self._dir_fd)
+            return self
+
+        def __next__(self):
+            dent = self._client.readdir(self._dirp)
+            if dent is None:
+                raise StopIteration()
+            return DirEntry(self._path, dent.d_name, dent.d_type, self._fd, client=self._client)
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.close()
+            return False
+    
+    if path is None:
+        path = '.'
+
+    return DirEntryIter(path, client, dir_fd)
+
+@withClient
+def walk2(top, topdown=True, onerror=None, followlinks=False, dir_fd=None, curr_dir=None, client=None):
+    try:
+        with scandir(curr_dir if curr_dir is not None else top, dir_fd, client=client) as sd:
+            dirnames = []
+            filenames = []
+            for dent in sd:
+                name = dent.name
+                if dent.is_dir(followlinks):
+                    dirnames.append(name)
+                    if not topdown:
+                        yield from walk2(dent.path, False, onerror, followlinks, sd.dir_fd, name, client=client)
+                else:
+                    filenames.append(name)
+
+            yield (top, dirnames, filenames, sd.dir_fd)
+            if topdown:
+                topp = PurePosixPath(top)
+                for dirname in dirnames:
+                    yield from walk2(str(topp / dirname), True, onerror, followlinks, sd.dir_fd, dirname, client=client)
+    except OSError as e:
+        if onerror is not None:
+            onerror(e)
+
+@withClient
+def walk(top, topdown=True, onerror=None, followlinks=False, dir_fd=None, client=None):
+    for dp, dns, fns, dfd in walk2(top, topdown, onerror, followlinks, dir_fd, None, client=client):
+        yield (dp, dns, fns)
+
+class BinaryFile(AbstractContextManager):
+    @withClient
+    def __init__(self, path, mode, dir_fd=None, client=None, ignore_cache=False):
+        self.client = client
+
+        if mode == 'r':
+            flags = os.O_RDONLY
+        elif mode == 'r+':
+            flags = os.O_RDWR
+        elif mode == 'r+c':
+            flags = os.O_RDWR | os.O_CREAT
+        elif mode == 'w':
+            flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
+        elif mode == 'w+':
+            flags = os.O_RDWR | os.O_CREAT | os.O_TRUNC
+        else:
+            raise ValueError(f'invalid mode {mode}')
+
+        if ignore_cache:
+            flags |= os.O_NONBLOCK
+
+        self._fd = None
+        self._off = 0
+        
+        self._fd = client.open(path, flags, 0o644, dir_fd=dir_fd)
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def close(self):
+#        print('to close fd', self._fd, flush=True)
+        if self._fd is not None:
+            self.client.close(self._fd)
+            self._fd = None
+
+    def fileno(self):
+        return self._fd
+
+    def seek(self, pos, how=os.SEEK_SET, readahead=None):
+        self._off = self.client.lseek(self._fd, pos, how, readahead=readahead)
+        return self._off
+
+    def tell(self):
+        return self._off
+
+    def _bytesLeft(self):
+        off = self._off
+        flen = self.seek(0, os.SEEK_END)
+        self.seek(off)
+        return flen - off
+
+    def read(self, size=None, readahead=None):
+        if size is None:
+            size = self._bytesLeft()
+
+        buf = memoryview(bytearray(size))
+        red = self.readinto(buf, readahead=readahead)
+        return buf[:red]
+
+    def readinto(self, buf, readahead=None):
+        red = self.client.read(self._fd, buf, readahead=readahead)
+        self._off += red
+        return red
+
+    def write(self, buf, flush=False):
+        writ = self.client.write(self._fd, buf, flush=flush)
+        self._off += writ
+        return writ

hf3fs/fuse.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import PosixPath
+
+from hf3fs_py_usrbio import HF3FS_SUPER_MAGIC
+
+HF3FS_IOC_GET_MOUNT_NAME = 2149607424
+HF3FS_IOC_GET_PATH_OFFSET = 2147772417
+HF3FS_IOC_GET_MAGIC_NUM = 2147772418
+
+HF3FS_IOC_RECURSIVE_RM = 2147772426
+
+def serverPath(p):
+    '''
+    从完整路径获取 client 接受的路径名
+
+    Args:
+        p: 待解析的路径名
+
+    Examples:
+
+    .. code-block:: python
+
+        import hf3fs.fuse
+        hf3fs.fuse.serverPath('/hf3fs-cluster/aaa/../cpu/abc/def')
+    '''
+    np = os.path.normpath(os.path.realpath(p))
+    return os.path.join('/', *PosixPath(np).parts[3:])
+
+def mountName(p):
+    '''
+    从完整路径获取 mount name
+
+    Args:
+        p: 待解析的路径名
+
+    Examples:
+    
+    .. code-block:: python
+    
+        import hf3fs.fuse
+        hf3fs.fuse.mountName('/hf3fs-cluster/aaa/../cpu/abc/def')
+    '''
+    np = os.path.normpath(os.path.realpath(p))
+    return PosixPath(np).parts[2]