Remove obsolete jobs arg and accept 'processes' arg when handling distribution args (#73)

* remove obsolete jobs arg and also accept 'processes' arg when handling distribution args

* fix formatting

* fix jobs arg handling

Author: philippotto

wkcuber/__main__.py

@@ -83,12 +83,11 @@ def create_parser():
         args.layer_name,
         args.dtype,
         args.batch_size,
-        args.jobs,
         args,
     )
 
     if not args.no_compress:
-        compress_mag_inplace(args.target_path, args.layer_name, Mag(1), args.jobs, args)
+        compress_mag_inplace(args.target_path, args.layer_name, Mag(1), args)
 
     downsample_mags(
         args.target_path,
@@ -97,7 +96,6 @@ def create_parser():
         Mag(args.max_mag),
         "default",
         DEFAULT_EDGE_LEN,
-        args.jobs,
         not args.no_compress,
     )
 

wkcuber/compress.py

@@ -72,7 +72,7 @@ def compress_file_job(source_path, target_path):
         raise exc
 
 
-def compress_mag(source_path, layer_name, target_path, mag: Mag, jobs, args=None):
+def compress_mag(source_path, layer_name, target_path, mag: Mag, args=None):
     if path.exists(path.join(target_path, layer_name, str(mag))):
         logging.error("Target path '{}' already exists".format(target_path))
         exit(1)
@@ -98,9 +98,9 @@ def compress_mag(source_path, layer_name, target_path, mag: Mag, jobs, args=None
     logging.info("Mag {0} successfully compressed".format(str(mag)))
 
 
-def compress_mag_inplace(target_path, layer_name, mag: Mag, jobs, args=None):
+def compress_mag_inplace(target_path, layer_name, mag: Mag, args=None):
     compress_target_path = "{}.compress-{}".format(target_path, uuid4())
-    compress_mag(target_path, layer_name, compress_target_path, mag, jobs, args)
+    compress_mag(target_path, layer_name, compress_target_path, mag, args)
 
     shutil.rmtree(path.join(target_path, layer_name, str(mag)))
     shutil.move(
@@ -111,7 +111,7 @@ def compress_mag_inplace(target_path, layer_name, mag: Mag, jobs, args=None):
 
 
 def compress_mags(
-    source_path, layer_name, target_path=None, mags: List[Mag] = None, jobs=1, args=None
+    source_path, layer_name, target_path=None, mags: List[Mag] = None, args=None
 ):
     with_tmp_dir = target_path is None
     target_path = source_path + ".tmp" if with_tmp_dir else target_path
@@ -121,7 +121,7 @@ def compress_mags(
     mags.sort()
 
     for mag in mags:
-        compress_mag(source_path, layer_name, target_path, mag, jobs, args)
+        compress_mag(source_path, layer_name, target_path, mag, args)
 
     if with_tmp_dir:
         makedirs(path.join(source_path + ".bak", layer_name), exist_ok=True)
@@ -146,11 +146,4 @@ def compress_mags(
     args = create_parser().parse_args()
     if args.verbose:
         logging.basicConfig(level=logging.DEBUG)
-    compress_mags(
-        args.source_path,
-        args.layer_name,
-        args.target_path,
-        args.mag,
-        int(args.jobs),
-        args,
-    )
+    compress_mags(args.source_path, args.layer_name, args.target_path, args.mag, args)

wkcuber/cubing.py

@@ -144,9 +144,7 @@ def cubing_job(
                 raise exc
 
 
-def cubing(
-    source_path, target_path, layer_name, dtype, batch_size, jobs, args=None
-) -> dict:
+def cubing(source_path, target_path, layer_name, dtype, batch_size, args=None) -> dict:
 
     target_wkw_info = WkwDatasetInfo(target_path, layer_name, dtype, 1)
     source_files = find_source_filenames(source_path)
@@ -198,6 +196,5 @@ def cubing(
         args.layer_name,
         args.dtype,
         args.batch_size,
-        args.jobs,
         args=args,
     )

wkcuber/downsampling.py

@@ -114,7 +114,6 @@ def downsample(
     target_mag: Mag,
     interpolation_mode,
     cube_edge_len,
-    jobs,
     compress,
     args=None,
 ):
@@ -401,7 +400,6 @@ def downsample_mag(
     target_mag: Mag,
     interpolation_mode="default",
     cube_edge_len=DEFAULT_EDGE_LEN,
-    jobs=1,
     compress=False,
     args=None,
 ):
@@ -426,7 +424,6 @@ def downsample_mag(
         target_mag,
         interpolation_mode,
         cube_edge_len,
-        jobs,
         compress,
         args,
     )
@@ -439,7 +436,6 @@ def downsample_mags(
     max_mag: Mag,
     interpolation_mode,
     cube_edge_len,
-    jobs,
     compress,
     args=None,
 ):
@@ -453,7 +449,6 @@ def downsample_mags(
             target_mag,
             interpolation_mode,
             cube_edge_len,
-            jobs,
             compress,
             args,
         )
@@ -478,7 +473,6 @@ def downsample_mags(
             anisotropic_target_mag,
             args.interpolation_mode,
             args.buffer_cube_size,
-            args.jobs,
             args.compress,
             args,
         )
@@ -490,7 +484,6 @@ def downsample_mags(
             max_mag,
             args.interpolation_mode,
             args.buffer_cube_size,
-            args.jobs,
             args.compress,
             args,
         )

wkcuber/utils.py

@@ -112,8 +112,17 @@ def get_executor_for_args(args):
     executor = None
 
     if args.distribution_strategy == "multiprocessing":
-        executor = cluster_tools.get_executor("multiprocessing", args.jobs)
-        logging.info("Using pool of {} workers.".format(args.jobs))
+        # Also accept "processes" instead of job to be compatible with segmentation-tools.
+        # In the long run, the args should be unified and provided by the clustertools.
+        if "jobs" in args:
+            jobs = args.jobs
+        elif "processes" in args:
+            jobs = args.processes
+        else:
+            jobs = cpu_count()
+
+        executor = cluster_tools.get_executor("multiprocessing", jobs)
+        logging.info("Using pool of {} workers.".format(jobs))
     elif args.distribution_strategy == "slurm":
         if args.job_resources is None:
             raise argparse.ArgumentTypeError(