opsml.cards.base

  1# Copyright (c) Shipt, Inc.
  2# This source code is licensed under the MIT license found in the
  3# LICENSE file in the root directory of this source tree.
  4import os
  5from pathlib import Path
  6from typing import Any, Dict, Optional
  7
  8from pydantic import BaseModel, ConfigDict, model_validator
  9
 10from opsml.helpers.logging import ArtifactLogger
 11from opsml.helpers.utils import clean_string, validate_name_repository_pattern
 12from opsml.settings.config import config
 13from opsml.types import CardInfo, CommonKwargs, RegistryTableNames
 14
 15logger = ArtifactLogger.get_logger()
 16
 17
 18class ArtifactCard(BaseModel):
 19    """Base pydantic class for artifact cards"""
 20
 21    model_config = ConfigDict(
 22        arbitrary_types_allowed=True,
 23        validate_assignment=False,
 24        validate_default=True,
 25    )
 26
 27    name: str = CommonKwargs.UNDEFINED.value
 28    repository: str = CommonKwargs.UNDEFINED.value
 29    contact: str = CommonKwargs.UNDEFINED.value
 30    version: str = CommonKwargs.BASE_VERSION.value
 31    uid: Optional[str] = None
 32    info: Optional[CardInfo] = None
 33    tags: Dict[str, str] = {}
 34
 35    @model_validator(mode="before")
 36    @classmethod
 37    def validate_args(cls, card_args: Dict[str, Any]) -> Dict[str, Any]:  # pylint: disable=arguments-renamed
 38        """Validate base args and Lowercase name and repository
 39
 40        Args:
 41            card_args:
 42                named args passed to card
 43
 44        Returns:
 45            validated card_args
 46        """
 47        card_info = card_args.get("info")
 48
 49        for key in ["name", "repository", "contact", "version", "uid"]:
 50            # check card args
 51            val = card_args.get(key)
 52
 53            # check card info
 54            if card_info is not None:
 55                val = val or getattr(card_info, key)
 56
 57            # check runtime env vars
 58            if val is None:
 59                val = os.environ.get(f"OPSML_RUNTIME_{key.upper()}")
 60
 61            if key in ["name", "repository"]:
 62                if val is not None:
 63                    val = clean_string(val)
 64
 65            if key == "version" and val is None:
 66                val = CommonKwargs.BASE_VERSION.value
 67
 68            card_args[key] = val
 69
 70        # need to check that name, repository and contact are set
 71        if not all(card_args[key] for key in ["name", "repository", "contact"]):
 72            raise ValueError("name, repository and contact must be set either as named arguments or through CardInfo")
 73
 74        # validate name and repository for pattern
 75        validate_name_repository_pattern(name=card_args["name"], repository=card_args["repository"])
 76
 77        return card_args
 78
 79    def create_registry_record(self) -> Dict[str, Any]:
 80        """Creates a registry record from self attributes"""
 81        raise NotImplementedError
 82
 83    def add_tag(self, key: str, value: str) -> None:
 84        self.tags[key] = str(value)
 85
 86    @property
 87    def uri(self) -> Path:
 88        """The base URI to use for the card and it's artifacts."""
 89        if self.version == CommonKwargs.BASE_VERSION.value:
 90            raise ValueError("Could not create card uri - version is not set")
 91
 92        assert self.repository is not None, "Repository must be set"
 93        assert self.name is not None, "Name must be set"
 94
 95        return Path(
 96            config.storage_root,
 97            RegistryTableNames.from_str(self.card_type).value,
 98            self.repository,
 99            self.name,
100            f"v{self.version}",
101        )
102
103    @property
104    def artifact_uri(self) -> Path:
105        """Returns the root URI to which artifacts associated with this card should be saved."""
106        return self.uri / "artifacts"
107
108    @property
109    def card_type(self) -> str:
110        raise NotImplementedError