From 00c6589c1a9c136d660bd05ec3d3ab24d261fadf Mon Sep 17 00:00:00 2001 From: Romain Date: Mon, 4 Dec 2023 13:41:23 -0800 Subject: [PATCH] Fixes from upstream: (#30) - fix an issue when resolving multiple mixed environment in parallel - allow the pypi-indices section in environment.yml - improved parsing of versions in environment.yml --- .../cmd/environment/environment_cmd.py | 55 ++++++++++++++++++- .../netflix_ext/plugins/conda/conda.py | 2 + .../conda/resolvers/conda_lock_resolver.py | 27 +++++---- .../netflix_ext/plugins/conda/utils.py | 2 + 4 files changed, 69 insertions(+), 17 deletions(-) diff --git a/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py b/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py index d96e1c7..7310e91 100644 --- a/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py +++ b/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py @@ -55,7 +55,14 @@ REQ_SPLIT_LINE = re.compile(r"([^~<=>]*)([~<=>]+.*)?") -YML_SPLIT_LINE = re.compile(r"(<=|>=|=>|=<|~=|==|<|>|=)") + +# Allows things like: +# pkg = <= version +# pkg <= version +# pkg = version +# pkg = ==version or pkg = =version +# In other words, the = is optional but possible +YML_SPLIT_LINE = re.compile(r"(?:=\s)?(<=|>=|~=|==|<|>|=)") class CommandObj: @@ -152,6 +159,13 @@ def cli(ctx): help="Root path for Conda cached information. If not set, " "looks for METAFLOW_CONDA_S3ROOT (for S3)", ) +@click.option( + "--quiet-file-output", + default=None, + hidden=True, + type=str, + help="Output the quiet output to this file; used for scripting", +) @click.pass_context def environment( ctx: Any, @@ -160,6 +174,7 @@ def environment( datastore: str, environment: str, conda_root: Optional[str], + quiet_file_output: Optional[str], ): if quiet: echo = echo_dev_null @@ -171,6 +186,7 @@ def environment( obj.echo = echo obj.echo_always = echo_always obj.datastore_type = datastore + obj.quiet_file_output = quiet_file_output if conda_root: if obj.datastore_type == "s3": @@ -453,6 +469,10 @@ def create( if obj.quiet: obj.echo_always(python_bin) + if obj.quiet_file_output: + with open(obj.quiet_file_output, "w") as f: + f.write(python_bin) + f.write("\n") else: obj.echo( "Created environment '%s' locally, activate with `%s activate %s`" @@ -724,6 +744,12 @@ def resolve( if obj.quiet: obj.echo_always(env_id.arch) obj.echo_always(env.quiet_print(existing_envs.get(env.env_id))) + if obj.quiet_file_output: + with open(obj.quiet_file_output, "w") as f: + f.write(env_id.arch) + f.write("\n") + f.write(env.quiet_print(existing_envs.get(env.env_id))) + f.write("\n") else: obj.echo("### Environment for architecture %s" % env_id.arch) obj.echo(env.pretty_print(existing_envs.get(env.env_id))) @@ -809,6 +835,10 @@ def show(obj, local_only: bool, arch: str, pathspec: bool, envs: Tuple[str]): if obj.quiet: obj.echo_always("%s@%s" % (env_name, arch)) obj.echo_always("NOT_FOUND") + if obj.quiet_file_output: + with open(obj.quiet_file_output, "w") as f: + f.write("%s@%s\n" % (env_name, arch)) + f.write("NOT_FOUND\n") else: obj.echo( "### Environment for '%s' (arch '%s') was not found" @@ -825,6 +855,17 @@ def show(obj, local_only: bool, arch: str, pathspec: bool, envs: Tuple[str]): ) ) ) + if obj.quiet_file_output: + with open(obj.quiet_file_output, "w") as f: + f.write("%s@%s\n" % (env_name, arch)) + f.write( + resolved_env.quiet_print( + all_envs.get(resolved_env.env_id.req_id, {}).get( + resolved_env.env_id.full_id + ) + ) + ) + f.write("\n") else: obj.echo("### Environment for '%s' (arch '%s'):" % (env_name, arch)) obj.echo( @@ -909,6 +950,10 @@ def get(obj, default: bool, arch: Optional[str], pathspec: bool, source_env: str existing_envs = [] if obj.quiet: obj.echo_always(env.quiet_print(existing_envs)) + if obj.quiet_file_output: + with open(obj.quiet_file_output, "w") as f: + f.write(env.quiet_print(existing_envs)) + f.write("\n") else: obj.echo(env.pretty_print(existing_envs)) cast(Conda, obj.conda).write_out_environments() @@ -1007,11 +1052,15 @@ def _parse_yml_file( mode = "sources" elif line == "dependencies:": mode = "deps" + elif line == "pypi-indices:": + mode = "pypi_sources" else: mode = "ignore" - elif mode == "sources": + elif mode == "sources" or mode == "pypi_sources": line = line.lstrip(" -").rstrip() - sources.setdefault("conda", []).append(line) + sources.setdefault("conda" if mode == "sources" else "pypi", []).append( + line + ) elif mode == "deps" or mode == "pypi_deps": line = line.lstrip(" -").rstrip() if line == "pip:": diff --git a/metaflow_extensions/netflix_ext/plugins/conda/conda.py b/metaflow_extensions/netflix_ext/plugins/conda/conda.py index bb364d8..098f681 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/conda.py +++ b/metaflow_extensions/netflix_ext/plugins/conda/conda.py @@ -317,6 +317,7 @@ def call_binary( args: List[str], binary: str, addl_env: Optional[Mapping[str, str]] = None, + cwd: Optional[str] = None, pretty_print_exception: bool = True, ) -> bytes: if binary in _CONDA_DEP_RESOLVERS: @@ -333,6 +334,7 @@ def call_binary( [binary] + args, stderr=subprocess.PIPE, env=dict(os.environ, **addl_env), + cwd=cwd, ).strip() except subprocess.CalledProcessError as e: print( diff --git a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_lock_resolver.py b/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_lock_resolver.py index 1c5d579..2c1d255 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_lock_resolver.py +++ b/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_lock_resolver.py @@ -243,20 +243,19 @@ def resolve( virtual_yml.writelines(lines) args.extend(["--virtual-package-spec", "virtual_yml.spec"]) - with WithDir(conda_lock_dir): - debug.conda_exec("Build directory: %s" % conda_lock_dir) - # conda-lock will only consider a `pyproject.toml` as a TOML file which - # is somewhat annoying. - with open( - "pyproject.toml", mode="w", encoding="ascii" - ) as input_toml: - input_toml.writelines(toml_lines) - debug.conda_exec( - "TOML configuration:\n%s" % "".join(toml_lines) - ) - self._conda.call_binary( - args, binary="conda-lock", addl_env=addl_env - ) + debug.conda_exec("Build directory: %s" % conda_lock_dir) + # conda-lock will only consider a `pyproject.toml` as a TOML file which + # is somewhat annoying. + with open( + os.path.join(conda_lock_dir, "pyproject.toml"), + mode="w", + encoding="ascii", + ) as input_toml: + input_toml.writelines(toml_lines) + debug.conda_exec("TOML configuration:\n%s" % "".join(toml_lines)) + self._conda.call_binary( + args, binary="conda-lock", addl_env=addl_env, cwd=conda_lock_dir + ) # At this point, we need to read the explicit dependencies in the file created emit = False packages = [] # type: List[PackageSpecification] diff --git a/metaflow_extensions/netflix_ext/plugins/conda/utils.py b/metaflow_extensions/netflix_ext/plugins/conda/utils.py index d0b631c..b9c8fde 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/utils.py +++ b/metaflow_extensions/netflix_ext/plugins/conda/utils.py @@ -661,6 +661,8 @@ def change_pypi_package_version( class WithDir: + # WARNING: os.chdir is not compatible with thread processing so do not use in + # a context where multiple threads can exist. def __init__(self, new_dir: str): self._current_dir = os.getcwd() self._new_dir = new_dir