Update from upstream. (#59)

* Update from upstream.

Some changes:
  - compatibility with configs and Metaflow 2.16+ (new packaging)
  - updates to better support Mamba 2.0
  - lots of bug fixes around debugging command

* Forgot one change

* Trying to fix tests/installs

* Change to prefix.dev mirror. Remove borked test

* Cleanup yml test file

* Wrong channel alias

* Fix tests

* More test fixes; hopefully for real

* Fix and can now use 2.3.2

* Strip out build info for glibc

* Remove broken test for 3.8

* Remove borked test

Author: romain-intel

.github/workflows/test.yml

@@ -19,7 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest]
-        python-version: ["3.8", "3.10", "3.12"]
+        python-version: ["3.8", "3.10", "3.12", "3.13"]
         resolver: [mamba, conda, micromamba]
         micromamba-version: ["1.5.10-0", "latest"]
         mamba-version: ["mamba=1.5.10", "mamba"]
@@ -34,22 +34,29 @@ jobs:
             mamba-version: "mamba=1.5.10"
           - python-version: "3.8"
             mamba-version: "mamba=1.5.10"
+          - python-version: "3.13"
+            mamba-version: "mamba=1.5.10"
     env:
       METAFLOW_CONDA_DEPENDENCY_RESOLVER: ${{ matrix.resolver }}
       METAFLOW_CONDA_TEST: 1
       METAFLOW_DATASTORE_SYSROOT_LOCAL: .metaflow
+      METAFLOW_HACK_CONDA_CHANNEL_ALIAS: prefix.dev
     steps:
       - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
 
-      - uses: mamba-org/setup-micromamba@f8b8a1e23a26f60a44c853292711bacfd3eac822 # v1.9.0
+      - uses: mamba-org/setup-micromamba@b09ef9b599704322748535812ca03efb2625677b # v2.0.5
         with:
           micromamba-version: ${{ matrix.micromamba-version }}
           environment-file: dev-env.yml
           init-shell: bash
+          post-cleanup: all
+          condarc: |
+            channels:
+              - conda-forge
+            channel_alias: https://prefix.dev
           create-args: >-
             python=${{ matrix.python-version }}
             ${{ matrix.mamba-version }}
-
       - name: install nflx-extension
         shell: bash -eo pipefail -l {0}
         run: |

metaflow_extensions/netflix_ext/cmd/debug/constants.py

@@ -5,7 +5,10 @@
 
 class Constants(object):
     # Default Python version
-    DEFAULT_PYTHON_VERSION = "3.7"
+    DEFAULT_PYTHON_VERSION = "3.10"
+
+    # Name of the main debug file
+    DEBUG_SCRIPT_NAME = "00_DEBUG_ME.ipynb"
 
     # Stubbed class for Metaflow artifacts
     DEBUG_INPUT_ITEM_STUB = """
@@ -115,6 +118,12 @@ def __getattr__(self, name):
             
             raise AttributeError(f"No such attribute: {{name}}")
     
+        def __delattr__(self, name):
+            if name in self._artifact_data:
+                del self._artifact_data[name]
+                return
+            return super().__delattr__(name)
+
         def load_artifact_data(self, name):
             for artifact in self._prev_task.artifacts:
                 if artifact.id == name:
@@ -236,9 +245,10 @@ def merge_artifacts(
 
     # Escape hatch stub
     ESCAPE_HATCH_STUB = """
-    # Uncomment the following lines if os.environ.get("PYTHONPATH") is not set
-    # import sys
-    # sys.path.insert(0, "{MetaflowEnvEscapeDir}")    
+    import os
+    import sys
+    if os.environ.get("PYTHONPATH") is None:
+        sys.path.insert(0, "{MetaflowEnvEscapeDir}")
     """
 
     # Imports needed to define stubbed classes & Debug steps

metaflow_extensions/netflix_ext/cmd/debug/debug_cmd.py

@@ -1,16 +1,15 @@
 import os
 import json
-import tarfile
-import datetime
 
-from metaflow import namespace, Flow, Run, Step, Task
+from metaflow import Flow, Run, Step, Task
 from metaflow._vendor import click
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 from metaflow.exception import CommandException
 from metaflow.cli import echo_dev_null, echo_always
+from metaflow.packaging_sys import MetaflowCodeContent
+from .constants import Constants
 from .debug_script_generator import DebugScriptGenerator
 from .debug_utils import (
-    fetch_environment_type,
     copy_stub_generator_to_metaflow_root_dir,
     _set_python_environment,
     _extract_code_package,
@@ -277,7 +276,8 @@ def _generate_debug_scripts(
     os.makedirs(metaflow_root_dir, exist_ok=True)
     python_executable = python_executable.strip() if python_executable else None
     flow_name = task_pathspec.split("/")[0]
-    debug_file_name = task_pathspec.replace("/", "_").replace("-", "_") + "_debug.py"
+    translation_table = str.maketrans("/.-+=", "_____")
+    debug_file_name = task_pathspec.translate(translation_table) + "_debug.py"
     debug_script_generator = DebugScriptGenerator(task_pathspec, inspect)
     script = debug_script_generator.generate_debug_script(metaflow_root_dir, flow_name)
     debug_script_generator.write_debug_script(
@@ -301,7 +301,8 @@ def _generate_debug_scripts(
         debug_script_generator.write_debug_notebook(metaflow_root_dir, notebook_json)
 
     obj.echo(
-        f"Debug scripts and notebook generated at {metaflow_root_dir}",
+        f"Debug scripts and notebook generated at {metaflow_root_dir}; "
+        f"launch {Constants.DEBUG_SCRIPT_NAME} to debug the task",
         fg="green",
         bold=True,
     )
@@ -313,7 +314,9 @@ def _generate_debug_scripts(
 
 def _update_kernel_pythonpath(kernelspec_path, metaflow_root_dir):
     """
-    Updates the kernelspec with the escape trampolines added to the PYTHONPATH.
+    Updates the kernelspec with the escape trampolines added to the PYTHONPATH
+    as well as any other environment variables that may be needed by the
+    version of Metaflow installed.
 
     Parameters
     ----------
@@ -329,5 +332,18 @@ def _update_kernel_pythonpath(kernelspec_path, metaflow_root_dir):
     _ = kernel_json.setdefault("env", {})["PYTHONPATH"] = os.path.abspath(
         os.path.join(metaflow_root_dir, "_escape_trampolines")
     )
+
+    for key, value in MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
+        metaflow_root_dir
+    ).items():
+        if key.endswith(":"):
+            # If the key ends with a colon, we override the existing value
+            kernel_json["env"][key[:-1]] = value
+        elif key not in kernel_json["env"]:
+            kernel_json["env"][key] = value
+        else:
+            # If the key already exists, we prepend the value to the existing one
+            kernel_json["env"][key] = f"{value}:{kernel_json['env'][key]}"
+
     with open(kernel_json_path, "w") as f:
         json.dump(kernel_json, f, indent=4)

metaflow_extensions/netflix_ext/cmd/debug/debug_script_generator.py

@@ -336,7 +336,7 @@ def write_debug_notebook(
             The name of the debug notebook.
         """
         # We then write out the notebook JSON
-        debug_notebook_name = "00_DEBUG_ME.ipynb"
+        debug_notebook_name = Constants.DEBUG_SCRIPT_NAME
         debug_notebook_path = os.path.join(metaflow_root_dir, debug_notebook_name)
         with open(debug_notebook_path, "w") as f:
             json.dump(notebook_json, f, indent=4)

metaflow_extensions/netflix_ext/cmd/debug/debug_stub_generator.py

@@ -21,7 +21,12 @@ def __init__(self, task_pathspec):
         self.run_id = self.task.parent.parent.id
         self.flow_name = self.task.parent.parent.parent.id
         self.is_new_conda_step = self.is_new_conda_step()
-        self.workflow_dag = self.task["_graph_info"].data
+
+        parameters_task = Step(
+            "%s/_parameters" % self.task.parent.parent.pathspec, _namespace_check=False
+        ).task
+
+        self.workflow_dag = parameters_task["_graph_info"].data
         self.file_name = self.workflow_dag["file"]
         self._dag_structure = self.get_dag_structure(self.workflow_dag["steps"])
         self.step_type = self.get_step_type(self.step_name)