import base64 import datetime import io import json import os import re import tempfile import time from typing import Any, Dict, List, Optional, Tuple, Union from packaging.version import parse import wandb from wandb import util from wandb.data_types import Table from wandb.sdk.lib import telemetry openai = util.get_module( name="openai", required="This integration requires `openai`. To install, please run `pip install openai`", lazy=False, ) if parse(openai.__version__) < parse("1.12.0"): raise wandb.Error( f"This integration requires openai version 1.12.0 and above. Your current version is {openai.__version__} " "To fix, please `pip install -U openai`" ) from openai import OpenAI # noqa: E402 from openai.types.fine_tuning import FineTuningJob # noqa: E402 from openai.types.fine_tuning.fine_tuning_job import ( # noqa: E402 Error, Hyperparameters, ) np = util.get_module( name="numpy", required="`numpy` not installed >> This integration requires numpy! To fix, please `pip install numpy`", lazy=False, ) pd = util.get_module( name="pandas", required="`pandas` not installed >> This integration requires pandas! To fix, please `pip install pandas`", lazy=False, ) class WandbLogger: """Log OpenAI fine-tunes to [Weights & Biases](https://wandb.me/openai-docs).""" _wandb_api: Optional[wandb.Api] = None _logged_in: bool = False openai_client: Optional[OpenAI] = None _run: Optional[wandb.Run] = None @classmethod def sync( cls, fine_tune_job_id: Optional[str] = None, openai_client: Optional[OpenAI] = None, num_fine_tunes: Optional[int] = None, project: str = "OpenAI-Fine-Tune", entity: Optional[str] = None, overwrite: bool = False, wait_for_job_success: bool = True, log_datasets: bool = True, model_artifact_name: str = "model-metadata", model_artifact_type: str = "model", **kwargs_wandb_init: Dict[str, Any], ) -> str: """Sync fine-tunes to Weights & Biases. :param fine_tune_job_id: The id of the fine-tune (optional) :param openai_client: Pass the `OpenAI()` client (optional) :param num_fine_tunes: Number of most recent fine-tunes to log when an fine_tune_job_id is not provided. By default, every fine-tune is synced. :param project: Name of the project where you're sending runs. By default, it is "GPT-3". :param entity: Username or team name where you're sending runs. By default, your default entity is used, which is usually your username. :param overwrite: Forces logging and overwrite existing wandb run of the same fine-tune. :param wait_for_job_success: Waits for the fine-tune to be complete and then log metrics to W&B. By default, it is True. :param model_artifact_name: Name of the model artifact that is logged :param model_artifact_type: Type of the model artifact that is logged """ if openai_client is None: openai_client = OpenAI() cls.openai_client = openai_client if fine_tune_job_id: wandb.termlog("Retrieving fine-tune job...") fine_tune = openai_client.fine_tuning.jobs.retrieve( fine_tuning_job_id=fine_tune_job_id ) fine_tunes = [fine_tune] else: # get list of fine_tune to log fine_tunes = openai_client.fine_tuning.jobs.list() if not fine_tunes or fine_tunes.data is None: wandb.termwarn("No fine-tune has been retrieved") return # Select the `num_fine_tunes` from the `fine_tunes.data` list. # If `num_fine_tunes` is None, it selects all items in the list (from start to end). # If for example, `num_fine_tunes` is 5, it selects the last 5 items in the list. # Note that the last items in the list are the latest fine-tune jobs. fine_tunes = fine_tunes.data[ -num_fine_tunes if num_fine_tunes is not None else None : ] # log starting from oldest fine_tune show_individual_warnings = ( fine_tune_job_id is not None or num_fine_tunes is not None ) fine_tune_logged = [] for fine_tune in fine_tunes: fine_tune_id = fine_tune.id # check run with the given `fine_tune_id` has not been logged already run_path = f"{project}/{fine_tune_id}" if entity is not None: run_path = f"{entity}/{run_path}" wandb_run = cls._get_wandb_run(run_path) if wandb_run: wandb_status = wandb_run.summary.get("status") if show_individual_warnings: if wandb_status == "succeeded" and not overwrite: wandb.termwarn( f"Fine-tune {fine_tune_id} has already been logged successfully at {wandb_run.url}. " "Use `overwrite=True` if you want to overwrite previous run" ) elif wandb_status != "succeeded" or overwrite: if wandb_status != "succeeded": wandb.termwarn( f"A run for fine-tune {fine_tune_id} was previously created but didn't end successfully" ) wandb.termlog( f"A new wandb run will be created for fine-tune {fine_tune_id} and previous run will be overwritten" ) overwrite = True if wandb_status == "succeeded" and not overwrite: return # check if the user has not created a wandb run externally if wandb.run is None: cls._run = wandb.init( job_type="fine-tune", project=project, entity=entity, name=fine_tune_id, id=fine_tune_id, **kwargs_wandb_init, ) else: # if a run exits - created externally cls._run = wandb.run if wait_for_job_success: fine_tune = cls._wait_for_job_success(fine_tune) cls._log_fine_tune( fine_tune, project, entity, overwrite, show_individual_warnings, log_datasets, model_artifact_name, model_artifact_type, **kwargs_wandb_init, ) if not show_individual_warnings and not any(fine_tune_logged): wandb.termwarn("No new successful fine-tunes were found") return "🎉 wandb sync completed successfully" @classmethod def _wait_for_job_success(cls, fine_tune: FineTuningJob) -> FineTuningJob: wandb.termlog("Waiting for the OpenAI fine-tuning job to finish training...") wandb.termlog( "To avoid blocking, you can call `WandbLogger.sync` with `wait_for_job_success=False` after OpenAI training completes." ) while True: if fine_tune.status == "succeeded": wandb.termlog( "Fine-tuning finished, logging metrics, model metadata, and run metadata to Weights & Biases" ) return fine_tune if fine_tune.status == "failed": wandb.termwarn( f"Fine-tune {fine_tune.id} has failed and will not be logged" ) return fine_tune if fine_tune.status == "cancelled": wandb.termwarn( f"Fine-tune {fine_tune.id} was cancelled and will not be logged" ) return fine_tune time.sleep(10) fine_tune = cls.openai_client.fine_tuning.jobs.retrieve( fine_tuning_job_id=fine_tune.id ) @classmethod def _log_fine_tune( cls, fine_tune: FineTuningJob, project: str, entity: Optional[str], overwrite: bool, show_individual_warnings: bool, log_datasets: bool, model_artifact_name: str, model_artifact_type: str, **kwargs_wandb_init: Dict[str, Any], ): fine_tune_id = fine_tune.id status = fine_tune.status with telemetry.context(run=cls._run) as tel: tel.feature.openai_finetuning = True # check run completed successfully if status != "succeeded": if show_individual_warnings: wandb.termwarn( f'Fine-tune {fine_tune_id} has the status "{status}" and will not be logged' ) return # check results are present try: results_id = fine_tune.result_files[0] try: encoded_results = cls.openai_client.files.content( file_id=results_id ).read() results = base64.b64decode(encoded_results).decode("utf-8") except Exception: # attempt to read as text, works for older jobs results = cls.openai_client.files.content(file_id=results_id).text except openai.NotFoundError: if show_individual_warnings: wandb.termwarn( f"Fine-tune {fine_tune_id} has no results and will not be logged" ) return # update the config cls._run.config.update(cls._get_config(fine_tune)) # log results df_results = pd.read_csv(io.StringIO(results)) for _, row in df_results.iterrows(): metrics = {k: v for k, v in row.items() if not np.isnan(v)} step = metrics.pop("step") if step is not None: step = int(step) cls._run.log(metrics, step=step) fine_tuned_model = fine_tune.fine_tuned_model if fine_tuned_model is not None: cls._run.summary["fine_tuned_model"] = fine_tuned_model # training/validation files and fine-tune details cls._log_artifacts( fine_tune, project, entity, log_datasets, overwrite, model_artifact_name, model_artifact_type, ) # mark run as complete cls._run.summary["status"] = "succeeded" cls._run.finish() return True @classmethod def _ensure_logged_in(cls): if not cls._logged_in: if wandb.login(): cls._logged_in = True else: raise Exception( "It appears you are not currently logged in to Weights & Biases. " "Please run `wandb login` in your terminal or `wandb.login()` in a notebook. " "Create a new API key at https://wandb.ai/settings and store it securely." ) @classmethod def _get_wandb_run(cls, run_path: str): cls._ensure_logged_in() try: if cls._wandb_api is None: cls._wandb_api = wandb.Api() return cls._wandb_api.run(run_path) except Exception: return None @classmethod def _get_wandb_artifact(cls, artifact_path: str): cls._ensure_logged_in() try: if cls._wandb_api is None: cls._wandb_api = wandb.Api() return cls._wandb_api.artifact(artifact_path) except Exception: return None @classmethod def _get_config(cls, fine_tune: FineTuningJob) -> Dict[str, Any]: config = dict(fine_tune) config["result_files"] = config["result_files"][0] if config.get("created_at"): config["created_at"] = datetime.datetime.fromtimestamp( config["created_at"] ).strftime("%Y-%m-%d %H:%M:%S") if config.get("finished_at"): config["finished_at"] = datetime.datetime.fromtimestamp( config["finished_at"] ).strftime("%Y-%m-%d %H:%M:%S") if config.get("hyperparameters"): config["hyperparameters"] = cls.sanitize(config["hyperparameters"]) if config.get("error"): config["error"] = cls.sanitize(config["error"]) return config @classmethod def _unpack_hyperparameters(cls, hyperparameters: Hyperparameters): # `Hyperparameters` object is not unpacking properly using `vars` or `__dict__`, # vars(hyperparameters) return {n_epochs: n} only. hyperparams = {} try: hyperparams["n_epochs"] = hyperparameters.n_epochs hyperparams["batch_size"] = hyperparameters.batch_size hyperparams["learning_rate_multiplier"] = ( hyperparameters.learning_rate_multiplier ) except Exception: # If unpacking fails, return the object to be logged as config return None return hyperparams @staticmethod def sanitize(input: Any) -> Union[Dict, List, str]: valid_types = [bool, int, float, str] if isinstance(input, (Hyperparameters, Error)): return dict(input) if isinstance(input, dict): return { k: v if type(v) in valid_types else str(v) for k, v in input.items() } elif isinstance(input, list): return [v if type(v) in valid_types else str(v) for v in input] else: return str(input) @classmethod def _log_artifacts( cls, fine_tune: FineTuningJob, project: str, entity: Optional[str], log_datasets: bool, overwrite: bool, model_artifact_name: str, model_artifact_type: str, ) -> None: if log_datasets: wandb.termlog("Logging training/validation files...") # training/validation files training_file = fine_tune.training_file if fine_tune.training_file else None validation_file = ( fine_tune.validation_file if fine_tune.validation_file else None ) for file, prefix, artifact_type in ( (training_file, "train", "training_files"), (validation_file, "valid", "validation_files"), ): if file is not None: cls._log_artifact_inputs( file, prefix, artifact_type, project, entity, overwrite ) # fine-tune details fine_tune_id = fine_tune.id artifact = wandb.Artifact( model_artifact_name, type=model_artifact_type, metadata=dict(fine_tune), ) with artifact.new_file("model_metadata.json", mode="w", encoding="utf-8") as f: dict_fine_tune = dict(fine_tune) dict_fine_tune["hyperparameters"] = cls.sanitize( dict_fine_tune["hyperparameters"] ) dict_fine_tune["error"] = cls.sanitize(dict_fine_tune["error"]) dict_fine_tune = cls.sanitize(dict_fine_tune) json.dump(dict_fine_tune, f, indent=2) cls._run.log_artifact( artifact, aliases=["latest", fine_tune_id], ) @classmethod def _log_artifact_inputs( cls, file_id: Optional[str], prefix: str, artifact_type: str, project: str, entity: Optional[str], overwrite: bool, ) -> None: # get input artifact artifact_name = f"{prefix}-{file_id}" # sanitize name to valid wandb artifact name artifact_name = re.sub(r"[^a-zA-Z0-9_\-.]", "_", artifact_name) artifact_alias = file_id artifact_path = f"{project}/{artifact_name}:{artifact_alias}" if entity is not None: artifact_path = f"{entity}/{artifact_path}" artifact = cls._get_wandb_artifact(artifact_path) # create artifact if file not already logged previously if artifact is None or overwrite: # get file content try: file_content = cls.openai_client.files.content(file_id=file_id) except openai.NotFoundError: wandb.termerror( f"File {file_id} could not be retrieved. Make sure you have OpenAI permissions to download training/validation files" ) return artifact = wandb.Artifact(artifact_name, type=artifact_type) with tempfile.NamedTemporaryFile(delete=False) as tmp_file: tmp_file.write(file_content.content) tmp_file_path = tmp_file.name artifact.add_file(tmp_file_path, file_id) os.unlink(tmp_file_path) # create a Table try: table, n_items = cls._make_table(file_content.text) # Add table to the artifact. artifact.add(table, file_id) # Add the same table to the workspace. cls._run.log({f"{prefix}_data": table}) # Update the run config and artifact metadata cls._run.config.update({f"n_{prefix}": n_items}) artifact.metadata["items"] = n_items except Exception as e: wandb.termerror( f"Issue saving {file_id} as a Table to Artifacts, exception:\n '{e}'" ) else: # log number of items cls._run.config.update({f"n_{prefix}": artifact.metadata.get("items")}) cls._run.use_artifact(artifact, aliases=["latest", artifact_alias]) @classmethod def _make_table(cls, file_content: str) -> Tuple[Table, int]: table = wandb.Table(columns=["role: system", "role: user", "role: assistant"]) df = pd.read_json(io.StringIO(file_content), orient="records", lines=True) for _idx, message in df.iterrows(): messages = message.messages assert len(messages) == 3 table.add_data( messages[0]["content"], messages[1]["content"], messages[2]["content"], ) return table, len(df)