QuickLab-MCP 0.1.0__py3-none-any.whl

ml_mcp/workers/preprocessing.py

@@ -0,0 +1,170 @@
+"""Leak-free feature preprocessing and dataset splitting pipeline."""
+
+import numpy as np
+import pandas as pd
+from sklearn.compose import ColumnTransformer
+from sklearn.impute import SimpleImputer
+from sklearn.model_selection import (
+    KFold,
+    StratifiedKFold,
+    TimeSeriesSplit,
+    train_test_split,
+)
+from sklearn.preprocessing import OneHotEncoder, StandardScaler
+
+from ml_mcp.domain.value_objects import (
+    SplitStrategy,
+    SplitStrategyType,
+    TaskType,
+)
+
+
+class PreprocessingPipeline:
+    """Encapsulates dataset partitioning and strict leak-free feature transformations."""
+
+    def __init__(
+        self,
+        target_column: str,
+        task_type: TaskType,
+        feature_columns: list[str] | None = None,
+        split_strategy: SplitStrategy | None = None,
+    ) -> None:
+        self.target_column = target_column
+        self.task_type = task_type
+        self.feature_columns = feature_columns
+        self.split_strategy = split_strategy or SplitStrategy()
+        self.transformer: ColumnTransformer | None = None
+        self.numeric_features: list[str] = []
+        self.categorical_features: list[str] = []
+
+    def prepare_data(self, df: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
+        """Validate presence of target and features and separate X and y."""
+        if self.target_column not in df.columns:
+            raise ValueError(
+                f"Target column '{self.target_column}' not found in dataset columns: {list(df.columns)}"
+            )
+
+        y = df[self.target_column]
+        if self.feature_columns:
+            missing = [c for c in self.feature_columns if c not in df.columns]
+            if missing:
+                raise ValueError(f"Feature columns {missing} missing from dataset")
+            X = df[self.feature_columns].copy()
+        else:
+            X = df.drop(columns=[self.target_column]).copy()
+
+        # Identify numeric vs categorical columns
+        self.numeric_features = list(X.select_dtypes(include=[np.number]).columns)
+        self.categorical_features = [c for c in X.columns if c not in self.numeric_features]
+
+        return X, y
+
+    def build_transformer(self) -> ColumnTransformer:
+        """Construct scikit-learn ColumnTransformer for leak-free transformations."""
+        transformers = []
+        if self.numeric_features:
+            from sklearn.pipeline import Pipeline
+
+            num_pipe = Pipeline(
+                [
+                    ("imputer", SimpleImputer(strategy="median")),
+                    ("scaler", StandardScaler()),
+                ]
+            )
+            transformers.append(("numeric", num_pipe, self.numeric_features))
+
+        if self.categorical_features:
+            from sklearn.pipeline import Pipeline
+
+            cat_pipe = Pipeline(
+                [
+                    ("imputer", SimpleImputer(strategy="most_frequent")),
+                    ("encoder", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
+                ]
+            )
+            transformers.append(("categorical", cat_pipe, self.categorical_features))
+
+        return ColumnTransformer(transformers=transformers, remainder="drop")
+
+    def split_data(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+    ) -> list[tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
+        """Split data into (X_train, X_val, y_train, y_val) folds based on SplitStrategy."""
+        strat = self.split_strategy
+        seed = strat.random_seed
+
+        if strat.strategy == SplitStrategyType.TRAIN_TEST_SPLIT:
+            stratify = (
+                y
+                if self.task_type
+                in (TaskType.BINARY_CLASSIFICATION, TaskType.MULTICLASS_CLASSIFICATION)
+                else None
+            )
+            X_tr, X_val, y_tr, y_val = train_test_split(
+                X,
+                y,
+                test_size=strat.test_size,
+                random_state=seed,
+                shuffle=strat.shuffle,
+                stratify=stratify,
+            )
+            return [(X_tr, X_val, y_tr, y_val)]
+
+        elif strat.strategy == SplitStrategyType.STRATIFIED_K_FOLD:
+            skf = StratifiedKFold(n_splits=strat.n_splits, shuffle=strat.shuffle, random_state=seed)
+            folds = []
+            for train_idx, val_idx in skf.split(X, y):
+                folds.append(
+                    (X.iloc[train_idx], X.iloc[val_idx], y.iloc[train_idx], y.iloc[val_idx])
+                )
+            return folds
+
+        elif strat.strategy == SplitStrategyType.K_FOLD:
+            kf = KFold(n_splits=strat.n_splits, shuffle=strat.shuffle, random_state=seed)
+            folds = []
+            for train_idx, val_idx in kf.split(X):
+                folds.append(
+                    (X.iloc[train_idx], X.iloc[val_idx], y.iloc[train_idx], y.iloc[val_idx])
+                )
+            return folds
+
+        elif strat.strategy == SplitStrategyType.TIME_SERIES_SPLIT:
+            tscv = TimeSeriesSplit(n_splits=strat.n_splits)
+            folds = []
+            for train_idx, val_idx in tscv.split(X):
+                folds.append(
+                    (X.iloc[train_idx], X.iloc[val_idx], y.iloc[train_idx], y.iloc[val_idx])
+                )
+            return folds
+
+        else:
+            raise ValueError(f"Unknown split strategy: {strat.strategy}")
+
+    def fit_transform_folds(
+        self,
+        df: pd.DataFrame,
+    ) -> list[tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, ColumnTransformer]]:
+        """Separate, split, and fit transformers strictly on train folds to prevent data leakage."""
+        X, y = self.prepare_data(df)
+        raw_folds = self.split_data(X, y)
+
+        transformed_folds = []
+        for X_tr_raw, X_val_raw, y_tr_raw, y_val_raw in raw_folds:
+            transformer = self.build_transformer()
+            # STRICT LEAK-FREE RULE: fit strictly on train fold
+            X_tr_proc = transformer.fit_transform(X_tr_raw)
+            # transform validation fold using fitted statistics
+            X_val_proc = transformer.transform(X_val_raw)
+
+            y_tr_arr = np.asarray(y_tr_raw)
+            y_val_arr = np.asarray(y_val_raw)
+
+            transformed_folds.append((X_tr_proc, X_val_proc, y_tr_arr, y_val_arr, transformer))
+
+        # Retain last transformer as primary fitted reference
+        if transformed_folds:
+            self.transformer = transformed_folds[-1][4]
+
+        return transformed_folds

ml_mcp/workers/runner.py

@@ -0,0 +1,250 @@
+"""Worker runner executing end-to-end ML model training and artifact persistence."""
+
+import asyncio
+import io
+import logging
+import time
+from typing import Any
+
+import pandas as pd
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ml_mcp.domain.value_objects import ExperimentSpec
+from ml_mcp.infrastructure.object_storage.s3 import S3StorageService, get_storage_service
+from ml_mcp.infrastructure.postgres.base import generate_uuid7
+from ml_mcp.infrastructure.postgres.models import (
+    ArtifactOrm,
+    ExperimentRunOrm,
+    MetricOrm,
+)
+from ml_mcp.infrastructure.postgres.repositories import (
+    ArtifactRepository,
+    DatasetRepository,
+    ExperimentRepository,
+    MetricRepository,
+    ModelRepository,
+)
+from ml_mcp.workers.preprocessing import PreprocessingPipeline
+from ml_mcp.workers.sandbox import WorkerSandbox
+from ml_mcp.workers.trainers import get_trainer
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerRunner:
+    """Orchestrates ML experiment execution, metric logging, and artifact persistence."""
+
+    def __init__(
+        self,
+        session: AsyncSession,
+        storage_service: S3StorageService | None = None,
+    ) -> None:
+        self.session = session
+        self.storage = storage_service or get_storage_service()
+        self.exp_repo = ExperimentRepository(session)
+        self.ds_repo = DatasetRepository(session)
+        self.model_repo = ModelRepository(session)
+        self.metric_repo = MetricRepository(session)
+        self.artifact_repo = ArtifactRepository(session)
+
+    async def execute_experiment(
+        self,
+        experiment_id: str,
+        tenant_id: str,
+    ) -> dict[str, Any]:
+        """Execute experiment end-to-end within isolated sandbox."""
+        start_time = time.time()
+
+        # 1. Fetch experiment record
+        exp = await self.exp_repo.get_experiment(tenant_id, experiment_id)
+        if not exp:
+            raise ValueError(f"Experiment '{experiment_id}' not found for tenant '{tenant_id}'")
+
+        spec = ExperimentSpec(**exp.spec_json)
+
+        # 2. Mark running & create run record
+        await self.exp_repo.update_status(tenant_id, experiment_id, "RUNNING")
+        run_id = generate_uuid7()
+        run_orm = ExperimentRunOrm(
+            id=run_id,
+            experiment_id=experiment_id,
+            run_number=len(exp.runs) + 1,
+            status="RUNNING",
+            worker_id="worker-cpu-01",
+        )
+        await self.exp_repo.create_run(run_orm)
+
+        sandbox = WorkerSandbox(
+            timeout_seconds=spec.resource_policy.max_runtime_seconds,
+            max_memory_mb=spec.resource_policy.max_memory_mb,
+        )
+
+        try:
+            result = await sandbox.execute(self._execute_training_job, exp, spec, run_id)
+            duration = time.time() - start_time
+            await self.exp_repo.update_run(
+                run_id, "SUCCEEDED", duration_seconds=duration, end_time=True
+            )
+            await self.exp_repo.update_status(tenant_id, experiment_id, "SUCCEEDED")
+            return result
+
+        except Exception as exc:
+            duration = time.time() - start_time
+            logger.error("Experiment %s execution failed: %s", experiment_id, exc)
+            await self.exp_repo.update_run(
+                run_id, "FAILED", duration_seconds=duration, failure_reason=str(exc), end_time=True
+            )
+            await self.exp_repo.update_status(tenant_id, experiment_id, "FAILED")
+            raise
+
+    async def _execute_training_job(
+        self,
+        exp: Any,
+        spec: ExperimentSpec,
+        run_id: str,
+    ) -> dict[str, Any]:
+        tenant_id = exp.tenant_id
+        experiment_id = exp.id
+
+        # 1. Download dataset Parquet from S3
+        dsv = await self.ds_repo.get_version_by_id(tenant_id, spec.dataset_version_id)
+        if not dsv:
+            raise ValueError(f"DatasetVersion '{spec.dataset_version_id}' not found")
+
+        dataset_key = self.storage.get_dataset_key(
+            tenant_id, dsv.dataset_id, dsv.version, "data.parquet"
+        )
+        dataset_bytes = await asyncio.to_thread(self.storage.get_object, dataset_key)
+        df = await asyncio.to_thread(pd.read_parquet, io.BytesIO(dataset_bytes))
+
+        # 2. Preprocess features and create leak-free splits
+        prep = PreprocessingPipeline(
+            target_column=spec.target_column,
+            task_type=spec.task_type,
+            feature_columns=spec.feature_columns,
+            split_strategy=spec.split_strategy,
+        )
+        folds = await asyncio.to_thread(prep.fit_transform_folds, df)
+        X_tr, X_val, y_tr, y_val, fitted_transformer = folds[0]
+
+        # 3. Fetch model metadata and instantiate trainer
+        model_version = exp.model_version
+        model_family = model_version.model.family
+        trainer = get_trainer(model_family)
+
+        # 4. Train model
+        train_result = await asyncio.to_thread(
+            trainer.train,
+            X_tr=X_tr,
+            y_tr=y_tr,
+            X_val=X_val,
+            y_val=y_val,
+            task_type=spec.task_type,
+            hyperparameters=spec.hyperparameters,
+            random_seed=spec.random_seed,
+            max_cpu_cores=spec.resource_policy.max_cpu_cores,
+        )
+
+        # 5. Save model artifact to S3
+        model_key = self.storage.get_experiment_artifact_key(
+            tenant_id, experiment_id, "model", "model.joblib"
+        )
+        model_bytes = await asyncio.to_thread(train_result.serialize_artifact)
+        model_uri, model_hash = await asyncio.to_thread(
+            self.storage.put_object,
+            key=model_key,
+            data=model_bytes,
+            content_type="application/octet-stream",
+        )
+        await self.artifact_repo.record_artifact(
+            ArtifactOrm(
+                id=generate_uuid7(),
+                tenant_id=tenant_id,
+                experiment_id=experiment_id,
+                artifact_type="model",
+                storage_uri=model_uri,
+                size_bytes=len(train_result.model_artifact_bytes),
+                content_hash=model_hash,
+                metadata_json={"model_family": model_family, "task_type": spec.task_type.value},
+            )
+        )
+
+        # 6. Save predictions to S3
+        pred_df = pd.DataFrame(
+            {
+                "y_true": y_val,
+                "y_pred": train_result.val_predictions,
+            }
+        )
+        if train_result.val_probabilities is not None:
+            if train_result.val_probabilities.ndim == 1:
+                pred_df["y_proba"] = train_result.val_probabilities
+            else:
+                for idx in range(train_result.val_probabilities.shape[1]):
+                    pred_df[f"y_proba_class_{idx}"] = train_result.val_probabilities[:, idx]
+
+        pred_buf = io.BytesIO()
+        await asyncio.to_thread(pred_df.to_parquet, pred_buf, index=False)
+        pred_bytes = pred_buf.getvalue()
+
+        pred_key = self.storage.get_experiment_artifact_key(
+            tenant_id, experiment_id, "predictions", "predictions.parquet"
+        )
+        pred_uri, pred_hash = await asyncio.to_thread(
+            self.storage.put_object,
+            key=pred_key,
+            data=pred_bytes,
+            content_type="application/vnd.apache.parquet",
+        )
+        await self.artifact_repo.record_artifact(
+            ArtifactOrm(
+                id=generate_uuid7(),
+                tenant_id=tenant_id,
+                experiment_id=experiment_id,
+                artifact_type="predictions",
+                storage_uri=pred_uri,
+                size_bytes=len(pred_bytes),
+                content_hash=pred_hash,
+                metadata_json={"row_count": len(pred_df)},
+            )
+        )
+
+        # 7. Record metrics in DB
+        metric_records: list[MetricOrm] = []
+        for metric_name, val in train_result.validation_metrics.items():
+            if isinstance(val, (int, float)):
+                metric_records.append(
+                    MetricOrm(
+                        id=generate_uuid7(),
+                        experiment_id=experiment_id,
+                        run_id=run_id,
+                        metric_name=metric_name,
+                        metric_value=float(val),
+                        split="validation",
+                    )
+                )
+
+        for metric_name, val in train_result.train_metrics.items():
+            if isinstance(val, (int, float)):
+                metric_records.append(
+                    MetricOrm(
+                        id=generate_uuid7(),
+                        experiment_id=experiment_id,
+                        run_id=run_id,
+                        metric_name=metric_name,
+                        metric_value=float(val),
+                        split="train",
+                    )
+                )
+
+        await self.metric_repo.record_metrics(metric_records)
+
+        return {
+            "experiment_id": experiment_id,
+            "status": "SUCCEEDED",
+            "metrics": train_result.validation_metrics,
+            "artifacts": {
+                "model_uri": model_uri,
+                "predictions_uri": pred_uri,
+            },
+        }

ml_mcp/workers/sandbox.py

@@ -0,0 +1,35 @@
+"""Worker sandboxing and resource boundary enforcement."""
+
+import asyncio
+import logging
+from collections.abc import Callable, Coroutine
+from typing import Any, TypeVar
+
+from ml_mcp.domain.errors import ResourceLimitExceededError
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class WorkerSandbox:
+    """Safeguards worker execution with wall-clock timeouts and error containment."""
+
+    def __init__(self, timeout_seconds: int = 1800, max_memory_mb: int = 4096) -> None:
+        self.timeout_seconds = timeout_seconds
+        self.max_memory_mb = max_memory_mb
+
+    async def execute(
+        self,
+        coro_fn: Callable[..., Coroutine[Any, Any, T]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> T:
+        """Execute async function wrapped with wall-clock timeout."""
+        try:
+            return await asyncio.wait_for(coro_fn(*args, **kwargs), timeout=self.timeout_seconds)
+        except TimeoutError as exc:
+            logger.error("Worker execution timed out after %d seconds", self.timeout_seconds)
+            raise ResourceLimitExceededError(
+                "execution_time", f"{self.timeout_seconds}s", "timeout"
+            ) from exc

ml_mcp/workers/trainers/__init__.py

@@ -0,0 +1,52 @@
+"""Model trainers registry and factory."""
+
+from ml_mcp.domain.value_objects import ModelFamily
+from ml_mcp.workers.trainers.base import BaseModelTrainer, TrainingResult
+from ml_mcp.workers.trainers.catboost_model import CatBoostTrainer
+from ml_mcp.workers.trainers.lightgbm_model import LightGBMTrainer
+from ml_mcp.workers.trainers.sklearn_models import (
+    LinearSVMTrainer,
+    LogisticRegressionTrainer,
+    MLPTrainer,
+    RandomForestTrainer,
+)
+from ml_mcp.workers.trainers.xgboost_model import XGBoostTrainer
+
+TRAINER_MAP: dict[ModelFamily | str, type[BaseModelTrainer]] = {
+    ModelFamily.LOGISTIC_REGRESSION: LogisticRegressionTrainer,
+    ModelFamily.LOGISTIC_REGRESSION.value: LogisticRegressionTrainer,
+    ModelFamily.RANDOM_FOREST: RandomForestTrainer,
+    ModelFamily.RANDOM_FOREST.value: RandomForestTrainer,
+    ModelFamily.XGBOOST: XGBoostTrainer,
+    ModelFamily.XGBOOST.value: XGBoostTrainer,
+    ModelFamily.LIGHTGBM: LightGBMTrainer,
+    ModelFamily.LIGHTGBM.value: LightGBMTrainer,
+    ModelFamily.CATBOOST: CatBoostTrainer,
+    ModelFamily.CATBOOST.value: CatBoostTrainer,
+    ModelFamily.LINEAR_SVM: LinearSVMTrainer,
+    ModelFamily.LINEAR_SVM.value: LinearSVMTrainer,
+    ModelFamily.MLP: MLPTrainer,
+    ModelFamily.MLP.value: MLPTrainer,
+}
+
+
+def get_trainer(family: ModelFamily | str) -> BaseModelTrainer:
+    """Return concrete trainer instance for a given model family."""
+    trainer_cls = TRAINER_MAP.get(family)
+    if not trainer_cls:
+        raise ValueError(f"No trainer registered for model family '{family}'")
+    return trainer_cls()
+
+
+__all__ = [
+    "BaseModelTrainer",
+    "TrainingResult",
+    "LogisticRegressionTrainer",
+    "RandomForestTrainer",
+    "LinearSVMTrainer",
+    "MLPTrainer",
+    "XGBoostTrainer",
+    "LightGBMTrainer",
+    "CatBoostTrainer",
+    "get_trainer",
+]

ml_mcp/workers/trainers/base.py

@@ -0,0 +1,52 @@
+"""Base trainer contract and training result container."""
+
+import io
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+import joblib
+import numpy as np
+
+from ml_mcp.domain.value_objects import TaskType
+
+
+@dataclass
+class TrainingResult:
+    """Artifacts, predictions, and metrics emitted from model training."""
+
+    model: Any
+    task_type: TaskType
+    train_metrics: dict[str, Any]
+    validation_metrics: dict[str, Any]
+    feature_importances: dict[str, float]
+    val_predictions: np.ndarray
+    val_probabilities: np.ndarray | None
+    model_artifact_bytes: bytes
+
+    def serialize_artifact(self) -> bytes:
+        """Serialize model object using joblib."""
+        if self.model_artifact_bytes:
+            return self.model_artifact_bytes
+        bio = io.BytesIO()
+        joblib.dump(self.model, bio)
+        return bio.getvalue()
+
+
+class BaseModelTrainer(ABC):
+    """Abstract base class for all approved tabular model trainers."""
+
+    @abstractmethod
+    def train(
+        self,
+        X_tr: np.ndarray,
+        y_tr: np.ndarray,
+        X_val: np.ndarray,
+        y_val: np.ndarray,
+        task_type: TaskType,
+        hyperparameters: dict[str, Any],
+        random_seed: int = 42,
+        max_cpu_cores: int | None = None,
+    ) -> TrainingResult:
+        """Execute model training, evaluate metrics, and package artifacts."""
+        pass

ml_mcp/workers/trainers/catboost_model.py

@@ -0,0 +1,85 @@
+"""CatBoost model trainer."""
+
+import io
+from typing import Any
+
+import catboost as cb
+import joblib
+import numpy as np
+
+from ml_mcp.domain.value_objects import TaskType
+from ml_mcp.workers.metrics import MetricsCalculator
+from ml_mcp.workers.trainers.base import BaseModelTrainer, TrainingResult
+
+
+class CatBoostTrainer(BaseModelTrainer):
+    """Trainer for CatBoost gradient boosted trees."""
+
+    def train(
+        self,
+        X_tr: np.ndarray,
+        y_tr: np.ndarray,
+        X_val: np.ndarray,
+        y_val: np.ndarray,
+        task_type: TaskType,
+        hyperparameters: dict[str, Any],
+        random_seed: int = 42,
+        max_cpu_cores: int | None = None,
+    ) -> TrainingResult:
+        thread_count = (
+            max_cpu_cores if max_cpu_cores is not None else hyperparameters.get("thread_count", -1)
+        )
+        params = {
+            "random_seed": random_seed,
+            "iterations": hyperparameters.get("iterations", 100),
+            "depth": hyperparameters.get("depth", 6),
+            "learning_rate": hyperparameters.get("learning_rate", 0.1),
+            "thread_count": thread_count,
+            "verbose": False,
+            "allow_writing_files": False,
+        }
+
+        if task_type in (TaskType.BINARY_CLASSIFICATION, TaskType.MULTICLASS_CLASSIFICATION):
+            model = cb.CatBoostClassifier(**params)
+            model.fit(X_tr, y_tr, eval_set=(X_val, y_val), verbose=False)
+
+            train_pred = model.predict(X_tr)
+            train_proba = model.predict_proba(X_tr)
+            val_pred = model.predict(X_val)
+            val_proba = model.predict_proba(X_val)
+
+            train_metrics = MetricsCalculator.compute_classification_metrics(
+                y_tr, train_pred, train_proba, task_type
+            )
+            val_metrics = MetricsCalculator.compute_classification_metrics(
+                y_val, val_pred, val_proba, task_type
+            )
+        else:
+            model = cb.CatBoostRegressor(**params)
+            model.fit(X_tr, y_tr, eval_set=(X_val, y_val), verbose=False)
+
+            train_pred = model.predict(X_tr)
+            val_pred = model.predict(X_val)
+            val_proba = None
+
+            train_metrics = MetricsCalculator.compute_regression_metrics(y_tr, train_pred)
+            val_metrics = MetricsCalculator.compute_regression_metrics(y_val, val_pred)
+
+        feat_imp: dict[str, float] = {}
+        if hasattr(model, "get_feature_importance"):
+            importances = model.get_feature_importance()
+            feat_imp = {f"feature_{i}": float(v) for i, v in enumerate(importances)}
+
+        bio = io.BytesIO()
+        joblib.dump(model, bio)
+
+        return TrainingResult(
+            model=model,
+            task_type=task_type,
+            train_metrics=train_metrics,
+            validation_metrics=val_metrics,
+            feature_importances=feat_imp,
+            val_predictions=val_pred,
+            val_probabilities=val_proba,
+            model_artifact_bytes=bio.getvalue(),
+        )