# train_flow.py
from metaflow import FlowSpec, step, pypi_base
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from pathlib import Path
import joblib

@pypi_base(python="3.10.12", packages={"scikit-learn": "1.3.2"})
class TrainFlow(FlowSpec):

    @step
    def start(self):
		self.next(self.train)

    @step
    def train(self):
        X, y = load_iris(return_X_y=True)
        model = RandomForestClassifier().fit(X, y)
        joblib.dump(model, "model.joblib")
        self.model_bytes = Path("model.joblib").read_bytes()
		self.next(self.end)

    @step
    def end(self):
		print("Training complete.")

if __name__ == "__main__":
    TrainFlow()
	

from metaflow import FlowSpec, step, pypi_base, Parameter, Flow
import joblib
import io

@pypi_base(python="3.10.12", packages={"scikit-learn": "1.3.2"})
class PredictFlow(FlowSpec):
	run_id = Parameter("run_id")

    @step
    def start(self):
		self.next(self.predict)

    @step
    def predict(self):
        model_bytes = Flow("TrainFlow").latest_run.data.model_bytes
        model = joblib.load(io.BytesIO(model_bytes))
        self.prediction = model.predict([[5.1, 3.5, 1.4, 0.2]]).tolist()
		self.next(self.end)

    @step
    def end(self):
		print(f"Prediction: {self.prediction}")

if __name__ == "__main__":
    PredictFlow()