fix: log MOJO model as directory

capitalone · Sep 26, 2024 · 240b254 · 240b254
1 parent 8227c45
commit 240b254
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/rubicon_ml/client/mixin.py b/rubicon_ml/client/mixin.py
@@ -275,6 +275,12 @@ def log_h2o_model(
             if use_mojo:
                 model_data_path = f"{temp_dir_name}/{artifact_name}.zip"
                 h2o_model.save_mojo(path=model_data_path)
+
+                artifact = self.log_artifact(
+                    name=artifact_name,
+                    data_directory=model_data_path,
+                    **log_artifact_kwargs,
+                )
             else:
                 model_data_path = h2o.save_model(
                     h2o_model,
@@ -283,11 +289,11 @@ def log_h2o_model(
                     path=temp_dir_name,
                 )
 
-            artifact = self.log_artifact(
-                name=artifact_name,
-                data_path=model_data_path,
-                **log_artifact_kwargs,
-            )
+                artifact = self.log_artifact(
+                    name=artifact_name,
+                    data_path=model_data_path,
+                    **log_artifact_kwargs,
+                )
 
         return artifact