Add weighted_choice helper for weighted random selection in TPV configs

docs/topics/advanced_topics.rst

@@ -277,3 +277,43 @@ As this may often be too simplistic, the rank function can be overridden by spec
 rank clause. The rank clause can contain an arbitrary code block, which can do the desired sorting,
 for example by determining destination load by querying the job manager, influx statistics etc.
 The final statement in the rank clause must be the list of sorted destinations.
+
+Helper functions
+================
+TPV exposes a ``helpers`` module in the evaluation context, which provides utility functions
+that can be used in rules, rank functions, params, and other code blocks.
+
++------------------------------------+--------------------------------------------------------------------------+
+| Helper                             | Description                                                              |
++====================================+==========================================================================+
+| ``helpers.job_args_match(``        | Checks whether a dict of argument key/value pairs matches the job's      |
+| ``job, app, args)``                | input parameters. Useful for routing based on specific tool argument     |
+|                                    | values, similar to Galaxy's dynamic tool destination matching.           |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.weighted_random_``       | Returns a shuffled list of all destinations, weighted by each            |
+| ``sampling(destinations)``         | destination's optional ``params.weight`` value. Used in rank functions   |
+|                                    | to break ties or provide a fallback when load-based ranking fails.       |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.weighted_choice(items)`` | Selects a single value from a weighted pool of ``{value, weight}`` dicts |
+|                                    | and returns the chosen ``value`` string. Generic over the meaning of     |
+|                                    | ``value``; primary use case is distributing jobs across multiple job     |
+|                                    | working directory roots. See the recipe in :doc:`tpv_by_example`.        |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.input_size(job)``        | Returns the total input dataset size in GB for the given job.            |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.concurrent_job_count_``  | Returns the number of queued/running jobs for the given tool (and        |
+| ``for_tool(app, tool, user)``      | optional user). Useful for limiting concurrent executions per tool.      |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.tag_values_match(``      | Returns ``True`` if an entity has all ``match_tag_values`` tags and none |
+| ``entity, match_tag_values,``      | of the ``exclude_tag_values`` tags.                                      |
+| ``exclude_tag_values)``            |                                                                          |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.tool_version_eq/lte/``   | Compare the tool's version against a given version string using the      |
+| ``lt/gte/gt(tool, version)``       | specified comparator.                                                    |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.get_tool_resource_``     | Extracts a specific resource field (cores, mem, gpus) from a tool's      |
+| ``field(tool, field_name)``        | resource requirements.                                                   |
++------------------------------------+--------------------------------------------------------------------------+
+| ``helpers.get_dataset_``           | Returns a dict mapping dataset IDs to their object store ID and file     |
+| ``attributes(datasets)``           | size in bytes.                                                           |
++------------------------------------+--------------------------------------------------------------------------+

docs/topics/tpv_by_example.rst

@@ -583,3 +583,47 @@ property.
      slurm:
        runner: slurm
        destination_name_override: "my-dest-with-{cores}-cores-{mem}-mem"
+
+Selecting a job working directory from a weighted pool
+------------------------------------------------------
+The ``helpers.weighted_choice`` helper selects a single value from a weighted
+pool of ``{value, weight}`` dictionaries.  One of the interesting use cases is distributing
+jobs across multiple job working directory roots with configurable weights —
+useful when you have multiple storage paths for job working directories and
+want to steer the majority of jobs to a particular path while keeping a fallback
+available.  The helper is generic over the meaning of ``value``, so any
+string-valued pool (cache paths, runner URLs, etc.) works the same way.
+
+.. code-block:: yaml
+   :linenos:
+   :emphasize-lines: 4-9,14
+
+   global:
+     context:
+       jwd_pool:
+         - value: /fast/jobs
+           weight: 3
+         - value: /slow/jobs
+           weight: 1
+
+   tools:
+     default:
+       cores: 2
+       mem: 4
+       params:
+         job_working_directory: "{helpers.weighted_choice(jwd_pool)}"
+
+In this example, the ``jwd_pool`` context variable is a list of dictionaries,
+each with a ``value`` (the string to use) and an optional ``weight`` (defaults
+to 1).  ``helpers.weighted_choice`` returns the ``value`` of a randomly selected
+item, weighted by the ``weight`` value.  With the weights above, ``/fast/jobs``
+is selected three times as often as ``/slow/jobs``.
+
+To drain a directory, set its weight to ``0``.  If all weights are zero, the
+helper falls back to an unweighted random selection so the configuration always
+resolves.
+
+This recipe requires Galaxy's ``job_working_directory`` job-concern support
+(Galaxy PR :issue:`23133` / :issue:`15616` / :issue:`20062`).  Without that
+change, the param is still set by TPV but Galaxy will use the object-store
+derived path instead.

tests/fixtures/mapping-weighted-choice.yml

@@ -0,0 +1,34 @@
+global:
+  default_inherits: default
+  context:
+    jwd_pool:
+      - value: /fast/jobs
+        weight: 3
+      - value: /slow/jobs
+        weight: 1
+
+tools:
+  default:
+    cores: 2
+    mem: cores * 3
+    params:
+      job_working_directory: "{helpers.weighted_choice(jwd_pool)}"
+    scheduling:
+      require: []
+      prefer:
+        - general
+      accept:
+      reject: []
+
+  bwa:
+    cores: 4
+    mem: cores * 4
+
+destinations:
+  local:
+    runner: local
+    max_accepted_cores: 16
+    max_accepted_mem: 64
+    scheduling:
+      prefer:
+        - general

tests/test_helpers.py

@@ -5,7 +5,7 @@
 from unittest.mock import patch
 
 from tpv.commands.test import mock_galaxy
-from tpv.core.helpers import get_dataset_attributes, weighted_random_sampling
+from tpv.core.helpers import get_dataset_attributes, weighted_choice, weighted_random_sampling
 
 
 class TestHelpers(unittest.TestCase):
@@ -57,3 +57,103 @@ def test_weighted_random_sampling_with_weights_uses_weighted_choices(self):
         self.assertEqual(result, sampled_destinations)
         choices_mock.assert_called_once_with(destinations, weights=[5, 1, 1], k=3)
         sample_mock.assert_not_called()
+
+    def test_weighted_choice_without_weights_uses_unweighted_choice(self):
+        """When no item defines weight, use unweighted random choice."""
+        items = [
+            {"value": "/fast/jobs"},
+            {"value": "/slow/jobs"},
+            {"value": "/backup/jobs", "foo": "bar"},
+        ]
+
+        with patch("tpv.core.helpers.random.choice", return_value=items[1]) as choice_mock:
+            with patch("tpv.core.helpers.random.choices") as choices_mock:
+                result = weighted_choice(items)
+
+        self.assertEqual(result, "/slow/jobs")
+        choice_mock.assert_called_once_with(items)
+        choices_mock.assert_not_called()
+
+    def test_weighted_choice_with_weights_uses_weighted_choices(self):
+        """When any item defines weight, use weighted random choices."""
+        items = [
+            {"value": "/fast/jobs", "weight": 3},
+            {"value": "/slow/jobs"},
+            {"value": "/backup/jobs"},
+        ]
+        with patch("tpv.core.helpers.random.choices", return_value=[items[0]]) as choices_mock:
+            with patch("tpv.core.helpers.random.choice") as choice_mock:
+                result = weighted_choice(items)
+
+        self.assertEqual(result, "/fast/jobs")
+        choices_mock.assert_called_once_with(items, weights=[3, 1, 1], k=1)
+        choice_mock.assert_not_called()
+
+    def test_weighted_choice_missing_weight_defaults_to_one(self):
+        """Items without a weight key should default to weight 1."""
+        items = [
+            {"value": "/a", "weight": 5},
+            {"value": "/b"},
+            {"value": "/c"},
+        ]
+        with patch("tpv.core.helpers.random.choices", return_value=[items[1]]) as choices_mock:
+            weighted_choice(items)
+
+        choices_mock.assert_called_once_with(items, weights=[5, 1, 1], k=1)
+
+    def test_weighted_choice_zero_weights_falls_back_to_unweighted(self):
+        """If all weights are zero or negative, fall back to unweighted choice."""
+        items = [
+            {"value": "/drained/jobs", "weight": 0},
+            {"value": "/also-drained/jobs", "weight": -1},
+        ]
+        with patch("tpv.core.helpers.random.choice", return_value=items[0]) as choice_mock:
+            with patch("tpv.core.helpers.random.choices") as choices_mock:
+                result = weighted_choice(items)
+
+        self.assertEqual(result, "/drained/jobs")
+        choice_mock.assert_called_once_with(items)
+        choices_mock.assert_not_called()
+
+    def test_weighted_choice_negative_weight_is_clamped_to_zero(self):
+        """A negative weight is clamped to 0 while positive weights are preserved."""
+        items = [
+            {"value": "/a", "weight": 5},
+            {"value": "/b", "weight": -3},
+        ]
+        with patch("tpv.core.helpers.random.choices", return_value=[items[0]]) as choices_mock:
+            with patch("tpv.core.helpers.random.choice") as choice_mock:
+                weighted_choice(items)
+
+        choices_mock.assert_called_once_with(items, weights=[5, 0], k=1)
+        choice_mock.assert_not_called()
+
+    def test_weighted_choice_default_weight_keys_use_unweighted(self):
+        """If every item has weight: 1 (the default), use unweighted choice."""
+        items = [
+            {"value": "/a", "weight": 1},
+            {"value": "/b", "weight": 1},
+        ]
+        with patch("tpv.core.helpers.random.choice", return_value=items[0]) as choice_mock:
+            with patch("tpv.core.helpers.random.choices") as choices_mock:
+                weighted_choice(items)
+
+        choice_mock.assert_called_once_with(items)
+        choices_mock.assert_not_called()
+
+    def test_weighted_choice_empty_list_raises_value_error(self):
+        """An empty list should raise ValueError."""
+        with self.assertRaises(ValueError):
+            weighted_choice([])
+
+    def test_weighted_choice_returns_value_string(self):
+        """The helper should return the value string, not the whole dict."""
+        items = [
+            {"value": "/primary/jobs", "weight": 10},
+            {"value": "/secondary/jobs", "weight": 1},
+        ]
+        with patch("tpv.core.helpers.random.choices", return_value=[items[1]]):
+            result = weighted_choice(items)
+
+        self.assertIsInstance(result, str)
+        self.assertEqual(result, "/secondary/jobs")

tests/test_mapper_weighted_choice.py

@@ -0,0 +1,59 @@
+"""Integration tests for the weighted_choice helper in TPV config files."""
+
+import os
+import unittest
+from unittest.mock import patch
+
+from galaxy.jobs import JobDestination
+
+from tpv.commands.test import mock_galaxy
+from tpv.rules import gateway
+
+
+class TestMapperWeightedChoice(unittest.TestCase):
+
+    @staticmethod
+    def _map_to_destination(tool: mock_galaxy.Tool, user: mock_galaxy.User, tpv_config_path: str) -> JobDestination:
+        galaxy_app = mock_galaxy.App(job_conf=os.path.join(os.path.dirname(__file__), "fixtures/job_conf.yml"))
+        job = mock_galaxy.Job()
+        gateway.ACTIVE_DESTINATION_MAPPERS = {}
+        return gateway.map_tool_to_destination(galaxy_app, job, tool, user, tpv_config_files=[tpv_config_path])  # type: ignore[arg-type]
+
+    def test_weighted_choice_sets_job_working_directory_param(self) -> None:
+        """helpers.weighted_choice in a params f-string should set job_working_directory."""
+        fixture = os.path.join(
+            os.path.dirname(__file__),
+            "fixtures/mapping-weighted-choice.yml",
+        )
+        tool = mock_galaxy.Tool("bwa")
+        user = mock_galaxy.User("ford", "prefect@vortex.org")
+
+        # Patch random.choices to always return the first item (weight 3, /fast/jobs)
+        pool = [
+            {"value": "/fast/jobs", "weight": 3},
+            {"value": "/slow/jobs", "weight": 1},
+        ]
+        with patch("tpv.core.helpers.random.choices", return_value=[pool[0]]):
+            destination = self._map_to_destination(tool, user, tpv_config_path=fixture)
+
+        self.assertIn("job_working_directory", destination.params)
+        self.assertEqual(destination.params["job_working_directory"], "/fast/jobs")
+
+    def test_weighted_choice_selects_different_path_when_patched(self) -> None:
+        """When random.choices returns a different item, the param should reflect it."""
+        fixture = os.path.join(
+            os.path.dirname(__file__),
+            "fixtures/mapping-weighted-choice.yml",
+        )
+        tool = mock_galaxy.Tool("bwa")
+        user = mock_galaxy.User("ford", "prefect@vortex.org")
+
+        pool = [
+            {"value": "/fast/jobs", "weight": 3},
+            {"value": "/slow/jobs", "weight": 1},
+        ]
+        # Patch to return the second item
+        with patch("tpv.core.helpers.random.choices", return_value=[pool[1]]):
+            destination = self._map_to_destination(tool, user, tpv_config_path=fixture)
+
+        self.assertEqual(destination.params["job_working_directory"], "/slow/jobs")

tpv/core/helpers.py

@@ -9,7 +9,9 @@
 import random
 from collections.abc import Callable
 from functools import reduce
-from typing import Any
+from typing import Any, TypedDict, TypeVar
+
+T = TypeVar("T")
 
 from galaxy import model
 from galaxy.app import UniverseApplication
@@ -47,14 +49,53 @@ def input_size(job: Job) -> float:
     return calculate_dataset_total(job.input_datasets) / GIGABYTES
 
 
+def _resolve_weights(items: list[T], get_weight: Callable[[T], float]) -> list[float]:
+    """Return clamped, non-negative weights for *items*.
+
+    ``get_weight`` extracts the weight from each item (defaulting to 1).
+    Negative weights are clamped to 0 so they are excluded from selection.
+    """
+    return [max(get_weight(item), 0) for item in items]
+
+
+def _has_weighted_selection(weights: list[float]) -> bool:
+    """True if *weights* warrant weighted selection: at least one non-default
+    weight and at least one positive weight."""
+    return any(w != 1 for w in weights) and any(w > 0 for w in weights)
+
+
 def weighted_random_sampling(destinations: list[Destination]) -> list[Destination]:
     if not destinations:
         return []
-    has_explicit_weight = any(d.params and "weight" in d.params for d in destinations)
-    if not has_explicit_weight:
+    weights = _resolve_weights(destinations, lambda d: d.params.get("weight", 1) if d.params else 1)
+    if not _has_weighted_selection(weights):
         return random.sample(destinations, k=len(destinations))
-    rankings = [(d.params.get("weight", 1) if d.params else 1) for d in destinations]
-    return random.choices(destinations, weights=rankings, k=len(destinations))
+    return random.choices(destinations, weights=weights, k=len(destinations))
+
+
+class WeightedItem(TypedDict, total=False):
+    value: str
+    weight: float
+
+
+def weighted_choice(items: list[WeightedItem]) -> str:
+    """Select a single ``value`` from a weighted pool of ``{value, weight}`` dicts.
+
+    ``weight`` is optional (defaults to 1); a weight of 0 or less excludes the
+    entry.  If no item defines a non-default weight, or all weights are zero,
+    an unweighted random selection is used.  Primary use case: distributing
+    jobs across multiple job working directory roots.
+
+    Raises:
+        ValueError: If *items* is empty.
+    """
+    if not items:
+        raise ValueError("weighted_choice requires a non-empty list of items")
+
+    weights = _resolve_weights(items, lambda item: item.get("weight", 1))
+    if not _has_weighted_selection(weights):
+        return random.choice(items)["value"]
+    return random.choices(items, weights=weights, k=1)[0]["value"]
 
 
 def __get_keys_from_dict(dl: Any, keys_list: list[str]) -> None: