genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/name_utils.py

@@ -0,0 +1,346 @@
+import random
+import uuid
+
+_EXPERIMENT_ID_FIXED_WIDTH = 18
+
+
+def _generate_unique_integer_id():
+    """Utility function for generating a random fixed-length integer
+
+    Args:
+        id_length: The target length of the string representation of the integer without
+            leading zeros
+
+    Returns:
+        a fixed-width integer
+    """
+
+    random_int = uuid.uuid4().int
+    # Cast to string to get a fixed length
+    random_str = str(random_int)[-_EXPERIMENT_ID_FIXED_WIDTH:]
+    # append a random int as string to the end of the generated string for as many
+    # leading zeros exist in the generated string in order to preserve the total length
+    # once cast back to int
+    for s in random_str:
+        if s == "0":
+            random_str = random_str + str(random.randint(0, 9))
+        else:
+            break
+    return int(random_str)
+
+
+def _generate_string(sep, integer_scale):
+    predicate = random.choice(_GENERATOR_PREDICATES).lower()
+    noun = random.choice(_GENERATOR_NOUNS).lower()
+    num = random.randint(0, 10**integer_scale)
+    return f"{predicate}{sep}{noun}{sep}{num}"
+
+
+def _generate_random_name(sep="-", integer_scale=3, max_length=20):
+    """Helper function for generating a random predicate, noun, and integer combination
+
+    Args:
+        sep: String separator for word spacing.
+        integer_scale: Dictates the maximum scale range for random integer sampling (power of 10).
+        max_length: Maximum allowable string length.
+
+    Returns:
+        A random string phrase comprised of a predicate, noun, and random integer.
+
+    """
+    name = None
+    for _ in range(10):
+        name = _generate_string(sep, integer_scale)
+        if len(name) <= max_length:
+            return name
+    # If the combined length isn't below the threshold after 10 iterations, truncate it.
+    return name[:max_length]
+
+
+_GENERATOR_NOUNS = [
+    "ant",
+    "ape",
+    "asp",
+    "auk",
+    "bass",
+    "bat",
+    "bear",
+    "bee",
+    "bird",
+    "boar",
+    "bug",
+    "calf",
+    "carp",
+    "cat",
+    "chimp",
+    "cod",
+    "colt",
+    "conch",
+    "cow",
+    "crab",
+    "crane",
+    "croc",
+    "crow",
+    "cub",
+    "deer",
+    "doe",
+    "dog",
+    "dolphin",
+    "donkey",
+    "dove",
+    "duck",
+    "eel",
+    "elk",
+    "fawn",
+    "finch",
+    "fish",
+    "flea",
+    "fly",
+    "foal",
+    "fowl",
+    "fox",
+    "frog",
+    "gnat",
+    "gnu",
+    "goat",
+    "goose",
+    "grouse",
+    "grub",
+    "gull",
+    "hare",
+    "hawk",
+    "hen",
+    "hog",
+    "horse",
+    "hound",
+    "jay",
+    "kit",
+    "kite",
+    "koi",
+    "lamb",
+    "lark",
+    "loon",
+    "lynx",
+    "mare",
+    "midge",
+    "mink",
+    "mole",
+    "moose",
+    "moth",
+    "mouse",
+    "mule",
+    "newt",
+    "owl",
+    "ox",
+    "panda",
+    "penguin",
+    "perch",
+    "pig",
+    "pug",
+    "quail",
+    "ram",
+    "rat",
+    "ray",
+    "robin",
+    "roo",
+    "rook",
+    "seal",
+    "shad",
+    "shark",
+    "sheep",
+    "shoat",
+    "shrew",
+    "shrike",
+    "shrimp",
+    "skink",
+    "skunk",
+    "sloth",
+    "slug",
+    "smelt",
+    "snail",
+    "snake",
+    "snipe",
+    "sow",
+    "sponge",
+    "squid",
+    "squirrel",
+    "stag",
+    "steed",
+    "stoat",
+    "stork",
+    "swan",
+    "tern",
+    "toad",
+    "trout",
+    "turtle",
+    "vole",
+    "wasp",
+    "whale",
+    "wolf",
+    "worm",
+    "wren",
+    "yak",
+    "zebra",
+]
+
+_GENERATOR_PREDICATES = [
+    "abundant",
+    "able",
+    "abrasive",
+    "adorable",
+    "adaptable",
+    "adventurous",
+    "aged",
+    "agreeable",
+    "ambitious",
+    "amazing",
+    "amusing",
+    "angry",
+    "auspicious",
+    "awesome",
+    "bald",
+    "beautiful",
+    "bemused",
+    "bedecked",
+    "big",
+    "bittersweet",
+    "blushing",
+    "bold",
+    "bouncy",
+    "brawny",
+    "bright",
+    "burly",
+    "bustling",
+    "calm",
+    "capable",
+    "carefree",
+    "capricious",
+    "caring",
+    "casual",
+    "charming",
+    "chill",
+    "classy",
+    "clean",
+    "clumsy",
+    "colorful",
+    "crawling",
+    "dapper",
+    "debonair",
+    "dashing",
+    "defiant",
+    "delicate",
+    "delightful",
+    "dazzling",
+    "efficient",
+    "enchanting",
+    "entertaining",
+    "enthused",
+    "exultant",
+    "fearless",
+    "flawless",
+    "fortunate",
+    "fun",
+    "funny",
+    "gaudy",
+    "gentle",
+    "gifted",
+    "glamorous",
+    "grandiose",
+    "gregarious",
+    "handsome",
+    "hilarious",
+    "honorable",
+    "illustrious",
+    "incongruous",
+    "indecisive",
+    "industrious",
+    "intelligent",
+    "inquisitive",
+    "intrigued",
+    "invincible",
+    "judicious",
+    "kindly",
+    "languid",
+    "learned",
+    "legendary",
+    "likeable",
+    "loud",
+    "luminous",
+    "luxuriant",
+    "lyrical",
+    "magnificent",
+    "marvelous",
+    "masked",
+    "melodic",
+    "merciful",
+    "mercurial",
+    "monumental",
+    "mysterious",
+    "nebulous",
+    "nervous",
+    "nimble",
+    "nosy",
+    "omniscient",
+    "orderly",
+    "overjoyed",
+    "peaceful",
+    "painted",
+    "persistent",
+    "placid",
+    "polite",
+    "popular",
+    "powerful",
+    "puzzled",
+    "rambunctious",
+    "rare",
+    "rebellious",
+    "respected",
+    "resilient",
+    "righteous",
+    "receptive",
+    "redolent",
+    "resilient",
+    "rogue",
+    "rumbling",
+    "salty",
+    "sassy",
+    "secretive",
+    "selective",
+    "sedate",
+    "serious",
+    "shivering",
+    "skillful",
+    "sincere",
+    "skittish",
+    "silent",
+    "smiling",
+    "sneaky",
+    "sophisticated",
+    "spiffy",
+    "stately",
+    "suave",
+    "stylish",
+    "tasteful",
+    "thoughtful",
+    "thundering",
+    "traveling",
+    "treasured",
+    "trusting",
+    "unequaled",
+    "upset",
+    "unique",
+    "unleashed",
+    "useful",
+    "upbeat",
+    "unruly",
+    "valuable",
+    "vaunted",
+    "victorious",
+    "welcoming",
+    "whimsical",
+    "wistful",
+    "wise",
+    "worried",
+    "youthful",
+    "zealous",
+]

mlflow/utils/nfs_on_spark.py

@@ -0,0 +1,62 @@
+import os
+import shutil
+import uuid
+
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils.databricks_utils import (
+    get_databricks_nfs_temp_dir,
+    is_databricks_connect,
+    is_in_databricks_runtime,
+    is_in_databricks_serverless_runtime,
+)
+
+# Set spark config "spark.mlflow.nfs.rootDir" to specify a NFS (network file system) directory
+# which is shared with all spark cluster nodes.
+# This will help optimize routine of distributing spark driver files to remote workers.
+# None represent no NFS directory available.
+# Note:
+#  1. If NFS directory set, you must ensure all spark cluster nodes using the same hardware and
+#  installed the same OS with the same environment configured, because mlflow uses NFS directory
+#  to distribute driver side virtual environment to remote workers if NFS available, heterogeneous
+#  cluster nodes might cause issues under the case.
+#  2. The NFS directory must be mounted before importing mlflow.
+#  3. For databricks users, don't set this config, databricks product sets up internal NFS service
+#  automatically.
+_NFS_CACHE_ROOT_DIR = None
+
+
+def get_nfs_cache_root_dir():
+    if is_in_databricks_runtime():
+        spark_sess = _get_active_spark_session()
+        if is_in_databricks_serverless_runtime():
+            # Databricks Serverless runtime VM can't access NFS.
+            nfs_enabled = False
+        else:
+            nfs_enabled = spark_sess and (
+                spark_sess.conf.get("spark.databricks.mlflow.nfs.enabled", "true").lower() == "true"
+            )
+        if nfs_enabled:
+            try:
+                return get_databricks_nfs_temp_dir()
+            except Exception:
+                nfs_root_dir = "/local_disk0/.ephemeral_nfs"
+                # Test whether the NFS directory is writable.
+                test_path = os.path.join(nfs_root_dir, uuid.uuid4().hex)
+                try:
+                    os.makedirs(test_path)
+                    return nfs_root_dir
+                except Exception:
+                    # For databricks cluster enabled Table ACL, we have no permission to access NFS
+                    # directory, in this case, return None, meaning NFS is not available.
+                    return None
+                finally:
+                    shutil.rmtree(test_path, ignore_errors=True)
+        else:
+            return None
+    else:
+        spark_session = _get_active_spark_session()
+        if is_databricks_connect(spark_session):
+            # Remote spark connect client can't access Databricks Serverless cluster NFS.
+            return None
+        if spark_session is not None:
+            return spark_session.conf.get("spark.mlflow.nfs.rootDir", None)

mlflow/utils/openai_utils.py

@@ -0,0 +1,164 @@
+import os
+import time
+from enum import Enum
+from typing import NamedTuple, Optional
+
+import mlflow
+
+REQUEST_URL_CHAT = "https://api.openai.com/v1/chat/completions"
+REQUEST_URL_COMPLETIONS = "https://api.openai.com/v1/completions"
+REQUEST_URL_EMBEDDINGS = "https://api.openai.com/v1/embeddings"
+
+REQUEST_FIELDS_CHAT = {
+    "model",
+    "messages",
+    "frequency_penalty",
+    "logit_bias",
+    "max_tokens",
+    "n",
+    "presence_penalty",
+    "response_format",
+    "seed",
+    "stop",
+    "stream",
+    "temperature",
+    "top_p",
+    "tools",
+    "tool_choice",
+    "user",
+    "function_call",
+    "functions",
+}
+REQUEST_FIELDS_COMPLETIONS = {
+    "model",
+    "prompt",
+    "best_of",
+    "echo",
+    "frequency_penalty",
+    "logit_bias",
+    "logprobs",
+    "max_tokens",
+    "n",
+    "presence_penalty",
+    "seed",
+    "stop",
+    "stream",
+    "suffix",
+    "temperature",
+    "top_p",
+    "user",
+}
+REQUEST_FIELDS_EMBEDDINGS = {"input", "model", "encoding_format", "user"}
+REQUEST_FIELDS = REQUEST_FIELDS_CHAT | REQUEST_FIELDS_COMPLETIONS | REQUEST_FIELDS_EMBEDDINGS
+
+
+def _validate_model_params(task, model, params):
+    if not params:
+        return
+
+    if any(key in model for key in params):
+        raise mlflow.MlflowException.invalid_parameter_value(
+            f"Providing any of {list(model.keys())} as parameters in the signature is not "
+            "allowed because they were indicated as part of the OpenAI model. Either remove "
+            "the argument when logging the model or remove the parameter from the signature.",
+        )
+    if "batch_size" in params and task == "chat.completions":
+        raise mlflow.MlflowException.invalid_parameter_value(
+            "Parameter `batch_size` is not supported for task `chat.completions`"
+        )
+
+
+class _OAITokenHolder:
+    def __init__(self, api_type):
+        self._credential = None
+        self._api_type = api_type
+        self._is_azure_ad = api_type in ("azure_ad", "azuread")
+        self._azure_ad_token = None
+        self._api_token_env = os.environ.get("OPENAI_API_KEY")
+
+        if self._is_azure_ad and not self._api_token_env:
+            try:
+                from azure.identity import DefaultAzureCredential
+            except ImportError:
+                raise mlflow.MlflowException(
+                    "Using API type `azure_ad` or `azuread` requires the package"
+                    " `azure-identity` to be installed."
+                )
+            self._credential = DefaultAzureCredential()
+
+    @property
+    def token(self):
+        return self._api_token_env or self._azure_ad_token.token
+
+    def refresh(self, logger=None):
+        """Validates the token or API key configured for accessing the OpenAI resource."""
+
+        if self._api_token_env is not None:
+            return
+
+        if self._is_azure_ad:
+            if not self._azure_ad_token or self._azure_ad_token.expires_on < time.time() + 60:
+                from azure.core.exceptions import ClientAuthenticationError
+
+                if logger:
+                    logger.debug(
+                        "Token for Azure AD is either expired or unset. Attempting to "
+                        "acquire a new token."
+                    )
+                try:
+                    self._azure_ad_token = self._credential.get_token(
+                        "https://cognitiveservices.azure.com/.default"
+                    )
+                except ClientAuthenticationError as err:
+                    raise mlflow.MlflowException(
+                        "Unable to acquire a valid Azure AD token for the resource due to "
+                        f"the following error: {err.message}"
+                    ) from err
+
+            if logger:
+                logger.debug("Token refreshed successfully")
+        else:
+            raise mlflow.MlflowException(
+                "OpenAI API key must be set in the ``OPENAI_API_KEY`` environment variable."
+            )
+
+
+class _OpenAIApiConfig(NamedTuple):
+    api_type: str
+    batch_size: int
+    max_requests_per_minute: int
+    max_tokens_per_minute: int
+    api_version: Optional[str]
+    api_base: str
+    deployment_id: Optional[str]
+    organization: Optional[str] = None
+    max_retries: int = 5
+    timeout: float = 60.0
+
+
+# See https://github.com/openai/openai-python/blob/cf03fe16a92cd01f2a8867537399c12e183ba58e/openai/__init__.py#L30-L38
+# for the list of environment variables that openai-python uses
+class _OpenAIEnvVar(str, Enum):
+    OPENAI_API_TYPE = "OPENAI_API_TYPE"
+    OPENAI_BASE_URL = "OPENAI_BASE_URL"
+    OPENAI_API_BASE = "OPENAI_API_BASE"
+    OPENAI_API_KEY = "OPENAI_API_KEY"
+    OPENAI_API_KEY_PATH = "OPENAI_API_KEY_PATH"
+    OPENAI_API_VERSION = "OPENAI_API_VERSION"
+    OPENAI_ORGANIZATION = "OPENAI_ORGANIZATION"
+    OPENAI_ENGINE = "OPENAI_ENGINE"
+    # use deployment_name instead of deployment_id to be
+    # consistent with gateway
+    OPENAI_DEPLOYMENT_NAME = "OPENAI_DEPLOYMENT_NAME"
+
+    @property
+    def secret_key(self):
+        return self.value.lower()
+
+    @classmethod
+    def read_environ(cls):
+        env_vars = {}
+        for e in _OpenAIEnvVar:
+            if value := os.getenv(e.value):
+                env_vars[e.value] = value
+        return env_vars

mlflow/utils/os.py

@@ -0,0 +1,12 @@
+import os
+
+
+def is_windows():
+    """
+    Returns true if the local system/OS name is Windows.
+
+    Returns:
+        True if the local system/OS name is Windows.
+
+    """
+    return os.name == "nt"

mlflow/utils/oss_registry_utils.py

@@ -0,0 +1,29 @@
+import urllib.parse
+
+from mlflow.environment_variables import MLFLOW_UC_OSS_TOKEN
+from mlflow.exceptions import MlflowException
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.rest_utils import MlflowHostCreds
+from mlflow.utils.uri import (
+    _DATABRICKS_UNITY_CATALOG_SCHEME,
+)
+
+
+def get_oss_host_creds(server_uri=None):
+    """
+    Retrieve the host credentials for the OSS server.
+
+    Args:
+        server_uri (str): The URI of the server.
+
+    Returns:
+        MlflowHostCreds: The host credentials for the OSS server.
+    """
+    parsed_uri = urllib.parse.urlparse(server_uri)
+
+    if parsed_uri.scheme != "uc":
+        raise MlflowException("The scheme of the server_uri should be 'uc'")
+
+    if parsed_uri.path == _DATABRICKS_UNITY_CATALOG_SCHEME:
+        return get_databricks_host_creds(parsed_uri.path)
+    return MlflowHostCreds(host=parsed_uri.path, token=MLFLOW_UC_OSS_TOKEN.get())

mlflow/utils/plugins.py

@@ -0,0 +1,17 @@
+import importlib.metadata
+import sys
+
+
+def _get_entry_points(group: str) -> list[importlib.metadata.EntryPoint]:
+    if sys.version_info >= (3, 10):
+        return importlib.metadata.entry_points(group=group)
+
+    entrypoints = importlib.metadata.entry_points()
+    try:
+        return entrypoints.get(group, [])
+    except AttributeError:
+        return entrypoints.select(group=group)
+
+
+def get_entry_points(group: str) -> list[importlib.metadata.EntryPoint]:
+    return _get_entry_points(group)