MindsDB 25.9.2.0a1__py3-none-any.whl → 25.10.0__py3-none-any.whl

mindsdb/integrations/handlers/lightwood_handler/functions.py

@@ -27,51 +27,40 @@ logger = log.getLogger(__name__)
 
 
 def create_learn_mark():
-    if os.name == 'posix':
-        p = Path(tempfile.gettempdir()).joinpath('mindsdb/learn_processes/')
+    if os.name == "posix":
+        p = Path(tempfile.gettempdir()).joinpath("mindsdb/learn_processes/")
         p.mkdir(parents=True, exist_ok=True)
-        p.joinpath(f'{os.getpid()}').touch()
+        p.joinpath(f"{os.getpid()}").touch()
 
 
 def delete_learn_mark():
-    if os.name == 'posix':
-        p = (
-            Path(tempfile.gettempdir())
-            .joinpath('mindsdb/learn_processes/')
-            .joinpath(f'{os.getpid()}')
-        )
+    if os.name == "posix":
+        p = Path(tempfile.gettempdir()).joinpath("mindsdb/learn_processes/").joinpath(f"{os.getpid()}")
         if p.exists():
             p.unlink()
 
 
-@mark_process(name='learn')
+@mark_process(name="learn")
 @profiler.profile()
 def run_generate(df: DataFrame, predictor_id: int, model_storage, args: dict = None):
+    model_storage.training_state_set(current_state_num=1, total_states=5, state_name="Generating problem definition")
+    json_ai_override = args.pop("using", {})
 
-    model_storage.training_state_set(
-        current_state_num=1, total_states=5, state_name='Generating problem definition'
-    )
-    json_ai_override = args.pop('using', {})
-
-    if 'dtype_dict' in json_ai_override:
-        args['dtype_dict'] = json_ai_override.pop('dtype_dict')
+    if "dtype_dict" in json_ai_override:
+        args["dtype_dict"] = json_ai_override.pop("dtype_dict")
 
-    if 'problem_definition' in json_ai_override:
-        args = {**args, **json_ai_override['problem_definition']}
+    if "problem_definition" in json_ai_override:
+        args = {**args, **json_ai_override["problem_definition"]}
 
-    if 'timeseries_settings' in args:
-        for tss_key in [
-            f.name for f in dataclasses.fields(lightwood.api.TimeseriesSettings)
-        ]:
-            k = f'timeseries_settings.{tss_key}'
+    if "timeseries_settings" in args:
+        for tss_key in [f.name for f in dataclasses.fields(lightwood.api.TimeseriesSettings)]:
+            k = f"timeseries_settings.{tss_key}"
             if k in json_ai_override:
-                args['timeseries_settings'][tss_key] = json_ai_override.pop(k)
+                args["timeseries_settings"][tss_key] = json_ai_override.pop(k)
 
     problem_definition = lightwood.ProblemDefinition.from_dict(args)
 
-    model_storage.training_state_set(
-        current_state_num=2, total_states=5, state_name='Generating JsonAI'
-    )
+    model_storage.training_state_set(current_state_num=2, total_states=5, state_name="Generating JsonAI")
     json_ai = lightwood.json_ai_from_problem(df, problem_definition)
     json_ai = json_ai.to_dict()
     unpack_jsonai_old_args(json_ai_override)
@@ -79,9 +68,7 @@ def run_generate(df: DataFrame, predictor_id: int, model_storage, args: dict = N
     rep_recur(json_ai, json_ai_override)
     json_ai = JsonAI.from_dict(json_ai)
 
-    model_storage.training_state_set(
-        current_state_num=3, total_states=5, state_name='Generating code'
-    )
+    model_storage.training_state_set(current_state_num=3, total_states=5, state_name="Generating code")
     code = lightwood.code_from_json_ai(json_ai)
 
     predictor_record = db.Predictor.query.with_for_update().get(predictor_id)
@@ -89,33 +76,27 @@ def run_generate(df: DataFrame, predictor_id: int, model_storage, args: dict = N
     db.session.commit()
 
     json_storage = get_json_storage(resource_id=predictor_id)
-    json_storage.set('json_ai', json_ai.to_dict())
+    json_storage.set("json_ai", json_ai.to_dict())
 
 
-@mark_process(name='learn')
+@mark_process(name="learn")
 @profiler.profile()
 def run_fit(predictor_id: int, df: pd.DataFrame, model_storage) -> None:
     try:
         predictor_record = db.Predictor.query.with_for_update().get(predictor_id)
         assert predictor_record is not None
 
-        predictor_record.data = {'training_log': 'training'}
+        predictor_record.data = {"training_log": "training"}
         predictor_record.status = PREDICTOR_STATUS.TRAINING
         db.session.commit()
 
-        model_storage.training_state_set(
-            current_state_num=4, total_states=5, state_name='Training model'
-        )
-        predictor: lightwood.PredictorInterface = lightwood.predictor_from_code(
-            predictor_record.code
-        )
+        model_storage.training_state_set(current_state_num=4, total_states=5, state_name="Training model")
+        predictor: lightwood.PredictorInterface = lightwood.predictor_from_code(predictor_record.code)
         predictor.learn(df)
 
         db.session.refresh(predictor_record)
 
-        fs = FileStorage(
-            resource_group=RESOURCE_GROUP.PREDICTOR, resource_id=predictor_id, sync=True
-        )
+        fs = FileStorage(resource_group=RESOURCE_GROUP.PREDICTOR, resource_id=predictor_id, sync=True)
         predictor.save(fs.folder_path / fs.folder_name)
         fs.push(compression_level=0)
 
@@ -124,9 +105,7 @@ def run_fit(predictor_id: int, df: pd.DataFrame, model_storage) -> None:
         # getting training time for each tried model. it is possible to do
         # after training only
         fit_mixers = list(
-            predictor.runtime_log[x]
-            for x in predictor.runtime_log
-            if isinstance(x, tuple) and x[0] == "fit_mixer"
+            predictor.runtime_log[x] for x in predictor.runtime_log if isinstance(x, tuple) and x[0] == "fit_mixer"
         )
         submodel_data = predictor_record.data.get("submodel_data", [])
         # add training time to other mixers info
@@ -135,43 +114,39 @@ def run_fit(predictor_id: int, df: pd.DataFrame, model_storage) -> None:
                 submodel_data[i]["training_time"] = tr_time
         predictor_record.data["submodel_data"] = submodel_data
 
-        model_storage.training_state_set(
-            current_state_num=5, total_states=5, state_name='Complete'
-        )
+        model_storage.training_state_set(current_state_num=5, total_states=5, state_name="Complete")
         predictor_record.dtype_dict = predictor.dtype_dict
         db.session.commit()
     except Exception as e:
         db.session.refresh(predictor_record)
-        predictor_record.data = {'error': f'{traceback.format_exc()}\nMain error: {e}'}
+        predictor_record.data = {"error": f"{traceback.format_exc()}\nMain error: {e}"}
         db.session.commit()
         raise e
 
 
-@mark_process(name='learn')
+@mark_process(name="learn")
 def run_learn_remote(df: DataFrame, predictor_id: int) -> None:
     try:
         serialized_df = json.dumps(df.to_dict())
         predictor_record = db.Predictor.query.with_for_update().get(predictor_id)
         resp = requests.post(
-            predictor_record.data['train_url'],
-            json={'df': serialized_df, 'target': predictor_record.to_predict[0]},
+            predictor_record.data["train_url"],
+            json={"df": serialized_df, "target": predictor_record.to_predict[0]},
         )
 
         assert resp.status_code == 200
-        predictor_record.data['status'] = 'complete'
+        predictor_record.data["status"] = "complete"
     except Exception:
-        predictor_record.data['status'] = 'error'
-        predictor_record.data['error'] = str(resp.text)
+        predictor_record.data["status"] = "error"
+        predictor_record.data["error"] = str(resp.text)
 
     db.session.commit()
 
 
-@mark_process(name='learn')
+@mark_process(name="learn")
 def run_learn(df: DataFrame, args: dict, model_storage) -> None:
     if df is None or df.shape[0] == 0:
-        raise Exception(
-            'No input data. Ensure the data source is healthy and try again.'
-        )
+        raise Exception("No input data. Ensure the data source is healthy and try again.")
 
     predictor_id = model_storage.predictor_id
 
@@ -187,15 +162,13 @@ def run_learn(df: DataFrame, args: dict, model_storage) -> None:
     db.session.commit()
 
 
-@mark_process(name='finetune')
+@mark_process(name="finetune")
 def run_finetune(df: DataFrame, args: dict, model_storage):
     try:
         if df is None or df.shape[0] == 0:
-            raise Exception(
-                'No input data. Ensure the data source is healthy and try again.'
-            )
+            raise Exception("No input data. Ensure the data source is healthy and try again.")
 
-        base_predictor_id = args['base_model_id']
+        base_predictor_id = args["base_model_id"]
         base_predictor_record = db.Predictor.query.get(base_predictor_id)
         if base_predictor_record.status != PREDICTOR_STATUS.COMPLETE:
             raise Exception("Base model must be in status 'complete'")
@@ -204,11 +177,9 @@ def run_finetune(df: DataFrame, args: dict, model_storage):
         predictor_record = db.Predictor.query.get(predictor_id)
 
         # TODO move this to ModelStorage (don't work with database directly)
-        predictor_record.data = {'training_log': 'training'}
+        predictor_record.data = {"training_log": "training"}
         predictor_record.training_start_at = datetime.now()
-        predictor_record.status = (
-            PREDICTOR_STATUS.FINETUNING
-        )  # TODO: parallel execution block
+        predictor_record.status = PREDICTOR_STATUS.FINETUNING  # TODO: parallel execution block
         db.session.commit()
 
         base_fs = FileStorage(
@@ -219,28 +190,23 @@ def run_finetune(df: DataFrame, args: dict, model_storage):
         predictor = lightwood.predictor_from_state(
             base_fs.folder_path / base_fs.folder_name, base_predictor_record.code
         )
-        predictor.adjust(df, adjust_args=args.get('using', {}))
+        predictor.adjust(df, adjust_args=args.get("using", {}))
 
-        fs = FileStorage(
-            resource_group=RESOURCE_GROUP.PREDICTOR, resource_id=predictor_id, sync=True
-        )
+        fs = FileStorage(resource_group=RESOURCE_GROUP.PREDICTOR, resource_id=predictor_id, sync=True)
         predictor.save(fs.folder_path / fs.folder_name)
         fs.push(compression_level=0)
 
-        predictor_record.data = (
-            predictor.model_analysis.to_dict()
-        )  # todo: update accuracy in LW as post-finetune hook
+        predictor_record.data = predictor.model_analysis.to_dict()  # todo: update accuracy in LW as post-finetune hook
         predictor_record.code = base_predictor_record.code
-        predictor_record.update_status = 'up_to_date'
+        predictor_record.update_status = "up_to_date"
         predictor_record.status = PREDICTOR_STATUS.COMPLETE
         predictor_record.training_stop_at = datetime.now()
         db.session.commit()
 
     except Exception as e:
-        logger.error(e)
+        logger.error("Unexpected error during Lightwood model finetune:", exc_info=True)
         predictor_id = model_storage.predictor_id
         predictor_record = db.Predictor.query.with_for_update().get(predictor_id)
-        logger.error(traceback.format_exc())
         error_message = format_exception_error(e)
         predictor_record.data = {"error": error_message}
         predictor_record.status = PREDICTOR_STATUS.ERROR