diff --git a/src/axolotl/cli/config.py b/src/axolotl/cli/config.py index e24e3fe61..c88cf44b5 100644 --- a/src/axolotl/cli/config.py +++ b/src/axolotl/cli/config.py @@ -14,8 +14,8 @@ import yaml from transformers.utils import is_torch_bf16_gpu_available from axolotl.integrations.base import PluginManager -from axolotl.telemetry import TelemetryManager from axolotl.telemetry.errors import send_errors +from axolotl.telemetry.manager import TelemetryManager from axolotl.utils.comet_ import setup_comet_env_vars from axolotl.utils.config import ( normalize_cfg_datasets, diff --git a/src/axolotl/telemetry/callbacks.py b/src/axolotl/telemetry/callbacks.py index ca0aaae92..d664931e0 100644 --- a/src/axolotl/telemetry/callbacks.py +++ b/src/axolotl/telemetry/callbacks.py @@ -45,7 +45,7 @@ class TelemetryCallback(TrainerCallback): **kwargs, # pylint: disable=unused-argument ): """Handle training start.""" - self.telemetry_manager.send_event(event_type="train-start") + self.telemetry_manager.send_event(event_type="train-started") def on_train_end( self, @@ -57,7 +57,7 @@ class TelemetryCallback(TrainerCallback): """Handle training end.""" # Send training completion event self.telemetry_manager.send_event( - event_type="train-end", + event_type="train-ended", properties={ "loss": state.log_history[-1].get("loss", 0) if state.log_history diff --git a/src/axolotl/telemetry/errors.py b/src/axolotl/telemetry/errors.py index 98acd6a2c..42cd5aaf3 100644 --- a/src/axolotl/telemetry/errors.py +++ b/src/axolotl/telemetry/errors.py @@ -148,7 +148,7 @@ def send_errors(func: Callable) -> Callable: # Send error telemetry telemetry_manager.send_event( - event_type=f"{module_path}-error", + event_type=f"{module_path}-errored", properties={ "exception": str(exception), "stack_trace": stack_trace, diff --git a/src/axolotl/train.py b/src/axolotl/train.py index 943b58f22..b6d4304c7 100644 --- a/src/axolotl/train.py +++ b/src/axolotl/train.py @@ -22,8 +22,8 @@ from axolotl.contribs.lgpl.unsloth import ( # pylint: disable = no-name-in-modu fix_untrained_tokens, ) from axolotl.logging_config import configure_logging -from axolotl.telemetry import TelemetryManager from axolotl.telemetry.errors import send_errors +from axolotl.telemetry.manager import TelemetryManager from axolotl.utils.dict import DictDefault from axolotl.utils.freeze import freeze_layers_except from axolotl.utils.models import load_model, load_processor, load_tokenizer @@ -90,11 +90,11 @@ def train( model.generation_config.do_sample = True TELEMETRY_MANAGER.send_event( - event_type="model-load", properties=model.config.to_dict() + event_type="model-loaded", properties=model.config.to_dict() ) if peft_config: TELEMETRY_MANAGER.send_event( - event_type="peft-config-load", properties=peft_config.to_dict() + event_type="peft-config-loaded", properties=peft_config.to_dict() ) model_ref = None