pixi_plugin.py

import json
import logging

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Literal

from ray._private.runtime_env.conda_utils import exec_cmd
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
from ray.exceptions import RuntimeEnvSetupError


if TYPE_CHECKING:
    from ray._private.runtime_env.context import RuntimeEnvContext
    from ray.runtime_env.runtime_env import RuntimeEnv


def get_pixi_env_list() -> list:
    """
    Get pixi env list.
    """
    try:
        exec_cmd(["pixi", "--help"], throw_on_error=False)
    except OSError as err:
        raise ValueError("Could not find pixi executable.") from err

    _, stdout, _ = exec_cmd(["pixi", "info", "--json"])
    return json.loads(stdout)["environments_info"]


@dataclass
class PixiRuntimeEnvConfig:
    env: str
    execution_method: Literal["pixi-run", "pixi-prefix"] = "pixi-prefix"


class PixiPlugin(RuntimeEnvPlugin):
    name = "pixi"

    def modify_context(
        self,
        uris: list[str],
        runtime_env: "RuntimeEnv",
        context: "RuntimeEnvContext",
        logger: logging.Logger,
    ):
        if "pixi" not in runtime_env:
            return
        runtime_env_pixi = runtime_env["pixi"]

        pixi_config = PixiRuntimeEnvConfig(
            env=runtime_env_pixi["env"],
            execution_method=runtime_env_pixi.get("execution_method", "pixi-prefix"),
        )

        # Dictionary of all pixi environments in current folder
        envs_prefix = {e["name"]: e["prefix"] for e in get_pixi_env_list()}

        if pixi_config.env not in envs_prefix:
            msg = f"The pixi environment {pixi_config.env} does not exist."
            raise RuntimeEnvSetupError(msg)

        if pixi_config.execution_method == "pixi-run":
            logger.info("Using `pixi run` to execute Python code in the %s environment.", pixi_config.env)
            context.py_executable = f"pixi run --frozen --no-progress -e {pixi_config.env} python"
        elif pixi_config.execution_method == "pixi-prefix":
            logger.info("Using the %s environment prefix to execute Python code.", pixi_config.env)
            context.py_executable = str(Path(envs_prefix[pixi_config.env]) / "bin" / "python")
        else:
            msg = (
                f"Unsupported execution method: {pixi_config.execution_method}. "
                "The supported methods are 'pixi-run' and 'pixi-prefix'."
            )
            raise RuntimeEnvSetupError(msg)

@cynddl Can you add a license? 🙏