script-2 missing connector feature_table_name (#529)

rudderlabs · Dec 12, 2024 · d4b2138 · d4b2138
1 parent e882104
commit d4b2138
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/src/predictions/profiles_mlcorelib/ml_core/preprocess_and_train.py b/src/predictions/profiles_mlcorelib/ml_core/preprocess_and_train.py
@@ -267,6 +267,7 @@ def preprocess_and_train(
     parser.add_argument("--merged_config", type=json.loads)
     parser.add_argument("--input_column_types", type=json.loads)
     parser.add_argument("--input_columns", type=json.loads)
+    parser.add_argument("--connector_feature_table_name", type=str)
     parser.add_argument("--wh_creds", type=json.loads)
     parser.add_argument("--output_path", type=str)
     parser.add_argument("--mode", type=str)
@@ -298,6 +299,7 @@ def preprocess_and_train(
     warehouse = wh_creds["type"]
     train_procedure = train_and_store_model_results
     connector = ConnectorFactory.create(wh_creds, output_dir)
+    connector.feature_table_name = args.connector_feature_table_name
     local_folder = connector.get_local_dir()
 
     material_info_ = args.material_names

diff --git a/src/predictions/profiles_mlcorelib/processors/K8sProcessor.py b/src/predictions/profiles_mlcorelib/processors/K8sProcessor.py
@@ -210,6 +210,8 @@ def train(
             json.dumps(input_column_types),
             "--input_columns",
             json.dumps(input_columns),
+            "--connector_feature_table_name",
+            self.connector.feature_table_name,
             "--metrics_table",
             metrics_table,
         ]

diff --git a/src/predictions/profiles_mlcorelib/processors/LocalProcessor.py b/src/predictions/profiles_mlcorelib/processors/LocalProcessor.py
@@ -38,6 +38,8 @@ def train(
             json.dumps(input_column_types),
             "--input_columns",
             json.dumps(input_columns),
+            "--connector_feature_table_name",
+            self.connector.feature_table_name,
             "--wh_creds",
             json.dumps(wh_creds),
             "--output_path",