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: Remove PyTorch from supported model flavors as it was never functional due to missing runtime dependencies.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-27 14:58:43 -03:00
parent 853c83e662
commit dfb9deae1a
7 changed files with 70 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -45,7 +45,7 @@ ARTIFACTS_PATH = './tmp/artifacts'
TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch' or 'joblib'."
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'."
def force_memory_release(logger: Logger):
@@ -550,7 +550,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_name (str): The name of the model to download from the registry.
metadata (dict[str, Any]): Metadata used for structured logging.
flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch', 'joblib')
flavor (str): Model flavor ('pyfunc', 'sklearn', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -564,7 +564,7 @@ class MLFlowRepository(SientiaMonitoring):
`joblib.load`, which needs no `MLmodel` manifest. The artifact's
`code` directory, when present, is put on `sys.path` first.
"""
if flavor not in ('pyfunc', 'sklearn', 'pytorch', 'joblib'):
if flavor not in ('pyfunc', 'sklearn', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
if flavor == 'joblib':
@@ -584,10 +584,8 @@ class MLFlowRepository(SientiaMonitoring):
try:
if flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
model = mlflow.sklearn.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
@@ -608,7 +606,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_name (str): The name of the model to download.
metadata (dict[str, Any]): Metadata used for structured logging.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -626,7 +624,7 @@ class MLFlowRepository(SientiaMonitoring):
directory, when present, is put on `sys.path` first.
"""
if flavor not in ('sklearn', 'pyfunc', 'pytorch', 'joblib'):
if flavor not in ('sklearn', 'pyfunc', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
@@ -650,10 +648,8 @@ class MLFlowRepository(SientiaMonitoring):
try:
if flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
model = mlflow.pyfunc.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
@@ -677,7 +673,7 @@ class MLFlowRepository(SientiaMonitoring):
model_name (str): Name of the model to download
metadata (dict[str, Any]): Metadata used for structured logging.
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib')
load_wrapper (bool): Whether to load wrapper
Returns:
@@ -862,7 +858,7 @@ class MLFlowRepository(SientiaMonitoring):
metadata (dict[str, Any]): Metadata used for structured logging.
retention (int): Cache retention time in minutes (0 = no cache).
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib')
Returns:
Any: Model object.
@@ -950,7 +946,7 @@ class MLFlowRepository(SientiaMonitoring):
model_name (str): Registered model name.
data (pd.DataFrame): Input data.
retention (int): Cache retention in minutes.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib').
metadata (dict[str, Any]): Metadata used for structured logging.
Returns:
pd.DataFrame | ndarray: Operation result.
@@ -1228,7 +1224,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_data (dict): Model holder with keys 'model' and optional 'artifact_path'.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib').
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib').
model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'.
metadata (dict): Metadata for structured logging.
@@ -1257,8 +1253,6 @@ class MLFlowRepository(SientiaMonitoring):
model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False)
self.debug('Model uploaded successfully', metadata)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
elif flavor == 'joblib':
with tempfile.TemporaryDirectory() as tmp_dir:
pkl_path = path.join(tmp_dir, 'model.pkl')