1# pylint: disable=invalid-envvar-value
  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.
  5# pylint: disable=protected-access
  6
  7from contextlib import contextmanager
  8from typing import Any, Dict, Iterator, List, Optional, Union, cast
  9
 10from opsml.cards.base import ArtifactCard
 11from opsml.cards.project import ProjectCard
 12from opsml.cards.run import RunCard
 13from opsml.helpers.logging import ArtifactLogger
 14from opsml.projects._run_manager import ActiveRunException, _RunManager
 15from opsml.projects.active_run import ActiveRun, CardHandler
 16from opsml.projects.types import ProjectInfo
 17from opsml.registry import CardRegistries
 18from opsml.registry.sql.base.client import ClientProjectCardRegistry
 19from opsml.types import CardInfo, CardType, Metric, Metrics, Param, Params
 20
 21logger = ArtifactLogger.get_logger()
 22
 23
 24class _ProjectRegistrar:
 25    def __init__(self, project_info: ProjectInfo):
 26        self._project_info = project_info
 27        self.registries = CardRegistries()
 28
 29    @property
 30    def project_registry(self) -> ClientProjectCardRegistry:
 31        # both server and client registries are the same for methods
 32        return cast(ClientProjectCardRegistry, self.registries.project._registry)
 33
 34    def register_project(self) -> int:
 35        """
 36        Checks if the project name exists in the project registry. A ProjectCard is created if it
 37        doesn't exist.
 38
 39        Returns:
 40            project_id: int
 41        """
 42
 43        card = ProjectCard(
 44            name=self._project_info.name,
 45            repository=self._project_info.repository,
 46            contact=self._project_info.contact,
 47        )
 48        self.registries.project.register_card(card=card)
 49
 50        return card.project_id
 51
 52
 53class OpsmlProject:
 54    def __init__(self, info: ProjectInfo):
 55        """
 56        Instantiates a project which creates cards, metrics and parameters to
 57        the opsml registry via a "run" object.
 58
 59        If info.run_id is set, that run_id will be loaded as read only. In read
 60        only mode, you can retrieve cards, metrics, and parameters, however you
 61        cannot write new data. If you wish to record data/create a new run, you will
 62        need to enter the run context.
 63
 64        In order to create new cards, you need to create a run using the `run`
 65        context manager.
 66
 67        Example:
 68
 69            project: OpsmlProject = OpsmlProject(
 70                ProjectInfo(
 71                    name="test-project",
 72                    # If run_id is omitted, a new run is created.
 73                    run_id="123ab123kaj8u8naskdfh813",
 74                )
 75            )
 76            # the project is in "read only" mode. all read operations will work
 77            for k, v in project.parameters:
 78                logger.info("{} = {}", k, v)
 79
 80            # creating a project run
 81            with project.run() as run:
 82                # Now that the run context is entered, it's in read/write mode
 83                # You can write cards, parameters, and metrics to the project.
 84                run.log_parameter(key="my_param", value="12.34")
 85
 86        Args:
 87            info:
 88                Run information. if a run_id is given, that run is set
 89                as the project's current run.
 90        """
 91        # Set the run manager and project_id (creates ProjectCard if project doesn't exist)
 92        registrar = _ProjectRegistrar(project_info=info)
 93
 94        # get project id or register new project
 95        info.project_id = registrar.register_project()
 96
 97        # crete run manager
 98        self._run_mgr = _RunManager(project_info=info, registries=registrar.registries)
 99
100    @property
101    def run_id(self) -> str:
102        """Current run id associated with project"""
103        if self._run_mgr.run_id is not None:
104            return self._run_mgr.run_id
105        raise ValueError("Run id not set for current project")
106
107    @run_id.setter
108    def run_id(self, run_id: str) -> None:
109        """Set the run_id to use with the active project"""
110        self._run_mgr.run_id = run_id
111
112    @property
113    def project_id(self) -> int:
114        return self._run_mgr.project_id
115
116    @property
117    def project_name(self) -> str:
118        return self._run_mgr._project_info.name  # pylint: disable=protected-access
119
120    @contextmanager
121    def run(self, run_name: Optional[str] = None) -> Iterator[ActiveRun]:
122        """
123        Starts a new run for the project
124
125        Args:
126            run_name:
127                Optional run name
128        """
129
130        try:
131            yield self._run_mgr.start_run(run_name=run_name)  # self._run_mgr.active_run
132
133        except ActiveRunException as error:
134            logger.error("Run already active. Ending run.")
135            raise error
136
137        except Exception as error:
138            logger.error("Error encountered. Ending run. {}", error)
139            self._run_mgr.end_run()
140            raise error
141
142        self._run_mgr.end_run()
143
144    def load_card(self, registry_name: str, info: CardInfo) -> ArtifactCard:
145        """
146        Loads an ArtifactCard.
147
148        Args:
149            registry_name:
150                Name of registry to load card from
151            info:
152                Card information to retrieve. `uid` takes precedence if it
153                exists. If the optional `version` is specified, that version
154                will be loaded. If it doesn't exist, the most recent ersion will
155                be loaded.
156
157        Returns
158            `ArtifactCard`
159        """
160        card_type = CardType(registry_name.lower()).value
161        return CardHandler.load_card(
162            registries=self._run_mgr.registries,
163            registry_name=card_type,
164            info=info,
165        )
166
167    def list_runs(self, limit: int = 100) -> List[Dict[str, Any]]:
168        """
169        Lists all runs for the current project, sorted by timestamp
170
171        Returns:
172            List of RunCard
173        """
174        logger.info("Listing runs for project {}", self.project_name)
175
176        project_runs = self._run_mgr.registries.run._registry.list_cards(  # pylint: disable=protected-access
177            limit=limit,
178            query_terms={"project": self.project_name},
179        )
180
181        return sorted(project_runs, key=lambda k: k["timestamp"], reverse=True)
182
183    @property
184    def runcard(self) -> RunCard:
185        return cast(RunCard, self._run_mgr.registries.run.load_card(uid=self.run_id))
186
187    @property
188    def metrics(self) -> Metrics:
189        runcard = self.runcard
190        runcard.load_metrics()
191        return runcard.metrics
192
193    def get_metric(self, name: str) -> Union[List[Metric], Metric]:
194        """
195        Get metric by name
196
197        Args:
198            name: str
199
200        Returns:
201            List of Metric or Metric
202
203        """
204        return self.runcard.get_metric(name=name)
205
206    @property
207    def parameters(self) -> Params:
208        return self.runcard.parameters
209
210    def get_parameter(self, name: str) -> Union[List[Param], Param]:
211        """
212        Get param by name
213
214        Args:
215            name: str
216
217        Returns:
218            List of Param or Param
219
220        """
221        return self.runcard.get_parameter(name=name)
222
223    @property
224    def tags(self) -> Dict[str, Union[str, int]]:
225        return self.runcard.tags
226
227    @property
228    def datacard_uids(self) -> List[str]:
229        """DataCards associated with the current run"""
230        return self.runcard.datacard_uids
231
232    @property
233    def modelcard_uids(self) -> List[str]:
234        """ModelCards associated with the current run"""
235        return self.runcard.modelcard_uids