SNOW-1995584: how to register a method with a ragged output

[deadsoul44]

Hello,

I am able to register a custom model but cannot use its predict_nodes method. When I use inference_api decorator, it fails due to ragged output of the method.

class PerpetualCustomRegressionModel(custom_model.CustomModel):
    def __init__(self, context: custom_model.ModelContext) -> None:
        super().__init__(context)
        with open(context["model_file"], "rb") as f:
            self.model = pickle.load(f)

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        return pd.DataFrame({"output": self.model.predict(input)})

    def predict_nodes(self, input: pd.DataFrame):
        return self.model.predict_nodes(input)

When I don't use the ìnference_api decorator, I get the following error:

ValueError: There is no method with name predict_nodes available in the model CAL_HOUSING.CAL_HOUSING_FS.MODEL_EXTRA_METHODS version NEW_WARTHOG_1

How can I use this method?

[sfc-gh-sdas]

Thanks for filing the issue.

Model Registry will not find a method that is not decorated with @ìnference_api. Model Registry also expects an uniform signature (including return value) at this moment. Workaround is what you have done - packed in a JSON OBJECT.

[Taooo-habitus]

I mean if you're okay with not using the built-in model serving function in Snowflake, you can also unwrap it outside of the class, i.e.

reg = Registry(session=session, database_name=database, schema_name=schema)

# Load model
model_version = reg.get_model(name).version(alias)
unwrapped_model = model_version.model
unwrapped_model.predict_nodes(input)