Repository: dask/dask
Pull request #8492 (2021-12-15 22:00:16)

Allow a callable scheduler for `bag.groupby`
This is related to https://github.com/dask/dask/pull/8472 and represents the only other place in the codebase that I could find where `dask.config.get("scheduler")` is assumed to be a string.


Patch:
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
new file mode 100644
index 000000000..b10ec6cba
--- /dev/null
+++ b/.github/workflows/conda.yml
@@ -0,0 +1,59 @@
+name: Conda build
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+# When this workflow is queued, automatically cancel any previous running
+# or pending jobs from the same branch
+concurrency:
+  group: conda-${{ github.head_ref }}
+  cancel-in-progress: true
+
+# Required shell entrypoint to have properly activated conda environments
+defaults:
+  run:
+    shell: bash -l {0}
+
+jobs:
+  conda:
+    name: Build (and upload)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+      - name: Set up Python
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          miniforge-variant: Mambaforge
+          use-mamba: true
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          mamba install boa conda-verify
+
+          which python
+          pip list
+          mamba list
+      - name: Build conda package
+        run: |
+          # suffix for nightly package versions
+          export VERSION_SUFFIX=a`date +%y%m%d`
+
+          conda mambabuild continuous_integration/recipe \
+                           --no-anaconda-upload \
+                           --output-folder .
+      - name: Upload conda package
+        if: |
+          github.event_name == 'push'
+          && github.ref == 'refs/heads/main'
+          && github.repository == 'dask/dask'
+        env:
+          ANACONDA_API_TOKEN: ${{ secrets.DASK_CONDA_TOKEN }}
+        run: |
+          # install anaconda for upload
+          mamba install anaconda-client
+
+          anaconda upload --label dev noarch/*.tar.bz2
diff --git a/continuous_integration/gpuci/axis.yaml b/continuous_integration/gpuci/axis.yaml
index a29aef177..45d53da74 100644
--- a/continuous_integration/gpuci/axis.yaml
+++ b/continuous_integration/gpuci/axis.yaml
@@ -2,7 +2,7 @@ PYTHON_VER:
 - "3.8"
 
 CUDA_VER:
-- "11.2"
+- "11.5"
 
 LINUX_VER:
 - ubuntu18.04
diff --git a/continuous_integration/recipe/meta.yaml b/continuous_integration/recipe/meta.yaml
new file mode 100644
index 000000000..66b0a0c21
--- /dev/null
+++ b/continuous_integration/recipe/meta.yaml
@@ -0,0 +1,47 @@
+{% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0.dev') + environ.get('VERSION_SUFFIX', '') %}
+{% set py_version=environ.get('CONDA_PY', 36) %}
+
+
+package:
+  name: dask-core
+  version: {{ version }}
+
+source:
+  git_url: ../..
+
+build:
+  number: {{ GIT_DESCRIBE_NUMBER }}
+  noarch: python
+  string: py{{ py_version }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
+  script: {{ PYTHON }} -m pip install . -vv
+
+requirements:
+  host:
+    - python >=3.7
+    - pip
+
+  run:
+    - python >=3.7
+    - cloudpickle >=1.1.1
+    - fsspec >=0.6.0
+    - packaging >=20.0
+    - partd >=0.3.10
+    - pyyaml
+    - toolz >=0.8.2
+
+test:
+  imports:
+    - dask
+  commands:
+    - pip check
+  requires:
+    - pip
+
+about:
+  home: https://github.com/dask/dask/
+  summary: Parallel PyData with Task Scheduling
+  license: BSD-3-Clause
+  license_file: LICENSE.txt
+  summary: Parallel Python with task scheduling
+  doc_url: https://dask.org/
+  dev_url: https://github.com/dask/dask
diff --git a/dask/array/core.py b/dask/array/core.py
index 9ac285915..095dfa447 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3413,9 +3413,9 @@ def to_zarr(
 
     if isinstance(url, zarr.Array):
         z = url
-        if isinstance(z.store, (dict, MutableMapping)) and "distributed" in config.get(
+        if isinstance(z.store, (dict, MutableMapping)) and config.get(
             "scheduler", ""
-        ):
+        ) in ("dask.distributed", "distributed"):
             raise RuntimeError(
                 "Cannot store into in memory Zarr Array using "
                 "the Distributed Scheduler."
diff --git a/dask/bag/core.py b/dask/bag/core.py
index b2e8f14a9..c1d57e95e 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -1514,7 +1514,7 @@ class Bag(DaskMethodsMixin):
         if shuffle is None:
             shuffle = config.get("shuffle", None)
         if shuffle is None:
-            if "distributed" in config.get("scheduler", ""):
+            if config.get("scheduler", None) in ("dask.distributed", "distributed"):
                 shuffle = "tasks"
             else:
                 shuffle = "disk"


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index a691bd01f..ea2cc9b46 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -793,6 +793,13 @@ def test_groupby_with_npartitions_changed():
     assert result.npartitions == 1
 
 
+def test_groupby_with_scheduler_func():
+    from dask.threaded import get
+
+    with dask.config.set(scheduler=get):
+        b.groupby(lambda x: x, npartitions=1).compute()
+
+
 def test_concat():
     a = db.from_sequence([1, 2, 3])
     b = db.from_sequence([4, 5, 6])