log models as artifacts

capitalone · Feb 29, 2024 · 3bdf255 · 3bdf255
1 parent ec18b9d
commit 3bdf255
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 6 deletions.
diff --git a/rubicon_ml/schema/logger.py b/rubicon_ml/schema/logger.py
@@ -181,9 +181,15 @@ def log_with_schema(
                 if artifact == "self":
                     experiment.log_artifact(name=obj.__class__.__name__, data_object=obj)
             elif isinstance(artifact, dict):
-                data_object = _get_data_object(obj, artifact)
-                if data_object is not None:
-                    experiment.log_artifact(name=artifact["name"], data_object=data_object)
+                if "self" in artifact:
+                    logging_func_name = artifact["self"]
+                    logging_func = getattr(experiment, logging_func_name)
+                    logging_func(obj)
+                else:
+                    data_object = _get_data_object(obj, artifact)
+
+                    if data_object is not None:
+                        experiment.log_artifact(name=artifact["name"], data_object=data_object)
 
         for dataframe in self.schema_.get("dataframes", []):
             df_value = _get_df(obj, dataframe)

diff --git a/rubicon_ml/schema/schema/h2o__H2OGeneralizedLinearEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2OGeneralizedLinearEstimator.yaml
@@ -7,6 +7,8 @@ compatibility:
     min_version: 3.44.0.1
 docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2ogeneralizedlinearestimator
 
+artifacts:
+  - self: log_h2o_model
 parameters:
   - name: alpha
     value_attr: alpha

diff --git a/rubicon_ml/schema/schema/h2o__H2OGradientBoostingEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2OGradientBoostingEstimator.yaml
@@ -7,6 +7,8 @@ compatibility:
     min_version: 3.44.0.1
 docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2ogradientboostingestimator
 
+artifacts:
+  - self: log_h2o_model
 parameters:
   - name: auc_type
     value_attr: auc_type

diff --git a/rubicon_ml/schema/schema/h2o__H2ORandomForestEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2ORandomForestEstimator.yaml
@@ -7,6 +7,8 @@ compatibility:
     min_version: 3.44.0.1
 docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2orandomforestestimator
 
+artifacts:
+  - self: log_h2o_model
 parameters:
   - name: auc_type
     value_attr: auc_type

diff --git a/rubicon_ml/schema/schema/h2o__H2OTargetEncoderEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2OTargetEncoderEstimator.yaml
@@ -7,6 +7,8 @@ compatibility:
     min_version: 3.44.0.1
 docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2otargetencoderestimator
 
+artifacts:
+  - self: log_h2o_model
 parameters:
   - name: blending
     value_attr: blending

diff --git a/rubicon_ml/schema/schema/h2o__H2OXGBoostEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2OXGBoostEstimator.yaml
@@ -7,6 +7,8 @@ compatibility:
     min_version: 3.44.0.1
 docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2oxgboostestimator
 
+artifacts:
+  - self: log_h2o_model
 parameters:
   - name: auc_type
     value_attr: auc_type

diff --git a/tests/integration/test_schema.py b/tests/integration/test_schema.py
@@ -97,10 +97,18 @@ def test_estimator_schema_fit_dask_df(
 
 @pytest.mark.integration
 @pytest.mark.parametrize("schema_cls", H2O_SCHEMA_CLS)
-def test_estimator_h2o_schema_train(schema_cls, make_classification_df, rubicon_project):
+def test_estimator_h2o_schema_train(
+    schema_cls,
+    make_classification_df,
+    rubicon_local_filesystem_client_with_project,
+):
+    _, project = rubicon_local_filesystem_client_with_project
+
     X, y = make_classification_df
     y = y > y.mean()
 
-    experiment = _train_and_log(X, y, schema_cls, rubicon_project)
+    experiment = _train_and_log(X, y, schema_cls, project)
+    model_artifact = experiment.artifact(name=schema_cls.__name__)
 
-    assert len(rubicon_project.schema_["parameters"]) == len(experiment.parameters())
+    assert len(project.schema_["parameters"]) == len(experiment.parameters())
+    assert model_artifact.get_data(deserialize="h2o").__class__.__name__ == schema_cls.__name__