Fix error in batch_run

JasonGross · Nov 7, 2024 · 9b6f505 · 9b6f505
1 parent 8830e1d
commit 9b6f505
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 32 deletions.
diff --git a/gbmi/utils/images.py b/gbmi/utils/images.py
@@ -107,40 +107,42 @@ def batch_run(
             for i in range(0, len(images), batchsize)
         ]
 
-    process = subprocess.Popen(
-        [*args, *map(str, images), *post_args],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        **kwargs,
-    )
-
     if stdout_write is None:
         stdout_write = partial(print, file=sys.stdout)
     if stderr_write is None:
         stderr_write = partial(print, file=sys.stderr)
 
-    stdout_thread = threading.Thread(
-        target=forward_output, args=(process.stdout, stdout_write, trim_printout)
-    )
-    stderr_thread = threading.Thread(
-        target=forward_output, args=(process.stderr, stderr_write, trim_printout)
-    )
+    try:
+        process = subprocess.Popen(
+            [*args, *map(str, images), *post_args],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            **kwargs,
+        )
+
+        stdout_thread = threading.Thread(
+            target=forward_output, args=(process.stdout, stdout_write, trim_printout)
+        )
+        stderr_thread = threading.Thread(
+            target=forward_output, args=(process.stderr, stderr_write, trim_printout)
+        )
 
-    stdout_thread.start()
-    stderr_thread.start()
+        stdout_thread.start()
+        stderr_thread.start()
 
-    stdout_thread.join()
-    stderr_thread.join()
+        stdout_thread.join()
+        stderr_thread.join()
 
-    process.wait()
+        process.wait()
 
-    if check and process.returncode != 0:
-        exn = subprocess.CalledProcessError(process.returncode, process.args)
+        if check and process.returncode != 0:
+            raise subprocess.CalledProcessError(process.returncode, process.args)
+    except (FileNotFoundError, subprocess.CalledProcessError, OSError) as e:
         if wrap_errs is not None:
-            stderr_write(f"Error: {exn}")
-            wrap_errs.append(exn)
+            stderr_write(f"Error: {e}")
+            wrap_errs.append(e)
         else:
-            raise exn
+            raise e
 
     return process
 

diff --git a/notebooks_jason/max_of_K_all_models.py b/notebooks_jason/max_of_K_all_models.py
@@ -300,13 +300,7 @@ def wrap_err(f, *args, return_bool: bool = False, **kwargs):
     try:
         result = f(*args, **kwargs)
         return True if return_bool else result
-    except FileNotFoundError as e:
-        print(f"Warning: {e}")
-        errs.append(e)
-    except subprocess.CalledProcessError as e:
-        print(f"Warning: {e}")
-        errs.append(e)
-    except OSError as e:
+    except (FileNotFoundError, subprocess.CalledProcessError, OSError) as e:
         print(f"Warning: {e}")
         errs.append(e)
     if return_bool:
@@ -330,13 +324,15 @@ def optimize_pngs(errs: list[Exception] = []):
     #     wrap_err(image_utils.optipng, f)
 
     new_errs = []
-    image_utils.optimize(
+    opt_success = wrap_err(
+        image_utils.optimize,
         *LATEX_FIGURE_PATH.glob("*.png"),
         exhaustive=True,
         trim_printout=COMPACT_IMAGE_OPTIMIZE_OUTPUT,
         wrap_errs=new_errs,
+        return_bool=True,
     )
-    opt_success = not new_errs
+    opt_success = opt_success and not new_errs
     errs.extend(new_errs)
 
     if not opt_success: