diff --git a/clearml_agent/helper/package/uv_api.py b/clearml_agent/helper/package/uv_api.py new file mode 100644 index 0000000..295d018 --- /dev/null +++ b/clearml_agent/helper/package/uv_api.py @@ -0,0 +1,223 @@ +from copy import deepcopy +from functools import wraps + +import attr +import sys +import os +from pathlib2 import Path + +from clearml_agent.definitions import ENV_AGENT_FORCE_UV +from clearml_agent.helper.process import Argv, DEVNULL, check_if_command_exists +from clearml_agent.session import Session, UV + + +def prop_guard(prop, log_prop=None): + assert isinstance(prop, property) + assert not log_prop or isinstance(log_prop, property) + + def decorator(func): + message = "%s:%s calling {}, {} = %s".format(func.__name__, prop.fget.__name__) + + @wraps(func) + def new_func(self, *args, **kwargs): + prop_value = prop.fget(self) + if log_prop: + log_prop.fget(self).debug( + message, + type(self).__name__, + "" if prop_value else " not", + prop_value, + ) + if prop_value: + return func(self, *args, **kwargs) + + return new_func + + return decorator + + +class UvConfig: + def __init__(self, session): + # type: (Session, str) -> None + self.session = session + self._log = session.get_logger(__name__) + self._python = ( + sys.executable + ) # default, overwritten from session config in initialize() + self._initialized = False + + @property + def log(self): + return self._log + + @property + def enabled(self): + return ( + ENV_AGENT_FORCE_UV.get() + or self.session.config["agent.package_manager.type"] == UV + ) + + _guard_enabled = prop_guard(enabled, log) + + def run(self, *args, **kwargs): + func = kwargs.pop("func", Argv.get_output) + kwargs.setdefault("stdin", DEVNULL) + kwargs["env"] = deepcopy(os.environ) + if "VIRTUAL_ENV" in kwargs["env"] or "CONDA_PREFIX" in kwargs["env"]: + kwargs["env"].pop("VIRTUAL_ENV", None) + kwargs["env"].pop("CONDA_PREFIX", None) + kwargs["env"].pop("PYTHONPATH", None) + if hasattr(sys, "real_prefix") and hasattr(sys, "base_prefix"): + path = ":" + kwargs["env"]["PATH"] + path = path.replace(":" + sys.base_prefix, ":" + sys.real_prefix, 1) + kwargs["env"]["PATH"] = path + + if self.session and self.session.config and args and args[0] == "sync": + extra_args = self.session.config.get( + "agent.package_manager.uv_install_extra_args", None + ) + if extra_args: + args = args + tuple(extra_args) + + if check_if_command_exists("uv"): + argv = Argv("uv", *args) + else: + argv = Argv(self._python, "-m", "uv", *args) + self.log.debug("running: %s", argv) + return func(argv, **kwargs) + + @_guard_enabled + def initialize(self, cwd=None): + if not self._initialized: + # use correct python version -- detected in Worker.install_virtualenv() and written to + # session + if self.session.config.get("agent.python_binary", None): + self._python = self.session.config.get("agent.python_binary") + + if ( + self.session.config.get("agent.package_manager.uv_version", None) + is not None + ): + version = str( + self.session.config.get("agent.package_manager.uv_version") + ) + + # get uv version + version = version.replace(" ", "") + if ( + ("=" in version) + or ("~" in version) + or ("<" in version) + or (">" in version) + ): + version = version + elif version: + version = "==" + version + # (we are not running it yet) + argv = Argv( + self._python, + "-m", + "pip", + "install", + "uv{}".format(version), + "--upgrade", + "--disable-pip-version-check", + ) + # this is just for beauty and checks, we already set the verion in the Argv + if not version: + version = "latest" + else: + # mark to install uv if not already installed (we are not running it yet) + argv = Argv( + self._python, + "-m", + "pip", + "install", + "uv", + "--disable-pip-version-check", + ) + version = "" + + # first upgrade pip if we need to + try: + from clearml_agent.helper.package.pip_api.venv import VirtualenvPip + + pip = VirtualenvPip( + session=self.session, + python=self._python, + requirements_manager=None, + path=None, + interpreter=self._python, + ) + pip.upgrade_pip() + except Exception as ex: + self.log.warning("failed upgrading pip: {}".format(ex)) + + # check if we do not have a specific version and uv is found skip installation + if not version and check_if_command_exists("uv"): + print( + "Notice: uv was found, no specific version required, skipping uv installation" + ) + else: + print("Installing / Upgrading uv package to {}".format(version)) + # now install uv + try: + print(argv.get_output()) + except Exception as ex: + self.log.warning("failed installing uv: {}".format(ex)) + + # all done. + self._initialized = True + + def get_api(self, path): + # type: (Path) -> UvAPI + return UvAPI(self, path) + + +@attr.s +class UvAPI(object): + config = attr.ib(type=UvConfig) + path = attr.ib(type=Path, converter=Path) + + INDICATOR_FILES = "pyproject.toml", "uv.lock" + + def install(self): + # type: () -> bool + if self.enabled: + self.config.run("sync", "--locked", cwd=str(self.path), func=Argv.check_call) + return True + return False + + @property + def enabled(self): + return self.config.enabled and ( + any((self.path / indicator).exists() for indicator in self.INDICATOR_FILES) + ) + + def freeze(self, freeze_full_environment=False): + lines = self.config.run("pip", "show", cwd=str(self.path)).splitlines() + lines = [[p for p in line.split(" ") if p] for line in lines] + return { + "pip": [ + parts[0] + "==" + parts[1] + " # " + " ".join(parts[2:]) + for parts in lines + ] + } + + def get_python_command(self, extra): + if check_if_command_exists("uv"): + return Argv("uv", "run", "python", *extra) + else: + return Argv(self.config._python, "-m", "uv", "run", "python", *extra) + + def upgrade_pip(self, *args, **kwargs): + pass + + def set_selected_package_manager(self, *args, **kwargs): + pass + + def out_of_scope_install_package(self, *args, **kwargs): + pass + + def install_from_file(self, *args, **kwargs): + pass diff --git a/clearml_agent/session.py b/clearml_agent/session.py index 7962cd5..f9ceb18 100644 --- a/clearml_agent/session.py +++ b/clearml_agent/session.py @@ -24,6 +24,7 @@ from clearml_agent.helper.docker_args import DockerArgsSanitizer, sanitize_urls from .version import __version__ POETRY = "poetry" +UV = "uv" @attr.s