import dask from distributed.client import Client, _get_global_client from distributed.worker import Worker from fsspec import filesystem from fsspec.spec import AbstractBufferedFile, AbstractFileSystem from fsspec.utils import infer_storage_options def _get_client(client): if client is None: return _get_global_client() elif isinstance(client, Client): return client else: # e.g., connection string return Client(client) def _in_worker(): return bool(Worker._instances) class DaskWorkerFileSystem(AbstractFileSystem): """View files accessible to a worker as any other remote file-system When instances are run on the worker, uses the real filesystem. When run on the client, they call the worker to provide information or data. **Warning** this implementation is experimental, and read-only for now. """ def __init__( self, target_protocol=None, target_options=None, fs=None, client=None, **kwargs ): super().__init__(**kwargs) if not (fs is None) ^ (target_protocol is None): raise ValueError( "Please provide one of filesystem instance (fs) or" " target_protocol, not both" ) self.target_protocol = target_protocol self.target_options = target_options self.worker = None self.client = client self.fs = fs self._determine_worker() @staticmethod def _get_kwargs_from_urls(path): so = infer_storage_options(path) if "host" in so and "port" in so: return {"client": f"{so['host']}:{so['port']}"} else: return {} def _determine_worker(self): if _in_worker(): self.worker = True if self.fs is None: self.fs = filesystem( self.target_protocol, **(self.target_options or {}) ) else: self.worker = False self.client = _get_client(self.client) self.rfs = dask.delayed(self) def mkdir(self, *args, **kwargs): if self.worker: self.fs.mkdir(*args, **kwargs) else: self.rfs.mkdir(*args, **kwargs).compute() def rm(self, *args, **kwargs): if self.worker: self.fs.rm(*args, **kwargs) else: self.rfs.rm(*args, **kwargs).compute() def copy(self, *args, **kwargs): if self.worker: self.fs.copy(*args, **kwargs) else: self.rfs.copy(*args, **kwargs).compute() def mv(self, *args, **kwargs): if self.worker: self.fs.mv(*args, **kwargs) else: self.rfs.mv(*args, **kwargs).compute() def ls(self, *args, **kwargs): if self.worker: return self.fs.ls(*args, **kwargs) else: return self.rfs.ls(*args, **kwargs).compute() def _open( self, path, mode="rb", block_size=None, autocommit=True, cache_options=None, **kwargs, ): if self.worker: return self.fs._open( path, mode=mode, block_size=block_size, autocommit=autocommit, cache_options=cache_options, **kwargs, ) else: return DaskFile( fs=self, path=path, mode=mode, block_size=block_size, autocommit=autocommit, cache_options=cache_options, **kwargs, ) def fetch_range(self, path, mode, start, end): if self.worker: with self._open(path, mode) as f: f.seek(start) return f.read(end - start) else: return self.rfs.fetch_range(path, mode, start, end).compute() class DaskFile(AbstractBufferedFile): def __init__(self, mode="rb", **kwargs): if mode != "rb": raise ValueError('Remote dask files can only be opened in "rb" mode') super().__init__(**kwargs) def _upload_chunk(self, final=False): pass def _initiate_upload(self): """Create remote file/upload""" pass def _fetch_range(self, start, end): """Get the specified set of bytes from remote""" return self.fs.fetch_range(self.path, self.mode, start, end)