opsml.cards.run

  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.
  4
  5# pylint: disable=invalid-name
  6
  7
  8import tempfile
  9import uuid
 10from functools import cached_property
 11from pathlib import Path
 12from typing import Any, Dict, List, Optional, Tuple, Union, cast
 13
 14import joblib
 15import numpy as np
 16from numpy.typing import NDArray
 17from pydantic import model_validator
 18
 19from opsml.cards.base import ArtifactCard
 20from opsml.helpers.logging import ArtifactLogger
 21from opsml.helpers.utils import TypeChecker
 22from opsml.settings.config import config
 23from opsml.storage import client
 24from opsml.types import (
 25    Artifact,
 26    ArtifactUris,
 27    CardType,
 28    CommonKwargs,
 29    GraphStyle,
 30    Metric,
 31    Metrics,
 32    Param,
 33    Params,
 34    RegistryTableNames,
 35    RegistryType,
 36    RunCardRegistry,
 37    RunGraph,
 38    SaveName,
 39)
 40
 41logger = ArtifactLogger.get_logger()
 42
 43_List = List[Union[float, int]]
 44_Dict = Dict[str, List[Union[float, int]]]
 45_YReturn = Union[_List, _Dict]
 46_ParseReturn = Tuple[_YReturn, str]
 47
 48
 49def _dump_graph_artifact(graph: RunGraph, name: str, uri: Path) -> Tuple[Path, Path]:
 50    """Helper method for saving graph artifacts to storage
 51
 52    Args:
 53        graph:
 54            RunGraph object
 55        name:
 56            Name of graph
 57        uri:
 58            Uri to store graph artifact
 59    """
 60    with tempfile.TemporaryDirectory() as tempdir:
 61        lpath = Path(tempdir) / f"{name}.joblib"
 62        rpath = (uri / SaveName.GRAPHS.value) / lpath.name
 63        joblib.dump(graph.model_dump(), lpath)
 64
 65        client.storage_client.put(lpath, rpath)
 66
 67        return lpath, rpath
 68
 69
 70def _decimate_list(array: List[Union[float, int]]) -> List[Union[float, int]]:
 71    """Decimates array to no more than 200,000 points
 72
 73    Args:
 74        array:
 75            List of floats or ints
 76
 77    Returns:
 78        Decimated array
 79    """
 80    length = len(array)
 81    if len(array) > 200_000:
 82        step = round(length / 200_000)
 83        return array[::step]
 84
 85    return array
 86
 87
 88def _parse_y_to_list(
 89    x_length: int,
 90    y: Union[List[Union[float, int]], NDArray[Any], Dict[str, Union[List[Union[float, int]], NDArray[Any]]]],
 91) -> _ParseReturn:
 92    """Helper method for parsing y to list when logging a graph
 93
 94    Args:
 95        x_length:
 96            Length of x
 97        y:
 98            Y values to parse. Can be a list, dictionary or numpy array
 99
100    Returns:
101        List or dictionary of y values
102
103    """
104    # if y is dictionary
105    if isinstance(y, dict):
106        _y: Dict[str, List[Union[float, int]]] = {}
107
108        # common sense constraint
109        if len(y.keys()) > 50:
110            raise ValueError("Too many keys in dictionary. A maximum of 50 keys for y is allowed.")
111
112        for k, v in y.items():
113            if isinstance(v, np.ndarray):
114                v = v.flatten().tolist()
115                assert isinstance(v, list), "y must be a list or dictionary"
116                v = _decimate_list(v)
117
118            assert x_length == len(v), "x and y must be the same length"
119            _y[k] = v
120
121        return _y, "multi"
122
123    # if y is ndarray
124    if isinstance(y, np.ndarray):
125        y = y.flatten().tolist()
126        assert isinstance(y, list), "y must be a list or dictionary"
127
128        y = _decimate_list(y)
129        assert x_length == len(y), "x and y must be the same length"
130
131        return y, "single"
132
133    # if y is list
134    assert isinstance(y, list), "y must be a list or dictionary"
135    y = _decimate_list(y)
136    assert x_length == len(y), "x and y must be the same length"
137    return y, "single"
138
139
140class RunCard(ArtifactCard):
141
142    """
143    Create a RunCard from specified arguments.
144
145    Apart from required args, a RunCard must be associated with one of
146    datacard_uid, modelcard_uids or pipelinecard_uid
147
148    Args:
149        name:
150            Run name
151        repository:
152            Repository that this card is associated with
153        contact:
154            Contact to associate with card
155        info:
156            `CardInfo` object containing additional metadata. If provided, it will override any
157            values provided for `name`, `repository`, `contact`, and `version`.
158
159            Name, repository, and contact are required arguments for all cards. They can be provided
160            directly or through a `CardInfo` object.
161
162        datacard_uids:
163            Optional DataCard uids associated with this run
164        modelcard_uids:
165            Optional List of ModelCard uids to associate with this run
166        pipelinecard_uid:
167            Optional PipelineCard uid to associate with this experiment
168        metrics:
169            Optional dictionary of key (str), value (int, float) metric paris.
170            Metrics can also be added via class methods.
171        parameters:
172            Parameters associated with a RunCard
173        artifact_uris:
174            Optional dictionary of artifact uris associated with artifacts.
175        uid:
176            Unique id (assigned if card has been registered)
177        version:
178            Current version (assigned if card has been registered)
179
180    """
181
182    datacard_uids: List[str] = []
183    modelcard_uids: List[str] = []
184    pipelinecard_uid: Optional[str] = None
185    metrics: Metrics = {}
186    parameters: Params = {}
187    artifact_uris: ArtifactUris = {}
188    tags: Dict[str, Union[str, int]] = {}
189    project: Optional[str] = None
190
191    @model_validator(mode="before")
192    @classmethod
193    def validate_defaults_args(cls, card_args: Dict[str, Any]) -> Dict[str, Any]:
194        # add default
195        contact = card_args.get("contact")
196
197        if contact is None:
198            card_args["contact"] = CommonKwargs.UNDEFINED.value
199
200        repository = card_args.get("repository")
201
202        if repository is None:
203            card_args["repository"] = "opsml"
204
205        return card_args
206
207    def add_tag(self, key: str, value: str) -> None:
208        """
209        Logs tags to current RunCard
210
211        Args:
212            key:
213                Key for tag
214            value:
215                value for tag
216        """
217        self.tags = {**{key: value}, **self.tags}
218
219    def add_tags(self, tags: Dict[str, str]) -> None:
220        """
221        Logs tags to current RunCard
222
223        Args:
224            tags:
225                Dictionary of tags
226        """
227        self.tags = {**tags, **self.tags}
228
229    def log_graph(
230        self,
231        name: str,
232        x: Union[List[Union[float, int]], NDArray[Any]],
233        y: Union[List[Union[float, int]], NDArray[Any], Dict[str, Union[List[Union[float, int]], NDArray[Any]]]],
234        y_label: str,
235        x_label: str,
236        graph_style: str,
237    ) -> None:
238        """Logs a graph to the RunCard, which will be rendered in the UI as a line graph
239
240        Args:
241            name:
242                Name of graph
243            x:
244                List or numpy array of x values
245
246            x_label:
247                Label for x axis
248            y:
249                Either a list or numpy array of y values or a dictionary of y values where key is the group label and
250                value is a list or numpy array of y values
251            y_label:
252                Label for y axis
253            graph_style:
254                Style of graph. Options are "line" or "scatter"
255
256        example:
257
258            ### single line graph
259            x = np.arange(1, 400, 0.5)
260            y = x * x
261            run.log_graph(name="graph1", x=x, y=y, x_label="x", y_label="y", graph_style="line")
262
263            ### multi line graph
264            x = np.arange(1, 1000, 0.5)
265            y1 = x * x
266            y2 = y1 * 1.1
267            y3 = y2 * 3
268            run.log_graph(
269                name="multiline",
270                x=x,
271                y={"y1": y1, "y2": y2, "y3": y3},
272                x_label="x",
273                y_label="y",
274                graph_style="line",
275            )
276
277        """
278
279        if isinstance(x, np.ndarray):
280            x = x.flatten().tolist()
281            assert isinstance(x, list), "x must be a list or dictionary"
282
283        x = _decimate_list(x)
284
285        parsed_y, graph_type = _parse_y_to_list(len(x), y)
286
287        logger.info(f"Logging graph {name} to RunCard")
288        graph = RunGraph(
289            name=name,
290            x=x,
291            x_label=x_label,
292            y=parsed_y,
293            y_label=y_label,
294            graph_type=graph_type,
295            graph_style=GraphStyle.from_str(graph_style).value,  # validate graph style
296        )
297
298        # save graph to storage so we can view in ui while run is active
299        lpath, rpath = _dump_graph_artifact(graph, name, self.uri)
300
301        self._add_artifact_uri(
302            name=name,
303            local_path=lpath.as_posix(),
304            remote_path=rpath.as_posix(),
305        )
306
307    def log_parameters(self, parameters: Dict[str, Union[float, int, str]]) -> None:
308        """
309        Logs parameters to current RunCard
310
311        Args:
312            parameters:
313                Dictionary of parameters
314        """
315
316        for key, value in parameters.items():
317            # check key
318            self.log_parameter(key, value)
319
320    def log_parameter(self, key: str, value: Union[int, float, str]) -> None:
321        """
322        Logs parameter to current RunCard
323
324        Args:
325            key:
326                Param name
327            value:
328                Param value
329        """
330
331        TypeChecker.check_param_type(param=value)
332        _key = TypeChecker.replace_spaces(key)
333
334        param = Param(name=key, value=value)
335
336        if self.parameters.get(_key) is not None:
337            self.parameters[_key].append(param)
338
339        else:
340            self.parameters[_key] = [param]
341
342    def log_metric(
343        self,
344        key: str,
345        value: Union[int, float],
346        timestamp: Optional[int] = None,
347        step: Optional[int] = None,
348    ) -> None:
349        """
350        Logs metric to the existing RunCard metric dictionary
351
352        Args:
353            key:
354                Metric name
355            value:
356                Metric value
357            timestamp:
358                Optional timestamp
359            step:
360                Optional step associated with name and value
361        """
362
363        TypeChecker.check_metric_type(metric=value)
364        _key = TypeChecker.replace_spaces(key)
365
366        metric = Metric(name=_key, value=value, timestamp=timestamp, step=step)
367
368        self._registry.insert_metric([{**metric.model_dump(), **{"run_uid": self.uid}}])
369
370        if self.metrics.get(_key) is not None:
371            self.metrics[_key].append(metric)
372        else:
373            self.metrics[_key] = [metric]
374
375    def log_metrics(self, metrics: Dict[str, Union[float, int]], step: Optional[int] = None) -> None:
376        """
377        Log metrics to the existing RunCard metric dictionary
378
379        Args:
380            metrics:
381                Dictionary containing key (str) and value (float or int) pairs
382                to add to the current metric set
383            step:
384                Optional step associated with metrics
385        """
386
387        for key, value in metrics.items():
388            self.log_metric(key, value, step)
389
390    def log_artifact_from_file(
391        self,
392        name: str,
393        local_path: Union[str, Path],
394        artifact_path: Optional[Union[str, Path]] = None,
395    ) -> None:
396        """
397        Log a local file or directory to the opsml server and associate with the current run.
398
399        Args:
400            name:
401                Name to assign to artifact(s)
402            local_path:
403                Local path to file or directory. Can be string or pathlike object
404            artifact_path:
405                Optional path to store artifact in opsml server. If not provided, 'artifacts' will be used
406        """
407
408        lpath = Path(local_path)
409        rpath = self.uri / (artifact_path or SaveName.ARTIFACTS.value)
410
411        if lpath.is_file():
412            rpath = rpath / lpath.name
413
414        client.storage_client.put(lpath, rpath)
415        self._add_artifact_uri(
416            name=name,
417            local_path=lpath.as_posix(),
418            remote_path=rpath.as_posix(),
419        )
420
421    def create_registry_record(self) -> Dict[str, Any]:
422        """Creates a registry record from the current RunCard"""
423
424        exclude_attr = {"parameters", "metrics"}
425
426        return self.model_dump(exclude=exclude_attr)
427
428    def _add_artifact_uri(self, name: str, local_path: str, remote_path: str) -> None:
429        """
430        Adds an artifact_uri to the runcard
431
432        Args:
433            name:
434                Name to associate with artifact
435            uri:
436                Uri where artifact is stored
437        """
438
439        self.artifact_uris[name] = Artifact(
440            name=name,
441            local_path=local_path,
442            remote_path=remote_path,
443        )
444
445    def add_card_uid(self, card_type: str, uid: str) -> None:
446        """
447        Adds a card uid to the appropriate card uid list for tracking
448
449        Args:
450            card_type:
451                ArtifactCard class name
452            uid:
453                Uid of registered ArtifactCard
454        """
455
456        if card_type == CardType.DATACARD:
457            self.datacard_uids = [uid, *self.datacard_uids]
458        elif card_type == CardType.MODELCARD:
459            self.modelcard_uids = [uid, *self.modelcard_uids]
460
461    def get_metric(self, name: str) -> Union[List[Metric], Metric]:
462        """
463        Gets a metric by name
464
465        Args:
466            name:
467                Name of metric
468
469        Returns:
470            List of dictionaries or dictionary containing value
471
472        """
473        _key = TypeChecker.replace_spaces(name)
474
475        metric = self.metrics.get(_key)
476
477        if metric is None:
478            # try to get metric from registry
479            assert self.uid is not None, "RunCard must be registered to get metric"
480            _metric = self._registry.get_metric(run_uid=self.uid, name=[_key])
481
482            if _metric is not None:
483                metric = [Metric(**i) for i in _metric]
484
485            else:
486                raise ValueError(f"Metric {metric} was not defined")
487
488        if len(metric) > 1:
489            return metric
490        if len(metric) == 1:
491            return metric[0]
492        return metric
493
494    def load_metrics(self) -> None:
495        """Reloads metrics from registry"""
496        assert self.uid is not None, "RunCard must be registered to load metrics"
497
498        metrics = self._registry.get_metric(run_uid=self.uid)
499
500        if metrics is None:
501            logger.info("No metrics found for RunCard")
502            return None
503
504        # reset metrics
505        self.metrics = {}
506        for metric in metrics:
507            _metric = Metric(**metric)
508            if _metric.name not in self.metrics:
509                self.metrics[_metric.name] = [_metric]
510            else:
511                self.metrics[_metric.name].append(_metric)
512        return None
513
514    def get_parameter(self, name: str) -> Union[List[Param], Param]:
515        """
516        Gets a parameter by name
517
518        Args:
519            name:
520                Name of parameter
521
522        Returns:
523            List of dictionaries or dictionary containing value
524
525        """
526        _key = TypeChecker.replace_spaces(name)
527        param = self.parameters.get(_key)
528        if param is not None:
529            if len(param) > 1:
530                return param
531            if len(param) == 1:
532                return param[0]
533            return param
534
535        raise ValueError(f"Param {param} is not defined")
536
537    def load_artifacts(self, name: Optional[str] = None) -> None:
538        """Loads artifacts from artifact_uris"""
539        if bool(self.artifact_uris) is False:
540            logger.info("No artifact uris associated with RunCard")
541            return None
542
543        if name is not None:
544            artifact = self.artifact_uris.get(name)
545            assert artifact is not None, f"Artifact {name} not found"
546            client.storage_client.get(
547                Path(artifact.remote_path),
548                Path(artifact.local_path),
549            )
550
551        else:
552            for _, artifact in self.artifact_uris.items():
553                client.storage_client.get(
554                    Path(artifact.remote_path),
555                    Path(artifact.local_path),
556                )
557        return None
558
559    @property
560    def uri(self) -> Path:
561        """The base URI to use for the card and it's artifacts."""
562
563        # when using runcard outside of run context
564        if self.version == CommonKwargs.BASE_VERSION.value:
565            if self.uid is None:
566                self.uid = uuid.uuid4().hex
567
568            end_path = self.uid
569        else:
570            end_path = f"v{self.version}"
571
572        return Path(
573            config.storage_root,
574            RegistryTableNames.from_str(self.card_type).value,
575            str(self.repository),
576            str(self.name),
577            end_path,
578        )
579
580    @cached_property
581    def _registry(self) -> RunCardRegistry:
582        from opsml.registry.backend import _set_registry
583
584        return cast(RunCardRegistry, _set_registry(RegistryType.RUN))
585
586    @property
587    def card_type(self) -> str:
588        return CardType.RUNCARD.value