#6387 [Core feature] Add possibility to define a default ImageSpec for a workflow Issue created by HansBambel Motivation: Why do you think this is important? I wanted to create a dummy project with several tasks and wondered if it is possible to define a default ImageSpec for a whole workflow. Having to add

container_image=my_image

to each task-decorator is tedious. import datetime from pathlib import Path from time import sleep import flytekit as fl from flytekit import task, workflow, FlyteDirectory import polars as pl image_polars = fl.ImageSpec(name="polars-image", requirements="uv.lock") @task(container_image=image_polars) def create_initial_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [1, 1, 1, 1], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @task(container_image=image_polars) def create_second_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [2, 2, 2, 2], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @task(container_image=image_polars) def combine_datasets_task(dataset_1: pl.DataFrame, dataset_2:pl.DataFrame) -> pl.DataFrame: """Combine datasets from a folder into a single parquet file.""" combined = pl.concat([dataset_1, dataset_2]) return combined @task(container_image=image_polars) def train_model_task(dataset: pl.DataFrame) -> dict: for i in range(10): print(f"Epoch {i}...") sleep(3) print("Model trained!") return {"status": "trained", "Other params": {"lr": 0.001, "drop-out": 0.5}} @task() def evaluate_model_task(model: dict) -> dict: print("Testing model...") lr = model["Other params"]["lr"] for i in range(10): print((i + 1) * lr) return {"datetime": datetime.datetime.now().strftime("%Y-%m-%d--%H-%M-%S"), "status": "evaluated", "my-metric": {"lr": lr, "accuracy": 0.9}} @task() def create_folder(name: str) -> FlyteDirectory: folder = Path(fl.current_context().working_directory) / name folder.mkdir(parents=True, exist_ok=True) return FlyteDirectory(path=str(folder)) @workflow def pipeline() -> None: ds1 = create_initial_dataset_task() ds2 = create_second_dataset_task() dataset = combine_datasets_task(ds1, ds2) model = train_model_task(dataset) metrics = evaluate_model_task(model) if name == '__main__': pipeline() Goal: What should the final outcome look like, ideally? Instead of adding it to the task, maybe adding it to the workflow. import datetime from pathlib import Path from time import sleep import flytekit as fl from flytekit import task, workflow, FlyteDirectory import polars as pl image_polars = fl.ImageSpec(name="polars-image", requirements="uv.lock") @task() def create_initial_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [1, 1, 1, 1], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @task() def create_second_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [2, 2, 2, 2], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @task() def combine_datasets_task(dataset_1: pl.DataFrame, dataset_2:pl.DataFrame) -> pl.DataFrame: """Combine datasets from a folder into a single parquet file.""" combined = pl.concat([dataset_1, dataset_2]) return combined @task() def train_model_task(dataset: pl.DataFrame) -> dict: for i in range(10): print(f"Epoch {i}...") sleep(3) print("Model trained!") return {"status": "trained", "Other params": {"lr": 0.001, "drop-out": 0.5}} @task() def evaluate_model_task(model: dict) -> dict: print("Testing model...") lr = model["Other params"]["lr"] for i in range(10): print((i + 1) * lr) return {"datetime": datetime.datetime.now().strftime("%Y-%m-%d--%H-%M-%S"), "status": "evaluated", "my-metric": {"lr": lr, "accuracy": 0.9}} @task() def create_folder(name: str) -> FlyteDirectory: folder = Path(fl.current_context().working_directory) / name folder.mkdir(parents=True, exist_ok=True) return FlyteDirectory(path=str(folder)) @workflow(default_image=image_polars) def pipeline() -> None: ds1 = create_initial_dataset_task() ds2 = create_second_dataset_task() dataset = combine_datasets_task(ds1, ds2) model = train_model_task(dataset) metrics = evaluate_model_task(model) if name == '__main__': pipeline() Describe alternatives you've considered I also thought about extending the task-decorator, but wanted to know whether this is the way it is intended: import datetime from pathlib import Path from time import sleep import flytekit as fl from flytekit import task, workflow, FlyteDirectory import polars as pl image_polars = fl.ImageSpec(name="polars-image", requirements="uv.lock") @wraps(task) def polars_task(func=None, *args, retries=3, **kwargs): def decorator(func): # Combine your pre-filled param with any extra ones passed in return task(*args, container_image=image_polars, retries=retries, **kwargs)(func) if func is None: return decorator return decorator(func) @polars_task def create_initial_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [1, 1, 1, 1], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @polars_task(retries=1) def create_second_dataset_task() -> pl.DataFrame: df = pl.DataFrame( { "learner_int_id": [2, 2, 2, 2], "content_int_id": [10, 20, 30, 40], "score": [0.5, 0.6, 0.7, 0.8], "outcome_int": [2, 0, 1, 1], "play_rank": [1, 2, 3, 4], } ) return df @polars_task def combine_datasets_task(dataset_1: pl.DataFrame, dataset_2:pl.DataFrame) -> pl.DataFrame: """Combine datasets from a folder into a single parquet file.""" combined = pl.concat([dataset_1, dataset_2]) return combined @polars_task def train_model_task(dataset: pl.DataFrame) -> dict: for i in range(10): print(f"Epoch {i}...") sleep(3) print("Model trained!") return {"status": "trained", "Other params": {"lr": 0.001, "drop-out": 0.5}} @task() def evaluate_model_task(model: dict) -> dict: print("Testing model...") lr = model["Other params"]["lr"] for i in range(10): print((i + 1) * lr) return {"datetime": datetime.datetime.now().strftime("%Y-%m-%d--%H-%M-%S"), "status": "evaluated", "my-metric": {"lr": lr, "accuracy": 0.9}} @task() def create_folder(name: str) -> FlyteDirectory: folder = Path(fl.current_context().working_directory) / name folder.mkdir(parents=True, exist_ok=True) return FlyteDirectory(path=str(folder)) @workflow def pipeline() -> None: ds1 = create_initial_dataset_task() ds2 = create_second_dataset_task() dataset = combine_datasets_task(ds1, ds2) model = train_model_task(dataset) metrics = evaluate_model_task(model) if name == '__main__': pipeline() Propose: Link/Inline OR Additional context No response Are you sure this issue hasn't been raised already? • Yes Have you read the Code of Conduct? • Yes flyteorg/flyte