Hi, is it possible to wrap a flyte task within a wrapper class implementation. When I tried calling the fit_classifier task from an object of the class I get an error message. Do you have any recommendation on how to wrap a flyte task in a class implementation?

class model_wrapper:
    @task
    def fit_classifier(self, train: pd.DataFrame,
                       val: pd.DataFrame) -> JoblibSerializedFile:

        # fetch the features and target columns from the train dataset
        x = train[train.columns[1:]]
        y = train[train.columns[0]]

        # fetch the features and target columns from the validation dataset
        eval_x = val[val.columns[1:]]
        eval_y = val[val.columns[0]]

        m = self.xgb_classifier()
        # fit the model to the train data
        m.fit(x, y, eval_set=[(eval_x, eval_y)])

        working_dir = flytekit.current_context().working_directory

        fname = os.path.join(working_dir, "model.joblib.dat")
        joblib.dump(m, fname)

        # return the serialized model
        return JoblibSerializedFile(path=fname)

Error message

...
raise _user_exceptions.FlyteAssertion("Input was not specified for: {} of type {}".format(k, var.type))
flytekit.exceptions.user.FlyteAssertion: Input was not specified for: self of type simple: NONE

@bright-florist-6681 we should sync tomorrow on what you want to do. There is probably a Flyte friendly alternative way to get what you want or a good opportunity to try UnionML, which is a conscience wrapper for ML Pipelines in Flyte.

I can help! You basically have to write a custom task resolver

Just for clarification…adedeji and I work together!

cool, lemme know if you want to set up a time to chat this week, you can invite me

<mailto:niels@union.ai|niels@union.ai>