Merge pull request #53 from togethercomputer/orangetin/handle-errors-fix

Handle errors gracefully

Author: orangetin

src/together/__init__.py

@@ -21,6 +21,10 @@
 default_image_model = "runwayml/stable-diffusion-v1-5"
 log_level = "WARNING"
 
+MISSING_API_KEY_MESSAGE = """TOGETHER_API_KEY not found.
+Please set it as an environment variable or set it as together.api_key
+Find your TOGETHER_API_KEY at https://api.together.xyz/settings/api-keys"""
+
 MAX_CONNECTION_RETRIES = 2
 BACKOFF_FACTOR = 0.2
 
@@ -49,6 +53,7 @@
     "Finetune",
     "Image",
     "MAX_CONNECTION_RETRIES",
+    "MISSING_API_KEY_MESSAGE",
     "BACKOFF_FACTOR",
     "min_samples",
 ]

src/together/commands/chat.py

@@ -6,6 +6,10 @@
 import together
 import together.tools.conversation as convo
 from together import Complete
+from together.utils import get_logger
+
+
+logger = get_logger(str(__name__))
 
 
 def add_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
@@ -99,18 +103,22 @@ def precmd(self, line: str) -> str:
     def do_say(self, arg: str) -> None:
         self._convo.push_human_turn(arg)
         output = ""
-        for token in self.infer.create_streaming(
-            prompt=self._convo.get_raw_prompt(),
-            model=self.args.model,
-            max_tokens=self.args.max_tokens,
-            stop=self.args.stop,
-            temperature=self.args.temperature,
-            top_p=self.args.top_p,
-            top_k=self.args.top_k,
-            repetition_penalty=self.args.repetition_penalty,
-        ):
-            print(token, end="", flush=True)
-            output += token
+        try:
+            for token in self.infer.create_streaming(
+                prompt=self._convo.get_raw_prompt(),
+                model=self.args.model,
+                max_tokens=self.args.max_tokens,
+                stop=self.args.stop,
+                temperature=self.args.temperature,
+                top_p=self.args.top_p,
+                top_k=self.args.top_k,
+                repetition_penalty=self.args.repetition_penalty,
+            ):
+                print(token, end="", flush=True)
+                output += token
+        except together.AuthenticationError:
+            logger.critical(together.MISSING_API_KEY_MESSAGE)
+            exit(0)
         print("\n")
         self._convo.push_model_response(output)
 

src/together/commands/complete.py

@@ -131,33 +131,41 @@ def _run_complete(args: argparse.Namespace) -> None:
     complete = Complete()
 
     if args.no_stream:
-        response = complete.create(
-            prompt=args.prompt,
-            model=args.model,
-            max_tokens=args.max_tokens,
-            stop=args.stop,
-            temperature=args.temperature,
-            top_p=args.top_p,
-            top_k=args.top_k,
-            repetition_penalty=args.repetition_penalty,
-            logprobs=args.logprobs,
-        )
+        try:
+            response = complete.create(
+                prompt=args.prompt,
+                model=args.model,
+                max_tokens=args.max_tokens,
+                stop=args.stop,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                repetition_penalty=args.repetition_penalty,
+                logprobs=args.logprobs,
+            )
+        except together.AuthenticationError:
+            logger.critical(together.MISSING_API_KEY_MESSAGE)
+            exit(0)
         no_streamer(args, response)
     else:
-        for text in complete.create_streaming(
-            prompt=args.prompt,
-            model=args.model,
-            max_tokens=args.max_tokens,
-            stop=args.stop,
-            temperature=args.temperature,
-            top_p=args.top_p,
-            top_k=args.top_k,
-            repetition_penalty=args.repetition_penalty,
-            raw=args.raw,
-        ):
-            if not args.raw:
-                print(text, end="", flush=True)
-            else:
-                print(text)
+        try:
+            for text in complete.create_streaming(
+                prompt=args.prompt,
+                model=args.model,
+                max_tokens=args.max_tokens,
+                stop=args.stop,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                repetition_penalty=args.repetition_penalty,
+                raw=args.raw,
+            ):
+                if not args.raw:
+                    print(text, end="", flush=True)
+                else:
+                    print(text)
+        except together.AuthenticationError:
+            logger.critical(together.MISSING_API_KEY_MESSAGE)
+            exit(0)
         if not args.raw:
             print("\n")

src/together/commands/files.py

@@ -5,8 +5,12 @@
 
 from tabulate import tabulate
 
+import together
 from together import Files
-from together.utils import bytes_to_human_readable, extract_time
+from together.utils import bytes_to_human_readable, extract_time, get_logger
+
+
+logger = get_logger(str(__name__))
 
 
 def add_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
@@ -123,7 +127,11 @@ def _add_retrieve_content(
 
 
 def _run_list(args: argparse.Namespace) -> None:
-    response = Files.list()
+    try:
+        response = Files.list()
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     response["data"].sort(key=extract_time)
     if args.raw:
         print(json.dumps(response, indent=4))
@@ -146,22 +154,40 @@ def _run_list(args: argparse.Namespace) -> None:
 
 
 def _run_check(args: argparse.Namespace) -> None:
-    response = Files.check(args.file)
+    try:
+        response = Files.check(args.file)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(json.dumps(response, indent=4))
 
 
 def _run_upload(args: argparse.Namespace) -> None:
-    response = Files.upload(file=args.file, check=not args.no_check, model=args.model)
+    try:
+        response = Files.upload(
+            file=args.file, check=not args.no_check, model=args.model
+        )
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(json.dumps(response, indent=4))
 
 
 def _run_delete(args: argparse.Namespace) -> None:
-    response = Files.delete(args.file_id)
+    try:
+        response = Files.delete(args.file_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(json.dumps(response, indent=4))
 
 
 def _run_retrieve(args: argparse.Namespace) -> None:
-    response = Files.retrieve(args.file_id)
+    try:
+        response = Files.retrieve(args.file_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     if args.raw:
         print(json.dumps(response, indent=4))
     else:
@@ -171,5 +197,9 @@ def _run_retrieve(args: argparse.Namespace) -> None:
 
 
 def _run_retrieve_content(args: argparse.Namespace) -> None:
-    output = Files.retrieve_content(args.file_id, args.output)
+    try:
+        output = Files.retrieve_content(args.file_id, args.output)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(output)

src/together/commands/finetune.py

@@ -6,8 +6,12 @@
 
 from tabulate import tabulate
 
+import together
 from together import Finetune
-from together.utils import finetune_price_to_dollars, parse_timestamp
+from together.utils import finetune_price_to_dollars, get_logger, parse_timestamp
+
+
+logger = get_logger(str(__name__))
 
 
 def add_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
@@ -24,7 +28,6 @@ def add_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser])
     _add_download(child_parsers)
     _add_status(child_parsers)
     _add_checkpoints(child_parsers)
-    # _add_delete_model(child_parsers)
 
 
 def _add_create(parser: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
@@ -252,25 +255,32 @@ def _run_create(args: argparse.Namespace) -> None:
             args.batch_size = 144
         else:
             args.batch_size = 32
-
-    response = finetune.create(
-        training_file=args.training_file,  # training file_id
-        model=args.model,
-        n_epochs=args.n_epochs,
-        n_checkpoints=args.n_checkpoints,
-        batch_size=args.batch_size,
-        learning_rate=args.learning_rate,
-        suffix=args.suffix,
-        estimate_price=args.estimate_price,
-        wandb_api_key=args.wandb_api_key if not args.no_wandb_api_key else None,
-        confirm_inputs=not args.quiet,
-    )
+    try:
+        response = finetune.create(
+            training_file=args.training_file,  # training file_id
+            model=args.model,
+            n_epochs=args.n_epochs,
+            n_checkpoints=args.n_checkpoints,
+            batch_size=args.batch_size,
+            learning_rate=args.learning_rate,
+            suffix=args.suffix,
+            estimate_price=args.estimate_price,
+            wandb_api_key=args.wandb_api_key if not args.no_wandb_api_key else None,
+            confirm_inputs=not args.quiet,
+        )
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
 
     print(json.dumps(response, indent=4))
 
 
 def _run_list(args: argparse.Namespace) -> None:
-    response = Finetune.list()
+    try:
+        response = Finetune.list()
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     response["data"].sort(key=lambda x: parse_timestamp(x["created_at"]))
     if args.raw:
         print(json.dumps(response, indent=4))
@@ -293,7 +303,11 @@ def _run_list(args: argparse.Namespace) -> None:
 
 
 def _run_retrieve(args: argparse.Namespace) -> None:
-    response = Finetune.retrieve(args.fine_tune_id)
+    try:
+        response = Finetune.retrieve(args.fine_tune_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     if args.raw:
         print(json.dumps(response, indent=4))
     else:
@@ -307,12 +321,20 @@ def _run_retrieve(args: argparse.Namespace) -> None:
 
 
 def _run_cancel(args: argparse.Namespace) -> None:
-    response = Finetune.cancel(args.fine_tune_id)
+    try:
+        response = Finetune.cancel(args.fine_tune_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(json.dumps(response, indent=4))
 
 
 def _run_list_events(args: argparse.Namespace) -> None:
-    response = Finetune.list_events(args.fine_tune_id)
+    try:
+        response = Finetune.list_events(args.fine_tune_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     if args.raw:
         print(json.dumps(response, indent=4))
     else:
@@ -330,16 +352,30 @@ def _run_list_events(args: argparse.Namespace) -> None:
 
 
 def _run_download(args: argparse.Namespace) -> None:
-    response = Finetune.download(args.fine_tune_id, args.output, args.checkpoint_step)
+    try:
+        response = Finetune.download(
+            args.fine_tune_id, args.output, args.checkpoint_step
+        )
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(response)
 
 
 def _run_status(args: argparse.Namespace) -> None:
-    response = Finetune.get_job_status(args.fine_tune_id)
+    try:
+        response = Finetune.get_job_status(args.fine_tune_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(response)
 
 
 def _run_checkpoint(args: argparse.Namespace) -> None:
-    checkpoints = Finetune.get_checkpoints(args.fine_tune_id)
+    try:
+        checkpoints = Finetune.get_checkpoints(args.fine_tune_id)
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
     print(json.dumps(checkpoints, indent=4))
     print(f"\n{len(checkpoints)} checkpoints found")

src/together/commands/image.py

@@ -124,16 +124,19 @@ def _save_image(args: argparse.Namespace, response: Dict[str, Any]) -> None:
 
 def _run_complete(args: argparse.Namespace) -> None:
     complete = Image()
-
-    response = complete.create(
-        prompt=args.prompt,
-        model=args.model,
-        steps=args.steps,
-        seed=args.seed,
-        results=args.results,
-        height=args.height,
-        width=args.width,
-        negative_prompt=args.negative_prompt,
-    )
+    try:
+        response = complete.create(
+            prompt=args.prompt,
+            model=args.model,
+            steps=args.steps,
+            seed=args.seed,
+            results=args.results,
+            height=args.height,
+            width=args.width,
+            negative_prompt=args.negative_prompt,
+        )
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)
 
     _save_image(args, response)