[compat] wandb

[asaiacai]

Adds pluto.compat.wandb module that lets users replace import wandb with import pluto.compat.wandb as wandb to route all logging through pluto with minimal code changes.

Module structure:

• init.py: Module-level API (init, log, finish, watch, config, summary, run)
• run.py: Run class wrapping pluto.Op with wandb.Run-compatible interface
• config.py: Dict-like Config object that syncs mutations to pluto
• summary.py: Dict-like Summary object tracking last-logged values
• data_types.py: Wrappers (Image, Audio, Video, Table, Histogram, Html, Artifact, AlertLevel) that convert to pluto equivalents

Key features:

• commit=False buffering (accumulate data across log calls)
• Nested dict flattening with / separator
• wandb env var fallbacks (WANDB_PROJECT, WANDB_MODE, WANDB_TAGS, etc.)
• Graceful degradation for unsupported features (define_metric, save, etc.)
• Context manager support
• Disabled-mode fallback when pluto.init() fails

https://claude.ai/code/session_01VTSZKK5UsMqjiADFX57SMY

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)

[gemini-code-assist]

Summary of Changes

Hello @asaiacai, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new pluto.compat.wandb module designed to simplify the migration of existing Weights & Biases (wandb) logging code to Pluto. By allowing users to replace their import wandb statement with import pluto.compat.wandb as wandb, the module routes all logging operations through Pluto with minimal code modifications. This ensures that projects leveraging wandb's API can easily transition to using Pluto for experiment tracking and data visualization, maintaining compatibility while benefiting from Pluto's backend.

Highlights

• Wandb Compatibility Module: Introduced pluto.compat.wandb to enable seamless integration for users migrating from wandb by allowing direct replacement of import statements.
• API Emulation: Provides compatible interfaces for core wandb functions like init, log, finish, watch, config, summary, and run, ensuring existing code works with minimal changes.
• Data Type Conversion: Includes wrappers for wandb.Image, Audio, Video, Table, Histogram, Html, Artifact, and AlertLevel to automatically convert them to their Pluto equivalents.
• Robustness and Usability Features: Implements commit=False buffering for accumulating data, nested dictionary flattening, environment variable fallbacks (e.g., WANDB_PROJECT), graceful degradation for unsupported features, and context manager support for wandb.init().
• Error Handling: Incorporates a fallback mechanism to create a disabled run that no-ops all operations if pluto.init() fails during initialization.

Changelog

• pluto/compat/wandb/init.py
  • Added the main compatibility module, exposing core wandb API functions and managing module-level state for run, config, and summary.
• pluto/compat/wandb/config.py
  • Added a Config class that mimics wandb.config, supporting dict-like and attribute access, and syncing changes to Pluto's underlying operation.
• pluto/compat/wandb/data_types.py
  • Added wrapper classes for various wandb data types (Image, Audio, Video, Table, Histogram, Html, Graph, Artifact, AlertLevel) to facilitate their conversion into Pluto's native data types.
• pluto/compat/wandb/run.py
  • Added a Run class that acts as a wandb.Run-compatible wrapper around a pluto.Op, handling logging, configuration, and run lifecycle events.
• pluto/compat/wandb/summary.py
  • Added a Summary class that emulates wandb.summary, automatically tracking scalar metrics from log calls and allowing manual overrides.
• tests/test_wandb_compat.py
  • Added comprehensive unit tests to verify the functionality and compatibility of the new wandb compatibility layer, covering Config, Summary, DataTypes, Run, and module-level API interactions.

Activity

• No human activity (comments, reviews) was detected for this pull request.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a wandb-compatible layer for Pluto, allowing users to easily switch from wandb to pluto with minimal code changes. The implementation covers core API functions like init, log, finish, and various data types. The module handles environment variable fallbacks and gracefully degrades for unsupported features, which is a good design choice for compatibility. The accompanying tests provide good coverage for the new functionality. I've identified a few areas for improvement related to redundant logic, unused parameters, and potential data handling nuances that could enhance correctness and maintainability.

[gemini-code-assist]

When np_histogram is provided (which is a tuple of (values, bin_edges)), the _to_pluto() method passes the entire tuple as the bins argument to PlutoHistogram. The PlutoHistogram constructor expects bins to be an integer or a sequence of bin edges. Passing the tuple directly might not correctly extract the bin edges, potentially leading to incorrect histogram representation. It should likely be bins=self.np_histogram[1] to pass only the bin edges.

            return PlutoHistogram(data=self.np_histogram, bins=self.np_histogram[1])

[gemini-code-assist]

Several parameters (grouping, classes, boxes, masks, file_type, normalize) are accepted by the Image constructor but are not stored or used within the class or during conversion to PlutoImage. This can create confusion for users expecting these parameters to have an effect. Consider removing them from the signature or adding a warning if they are provided.

    def __init__(
        self,
        data_or_path: Any = None,
        caption: Optional[str] = None,
    ) -> None:

[gemini-code-assist]

The configuration filtering logic could be more robust. If config is None or not a dict/object, config_dict could remain None, leading to errors if config_include_keys or config_exclude_keys are present. It's safer to initialize config_dict to an empty dictionary. Additionally, if both config_include_keys and config_exclude_keys are provided, the order of operations might lead to unexpected results. Consider clarifying the precedence or disallowing both simultaneously.

    config_dict: Dict[str, Any] = {}
    if config is not None:
        if isinstance(config, dict):
            config_dict = dict(config)
        elif hasattr(config, '__dict__'):
            config_dict = vars(config)

        if config_include_keys:
            config_dict = {
                k: v for k, v in config_dict.items() if k in config_include_keys
            }
        if config_exclude_keys:
            config_dict = {
                k: v for k, v in config_dict.items() if k not in config_exclude_keys
            }

[gemini-code-assist]

When pluto.init() fails, the exception e is logged directly in the warning message. For better debugging and to capture the full context of the error, it's recommended to log the exception with exc_info=True.

        logger.warning('%s: pluto.init() failed (%s), creating disabled run', tag, e, exc_info=True)

[gemini-code-assist]

The _allow_val_change attribute is initialized but never used within the Config class. Consider removing it if it's not intended for future functionality, to keep the codebase clean.

Suggested change

	object.__setattr__(self, '_allow_val_change', True)
	object.__setattr__(self, '_data', {})

[gemini-code-assist]

The checksum and max_objects parameters are accepted by the add_reference method but are not used. Since add_reference is already a no-op, these parameters further highlight unused functionality. They should be removed.

        self,
        uri: str,
        name: Optional[str] = None,
    ) -> 'Artifact':

[gemini-code-assist]

The start_time property currently returns time.time(), which is the current time, not the actual start time of the run. wandb.run.start_time typically refers to the timestamp when the run was initialized. This property should ideally retrieve the actual run start time from self._op.settings or a similar source to provide accurate information.

Suggested change

	import time
	return time.time()
	@property
	def start_time(self) -> float:
	return getattr(self._op.settings, 'start_time', time.time())

[gemini-code-assist]

When log_artifact is called with a string path and no name is provided, the log_name defaults to 'artifact'. If multiple string paths are logged without explicit names, they will all be logged under the same key, potentially overwriting previous logs or making them indistinguishable. It would be more robust to use the base filename of the path as the default name if name is None.

            art = PlutoArtifact(data=artifact_or_path, caption=name)
            log_name = name or os.path.basename(artifact_or_path)

[gemini-code-assist]

The _update_from_log method explicitly excludes boolean values from being stored in the summary (and not isinstance(v, bool)). While booleans are often treated as flags, wandb.summary might include them as scalar values in some contexts. If full wandb.summary compatibility is desired, this exclusion might lead to a slight behavioral difference. Consider if this exclusion is intentional and aligns with wandb's behavior, or adjust to include booleans if appropriate.

Suggested change

	if isinstance(v, (int, float)) and not isinstance(v, bool):
	store[k] = v
	elif hasattr(v, 'item') and callable(v.item):
	store[k] = v.item()
	for k, v in data.items():
	if isinstance(v, (int, float)):
	store[k] = v
	elif hasattr(v, 'item') and callable(v.item):
	store[k] = v.item()

[gemini-code-assist]

The test_tags_get_and_set test only asserts that op.add_tags was called, but it doesn't verify the arguments passed to add_tags or remove_tags. This makes the test less robust. It would be more effective to assert the specific tags that were added and removed.

Suggested change

	op.add_tags.assert_called()
	# Setting tags
	run.tags = ('tag1', 'tag2')
	op.remove_tags.assert_called_with(['tag1'])
	op.add_tags.assert_called_with(['tag2'])