ParallelTable Musings

daskms/parallel_table.py

@@ -0,0 +1,95 @@
+import logging
+import threading
+from pyrap.tables import table as Table
+from weakref import finalize
+from pathlib import Path
+
+
+log = logging.getLogger(__name__)
+
+
+def _parallel_table_finalizer(_cached_tables):
+
+    for table in _cached_tables.values():
+        link_path = Path(table.name())
+        table.close()
+        link_path.unlink()
+
+
+class ParallelTable(Table):
+
+    @classmethod
+    def _map_create_parallel_table(cls, args, kwargs):
+        """ Support pickling of kwargs in ParallelTable.__reduce__ """
+        return cls(*args, **kwargs)
+
+    def __init__(self, *args, **kwargs):
+
+        self._args = args
+        self._kwargs = kwargs
+
+        self._cached_tables = {}
+        self._table_path = args[0]  # TODO: This should be checked.
+
+        super().__init__(*args, **kwargs)
+
+        finalize(self, _parallel_table_finalizer, self._cached_tables)
+
+    def __reduce__(self):
+        """ Defer to _map_create_parallel_table to support kwarg pickling """
+        return (
+            ParallelTable._map_create_parallel_table,
+            (ParallelTable, self._args, self._kwargs)
+        )
+
+    def getcol(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcol(*args, **kwargs)
+
+    def getcolslice(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcolslice(*args, **kwargs)
+
+    def getcolnp(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcolnp(*args, **kwargs)
+
+    def getvarcol(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getvarcol(*args, **kwargs)
+
+    def getcell(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcell(*args, **kwargs)
+
+    def getcellslice(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcellslice(*args, **kwargs)
+
+    def getkeywords(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getkeywords(*args, **kwargs)
+
+    def getcolkeywords(self, *args, **kwargs):
+        table = self._get_table(threading.get_ident())
+        return table.getcolkeywords(*args, **kwargs)
+
+    def _get_table(self, thread_id):
+
+        try:
+            table = self._cached_tables[thread_id]
+        except KeyError:
+            print(f"opening for {thread_id}")
+
+            table_path = Path(self._table_path)
+            table_name = table_path.name
+            link_path = Path(f"/tmp/{thread_id}-{table_name}")
+
+            link_path.symlink_to(table_path)
+
+            self._cached_tables[thread_id] = table = Table(
+                str(link_path),
+                **self._kwargs
+            )
+
+        return table

daskms/reads.py

@@ -17,6 +17,7 @@
 from daskms.dataset import Dataset
 from daskms.table_executor import executor_key
 from daskms.table import table_exists
+from daskms.parallel_table import ParallelTable
 from daskms.table_proxy import TableProxy, READLOCK
 from daskms.table_schemas import lookup_table_schema
 from daskms.utils import table_path_split
@@ -305,7 +306,7 @@ def __init__(self, table, select_cols, group_cols, index_cols, **kwargs):
             raise ValueError(f"Unhandled kwargs: {kwargs}")
 
     def _table_proxy_factory(self):
-        return TableProxy(pt.table, self.table_path, ack=False,
+        return TableProxy(ParallelTable, self.table_path, ack=False,
                           readonly=True, lockoptions='user',
                           __executor_key__=executor_key(self.canonical_name))
 

daskms/table_executor.py

@@ -38,7 +38,7 @@ def __call__(cls, key=STANDARD_EXECUTOR):
 class Executor(object, metaclass=ExecutorMetaClass):
     def __init__(self, key=STANDARD_EXECUTOR):
         # Initialise a single thread
-        self.impl = impl = cf.ThreadPoolExecutor(1)
+        self.impl = impl = DummyThreadPoolExecutor(1)
         self.key = key
 
         # Register a finaliser shutting down the
@@ -54,6 +54,30 @@ def __repr__(self):
     __str__ = __repr__
 
 
+class DummyThreadPoolExecutor(object):
+
+    def __init__(self, nthread):
+        pass
+
+    def submit(self, fn, *args, **kwargs):
+
+        return DummyFuture(fn(*args, **kwargs))
+
+    def shutdown(self, wait=True):
+        pass
+
+
+class DummyFuture(object):
+
+    def __init__(self, value):
+
+        self.value = value
+
+    def result(self):
+
+        return self.value
+
+
 def executor_key(table_name):
     """
     Product an executor key from table_name

daskms/tests/test_table_proxy.py

@@ -222,3 +222,36 @@ def _getcol(tp, column):
 
     del futures, data, u, tab_fut
     assert_liveness(0, 0)
+
+
+@pytest.mark.parametrize("scheduler",
+                         ["sync", "threads", "processes", "distributed"])
+def test_softlinks(ms, scheduler):
+
+    import dask.array as da
+    from dask.distributed import Client, LocalCluster
+    from daskms import xds_from_ms
+
+    if scheduler == "distributed":
+
+        cluster = LocalCluster(
+            processes=True,
+            n_workers=2,
+            threads_per_worker=2,
+            memory_limit=0
+        )
+
+        client = Client(cluster)  # noqa
+
+    # ms = "/home/jonathan/reductions/3C147/msdir/C147_unflagged.MS"
+    ms = "/home/jonathan/recipes/proxy_experiments/vla_empty.ms"
+
+    xdsl = xds_from_ms(
+        ms,
+        group_cols=["DATA_DESC_ID", "SCAN_NUMBER", "FIELD_ID"],
+        chunks={"row": 30000, "chan": -1, "corr": -1}
+    )
+
+    result = [xds.DATA.data.map_blocks(lambda x: x[:1, :1, :1]) for xds in xdsl]
+
+    da.compute(result, scheduler=scheduler, num_workers=4)