fiddler-client 2.3.0.dev2__py3-none-any.whl → 2.4.0.dev1__py3-none-any.whl

fiddler3/entities/model.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import builtins
+import typing
+from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, Iterator
 from uuid import UUID
 
 import pandas as pd
@@ -12,20 +15,34 @@ from fiddler3.decorators import handle_api_error
 from fiddler3.entities.base import BaseEntity
 from fiddler3.entities.helpers import raise_not_found
 from fiddler3.entities.job import Job
-from fiddler3.schemas.job import JobCompact
-from fiddler3.schemas.model import ModelResponse
+from fiddler3.entities.model_artifact import ArtifactMixin
+from fiddler3.entities.model_surrogate import SurrogateMixin
+from fiddler3.entities.project import ProjectCompactMixin
+from fiddler3.entities.user import CreatedByMixin, UpdatedByMixin
+from fiddler3.entities.xai import XaiMixin
+from fiddler3.schemas.filter_query import OperatorType, QueryCondition, QueryRule
+from fiddler3.schemas.job import JobCompactResp
+from fiddler3.schemas.model import ModelResp
 from fiddler3.schemas.model_schema import Column, ModelSchema
 from fiddler3.schemas.model_spec import ModelSpec
 from fiddler3.schemas.model_task_params import ModelTaskParams
 from fiddler3.schemas.xai_params import XaiParams
-from fiddler3.utils.logger import get_logger
 from fiddler3.utils.model_schema_generator import SchemaGeneratorFactory
 
-logger = get_logger(__name__)
+if typing.TYPE_CHECKING:
+    from fiddler3.entities.dataset import Dataset
 
 
-class Model(BaseEntity):
-    def __init__(
+class Model(
+    BaseEntity,
+    ArtifactMixin,
+    CreatedByMixin,
+    ProjectCompactMixin,
+    UpdatedByMixin,
+    SurrogateMixin,
+    XaiMixin,
+):  # pylint: disable=too-many-ancestors
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         name: str,
         project_id: UUID,
@@ -40,9 +57,7 @@ class Model(BaseEntity):
         event_ts_format: str | None = None,
         xai_params: XaiParams | None = None,
     ) -> None:
-        """
-        Construct a model instance
-        """
+        """Construct a model instance"""
         self.name = name
         self.project_id = project_id
         self.schema = schema
@@ -69,20 +84,20 @@ class Model(BaseEntity):
         self.updated_at: datetime | None = None
 
         # Deserialized response object
-        self._resp: ModelResponse | None = None
+        self._resp: ModelResp | None = None
 
     @staticmethod
     def _get_url(id_: UUID | str | None = None) -> str:
-        """Get model resource/item url"""
+        """Get model resource/item url."""
         url = '/v3/models'
         return url if not id_ else f'{url}/{id_}'
 
     @classmethod
     def _from_dict(cls, data: dict) -> Model:
-        """Build entity object from the given dictionary"""
+        """Build entity object from the given dictionary."""
 
         # Deserialize the response
-        resp_obj = ModelResponse(**data)
+        resp_obj = ModelResp(**data)
 
         # Initialize
         instance = cls(
@@ -120,21 +135,69 @@ class Model(BaseEntity):
         instance._resp = resp_obj
         return instance
 
+    def _refresh(self, data: dict) -> None:
+        """Refresh the fields of this instance from the given response dictionary"""
+        # Deserialize the response
+        resp_obj = ModelResp(**data)
+
+        # Reset fields
+        self.schema = resp_obj.schema_
+        self.project_id = resp_obj.project.id
+
+        fields = [
+            'id',
+            'name',
+            'spec',
+            'input_type',
+            'task',
+            'task_params',
+            'description',
+            'event_id_col',
+            'event_ts_col',
+            'event_ts_format',
+            'xai_params',
+            'created_at',
+            'updated_at',
+            'artifact_status',
+            'artifact_files',
+            'input_cols',
+            'output_cols',
+            'target_cols',
+            'metadata_cols',
+            'decision_cols',
+            'is_binary_ranking_model',
+        ]
+        for field in fields:
+            setattr(self, field, getattr(resp_obj, field, None))
+
+        self._resp = resp_obj
+
     @classmethod
     @handle_api_error
     def get(cls, id_: UUID | str) -> Model:
-        """Get the model instance using model id"""
+        """Get the model instance using model id."""
         response = cls._client().get(url=cls._get_url(id_))
         return cls._from_response(response=response)
 
     @classmethod
     @handle_api_error
-    def from_name(cls, model_name: str, project_name: str) -> Model:
-        """Get the model instance using model name and project name"""
+    def from_name(cls, name: str, project_name: str) -> Model:
+        """Get the model instance using model name and project name."""
+        _filter = QueryCondition(
+            rules=[
+                QueryRule(field='name', operator=OperatorType.EQUAL, value=name),
+                QueryRule(
+                    field='project_name',
+                    operator=OperatorType.EQUAL,
+                    value=project_name,
+                ),
+            ]
+        )
+
         response = cls._client().get(
-            url=cls._get_url(),
-            params={'name': model_name, 'project_name': project_name},
+            url=cls._get_url(), params={'filter': _filter.json()}
         )
+
         if response.json()['data']['total'] == 0:
             raise_not_found('Model not found for the given identifier')
 
@@ -142,7 +205,7 @@ class Model(BaseEntity):
 
     @handle_api_error
     def create(self) -> Model:
-        """Create a new model"""
+        """Create a new model."""
         response = self._client().post(
             url=self._get_url(),
             data={
@@ -160,24 +223,68 @@ class Model(BaseEntity):
                 'xai_params': self.xai_params.dict(),
             },
         )
-        return self._from_response(response=response)
+        self._refresh_from_response(response=response)
+        return self
+
+    @handle_api_error
+    def update(self) -> None:
+        """Update an existing model."""
+        body: dict[str, Any] = {
+            'xai_params': self.xai_params.dict(),
+            'description': self.description,
+            'event_id_col': self.event_id_col,
+            'event_ts_col': self.event_ts_col,
+            'event_ts_format': self.event_ts_format,
+        }
+
+        response = self._client().patch(url=self._get_url(id_=self.id), data=body)
+        self._refresh_from_response(response=response)
+
+    @classmethod
+    @handle_api_error
+    def list(
+        cls,
+        project_id: UUID | None = None,
+    ) -> Iterator[ModelCompact]:
+        """
+        Get a list of all models with the given filters
+
+        :param project_id: Project identifier
+        :return: ModelCompact iterator
+        """
+        params = {}
+        if project_id:
+            params['project_id'] = project_id
+
+        for model in cls._paginate(url=cls._get_url(), params=params):
+            yield ModelCompact(id=model['id'], name=model['name'])
+
+    @property
+    def datasets(self) -> Iterator[Dataset]:
+        """Fetch all the datasets of this model"""
+        from fiddler3.entities.dataset import (  # pylint: disable=import-outside-toplevel
+            Dataset,
+        )
+
+        yield from Dataset.list(model_id=self.id)
 
     @classmethod
     @handle_api_error
     def generate_schema(
         cls,
-        source: pd.DataFrame | Path | list[dict[str, Any]] | str,
+        source: pd.DataFrame | Path | builtins.list[dict[str, Any]] | str,
         max_cardinality: int | None = None,
         sample_size: int | None = None,
         enrich: bool = True,
     ) -> ModelSchema:
         """
-        Generate model schema from the given data
+        Generate model schema from the given data.
+
         :param source: Data source - Dataframe or path to csv or parquet file.
-        :param max_cardinality: Max cardinality to detect categorical columns
-        :param sample_size: No. of samples to use for generating schema
-        :param enrich: Enrich the model schema client side by scanning all data
-        :return: Generated ModelSchema object
+        :param max_cardinality: Max cardinality to detect categorical columns.
+        :param sample_size: No. of samples to use for generating schema.
+        :param enrich: Enrich the model schema client side by scanning all data.
+        :return: Generated ModelSchema object.
         """
         schema_generator = SchemaGeneratorFactory.create(
             source=source,
@@ -190,9 +297,33 @@ class Model(BaseEntity):
 
     @handle_api_error
     def delete(self) -> Job:
-        """Delete a model and it's associated resources"""
+        """Delete a model and it's associated resources."""
         assert self.id is not None
         response = self._client().delete(url=self._get_url(id_=self.id))
 
-        job_compact = JobCompact(**response.json()['data']['job'])
+        job_compact = JobCompactResp(**response.json()['data']['job'])
         return Job.get(id_=job_compact.id)
+
+
+@dataclass
+class ModelCompact:
+    id: UUID
+    name: str
+
+    def fetch(self) -> Model:
+        """Fetch model instance"""
+        return Model.get(id_=self.id)
+
+
+class ModelCompactMixin:
+    @property
+    def model(self) -> ModelCompact:
+        """Model instance"""
+        response = getattr(self, '_resp', None)
+        if not response or not hasattr(response, 'model'):
+            raise AttributeError(
+                'This property is available only for objects generated from API '
+                'response.'
+            )
+
+        return ModelCompact(id=response.model.id, name=response.model.name)

fiddler3/entities/model_artifact.py

@@ -0,0 +1,350 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Callable
+from uuid import UUID
+
+from fiddler3.constants.common import (
+    CONTENT_TYPE_OCTET_STREAM_HEADER,
+    MULTI_PART_CHUNK_SIZE,
+)
+from fiddler3.decorators import handle_api_error
+from fiddler3.entities.job import Job
+from fiddler3.schemas.deployment_params import ArtifactType, DeploymentParams
+from fiddler3.schemas.job import JobCompactResp
+from fiddler3.schemas.model_artifact import ModelArtifactDeployMultiPartUploadResp
+from fiddler3.utils.logger import get_logger
+from fiddler3.utils.validations import validate_artifact_dir
+
+logger = get_logger(__name__)
+
+
+class ArtifactMixin:
+    id: UUID | None
+    _client: Callable
+
+    def _get_method(self, update: bool = False) -> Callable:
+        """Get HTTP method"""
+        return self._client().put if update else self._client().post
+
+    def _deploy_model_artifact(
+        self,
+        model_dir: str | Path,
+        deployment_params: DeploymentParams | None = None,
+        update: bool = False,
+    ) -> Job:
+        """
+        Upload and deploy model artifact for an existing model
+
+        :param model_dir: Model artifact directory
+        :param deployment_params: Model deployment parameters
+        :param update: Set True for updating artifact, False for adding artifact
+        :return: Async job
+        """
+        model_dir = Path(model_dir)
+        validate_artifact_dir(model_dir)
+
+        if (
+            deployment_params
+            and deployment_params.artifact_type.upper() == ArtifactType.SURROGATE
+        ):
+            raise ValueError(
+                f'{ArtifactType.SURROGATE} artifact_type is an invalid value for this '
+                f'method. Use {ArtifactType.PYTHON_PACKAGE} instead.'
+            )
+
+        with tempfile.TemporaryDirectory() as tmp:
+            # Archive model artifact directory
+            logger.info(
+                'Model[%s] - Tarring model artifact directory - %s',
+                self.id,
+                model_dir,
+            )
+            file_path = shutil.make_archive(
+                base_name=str(Path(tmp) / 'files'),
+                format='tar',
+                root_dir=str(model_dir),
+                base_dir='.',
+            )
+
+            logger.info(
+                'Model[%s] - Model artifact tar file created at %s',
+                self.id,
+                file_path,
+            )
+
+            # Choose deployer based on archive file size
+            if os.path.getsize(file_path) < MULTI_PART_CHUNK_SIZE:
+                job = self._artifact_deploy(
+                    file_path=Path(file_path),
+                    deployment_params=deployment_params,
+                    update=update,
+                )
+            else:
+                job = self._artifact_deploy_multipart(
+                    file_path=Path(file_path),
+                    deployment_params=deployment_params,
+                    update=update,
+                )
+
+        logger.info(
+            'Model[%s] - Submitted job (%s) for deploying model artifact',
+            self.id,
+            job.id,
+        )
+
+        return job
+
+    def _initialize_multi_part_upload(self, update: bool = False) -> str:
+        """
+        Initialize multi-part upload request
+
+        :return: Multi-part upload id
+        """
+        logger.info(
+            'Model[%s] - Initializing multi-part model upload',
+            self.id,
+        )
+
+        method = self._get_method(update)
+
+        response = method(url=f'/v3/models/{self.id}/deploy-artifact-multi-part-init')
+
+        response_data = response.json().get('data', {})
+
+        logger.info(
+            'Model[%s] - Multi-part model upload initialized',
+            self.id,
+        )
+        return response_data.get('upload_id', '')
+
+    def _upload_multi_part_chunk(
+        self, data: bytes, upload_id: str, part_number: int, update: bool
+    ) -> dict:
+        """
+        Upload data chunk
+
+        :param data: Data chunk
+        :param upload_id: Multi-part upload id
+        :param part_number: Chunk part number
+        :return: Part details
+        """
+        method = self._get_method(update)
+
+        response = method(
+            url=f'/v3/models/{self.id}/deploy-artifact-multi-part-upload',
+            params={
+                'upload_id': upload_id,
+                'part_number': part_number,
+            },
+            data=data,
+            headers=CONTENT_TYPE_OCTET_STREAM_HEADER,
+        )
+        response_data = ModelArtifactDeployMultiPartUploadResp(
+            **response.json()['data']
+        )
+
+        return {
+            'etag': response_data.etag,
+            'part_number': response_data.part_number,
+        }
+
+    def _complete_multi_part_upload(
+        self,
+        upload_id: str,
+        parts: list[dict],
+        update: bool,
+        deployment_params: DeploymentParams | None = None,
+    ) -> Job:
+        """
+        Complete multi-part upload request and deploy model artifact
+
+        :param upload_id: Multi-part upload id
+        :param parts: List of all parts
+        :param deployment_params: Deployment parameters
+        :return: Async job
+        """
+        method = self._get_method(update)
+        payload: dict[str, Any] = {
+            'upload_id': upload_id,
+            'parts': parts,
+        }
+
+        if deployment_params:
+            payload['deployment_params'] = deployment_params.dict(exclude_unset=True)
+
+        response = method(
+            url=f'/v3/models/{self.id}/deploy-artifact-multi-part-complete',
+            data=payload,
+        )
+
+        logger.info(
+            'Model[%s] - Multi-part model upload completed',
+            self.id,
+        )
+        job_compact = JobCompactResp(**response.json()['data']['job'])
+        return Job.get(id_=job_compact.id)
+
+    def _artifact_deploy_multipart(
+        self,
+        file_path: Path,
+        deployment_params: DeploymentParams | None = None,
+        update: bool = False,
+    ) -> Job:
+        """
+        Upload and deploy model artifact
+
+        :param file_path: Path to model artifact tar file
+        :param deployment_params: Model deployment parameters
+        :param update: Flag for add or update artifacts
+        :return: Async job uuid
+        """
+        # 1. Initialize multi-part upload request
+        upload_id = self._initialize_multi_part_upload(update)
+        part_number = 1
+        parts = []
+        file_size = os.path.getsize(file_path)
+        total_chunks = math.ceil(file_size / MULTI_PART_CHUNK_SIZE)
+
+        # 2. Chunk and upload
+        with open(file_path, 'rb') as f:
+            while True:
+                data = f.read(MULTI_PART_CHUNK_SIZE)
+                if not data:
+                    break
+
+                logger.info(
+                    'Model[%s] - Uploading multi-part chunk - %d/%d',
+                    self.id,
+                    part_number,
+                    total_chunks,
+                )
+
+                part = self._upload_multi_part_chunk(
+                    data=data,
+                    upload_id=upload_id,
+                    part_number=part_number,
+                    update=update,
+                )
+                parts.append(part)
+                logger.info(
+                    'Model[%s] - Uploaded multi-part chunk - %d/%d',
+                    self.id,
+                    part_number,
+                    total_chunks,
+                )
+                part_number += 1
+
+        # 3: Complete the upload
+        return self._complete_multi_part_upload(
+            upload_id=upload_id,
+            parts=parts,
+            update=update,
+            deployment_params=deployment_params,
+        )
+
+    def _artifact_deploy(
+        self,
+        file_path: Path,
+        deployment_params: DeploymentParams | None = None,
+        update: bool = False,
+    ) -> Job:
+        """Artifact deploy base method."""
+        method = self._get_method(update)
+        params = {}
+        if deployment_params:
+            params['deployment_params'] = json.dumps(
+                deployment_params.dict(exclude_unset=True)
+            )
+
+        with open(file_path, 'rb') as f:
+            data = f.read()
+            response = method(
+                url=f'/v3/models/{self.id}/deploy-artifact',
+                params=params,
+                data=data,
+                headers=CONTENT_TYPE_OCTET_STREAM_HEADER,
+            )
+
+        job_compact = JobCompactResp(**response.json()['data']['job'])
+        return Job.get(id_=job_compact.id)
+
+    @handle_api_error
+    def add_artifact(
+        self,
+        model_dir: str | Path,
+        deployment_params: DeploymentParams | None = None,
+    ) -> Job:
+        """
+        Upload and deploy model artifact.
+
+        :param model_dir: Path to model artifact tar file
+        :param deployment_params: Model deployment parameters
+        :return: Async job.
+        """
+        self._check_id_attributes()
+        job = self._deploy_model_artifact(
+            model_dir=model_dir, deployment_params=deployment_params
+        )
+        return job
+
+    @handle_api_error
+    def update_artifact(
+        self,
+        model_dir: str | Path,
+        deployment_params: DeploymentParams | None = None,
+    ) -> Job:
+        """
+        Update existing model artifact.
+
+        :param model_dir: Path to model artifact tar file
+        :param deployment_params: Model deployment parameters
+        :return: Async job.
+        """
+        self._check_id_attributes()
+        job = self._deploy_model_artifact(
+            model_dir=model_dir, deployment_params=deployment_params, update=True
+        )
+        return job
+
+    @handle_api_error
+    def download_artifact(
+        self,
+        output_dir: str | Path,
+    ) -> None:
+        """
+        Download existing model artifact.
+
+        :param output_dir: Path to download model artifact tar file
+        """
+        self._check_id_attributes()
+        output_dir = Path(output_dir)
+        if output_dir.exists():
+            raise ValueError(f'Output dir already exists {output_dir}')
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Download tar file
+            tar_file_path = os.path.join(tmp_dir, 'artifact.tar')
+
+            with self._client().get(
+                url=f'/v3/models/{self.id}/download-artifact'
+            ) as resp:
+                resp.raise_for_status()
+                with open(tar_file_path, mode='wb') as f:
+                    for chunk in resp.iter_content(chunk_size=8192):
+                        if chunk:
+                            f.write(chunk)
+            os.makedirs(output_dir, exist_ok=True)
+            shutil.unpack_archive(tar_file_path, extract_dir=output_dir, format='tar')
+
+    def _check_id_attributes(self) -> None:
+        if not self.id:
+            raise AttributeError(
+                'This method is available only for model object generated from '
+                'API response.'
+            )