Implement MonitoredChannel

[jchmura-sc]

Purpose
The purpose of this PR is to add telemetry to our ShmChannel.

We want to know how many SampleMessage's are in each of our channels at any given time, since an empty channel implies a blocking recv, which implies a dataloader bottleneck. Conversely, if our channels stay comfortably above some low watermark then any changes upstream of dataloader.__next__ is unlikely to improve throughput.

Context (Colocated-Mode)
Each of our N replicas launch N_local training workers. Each training worker initializes one or more data loaders (e.g. DistABLPLoader, DistNeighborLoader). Each dataloader is associated with a ShmChannel. For a given dataloader, multiple sampling workers (e.g. DistNeighborSampler, DistPPRNeighborSampler) push data into the queue.

The channel is roughly used like:

channel = ShmChannel()

# In dataloader process (single consumer)
def __next__(self):
    msg: SampleMessage = channel.recv() # Blocks if empty
    ... # Convert SampleMessage to PyG data

# In each`SamplingRuntime` worker process(multiple producers)
async def _send_adapter(self, ...):
    ... # Sample subgraphs    
    msg: SampleMessage = self._collate_fn(...) # Hydrate features over RPC
    channel.send(msg) # Blocks if channel is at capacity

Challenge
There is no direct way to gauge the queue size since GLT's ShmChannel does not expose the size to python.

PR Changes

1. New class SizedShmChannel(ShmChannel): adds qsize() method to ShmChannel by attaching to the channels memory region and inspecting the C++ struct layout. It calculates queue size as the delta between the circular buffer's write and read sequence numbers.
2. New class MonitoredShmChannel(SizedShmChannel): a subclass that automatically integrates with the GiGL metrics service, pushing the qsize as a gauge on every recv() call.
3. Dataloaders: Colocated mode setup creates a MonitoredShmChannel instead of a ShmChannel
4. Metrics Service Provider: As far as I am aware, this is the first metric that is being exported from within GiGL. We call initialize_metrics() in the GiGL runtime setup but this doesn't propogate to the processes which actually own the MonitoredChannel. I've made it such that get_metrics_service_instance attempts to implicitly call initialize_metrics(task_config_uri, service_name) if the relevant environment variables are set (which they are in the VAI workers). This enables us to integrate with an OpsMetricPublisher from the training code.

Things to Verify

1. Are we okay with the memory address inspection / struct pointer arithmetic to derive ShmChannel queue size.
2. Do we want to make it possible to turn off channel monitoring and fallback to standard ShmChannel.
3. More Large Scale Testing TBD

Updated Changelog.md? NO

Ready for code review?: YES

[jchmura-sc]

/unit_test

[github-actions]

GiGL Automation

@ 21:34:08UTC : 🔄 C++ Unit Test started.

@ 21:35:52UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:34:10UTC : 🔄 Scala Unit Test started.

@ 21:43:27UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:34:10UTC : 🔄 Python Unit Test started.

[jchmura-sc]

These are monotonically increasing sequence numbers, not wraparound array indices. Hence we can compute queue size as write_block_id_ - read_block_id_

[jchmura-sc]

It's actually "GIGL_APPLIED_TASK_IDENTIFIER" this serves as the appropriate service name based on frozen config. Should double check if "GBML_JOB_NAME" is set and if so then this can be simplified

[kmontemayor2-sc]

Why don't we call initialize_metrics first? I'm. a little confused as to why we're making this change

[kmontemayor2-sc]

BTW would it be easier to just copy over / re-implement our own CPP share memory channel?

We support cpp now - though I think there may be some fiddling required to setup GLT imports in our CPP code.

But it seems like a better long term solution to just add a qsize to the shm cpp directly, WDYT?

[jchmura-sc]

I agree that implementing our own CPP shared memory channel is the most robust long term solution. But this would be a rather large change, requiring a few different moving parts copied over. I suspect we are not too worried about upstream features for these queues on the GLT side, but worth considering too.

I think the options are something like:

1. Copy everything over from GLT. Gives us full control but large change.
2. Build GLT from source with minimal changes to support qsize. Not worth the hassle in my opinion.
3. Try to extend GLT using our own bindings to support qsize. This would be best but (a) queue meta is private field of the channel, and the fields we read to derive the qsize are private within meta (friendship would not be inherited to our subclasses).
4. Use the shared memory inspection as we currently have.

If it's okay, we can keep the implementation as is, leave a TODO, and consider full channel re-implementation later if we think it's worth it.

[kmontemayor2-sc]

Why don't we call initialize_metrics first? I'm. a little confused as to why we're making this change