make memory reporting optional

pfb/__init__.py

@@ -30,7 +30,7 @@ def set_envs(nthreads, ncpu):
     os.environ["NUMEXPR_NUM_THREADS"] = str(ne_threads)
 
 
-def set_client(nworkers, stack, log,
+def set_client(nworkers, log, stack=None,
                host_address=None, direct_to_workers=False):
     import dask
     dask.config.set({'distributed.comm.compression': 'lz4'})
@@ -54,10 +54,12 @@ def set_client(nworkers, stack, log,
                                threads_per_worker=1,
                                memory_limit=0,  # str(mem_limit/nworkers)+'GB'
                                asynchronous=False)
-        cluster = stack.enter_context(cluster)
+        if stack is not None:
+            cluster = stack.enter_context(cluster)
         client = Client(cluster,
                         direct_to_workers=direct_to_workers)
-        client = stack.enter_context(client)
+        if stack is not None:
+            client = stack.enter_context(client)
 
     client.wait_for_workers(nworkers)
     dashboard_url = client.dashboard_link

pfb/parser/hci.yaml

@@ -163,6 +163,11 @@ inputs:
     default: false
     info:
       Compute naural gradient
+  memory_reporting:
+    dtype: bool
+    default: false
+    info:
+      Report worker memory as tasks complete
 
   _include:
     - (.)gridding.yml

pfb/parser/init.yaml

@@ -121,6 +121,12 @@ inputs:
     info:
       Display progress.
       Use --no-progressbar to deactivate.
+  memory_reporting:
+    dtype: bool
+    default: false
+    info:
+      Report worker memory as tasks complete
+
   _include:
     - (.)out.yml
     - (.)dist.yml

pfb/workers/degrid.py

@@ -76,7 +76,7 @@ def degrid(**kw):
 
     with ExitStack() as stack:
         from pfb import set_client
-        client = set_client(opts.nworkers, stack, log)
+        client = set_client(opts.nworkers, log, stack)
 
         ti = time.time()
         _degrid(**opts)

pfb/workers/fluxmop.py

@@ -69,7 +69,7 @@ def fluxmop(**kw):
         from distributed import wait
         from pfb import set_client
         if opts.nworkers > 1:
-            client = set_client(opts.nworkers, stack, log)
+            client = set_client(opts.nworkers, log, stack)
             from distributed import as_completed
         else:
             print("Faking client", file=log)

pfb/workers/grid.py

@@ -81,7 +81,7 @@ def grid(**kw):
     with ExitStack() as stack:
         from pfb import set_client
         if opts.nworkers > 1:
-            client = set_client(opts.nworkers, stack, log)
+            client = set_client(opts.nworkers, log, stack)
             from distributed import as_completed
         else:
             print("Faking client", file=log)

pfb/workers/hci.py

@@ -87,7 +87,7 @@ def hci(**kw):
         dask.config.set(**{'array.slicing.split_large_chunks': False})
         from pfb import set_client
         from distributed import wait, get_client
-        client = set_client(opts.nworkers, stack, log)
+        client = set_client(opts.nworkers, log, stack)
 
         ti = time.time()
         _hci(**opts)
@@ -415,6 +415,11 @@ def _hci(**kw):
         associated_workers[future] = worker
         n_launched += 1
 
+        if opts.memory_reporting:
+            worker_info = client.scheduler_info()['workers']
+            print(f'Total memory {worker} MB = ',
+                  worker_info[worker]['metrics']['memory']/1e6, file=log)
+
         if opts.progressbar:
             print(f"\rProcessing: {n_launched}/{nds}", end='', flush=True)
     print("\n")  # after progressbar above

pfb/workers/init.py

@@ -82,17 +82,19 @@ def init(**kw):
     resize_thread_pool(opts.nthreads)
     set_envs(opts.nthreads, ncpu)
 
-    with ExitStack() as stack:
-        import dask
-        dask.config.set(**{'array.slicing.split_large_chunks': False})
-        from pfb import set_client
-        from distributed import wait, get_client
-        client = set_client(opts.nworkers, stack, log)
+    # with ExitStack() as stack:
+    import dask
+    dask.config.set(**{'array.slicing.split_large_chunks': False})
+    from pfb import set_client
+    from distributed import wait, get_client
+    client = set_client(opts.nworkers, log)
+
+    ti = time.time()
+    _init(**opts)
 
-        ti = time.time()
-        _init(**opts)
+    print(f"All done after {time.time() - ti}s", file=log)
 
-        print(f"All done after {time.time() - ti}s", file=log)
+    client.close()
 
 def _init(**kw):
     opts = OmegaConf.create(kw)
@@ -376,8 +378,10 @@ def _init(**kw):
         associated_workers[future] = worker
         n_launched += 1
 
-        worker_info = client.scheduler_info()['workers']
-        print(f'Total memory {worker} MB = ', worker_info[worker]['metrics']['memory'])
+        if opts.memory_reporting:
+            worker_info = client.scheduler_info()['workers']
+            print(f'Total memory {worker} MB = ',
+                  worker_info[worker]['metrics']['memory']/1e6, file=log)
 
         if opts.progressbar:
             print(f"\rProcessing: {n_launched}/{nds}", end='', flush=True)

pfb/workers/restore.py

@@ -55,7 +55,7 @@ def restore(**kw):
     with ExitStack() as stack:
         if opts.nworkers > 1:
             from pfb import set_client
-            client = set_client(opts.nworkers, stack, log)
+            client = set_client(opts.nworkers, log, stack)
         else:
             print("Faking client", file=log)
             from pfb.utils.dist import fake_client

pfb/workers/spotless.py

@@ -72,7 +72,7 @@ def spotless(**kw):
     with ExitStack() as stack:
         from pfb import set_client
         if opts.nworkers > 1:
-            client = set_client(opts.nworkers, stack, log,
+            client = set_client(opts.nworkers, log, stack,
                                 direct_to_workers=opts.direct_to_workers)
             from distributed import as_completed
         else: