genesis-flow 1.0.0__py3-none-any.whl

mlflow/evaluation/utils.py

@@ -0,0 +1,201 @@
+"""
+THE 'mlflow.evaluation` MODULE IS LEGACY AND WILL BE REMOVED SOON. PLEASE DO NOT USE THESE CLASSES
+IN NEW CODE. INSTEAD, USE `mlflow/entities/assessment.py` FOR ASSESSMENT CLASSES.
+"""
+
+import pandas as pd
+
+from mlflow.evaluation.evaluation import EvaluationEntity as EvaluationEntity
+from mlflow.utils.annotations import deprecated
+
+
+@deprecated(since="3.0.0")
+def evaluations_to_dataframes(
+    evaluations: list[EvaluationEntity],
+) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    """
+    Converts a list of Evaluation entities to four separate DataFrames: one for main evaluation
+    data (excluding assessments and metrics), one for metrics, one for assessments, and one for
+    tags.
+
+    Args:
+        evaluations (List[Evaluation]): List of Evaluation entities.
+
+    Returns:
+        Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]: A tuple of four DataFrames
+        containing evaluation data, metrics data, assessments data, and tags data.
+    """
+    evaluations_data = []
+    metrics_data = []
+    assessments_data = []
+    tags_data = []
+
+    for evaluation in evaluations:
+        eval_dict = evaluation.to_dictionary()
+
+        # Extract assessment and metrics
+        assessments_list = eval_dict.pop("assessments", [])
+        metrics_list = eval_dict.pop("metrics", [])
+        tags_list = eval_dict.pop("tags", [])
+
+        evaluations_data.append(eval_dict)
+
+        for metric_dict in metrics_list:
+            metric_dict["evaluation_id"] = eval_dict["evaluation_id"]
+            # Remove 'step' key if it exists, since it is not valid for evaluation metrics
+            metric_dict.pop("step", None)
+            metrics_data.append(metric_dict)
+
+        for assess_dict in assessments_list:
+            assess_dict["evaluation_id"] = eval_dict["evaluation_id"]
+            assessments_data.append(assess_dict)
+
+        for tag_dict in tags_list:
+            tag_dict["evaluation_id"] = eval_dict["evaluation_id"]
+            tags_data.append(tag_dict)
+
+    evaluations_df = (
+        _apply_schema_to_dataframe(
+            pd.DataFrame(evaluations_data), _get_evaluations_dataframe_schema()
+        )
+        if evaluations_data
+        else _get_empty_evaluations_dataframe()
+    )
+    metrics_df = (
+        _apply_schema_to_dataframe(pd.DataFrame(metrics_data), _get_metrics_dataframe_schema())
+        if metrics_data
+        else _get_empty_metrics_dataframe()
+    )
+    assessments_df = (
+        _apply_schema_to_dataframe(
+            pd.DataFrame(assessments_data), _get_assessments_dataframe_schema()
+        )
+        if assessments_data
+        else _get_empty_assessments_dataframe()
+    )
+    tags_df = (
+        _apply_schema_to_dataframe(pd.DataFrame(tags_data), _get_tags_dataframe_schema())
+        if tags_data
+        else _get_empty_tags_dataframe()
+    )
+
+    return evaluations_df, metrics_df, assessments_df, tags_df
+
+
+def _get_evaluations_dataframe_schema() -> dict[str, str]:
+    """
+    Returns the pandas schema for the evaluation DataFrame.
+    """
+    return {
+        "evaluation_id": "string",
+        "run_id": "string",
+        "inputs_id": "string",
+        "inputs": "object",
+        "outputs": "object",
+        "request_id": "object",
+        "targets": "object",
+        "error_code": "object",
+        "error_message": "object",
+    }
+
+
+def _get_empty_evaluations_dataframe() -> pd.DataFrame:
+    """
+    Creates an empty DataFrame with columns for evaluation data.
+    """
+    schema = _get_evaluations_dataframe_schema()
+    df = pd.DataFrame(columns=schema.keys())
+    return _apply_schema_to_dataframe(df, schema)
+
+
+def _get_assessments_dataframe_schema() -> dict[str, str]:
+    """
+    Returns the pandas schema for the assessments DataFrame.
+    """
+    return {
+        "evaluation_id": "string",
+        "name": "string",
+        "source": "object",
+        "timestamp": "int64",
+        "boolean_value": "object",
+        "numeric_value": "object",
+        "string_value": "object",
+        "rationale": "object",
+        "metadata": "object",
+        "error_code": "object",
+        "error_message": "object",
+        "span_id": "object",
+    }
+
+
+def _get_empty_assessments_dataframe() -> pd.DataFrame:
+    """
+    Creates an empty DataFrame with columns for evaluation assessments data.
+    """
+    schema = _get_assessments_dataframe_schema()
+    df = pd.DataFrame(columns=schema.keys())
+    return _apply_schema_to_dataframe(df, schema)
+
+
+def _get_metrics_dataframe_schema() -> dict[str, str]:
+    """
+    Returns the pandas schema for the metrics DataFrame.
+    """
+    return {
+        "evaluation_id": "string",
+        "key": "string",
+        "value": "float64",
+        "timestamp": "int64",
+        "model_id": "string",
+        "dataset_name": "string",
+        "dataset_digest": "string",
+        "run_id": "string",
+    }
+
+
+def _get_empty_metrics_dataframe() -> pd.DataFrame:
+    """
+    Creates an empty DataFrame with columns for evaluation metric data.
+    """
+    schema = _get_metrics_dataframe_schema()
+    df = pd.DataFrame(columns=schema.keys())
+    return _apply_schema_to_dataframe(df, schema)
+
+
+def _get_tags_dataframe_schema() -> dict[str, str]:
+    """
+    Returns the pandas schema for the tags DataFrame.
+    """
+    return {
+        "evaluation_id": "string",
+        "key": "string",
+        "value": "string",
+    }
+
+
+def _get_empty_tags_dataframe() -> pd.DataFrame:
+    """
+    Creates an empty DataFrame with columns for evaluation tags data.
+    """
+    schema = _get_tags_dataframe_schema()
+    df = pd.DataFrame(columns=schema.keys())
+    return _apply_schema_to_dataframe(df, schema)
+
+
+def _apply_schema_to_dataframe(df: pd.DataFrame, schema: dict[str, str]) -> pd.DataFrame:
+    """
+    Applies a schema to a DataFrame.
+
+    Args:
+        df (pd.DataFrame): DataFrame to apply the schema to.
+        schema (Dict[str, Any]): Schema to apply.
+
+    Returns:
+        pd.DataFrame: DataFrame with schema applied.
+    """
+    for column in df.columns:
+        df[column] = df[column].astype(schema[column])
+    # By default, null values are represented as `pd.NA` in pandas when reading a dataframe from
+    # JSON. However, MLflow entities use `None` to represent null values. Accordingly, we convert
+    # instances of pd.NA to None so that DataFrame rows can be parsed as MLflow entities
+    return df.replace(pd.NA, None)

mlflow/exceptions.py

@@ -0,0 +1,213 @@
+import json
+import logging
+from typing import Optional
+
+from mlflow.protos.databricks_pb2 import (
+    ABORTED,
+    ALREADY_EXISTS,
+    BAD_REQUEST,
+    CANCELLED,
+    CUSTOMER_UNAUTHORIZED,
+    DATA_LOSS,
+    DEADLINE_EXCEEDED,
+    ENDPOINT_NOT_FOUND,
+    INTERNAL_ERROR,
+    INVALID_PARAMETER_VALUE,
+    INVALID_STATE,
+    NOT_FOUND,
+    NOT_IMPLEMENTED,
+    PERMISSION_DENIED,
+    REQUEST_LIMIT_EXCEEDED,
+    RESOURCE_ALREADY_EXISTS,
+    RESOURCE_CONFLICT,
+    RESOURCE_DOES_NOT_EXIST,
+    RESOURCE_EXHAUSTED,
+    TEMPORARILY_UNAVAILABLE,
+    UNAUTHENTICATED,
+    ErrorCode,
+)
+
+ERROR_CODE_TO_HTTP_STATUS = {
+    ErrorCode.Name(INTERNAL_ERROR): 500,
+    ErrorCode.Name(INVALID_STATE): 500,
+    ErrorCode.Name(DATA_LOSS): 500,
+    ErrorCode.Name(NOT_IMPLEMENTED): 501,
+    ErrorCode.Name(TEMPORARILY_UNAVAILABLE): 503,
+    ErrorCode.Name(DEADLINE_EXCEEDED): 504,
+    ErrorCode.Name(REQUEST_LIMIT_EXCEEDED): 429,
+    ErrorCode.Name(CANCELLED): 499,
+    ErrorCode.Name(RESOURCE_EXHAUSTED): 429,
+    ErrorCode.Name(ABORTED): 409,
+    ErrorCode.Name(RESOURCE_CONFLICT): 409,
+    ErrorCode.Name(ALREADY_EXISTS): 409,
+    ErrorCode.Name(NOT_FOUND): 404,
+    ErrorCode.Name(ENDPOINT_NOT_FOUND): 404,
+    ErrorCode.Name(RESOURCE_DOES_NOT_EXIST): 404,
+    ErrorCode.Name(PERMISSION_DENIED): 403,
+    ErrorCode.Name(CUSTOMER_UNAUTHORIZED): 401,
+    ErrorCode.Name(UNAUTHENTICATED): 401,
+    ErrorCode.Name(BAD_REQUEST): 400,
+    ErrorCode.Name(RESOURCE_ALREADY_EXISTS): 400,
+    ErrorCode.Name(INVALID_PARAMETER_VALUE): 400,
+}
+
+HTTP_STATUS_TO_ERROR_CODE = {v: k for k, v in ERROR_CODE_TO_HTTP_STATUS.items()}
+HTTP_STATUS_TO_ERROR_CODE[400] = ErrorCode.Name(BAD_REQUEST)
+HTTP_STATUS_TO_ERROR_CODE[404] = ErrorCode.Name(ENDPOINT_NOT_FOUND)
+HTTP_STATUS_TO_ERROR_CODE[500] = ErrorCode.Name(INTERNAL_ERROR)
+
+_logger = logging.getLogger(__name__)
+
+
+def get_error_code(http_status):
+    return ErrorCode.Value(
+        HTTP_STATUS_TO_ERROR_CODE.get(http_status, ErrorCode.Name(INTERNAL_ERROR))
+    )
+
+
+class MlflowException(Exception):
+    """
+    Generic exception thrown to surface failure information about external-facing operations.
+    The error message associated with this exception may be exposed to clients in HTTP responses
+    for debugging purposes. If the error text is sensitive, raise a generic `Exception` object
+    instead.
+    """
+
+    def __init__(self, message, error_code=INTERNAL_ERROR, **kwargs):
+        """
+        Args:
+            message: The message or exception describing the error that occurred. This will be
+                included in the exception's serialized JSON representation.
+            error_code: An appropriate error code for the error that occurred; it will be
+                included in the exception's serialized JSON representation. This should
+                be one of the codes listed in the `mlflow.protos.databricks_pb2` proto.
+            kwargs: Additional key-value pairs to include in the serialized JSON representation
+                of the MlflowException.
+        """
+        try:
+            self.error_code = ErrorCode.Name(error_code)
+        except (ValueError, TypeError):
+            self.error_code = ErrorCode.Name(INTERNAL_ERROR)
+        message = str(message)
+        self.message = message
+        self.json_kwargs = kwargs
+        super().__init__(message)
+
+    def serialize_as_json(self):
+        exception_dict = {"error_code": self.error_code, "message": self.message}
+        exception_dict.update(self.json_kwargs)
+        return json.dumps(exception_dict)
+
+    def get_http_status_code(self):
+        return ERROR_CODE_TO_HTTP_STATUS.get(self.error_code, 500)
+
+    @classmethod
+    def invalid_parameter_value(cls, message, **kwargs):
+        """Constructs an `MlflowException` object with the `INVALID_PARAMETER_VALUE` error code.
+
+        Args:
+            message: The message describing the error that occurred. This will be included in the
+                exception's serialized JSON representation.
+            kwargs: Additional key-value pairs to include in the serialized JSON representation
+                of the MlflowException.
+        """
+        return cls(message, error_code=INVALID_PARAMETER_VALUE, **kwargs)
+
+
+class RestException(MlflowException):
+    """Exception thrown on non 200-level responses from the REST API"""
+
+    def __init__(self, json):
+        self.json = json
+
+        error_code = json.get("error_code", ErrorCode.Name(INTERNAL_ERROR))
+        message = "{}: {}".format(
+            error_code,
+            json["message"] if "message" in json else "Response: " + str(json),
+        )
+
+        try:
+            super().__init__(message, error_code=ErrorCode.Value(error_code))
+        except ValueError:
+            try:
+                # The `error_code` can be an http error code, in which case we convert it to the
+                # corresponding `ErrorCode`.
+                error_code = HTTP_STATUS_TO_ERROR_CODE[int(error_code)]
+                super().__init__(message, error_code=ErrorCode.Value(error_code))
+            except ValueError or KeyError:
+                _logger.warning(
+                    f"Received error code not recognized by MLflow: {error_code}, this may "
+                    "indicate your request encountered an error before reaching MLflow server, "
+                    "e.g., within a proxy server or authentication / authorization service."
+                )
+                super().__init__(message)
+
+    def __reduce__(self):
+        """
+        Overriding `__reduce__` to make `RestException` instance pickle-able.
+        """
+        return RestException, (self.json,)
+
+
+class ExecutionException(MlflowException):
+    """Exception thrown when executing a project fails"""
+
+
+class MissingConfigException(MlflowException):
+    """Exception thrown when expected configuration file/directory not found"""
+
+
+class InvalidUrlException(MlflowException):
+    """Exception thrown when a http request fails to send due to an invalid URL"""
+
+
+class _UnsupportedMultipartUploadException(MlflowException):
+    """Exception thrown when multipart upload is unsupported by an artifact repository"""
+
+    MESSAGE = "Multipart upload is not supported for the current artifact repository"
+
+    def __init__(self):
+        super().__init__(self.MESSAGE, error_code=NOT_IMPLEMENTED)
+
+
+class MlflowTracingException(MlflowException):
+    """
+    Exception thrown from tracing logic
+
+    Tracing logic should not block the main execution flow in general, hence this exception
+    is used to distinguish tracing related errors and handle them properly.
+    """
+
+    def __init__(self, message, error_code=INTERNAL_ERROR):
+        super().__init__(message, error_code=error_code)
+
+
+class MlflowTraceDataException(MlflowTracingException):
+    """Exception thrown for trace data related error"""
+
+    def __init__(
+        self, error_code: str, request_id: Optional[str] = None, artifact_path: Optional[str] = None
+    ):
+        if request_id:
+            self.ctx = f"request_id={request_id}"
+        elif artifact_path:
+            self.ctx = f"path={artifact_path}"
+
+        if error_code == NOT_FOUND:
+            super().__init__(f"Trace data not found for {self.ctx}", error_code=error_code)
+        elif error_code == INVALID_STATE:
+            super().__init__(f"Trace data is corrupted for {self.ctx}", error_code=error_code)
+
+
+class MlflowTraceDataNotFound(MlflowTraceDataException):
+    """Exception thrown when trace data is not found"""
+
+    def __init__(self, request_id: Optional[str] = None, artifact_path: Optional[str] = None):
+        super().__init__(NOT_FOUND, request_id, artifact_path)
+
+
+class MlflowTraceDataCorrupted(MlflowTraceDataException):
+    """Exception thrown when trace data is corrupted"""
+
+    def __init__(self, request_id: Optional[str] = None, artifact_path: Optional[str] = None):
+        super().__init__(INVALID_STATE, request_id, artifact_path)

mlflow/experiments.py

@@ -0,0 +1,140 @@
+import os
+
+import click
+
+import mlflow
+from mlflow.entities import ViewType
+from mlflow.tracking import _get_store, fluent
+from mlflow.utils.data_utils import is_uri
+from mlflow.utils.string_utils import _create_table
+
+EXPERIMENT_ID = click.option("--experiment-id", "-x", type=click.STRING, required=True)
+
+
+@click.group("experiments")
+def commands():
+    """
+    Manage experiments. To manage experiments associated with a tracking server, set the
+    MLFLOW_TRACKING_URI environment variable to the URL of the desired server.
+    """
+
+
+@commands.command()
+@click.option("--experiment-name", "-n", type=click.STRING, required=True)
+@click.option(
+    "--artifact-location",
+    "-l",
+    help="Base location for runs to store artifact results. Artifacts will be stored "
+    "at $artifact_location/$run_id/artifacts. See "
+    "https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded for "
+    "more info on the properties of artifact location. "
+    "If no location is provided, the tracking server will pick a default.",
+)
+def create(experiment_name, artifact_location):
+    """
+    Create an experiment.
+
+    All artifacts generated by runs related to this experiment will be stored under artifact
+    location, organized under specific run_id sub-directories.
+
+    Implementation of experiment and metadata store is dependent on backend storage. ``FileStore``
+    creates a folder for each experiment ID and stores metadata in ``meta.yaml``. Runs are stored
+    as subfolders.
+    """
+    store = _get_store()
+    exp_id = store.create_experiment(experiment_name, artifact_location)
+    click.echo(f"Created experiment '{experiment_name}' with id {exp_id}")
+
+
+@commands.command("search")
+@click.option(
+    "--view",
+    "-v",
+    default="active_only",
+    help="Select view type for experiments. Valid view types are "
+    "'active_only' (default), 'deleted_only', and 'all'.",
+)
+def search_experiments(view):
+    """
+    Search for experiments in the configured tracking server.
+    """
+    view_type = ViewType.from_string(view) if view else ViewType.ACTIVE_ONLY
+    experiments = mlflow.search_experiments(view_type=view_type)
+    table = [
+        [
+            exp.experiment_id,
+            exp.name,
+            exp.artifact_location
+            if is_uri(exp.artifact_location)
+            else os.path.abspath(exp.artifact_location),
+        ]
+        for exp in experiments
+    ]
+    click.echo(_create_table(sorted(table), headers=["Experiment Id", "Name", "Artifact Location"]))
+
+
+@commands.command("delete")
+@EXPERIMENT_ID
+def delete_experiment(experiment_id):
+    """
+    Mark an active experiment for deletion. This also applies to experiment's metadata, runs and
+    associated data, and artifacts if they are store in default location. Use ``list`` command to
+    view artifact location. Command will throw an error if experiment is not found or already
+    marked for deletion.
+
+    Experiments marked for deletion can be restored using ``restore`` command, unless they are
+    permanently deleted.
+
+    Specific implementation of deletion is dependent on backend stores. ``FileStore`` moves
+    experiments marked for deletion under a ``.trash`` folder under the main folder used to
+    instantiate ``FileStore``. Experiments marked for deletion can be permanently deleted by
+    clearing the ``.trash`` folder. It is recommended to use a ``cron`` job or an alternate
+    workflow mechanism to clear ``.trash`` folder.
+    """
+    store = _get_store()
+    store.delete_experiment(experiment_id)
+    click.echo(f"Experiment with ID {experiment_id} has been deleted.")
+
+
+@commands.command("restore")
+@EXPERIMENT_ID
+def restore_experiment(experiment_id):
+    """
+    Restore a deleted experiment. This also applies to experiment's metadata, runs and associated
+    data. The command throws an error if the experiment is already active, cannot be found, or
+    permanently deleted.
+    """
+    store = _get_store()
+    store.restore_experiment(experiment_id)
+    click.echo(f"Experiment with id {experiment_id} has been restored.")
+
+
+@commands.command("rename")
+@EXPERIMENT_ID
+@click.option("--new-name", type=click.STRING, required=True)
+def rename_experiment(experiment_id, new_name):
+    """
+    Renames an active experiment.
+    Returns an error if the experiment is inactive.
+    """
+    store = _get_store()
+    store.rename_experiment(experiment_id, new_name)
+    click.echo(f"Experiment with id {experiment_id} has been renamed to '{new_name}'.")
+
+
+@commands.command("csv")
+@EXPERIMENT_ID
+@click.option("--filename", "-o", type=click.STRING)
+def generate_csv_with_runs(experiment_id, filename):
+    # type: (str, str) -> None
+    """
+    Generate CSV with all runs for an experiment
+    """
+    runs = fluent.search_runs(experiment_ids=experiment_id)
+    if filename:
+        runs.to_csv(filename, index=False)
+        click.echo(
+            f"Experiment with ID {experiment_id} has been exported as a CSV to file: {filename}."
+        )
+    else:
+        click.echo(runs.to_csv(index=False))

mlflow/gemini/__init__.py

@@ -0,0 +1,81 @@
+"""
+The ``mlflow.gemini`` module provides an API for tracing the interaction with Gemini models.
+"""
+
+from mlflow.gemini.autolog import (
+    patched_class_call,
+    patched_module_call,
+)
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+FLAVOR_NAME = "gemini"
+
+
+@experimental(version="2.19.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Gemini to MLflow.
+    Currently, both legacy SDK google-generativeai and new SDK google-genai are supported.
+    Only synchronous calls are supported. Asynchronous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Gemini models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Gemini autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Gemini
+            autologging. If ``False``, show all events and warnings.
+    """
+    try:
+        from google import generativeai
+
+        for method in ["generate_content", "count_tokens"]:
+            safe_patch(
+                FLAVOR_NAME,
+                generativeai.GenerativeModel,
+                method,
+                patched_class_call,
+            )
+
+        safe_patch(
+            FLAVOR_NAME,
+            generativeai.ChatSession,
+            "send_message",
+            patched_class_call,
+        )
+
+        safe_patch(
+            FLAVOR_NAME,
+            generativeai,
+            "embed_content",
+            patched_module_call,
+        )
+    except ImportError:
+        pass
+
+    try:
+        from google import genai
+
+        # Since the genai SDK calls "_generate_content" iteratively within "generate_content",
+        # we need to patch both "generate_content" and "_generate_content".
+        for method in ["generate_content", "_generate_content", "count_tokens", "embed_content"]:
+            safe_patch(
+                FLAVOR_NAME,
+                genai.models.Models,
+                method,
+                patched_class_call,
+            )
+
+        safe_patch(
+            FLAVOR_NAME,
+            genai.chats.Chat,
+            "send_message",
+            patched_class_call,
+        )
+    except ImportError:
+        pass