[issue-2116] - Add target property to MlModel (#2134)

* Add target property to MlModel * Use pull_request_target for forks * Simplify str * Update ingestion-core setup * Bump ingestion-core version * Update install
2025-09-02 05:33:49 +00:00 · 2022-01-18 16:27:35 +01:00 · 2022-01-18 16:27:35 +01:00 · a9c840778a
commit a9c840778a
parent d4e66ecd7c
11 changed files with 85 additions and 12 deletions
--- a/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/MlModelRepository.java
+++ b/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/MlModelRepository.java
@ -44,9 +44,12 @@ public class MlModelRepository extends EntityRepository<MlModel> {
  private static final Logger LOG = LoggerFactory.getLogger(MlModelRepository.class);
  private static final Fields MODEL_UPDATE_FIELDS =
      new Fields(
-          MlModelResource.FIELD_LIST, "owner,algorithm,dashboard,mlHyperParameters,mlFeatures,mlStore,server,tags");
+          MlModelResource.FIELD_LIST,
+          "owner,algorithm,dashboard,mlHyperParameters,mlFeatures,mlStore,server,target,tags");
  private static final Fields MODEL_PATCH_FIELDS =
-      new Fields(MlModelResource.FIELD_LIST, "owner,algorithm,dashboard,mlHyperParameters,mlFeatures,tags");
+      new Fields(
+          MlModelResource.FIELD_LIST,
+          "owner,algorithm,dashboard,mlHyperParameters,mlFeatures,mlStore,server,target,tags");

  public MlModelRepository(CollectionDAO dao) {
    super(
@ -77,6 +80,7 @@ public class MlModelRepository extends EntityRepository<MlModel> {
    mlModel.setOwner(fields.contains("owner") ? getOwner(mlModel) : null);
    mlModel.setDashboard(fields.contains("dashboard") ? getDashboard(mlModel) : null);
    mlModel.setMlFeatures(fields.contains("mlFeatures") ? mlModel.getMlFeatures() : null);
+    mlModel.setTarget(fields.contains("target") ? mlModel.getTarget() : null);
    mlModel.setMlHyperParameters(fields.contains("mlHyperParameters") ? mlModel.getMlHyperParameters() : null);
    mlModel.setMlStore(fields.contains("mlStore") ? mlModel.getMlStore() : null);
    mlModel.setServer(fields.contains("server") ? mlModel.getServer() : null);
@ -396,6 +400,7 @@ public class MlModelRepository extends EntityRepository<MlModel> {
      updateMlHyperParameters(origMlModel, updatedMlModel);
      updateMlStore(origMlModel, updatedMlModel);
      updateServer(origMlModel, updatedMlModel);
+      updateTarget(origMlModel, updatedMlModel);
    }

    private void updateAlgorithm(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
@ -442,6 +447,13 @@ public class MlModelRepository extends EntityRepository<MlModel> {
      }
    }

+    private void updateTarget(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
+      // Updating the target changes the model response
+      if (recordChange("target", origModel.getTarget(), updatedModel.getTarget())) {
+        majorVersionChange = true;
+      }
+    }
+
    private void updateDashboard(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
      EntityReference origDashboard = origModel.getDashboard();
      EntityReference updatedDashboard = updatedModel.getDashboard();
--- a/catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/mlmodels/MlModelResource.java
+++ b/catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/mlmodels/MlModelResource.java
@ -103,7 +103,7 @@ public class MlModelResource {
  }

  static final String FIELDS =
-      "owner,dashboard,algorithm,mlFeatures,mlHyperParameters,mlStore,server," + "followers,tags,usageSummary";
+      "owner,dashboard,algorithm,mlFeatures,mlHyperParameters,mlStore,server,target,followers,tags,usageSummary";
  public static final List<String> FIELD_LIST = Arrays.asList(FIELDS.replace(" ", "").split(","));

  @GET
@ -423,6 +423,7 @@ public class MlModelResource {
        .withMlHyperParameters(create.getMlHyperParameters())
        .withMlStore(create.getMlStore())
        .withServer(create.getServer())
+        .withTarget(create.getTarget())
        .withTags(create.getTags())
        .withOwner(create.getOwner())
        .withUpdatedBy(securityContext.getUserPrincipal().getName())
--- a/catalog-rest-service/src/main/resources/json/schema/api/data/createMlModel.json
+++ b/catalog-rest-service/src/main/resources/json/schema/api/data/createMlModel.json
@ -31,6 +31,10 @@
      },
      "default" : null
    },
+    "target": {
+      "description": "For supervised ML Models, the value to estimate.",
+      "$ref": "../../entity/data/mlmodel.json#/definitions/featureName"
+    },
    "mlHyperParameters": {
      "description": "Hyper Parameters used to train the ML Model.",
      "type": "array",
--- a/catalog-rest-service/src/main/resources/json/schema/entity/data/mlmodel.json
+++ b/catalog-rest-service/src/main/resources/json/schema/entity/data/mlmodel.json
@ -222,6 +222,10 @@
      },
      "default": null
    },
+    "target": {
+      "description": "For supervised ML Models, the value to estimate.",
+      "$ref": "#/definitions/featureName"
+    },
    "dashboard" : {
      "description": "Performance Dashboard URL to track metric evolution.",
      "$ref" : "../../type/entityReference.json"
--- a/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/mlmodels/MlModelResourceTest.java
+++ b/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/mlmodels/MlModelResourceTest.java
@ -375,6 +375,30 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel> {
        request.withMlHyperParameters(ML_HYPERPARAMS), Status.OK, adminAuthHeaders(), MINOR_UPDATE, change);
  }

+  @Test
+  void put_MlModelAddTarget_200(TestInfo test) throws IOException {
+    CreateMlModel request = create(test);
+    MlModel model = createAndCheckEntity(request, adminAuthHeaders());
+
+    ChangeDescription change = getChangeDescription(model.getVersion());
+    change.getFieldsAdded().add(new FieldChange().withName("target").withNewValue("myTarget"));
+
+    updateAndCheckEntity(request.withTarget("myTarget"), Status.OK, adminAuthHeaders(), MAJOR_UPDATE, change);
+  }
+
+  @Test
+  void put_MlModelUpdateTarget_200(TestInfo test) throws IOException {
+    CreateMlModel request = create(test).withTarget("origTarget");
+    MlModel model = createAndCheckEntity(request, adminAuthHeaders());
+
+    ChangeDescription change = getChangeDescription(model.getVersion());
+    change
+        .getFieldsUpdated()
+        .add(new FieldChange().withName("target").withNewValue("newTarget").withOldValue("origTarget"));
+
+    updateAndCheckEntity(request.withTarget("newTarget"), Status.OK, adminAuthHeaders(), MAJOR_UPDATE, change);
+  }
+
  @Test
  void delete_MlModel_200_ok(TestInfo test) throws HttpResponseException {
    MlModel model = createMlModel(create(test), adminAuthHeaders());
--- a/ingestion-core/Makefile
+++ b/ingestion-core/Makefile
@ -1,7 +1,10 @@
+install_dev:
+	python -m pip install ".[dev]"
+
 generate:
 	mkdir -p src/metadata/generated
 	python3 -m pip install --upgrade pip setuptools
-	python3 -m pip install --upgrade -r requirements-dev.txt
+	make install_dev
 	datamodel-codegen --input ../catalog-rest-service/src/main/resources/json --input-file-type jsonschema --output src/metadata/generated

 clean:
--- a/ingestion-core/README.md
+++ b/ingestion-core/README.md
@ -1,10 +1,12 @@
 ---
-This guide will help you setup the Ingestion framework and connectors
+This guide will help you set up OpenMetadata Core Models
 ---

 ![Python version 3.8+](https://img.shields.io/badge/python-3.8%2B-blue)

-OpenMetadata Ingesiton is a simple framework to build connectors and ingest metadata of various systems through OpenMetadata APIs. It could be used in an orchestration framework(e.g. Apache Airflow) to ingest metadata.
+These models are `pydantic` models automatically generated from the
+central JSON Schemas that define our APIs and Entities.
+
 **Prerequisites**

 - Python &gt;= 3.8.x
@ -12,3 +14,18 @@ OpenMetadata Ingesiton is a simple framework to build connectors and ingest meta
 ### Docs

 Please refer to the documentation here https://docs.open-metadata.org/connectors
+
+### Contribution
+
+In order to contribute to this package:
+
+```bash
+cd ingestion-core
+python -m virtualenv venv
+source venv/bin/activate
+python -m pip install ".[dev]"
+```
+
+> OBS: During development we might need to treat this in a different
+  virtual environment if we are yet to update the reference to the core
+  package in `openmetadata-ingestion`.
--- a/ingestion-core/requirements-dev.txt
+++ b/ingestion-core/requirements-dev.txt
@ -1,5 +0,0 @@
-datamodel-code-generator==0.11.14
-incremental
-twine
-twisted
-wheel
--- a/ingestion-core/setup.py
+++ b/ingestion-core/setup.py
@ -20,6 +20,15 @@ def get_long_description():
    return description


+dev = {
+    "datamodel-code-generator==0.11.14",
+    "incremental",
+    "twine",
+    "twisted",
+    "wheel",
+    "click",
+}
+
 setup(
    name="openmetadata-ingestion-core",
    url="https://open-metadata.org/",
@ -39,4 +48,7 @@ setup(
        "Source": "https://github.com/open-metadata/OpenMetadata",
    },
    packages=find_namespace_packages(where="./src", exclude=["tests*"]),
+    extras_require={
+        "dev": list(dev),
+    },
 )
--- a/ingestion-core/src/metadata/_version.py
+++ b/ingestion-core/src/metadata/_version.py
@ -7,5 +7,5 @@ Provides metadata version information.

 from incremental import Version

-__version__ = Version("metadata", 0, 8, 0, dev=4)
+__version__ = Version("metadata", 0, 8, 0, dev=5)
 __all__ = ["__version__"]
--- a/ingestion/tests/integration/ometa/test_ometa_model_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_model_api.py
@ -265,6 +265,7 @@ class OMetaModelTest(TestCase):
                MlHyperParameter(name="regularisation", value="0.5"),
                MlHyperParameter(name="random", value="hello"),
            ],
+            target="myTarget",
        )

        res = self.metadata.create_or_update(data=model)