Accessing node and graph metadata/context from within a node

Comienza revisando los hooks del ciclo de vida pre_node_execute y post_node_execute y la ruta de ejecución asíncrona; el issue también enlaza con dmexs/hamilton-contextvar-repro. Compara el enfoque de contextvars con los ejemplos de propagación de parámetros en mlflow_tracker.py, dag.py y client.py. Se considera terminado cuando exista una forma documentada para que las funciones de los nodos accedan a los metadatos y al contexto proporcionado por los hooks, con cobertura del flujo asíncrono hook-a-nodo-a-hook.

I'd like to be able to access node metadata (task_id, node_name, and other arbitrary parameters) from within my node functions. My particular use case is writing an async compatible MlFlowTracker plugin. (I have had intermittent sucess with the the mlflow Fluent API trace decorator and would rather use the more explicit client API.)

From within a node function I'd like a way to access context passed in from plugin hooks.

Approach 1 - contextvars
Working with @elijahbenizzy, I tried contextvars but found that the context wasn't maintained between the pre_node_execute hook, node function and the post_node_execute hook. See reproduction here: https://github.com/dmexs/hamilton-contextvar-repro

Approach 2 - modifying kwargs and parameter drilling
Instead I took the approach of injecting a context object into node kwargs which works but requires parameter/prop drilling which I don't love. I'd prefer to be able to retrieve passed data from a contextvar anywhere within the node function execution stack. In my case I'm retrieving my mlflow span/request ids to create subspans on API calls.

Here's what I have for approach #2 which works:

mlflow_tracker.py

from hamilton.lifecycle import base
from hamilton import node, graph
import logging
from typing import Type, Any, Optional, Dict, List
import mlflow
import pickle
import contextvars
from pydantic import BaseModel
from contextlib import asynccontextmanager

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class NodeContext(BaseModel):
    mlflow_request_id: Optional[str] = None
    mlflow_root_span_id: Optional[str] = None

class MLFlowTrackerAsync(
    base.BasePreNodeExecuteAsync,
    base.BasePostNodeExecuteAsync,
    base.BasePreGraphExecuteAsync,
    base.BasePostGraphExecuteAsync,
):
    mlflow_run_info = None
    mlflow_client: Optional[mlflow.MlflowClient] = None
    mlflow_request_id: Optional[str] = None
    mlflow_root_span_id: Optional[str] = None
    mlflow_node_span_id_cache = {}

    def __init__(self, *args, **kwargs):
        if mlflow.active_run():
            self.mlflow_client = mlflow.MlflowClient()
            self.mlflow_run_info = mlflow.active_run().info
            logger.debug(f"MLFlowTrackerAsync: {self.mlflow_run_info.run_id}")

    async def pre_graph_execute(
        self,
        *,
        run_id: str,
        graph: "graph.FunctionGraph",
        final_vars: List[str],
        inputs: Dict[str, Any],
        overrides: Dict[str, Any],
    ):
        if self.mlflow_run_info:
            trace = self.mlflow_client.start_trace(
                name=f'hamilton',
                inputs=inputs,
                experiment_id=self.mlflow_run_info.experiment_id,
                span_type=mlflow.entities.SpanType.CHAIN,
            )
            self.mlflow_request_id = trace.request_id
            self.mlflow_root_span_id = trace.span_id

    async def post_graph_execute(
        self,
        *,
        run_id: str,
        graph: "graph.FunctionGraph",
        success: bool,
        error: Optional[Exception],
        results: Optional[Dict[str, Any]],
    ):
        if self.mlflow_request_id:
            self.mlflow_client.end_trace(
                request_id=self.mlflow_request_id,
                outputs=results
            )

    async def pre_node_execute(
        self,
        run_id: str,
        node_: node.Node,
        kwargs: Dict[str, Any],
        task_id: Optional[str] = None
    ) -> None:
        context = kwargs.get('__context', NodeContext())
        if self.mlflow_request_id and self.mlflow_root_span_id:
            span = self.mlflow_client.start_span(
                name=node_.name,
                span_type=mlflow.entities.SpanType.CHAIN,
                request_id=self.mlflow_request_id,
                parent_id=self.mlflow_root_span_id,
                inputs=kwargs,
            )
            self.mlflow_node_span_id_cache[task_id] = span.span_id
            context.mlflow_request_id = span.request_id
            context.mlflow_root_span_id = span.span_id

        if '__context' in node_.input_types:
            kwargs['__context'] = context

    async def post_node_execute(
        self,
        run_id: str,
        node_: node.Node,
        success: bool,
        error: Optional[Exception],
        result: Any,
        task_id: Optional[str] = None,
        **future_kwargs: dict,
    ) -> None:
        if self.mlflow_node_span_id_cache.get(task_id, None):
            self.mlflow_client.end_span(
                request_id=self.mlflow_request_id,
                span_id=self.mlflow_node_span_id_cache[task_id],
                outputs=result
            )

dag.py

async def patient_fhir_id(hero_api: hero_api_async.HeroAPI, uid: str, __context: Optional[NodeContext] = None) -> str:
    """
    Lookup patient
    """
    async with api as client:
        patient = await client.request(endpoint='fhir/patient', params={"id": uid, "idType": "UID"}, context=__context)
    fhir_id = patient['id']
    return fhir_id

client.py

    async def request(
        self,
        endpoint,
        method = 'GET',
        params: Optional[dict] = None,
        context: Optional[NodeContext] = None
    ):       
        span = None
        return_response = None
        if context and context.mlflow_request_id and context.mlflow_root_span_id:
            span = self.mlflow_client.start_span(
                name='API.request',
                span_type=mlflow.entities.SpanType.RETRIEVER,
                request_id=context.mlflow_request_id,
                parent_id=context.mlflow_root_span_id,
                inputs={
                    "method": method,
                    "endpoint": endpoint,
                    "params": params
                }
            )


        try:
            url = self.BASE_URL + endpoint
            async with self.session.request(
                method=method,
                url=url,
                params=params,
            ) as response:
                response.raise_for_status()  # Raise an HTTPError for bad responses
                return_response = await response.json(content_type=None)  # Return the response data as a dictionary
                
            
                return return_response

        except Exception as e:
            logger.error(f"{method} {response.status} {url} - {e}")
            return_response = e
            raise e
        
        finally:
            if span:
                self.mlflow_client.end_span(
                    request_id=span.request_id,
                    span_id=span.span_id,
                    outputs=return_response
                )