metaflow-netflixext 1.3.4.dev1__tar.gz → 1.3.4.dev3__tar.gz

{metaflow_netflixext-1.3.4.dev1/metaflow_netflixext.egg-info → metaflow_netflixext-1.3.4.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow-netflixext
-Version: 1.3.4.dev1
+Version: 1.3.4.dev3
 Summary: Metaflow extensions from Netflix
 Author: Netflix Metaflow Developers
 Author-email: metaflow-dev@netflix.com

{metaflow_netflixext-1.3.4.dev1 → metaflow_netflixext-1.3.4.dev3}/metaflow_extensions/netflix_ext/plugins/conda/conda.py

@@ -2743,7 +2743,7 @@ class Conda(object):
 
         entrypoint = os.path.join(env_dir, "bin", "python")
         if os.environ.get("METAFLOW_COVERAGE_S3_PATH"):
-            from metaflow_extensions.nflx.plugins.coverage.setup_coverage import (
+            from metaflow_extensions.netflix_ext.plugins.coverage.setup_coverage import (
                 setup_coverage,
             )
 

{metaflow_netflixext-1.3.4.dev1 → metaflow_netflixext-1.3.4.dev3}/metaflow_extensions/netflix_ext/plugins/conda/conda_step_decorator.py

@@ -661,7 +661,7 @@ class CondaEnvInternalDecorator(StepDecorator):
         cli_args.entrypoint[0] = entrypoint
 
         if os.environ.get("METAFLOW_COVERAGE_S3_PATH"):
-            from metaflow_extensions.nflx.plugins.coverage.setup_coverage import (
+            from metaflow_extensions.netflix_ext.plugins.coverage.setup_coverage import (
                 setup_coverage,
             )
 

metaflow_netflixext-1.3.4.dev3/metaflow_extensions/netflix_ext/plugins/coverage/coveragerc.py

@@ -0,0 +1,21 @@
+# this file is named coveragerc.py since metaflow does not have a way to package hidden files
+
+[run]
+branch = True
+parallel = True
+concurrency = multiprocessing,thread
+
+[html]
+show_contexts = True
+
+[report]
+omit =
+  /**/site-packages/*
+  /**/vendor/*
+  /**/vendored/*
+  /**/_vendor/*
+  /**/build/*
+
+include_namespace_packages = True
+show_missing = False
+ignore_errors = False

metaflow_netflixext-1.3.4.dev3/metaflow_extensions/netflix_ext/plugins/coverage/generate_coveragerc.py

@@ -0,0 +1,96 @@
+import configparser
+import coverage
+import re
+import os
+from metaflow_extensions.netflix_ext.plugins.coverage.setup_coverage import COVERAGE_RCFILE
+
+
+def get_tmp_dirs(pattern, coverage_rcfile, data_file=".coverage"):
+    cov = coverage.Coverage(data_file=data_file, config_file=coverage_rcfile)
+    cov.load()
+    cov_data = cov.get_data()
+    files_covered = cov_data.measured_files()
+    tmp_dirs = set()
+    pattern = re.compile(pattern)
+    for file_path in files_covered:
+        match = pattern.match(file_path)
+        if match:
+            tmp_dirs.add(match.group(1))
+    return list(tmp_dirs)
+
+
+def remap_tmp_dirs(cov, coverage_rcfile, data_file):
+    """
+    We need to update the [paths] config to remap all the temporary directories.
+    This function iterates over all the files and uses a regex match to
+    find the relevent paths that need remapping.
+    """
+    config = cov.config
+    trampoline_dirs = get_tmp_dirs(
+        r"^(/data/tmp/[^/]+/)([^/]+\.py)$", coverage_rcfile, data_file
+    )
+    config.paths["_escape_trampolines"] = ["_escape_trampolines"] + trampoline_dirs
+    ttk_dirs = get_tmp_dirs(
+        r"(?P<dir>.*?/_ttk_[^/]+)/ray_train_func\.py$", coverage_rcfile, data_file
+    )
+    print(f"DEBUG: Found {len(ttk_dirs)} TTK directories: {ttk_dirs}")
+    config.paths["trainingplatform"] = [
+        ".mf_code/metaflow_extensions/nflx/plugins/trainingplatform"
+    ] + ttk_dirs
+    ray_dirs = get_tmp_dirs(
+        r"(?P<dir>.*?/_ray_pkg_[^/]+)/.+", coverage_rcfile, data_file
+    )
+    config.paths["ray"] = ["."] + ray_dirs
+    # Initialize training_platform if it doesn't exist, then add ray_dirs
+    if "training_platform" not in config.paths:
+        config.paths["training_platform"] = []
+    config.paths["training_platform"] = config.paths["training_platform"] + ray_dirs
+    formatted_dirs = get_tmp_dirs(
+        r"(?P<dir>.*?)/[^/]+_(check|test)_flow\.py$", coverage_rcfile, data_file
+    )
+    config.paths["_formatted_flows"] = ["_formatted_flows"] + formatted_dirs
+
+
+def save_updated_coveragerc(config, coverage_rcfile, output_file):
+    # Read template as text to preserve all sections and comments
+    with open(coverage_rcfile, "r") as f:
+        template_content = f.read()
+
+    # Parse template to get existing paths
+    config_parser = configparser.ConfigParser()
+    config_parser.read_string(template_content)
+
+    # Update [paths] section: preserve template paths, add/update dynamic ones
+    if not config_parser.has_section("paths"):
+        config_parser.add_section("paths")
+
+    # Merge: Start with all template paths, then add/update dynamic ones
+    merged_paths = {}
+
+    # First, preserve all template paths
+    if config_parser.has_section("paths"):
+        for key in config_parser.options("paths"):
+            value = config_parser.get("paths", key)
+            # ConfigParser returns multi-line values with newlines already
+            merged_paths[key] = value
+
+    # Then add or override with dynamic paths from config.paths
+    for key, path_list in config.paths.items():
+        # Skip empty path lists to avoid IndexError in coverage.py
+        if path_list:
+            merged_paths[key] = "\n".join(path_list)
+
+    # Write merged paths back (skip empty values)
+    for key, value in merged_paths.items():
+        if value.strip():  # Only write non-empty paths
+            config_parser.set("paths", key, value)
+
+    with open(output_file, "w") as configfile:
+        config_parser.write(configfile)
+    print(f"Updated .coveragerc file written to {output_file}")
+
+
+if __name__ == "__main__":
+    cov = coverage.Coverage(data_file=".coverage", config_file=COVERAGE_RCFILE)
+    remap_tmp_dirs(cov, COVERAGE_RCFILE, ".coverage")
+    save_updated_coveragerc(cov.config, COVERAGE_RCFILE, ".coveragerc")

metaflow_netflixext-1.3.4.dev3/metaflow_extensions/netflix_ext/plugins/coverage/setup_coverage.py

@@ -0,0 +1,111 @@
+import os
+import shutil
+import subprocess
+import sys
+from typing import Union, TYPE_CHECKING
+from metaflow.util import which
+
+if TYPE_CHECKING:
+    import sh
+
+MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
+SITECUSTOMIZE_FILE = os.path.join(MODULE_DIR, "sitecustomize.py")
+
+# Allow internal extensions to provide their own coveragerc override
+try:
+    import metaflow_extensions.nflx.plugins.coverage as _nflx_cov
+
+    _nflx_coveragerc = os.path.join(os.path.dirname(_nflx_cov.__file__), "coveragerc.py")
+    COVERAGE_RCFILE = _nflx_coveragerc if os.path.exists(_nflx_coveragerc) else os.path.join(MODULE_DIR, "coveragerc.py")
+except ImportError:
+    COVERAGE_RCFILE = os.path.join(MODULE_DIR, "coveragerc.py")
+
+
+def get_local_coverage_dir():
+    METAFLOW_COVERAGE_S3_PATH = os.environ.get("METAFLOW_COVERAGE_S3_PATH")
+    assert METAFLOW_COVERAGE_S3_PATH, "METAFLOW_COVERAGE_S3_PATH is not set"
+    group_tag = os.path.basename(METAFLOW_COVERAGE_S3_PATH)
+    coverage_dir = os.path.join("/tmp", group_tag)
+    os.makedirs(coverage_dir, exist_ok=True)
+    return coverage_dir
+
+
+def get_sitepackages_dir(python_binary: str):
+    return (
+        subprocess.run(
+            [python_binary, "-c", "import site; print(site.getsitepackages()[0])"],
+            check=True,
+            stdout=subprocess.PIPE,
+        )
+        .stdout.decode()
+        .strip()
+    )
+
+
+def setup_sitecustomize(python_binary: str):
+    site_packages_dir = get_sitepackages_dir(python_binary)
+    sitecustomize_file = os.path.join(site_packages_dir, "sitecustomize.py")
+    local_coverage_dir = get_local_coverage_dir()
+
+    # We also copy the .coveragerc file to the site-packages directory.
+    # This keeps the .coveragerc file in the same directory as the sitecustomize.py file
+    # so that when the env is saved and restored we will have both.
+    coverage_rcfile = os.path.join(site_packages_dir, ".coveragerc")
+    if not os.path.exists(coverage_rcfile):
+        shutil.copy(COVERAGE_RCFILE, coverage_rcfile)
+
+    # We harden out the environment variables as the fallback so that it will work with
+    # subprocesses that are launched without inheriting the environment.
+    lines = [
+        "import os",
+        "import site",
+        "site_packages_dir = site.getsitepackages()[0]",
+        f"os.environ['COVERAGE_FILE'] = os.environ.get('COVERAGE_FILE') or os.path.join('{local_coverage_dir}', f'.coverage')",
+        f"os.environ['COVERAGE_PROCESS_START'] = os.environ.get('COVERAGE_PROCESS_START') or os.path.join(site_packages_dir, '.coveragerc')",
+        f"os.environ['COVERAGE_RCFILE'] = os.environ.get('COVERAGE_RCFILE') or os.path.join(site_packages_dir, '.coveragerc')",
+        f"os.environ['METAFLOW_COVERAGE_S3_PATH'] = os.environ.get('METAFLOW_COVERAGE_S3_PATH') or '{os.environ['METAFLOW_COVERAGE_S3_PATH']}'",
+        "import coverage",
+        "coverage.process_startup()",
+        "coverage_context = os.environ.get('METAFLOW_COVERAGE_CONTEXT')",
+        "if coverage_context and coverage.Coverage.current():",
+        "    coverage.Coverage.current().switch_context(coverage_context)",
+    ]
+    if not os.path.exists(sitecustomize_file):
+        with open(sitecustomize_file, "w") as target_file:
+            target_file.write("\n".join(lines) + "\n")
+
+
+def maybe_install_coverage(python_binary: str):
+    try:
+        subprocess.run(
+            [python_binary, "-c", "import coverage"],
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+    except subprocess.CalledProcessError:
+        subprocess.run([python_binary, "-m", "pip", "install", "coverage"], check=True)
+
+
+def setup_coverage(sh_python: Union[str, "sh.Command"]):
+    sh_python = str(sh_python)
+    assert os.environ.get(
+        "METAFLOW_COVERAGE_S3_PATH"
+    ), "METAFLOW_COVERAGE_S3_PATH is not set"
+
+    python_binaries = set(
+        ["python"]
+        + os.environ.get("MF_BDI_PYTHONPATH", "").split()
+        + os.environ.get("MFPY", "").split()
+        + [sh_python]
+    )
+    for python_binary in python_binaries:
+        if not python_binary or not which(python_binary):
+            continue
+        maybe_install_coverage(python_binary)
+        setup_sitecustomize(python_binary)
+
+
+if __name__ == "__main__":
+    sh_python = sys.argv[1] if len(sys.argv) > 1 else ""
+    setup_coverage(sh_python)

metaflow_netflixext-1.3.4.dev3/metaflow_extensions/netflix_ext/plugins/coverage/upload_coverage_files.py

@@ -0,0 +1,83 @@
+import coverage
+import glob
+import gzip
+import hashlib
+import os
+import shutil
+import subprocess
+import uuid
+
+from .setup_coverage import get_local_coverage_dir, COVERAGE_RCFILE
+
+
+def hash_file(file_path: str) -> str:
+    hasher = hashlib.sha256()
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+def gzip_file(source_path: str, destination_path: str) -> None:
+    with open(source_path, "rb") as f_in:
+        with gzip.open(destination_path, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def has_arcs(file_path: str) -> bool:
+    try:
+        cov = coverage.Coverage(data_file=file_path)
+        cov.load()
+        return cov.get_data().has_arcs
+    except:
+        # sometimes the coverage file is corrupted and cannot be loaded, eg:
+        #   - doesn't seem to be a coverage data file
+        #   - database disk image is malformed
+        # so we just ignore it
+        return False
+
+
+def combine_coverage_files(local_dir):
+    combined_coverage_file = f"{local_dir}/.coverage"
+    if os.path.exists(combined_coverage_file):
+        shutil.move(combined_coverage_file, f"{combined_coverage_file}.{uuid.uuid4()}")
+    files = [f for f in glob.glob(f"{local_dir}/.coverage.*") if has_arcs(f)]
+    if not files:
+        return
+    cov = coverage.Coverage(
+        data_file=combined_coverage_file, config_file=COVERAGE_RCFILE
+    )
+    try:
+        cov.combine(data_paths=files, keep=True)
+    except coverage.exceptions.NoDataError:
+        return
+    cov.save()
+    assert os.path.exists(combined_coverage_file)
+
+
+def upload_coverage_files(coverage_s3_path: str) -> None:
+    coverage_s3_path = coverage_s3_path.rstrip("/")
+    local_coverage_dir = get_local_coverage_dir()
+    combine_coverage_files(local_coverage_dir)
+
+    combined_coverage_file = f"{local_coverage_dir}/.coverage"
+    file_hash = hash_file(combined_coverage_file)
+
+    gzipped_file = f"{local_coverage_dir}/.coverage.{file_hash}.gz"
+    gzip_file(combined_coverage_file, gzipped_file)
+
+    subprocess.run(
+        f"aws s3 cp {gzipped_file} {coverage_s3_path}/",
+        shell=True,
+        check=True,
+        cwd=local_coverage_dir,
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) != 2:
+        print("Usage: python script.py <COVERAGE_S3_PATH>")
+    else:
+        upload_coverage_files(sys.argv[1])

{metaflow_netflixext-1.3.4.dev1 → metaflow_netflixext-1.3.4.dev3/metaflow_netflixext.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow-netflixext
-Version: 1.3.4.dev1
+Version: 1.3.4.dev3
 Summary: Metaflow extensions from Netflix
 Author: Netflix Metaflow Developers
 Author-email: metaflow-dev@netflix.com

{metaflow_netflixext-1.3.4.dev1 → metaflow_netflixext-1.3.4.dev3}/metaflow_netflixext.egg-info/SOURCES.txt

@@ -48,6 +48,11 @@ metaflow_extensions/netflix_ext/plugins/conda/resolvers/pylock_toml_resolver.py
 metaflow_extensions/netflix_ext/plugins/conda/resources/logo-32x32.png
 metaflow_extensions/netflix_ext/plugins/conda/resources/logo-64x64.png
 metaflow_extensions/netflix_ext/plugins/conda/resources/logo-svg.svg
+metaflow_extensions/netflix_ext/plugins/coverage/__init__.py
+metaflow_extensions/netflix_ext/plugins/coverage/coveragerc.py
+metaflow_extensions/netflix_ext/plugins/coverage/generate_coveragerc.py
+metaflow_extensions/netflix_ext/plugins/coverage/setup_coverage.py
+metaflow_extensions/netflix_ext/plugins/coverage/upload_coverage_files.py
 metaflow_extensions/netflix_ext/toplevel/mfextinit_netflixext.py
 metaflow_extensions/netflix_ext/toplevel/netflixext_toplevel.py
 metaflow_extensions/netflix_ext/toplevel/netflixext_version.py