Aluode/PerceptionLabPortable

1import dask2from distributed.client import Client, _get_global_client3from distributed.worker import Worker4 5from fsspec import filesystem6from fsspec.spec import AbstractBufferedFile, AbstractFileSystem7from fsspec.utils import infer_storage_options8 9 10def _get_client(client):11    if client is None:12        return _get_global_client()13    elif isinstance(client, Client):14        return client15    else:16        # e.g., connection string17        return Client(client)18 19 20def _in_worker():21    return bool(Worker._instances)22 23 24class DaskWorkerFileSystem(AbstractFileSystem):25    """View files accessible to a worker as any other remote file-system26 27    When instances are run on the worker, uses the real filesystem. When28    run on the client, they call the worker to provide information or data.29 30    **Warning** this implementation is experimental, and read-only for now.31    """32 33    def __init__(34        self, target_protocol=None, target_options=None, fs=None, client=None, **kwargs35    ):36        super().__init__(**kwargs)37        if not (fs is None) ^ (target_protocol is None):38            raise ValueError(39                "Please provide one of filesystem instance (fs) or"40                " target_protocol, not both"41            )42        self.target_protocol = target_protocol43        self.target_options = target_options44        self.worker = None45        self.client = client46        self.fs = fs47        self._determine_worker()48 49    @staticmethod50    def _get_kwargs_from_urls(path):51        so = infer_storage_options(path)52        if "host" in so and "port" in so:53            return {"client": f"{so['host']}:{so['port']}"}54        else:55            return {}56 57    def _determine_worker(self):58        if _in_worker():59            self.worker = True60            if self.fs is None:61                self.fs = filesystem(62                    self.target_protocol, **(self.target_options or {})63                )64        else:65            self.worker = False66            self.client = _get_client(self.client)67            self.rfs = dask.delayed(self)68 69    def mkdir(self, *args, **kwargs):70        if self.worker:71            self.fs.mkdir(*args, **kwargs)72        else:73            self.rfs.mkdir(*args, **kwargs).compute()74 75    def rm(self, *args, **kwargs):76        if self.worker:77            self.fs.rm(*args, **kwargs)78        else:79            self.rfs.rm(*args, **kwargs).compute()80 81    def copy(self, *args, **kwargs):82        if self.worker:83            self.fs.copy(*args, **kwargs)84        else:85            self.rfs.copy(*args, **kwargs).compute()86 87    def mv(self, *args, **kwargs):88        if self.worker:89            self.fs.mv(*args, **kwargs)90        else:91            self.rfs.mv(*args, **kwargs).compute()92 93    def ls(self, *args, **kwargs):94        if self.worker:95            return self.fs.ls(*args, **kwargs)96        else:97            return self.rfs.ls(*args, **kwargs).compute()98 99    def _open(100        self,101        path,102        mode="rb",103        block_size=None,104        autocommit=True,105        cache_options=None,106        **kwargs,107    ):108        if self.worker:109            return self.fs._open(110                path,111                mode=mode,112                block_size=block_size,113                autocommit=autocommit,114                cache_options=cache_options,115                **kwargs,116            )117        else:118            return DaskFile(119                fs=self,120                path=path,121                mode=mode,122                block_size=block_size,123                autocommit=autocommit,124                cache_options=cache_options,125                **kwargs,126            )127 128    def fetch_range(self, path, mode, start, end):129        if self.worker:130            with self._open(path, mode) as f:131                f.seek(start)132                return f.read(end - start)133        else:134            return self.rfs.fetch_range(path, mode, start, end).compute()135 136 137class DaskFile(AbstractBufferedFile):138    def __init__(self, mode="rb", **kwargs):139        if mode != "rb":140            raise ValueError('Remote dask files can only be opened in "rb" mode')141        super().__init__(**kwargs)142 143    def _upload_chunk(self, final=False):144        pass145 146    def _initiate_upload(self):147        """Create remote file/upload"""148        pass149 150    def _fetch_range(self, start, end):151        """Get the specified set of bytes from remote"""152        return self.fs.fetch_range(self.path, self.mode, start, end)153