Support peak_measured_mem on Windows

Author: tomwhite

cubed/tests/test_core.py

@@ -1,4 +1,3 @@
-import platform
 import random
 from functools import partial
 
@@ -595,7 +594,6 @@ def test_array_pickle(spec, executor):
     assert_array_equal(c.compute(executor=executor), expected)
 
 
-@pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
 def test_measure_reserved_mem(executor):
     pytest.importorskip("lithops")
 

cubed/tests/test_executor_features.py

@@ -1,6 +1,5 @@
 import contextlib
 import os
-import platform
 import re
 
 import fsspec
@@ -72,9 +71,6 @@ def mock_apply_blockwise(*args, **kwargs):
 
 
 # see tests/runtime for more tests for retries for other executors
-@pytest.mark.skipif(
-    platform.system() == "Windows", reason="measuring memory does not run on windows"
-)
 def test_retries(mocker, spec):
     # Use threads executor since single-threaded executor doesn't support retries
     executor = create_executor("threads")
@@ -91,9 +87,6 @@ def test_retries(mocker, spec):
     )
 
 
-@pytest.mark.skipif(
-    platform.system() == "Windows", reason="measuring memory does not run on windows"
-)
 def test_callbacks(spec, executor):
     task_counter = TaskCounter()
     # test following indirectly by checking they don't cause a failure
@@ -149,9 +142,6 @@ def test_callbacks_modal(spec, modal_executor):
         fs.rm(tmp_path, recursive=True)
 
 
-@pytest.mark.skipif(
-    platform.system() == "Windows", reason="measuring memory does not run on windows"
-)
 def test_mem_warn(tmp_path, executor):
     if executor.name not in ("processes", "lithops"):
         pytest.skip(f"{executor.name} executor does not support MemoryWarningCallback")

cubed/tests/test_utils.py

@@ -1,6 +1,5 @@
 import inspect
 import itertools
-import platform
 
 import numpy as np
 import pytest
@@ -92,7 +91,6 @@ def test_memory_repr():
         memory_repr(-1)
 
 
-@pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
 def test_peak_measured_mem():
     assert peak_measured_mem() > 0
 

cubed/tests/utils.py

@@ -1,4 +1,3 @@
-import platform
 from typing import Iterable
 
 import networkx as nx
@@ -18,18 +17,14 @@
     "localhost": {"version": 1},
 }
 
-ALL_EXECUTORS = [create_executor("single-threaded")]
+ALL_EXECUTORS = [
+    create_executor("single-threaded"),
+    create_executor("threads"),
+    create_executor("processes"),
+]
 
 # don't run all tests on every executor as it's too slow, so just have a subset
-MAIN_EXECUTORS = [create_executor("single-threaded")]
-
-
-if platform.system() != "Windows":
-    # ThreadsExecutor calls `peak_measured_mem` which is not supported on Windows
-    ALL_EXECUTORS.append(create_executor("threads"))
-
-    ALL_EXECUTORS.append(create_executor("processes"))
-    MAIN_EXECUTORS.append(create_executor("processes"))
+MAIN_EXECUTORS = [create_executor("single-threaded"), create_executor("processes")]
 
 try:
     ALL_EXECUTORS.append(create_executor("beam"))

cubed/utils.py

@@ -112,13 +112,12 @@ def memory_repr(num: int) -> str:
 
 
 def peak_measured_mem() -> int:
-    """Measures the peak memory usage in bytes.
-
-    Note: this function currently doesn't work on Windows.
-    """
+    """Measures the peak memory usage in bytes."""
 
     if platform.system() == "Windows":
-        raise NotImplementedError("`peak_measured_mem` is not implemented on Windows")
+        import psutil
+
+        return psutil.Process().memory_info().peak_wset
 
     from resource import RUSAGE_SELF, getrusage