wandb fix inside.

[djwashout]

Fix for wandb that has been messing up my colab notebooks.

The script performs three main actions:

Disables wandb via environment variable by setting WANDB_DISABLED=true
Replaces the PyTorch Lightning wandb logger with a mock implementation that provides all the expected interface methods but doesn't actually log anything
Fixes issues with server_features_query.py by:

Creating backups of the original files
Commenting out problematic model_rebuild() calls that might cause errors
Using regex to find and disable all instances of these calls

Creates a complete mock wandb module that:

Provides a fake wandb object with common methods like init(), log(), and finish()
Sets up all necessary submodules in the sys.modules dictionary
Adds missing attributes and functions to prevent import errors

This is a clever workaround for environments where you want to use code that depends on wandb but either:

Don't have internet access
Don't want to create wandb accounts
Want to avoid the overhead of wandb logging
Are experiencing issues with wandb dependencies

Paste the following code in to a googel colab cell:

# @title Complete Fix for Wandb Dependencies

import os
import sys
import types

# Disable wandb via environment variable
os.environ["WANDB_DISABLED"] = "true"

# 1. Fix pytorch_lightning's wandb logger
wandb_logger_file = '/usr/local/lib/python3.11/dist-packages/pytorch_lightning/loggers/wandb.py'

# Create mock wandb logger implementation
mock_wandb_logger = """
# Mock implementation of WandbLogger that doesn't import wandb
class WandbLogger:
    def __init__(self, *args, **kwargs):
        self.experiment = None
        self.version = None

    def log_metrics(self, *args, **kwargs):
        pass

    def log_hyperparams(self, *args, **kwargs):
        pass

    def watch(self, *args, **kwargs):
        pass

    def finalize(self, *args, **kwargs):
        pass

    def after_save_checkpoint(self, *args, **kwargs):
        pass

    @property
    def name(self):
        return "mock_wandb_logger"

    @property
    def version(self):
        return "0.0.0"
"""

# Backup and replace the pytorch_lightning file
if os.path.exists(wandb_logger_file):
    backup_file = wandb_logger_file + '.backup'
    if not os.path.exists(backup_file):
        with open(wandb_logger_file, 'r') as f:
            original_content = f.read()
        with open(backup_file, 'w') as f:
            f.write(original_content)
        print(f"Created backup of pytorch_lightning wandb logger")

    with open(wandb_logger_file, 'w') as f:
        f.write(mock_wandb_logger)
    print(f"Successfully replaced wandb logger with mock implementation")

# 2. Fix the server_features_query.py file
server_features_query_file = '/usr/local/lib/python3.11/dist-packages/wandb/apis/public/_generated/server_features_query.py'

if os.path.exists(server_features_query_file):
    # Create backup
    backup_file = server_features_query_file + '.backup'
    if not os.path.exists(backup_file):
        with open(server_features_query_file, 'r') as f:
            original_content = f.read()
        with open(backup_file, 'w') as f:
            f.write(original_content)
        print(f"Created backup of server_features_query.py")

    # Replace all model_rebuild() calls in the file
    with open(server_features_query_file, 'r') as f:
        content = f.read()

    modified_content = content.replace(
        "ServerFeaturesQuery.model_rebuild()",
        "# ServerFeaturesQuery.model_rebuild() - commented out to avoid errors"
    ).replace(
        "ServerFeaturesQueryServerInfo.model_rebuild()",
        "# ServerFeaturesQueryServerInfo.model_rebuild() - commented out to avoid errors"
    )

    # Add a more comprehensive approach to comment out all model_rebuild() calls
    import re
    modified_content = re.sub(
        r'([A-Za-z0-9_]+)\.model_rebuild\(\)',
        r'# \1.model_rebuild() - commented out to avoid errors',
        modified_content
    )

    with open(server_features_query_file, 'w') as f:
        f.write(modified_content)

    print("Fixed all model_rebuild() calls in server_features_query.py")

# 3. Create and install a complete mock wandb module
class MockWandb:
    def __init__(self):
        self.run = type('Run', (), {'id': 'mock_run'})
        self.config = {}
        self.summary = {}

    def init(self, *args, **kwargs):
        return self

    def log(self, *args, **kwargs):
        pass

    def finish(self, *args, **kwargs):
        pass

# Create mock classes with model_rebuild method
class MockModelClass:
    @classmethod
    def model_rebuild(cls):
        pass

# Install mock wandb module
sys.modules['wandb'] = MockWandb()

# Create all necessary submodules
for submodule in [
    'wandb.sdk', 'wandb.sdk.artifacts', 'wandb.sdk.artifacts.artifact',
    'wandb.apis', 'wandb.apis.normalize', 'wandb.apis.public',
    'wandb.apis.public.api', 'wandb.apis.public._generated',
    'wandb.apis.public._generated.base', 'wandb.apis.public._generated.server_features_query'
]:
    sys.modules[submodule] = types.ModuleType(submodule)

# Add the missing attributes and functions
sys.modules['wandb.apis.normalize'].normalize_exceptions = lambda x: x
sys.modules['wandb.apis.public._generated'].SERVER_FEATURES_QUERY_GQL = ""
sys.modules['wandb.apis.public._generated'].ServerFeaturesQuery = MockModelClass
sys.modules['wandb.apis.public._generated.server_features_query'].ServerFeaturesQuery = MockModelClass
sys.modules['wandb.apis.public._generated.server_features_query'].ServerFeaturesQueryServerInfo = MockModelClass

print("✅ Complete wandb dependency fix successfully applied")