Merge pull request #25 from togethercomputer/justusc/enable-wandb

Enable WandB option in the CLI

README.md

@@ -295,6 +295,7 @@ resp = together.Finetune.create(
   batch_size = 4,
   learning_rate = 1e-5,
   suffix = 'my-demo-finetune',
+  wandb_api_key = '1a2b3c4d5e.......',
 )
 
 fine_tune_id = resp['id']

src/together/commands/finetune.py

@@ -2,6 +2,7 @@
 
 import argparse
 import json
+import os
 
 from together import Finetune, extract_time
 
@@ -136,21 +137,21 @@ def _add_create(parser: argparse._SubParsersAction[argparse.ArgumentParser]) ->
         help="Up to 40 characters that will be added to your fine-tuned model name.",
         type=str,
     )
-    # subparser.add_argument(
-    #    "--wandb-api-key",
-    #    "-wb",
-    #    metavar="WANDB_API_KEY",
-    #    default=os.getenv("WANDB_API_KEY"),
-    #    help="Wandb API key to report metrics to wandb.ai. If not set WANDB_API_KEY environment variable is used.",
-    #    type=str,
-    # )
-    # subparser.add_argument(
-    #    "--no-wandb-api-key",
-    #    "-nwb",
-    #    default=False,
-    #    help="Do not report metrics to wandb.ai.",
-    #    action="store_true",
-    # )
+    subparser.add_argument(
+        "--wandb-api-key",
+        "-wb",
+        metavar="WANDB_API_KEY",
+        default=os.getenv("WANDB_API_KEY"),
+        help="Wandb API key to report metrics to wandb.ai. If not set WANDB_API_KEY environment variable is used.",
+        type=str,
+    )
+    subparser.add_argument(
+        "--no-wandb-api-key",
+        "-nwb",
+        default=False,
+        help="Do not report metrics to wandb.ai.",
+        action="store_true",
+    )
 
     subparser.set_defaults(func=_run_create)
 
@@ -296,6 +297,7 @@ def _run_create(args: argparse.Namespace) -> None:
         # checkpoint_steps=args.checkpoint_steps,
         suffix=args.suffix,
         estimate_price=args.estimate_price,
+        wandb_api_key=args.wandb_api_key if not args.no_wandb_api_key else None,
     )
 
     print(json.dumps(response, indent=4))

src/together/finetune.py

@@ -69,6 +69,7 @@ def create(
             str
         ] = None,  # resulting finetuned model name will include the suffix
         estimate_price: bool = False,
+        wandb_api_key: Optional[str] = None,
     ) -> Dict[Any, Any]:
         if n_epochs is None or n_epochs < 1:
             logger.fatal("The number of epochs must be specified")
@@ -102,7 +103,7 @@ def create(
             # "seed": seed,
             # "fp16": fp16,
             "suffix": suffix,
-            # "wandb_key": wandb_api_key,
+            "wandb_key": wandb_api_key,
         }
 
         # check if model name is one of the models available for finetuning