truefoundry 0.3.3__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

truefoundry/ml/log_types/plot.py

@@ -0,0 +1,281 @@
+import enum
+import json
+import os
+import posixpath
+import sys
+import tempfile
+from typing import TYPE_CHECKING, Union
+
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    ArtifactType,
+)
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types.artifacts.artifact import (
+    ArtifactVersionInternalMetadata,
+    _log_artifact_version_helper,
+)
+from truefoundry.ml.log_types.artifacts.constants import (
+    FILES_DIR,
+    INTERNAL_METADATA_PATH,
+)
+from truefoundry.ml.log_types.pydantic_base import PydanticBase
+from truefoundry.ml.log_types.utils import validate_key_name
+from truefoundry.pydantic_v1 import BaseModel
+
+if TYPE_CHECKING:
+    import matplotlib.figure
+    import matplotlib.pyplot
+    import plotly.graph_objects
+
+    from truefoundry.ml import MlFoundryRun
+
+PlotObjType = Union[
+    "matplotlib.figure.Figure",
+    "plotly.graph_objects.Figure",
+    "matplotlib.pyplot",
+]
+
+
+@enum.unique
+class Format(enum.Enum):
+    SVG = "SVG"
+    HTML = "HTML"
+    PNG = "PNG"
+
+
+def _is_matplotlib_figure(fig) -> bool:
+    if "matplotlib" not in sys.modules:
+        return False
+    import matplotlib
+
+    return isinstance(fig, matplotlib.figure.Figure)
+
+
+def _is_matplotlib_plt(plt) -> bool:
+    if "matplotlib" not in sys.modules:
+        return False
+    return getattr(plt, "__name__", "") == "matplotlib.pyplot"
+
+
+def _is_plotly_figure(fig) -> bool:
+    if "plotly" not in sys.modules:
+        return False
+
+    import plotly
+
+    return isinstance(fig, plotly.graph_objects.Figure)
+
+
+def get_plot_file_name(format: Format) -> str:
+    return f"plot.{format.value.lower()}"
+
+
+class PlotArtifact(BaseModel):
+    artifact_file: str
+    format: Format
+
+    class Config:
+        allow_mutation = False
+        use_enum_values = True
+
+
+def _save_matplotlib_figure(
+    figure: "matplotlib.figure.Figure",
+    key: str,
+    step: int,
+    local_dir: str,
+) -> PlotArtifact:
+    supported_formats = figure.canvas.get_supported_filetypes().keys()
+    if "svg" in supported_formats:
+        format_ = Format.SVG
+    elif "png" in supported_formats:
+        format_ = Format.PNG
+    else:
+        raise MlFoundryException(
+            f"Could not save {key} {figure} matplotlib figure"
+            "in either SVG or PNG format"
+        )
+    file_path = get_plot_file_name(format=format_)
+    local_path = os.path.join(local_dir, file_path)
+    figure.savefig(local_path)
+    return PlotArtifact(artifact_file=file_path, format=format_)
+
+
+def _save_matplotlib_plt(
+    plt: "matplotlib.pyplot",
+    key: str,
+    step: int,
+    local_dir: str,
+) -> PlotArtifact:
+    figure = plt.gcf()
+    return _save_matplotlib_figure(
+        figure=figure, key=key, step=step, local_dir=local_dir
+    )
+
+
+def _save_plotly_figure(
+    figure: "plotly.graph_objects.Figure",
+    key: str,
+    step: int,
+    local_dir: str,
+) -> PlotArtifact:
+    format_ = Format.HTML
+    file_path = get_plot_file_name(format=format_)
+    local_path = os.path.join(local_dir, file_path)
+    figure.write_html(local_path, include_plotlyjs="cdn", auto_open=False)
+    return PlotArtifact(artifact_file=file_path, format=format_)
+
+
+class Plot:
+    def __init__(self, plot_obj: PlotObjType):
+        self._plot_obj = plot_obj
+
+    def _save_plot(self, key: str, step: int, local_dir: str) -> PlotArtifact:
+        if _is_matplotlib_plt(self._plot_obj):
+            return _save_matplotlib_plt(
+                plt=self._plot_obj, key=key, step=step, local_dir=local_dir
+            )
+
+        if _is_matplotlib_figure(self._plot_obj):
+            return _save_matplotlib_figure(
+                figure=self._plot_obj, key=key, step=step, local_dir=local_dir
+            )
+
+        if _is_plotly_figure(self._plot_obj):
+            return _save_plotly_figure(
+                figure=self._plot_obj, key=key, step=step, local_dir=local_dir
+            )
+
+        raise MlFoundryException(
+            f"Unknown type: {type(self._plot_obj)}"
+            "Supported types are, matplotlib.figure.Figure, matplotlib.pyplot"
+            " and plotly.graph_objects.Figure"
+        )
+
+    def save(
+        self,
+        run: "MlFoundryRun",
+        key: str,
+        step: int = 0,
+    ):
+        validate_key_name(key)
+
+        # creating a temp dir which will be logged
+        temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
+        local_files_dir = os.path.join(temp_dir.name, FILES_DIR)
+        os.makedirs(local_files_dir, exist_ok=True)
+
+        # save plot locally
+        plot_artifact = self._save_plot(key=key, step=step, local_dir=local_files_dir)
+
+        # save internal metadata
+        internal_metadata = PlotVersionInternalMetadata(
+            files_dir=FILES_DIR,
+            plot_file=posixpath.join(FILES_DIR, plot_artifact.artifact_file),
+        )
+        local_internal_metadata_path = os.path.join(
+            temp_dir.name, INTERNAL_METADATA_PATH
+        )
+        os.makedirs(os.path.dirname(local_internal_metadata_path), exist_ok=True)
+        with open(local_internal_metadata_path, "w") as f:
+            json.dump(internal_metadata.dict(), f)
+
+        # log the artifact
+        _log_artifact_version_helper(
+            run=run,
+            name=key,
+            artifact_type=ArtifactType.PLOT,
+            artifact_dir=temp_dir,
+            internal_metadata=internal_metadata,
+            step=step,
+        )
+
+
+class PlotRunLogType(PydanticBase):
+    value: PlotArtifact
+
+    @staticmethod
+    def get_log_type() -> str:
+        return "plot"
+
+
+class PlotVersionInternalMetadata(ArtifactVersionInternalMetadata):
+    plot_file: str
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+    import numpy as np
+    import plotly.express as px
+
+    from truefoundry.ml import get_client
+
+    client = get_client()
+
+    run = client.create_run(ml_repo="plot-test1")
+
+    df = px.data.iris()
+    fig = px.scatter(
+        df,
+        x="sepal_width",
+        y="sepal_length",
+        color="species",
+        size="petal_length",
+        hover_data=["petal_width"],
+    )
+    Plot(fig).save(run, "foo")
+
+    df = px.data.tips()
+    fig = px.histogram(
+        df,
+        x="total_bill",
+        y="tip",
+        color="sex",
+        marginal="rug",
+        hover_data=df.columns,
+    )
+    Plot(fig).save(run, "foo", step=1)
+
+    names = ["group_a", "group_b", "group_c"]
+    values = [1, 10, 100]
+    plt.figure(figsize=(9, 3))
+    plt.subplot(131)
+    plt.bar(names, values)
+    plt.subplot(132)
+    plt.scatter(names, values)
+    plt.subplot(133)
+    plt.plot(names, values)
+    plt.suptitle("Categorical Plotting")
+
+    Plot(plt).save(run, "bar")
+
+    plt.clf()
+
+    data = {
+        "a": np.arange(50),
+        "c": np.random.randint(0, 50, 50),
+        "d": np.random.randn(50),
+    }
+    data["b"] = data["a"] + 10 * np.random.randn(50)
+    data["d"] = np.abs(data["d"]) * 100
+    plt.scatter("a", "b", c="c", s="d", data=data)
+    plt.xlabel("entry a")
+    plt.ylabel("entry b")
+    Plot(plt).save(run, "bar", 2)
+
+    plt.clf()
+
+    ax = plt.subplot()
+    t = np.arange(0.0, 5.0, 0.01)
+    s = np.cos(2 * np.pi * t)
+    (line,) = plt.plot(t, s, lw=2)
+
+    plt.annotate(
+        "local max",
+        xy=(2, 1),
+        xytext=(3, 1.5),
+        arrowprops={"facecolor": "black", "shrink": 0.05},
+    )
+
+    plt.ylim(-2, 2)
+    Plot(plt.gcf()).save(run, "bar", 3)

truefoundry/ml/log_types/pydantic_base.py

@@ -0,0 +1,10 @@
+from truefoundry.pydantic_v1 import BaseModel
+
+
+class PydanticBase(BaseModel):
+    # I can make this a property,
+    # but <3.9, it is difficult to access
+    # property from classmethod
+    @staticmethod
+    def get_log_type() -> str:
+        raise NotImplementedError()

truefoundry/ml/log_types/utils.py

@@ -0,0 +1,12 @@
+import re
+
+from truefoundry.ml.exceptions import MlFoundryException
+
+KEY_REGEX = re.compile(r"^[a-zA-Z0-9-_]+$")
+
+
+def validate_key_name(key: str):
+    if not key or not KEY_REGEX.match(key):
+        raise MlFoundryException(
+            f"Invalid run image key: {key} should only contain alphanumeric, hyphen or underscore"
+        )

truefoundry/ml/logger.py

@@ -0,0 +1,17 @@
+import logging
+import sys
+
+logger = logging.getLogger("truefoundry.ml")
+
+
+def init_logger(level=logging.INFO):
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+    formatter = logging.Formatter(
+        "[%(name)s] %(asctime)s %(levelname)s %(message)s",
+        datefmt="%Y-%m-%dT%H:%M:%S%z",
+    )
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False

truefoundry/ml/login.py

@@ -0,0 +1,241 @@
+import os
+import threading
+from functools import lru_cache, wraps
+from pathlib import Path
+from typing import Optional
+
+import click
+from filelock import FileLock, Timeout
+
+from truefoundry.ml.env_vars import API_KEY_GLOBAL, TRACKING_HOST_GLOBAL
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.logger import logger
+from truefoundry.ml.run_utils import resolve_tracking_uri
+from truefoundry.ml.services.auth_service import get_auth_service
+from truefoundry.ml.services.entities import Token
+from truefoundry.pydantic_v1 import BaseModel, Field, NonEmptyStr
+
+CREDENTIALS_DIR = Path.home() / ".truefoundry"
+CREDENTIALS_FILE = CREDENTIALS_DIR / "credentials.json"
+
+
+OLD_CREDENTIALS_DIR = Path.home() / ".mlfoundry"
+OLD_CREDENTIALS_FILE = OLD_CREDENTIALS_DIR / "credentials.netrc"
+
+if OLD_CREDENTIALS_FILE.exists():
+    logger.warning(
+        "%s file is deprecated. You can delete this file now.", OLD_CREDENTIALS_FILE
+    )
+
+
+class CredentialsFileContent(BaseModel):
+    access_token: NonEmptyStr = Field(repr=False)
+    refresh_token: Optional[NonEmptyStr] = Field(repr=False)
+    host: NonEmptyStr
+
+    class Config:
+        allow_mutation = False
+
+    def to_token(self) -> Token:
+        return Token(
+            access_token=self.access_token,  # type: ignore[call-arg]
+            refresh_token=self.refresh_token,
+        )
+
+
+def _ensure_lock_taken(method):
+    @wraps(method)
+    def lock_guard(self: "CredentialsFileManager", *method_args, **method_kwargs):
+        if not self.lock_taken():
+            raise Exception(
+                "Trying to write to credential file without using with block"
+            )
+        return method(self, *method_args, **method_kwargs)
+
+    return lock_guard
+
+
+CRED_FILE_THREAD_LOCK = threading.RLock()
+
+
+@lru_cache(maxsize=None)
+def get_file_lock(lock_file_path: str) -> FileLock:
+    return FileLock(lock_file_path)
+
+
+class CredentialsFileManager:
+    def __init__(
+        self,
+        credentials_file_path: Path = CREDENTIALS_FILE,
+        lock_timeout: float = 60.0,
+    ):
+        credentials_file_path = credentials_file_path.absolute()
+
+        logger.debug("credential file path %r", credentials_file_path)
+
+        credentials_lock_file_path = f"{credentials_file_path}.lock"
+
+        logger.debug("credential lock file path %r", credentials_lock_file_path)
+        self._credentials_file_path = credentials_file_path
+
+        cred_file_dir = credentials_file_path.parent
+        cred_file_dir.mkdir(exist_ok=True, parents=True)
+
+        self._file_lock = get_file_lock(credentials_lock_file_path)
+        self._lock_timeout = lock_timeout
+        self._lock_owner: Optional[int] = None
+
+    def __enter__(self) -> "CredentialsFileManager":
+        # The lock objects are recursive locks, which means that once acquired,
+        # they will not block on successive lock requests:
+        lock_aquired = CRED_FILE_THREAD_LOCK.acquire(timeout=self._lock_timeout)
+        if not lock_aquired:
+            raise MlFoundryException(
+                "Could not aquire CRED_FILE_THREAD_LOCK"
+                f" in {self._lock_timeout} seconds"
+            )
+        try:
+            self._file_lock.acquire(timeout=self._lock_timeout)
+        except Timeout as ex:
+            raise MlFoundryException(
+                f"Failed to aquire lock on credential file within {self._lock_timeout} seconds.\n"
+                "Is any other process trying to login?"
+            ) from ex
+        logger.debug("Acquired file and thread lock to access credential file")
+        self._lock_owner = threading.get_ident()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._file_lock.release()
+        CRED_FILE_THREAD_LOCK.release()
+        logger.debug("Released file and thread lock to access credential file")
+        self._lock_owner = None
+
+    def lock_taken(self) -> bool:
+        return self._lock_owner == threading.get_ident()
+
+    @_ensure_lock_taken
+    def read(self) -> CredentialsFileContent:
+        try:
+            return CredentialsFileContent.parse_file(self._credentials_file_path)
+        except Exception as ex:
+            raise MlFoundryException(
+                "Error while reading the credentials file "
+                f"{self._credentials_file_path}. Please login again "
+                "using `tfy login` command or `truefoundry.login()` function"
+            ) from ex
+
+    @_ensure_lock_taken
+    def write(self, credentials_file_content: CredentialsFileContent):
+        if not isinstance(credentials_file_content, CredentialsFileContent):
+            raise MlFoundryException(
+                "Only object of type `CredentialsFileContent` is allowed. "
+                f"Got {type(credentials_file_content)}"
+            )
+        logger.debug("Updating the credential file content")
+        with open(self._credentials_file_path, "w", encoding="utf8") as file:
+            file.write(credentials_file_content.json())
+
+    @_ensure_lock_taken
+    def exists(self) -> bool:
+        return self._credentials_file_path.exists()
+
+
+def login(
+    tracking_uri: Optional[str] = None,
+    relogin: bool = False,
+    api_key: Optional[str] = None,
+) -> bool:
+    """Save API key in local file system for a given `tracking_uri`.
+
+    Args:
+        tracking_uri (Optional[str], optional): tracking_uri for the given API key
+        relogin (bool, optional): Overwrites the existing API key for the `tracking_uri` if
+            set to `True`. If set to `False` and an API key is already present for
+            the given `tracking_uri`, then the existing API key is kept untouched.
+            Default is `False`.
+        api_key (Optional[str], optional): The API key for the given `tracking_uri`.
+            If `api_key` is not passed, this function prompts for the API key.
+
+    Returns:
+        bool: Returns `True` if any credential was persisted.
+    """
+    from truefoundry.ml.session import EnvCredentialProvider
+
+    if API_KEY_GLOBAL in os.environ and TRACKING_HOST_GLOBAL in os.environ:
+        logger.warning(
+            "Skipping login because environment variables %s and "
+            "%s are set and will be used when running truefoundry. "
+            "If you want to relogin then unset these environment keys.",
+            TRACKING_HOST_GLOBAL,
+            API_KEY_GLOBAL,
+        )
+        return False
+
+    if EnvCredentialProvider.can_provide():
+        logger.warning(
+            "TFY_API_KEY env var is already set. "
+            "When running mlfoundry, it will use the api key to authorize.\n"
+            "Login will just save the credentials on disk."
+        )
+
+    tracking_uri = resolve_tracking_uri(tracking_uri).strip("/")
+    auth_service = get_auth_service(tracking_uri=tracking_uri)
+
+    cred_file = CredentialsFileManager()
+
+    with cred_file:
+        if not relogin and cred_file.exists():
+            cred_file_content = cred_file.read()
+            if tracking_uri != cred_file_content.host:
+                if click.confirm(
+                    f"Already logged in to {cred_file_content.host!r}\n"
+                    f"Do you want to relogin to {tracking_uri!r}?"
+                ):
+                    return login(
+                        tracking_uri=tracking_uri, relogin=True, api_key=api_key
+                    )
+            user_info = cred_file_content.to_token().to_user_info()
+            user_name_display_info = user_info.email or user_info.user_type.value
+            print(
+                f"Already logged in to {cred_file_content.host!r} as "
+                f"{user_info.user_id!r} ({user_name_display_info!r})\n"
+                "Please use `tfy login --relogin` or `truefoundry.login(relogin=True)` "
+                "to force relogin"
+            )
+            return False
+
+        if api_key:
+            logger.debug("Logging in with api key")
+            token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
+            cred_file_content = CredentialsFileContent(
+                access_token=token.access_token,
+                refresh_token=token.refresh_token,
+                host=tracking_uri,
+            )
+            cred_file.write(cred_file_content)
+        else:
+            device_code = auth_service.get_device_code()
+            url_to_go = device_code.get_user_clickable_url(tracking_uri=tracking_uri)
+            print(f"Opening:- {url_to_go}")
+            print(
+                "Please click on the above link if it is not "
+                "automatically opened in a browser window."
+            )
+            click.launch(url_to_go)
+            token = auth_service.get_token_from_device_code(
+                device_code=device_code.device_code, timeout=120
+            )
+            cred_file_content = CredentialsFileContent(
+                access_token=token.access_token,
+                refresh_token=token.refresh_token,
+                host=tracking_uri,
+            )
+            cred_file.write(cred_file_content)
+
+    user_info = token.to_user_info()
+    user_name_display_info = user_info.email or user_info.user_type.value
+    print(
+        f"Logged in to {cred_file_content.host!r} as {user_info.user_id!r} ({user_name_display_info})"
+    )
+    return True