gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
QTZPOC-20: Enable code_paths for sklearn models; improve worker shutdown handling

The MLFlowRepository now resolves and forwards code_paths to mlflow.sklearn.log_model, ensuring custom code dependencies for sklearn models are preserved and correctly reloaded. A new integration test validates this end-to-end functionality. Additionally, worker exception handling is refined to distinguish between deliberate shutdowns (SystemExit, KeyboardInterrupt) and unhandled exceptions, preventing unnecessary error logging for expected exits.
This commit is contained in:
Vitor Pimentel dos Santos
2026-09-10 11:46:29 -03:00
parent bda83e7943
commit 512a49f78f
4 changed files with 147 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
laborious/utils/repository/model_repository.py
View File

@@ -1330,7 +1330,7 @@ class MLFlowRepository(SientiaMonitoring):
return code_dir if path.isdir(code_dir) else None
def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]:
"""Resolve the code dependencies to forward when re-logging a pyfunc model.
"""Resolve the code dependencies to forward when re-logging a model.
MLflow flattens whatever was given as `code_paths` into `<artifact>/code/`,
so the layout differs per model: some ship a `utils` package, others a
@@ -1390,15 +1390,21 @@ class MLFlowRepository(SientiaMonitoring):
core_labels = self.get_core_labels(metadata, operation_type='log_model')
start_time = time.time()
code_paths = (
self.resolve_code_paths(model_data.get('artifact_path'), metadata)
if flavor in ('sklearn', 'pyfunc')
else None
)
try:
if flavor == 'sklearn':
mlflow.sklearn.log_model(model, model_type)
self.debug(f'Code paths: {code_paths}', metadata)
mlflow.sklearn.log_model(model, model_type, code_paths=code_paths or None)
elif flavor == 'pyfunc':
code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata)
self.debug(f'Code paths: {code_paths}', metadata)
self.debug(f'Code path: {code_path}', metadata)
model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False)
model.store_model(artifact_path=model_type, code_path=code_paths, to_disk=False)
self.debug('Model uploaded successfully', metadata)
elif flavor == 'joblib':