diff --git a/benchmarks/test_benchmark_chain_100.py b/benchmarks/test_benchmark_chain_100.py index 0359f29..000aae7 100644 --- a/benchmarks/test_benchmark_chain_100.py +++ b/benchmarks/test_benchmark_chain_100.py @@ -1,10 +1,14 @@ +import itertools from abc import ABC import pytest from cosy.maestro import Maestro +from luigi.mock import MockTarget from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter +counter = itertools.count() + class ChainLink(CoSyLuigiTask, ABC): chain_link: CoSyLuigiTaskParameter | None @@ -17,14 +21,13 @@ class StartingLink(ChainLink): class RepeatingLink(ChainLink): chain_link = CoSyLuigiTaskParameter(ChainLink) - -class FinalLink(CoSyLuigiTask): - chain_link = CoSyLuigiTaskParameter(ChainLink) + def output(self): + return {"counter": MockTarget(str(next(counter)))} @pytest.fixture def repo(): - return CoSyLuigiRepo(ChainLink, FinalLink) + return CoSyLuigiRepo(ChainLink) def create_infinite_chain(repo): @@ -32,7 +35,7 @@ def create_infinite_chain(repo): repo.cls_repo, repo.taxonomy, ) - list(maestro.query(FinalLink.target(), max_count=100)) + list(maestro.query(RepeatingLink.target(), max_count=100)) def test_benchmark_chain_creation(repo, benchmark): diff --git a/examples/ml_blood_sugar_level/ml_blood_sugar_level.py b/examples/ml_blood_sugar_level/ml_blood_sugar_level.py index 60df493..8228880 100644 --- a/examples/ml_blood_sugar_level/ml_blood_sugar_level.py +++ b/examples/ml_blood_sugar_level/ml_blood_sugar_level.py @@ -1,3 +1,4 @@ +import os import textwrap from abc import ABC from pathlib import Path @@ -9,7 +10,7 @@ from cosy.maestro import Maestro from sklearn.base import RegressorMixin from sklearn.datasets import load_diabetes -from sklearn.linear_model import LassoLars +from sklearn.linear_model import LassoLars, LinearRegression from sklearn.metrics import root_mean_squared_error from sklearn.model_selection import train_test_split from sklearn.preprocessing import MinMaxScaler, RobustScaler @@ -21,14 +22,14 @@ class LoadDiabetesData(CoSyLuigiTask): def output(self): - return {"diabetes_data": luigi.LocalTarget("diabetes.json")} + return {"diabetes_data": luigi.LocalTarget("data/diabetes.json")} def run(self): diabetes = load_diabetes() df = pd.DataFrame( data=np.c_[diabetes["data"], diabetes["target"]], columns=diabetes["feature_names"] + ["target"] ) - + os.makedirs("data", exist_ok=True) df.to_json(self.output()["diabetes_data"].path) @@ -37,10 +38,10 @@ class TrainTestSplit(CoSyLuigiTask): def output(self): return { - "x_train": luigi.LocalTarget("x_train.json"), - "x_test": luigi.LocalTarget("x_test.json"), - "y_train": luigi.LocalTarget("y_train.json"), - "y_test": luigi.LocalTarget("y_test.json"), + "x_train": luigi.LocalTarget("data/x_train.json"), + "x_test": luigi.LocalTarget("data/x_test.json"), + "y_train": luigi.LocalTarget("data/y_train.json"), + "y_test": luigi.LocalTarget("data/y_test.json"), } def run(self): @@ -62,9 +63,9 @@ class FitTransformScaler(CoSyLuigiTask, ABC): def output(self): return { - "scaled_x_train": luigi.LocalTarget(f"{self.scaler_name}_scaled_x_train.json"), - "scaled_x_test": luigi.LocalTarget(f"{self.scaler_name}_scaled_x_test.json"), - "scaler": luigi.LocalTarget(f"{self.scaler_name}_scaler.skops"), + "scaled_x_train": luigi.LocalTarget(f"data/{self.scaler_name}_scaled_x_train.json"), + "scaled_x_test": luigi.LocalTarget(f"data/{self.scaler_name}_scaled_x_test.json"), + "scaler": luigi.LocalTarget(f"data/{self.scaler_name}_scaler.skops"), } def scale(self, data_identifier: str): @@ -98,7 +99,7 @@ class TrainRegressionModel(CoSyLuigiTask, ABC): model: RegressorMixin def _get_variant_label(self): - return f"{self.model_name}-{Path(self.input()['scaled_feats']['scaled_x_train'].path).stem}" + return f"data/{self.model_name}-{Path(self.input()['scaled_feats']['scaled_x_train'].path).stem}" def output(self): return {"model": luigi.LocalTarget(self._get_variant_label() + ".skops")} @@ -114,18 +115,12 @@ def run(self): class TrainLinearRegressionModel(TrainRegressionModel): model_name = "linear_reg" + model = LinearRegression() class TrainLassoLarsModel(TrainRegressionModel): model_name = "lasso_lars" - - def run(self): - x_train = pd.read_json(self.input()["scaled_feats"]["scaled_x_train"].path) - y_train = pd.read_json(self.input()["splitted_data"]["y_train"].path) - - reg = LassoLars() - reg.fit(x_train, y_train) - sio.dump(reg, self.output()["model"].path) + model = LassoLars() class EvaluateRegressionModel(CoSyLuigiTask): @@ -137,7 +132,7 @@ def _get_variant_label(self): return Path(self.input()["regressor"]["model"].path).stem def output(self): - return luigi.LocalTarget("y_pred" + "-" + self._get_variant_label() + ".json") + return {"evaluation": luigi.LocalTarget("data/y_pred" + "-" + self._get_variant_label() + ".json")} def run(self): unknown_types = sio.get_untrusted_types(file=self.input()["regressor"]["model"].path) @@ -152,7 +147,7 @@ def run(self): print(self._get_variant_label()) print(f"RMSE: {rmse}") - y_pred.to_json(self.output().path) + y_pred.to_json(self.output()["evaluation"].path) def main(): diff --git a/src/cosy_luigi/core/combinatorics.py b/src/cosy_luigi/core/combinatorics.py index 69b5e8f..f3a023b 100644 --- a/src/cosy_luigi/core/combinatorics.py +++ b/src/cosy_luigi/core/combinatorics.py @@ -3,15 +3,15 @@ import logging import textwrap from collections import defaultdict +from collections.abc import Mapping 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, Mapping, Sequence + from collections.abc import Callable, Sequence from cosy.core.synthesizer import Specification @@ -26,6 +26,19 @@ def __init__(self, required_task: type[CoSyLuigiTask], *, unique_across_prior_ta class CoSyLuigiTask(luigi.Task): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + output = self.output() + if not output: + return + if not isinstance(output, Mapping): + msg = f"{self.__class__.__name__}'s output method does not return a Mapping. Unlike regular LuigiTasks, CoSyLuigiTasks must return None or a Mapping, i.e. a dict." + raise TypeError(msg) + # Map to filenames, str method of FileSystemTargets is path + # We do not check for value type, as Luigi will throw Exception if its not a FileSystemTarget already + self.task_id += "_" + "_".join(map(str, output.values())) + self.__hash = hash(self.task_id) + @classmethod @cache def get_all_variants(cls): @@ -122,8 +135,6 @@ def combinator(cls): class CoSyLuigiRepo: def __init__(self, *tasks: type[CoSyLuigiTask] | Sequence[type[CoSyLuigiTask]]): - Register.disable_instance_cache() - # Accepts completely heterogeneous nested collections # This doesn't technically need to unpack as flatten could be typed to accept packed tuples diff --git a/tests/test_infinite_chain.py b/tests/test_infinite_chain.py index 7e3bc88..fe4fb14 100644 --- a/tests/test_infinite_chain.py +++ b/tests/test_infinite_chain.py @@ -1,10 +1,14 @@ +import itertools from abc import ABC import pytest from cosy.maestro import Maestro +from luigi.mock import MockTarget from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter +counter = itertools.count() + class ChainLink(CoSyLuigiTask, ABC): chain_link: CoSyLuigiTaskParameter | None @@ -17,14 +21,13 @@ class StartingLink(ChainLink): class RepeatingLink(ChainLink): chain_link = CoSyLuigiTaskParameter(ChainLink) - -class FinalLink(CoSyLuigiTask): - chain_link = CoSyLuigiTaskParameter(ChainLink) + def output(self): + return {"counter": MockTarget(str(next(counter)))} @pytest.fixture def repo(): - return CoSyLuigiRepo(ChainLink, FinalLink) + return CoSyLuigiRepo(ChainLink) def test_infinite_chain(repo): @@ -32,14 +35,13 @@ def test_infinite_chain(repo): repo.cls_repo, repo.taxonomy, ) - results = list(maestro.query(FinalLink.target(), max_count=10)) + results = list(maestro.query(RepeatingLink.target(), max_count=10)) # Check for shapes of the pipelines for i, result in enumerate(results): current_link = result - assert isinstance(current_link, FinalLink) - current_link = current_link.chain_link for _ in range(i): assert isinstance(current_link, RepeatingLink) current_link = current_link.chain_link + current_link = current_link.chain_link assert isinstance(current_link, StartingLink) diff --git a/tests/test_task_id_generation.py b/tests/test_task_id_generation.py new file mode 100644 index 0000000..dc8936d --- /dev/null +++ b/tests/test_task_id_generation.py @@ -0,0 +1,107 @@ +from abc import ABC + +import luigi +import pytest +from cosy.maestro import Maestro +from luigi.mock import MockTarget +from luigi.task import task_id_str + +from cosy_luigi import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter + +target_a = MockTarget("A") +target_b = MockTarget("B") +fs = target_a.fs + + +class Shaded(CoSyLuigiTask, ABC): + identifier: str + + +class ShadedA(Shaded): + identifier = "A" + + +class ShadedB(Shaded): + identifier = "B" + + +class Shade(CoSyLuigiTask): + shaded = CoSyLuigiTaskParameter(Shaded) + + def complete(self): + return True + + +class Evaluate(CoSyLuigiTask): + shade = CoSyLuigiTaskParameter(Shade) + + def run(self): + with self.output()["output"].open("w") as f: + f.write("OK.") + + def output(self): + return {"output": MockTarget(self.shade.shaded.identifier)} + + +class EvaluateWithPotentialToShade(Evaluate): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Copy the behaviour of regular Luigi + self.task_id = task_id_str(self.get_task_family(), self.to_str_params(only_significant=True, only_public=True)) + self.__hash = hash(self.task_id) + + +@pytest.fixture +def repo(): + return CoSyLuigiRepo(Evaluate, Shade, Shaded) + + +@pytest.fixture +def shadeable_repo(): + return CoSyLuigiRepo(EvaluateWithPotentialToShade, Shade, Shaded) + + +def test_shading_not_possible(repo): + fs.clear() + assert not target_a.exists() + assert not target_b.exists() + maestro = Maestro( + repo.cls_repo, + repo.taxonomy, + ) + luigi.build(list(maestro.query(Evaluate.target())), local_scheduler=True, detailed_summary=True) + assert target_a.exists() + assert target_b.exists() + + +def test_shading_would_be_possible(shadeable_repo): + fs.clear() + assert not target_a.exists() + assert not target_b.exists() + maestro = Maestro( + shadeable_repo.cls_repo, + shadeable_repo.taxonomy, + ) + luigi.build(list(maestro.query(EvaluateWithPotentialToShade.target())), local_scheduler=True, detailed_summary=True) + assert not (target_a.exists() and target_b.exists()) + + +def test_output_mapping_is_enforced(): + class TaskWithWrongOutputA(CoSyLuigiTask): + def output(self): + return MockTarget("") + + with pytest.raises(TypeError): + TaskWithWrongOutputA() + + class TaskWithWrongOutputB(CoSyLuigiTask): + def output(self): + return [MockTarget("")] + + with pytest.raises(TypeError): + TaskWithWrongOutputB() + + class TaskWithNoneOutput(CoSyLuigiTask): + pass + + TaskWithNoneOutput()