Add option to mark a CoSyLuigiTaskParameter as unique across all prior tasks

.github/workflows/benchmarks.yml

@@ -84,7 +84,7 @@ jobs:
         output-file-path: cosy-luigi-nested/benchmarks.json
         github-token: ${{ steps.cls-python-workflow-token.outputs.token }}
         auto-push: ${{ env.PUSH_TO_BENCHMARKS }}
-        alert-threshold: '300%'
+        alert-threshold: '50%'
         comment-on-alert: true
         fail-on-alert: true
         comment-always: ${{ env.COMMENT_ALWAYS }}

benchmarks/test_benchmark_chain_100.py

@@ -4,7 +4,7 @@
 from cosy.maestro import Maestro
 from luigi.mock import MockTarget
 
-from cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 
 class ChainLink(CoSyLuigiTask, ABC):
@@ -50,3 +50,7 @@ def create_infinite_chain(repo):
 
 def test_benchmark_chain_creation(repo, benchmark):
     benchmark(create_infinite_chain, repo)
+
+
+if __name__ == "__main__":
+    create_infinite_chain(repo)

examples/getting_started/basic_example.py

@@ -3,7 +3,7 @@
 import luigi
 from cosy.maestro import Maestro
 
-from src.cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 
 class TaskA(CoSyLuigiTask):

examples/getting_started/variation_example.py

@@ -5,7 +5,7 @@
 import luigi
 from cosy.maestro import Maestro
 
-from src.cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 
 class WriteTemplateTask(CoSyLuigiTask):

examples/lot_sizing/lot_sizing_pipeline.py

@@ -13,7 +13,7 @@
 from lot_optimizers.silver_meal_heuristic import SilverMeal
 from lot_optimizers.wagner_whitin import WagnerWhitin
 
-from cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 
 class GetCosts(CoSyLuigiTask):

examples/ml_blood_sugar_level/ml_blood_sugar_level.py

@@ -1,7 +1,6 @@
 import json
 import textwrap
 from abc import ABC
-from collections.abc import Callable, Iterable, Mapping, Sequence
 from pathlib import Path
 
 import luigi
@@ -15,7 +14,7 @@
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import MinMaxScaler, RobustScaler
 
-from cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 ninetydegaisle = True
 
@@ -146,7 +145,7 @@ def run(self):
 
 class EvaluateRegressionModel(CoSyLuigiTask):
     regressor = CoSyLuigiTaskParameter(TrainRegressionModel)
-    scaled_feats = CoSyLuigiTaskParameter(FitTransformScaler)
+    scaled_feats = CoSyLuigiTaskParameter(FitTransformScaler, unique_across_prior_tasks=True)
     splitted_data = CoSyLuigiTaskParameter(TrainTestSplit)
 
     def _get_variant_label(self):
@@ -170,31 +169,6 @@ def run(self):
 
         y_pred.to_json(self.output().path)
 
-    @classmethod
-    def constraints(cls) -> Sequence[Callable[..., bool]]:
-        return [lambda vs: check_unique(vs, [FitTransformScaler]), lambda _: ninetydegaisle]
-
-
-def traverse_pipeline(vs: Iterable[CoSyLuigiTask]) -> Iterable[CoSyLuigiTask]:
-    result = [*vs]
-    for v in result:
-        result.extend(traverse_pipeline(v.requires().values()))
-    return result
-
-
-def check_unique(vs: Mapping[str, CoSyLuigiTask], required_to_be_unique: Iterable[type[CoSyLuigiTask]]) -> bool:
-    classes = [pc.__class__ for pc in traverse_pipeline(vs.values())]
-    seen_subclasses = {}
-    for c in classes:
-        for unique in required_to_be_unique:
-            if issubclass(c, unique):
-                if unique in seen_subclasses:
-                    if seen_subclasses[unique] != c:
-                        return False
-                else:
-                    seen_subclasses[unique] = c
-    return True
-
 
 def main():
     repo = CoSyLuigiRepo(

pyproject.toml

@@ -76,9 +76,12 @@ exclude_lines = [
   "no cov",
   "if __name__ == .__main__.:",
   "if TYPE_CHECKING:",
-  "@abstractmethod",
-  "@abc.abstractmethod"
 ]
+exclude_also = [
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "@(abc\\.)?abstractmethod",
+    ]
 
 
 [tool.ruff.lint.extend-per-file-ignores]

src/cosy_luigi/__init__.py

@@ -0,0 +1,3 @@
+from cosy_luigi.core import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+
+__all__ = ["CoSyLuigiRepo", "CoSyLuigiTask", "CoSyLuigiTaskParameter"]

src/cosy_luigi/constraints/__init__.py

@@ -0,0 +1,3 @@
+from cosy_luigi.constraints.unique import is_unique_in_prior_tasks
+
+__all__ = ["is_unique_in_prior_tasks"]

src/cosy_luigi/constraints/unique.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from cosy_luigi.utils import traverse_pipeline
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from cosy_luigi import CoSyLuigiTask
+
+
+def _is_unique_in_prior_tasks(
+    vs: Mapping[str, CoSyLuigiTask], required_to_be_unique: Sequence[type[CoSyLuigiTask]]
+) -> bool:
+    classes = [pc.__class__ for pc in traverse_pipeline(vs.values())]
+    seen_subclasses: dict[type[CoSyLuigiTask], type[CoSyLuigiTask]] = {}
+    for c in classes:
+        for unique in required_to_be_unique:
+            if issubclass(c, unique):
+                if unique in seen_subclasses:
+                    if seen_subclasses[unique] != c:
+                        return False
+                else:
+                    seen_subclasses[unique] = c
+    return True
+
+
+def is_unique_in_prior_tasks(
+    vs: Mapping[str, CoSyLuigiTask], required_to_be_unique: type[CoSyLuigiTask] | Sequence[type[CoSyLuigiTask]]
+) -> bool:
+    return _is_unique_in_prior_tasks(
+        vs,
+        required_to_be_unique if isinstance(required_to_be_unique, Sequence) else [required_to_be_unique],
+    )

src/cosy_luigi/core/__init__.py

@@ -0,0 +1,3 @@
+from cosy_luigi.core.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+
+__all__ = ["CoSyLuigiRepo", "CoSyLuigiTask", "CoSyLuigiTaskParameter"]

src/cosy_luigi/combinatorics.py → src/cosy_luigi/core/combinatorics.py

@@ -1,25 +1,28 @@
 from __future__ import annotations
 
-import inspect
-from abc import ABC
+import logging
+import textwrap
 from collections import defaultdict
-from functools import cache
-from typing import TYPE_CHECKING, cast
+from functools import cache, partial
+from typing import TYPE_CHECKING
 
 import luigi
 from cosy.core import Constructor, SpecificationBuilder
 from luigi.task_register import Register
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable, Mapping, Sequence
+    from collections.abc import Callable, Mapping, Sequence
 
     from cosy.core.synthesizer import Specification
 
+logger = logging.getLogger(__name__)
+
 
 class CoSyLuigiTaskParameter(luigi.TaskParameter):
-    def __init__(self, required_task: type[CoSyLuigiTask]):
+    def __init__(self, required_task: type[CoSyLuigiTask], *, unique_across_prior_tasks: bool = False):
         super().__init__()
         self.required_task = required_task
+        self.unique_across_prior_tasks = unique_across_prior_tasks
 
 
 class CoSyLuigiTask(luigi.Task):
@@ -60,13 +63,32 @@ def requires(self):
 
     @classmethod
     @cache
-    def _requirements(cls):
+    def _requirements(cls) -> Mapping[str, CoSyLuigiTaskParameter]:
         return {
             k: v
             for k, v in cls.get_all_class_attributes().items()
             if not k.startswith("__") and not callable(v) and issubclass(v.__class__, CoSyLuigiTaskParameter)
         }
 
+    @classmethod
+    @cache
+    def get_params(cls):
+        return list(cls._requirements().items())
+
+    @classmethod
+    @cache
+    def requirements_unique_in_prior_tasks(cls) -> Mapping[str, CoSyLuigiTaskParameter]:
+        return {
+            k: task_parameter
+            for k, task_parameter in cls._requirements().items()
+            if task_parameter.unique_across_prior_tasks
+        }
+
+    @classmethod
+    @cache
+    def unique_required_tasks_in_prior(cls) -> Sequence[type[CoSyLuigiTask]]:
+        return [task_parameter.required_task for task_parameter in cls.requirements_unique_in_prior_tasks().values()]
+
     @classmethod
     @cache
     def target(cls):
@@ -76,11 +98,21 @@ def target(cls):
     def constraints(cls) -> Sequence[Callable[..., bool]]:
         return []
 
+    @classmethod
+    def __constraints(cls) -> Sequence[Callable[..., bool]]:
+        from cosy_luigi.constraints.unique import _is_unique_in_prior_tasks  # noqa: PLC0415
+
+        if cls.requirements_unique_in_prior_tasks():
+            return [partial(_is_unique_in_prior_tasks, required_to_be_unique=cls.unique_required_tasks_in_prior())]
+        return []
+
     @classmethod
     def combinator_type(cls):
         sp = SpecificationBuilder()
         for name in [v.required_task.__name__ for v in cls._requirements().values()]:
             sp = sp.argument(name, Constructor(name))
+        for constraint in cls.__constraints():
+            sp = sp.constraint(constraint)
         for constraint in cls.constraints():
             sp = sp.constraint(constraint)
         return sp.suffix(cls.target())
@@ -93,32 +125,53 @@ def combinator(cls):
 
 
 class CoSyLuigiRepo:
-    def __init__(self, *tasks: type[CoSyLuigiTask] | Iterable[type[CoSyLuigiTask]]):
+    def __init__(self, *tasks: type[CoSyLuigiTask] | Sequence[type[CoSyLuigiTask]]):
         Register.disable_instance_cache()
 
         # Accepts completely heterogeneous nested collections
-        def flatten(*heterogeneous_task_collection: type[CoSyLuigiTask] | Iterable[type[CoSyLuigiTask]]):
-            return (
-                task
-                for task_or_task_collection in heterogeneous_task_collection
-                for task in (
-                    flatten(*cast("Iterable[type[CoSyLuigiTask]]", task_or_task_collection))
-                    if isinstance(task_or_task_collection, (tuple, list))
-                    else cast("type[CoSyLuigiTask]", task_or_task_collection).get_all_variants()
-                    if inspect.isabstract(task_or_task_collection)
-                    or ABC in cast("type[CoSyLuigiTask]", task_or_task_collection).__bases__
-                    else (task_or_task_collection,)
-                )
-            )
 
         # This doesn't technically need to unpack as flatten could be typed to accept packed tuples
         # But performance is equivalent/faster because the first layer doesn't need to be checked this way
+        from cosy_luigi.utils import flatten  # noqa: PLC0415
+
         self.luigi_repo: set[type[CoSyLuigiTask]] = set(flatten(*tasks))
+        self.check_unique_in_prior_tasks_sanity()
         self.taxonomy: Mapping[str, set[str]] = defaultdict(set)
-        self.cls_repo: set[tuple[str, Callable, Specification]] = set()
+        self.cls_repo: list[tuple[str, Callable, Specification]] = []
         for task in self.luigi_repo:
-            self.cls_repo.add(task.combinator())
+            self.cls_repo.append(task.combinator())
             for tpe in task.mro()[1:]:
                 if issubclass(tpe, CoSyLuigiTask):
                     # Is a subclass of CosyLuigiTask, but a superclass of task
                     self.taxonomy[task.__name__].add(tpe.__name__)
+
+    def check_unique_in_prior_tasks_sanity(self):
+        for source_task, param_name, required_type in [
+            (task, k, required_unique_task.required_task)
+            for task in self.luigi_repo
+            for k, required_unique_task in task.requirements_unique_in_prior_tasks().items()
+            if not any(
+                issubclass(task, required_unique_task.required_task) and task is not required_unique_task.required_task
+                for task in self.luigi_repo
+            )
+        ]:
+            logger.warning(
+                textwrap.dedent(
+                    f"""
+                        =================================================================
+                            WARNING ABOUT POTENTIALLY INCORRECT MODEL
+
+                            Class:      {source_task.__name__}
+                            Parameter:  {param_name}
+                            Type:       {required_type.__name__}
+
+                                        is required to be unique, but there are no sub-
+                                        classes of it present in the repository. Either
+                                        you forgot adding sub-classes of it to the rep-
+                                        ository or this will not behave as expected.
+
+                                        Please head over to the documentation:
+                        =================================================================
+                    """
+                )
+            )

src/cosy_luigi/utils/__init__.py

@@ -0,0 +1,3 @@
+from cosy_luigi.utils.traversals import flatten, traverse_pipeline
+
+__all__ = ["flatten", "traverse_pipeline"]

src/cosy_luigi/utils/traversals.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import inspect
+from abc import ABC
+from typing import TYPE_CHECKING, cast
+
+from cosy_luigi import CoSyLuigiTask
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+
+
+def flatten(*heterogeneous_task_collection: type[CoSyLuigiTask] | Sequence[type[CoSyLuigiTask]]):
+    return (
+        task
+        for task_or_task_collection in heterogeneous_task_collection
+        for task in (
+            flatten(*cast("Sequence[type[CoSyLuigiTask]]", task_or_task_collection))
+            if isinstance(task_or_task_collection, (tuple, list))
+            else cast("type[CoSyLuigiTask]", task_or_task_collection).get_all_variants()
+            if inspect.isabstract(task_or_task_collection)
+            or ABC in cast("type[CoSyLuigiTask]", task_or_task_collection).__bases__
+            else (task_or_task_collection,)
+        )
+    )
+
+
+def _traverse_pipeline(vs: Sequence[CoSyLuigiTask] | Iterable[CoSyLuigiTask]) -> Sequence[CoSyLuigiTask]:
+    result = [*vs]
+    for v in vs:
+        result.extend(traverse_pipeline(v.requires().values()))
+    return result
+
+
+def traverse_pipeline(
+    to_traverse: CoSyLuigiTask | Sequence[CoSyLuigiTask] | Iterable[CoSyLuigiTask],
+) -> Sequence[CoSyLuigiTask]:
+    return (
+        _traverse_pipeline([to_traverse]) if isinstance(to_traverse, CoSyLuigiTask) else _traverse_pipeline(to_traverse)
+    )