genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/process.py

@@ -0,0 +1,182 @@
+import functools
+import os
+import subprocess
+import sys
+
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.os import is_windows
+
+
+class ShellCommandException(Exception):
+    @classmethod
+    def from_completed_process(cls, process):
+        lines = [
+            f"Non-zero exit code: {process.returncode}",
+            f"Command: {process.args}",
+        ]
+        if process.stdout:
+            lines += [
+                "",
+                "STDOUT:",
+                process.stdout,
+            ]
+        if process.stderr:
+            lines += [
+                "",
+                "STDERR:",
+                process.stderr,
+            ]
+        return cls("\n".join(lines))
+
+
+def _remove_inaccessible_python_path(env):
+    """
+    Remove inaccessible path from PYTHONPATH environment variable.
+    """
+    if python_path := env.get("PYTHONPATH"):
+        paths = [p for p in python_path.split(":") if os.access(p, os.R_OK)]
+        env["PYTHONPATH"] = ":".join(paths)
+    return env
+
+
+def _exec_cmd(
+    cmd,
+    *,
+    throw_on_error=True,
+    extra_env=None,
+    capture_output=True,
+    synchronous=True,
+    stream_output=False,
+    **kwargs,
+):
+    """A convenience wrapper of `subprocess.Popen` for running a command from a Python script.
+
+    Args:
+        cmd: The command to run, as a string or a list of strings.
+        throw_on_error: If True, raises an Exception if the exit code of the program is nonzero.
+        extra_env: Extra environment variables to be defined when running the child process.
+            If this argument is specified, `kwargs` cannot contain `env`.
+        capture_output: If True, stdout and stderr will be captured and included in an exception
+            message on failure; if False, these streams won't be captured.
+        synchronous: If True, wait for the command to complete and return a CompletedProcess
+            instance, If False, does not wait for the command to complete and return
+            a Popen instance, and ignore the `throw_on_error` argument.
+        stream_output: If True, stream the command's stdout and stderr to `sys.stdout`
+            as a unified stream during execution.
+            If False, do not stream the command's stdout and stderr to `sys.stdout`.
+        kwargs: Keyword arguments (except `text`) passed to `subprocess.Popen`.
+
+    Returns:
+        If synchronous is True, return a `subprocess.CompletedProcess` instance,
+        otherwise return a Popen instance.
+
+    """
+
+    illegal_kwargs = set(kwargs.keys()).intersection({"text"})
+    if illegal_kwargs:
+        raise ValueError(f"`kwargs` cannot contain {list(illegal_kwargs)}")
+
+    env = kwargs.pop("env", None)
+    if extra_env is not None and env is not None:
+        raise ValueError("`extra_env` and `env` cannot be used at the same time")
+
+    if capture_output and stream_output:
+        raise ValueError(
+            "`capture_output=True` and `stream_output=True` cannot be specified at the same time"
+        )
+
+    # Copy current `os.environ` or passed in `env` to avoid mutating it.
+    env = env or os.environ.copy()
+    if extra_env is not None:
+        env.update(extra_env)
+
+    if is_in_databricks_runtime():
+        # in databricks runtime, the PYTHONPATH might contain inaccessible path
+        # which causes virtualenv python environment creation subprocess failure.
+        # as a workaround, we remove inaccessible path out of python path.
+        env = _remove_inaccessible_python_path(env)
+
+    # In Python < 3.8, `subprocess.Popen` doesn't accept a command containing path-like
+    # objects (e.g. `["ls", pathlib.Path("abc")]`) on Windows. To avoid this issue,
+    # stringify all elements in `cmd`. Note `str(pathlib.Path("abc"))` returns 'abc'.
+    if isinstance(cmd, list):
+        cmd = list(map(str, cmd))
+
+    if capture_output or stream_output:
+        if kwargs.get("stdout") is not None or kwargs.get("stderr") is not None:
+            raise ValueError(
+                "stdout and stderr arguments may not be used with capture_output or stream_output"
+            )
+        kwargs["stdout"] = subprocess.PIPE
+        if capture_output:
+            kwargs["stderr"] = subprocess.PIPE
+        elif stream_output:
+            # Redirect stderr to stdout in order to combine the streams for unified printing to
+            # `sys.stdout`, as documented in
+            # https://docs.python.org/3/library/subprocess.html#subprocess.run
+            kwargs["stderr"] = subprocess.STDOUT
+
+    process = subprocess.Popen(
+        cmd,
+        env=env,
+        text=True,
+        **kwargs,
+    )
+    if not synchronous:
+        return process
+
+    if stream_output:
+        for output_char in iter(lambda: process.stdout.read(1), ""):
+            sys.stdout.write(output_char)
+
+    stdout, stderr = process.communicate()
+    returncode = process.poll()
+    comp_process = subprocess.CompletedProcess(
+        process.args,
+        returncode=returncode,
+        stdout=stdout,
+        stderr=stderr,
+    )
+    if throw_on_error and returncode != 0:
+        raise ShellCommandException.from_completed_process(comp_process)
+    return comp_process
+
+
+def _join_commands(*commands):
+    entry_point = ["bash", "-c"] if not is_windows() else ["cmd", "/c"]
+    sep = " && " if not is_windows() else " & "
+    return [*entry_point, sep.join(map(str, commands))]
+
+
+# A global map storing (function, args_tuple) --> (value, pid)
+_per_process_value_cache_map = {}
+
+
+def cache_return_value_per_process(fn):
+    """
+    A decorator which globally caches the return value of the decorated function.
+    But if current process forked out a new child process, in child process,
+    old cache values are invalidated.
+
+    Restrictions: The decorated function must be called with only positional arguments,
+    and all the argument values must be hashable.
+    """
+
+    @functools.wraps(fn)
+    def wrapped_fn(*args, **kwargs):
+        if len(kwargs) > 0:
+            raise ValueError(
+                "The function decorated by `cache_return_value_per_process` is not allowed to be "
+                "called with key-word style arguments."
+            )
+        if (fn, args) in _per_process_value_cache_map:
+            prev_value, prev_pid = _per_process_value_cache_map.get((fn, args))
+            if os.getpid() == prev_pid:
+                return prev_value
+
+        new_value = fn(*args)
+        new_pid = os.getpid()
+        _per_process_value_cache_map[(fn, args)] = (new_value, new_pid)
+        return new_value
+
+    return wrapped_fn

mlflow/utils/promptlab_utils.py

@@ -0,0 +1,146 @@
+import json
+import os
+import tempfile
+import time
+from datetime import datetime
+
+from mlflow.entities.param import Param
+from mlflow.entities.run_status import RunStatus
+from mlflow.entities.run_tag import RunTag
+from mlflow.utils.file_utils import make_containing_dirs, write_to
+from mlflow.utils.mlflow_tags import MLFLOW_LOGGED_ARTIFACTS, MLFLOW_RUN_SOURCE_TYPE
+from mlflow.version import VERSION as __version__
+
+
+def create_eval_results_json(prompt_parameters, model_input, model_output_parameters, model_output):
+    columns = [param.key for param in prompt_parameters] + ["prompt", "output"]
+    data = [param.value for param in prompt_parameters] + [model_input, model_output]
+
+    updated_columns = columns + [param.key for param in model_output_parameters]
+    updated_data = data + [param.value for param in model_output_parameters]
+
+    eval_results = {"columns": updated_columns, "data": [updated_data]}
+
+    return json.dumps(eval_results)
+
+
+def _create_promptlab_run_impl(
+    store,
+    experiment_id: str,
+    run_name: str,
+    tags: list[RunTag],
+    prompt_template: str,
+    prompt_parameters: list[Param],
+    model_route: str,
+    model_parameters: list[Param],
+    model_input: str,
+    model_output_parameters: list[Param],
+    model_output: str,
+    mlflow_version: str,
+    user_id: str,
+    start_time: str,
+):
+    run = store.create_run(experiment_id, user_id, start_time, tags, run_name)
+    run_id = run.info.run_id
+
+    try:
+        prompt_parameters = [
+            Param(key=param.key, value=str(param.value)) for param in prompt_parameters
+        ]
+        model_parameters = [
+            Param(key=param.key, value=str(param.value)) for param in model_parameters
+        ]
+        model_output_parameters = [
+            Param(key=param.key, value=str(param.value)) for param in model_output_parameters
+        ]
+
+        # log model parameters
+        parameters_to_log = [
+            *model_parameters,
+            Param("model_route", model_route),
+            Param("prompt_template", prompt_template),
+        ]
+
+        tags_to_log = [
+            RunTag(
+                MLFLOW_LOGGED_ARTIFACTS,
+                json.dumps([{"path": "eval_results_table.json", "type": "table"}]),
+            ),
+            RunTag(MLFLOW_RUN_SOURCE_TYPE, "PROMPT_ENGINEERING"),
+        ]
+
+        store.log_batch(run_id, [], parameters_to_log, tags_to_log)
+
+        # log model
+        from mlflow.models import Model
+
+        artifact_dir = store.get_run(run_id).info.artifact_uri
+
+        utc_time_created = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S.%f")
+        promptlab_model = Model(
+            artifact_path="model",
+            run_id=run_id,
+            utc_time_created=utc_time_created,
+        )
+        store.record_logged_model(run_id, promptlab_model)
+
+        try:
+            from mlflow.models.signature import ModelSignature
+            from mlflow.types.schema import ColSpec, DataType, Schema
+        except ImportError:
+            signature = None
+        else:
+            inputs_colspecs = [ColSpec(DataType.string, param.key) for param in prompt_parameters]
+            outputs_colspecs = [ColSpec(DataType.string, "output")]
+            signature = ModelSignature(
+                inputs=Schema(inputs_colspecs),
+                outputs=Schema(outputs_colspecs),
+            )
+
+        from mlflow.prompt.promptlab_model import save_model
+        from mlflow.server.handlers import (
+            _get_artifact_repo_mlflow_artifacts,
+            _get_proxied_run_artifact_destination_path,
+            _is_servable_proxied_run_artifact_root,
+        )
+
+        # write artifact files
+        from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+
+        with tempfile.TemporaryDirectory() as local_dir:
+            save_model(
+                mlflow_model=promptlab_model,
+                path=os.path.join(local_dir, "model"),
+                signature=signature,
+                input_example={"inputs": [param.value for param in prompt_parameters]},
+                prompt_template=prompt_template,
+                prompt_parameters=prompt_parameters,
+                model_parameters=model_parameters,
+                model_route=model_route,
+                pip_requirements=[f"mlflow[gateway]=={__version__}"],
+            )
+
+            eval_results_json = create_eval_results_json(
+                prompt_parameters, model_input, model_output_parameters, model_output
+            )
+            eval_results_json_file_path = os.path.join(local_dir, "eval_results_table.json")
+            make_containing_dirs(eval_results_json_file_path)
+            write_to(eval_results_json_file_path, eval_results_json)
+
+            if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
+                artifact_repo = _get_artifact_repo_mlflow_artifacts()
+                artifact_path = _get_proxied_run_artifact_destination_path(
+                    proxied_artifact_root=run.info.artifact_uri,
+                )
+                artifact_repo.log_artifacts(local_dir, artifact_path=artifact_path)
+            else:
+                artifact_repo = get_artifact_repository(artifact_dir)
+                artifact_repo.log_artifacts(local_dir)
+
+    except Exception:
+        store.update_run_info(run_id, RunStatus.FAILED, int(time.time() * 1000), run_name)
+    else:
+        # end time is the current number of milliseconds since the UNIX epoch.
+        store.update_run_info(run_id, RunStatus.FINISHED, int(time.time() * 1000), run_name)
+
+    return store.get_run(run_id=run_id)