From b31d0f6c988c8fdaaccd96ea824477c72b27f17d Mon Sep 17 00:00:00 2001
From: Ryota Egashira <ryota.egashira@acryl.io>
Date: Tue, 4 Mar 2025 13:17:09 -0800
Subject: [PATCH] fix comment

---
 .../src/datahub/ingestion/source/vertexai.py  | 22 +++++++++----------
 .../tests/unit/test_vertexai_source.py        |  1 +
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/metadata-ingestion/src/datahub/ingestion/source/vertexai.py b/metadata-ingestion/src/datahub/ingestion/source/vertexai.py
index 97a4b777fd8f59..0bb4b360b62eeb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/vertexai.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/vertexai.py
@@ -143,7 +143,7 @@ def get_report(self) -> SourceReport:
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         """
-        Main Function to fetch and yields work units for various VertexAI resources.
+        Main Function to fetch and yields mcps for various VertexAI resources.
         - Models and Model Versions from the Model Registry
         - Training Jobs
         """
@@ -165,15 +165,15 @@ def _gen_project_workunits(self) -> Iterable[MetadataWorkUnit]:
 
     def _get_ml_models_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Fetch List of Models in Model Registry and generate a corresponding work unit.
+        Fetch List of Models in Model Registry and generate a corresponding mcp.
         """
         registered_models = self.client.Model.list()
         for model in registered_models:
-            # create work unit for Model Group (= Model in VertexAI)
+            # create mcp for Model Group (= Model in VertexAI)
             yield from self._gen_ml_group_mcps(model)
             model_versions = model.versioning_registry.list_versions()
             for model_version in model_versions:
-                # create work unit for Model (= Model Version in VertexAI)
+                # create mcp for Model (= Model Version in VertexAI)
                 logger.info(
                     f"Ingesting a model (name: {model.display_name} id:{model.name})"
                 )
@@ -183,11 +183,11 @@ def _get_ml_models_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
 
     def _get_training_jobs_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Fetches training jobs from Vertex AI and generates corresponding work units.
+        Fetches training jobs from Vertex AI and generates corresponding mcps.
         This method retrieves various types of training jobs from Vertex AI, including
         CustomJob, CustomTrainingJob, CustomContainerTrainingJob, CustomPythonPackageTrainingJob,
         AutoMLTabularTrainingJob, AutoMLTextTrainingJob, AutoMLImageTrainingJob, AutoMLVideoTrainingJob,
-        and AutoMLForecastingTrainingJob. For each job, it generates work units containing metadata
+        and AutoMLForecastingTrainingJob. For each job, it generates mcps containing metadata
         about the job, its inputs, and its outputs.
         """
         class_names = [
@@ -235,7 +235,7 @@ def _gen_training_job_mcps(
         self, job_meta: TrainingJobMetadata
     ) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Generate a work unit for VertexAI Training Job
+        Generate a mcp for VertexAI Training Job
         """
         job = job_meta.job
         job_id = self._make_vertexai_job_name(entity_id=job.name)
@@ -293,7 +293,7 @@ def _gen_ml_group_mcps(
         model: Model,
     ) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Generate an MLModelGroup work unit for a VertexAI  Model.
+        Generate an MLModelGroup mcp for a VertexAI  Model.
         """
         ml_model_group_urn = self._make_ml_model_group_urn(model)
 
@@ -518,7 +518,7 @@ def _gen_ml_model_endpoint_mcps(
         training_job_urn: Optional[str] = None,
     ) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Generate an MLModel and Endpoint work unit for an VertexAI Model Version.
+        Generate an MLModel and Endpoint mcp for an VertexAI Model Version.
         """
 
         endpoint: Optional[Endpoint] = self._search_endpoint(model)
@@ -539,11 +539,11 @@ def _gen_ml_model_mcps(
         endpoint_urn: Optional[str] = None,
     ) -> Iterable[MetadataChangeProposalWrapper]:
         """
-        Generate an MLModel workunit for an VertexAI Model Version.
+        Generate an MLModel mcp for an VertexAI Model Version.
         Every Model Version is a DataHub MLModel entity associated with an MLModelGroup
         corresponding to a registered Model in VertexAI Model Registry.
         """
-        logging.info(f"starting model work unit for model {model.name}")
+        logging.info(f"generating model mcp for {model.name}")
 
         model_group_urn = self._make_ml_model_group_urn(model)
         model_name = self._make_vertexai_model_name(entity_id=model.name)
diff --git a/metadata-ingestion/tests/unit/test_vertexai_source.py b/metadata-ingestion/tests/unit/test_vertexai_source.py
index 27784882b429bf..566884ee42bc94 100644
--- a/metadata-ingestion/tests/unit/test_vertexai_source.py
+++ b/metadata-ingestion/tests/unit/test_vertexai_source.py
@@ -193,6 +193,7 @@ def test_get_endpoint_mcps(
         elif isinstance(aspect, SubTypesClass):
             assert aspect.typeNames == ["Endpoint"]
 
+
 def test_get_training_jobs_mcps(
     source: VertexAISource,
 ) -> None: