SMDebug RulesConfig

from sagemaker.debugger import Rule, CollectionConfig, rule_configs
from sagemaker.estimator import Estimator # Placeholder for a SageMaker Estimator
import os

# In a real scenario, replace this with your actual SageMaker Estimator setup
# and ensure AWS credentials are configured (e.g., via environment variables or AWS CLI)
# For example: estimator = TensorFlow(role=os.environ.get('SAGEMAKER_ROLE', 'arn:aws:iam::123456789012:role/SageMakerRole'), ...)
# We use a placeholder here for quickstart reproducibility.

# Example: Vanilla built-in rule without customization
# This would typically be passed to a SageMaker Estimator's 'rules' parameter.
rule_vanishing_gradient = Rule.sagemaker(rule_configs.vanishing_gradient())
print(f"Vanishing Gradient Rule: {rule_vanishing_gradient.name}")

# Example: Built-in rule with customization
# This demonstrates how to customize a rule's parameters and collections.
rule_customized_weight_update = Rule.sagemaker(
    base_config=rule_configs.weight_update_ratio(),
    name="my_custom_wup_rule", # Optional
    # container_local_path="/local/path", # Optional, if running locally or specific container path
    # s3_output_path="s3://your-s3-bucket/debug-output/", # Optional, overrides default
    rule_parameters={
        "threshold": "0.001" # Example parameter customization
    },
    collections_to_save=[
        CollectionConfig(
            name="weights", # Required. Debugger will collect tensors for this collection.
            parameters={
                "save_interval": "100" # Example collection parameter
            }
        )
    ]
)
print(f"Customized Weight Update Ratio Rule: {rule_customized_weight_update.name}")
print(f"  Rule parameters: {rule_customized_weight_update.rule_parameters}")
print(f"  Collections to save: {[c.name for c in rule_customized_weight_update.collections_to_save]}")