mirror of
https://github.com/clearml/clearml-agent
synced 2025-04-08 14:34:19 +00:00
add uv as a package manager
This commit is contained in:
parent
3273f76b46
commit
11bf855547
223
clearml_agent/helper/package/uv_api.py
Normal file
223
clearml_agent/helper/package/uv_api.py
Normal file
@ -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
|
@ -24,6 +24,7 @@ from clearml_agent.helper.docker_args import DockerArgsSanitizer, sanitize_urls
|
||||
from .version import __version__
|
||||
|
||||
POETRY = "poetry"
|
||||
UV = "uv"
|
||||
|
||||
|
||||
@attr.s
|
||||
|
Loading…
Reference in New Issue
Block a user