add project label (#127)

khorshuheng · web-flow · commit 1f6b8eb42267 · 2022-03-28T15:14:29.000+08:00
Signed-off-by: Khor Shu Heng &lt;khor.heng@gojek.com&gt;

Co-authored-by: Khor Shu Heng &lt;khor.heng@gojek.com&gt;
diff --git a/python/feast_spark/pyspark/abc.py b/python/feast_spark/pyspark/abc.py
@@ -146,6 +146,7 @@ def get_arguments(self) -> List[str]:
 class RetrievalJobParameters(SparkJobParameters):
     def __init__(
         self,
+        project: str,
         feature_tables: List[Dict],
         feature_tables_sources: List[Dict],
         entity_source: Dict,
@@ -155,6 +156,7 @@ def __init__(
     ):
         """
         Args:
+            project (str): Client project
             entity_source (Dict): Entity data source configuration.
             feature_tables_sources (List[Dict]): List of feature tables data sources configurations.
             feature_tables (List[Dict]): List of feature table specification.
@@ -261,13 +263,17 @@ def __init__(
                 }
 
         """
+        self._project = project
         self._feature_tables = feature_tables
         self._feature_tables_sources = feature_tables_sources
         self._entity_source = entity_source
         self._destination = destination
         self._extra_packages = extra_packages if extra_packages else []
         self._checkpoint_path = checkpoint_path
 
+    def get_project(self) -> str:
+        return self._project
+
     def get_name(self) -> str:
         all_feature_tables_names = [ft["name"] for ft in self._feature_tables]
         return f"{self.get_job_type().to_pascal_case()}-{'-'.join(all_feature_tables_names)}"
diff --git a/python/feast_spark/pyspark/launcher.py b/python/feast_spark/pyspark/launcher.py
@@ -233,6 +233,7 @@ def start_historical_feature_retrieval_job(
 
     return launcher.historical_feature_retrieval(
         RetrievalJobParameters(
+            project=project,
             entity_source=_source_to_argument(entity_source, client.config),
             feature_tables_sources=feature_sources,
             feature_tables=[
diff --git a/python/feast_spark/pyspark/launchers/k8s/k8s.py b/python/feast_spark/pyspark/launchers/k8s/k8s.py
@@ -314,6 +314,7 @@ def historical_feature_retrieval(
             azure_credentials=self._get_azure_credentials(),
             arguments=job_params.get_arguments(),
             namespace=self._namespace,
+            extra_labels={LABEL_PROJECT: job_params.get_project()},
         )
 
         job_info = _submit_job(

Original file line number	Diff line number	Diff line change
`@@ -314,6 +314,7 @@ def historical_feature_retrieval(`
`314`	`314`	`azure_credentials=self._get_azure_credentials(),`
`315`	`315`	`arguments=job_params.get_arguments(),`
`316`	`316`	`namespace=self._namespace,`
	`317`	`+ extra_labels={LABEL_PROJECT: job_params.get_project()},`
`317`	`318`	`)`
`318`	`319`
`319`	`320`	`job_info = _submit_job(`