Make memory limit tests less brittle (#511)

format_bytes has changed in dask 2021.4
dask · Jul 22, 2021 · ccde944 · ccde944
1 parent dffe134
commit ccde944
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 17 deletions.
diff --git a/dask_jobqueue/tests/test_htcondor.py b/dask_jobqueue/tests/test_htcondor.py
@@ -5,6 +5,7 @@
 from distributed import Client
 
 import dask
+from dask.utils import format_bytes, parse_bytes
 
 from dask_jobqueue import HTCondorCluster
 
@@ -44,7 +45,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--memory-limit 50.00MB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("50MB")).replace(" ", "")
+        assert f"--memory-limit {formatted_bytes}" in job_script
         assert "--nthreads 2" in job_script
         assert "--nprocs 2" in job_script
 

diff --git a/dask_jobqueue/tests/test_lsf.py b/dask_jobqueue/tests/test_lsf.py
@@ -8,7 +8,7 @@
 import dask
 import pytest
 from dask.distributed import Client
-from distributed.utils import parse_bytes
+from dask.utils import format_bytes, parse_bytes
 
 from dask_jobqueue import LSFCluster, lsf
 
@@ -100,7 +100,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
     with LSFCluster(
         queue="general",
@@ -126,7 +127,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
     with LSFCluster(
         walltime="1:00",
@@ -149,7 +151,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 1 --memory-limit 16.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("16GB")).replace(" ", "")
+        assert f"--nthreads 1 --memory-limit {formatted_bytes}" in job_script
 
 
 @pytest.mark.env("lsf")

diff --git a/dask_jobqueue/tests/test_oar.py b/dask_jobqueue/tests/test_oar.py
@@ -1,7 +1,9 @@
 import sys
 
-from dask_jobqueue import OARCluster
 import dask
+from dask.utils import format_bytes, parse_bytes
+
+from dask_jobqueue import OARCluster
 
 
 def test_header():
@@ -40,7 +42,8 @@ def test_job_script():
         job_script = cluster.job_script()
         assert "#OAR" in job_script
         assert "#OAR -n dask-worker" in job_script
-        assert "--memory-limit 7.00GB " in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--memory-limit {formatted_bytes}" in job_script
         assert "#OAR -l /nodes=1/core=8,walltime=00:02:00" in job_script
         assert "#OAR --project" not in job_script
         assert "#OAR -q" not in job_script
@@ -51,7 +54,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
     with OARCluster(
         walltime="00:02:00",
@@ -67,7 +71,8 @@ def test_job_script():
         job_script = cluster.job_script()
         assert "#OAR" in job_script
         assert "#OAR -n dask-worker" in job_script
-        assert "--memory-limit 7.00GB " in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--memory-limit {formatted_bytes}" in job_script
         assert "#OAR -l /nodes=1/core=8,walltime=00:02:00" in job_script
         assert "#OAR --project" not in job_script
         assert "#OAR -q" not in job_script
@@ -80,7 +85,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
 
 def test_config_name_oar_takes_custom_config():

diff --git a/dask_jobqueue/tests/test_pbs.py b/dask_jobqueue/tests/test_pbs.py
@@ -1,8 +1,11 @@
 import sys
 from time import sleep, time
 
-import dask
 import pytest
+
+import dask
+from dask.utils import format_bytes, parse_bytes
+
 from dask.distributed import Client
 
 from dask_jobqueue import MoabCluster, PBSCluster
@@ -72,7 +75,8 @@ def test_job_script(Cluster):
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
     with Cluster(
         queue="regular",
@@ -95,7 +99,8 @@ def test_job_script(Cluster):
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
 
 @pytest.mark.env("pbs")

diff --git a/dask_jobqueue/tests/test_sge.py b/dask_jobqueue/tests/test_sge.py
@@ -5,6 +5,7 @@
 
 from dask_jobqueue import SGECluster
 import dask
+from dask.utils import format_bytes, parse_bytes
 
 from . import QUEUE_WAIT
 
@@ -83,10 +84,12 @@ def test_job_script(tmpdir):
         resource_spec="h_vmem=12G,mem_req=12G",
     ) as cluster:
         job_script = cluster.job_script()
+        formatted_bytes = format_bytes(parse_bytes("6GB")).replace(" ", "")
+
         for each in [
             "--nprocs 2",
             "--nthreads 3",
-            "--memory-limit 6.00GB",
+            f"--memory-limit {formatted_bytes}",
             "-q my-queue",
             "-P my-project",
             "-l h_rt=02:00:00",

diff --git a/dask_jobqueue/tests/test_slurm.py b/dask_jobqueue/tests/test_slurm.py
@@ -5,6 +5,7 @@
 from distributed import Client
 
 import dask
+from dask.utils import format_bytes, parse_bytes
 
 from dask_jobqueue import SLURMCluster
 
@@ -60,7 +61,8 @@ def test_job_script():
         job_script = cluster.job_script()
         assert "#SBATCH" in job_script
         assert "#SBATCH -J dask-worker" in job_script
-        assert "--memory-limit 7.00GB " in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--memory-limit {formatted_bytes}" in job_script
         assert "#SBATCH -n 1" in job_script
         assert "#SBATCH --cpus-per-task=8" in job_script
         assert "#SBATCH --mem=27G" in job_script
@@ -74,7 +76,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
     with SLURMCluster(
         walltime="00:02:00",
@@ -105,7 +108,8 @@ def test_job_script():
             "{} -m distributed.cli.dask_worker tcp://".format(sys.executable)
             in job_script
         )
-        assert "--nthreads 2 --nprocs 4 --memory-limit 7.00GB" in job_script
+        formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
+        assert f"--nthreads 2 --nprocs 4 --memory-limit {formatted_bytes}" in job_script
 
 
 @pytest.mark.env("slurm")