fundamental-client 0.2.3__py3-none-any.whl

fundamental/services/feature_importance.py

@@ -0,0 +1,172 @@
+"""
+Feature importance services for NEXUS client.
+Handles feature importance computation with the NEXUS service.
+"""
+
+from typing import Optional
+
+import numpy as np
+
+from fundamental.clients.base import BaseClient
+from fundamental.constants import (
+    DEFAULT_POLLING_INTERVAL_SECONDS,
+    DEFAULT_SUBMIT_REQUEST_TIMEOUT_SECONDS,
+)
+from fundamental.exceptions import ServerError
+from fundamental.models import TaskStatus
+from fundamental.utils.data import (
+    XType,
+    api_call,
+    create_feature_importance_task_metadata,
+    download_result_from_url,
+    serialize_df_to_parquet_bytes,
+    upload_feature_importance_data,
+)
+from fundamental.utils.polling import wait_for_task_status
+
+
+def submit_feature_importance_task(
+    X: XType,
+    trained_model_id: str,
+    client: BaseClient,
+) -> str:
+    """
+    Submit a feature importance computation task without waiting for completion.
+
+    Parameters
+    ----------
+    X : XType
+        Input features for feature importance computation.
+    trained_model_id : str
+        The trained model ID.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    str
+        The task ID to use with poll_feature_importance_result.
+    """
+    X_serialized = serialize_df_to_parquet_bytes(data=X)
+
+    metadata = create_feature_importance_task_metadata(
+        trained_model_id=trained_model_id,
+        x_size=len(X_serialized),
+        client=client,
+    )
+
+    upload_feature_importance_data(
+        X_serialized=X_serialized,
+        metadata=metadata,
+        trained_model_id=trained_model_id,
+        client=client,
+    )
+
+    json_data = {
+        "trained_model_id": trained_model_id,
+        "request_id": metadata.request_id,
+        "timeout": client.config.feature_importance_timeout,
+    }
+
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_feature_importance_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_SUBMIT_REQUEST_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+    task_id: str = data["task_id"]
+    return task_id
+
+
+def remote_get_feature_importance(
+    X: XType,
+    trained_model_id: str,
+    client: BaseClient,
+) -> np.ndarray:
+    """
+    Get feature importance for a trained model.
+
+    Submits the task and waits for completion.
+
+    Parameters
+    ----------
+    X : XType
+        Input features for feature importance computation.
+    trained_model_id : str
+        The trained model ID.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    np.ndarray
+        Feature importance values.
+    """
+    task_id = submit_feature_importance_task(
+        X=X,
+        trained_model_id=trained_model_id,
+        client=client,
+    )
+
+    status_response = wait_for_task_status(
+        client=client,
+        status_url=f"{client.config.get_full_feature_importance_status_url()}/{trained_model_id}/{task_id}",
+        timeout=client.config.feature_importance_timeout,
+        polling_interval=DEFAULT_POLLING_INTERVAL_SECONDS,
+    )
+
+    if not status_response.result:
+        raise ServerError("Request failed: Internal Server Error")
+
+    downloaded_result = download_result_from_url(
+        download_url=status_response.result.download_url,
+        client=client,
+        timeout=client.config.download_feature_importance_result_timeout,
+    )
+
+    return np.array(downloaded_result)
+
+
+def poll_feature_importance_result(
+    task_id: str,
+    trained_model_id: str,
+    client: BaseClient,
+) -> Optional[np.ndarray]:
+    """
+    Check the status of a feature importance task.
+
+    Parameters
+    ----------
+    task_id : str
+        The task ID returned by submit_feature_importance_task.
+    trained_model_id : str
+        The trained model ID.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    Optional[np.ndarray]
+        Feature importance values if completed, None if still in progress.
+    """
+    status_response = wait_for_task_status(
+        client=client,
+        status_url=f"{client.config.get_full_feature_importance_status_url()}/{trained_model_id}/{task_id}",
+        timeout=client.config.feature_importance_timeout,
+        polling_interval=DEFAULT_POLLING_INTERVAL_SECONDS,
+        wait_for_completion=False,
+    )
+
+    if status_response.status == TaskStatus.SUCCESS:
+        if not status_response.result:
+            raise ServerError("Request failed: Internal Server Error")
+        downloaded_result = download_result_from_url(
+            download_url=status_response.result.download_url,
+            client=client,
+            timeout=client.config.download_feature_importance_result_timeout,
+        )
+        return np.array(downloaded_result)
+
+    return None

fundamental/services/inference.py

@@ -0,0 +1,283 @@
+"""
+Model inference services for NEXUS client.
+Handles fit and predict operations with the NEXUS service.
+"""
+
+import logging
+from typing import Literal, Optional
+
+import numpy as np
+from pydantic import BaseModel
+
+from fundamental.clients.base import BaseClient
+from fundamental.constants import (
+    DEFAULT_POLLING_INTERVAL_SECONDS,
+    DEFAULT_PREDICT_POLLING_REQUESTS_WITHOUT_DELAY,
+    DEFAULT_SUBMIT_REQUEST_TIMEOUT_SECONDS,
+)
+from fundamental.models import TaskStatus
+from fundamental.services.models import ModelsService
+from fundamental.utils.data import (
+    XType,
+    YType,
+    api_call,
+    create_fit_task_metadata,
+    create_predict_task_metadata,
+    download_result_from_url,
+    serialize_df_to_parquet_bytes,
+    upload_fit_data,
+    upload_predict_data,
+)
+from fundamental.utils.polling import wait_for_task_status
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteFitResponse(BaseModel):
+    trained_model_id: str
+    estimator_fields: dict
+
+
+class SubmitFitTaskResult(BaseModel):
+    """Result of submitting a fit task."""
+
+    task_id: str
+    trained_model_id: str
+
+
+def submit_fit_task(
+    X: XType,
+    y: YType,
+    task: Literal["classification", "regression"],
+    mode: Literal["quality", "speed"],
+    client: BaseClient,
+) -> SubmitFitTaskResult:
+    """
+    Submit a fit task without waiting for completion.
+
+    Parameters
+    ----------
+    X : XType
+        Training features.
+    y : YType
+        Training targets.
+    task : {"classification", "regression"}
+        Task type.
+    mode : {"quality", "speed"}
+        Model fit mode.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    SubmitFitTaskResult
+        Result containing task_id and trained_model_id.
+    """
+    X_serialized = serialize_df_to_parquet_bytes(data=X)
+    y_serialized = serialize_df_to_parquet_bytes(data=y)
+
+    metadata = create_fit_task_metadata(
+        x_train_size=len(X_serialized),
+        y_train_size=len(y_serialized),
+        client=client,
+    )
+
+    upload_fit_data(
+        X_serialized=X_serialized,
+        y_serialized=y_serialized,
+        metadata=metadata,
+        client=client,
+    )
+
+    json_data = {
+        "task": task,
+        "mode": mode,
+        "trained_model_id": metadata.trained_model_id,
+        "timeout": client.config.fit_timeout,
+    }
+
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_fit_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_SUBMIT_REQUEST_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+    task_id: str = data["task_id"]
+
+    return SubmitFitTaskResult(
+        task_id=task_id,
+        trained_model_id=metadata.trained_model_id,
+    )
+
+
+def poll_fit_result(
+    task_id: str,
+    trained_model_id: str,
+    client: BaseClient,
+) -> Optional[RemoteFitResponse]:
+    """
+    Check the status of a fit task.
+
+    Parameters
+    ----------
+    task_id : str
+        The task ID returned by submit_fit_task.
+    trained_model_id : str
+        The trained model ID from the submit result.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    Optional[RemoteFitResponse]
+        RemoteFitResponse with trained_model_id and estimator_fields if completed,
+        None if still in progress.
+    """
+    status_response = wait_for_task_status(
+        client=client,
+        status_url=f"{client.config.get_full_fit_status_url()}/{task_id}",
+        timeout=client.config.fit_timeout,
+        polling_interval=DEFAULT_POLLING_INTERVAL_SECONDS,
+        wait_for_completion=False,
+    )
+
+    if status_response.status == TaskStatus.SUCCESS:
+        logger.debug("Loading trained model metadata")
+        models_service = ModelsService(client=client)
+        loaded_model = models_service.load(trained_model_id=trained_model_id)
+
+        return RemoteFitResponse(
+            trained_model_id=trained_model_id,
+            estimator_fields=loaded_model.estimator_fields,
+        )
+
+    return None
+
+
+def remote_fit(
+    X: XType,
+    y: YType,
+    task: Literal["classification", "regression"],
+    mode: Literal["quality", "speed"],
+    client: BaseClient,
+) -> RemoteFitResponse:
+    """
+    Fit a model (blocking).
+
+    Submits the task and waits for completion.
+
+    Parameters
+    ----------
+    X : XType
+        Training features.
+    y : YType
+        Training targets.
+    task : {"classification", "regression"}
+        Task type.
+    mode : {"quality", "speed"}
+        Model fit mode.
+    client : BaseClient
+        The client instance.
+
+    Returns
+    -------
+    RemoteFitResponse
+        Service response with trained_model_id and estimator_fields.
+    """
+    submit_result = submit_fit_task(
+        X=X,
+        y=y,
+        task=task,
+        mode=mode,
+        client=client,
+    )
+
+    wait_for_task_status(
+        client=client,
+        status_url=f"{client.config.get_full_fit_status_url()}/{submit_result.task_id}",
+        timeout=client.config.fit_timeout,
+        polling_interval=DEFAULT_POLLING_INTERVAL_SECONDS,
+    )
+
+    logger.debug("Loading trained model metadata")
+    models_service = ModelsService(client=client)
+    loaded_model = models_service.load(trained_model_id=submit_result.trained_model_id)
+
+    return RemoteFitResponse(
+        trained_model_id=submit_result.trained_model_id,
+        estimator_fields=loaded_model.estimator_fields,
+    )
+
+
+def remote_predict(
+    X: XType,
+    output_type: Literal["preds", "probas"],
+    trained_model_id: str,
+    client: BaseClient,
+) -> np.ndarray:
+    """
+    Make predictions using a trained model identified by trained_model_id.
+
+    Parameters
+    ----------
+    X : XType
+        Input features for prediction.
+    output_type : {"preds", "probas"}
+        Output type.
+    trained_model_id : str
+        The model ID generated by after the fit operation.
+
+    Returns
+    -------
+    np.ndarray
+        Prediction results.
+    """
+    X_serialized = serialize_df_to_parquet_bytes(data=X)
+
+    metadata = create_predict_task_metadata(
+        trained_model_id=trained_model_id,
+        x_test_size=len(X_serialized),
+        client=client,
+    )
+
+    upload_predict_data(
+        X_serialized=X_serialized,
+        metadata=metadata,
+        trained_model_id=trained_model_id,
+        client=client,
+    )
+
+    json_data = {
+        "output_type": output_type,
+        "trained_model_id": trained_model_id,
+        "request_id": metadata.request_id,
+        "timeout": client.config.predict_timeout,
+    }
+
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_predict_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_SUBMIT_REQUEST_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+    task_id = data["task_id"]
+
+    status_response = wait_for_task_status(
+        client=client,
+        status_url=f"{client.config.get_full_predict_status_url()}/{trained_model_id}/{task_id}",
+        timeout=client.config.predict_timeout,
+        polling_interval=DEFAULT_POLLING_INTERVAL_SECONDS,
+        polling_requests_without_delay=DEFAULT_PREDICT_POLLING_REQUESTS_WITHOUT_DELAY,
+    )
+
+    preds = download_result_from_url(
+        download_url=status_response.result.download_url,  # type: ignore[union-attr]
+        client=client,
+        timeout=client.config.download_prediction_result_timeout,
+    )
+
+    return np.array(preds)

fundamental/services/models.py

@@ -0,0 +1,186 @@
+"""Model management services for NEXUS client."""
+
+import base64
+import logging
+from typing import Any, Dict
+
+from pydantic import BaseModel
+from typing_extensions import TypeAlias
+
+from fundamental.clients.base import BaseClient
+from fundamental.exceptions import ValidationError
+from fundamental.models import (
+    DeleteTrainedModelResponse,
+    TrainedModelMetadata,
+    UpdateAttributesResponse,
+)
+from fundamental.utils.http import api_call
+from fundamental.utils.safetensors_deserialize import load_estimator_fields_from_bytes
+
+logger = logging.getLogger(__name__)
+
+TrainedModelsListResponse: TypeAlias = list[str]
+
+
+class LoadedModelResponse(BaseModel):
+    estimator_fields: Dict[str, Any]
+
+
+class ModelsService:
+    """Service for model management operations."""
+
+    def __init__(self, client: BaseClient) -> None:
+        """Initialize the models service.
+
+        Args:
+            client: Client instance
+        """
+        self.client = client
+
+    def _validate_model_id(self, model_id: str) -> None:
+        if not model_id or not model_id.strip():
+            raise ValidationError("model_id cannot be empty. Please provide a valid model_id.")
+
+    def list(self) -> TrainedModelsListResponse:
+        """
+        List all trained models.
+
+        Returns
+        -------
+        TrainedModelsListResponse
+            List of trained model IDs.
+        """
+        response = api_call(
+            method="GET",
+            full_url=self.client.config.get_full_model_management_url(),
+            client=self.client,
+        )
+
+        res: list[str] = response.json()
+        return res
+
+    def delete(self, model_id: str) -> DeleteTrainedModelResponse:
+        """
+        Delete a specific trained model.
+
+        Parameters
+        ----------
+        model_id : str
+            The ID of the model to delete.
+
+        Returns
+        -------
+        DeleteTrainedModelResponse
+            Response from the deletion operation.
+
+        Raises
+        ------
+        ValidationError
+            If model_id is empty or invalid.
+        """
+        self._validate_model_id(model_id)
+
+        response = api_call(
+            method="DELETE",
+            full_url=f"{self.client.config.get_full_model_management_url()}/{model_id}",
+            client=self.client,
+        )
+
+        return DeleteTrainedModelResponse(**response.json())
+
+    def get(self, model_id: str) -> TrainedModelMetadata:
+        """
+        Get information about a specific trained model.
+
+        Parameters
+        ----------
+        model_id : str
+            The ID of the model to retrieve.
+
+        Returns
+        -------
+        TrainedModelMetadata
+            Model information dictionary.
+
+        Raises
+        ------
+        ValidationError
+            If model_id is empty or invalid.
+        """
+        self._validate_model_id(model_id)
+
+        response = api_call(
+            method="GET",
+            full_url=f"{self.client.config.get_full_model_management_url()}/{model_id}",
+            client=self.client,
+            max_retries=3,
+        )
+
+        return TrainedModelMetadata(**response.json())
+
+    def set_attributes(
+        self,
+        model_id: str,
+        attributes: dict[str, str],
+    ) -> UpdateAttributesResponse:
+        """
+        Set attributes for a specific trained model.
+
+        Parameters
+        ----------
+        model_id : str
+            The ID of the model to update.
+        attributes : dict[str, str]
+            The attributes to set.
+
+        Returns
+        -------
+        UpdateAttributesResponse
+            Response containing the updated attributes.
+
+        Raises
+        ------
+        ValidationError
+            If model_id is empty or invalid.
+        """
+        self._validate_model_id(model_id)
+
+        response = api_call(
+            method="PATCH",
+            full_url=f"{self.client.config.get_full_model_management_url()}/{model_id}/attributes",
+            client=self.client,
+            json={"attributes": attributes},
+        )
+
+        return UpdateAttributesResponse(**response.json())
+
+    def load(
+        self,
+        trained_model_id: str,
+    ) -> LoadedModelResponse:
+        """
+        Load a specific trained model.
+
+        Parameters
+        ----------
+        trained_model_id : str
+            The ID of the model to load.
+
+        Returns
+        -------
+        LoadedModelResponse
+            Model information dictionary.
+        """
+        self._validate_model_id(trained_model_id)
+
+        response = api_call(
+            method="GET",
+            full_url=f"{self.client.config.get_full_model_management_url()}/{trained_model_id}/load_model",
+            client=self.client,
+        )
+        response_json = response.json()
+        raw_bytes = base64.b64decode(response_json["estimator_fields"], validate=True)
+        estimator_fields = load_estimator_fields_from_bytes(raw_bytes)
+        return LoadedModelResponse(
+            estimator_fields=estimator_fields,
+        )