""" Common classes and functions for all trainers. """ import json import os import shutil import time import traceback import requests from accelerate import PartialState from huggingface_hub import HfApi from pydantic import BaseModel from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments from autotrain import is_colab, logger ALLOW_REMOTE_CODE = os.environ.get("ALLOW_REMOTE_CODE", "true").lower() == "true" def get_file_sizes(directory): """ Calculate the sizes of all files in a given directory and its subdirectories. Args: directory (str): The path to the directory to scan for files. Returns: dict: A dictionary where the keys are the file paths and the values are the file sizes in gigabytes (GB). """ file_sizes = {} for root, _, files in os.walk(directory): for file in files: file_path = os.path.join(root, file) file_size = os.path.getsize(file_path) file_size_gb = file_size / (1024**3) # Convert bytes to GB file_sizes[file_path] = file_size_gb return file_sizes def remove_global_step(directory): """ Removes directories that start with 'global_step' within the specified directory. This function traverses the given directory and its subdirectories in a bottom-up manner. If it finds any directory whose name starts with 'global_step', it deletes that directory and all its contents. Args: directory (str): The path to the directory to be traversed and cleaned. Returns: None """ for root, dirs, _ in os.walk(directory, topdown=False): for name in dirs: if name.startswith("global_step"): folder_path = os.path.join(root, name) print(f"Removing folder: {folder_path}") shutil.rmtree(folder_path) def remove_autotrain_data(config): """ Removes the AutoTrain data directory and global step for a given project. Args: config (object): Configuration object that contains the project name. Raises: OSError: If the removal of the directory fails. """ os.system(f"rm -rf {config.project_name}/autotrain-data") remove_global_step(config.project_name) def save_training_params(config): """ Saves the training parameters to a JSON file, excluding the "token" key if it exists. Args: config (object): Configuration object that contains the project name. The function checks if a file named 'training_params.json' exists in the directory specified by `config.project_name`. If the file exists, it loads the JSON content, removes the "token" key if present, and then writes the updated content back to the file. """ if os.path.exists(f"{config.project_name}/training_params.json"): training_params = json.load(open(f"{config.project_name}/training_params.json")) if "token" in training_params: training_params.pop("token") json.dump( training_params, open(f"{config.project_name}/training_params.json", "w"), indent=4, ) def pause_endpoint(params): """ Pauses a Hugging Face endpoint using the provided parameters. Args: params (dict or object): Parameters containing the token required for authorization. If a dictionary is provided, it should have a key "token" with the authorization token. If an object is provided, it should have an attribute `token` with the authorization token. Returns: dict: The JSON response from the API call to pause the endpoint. Raises: KeyError: If the "token" key is missing in the params dictionary. requests.exceptions.RequestException: If there is an issue with the API request. Environment Variables: ENDPOINT_ID: Should be set to the endpoint identifier in the format "username/project_name". """ if isinstance(params, dict): token = params["token"] else: token = params.token endpoint_id = os.environ["ENDPOINT_ID"] username = endpoint_id.split("/")[0] project_name = endpoint_id.split("/")[1] api_url = f"https://api.endpoints.huggingface.cloud/v2/endpoint/{username}/{project_name}/pause" headers = {"Authorization": f"Bearer {token}"} r = requests.post(api_url, headers=headers, timeout=120) return r.json() def pause_space(params, is_failure=False): """ Pauses the Hugging Face space and optionally shuts down the endpoint. This function checks for the presence of "SPACE_ID" and "ENDPOINT_ID" in the environment variables. If "SPACE_ID" is found, it pauses the space and creates a discussion on the Hugging Face platform to notify the user about the status of the training run (success or failure). If "ENDPOINT_ID" is found, it pauses the endpoint. Args: params (object): An object containing the necessary parameters, including the token, username, and project name. is_failure (bool, optional): A flag indicating whether the training run failed. Defaults to False. Raises: Exception: If there is an error while creating the discussion on the Hugging Face platform. Logs: Info: Logs the status of pausing the space and endpoint. Warning: Logs any issues encountered while creating the discussion. Error: Logs if the model failed to train and the discussion was not created. """ if "SPACE_ID" in os.environ: # shut down the space logger.info("Pausing space...") api = HfApi(token=params.token) if is_failure: msg = "Your training run has failed! Please check the logs for more details" title = "Your training has failed ❌" else: msg = "Your training run was successful! [Check out your trained model here]" msg += f"(https://huggingface.co/{params.username}/{params.project_name})" title = "Your training has finished successfully ✅" if not params.token.startswith("hf_oauth_"): try: api.create_discussion( repo_id=os.environ["SPACE_ID"], title=title, description=msg, repo_type="space", ) except Exception as e: logger.warning(f"Failed to create discussion: {e}") if is_failure: logger.error("Model failed to train and discussion was not created.") else: logger.warning("Model trained successfully but discussion was not created.") api.pause_space(repo_id=os.environ["SPACE_ID"]) if "ENDPOINT_ID" in os.environ: # shut down the endpoint logger.info("Pausing endpoint...") pause_endpoint(params) def monitor(func): """ A decorator that wraps a function to monitor its execution and handle exceptions. This decorator performs the following actions: 1. Retrieves the 'config' parameter from the function's keyword arguments or positional arguments. 2. Executes the wrapped function. 3. If an exception occurs during the execution of the wrapped function, logs the error message and stack trace. 4. Optionally pauses the execution if the environment variable 'PAUSE_ON_FAILURE' is set to 1. Args: func (callable): The function to be wrapped by the decorator. Returns: callable: The wrapped function with monitoring capabilities. """ def wrapper(*args, **kwargs): config = kwargs.get("config", None) if config is None and len(args) > 0: config = args[0] try: return func(*args, **kwargs) except Exception as e: error_message = f"""{func.__name__} has failed due to an exception: {traceback.format_exc()}""" logger.error(error_message) logger.error(str(e)) if int(os.environ.get("PAUSE_ON_FAILURE", 1)) == 1: pause_space(config, is_failure=True) return wrapper class AutoTrainParams(BaseModel): """ AutoTrainParams is a base class for all AutoTrain parameters. Attributes: Config (class): Configuration class for Pydantic model. protected_namespaces (tuple): Protected namespaces for the model. Methods: save(output_dir): Save parameters to a JSON file in the specified output directory. __str__(): Return a string representation of the parameters, masking the token if present. __init__(**data): Initialize the parameters, check for unused/extra parameters, and warn the user if necessary. Raises ValueError if project_name is not alphanumeric (with hyphens allowed) or exceeds 50 characters. """ class Config: protected_namespaces = () def save(self, output_dir): """ Save parameters to a json file. """ os.makedirs(output_dir, exist_ok=True) path = os.path.join(output_dir, "training_params.json") # save formatted json with open(path, "w", encoding="utf-8") as f: f.write(self.model_dump_json(indent=4)) def __str__(self): """ String representation of the parameters. """ data = self.model_dump() data["token"] = "*****" if data.get("token") else None return str(data) def __init__(self, **data): """ Initialize the parameters, check for unused/extra parameters and warn the user. """ super().__init__(**data) if len(self.project_name) > 0: # make sure project_name is always alphanumeric but can have hyphens. if not, raise ValueError if not self.project_name.replace("-", "").isalnum(): raise ValueError("project_name must be alphanumeric but can contain hyphens") # project name cannot be more than 50 characters if len(self.project_name) > 50: raise ValueError("project_name cannot be more than 50 characters") # Parameters not supplied by the user defaults = set(self.model_fields.keys()) supplied = set(data.keys()) not_supplied = defaults - supplied if not_supplied and not is_colab: logger.warning(f"Parameters not supplied by user and set to default: {', '.join(not_supplied)}") # Parameters that were supplied but not used # This is a naive implementation. It might catch some internal Pydantic params. unused = supplied - set(self.model_fields) if unused: logger.warning(f"Parameters supplied but not used: {', '.join(unused)}") class UploadLogs(TrainerCallback): """ A callback to upload training logs to the Hugging Face Hub. Args: config (object): Configuration object containing necessary parameters. Attributes: config (object): Configuration object containing necessary parameters. api (HfApi or None): Instance of HfApi for interacting with the Hugging Face Hub. last_upload_time (float): Timestamp of the last upload. Methods: on_step_end(args, state, control, **kwargs): Called at the end of each training step. Uploads logs to the Hugging Face Hub if conditions are met. """ def __init__(self, config): self.config = config self.api = None self.last_upload_time = 0 if self.config.push_to_hub: if PartialState().process_index == 0: self.api = HfApi(token=config.token) self.api.create_repo( repo_id=f"{self.config.username}/{self.config.project_name}", repo_type="model", private=True ) def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): if self.config.push_to_hub is False: return control if not os.path.exists(os.path.join(self.config.project_name, "runs")): return control if (state.global_step + 1) % self.config.logging_steps == 0 and self.config.log == "tensorboard": if PartialState().process_index == 0: current_time = time.time() if current_time - self.last_upload_time >= 600: try: self.api.upload_folder( folder_path=os.path.join(self.config.project_name, "runs"), repo_id=f"{self.config.username}/{self.config.project_name}", path_in_repo="runs", ) except Exception as e: logger.warning(f"Failed to upload logs: {e}") logger.warning("Continuing training...") self.last_upload_time = current_time return control class LossLoggingCallback(TrainerCallback): """ LossLoggingCallback is a custom callback for logging loss during training. This callback inherits from `TrainerCallback` and overrides the `on_log` method to remove the "total_flos" key from the logs and log the remaining information if the current process is the local process zero. Methods: on_log(args, state, control, logs=None, **kwargs): Called when the logs are updated. Removes the "total_flos" key from the logs and logs the remaining information if the current process is the local process zero. Args: args: The training arguments. state: The current state of the Trainer. control: The control object for the Trainer. logs (dict, optional): The logs dictionary containing the training metrics. **kwargs: Additional keyword arguments. """ def on_log(self, args, state, control, logs=None, **kwargs): _ = logs.pop("total_flos", None) if state.is_local_process_zero: logger.info(logs) class TrainStartCallback(TrainerCallback): """ TrainStartCallback is a custom callback for the Trainer class that logs a message when training begins. Methods: on_train_begin(args, state, control, **kwargs): Logs a message indicating that training is starting. Args: args: The training arguments. state: The current state of the Trainer. control: The control object for the Trainer. **kwargs: Additional keyword arguments. """ def on_train_begin(self, args, state, control, **kwargs): logger.info("Starting to train...")