Accessing node and graph metadata/context from within a node

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
```python
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
```python
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
```python
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
)

```