Reduce boilerplate around MultiStorageClient feature checks

[Randl]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

Use a proxy object to handle MultiStorageClient.

Issue tracking

Linked issue: Fixes #4507

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[maanug-nv]

Hi @Randl , thanks for tackling this. MSC accesses seem a lot cleaner now!
The implementation of maybe_msc seems problematic though, iiuc it will always be disabled. Since MaybeMultiStorageClient's init is running once the first time this file is imported, it will get initialized to maybe_msc before a user even has the chance to run MultiStorageClientFeature.enable().

Some ideas off the top of my head that might work instead:

• move MaybeMultiStorageClient and maybe_msc to a different file. If this gets added to any 'init.py' files on accident though, might run into the same issue.
• construct a MaybeMultiStorageClient each time. ie in the other files, you'd do MaybeMultiStorageClient().os or similar.
• replace MaybeMultiStorageClient with a function.
• use lazy initialization of self.open, self.os, etc. this would allow you to still construct it once, and the 'is_enabled()' checks will happen at runtime lazily.
• add self.open , self.os to MultiStorageClientFeature , and set appropriately in enable() and disable().

Shouldn't take too long to implement any of those I think. Once it's fixed, we can push this through review to get merged.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Randl]

Yep, you're right @maanug-nv
I've replaced it with dynamic __getattr__

[maanug-nv]

/ok to test 07c463d

[maanug-nv]

Code review

Found 4 issues:

1. load_common discards its return value — always returns None

The refactored line dropped the return that both original branches had (return msc.torch.load(...) / return torch.load(...)). load_common now silently returns None on the happy path, so every caller that consumes the loaded common StateDict gets None and fails downstream. The docstring still says Returns: StateDict.

Megatron-LM/megatron/core/dist_checkpointing/strategies/common.py

Lines 36 to 43 in 07c463d

	load_path = os.path.join(checkpoint_dir, COMMON_STATE_FNAME)
	try:
	maybe_msc.torch.load(load_path, map_location='cpu')
	except FileNotFoundError as e:
	err_msg = f'Common file {load_path} does not exist'
	ckpt_files = [f.name for f in maybe_msc.Path(checkpoint_dir).iterdir()]
	logger.debug(f'{err_msg}. Checkpoint directory content: {ckpt_files}')
	raise CheckpointingException(err_msg) from e

2. verify_checkpoint passes strict=False to stdlib os.path.isdir when MSC is disabled

strict=False is an MSC-specific extension; in the original code it was only used in the MultiStorageClientFeature.is_enabled() branch, while the non-MSC branch called plain os.path.isdir(checkpoint_dir). The unified call always passes strict=False, so when MSC is disabled (the default) maybe_msc.os resolves to stdlib os and os.path.isdir(path, strict=False) raises TypeError: isdir() takes ... no keyword arguments, crashing verify_checkpoint.

Megatron-LM/megatron/core/dist_checkpointing/validation.py

Lines 206 to 211 in 07c463d

	Args:
	checkpoint_dir (str): checkpoint directory
	"""
	isdir = maybe_msc.os.path.isdir(str(checkpoint_dir), strict=False)
	if not isdir:
	raise CheckpointingException(f'Checkpoint directory {checkpoint_dir} does not exist')

3. to_yaml drops the safe_yaml_representers() context manager

The original to_yaml wrapped yaml.safe_dump in with safe_yaml_representers(): (still present in print_yaml just below). Without it, configs containing enums, functools.partial, torch dtypes, etc. raise yaml.representer.RepresenterError at serialization time.

Megatron-LM/megatron/training/config/container.py

Lines 189 to 202 in 07c463d

	yaml_path: Path where to save the YAML file.
	"""
	config_dict = self.to_dict()
	with maybe_msc.open(yaml_path, "w") as f:
	yaml.safe_dump(config_dict, f, default_flow_style=False)
	def print_yaml(self) -> None:
	"""
	Print the config container to the console in YAML format.
	"""
	config_dict = self.to_dict()
	with safe_yaml_representers():
	print(yaml.safe_dump(config_dict, default_flow_style=False))

4. save_integrity_manifest hashes subdirectories on local filesystems

The original local-path branch guarded each entry with entry.is_file() before calling _compute_file_hash; the MSC branch omitted it. The unified loop adopts the MSC behavior for both, so on a local checkpoint directory containing subdirectories, _compute_file_hash opens a directory and raises IsADirectoryError.

Megatron-LM/megatron/core/dist_checkpointing/validation.py

Lines 518 to 523 in 07c463d

	manifest: Dict[str, str] = {}
	ckpt_path = maybe_msc.Path(checkpoint_dir)
	for entry in sorted(ckpt_path.iterdir(), key=lambda p: str(p)):
	if entry.name != INTEGRITY_FNAME:
	manifest[entry.name] = _compute_file_hash(str(entry))

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[Randl]

Fair enough. isdir is a bit annoying, everything else is quite trivial

[maanug-nv]

/ok to test 1a55265

[maanug-nv]

/ok to test da12d84

[maanug-nv]

Hi @Randl , sorry this is taking so long.
There were some merge conflicts, I manually tried to resolve from Github UI, but looks like I re-introduced import of open_file. There might be some other failures as well.
Can you fix those up?

[Randl]

@maanug-nv I've pushed the update, do you need to rerun the CI manually?

[maanug-nv]

/ok to test 4dc5816

[maanug-nv]

/ok to test 1ecd38e

[Phlip79]

/ok to test bbf91bb

[dimapihtar]

/ok to test 6578f93

[Randl]

Looks like attempted merge broke something, I fixed and pushed again

[Phlip79]

/ok to test 6cc1dde

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29775113671