update the meta information when save_values (#104)

* refactor all path of result files to result_context

* update the meta information when save_values

* complement result handler

Author: cyjseagull

python/wedpr_ml_toolkit/setup.py

@@ -20,7 +20,7 @@ def run(self):
 setup_args = dict(
     name='wedpr_ml_toolkit',
     packages=find_packages(),
-    version="1.0.0.dev-20241125",
+    version="1.0.0.dev-20241126",
     description="wedpr-ml-toolkit: The ML toolkit for WeDPR",
     long_description_content_type="text/markdown",
     author="WeDPR Development Team",

python/wedpr_ml_toolkit/test/test_ml_toolkit.py

@@ -2,13 +2,13 @@
 import unittest
 import numpy as np
 import pandas as pd
-from sklearn import metrics
 from wedpr_ml_toolkit.config.wedpr_ml_config import WeDPRMlConfigBuilder
 from wedpr_ml_toolkit.wedpr_ml_toolkit import WeDPRMlToolkit
 from wedpr_ml_toolkit.context.dataset_context import DatasetContext
 from wedpr_ml_toolkit.context.data_context import DataContext
 from wedpr_ml_toolkit.context.job_context import JobType
 from wedpr_ml_toolkit.context.model_setting import PreprocessingSetting
+from wedpr_ml_toolkit.context.model_setting import ModelSetting
 
 
 class WeDPRMlToolkitTestWrapper:
@@ -25,28 +25,19 @@ def test_submit_job(self):
             # x1到x10列，随机数
             **{f'x{i}': np.random.rand(100) for i in range(1, 11)}
         })
+        # the dataset
         dataset1 = DatasetContext(storage_entrypoint=self.wedpr_ml_toolkit.get_storage_entry_point(),
                                   dataset_client=self.wedpr_ml_toolkit.get_dataset_client(),
                                   storage_workspace=self.wedpr_config.user_config.get_workspace_path(),
                                   dataset_id="d-9743660607744005",
                                   is_label_holder=True)
         dataset1.save_values(df, path='d-101')
 
-        # hdfs_path
+        # the dataset
         dataset2 = DatasetContext(storage_entrypoint=self.wedpr_ml_toolkit.get_storage_entry_point(),
                                   dataset_client=self.wedpr_ml_toolkit.get_dataset_client(),
                                   dataset_id="d-9743674298214405")
-
-        dataset2.storage_client = None
-        # dataset2.load_values()
-        if dataset2.storage_client is None:
-            # 支持更新dataset的values数据
-            df2 = pd.DataFrame({
-                'id': np.arange(0, 100),  # id列，顺序整数
-                # x1到x10列，随机数
-                **{f'z{i}': np.random.rand(100) for i in range(1, 11)}
-            })
-            dataset2.save_values(values=df2)
+        print(f"### dataset2 meta: {dataset2.dataset_meta}")
         if dataset1.storage_client is not None:
             # save values to dataset1
             dataset1.save_values(df)
@@ -70,20 +61,67 @@ def test_submit_job(self):
         psi_result = psi_job_context.fetch_job_result(psi_job_id, True)
         print(
             f"* fetch_job_result for psi job {psi_job_id} success, result: {psi_result}")
+        # build the psi result:
+        psi_result_ctx = self.wedpr_ml_toolkit.build_result_context(
+            psi_job_context, psi_result)
+        print(f"* psi_result_ctx: {psi_result_ctx}")
+        (psi_result_values, psi_result_columns,
+         psi_result_shape) = psi_result_ctx.result_dataset.load_values()
+        # obtain the intersection
+        print(
+            f"* psi result, psi_result_columns: {psi_result_columns}, "
+            f"psi_result_shape: {psi_result_shape}, psi_result_values: {psi_result_values}")
         # 初始化
         print(f"* build pre-processing data-context")
         preprocessing_data = DataContext(dataset1, dataset2)
         preprocessing_job_context = self.wedpr_ml_toolkit.build_job_context(
             JobType.PREPROCESSING, project_id, preprocessing_data, PreprocessingSetting())
         # 执行预处理任务
         print(f"* submit pre-processing job")
-        fe_job_id = preprocessing_job_context.submit()
-        print(f"* submit pre-processing job success, job_id: {fe_job_id}")
-        fe_result = preprocessing_job_context.fetch_job_result(fe_job_id, True)
+        preprocessing_job_id = preprocessing_job_context.submit()
+        print(
+            f"* submit pre-processing job success, job_id: {preprocessing_job_id}")
+        preprocessing_result = preprocessing_job_context.fetch_job_result(
+            preprocessing_job_id, True)
         print(
-            f"* fetch pre-processing job result success, job_id: {fe_job_id}, result: {fe_result}")
+            f"* fetch pre-processing job result success, job_id: {preprocessing_job_id}, result: {preprocessing_result}")
         print(preprocessing_job_context.participant_id_list,
               preprocessing_job_context.result_receiver_id_list)
+        # build the context
+        preprocessing_result_ctx = self.wedpr_ml_toolkit.build_result_context(preprocessing_job_context,
+                                                                              preprocessing_result)
+        print(
+            f"* preprocessing_result_ctx: {preprocessing_result_ctx.preprocessing_dataset}")
+        preprocessing_values, columns, shape = preprocessing_result_ctx.preprocessing_dataset.load_values()
+        print(
+            f"* preprocessing_result_dataset, columns: {columns}, shape: {shape}")
+        # test xgb job
+        xgb_data = DataContext(dataset1, dataset2)
+        model_setting = ModelSetting()
+        model_setting.use_psi = True
+        xgb_job_context = self.wedpr_ml_toolkit.build_job_context(
+            job_type=JobType.XGB_TRAINING, project_id=project_id,
+            dataset=xgb_data,
+            model_setting=model_setting, id_fields="id")
+        print(f"* construct xgb job context: participant_id_list: {xgb_job_context.participant_id_list}, "
+              f"result_receiver_id_list: {xgb_job_context.result_receiver_id_list}")
+        xgb_job_id = xgb_job_context.submit()
+        print(f"* submit xgb job success, {xgb_job_id}")
+        xgb_job_result = xgb_job_context.fetch_job_result(xgb_job_id, True)
+        print(f"* xgb job result: {xgb_job_result}")
+        xgb_job_context = self.wedpr_ml_toolkit.build_result_context(
+            job_context=xgb_job_context, job_result_detail=xgb_job_result)
+        print(f"* xgb job result: {xgb_job_context}")
+        # load the feature_importance information
+        (feature_importance_value, feature_importance_cols, feature_importance_shape) = \
+            xgb_job_context.feature_importance_dataset.load_values()
+        print(f"* xgb feature importance information: {feature_importance_cols}, "
+              f"{feature_importance_shape}, {feature_importance_value}")
+        # load the evaluation information
+        (evaluation_value, evaluation_cols, evaluation_shape) = \
+            xgb_job_context.evaluation_dataset.load_values()
+        print(f"* xgb evaluation information: {evaluation_cols}, "
+              f"{evaluation_shape}, {evaluation_value}")
 
     def test_query_job(self, job_id: str, block_until_finish):
         job_result = self.wedpr_ml_toolkit.query_job_status(

python/wedpr_ml_toolkit/wedpr_ml_toolkit/common/utils/constant.py

@@ -7,11 +7,17 @@ class Constant:
     DEFAULT_SUBMIT_JOB_URI = f'{WEDPR_API_PREFIX}project/submitJob'
     DEFAULT_QUERY_JOB_STATUS_URL = f'{WEDPR_API_PREFIX}project/queryJobByCondition'
     DEFAULT_QUERY_JOB_DETAIL_URL = f'{WEDPR_API_PREFIX}scheduler/queryJobDetail'
+    # the dataset related url
     DEFAULT_QUERY_DATASET_URL = f'{WEDPR_API_PREFIX}dataset/queryDataset'
+    DEFAULT_UPDATED_DATASET_URL = f'{WEDPR_API_PREFIX}dataset/updateDatasetMeta'
     PSI_RESULT_FILE = "psi_result.csv"
 
     FEATURE_BIN_FILE = "feature_bin.json"
-    TEST_MODEL_OUTPUT_FILE = "test_output.csv"
-    TRAIN_MODEL_OUTPUT_FILE = "train_output.csv"
-
-    FE_RESULT_FILE = "fe_result.csv"
+    XGB_TREE_PREFIX = "xgb_tree"
+    MODEL_RESULT_FILE = XGB_TREE_PREFIX + '.json'
+    PREPROCESSING_RESULT_FILE = "preprocessing_result.csv"
+    EVALUATION_TABLE_FILE = "mpc_xgb_evaluation_table.csv"
+    FEATURE_IMPORTANCE_FILE = "xgb_result_feature_importance_table.csv"
+    FEATURE_SELECTION_FILE = "xgb_result_column_info_selected.csv"
+    MODEL_FILE = "model_enc.kpl"
+    WOE_IV_FILE = "woe_iv.csv"

python/wedpr_ml_toolkit/wedpr_ml_toolkit/config/wedpr_ml_config.py

@@ -32,8 +32,11 @@ def __init__(self, polling_interval_s: int = 5, max_retries: int = 2, retry_dela
 
 
 class DatasetConfig(BaseObject):
-    def __init__(self, query_dataset_uri=Constant.DEFAULT_QUERY_DATASET_URL):
+    def __init__(self,
+                 query_dataset_uri=Constant.DEFAULT_QUERY_DATASET_URL,
+                 update_dataset_uri=Constant.DEFAULT_UPDATED_DATASET_URL):
         self.query_dataset_uri = query_dataset_uri
+        self.update_dataset_uri = update_dataset_uri
 
 
 class StorageConfig(BaseObject):
@@ -56,11 +59,6 @@ def __init__(self, timeout_seconds=3):
         self.timeout_seconds = timeout_seconds
 
 
-class AgencyConfig(BaseObject):
-    def __init__(self, agency_name=None):
-        self.agency_name = agency_name
-
-
 class WeDPRMlConfig:
     def __init__(self, config_dict):
         self.auth_config = AuthConfig()
@@ -73,8 +71,6 @@ def __init__(self, config_dict):
         self.user_config.set_params(**config_dict)
         self.http_config = HttpConfig()
         self.http_config.set_params(**config_dict)
-        self.agency_config = AgencyConfig()
-        self.agency_config.set_params(**config_dict)
         self.dataset_config = DatasetConfig()
 
 

python/wedpr_ml_toolkit/wedpr_ml_toolkit/context/data_context.py

@@ -9,23 +9,6 @@ class DataContext:
     def __init__(self, *datasets):
         self.datasets = list(datasets)
 
-        self._check_datasets()
-
-    def _save_dataset(self, dataset: DatasetContext):
-        file_path = dataset.dataset_meta.file_path
-        if file_path is None:
-            dataset.dataset_id = utils.make_id(
-                utils.IdPrefixEnum.DATASET.value)
-            file_path = os.path.join(
-                dataset.storage_workspace, dataset.dataset_id)
-            if dataset.storage_client is not None:
-                dataset.storage_client.upload(
-                    dataset.values, file_path)
-
-    def _check_datasets(self):
-        for dataset in self.datasets:
-            self._save_dataset(dataset)
-
     def to_psi_format(self, merge_filed, result_receiver_id_list):
         dataset_psi = []
         for dataset in self.datasets:
@@ -40,14 +23,10 @@ def to_psi_format(self, merge_filed, result_receiver_id_list):
 
     def __generate_dataset_info__(self, id_field: str, receive_result: bool, label_provider: bool, dataset: DatasetContext):
         return {"idFields": [id_field],
-                "dataset": {"owner": dataset.dataset_meta.ownerUserName,
-                            "ownerAgency": dataset.dataset_meta.ownerAgencyName,
-                            "path": dataset.dataset_meta.file_path,
-                            "storageTypeStr": "HDFS",
+                "dataset": {"ownerAgency": dataset.dataset_meta.ownerAgencyName,
                             "datasetID": dataset.dataset_id},
                 "receiveResult": receive_result,
-                "labelProvider": label_provider
-                }
+                "labelProvider": label_provider}
 
     def to_model_formort(self, merge_filed, result_receiver_id_list):
         dataset_model = []

python/wedpr_ml_toolkit/wedpr_ml_toolkit/context/dataset_context.py

@@ -3,6 +3,7 @@
 from wedpr_ml_toolkit.transport.storage_entrypoint import StorageEntryPoint
 from wedpr_ml_toolkit.transport.wedpr_remote_dataset_client import WeDPRDatasetClient
 from wedpr_ml_toolkit.transport.wedpr_remote_dataset_client import DatasetMeta
+import io
 
 
 class DatasetContext:
@@ -27,6 +28,10 @@ def __init__(self,
         # the storage workspace
         self.storage_workspace = storage_workspace
 
+    def __repr__(self):
+        return f"dataset_id: {self.dataset_id}, " \
+               f"dataset_meta: {self.dataset_meta}"
+
     def load_values(self, header=None):
         # 加载hdfs的数据集
         if self.storage_client is not None:
@@ -37,6 +42,25 @@ def load_values(self, header=None):
             return values, values.columns, values.shape
 
     def save_values(self, values: pd.DataFrame = None, path=None):
+        # no values to save
+        if values is None:
+            return
+        csv_buffer = io.StringIO()
+        values.to_csv(csv_buffer, index=False)
+        value_bytes = csv_buffer.getvalue()
+        # update the meta firstly
+        if path is None and self.dataset_meta is not None and self.dataset_meta.datasetId is not None:
+            columns = values.columns.to_list()
+            dataset_meta = DatasetMeta(dataset_id=self.dataset_meta.datasetId,
+                                       dataset_fields=','.join(columns),
+                                       dataset_size=len(value_bytes),
+                                       dataset_record_count=len(values),
+                                       dataset_column_count=len(columns))
+            self.dataset_client.update_dataset(dataset_meta)
+            self.dataset_meta.datasetFields = ','.join(columns)
+            self.dataset_meta.dataset_record_count = len(values)
+            self.dataset_meta.columnCount = len(columns)
+        # update the content
         target_path = self.dataset_meta.file_path
         # 保存数据到hdfs目录
         if path is not None:
@@ -47,7 +71,7 @@ def save_values(self, values: pd.DataFrame = None, path=None):
             target_path = os.path.join(
                 self.storage_workspace, target_path)
         if self.storage_client is not None:
-            self.storage_client.upload(values, target_path)
+            self.storage_client.upload_bytes(value_bytes, target_path)
 
     def update_path(self, path: str = None):
         # 将数据集存入hdfs相同路径，替换旧数据集