opsml.registry.registry

  1# pylint: disable=protected-access
  2# Copyright (c) Shipt, Inc.
  3# This source code is licensed under the MIT license found in the
  4# LICENSE file in the root directory of this source tree.
  5import textwrap
  6from typing import Any, Dict, List, Optional, Type, Union
  7
  8from opsml.cards import ArtifactCard, CardInfo
  9from opsml.data import DataInterface
 10from opsml.helpers.logging import ArtifactLogger
 11from opsml.helpers.utils import clean_string
 12from opsml.model import ModelInterface
 13from opsml.registry.backend import _set_registry
 14from opsml.registry.semver import VersionType
 15from opsml.storage.card_loader import CardLoader
 16from opsml.types import CommonKwargs, RegistryType
 17
 18logger = ArtifactLogger.get_logger()
 19
 20
 21class CardRegistry:
 22    def __init__(self, registry_type: Union[RegistryType, str]):
 23        """
 24        Interface for connecting to any of the ArtifactCard registries
 25
 26        Args:
 27            registry_type:
 28                Type of card registry to create
 29            settings:
 30                Storage settings
 31
 32        Returns:
 33            Instantiated connection to specific Card registry
 34
 35        Example:
 36            data_registry = CardRegistry(RegistryType.DATA)
 37            data_registry.list_cards()
 38
 39            or
 40            data_registry = CardRegistry("data")
 41            data_registry.list_cards()
 42        """
 43
 44        _registry_type = (
 45            registry_type if isinstance(registry_type, RegistryType) else RegistryType.from_str(registry_type)
 46        )
 47
 48        self._registry = _set_registry(_registry_type)
 49        self.table_name = self._registry.table_name
 50
 51    @property
 52    def registry_type(self) -> RegistryType:
 53        "Registry type for card registry"
 54        return self._registry.registry_type
 55
 56    def list_cards(
 57        self,
 58        uid: Optional[str] = None,
 59        name: Optional[str] = None,
 60        repository: Optional[str] = None,
 61        version: Optional[str] = None,
 62        tags: Optional[Dict[str, str]] = None,
 63        info: Optional[CardInfo] = None,
 64        max_date: Optional[str] = None,
 65        limit: Optional[int] = None,
 66        ignore_release_candidates: bool = False,
 67    ) -> List[Dict[str, Any]]:
 68        """Retrieves records from registry
 69
 70        Args:
 71            name:
 72                Card name
 73            repository:
 74                Repository associated with card
 75            version:
 76                Optional version number of existing data. If not specified, the
 77                most recent version will be used
 78            tags:
 79                Dictionary of key, value tags to search for
 80            uid:
 81                Unique identifier for Card. If present, the uid takes precedence
 82            max_date:
 83                Max date to search. (e.g. "2023-05-01" would search for cards up to and including "2023-05-01")
 84            limit:
 85                Places a limit on result list. Results are sorted by SemVer
 86            info:
 87                CardInfo object. If present, the info object takes precedence
 88            ignore_release_candidates:
 89                If True, ignores release candidates
 90
 91        Returns:
 92            pandas dataframe of records or list of dictionaries
 93        """
 94
 95        if info is not None:
 96            name = name or info.name
 97            repository = repository or info.repository
 98            uid = uid or info.uid
 99            version = version or info.version
100            tags = tags or info.tags
101
102        if name is not None:
103            name = name.lower()
104
105        if repository is not None:
106            repository = repository.lower()
107
108        if all(not bool(var) for var in [name, repository, version, uid, tags]):
109            limit = limit or 25
110
111        card_list = self._registry.list_cards(
112            uid=uid,
113            name=name,
114            repository=repository,
115            version=version,
116            max_date=max_date,
117            limit=limit,
118            tags=tags,
119            ignore_release_candidates=ignore_release_candidates,
120        )
121
122        return card_list
123
124    def load_card(
125        self,
126        name: Optional[str] = None,
127        repository: Optional[str] = None,
128        uid: Optional[str] = None,
129        tags: Optional[Dict[str, str]] = None,
130        version: Optional[str] = None,
131        info: Optional[CardInfo] = None,
132        ignore_release_candidates: bool = False,
133        interface: Optional[Union[Type[ModelInterface], Type[DataInterface]]] = None,
134    ) -> ArtifactCard:
135        """Loads a specific card
136
137        Args:
138            name:
139                Optional Card name
140            uid:
141                Unique identifier for card. If present, the uid takes
142                precedence.
143            tags:
144                Optional tags associated with model.
145            repository:
146                Optional repository associated with card
147            version:
148                Optional version number of existing data. If not specified, the
149                most recent version will be used
150            info:
151                Optional CardInfo object. If present, the info takes precedence
152            ignore_release_candidates:
153                If True, ignores release candidates
154            interface:
155                Optional interface to use for loading card. This is required for when using
156                subclassed interfaces.
157
158        Returns
159            ArtifactCard
160        """
161
162        # find better way to do this later
163        if info is not None:
164            name = name or info.name
165            uid = uid or info.uid
166            version = version or info.version
167            tags = tags or info.tags
168
169        name = clean_string(name)
170
171        records = self.list_cards(
172            uid=uid,
173            name=name,
174            repository=repository,
175            version=version,
176            tags=tags,
177            ignore_release_candidates=ignore_release_candidates,
178            limit=1,
179        )
180
181        return CardLoader(
182            card_args=records[0],
183            registry_type=self.registry_type,
184        ).load_card(interface=interface)
185
186    def register_card(
187        self,
188        card: ArtifactCard,
189        version_type: Union[VersionType, str] = VersionType.MINOR,
190        pre_tag: str = "rc",
191        build_tag: str = "build",
192    ) -> None:
193        """
194        Adds a new `Card` record to registry. Registration will be skipped if the card already exists.
195
196        Args:
197            card:
198                card to register
199            version_type:
200                Version type for increment. Options are "major", "minor" and
201                "patch". Defaults to "minor".
202            pre_tag:
203                pre-release tag to add to card version
204            build_tag:
205                build tag to add to card version
206        """
207
208        _version_type = version_type if isinstance(version_type, VersionType) else VersionType.from_str(version_type)
209
210        if card.uid is not None and card.version != CommonKwargs.BASE_VERSION.value:
211            logger.info(
212                textwrap.dedent(
213                    f"""
214                Card {card.uid} already exists. Skipping registration. If you'd like to register
215                a new card, please instantiate a new Card object. If you'd like to update the
216                existing card, please use the update_card method.
217                """
218                )
219            )
220
221        else:
222            self._registry.register_card(
223                card=card,
224                version_type=_version_type,
225                pre_tag=pre_tag,
226                build_tag=build_tag,
227            )
228
229    def update_card(self, card: ArtifactCard) -> None:
230        """
231        Update an artifact card based on current registry
232
233        Args:
234            card:
235                Card to register
236        """
237        return self._registry.update_card(card=card)
238
239    def query_value_from_card(self, uid: str, columns: List[str]) -> Dict[str, Any]:
240        """
241        Query column values from a specific Card
242
243        Args:
244            uid:
245                Uid of Card
246            columns:
247                List of columns to query
248
249        Returns:
250            Dictionary of column, values pairs
251        """
252        results = self._registry.list_cards(uid=uid)[0]
253        return {col: results[col] for col in columns}
254
255    def delete_card(self, card: ArtifactCard) -> None:
256        """
257        Delete a specific Card
258
259        Args:
260            card:
261                Card to delete
262        """
263        return self._registry.delete_card(card)
264
265
266class CardRegistries:
267    def __init__(self) -> None:
268        """Instantiates class that contains all registries"""
269
270        self.data = CardRegistry(registry_type=RegistryType.DATA)
271        self.model = CardRegistry(registry_type=RegistryType.MODEL)
272        self.run = CardRegistry(registry_type=RegistryType.RUN)
273        self.pipeline = CardRegistry(registry_type=RegistryType.PIPELINE)
274        self.project = CardRegistry(registry_type=RegistryType.PROJECT)
275        self.audit = CardRegistry(registry_type=RegistryType.AUDIT)