Fix benchmark workflow

.github/workflows/benchmark.yml

@@ -28,7 +28,7 @@ jobs:
         run: |
           sudo apt update -y
           python -m pip install --upgrade pip
-          pip install asv
+          pip install asv==0.6.6
 
       - name: Set and log asv machine configuration
         run: |

benchmarks/asv.conf.json

@@ -27,10 +27,9 @@
        "in-dir={build_dir} python -m pip install ."
      ],
     // "uninstall_command": ["return-code=any python -mpip uninstall -y {project}"],
-    // "build_command": [
-    //     "python setup.py build",
-    //     "PIP_NO_BUILD_ISOLATION=false python -mpip wheel --no-deps --no-index -w {build_cache_dir} {build_dir}"
-    // ],
+     "build_command": [
+       "python -m pip wheel --no-deps -w {build_cache_dir} {build_dir}"
+     ],
 
     // List of branches to benchmark. If not provided, defaults to "master"
     // (for git) or "default" (for mercurial).

benchmarks/benchmarks_stats.py

@@ -19,14 +19,21 @@ def setup(self) -> None:
             n_variant=100_000, n_sample=1000
         )
         sample_cohort = np.repeat(
-            [0, 1], self.count_cohort_alleles_ds.dims["samples"] // 2
+            [0, 1], self.count_cohort_alleles_ds.sizes["samples"] // 2
         )
         self.count_cohort_alleles_ds["sample_cohort"] = xr.DataArray(
             sample_cohort, dims="samples"
         )
+        # jit warmup
+        warmup_ds = simulate_genotype_call_dataset(n_variant=10, n_sample=10)
+        warmup_ds["sample_cohort"] = xr.DataArray(
+            np.repeat([0, 1], warmup_ds.sizes["samples"] // 2), dims="samples"
+        )
+        count_call_alleles(warmup_ds).call_allele_count.compute()
+        count_cohort_alleles(warmup_ds).cohort_allele_count.compute()
 
     def time_count_call_alleles(self) -> None:
-        count_call_alleles(self.count_call_alleles_ds)
+        count_call_alleles(self.count_call_alleles_ds).call_allele_count.compute()
 
     def time_count_cohort_alleles(self) -> None:
-        count_cohort_alleles(self.count_cohort_alleles_ds)
+        count_cohort_alleles(self.count_cohort_alleles_ds).cohort_allele_count.compute()

requirements-dev.txt

@@ -16,7 +16,7 @@ cbgen > 1.0.5
 bio2zarr[vcf]; platform_system != "Windows"
 yarl
 matplotlib
-asv
+asv >= 0.6.6
 networkx
 aiohttp
 requests