Add ImageSpec to wine classification example (#48)

* Add imagespec

Signed-off-by: Peeter Piegaze <1153481+ppiegaze@users.noreply.github.com>

* change package flytekit to unionai

Signed-off-by: Peeter Piegaze <1153481+ppiegaze@users.noreply.github.com>

---------

Signed-off-by: Peeter Piegaze <1153481+ppiegaze@users.noreply.github.com>

.gitignore

@@ -1,2 +1,3 @@
 .idea
 .venv
+.DS_Store

wine-classification/{{cookiecutter.project_name}}/image-requirements.txt

@@ -0,0 +1,2 @@
+pandas
+scikit-learn

wine-classification/{{cookiecutter.project_name}}/local-requirements.txt

@@ -0,0 +1,3 @@
+unionai
+flytekitplugins-envd
+-r image-requirements.txt

wine-classification/{{cookiecutter.project_name}}/workflows/wine_classification_example.py

@@ -2,19 +2,26 @@
 
 from sklearn.datasets import load_wine
 from sklearn.linear_model import LogisticRegression
-from flytekit import task, workflow
+from flytekit import ImageSpec, task, workflow
 
-@task
+image_spec = ImageSpec(
+    registry="ghcr.io/<my-github-org>",
+    name="wine-classification-image",
+    base_image="ghcr.io/flyteorg/flytekit:py3.11-latest",
+    requirements="image-requirements.txt"
+)
+
+@task(container_image=image_spec)
 def get_data() -> pd.DataFrame:
     """Get the wine dataset."""
     return load_wine(as_frame=True).frame
 
-@task
+@task(container_image=image_spec)
 def process_data(data: pd.DataFrame) -> pd.DataFrame:
     """Simplify the task from a 3-class to a binary classification problem."""
     return data.assign(target=lambda x: x["target"].where(x["target"] == 0, 1))
 
-@task
+@task(container_image=image_spec)
 def train_model(data: pd.DataFrame, hyperparameters: dict) -> LogisticRegression:
     """Train a model on the wine dataset."""
     features = data.drop("target", axis="columns")
@@ -30,6 +37,3 @@ def training_workflow(hyperparameters: dict = {"C": 0.1}) -> LogisticRegression:
         data=processed_data,
         hyperparameters=hyperparameters,
     )
-
-if __name__ == "__main__":
-    training_workflow(hyperparameters={"C": 0.1})