Adjust h2o model handling to use the h2o functions to save a binary or mojo model.

Author: smlindauer

src/sasctl/pzmm/pickle_model.py

@@ -106,14 +106,22 @@ def pickle_trained_model(
                         )
                 else:
                     return {model_prefix + PICKLE: pickle.dumps(trained_model)}
-            # For binary H2O models, rename the binary file as a pickle file
+            # For binary H2O models, save the binary file as a "pickle" file
             elif is_binary_model and pickle_path:
-                binary_file = Path(pickle_path) / model_prefix
-                binary_file.rename(binary_file.with_suffix(PICKLE))
+                h2o.save_model(
+                    model=trained_model,
+                    force=True,
+                    path=pickle_path,
+                    filename=f"{model_prefix}.pickle"
+                )
             # For MOJO H2O models, gzip the model file and adjust the file extension
             elif is_h2o_model and pickle_path:
-                h2o.save_model(model=trained_model, force=True, path=pickle_path, filename=f'{model_prefix}.mojo')
-            else:
+                trained_model.save_mojo(
+                    force=True,
+                    path=pickle_path,
+                    filename=f"{model_prefix}.mojo"
+                )
+            elif is_binary_model or is_h2o_model:
                 raise ValueError(
                     "There is currently no support for file-less H2O.ai model handling."
                     " Please include a value for the pickle_path argument."