Various fixes and cleanups

Author: romain-intel

README.md

@@ -33,6 +33,9 @@ This extension currently contains:
 This functionality is currently being actively tested within Netflix but has not yet
 been deployed in production.
 
+It is likely to evolve primarily in its implementation as we do further testing. Feedback
+on what is working and what is not is most welcome.
+
 ### Improvements over the included Conda decorator
 This decorator improves several aspects of the included Conda decorator:
 - it has significant performance gains:
@@ -93,13 +96,14 @@ needs to satisfy the following requirements:
 - `lockfile`
 - (optional but recommended) `mamba>=1.1.0`
 
-In addition, and only if you want to support `pypi` packages, you need to apply the
+In addition, and only if you want to support `pypi` packages, it is best to apply the
 PR `https://github.com/conda-incubator/conda-lock/pull/290` to `conda-lock`. This is
 the unfortunate result of a bug in how `conda-lock` handles packages that are both
 present in the `conda` environment and `pypi` one.
 
-Due to bugs in `conda`, if your resolved environment contains `.conda` packages and you
-do not have `micromamba` installed, the environment creation will fail.
+Due to bugs in `conda` and the way we use it, if your resolved environment
+contains `.conda` packages and you do not have `micromamba` installed, the
+environment creation will fail.
 
 #### Uninstallation
 Uninstalling this package will revert the behavior of the conda decorator to the one

metaflow_extensions/netflix_ext/cmd/environment/utils.py

@@ -14,7 +14,7 @@
 from metaflow_extensions.netflix_ext.plugins.conda.utils import arch_id
 
 _deps_parse = re.compile(r"([^<>=!~]+)(.*)")
-_ext_parse = re.compile(r"(\w+)\(([^)]+)\)")
+_ext_parse = re.compile(r"([-_\w]+)\(([^)]+)\)")
 
 name_to_pkg = {"netflix-ext": "metaflow-netflixext"}
 

metaflow_extensions/netflix_ext/plugins/conda/conda.py

@@ -66,6 +66,7 @@
     get_conda_root,
     parse_explicit_url_conda,
     parse_explicit_url_pip,
+    plural_marker,
 )
 
 from .env_descr import (
@@ -114,7 +115,9 @@ def _modified_logger(*args: Any, **kwargs: Any):
         self._use_conda_lock_to_resolve = (
             CONDA_PREFERRED_RESOLVER == "conda-lock"
         )  # type: bool
-        self._find_conda_binary()
+        self._found_binaries = False  # We delay resolving binaries until we need them
+        # because in the remote case, in conda_environment.py
+        # we create this object but don't use it.
 
         # Figure out what environments we know about locally
         self._local_root = LocalStorage.get_datastore_root_from_config(
@@ -136,6 +139,8 @@ def _modified_logger(*args: Any, **kwargs: Any):
             self._storage = None
 
     def binary(self, binary: str) -> Optional[str]:
+        if not self._found_binaries:
+            self._find_conda_binary()
         if self._bins:
             return self._bins.get(binary)
         return None
@@ -153,6 +158,9 @@ def resolve(
             # is a "lighter" conda.
             raise CondaException("Cannot resolve environments in a remote environment")
 
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         have_pip_deps = any([d.category == "pip" for d in deps])
         if have_pip_deps and not self._use_conda_lock_to_resolve:
             raise CondaException(
@@ -189,6 +197,9 @@ def add_to_resolved_env(
             # is a "lighter" conda.
             raise CondaException("Cannot resolve environments in a remote environment")
 
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         have_pip_deps = any(
             chain(
                 [p.TYPE == "pip" for p in cur_env.packages],
@@ -252,6 +263,10 @@ def create_for_step(
         env: ResolvedEnvironment,
         do_symlink: bool = False,
     ):
+
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         try:
             # I am not 100% sure the lock is required but since the environments share
             # a common package cache, we will keep it for now
@@ -263,6 +278,10 @@ def create_for_step(
     def create_for_name(
         self, name: str, env: ResolvedEnvironment, do_symlink: bool = False
     ):
+
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         with CondaLock(self._env_lock_file(name)):
             self._create(env, name)
         if do_symlink:
@@ -272,6 +291,9 @@ def create_for_name(
 
     def remove_for_step(self, step_name: str, env_id: EnvID):
         # Remove the conda environment
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         try:
             env_name = self._env_directory_from_envid(env_id)
             return self.remove_for_name(env_name)
@@ -280,11 +302,15 @@ def remove_for_step(self, step_name: str, env_id: EnvID):
             raise CondaStepException(e, [step_name])
 
     def remove_for_name(self, name: str):
+        if not self._found_binaries:
+            self._find_conda_binary()
         with CondaLock(self._env_lock_file(name)):
             self._remove(name)
 
     def python(self, env_desc: Union[EnvID, str]) -> Optional[str]:
         # Get Python interpreter for the conda environment
+        if not self._found_binaries:
+            self._find_conda_binary()
         env_path = None
         if isinstance(env_desc, EnvID):
             env_path = self.created_environment(env_desc)
@@ -305,6 +331,9 @@ def python(self, env_desc: Union[EnvID, str]) -> Optional[str]:
     def created_environment(
         self, env_desc: Union[EnvID, str]
     ) -> Optional[Tuple[EnvID, str]]:
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         if isinstance(env_desc, EnvID):
             prefix = "metaflow_%s_%s" % (env_desc.req_id, env_desc.full_id)
         else:
@@ -320,6 +349,8 @@ def created_environments(
         # List all existing metaflow environments; this can include environments that
         # were created with the `environment` command and therefore have a different
         # name
+        if not self._found_binaries:
+            self._find_conda_binary()
         prefix = "metaflow_%s_" % req_id if req_id else ""
         return {
             k: v
@@ -400,6 +431,8 @@ def cache_environments(
             raise CondaException(
                 "Cannot cache environments since no datastore configured"
             )
+        if not self._found_binaries:
+            self._find_conda_binary()
 
         cache_paths_to_check = []  # type: List[Tuple[str, str, str]]
         my_arch_id = arch_id()
@@ -633,12 +666,19 @@ def _cache_pkg(pkg: PackageSpecification, pkg_fmt: str, local_path: str) -> str:
             if upload_files:
                 start = time.time()
                 self._echo(
-                    "    Caching %d items to %s ..."
-                    % (len(upload_files), self._datastore_type),
+                    "    Caching %d item%s to %s ..."
+                    % (
+                        len(upload_files),
+                        plural_marker(len(upload_files)),
+                        self._datastore_type,
+                    ),
                     nl=False,
                 )
                 self._upload_to_ds(upload_files)
-                self._echo(" done in %d seconds." % int(time.time() - start))
+                delta_time = int(time.time() - start)
+                self._echo(
+                    " done in %d second%s." % (delta_time, plural_marker(delta_time))
+                )
 
                 # If this is successful, we cache the environments. We do this *after*
                 # in case some packages fail to upload so we don't write corrupt
@@ -656,12 +696,19 @@ def _cache_pkg(pkg: PackageSpecification, pkg_fmt: str, local_path: str) -> str:
                     )
                 start = time.time()
                 self._echo(
-                    "    Caching %d environments to %s ..."
-                    % (len(upload_files), self._datastore_type),
+                    "    Caching %d environment%s to %s ..."
+                    % (
+                        len(upload_files),
+                        plural_marker(len(upload_files)),
+                        self._datastore_type,
+                    ),
                     nl=False,
                 )
                 self._upload_to_ds(upload_files)
-                self._echo(" done in %d seconds." % int(time.time() - start))
+                delta_time = int(time.time() - start)
+                self._echo(
+                    " done in %d second%s." % (delta_time, plural_marker(delta_time))
+                )
             else:
                 self._echo("    All items already cached in %s." % self._datastore_type)
 
@@ -691,6 +738,9 @@ def lazy_fetch_packages(
         # transmutation also happens if the requested format is not found (only for
         # conda packages))
 
+        if not self._found_binaries:
+            self._find_conda_binary()
+
         if require_conda_format is None:
             require_conda_format = []
         use_package_dirs = True
@@ -961,8 +1011,12 @@ def _micromamba_transmute(src_file: str, dst_file: str, dst_format: str):
         if do_download:
             start = time.time()
             self._echo(
-                "    Downloading %d(web) + %d(cache) packages ..."
-                % (len(web_downloads), len(cache_downloads)),
+                "    Downloading %d(web) + %d(cache) package%s ..."
+                % (
+                    len(web_downloads),
+                    len(cache_downloads),
+                    plural_marker(len(web_downloads) + len(cache_downloads)),
+                ),
                 nl=False,
             )
 
@@ -1041,12 +1095,18 @@ def _micromamba_transmute(src_file: str, dst_file: str, dst_format: str):
                         )
 
         if do_download:
+            delta_time = int(time.time() - start)
             self._echo(
-                " done in %d seconds." % int(time.time() - start), timestamp=False
+                " done in %d second%s." % (delta_time, plural_marker(delta_time)),
+                timestamp=False,
             )
         if not pending_errors and transmutes:
             start = time.time()
-            self._echo("    Transmuting %d packages ..." % len(transmutes), nl=False)
+            self._echo(
+                "    Transmuting %d package%s ..."
+                % (len(transmutes), plural_marker(len(transmutes))),
+                nl=False,
+            )
             with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
                 transmut_results = [
                     executor.submit(_transmute, entry) for entry in transmutes
@@ -1064,8 +1124,10 @@ def _micromamba_transmute(src_file: str, dst_file: str, dst_format: str):
 
                         if new_url not in known_urls:
                             url_adds.append(new_url)
+            delta_time = int(time.time() - start)
             self._echo(
-                " done in %d seconds." % int(time.time() - start), timestamp=False
+                " done in %d second%s." % (delta_time, plural_marker(delta_time)),
+                timestamp=False,
             )
         if url_adds:
             # Update the urls file in the packages directory so that Conda knows that the
@@ -1387,6 +1449,7 @@ def _find_conda_binary(self):
         err = self._validate_conda_installation()
         if err:
             raise err
+        self._found_binaries = True
 
     def _ensure_local_conda(self):
         if CONDA_LOCAL_PATH is not None:
@@ -1458,7 +1521,11 @@ def _install_local_conda(self):
                 raise InvalidEnvironmentException(
                     msg="Could not extract environment: %s" % str(e)
                 )
-        self._echo(" done in %d seconds." % int(time.time() - start), timestamp=False)
+        delta_time = int(time.time() - start)
+        self._echo(
+            " done in %d second%s." % (delta_time, plural_marker(delta_time)),
+            timestamp=False,
+        )
 
     def _ensure_remote_conda(self):
         if CONDA_REMOTE_INSTALLER is not None:
@@ -1467,14 +1534,17 @@ def _ensure_remote_conda(self):
             # If we don't have a REMOTE_INSTALLER, we check if we need to install one
             args = [
                 "/bin/bash",
-                "-c"
-                "if ! type %s  >/dev/null 2>&1; "
-                "then wget --no-check-certificate "
-                "https://micro.mamba.pm/install.sh -O micromamba.sh >/dev/null 2>&1; "
-                "./micromamba.sh >/dev/null 2>&1; echo $HOME/.local/bin/micromamba; "
-                "else which %s; fi" % ("micromamba", "micromamba"),
+                "-c",
+                "if ! type micromamba  >/dev/null 2>&1; then "
+                "mkdir -p ~/.local/bin >/dev/null 2>&1; "
+                "curl -Ls https://micro.mamba.pm/api/micromamba/%s/latest | "
+                "tar -xvj -C ~/.local/bin/ --strip-components=1 bin/micromamba >/dev/null 2>&1; "
+                "echo $HOME/.local/bin/micromamba; "
+                "else which micromamba; fi" % arch_id(),
             ]
-            self._bins = {"micromamba": subprocess.check_output(args).decode("utf-8")}
+            self._bins = {
+                "conda": subprocess.check_output(args).decode("utf-8").strip()
+            }
 
     def _install_remote_conda(self):
         from metaflow.plugins import DATASTORES
@@ -1541,7 +1611,8 @@ def _validate_conda_installation(self) -> Optional[Exception]:
                     to_remove.append(k)
                 else:
                     return InvalidEnvironmentException(
-                        "Required binary '%s' found. Install using `%s install -n base %s`"
+                        "Required binary '%s' not found. "
+                        "Install using `%s install -n base %s`"
                         % (k, self._dependency_solver, k)
                     )
         if to_remove:
@@ -1847,7 +1918,11 @@ def _create(self, env: ResolvedEnvironment, env_name: str) -> None:
         ) as f:
             json.dump(env.env_id, f)
 
-        self._echo(" done in %s seconds." % int(time.time() - start), timestamp=False)
+        delta_time = int(time.time() - start)
+        self._echo(
+            " done in %s second%s." % (delta_time, plural_marker(delta_time)),
+            timestamp=False,
+        )
 
     def _remove(self, env_name: str):
         # TODO: Verify that this is a proper metaflow environment to remove

metaflow_extensions/netflix_ext/plugins/conda/conda_environment.py

@@ -31,7 +31,7 @@
 
 from metaflow.metaflow_environment import MetaflowEnvironment
 
-from .utils import get_conda_manifest_path
+from .utils import get_conda_manifest_path, plural_marker
 
 from .env_descr import CachedEnvironmentInfo, EnvID, ResolvedEnvironment
 from .conda import Conda
@@ -84,9 +84,20 @@ def init_environment(self, echo: Callable[..., None]):
             # Figure out the environments that we need to resolve for all steps
             # We will resolve all unique environments in parallel
             step_conda_dec = get_conda_decorator(self._flow, step.__name__)
-            env_ids = step_conda_dec.env_ids
-            for env_id in env_ids:
-                resolved_env = self._conda.environment(env_id)
+            my_arch_env_id = step_conda_dec.env_id
+            requested_archs = step_conda_dec.requested_architectures
+            # First look up the env_id (it's the default one)
+            resolved_env = self._conda.environment(my_arch_env_id)
+            if resolved_env is not None:
+                # Check to make sure it supports all requested architectures
+                if not set(requested_archs).issubset(resolved_env.co_resolved_archs):
+                    resolved_env = None
+            for arch in requested_archs:
+                env_id = EnvID(
+                    req_id=my_arch_env_id.req_id,
+                    full_id=my_arch_env_id.full_id,
+                    arch=arch,
+                )
                 if env_id not in self._requested_envs:
                     self._requested_envs[env_id] = {
                         "id": env_id,
@@ -177,11 +188,19 @@ def _resolve_environments(
     ):
         start = time.time()
         if len(env_ids) == len(self._requested_envs):
-            echo("    Resolving %d environments in flow ..." % len(env_ids), nl=False)
+            echo(
+                "    Resolving %d environment%s in flow ..."
+                % (len(env_ids), plural_marker(len(env_ids))),
+                nl=False,
+            )
         else:
             echo(
-                "    Resolving %d of %d environments in flows (others are cached) ..."
-                % (len(env_ids), len(self._requested_envs)),
+                "    Resolving %d of %d environment%s in flow (others are cached) ..."
+                % (
+                    len(env_ids),
+                    len(self._requested_envs),
+                    plural_marker(len(self._requested_envs)),
+                ),
                 nl=False,
             )
 
@@ -262,7 +281,7 @@ def _resolve(env_desc: Mapping[str, Any]) -> Tuple[EnvID, ResolvedEnvironment]:
                     )
 
         duration = int(time.time() - start)
-        echo(" done in %d seconds." % duration)
+        echo(" done in %d second%s." % (duration, plural_marker(duration)))
 
     def _get_env_id(self, step_name: str) -> Optional[EnvID]:
         conda_decorator = get_conda_decorator(self._flow, step_name)