Fix CI

coiled · Nov 13, 2024 · 8466b6c · 8466b6c
1 parent 08dc5d5
commit 8466b6c
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 113 deletions.
diff --git a/.github/workflows/geospatial.yml b/.github/workflows/geospatial.yml
@@ -71,6 +71,7 @@ jobs:
             tests/geospatial -m geo_execution \
             -n 4 --dist loadscope \
             --scale ${{ inputs.scale }} \
+             --memray ${{ env.MEMRAY_PROFILE }} \
 
     - name: Upload benchmark results
       uses: actions/upload-artifact@v4

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -138,7 +138,7 @@ jobs:
           DB_NAME: ${{ matrix.name_prefix }}-${{ matrix.os }}-py${{ matrix.python_version }}.db
           MEMRAY_PROFILE: "none"
         run: |
-          pytest --benchmark -n 4 --dist loadscope --memray $MEMRAY_PROFILE ${{ env.PYTEST_MARKERS }} ${{ matrix.pytest_args }}
+          pytest --benchmark -n 4 --dist loadscope --memray ${{ env.MEMRAY_PROFILE }} ${{ env.PYTEST_MARKERS }} ${{ matrix.pytest_args }}
 
       - name: Dump coiled.Cluster kwargs
         run: cat cluster_kwargs.merged.yaml || true

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -9,7 +9,6 @@
 import pickle
 import subprocess
 import sys
-import tarfile
 import threading
 import time
 import uuid
@@ -29,7 +28,6 @@
 import yaml
 from coiled import Cluster
 from dask.distributed import Client, WorkerPlugin
-from dask.distributed.diagnostics import memray
 from dask.distributed.diagnostics.memory_sampler import MemorySampler
 from dask.distributed.scheduler import logger as scheduler_logger
 from dotenv import load_dotenv
@@ -886,116 +884,6 @@ def memory_multiplier(request):
     return request.param
 
 
-@pytest.fixture
-def memray_profile(
-    pytestconfig,
-    s3,
-    s3_performance_url,
-    s3_storage_options,
-    test_run_benchmark,
-    tmp_path,
-):
-    if not test_run_benchmark:
-        yield contextlib.nullcontext
-    else:
-        memray_option = pytestconfig.getoption("--memray")
-
-        if memray_option == "none":
-            yield contextlib.nullcontext
-            return
-
-        if memray_option != "scheduler":
-            raise ValueError(f"Unhandled value for --memray: {memray_option}")
-
-        @contextlib.contextmanager
-        def _memray_profile(client):
-            local_directory = tmp_path / "profiles" / "memray"
-            local_directory.mkdir(parents=True)
-            try:
-                with memray.memray_scheduler(directory=local_directory):
-                    yield
-            finally:
-                archive_name = "memray.tar.gz"
-                archive = tmp_path / archive_name
-                with tarfile.open(archive, mode="w:gz") as tar:
-                    for item in local_directory.iterdir():
-                        tar.add(item, arcname=item.name)
-                destination = f"{s3_performance_url}/{archive_name}"
-                test_run_benchmark.memray_profiles_url = destination
-                s3.put_file(archive, destination)
-
-        yield _memray_profile
-
-
-@pytest.fixture
-def py_spy_profile(
-    pytestconfig,
-    s3,
-    s3_performance_url,
-    s3_storage_options,
-    test_run_benchmark,
-    tmp_path,
-):
-    if not test_run_benchmark:
-        yield contextlib.nullcontext
-        return
-
-    py_spy_option = pytestconfig.getoption("--py-spy")
-    if py_spy_option == "none":
-        yield contextlib.nullcontext
-        return
-
-    profile_scheduler = False
-    profile_workers = False
-
-    if py_spy_option == "scheduler":
-        profile_scheduler = True
-    elif py_spy_option == "workers":
-        profile_workers = True
-    elif py_spy_option == "all":
-        profile_scheduler = True
-        profile_workers = True
-    else:
-        raise ValueError(f"Unhandled value for --py-spy: {py_spy_option}")
-
-    try:
-        from dask_pyspy import pyspy, pyspy_on_scheduler
-    except ModuleNotFoundError as e:
-        raise ModuleNotFoundError(
-            "py-spy profiling benchmarks requires dask-pyspy to be installed."
-        ) from e
-
-    @contextlib.contextmanager
-    def _py_spy_profile(client):
-        local_directory = tmp_path / "profiles" / "py-spy"
-        local_directory.mkdir(parents=True)
-
-        worker_ctx = contextlib.nullcontext()
-        if profile_workers:
-            worker_ctx = pyspy(local_directory, client=client)
-
-        scheduler_ctx = contextlib.nullcontext()
-        if profile_scheduler:
-            scheduler_ctx = pyspy_on_scheduler(
-                local_directory / "scheduler.json", client=client
-            )
-
-        try:
-            with worker_ctx, scheduler_ctx:
-                yield
-        finally:
-            archive_name = "py-spy.tar.gz"
-            archive = tmp_path / archive_name
-            with tarfile.open(archive, mode="w:gz") as tar:
-                for item in local_directory.iterdir():
-                    tar.add(item, arcname=item.name)
-            destination = f"{s3_performance_url}/{archive_name}"
-            test_run_benchmark.py_spy_profiles_url = destination
-            s3.put_file(archive, destination)
-
-    yield _py_spy_profile
-
-
 @pytest.fixture
 def performance_report(
     pytestconfig,

diff --git a/tests/geospatial/conftest.py b/tests/geospatial/conftest.py
@@ -1,4 +1,6 @@
+import contextlib
 import os
+import tarfile
 import time
 import uuid
 
@@ -37,6 +39,123 @@ def benchmark_type(request):
     return request.param
 
 
+@pytest.fixture
+def memray_profile(
+    pytestconfig,
+    s3,
+    s3_performance_url,
+    s3_storage_options,
+    test_run_benchmark,
+    tmp_path,
+):
+    if not test_run_benchmark:
+        yield contextlib.nullcontext
+    else:
+        memray_option = pytestconfig.getoption("--memray")
+
+        if memray_option == "none":
+            yield contextlib.nullcontext
+            return
+
+        if memray_option != "scheduler":
+            raise ValueError(f"Unhandled value for --memray: {memray_option}")
+
+        try:
+            from dask.distributed.diagnostics import memray
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                "memray profiling benchmarks requires memray to be installed."
+            ) from e
+
+        @contextlib.contextmanager
+        def _memray_profile(client):
+            local_directory = tmp_path / "profiles" / "memray"
+            local_directory.mkdir(parents=True)
+            try:
+                with memray.memray_scheduler(directory=local_directory):
+                    yield
+            finally:
+                archive_name = "memray.tar.gz"
+                archive = tmp_path / archive_name
+                with tarfile.open(archive, mode="w:gz") as tar:
+                    for item in local_directory.iterdir():
+                        tar.add(item, arcname=item.name)
+                destination = f"{s3_performance_url}/{archive_name}"
+                test_run_benchmark.memray_profiles_url = destination
+                s3.put_file(archive, destination)
+
+        yield _memray_profile
+
+
+@pytest.fixture
+def py_spy_profile(
+    pytestconfig,
+    s3,
+    s3_performance_url,
+    s3_storage_options,
+    test_run_benchmark,
+    tmp_path,
+):
+    if not test_run_benchmark:
+        yield contextlib.nullcontext
+        return
+
+    py_spy_option = pytestconfig.getoption("--py-spy")
+    if py_spy_option == "none":
+        yield contextlib.nullcontext
+        return
+
+    profile_scheduler = False
+    profile_workers = False
+
+    if py_spy_option == "scheduler":
+        profile_scheduler = True
+    elif py_spy_option == "workers":
+        profile_workers = True
+    elif py_spy_option == "all":
+        profile_scheduler = True
+        profile_workers = True
+    else:
+        raise ValueError(f"Unhandled value for --py-spy: {py_spy_option}")
+
+    try:
+        from dask_pyspy import pyspy, pyspy_on_scheduler
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            "py-spy profiling benchmarks requires dask-pyspy to be installed."
+        ) from e
+
+    @contextlib.contextmanager
+    def _py_spy_profile(client):
+        local_directory = tmp_path / "profiles" / "py-spy"
+        local_directory.mkdir(parents=True)
+
+        worker_ctx = contextlib.nullcontext()
+        if profile_workers:
+            worker_ctx = pyspy(local_directory, client=client)
+
+        scheduler_ctx = contextlib.nullcontext()
+        if profile_scheduler:
+            scheduler_ctx = pyspy_on_scheduler(
+                local_directory / "scheduler.json", client=client
+            )
+
+        try:
+            with worker_ctx, scheduler_ctx:
+                yield
+        finally:
+            archive_name = "py-spy.tar.gz"
+            archive = tmp_path / archive_name
+            with tarfile.open(archive, mode="w:gz") as tar:
+                for item in local_directory.iterdir():
+                    tar.add(item, arcname=item.name)
+            destination = f"{s3_performance_url}/{archive_name}"
+            test_run_benchmark.py_spy_profiles_url = destination
+            s3.put_file(archive, destination)
+
+    yield _py_spy_profile
+
+
 @pytest.fixture
 def setup_benchmark(
     benchmark_type,