tasks/utils/emissions.py

from codecarbon import EmissionsTracker
import os

# Initialize tracker with basic configuration
tracker = EmissionsTracker(
    project_name="climate-guard",
    output_dir=os.getenv("OUTPUT_DIR", "./"),
    log_level='error'
)

class EmissionsData:
    def __init__(self, energy_consumed: float, emissions: float):
        self.energy_consumed = energy_consumed
        self.emissions = emissions
        self.timestamp = None
        self.project_name = None
        self.experiment_id = None
        self.latitude = None
        self.longitude = None

def clean_emissions_data(emissions_data):
    """Remove unwanted fields from emissions data"""
    if hasattr(emissions_data, '__dict__'):
        data_dict = emissions_data.__dict__
    else:
        # If emissions_data is not an object with __dict__
        data_dict = {
            'energy_consumed': getattr(emissions_data, 'energy_consumed', 0),
            'emissions': getattr(emissions_data, 'emissions', 0)
        }
    
    fields_to_remove = ['timestamp', 'project_name', 'experiment_id', 'latitude', 'longitude']
    return {k: v for k, v in data_dict.items() if k not in fields_to_remove}

def get_space_info():
    """Get the space username and URL from environment variables"""
    space_name = os.getenv("SPACE_ID", "")
    if space_name:
        try:
            username = space_name.split("/")[0]
            space_url = f"https://huggingface.co/spaces/{space_name}"
            return username, space_url
        except Exception as e:
            print(f"Error getting space info: {e}")
    return "local-user", "local-development"

def start_tracking():
    """Safely start the emissions tracking"""
    try:
        if not tracker._tracking:
            tracker.start()
        return True
    except Exception as e:
        print(f"Error starting emissions tracking: {e}")
        return False

def stop_tracking():
    """Safely stop the emissions tracking and return data"""
    try:
        if tracker._tracking:
            emissions = tracker.stop()
            return EmissionsData(
                energy_consumed=getattr(emissions, 'energy_consumed', 0),
                emissions=getattr(emissions, 'emissions', 0)
            )
    except Exception as e:
        print(f"Error stopping emissions tracking: {e}")
        return EmissionsData(energy_consumed=0, emissions=0)