From a30c14ec023f0c420e5cb65514c7dd5d7acf748d Mon Sep 17 00:00:00 2001 From: K-rolls Date: Tue, 14 Jan 2025 13:01:53 -0330 Subject: [PATCH 01/12] add create endpoints --- src/api/machine_api.py | 26 +++++++++++++++++++++++++- src/commands/machine.py | 10 ++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/api/machine_api.py b/src/api/machine_api.py index d70c191..dcdb7f7 100644 --- a/src/api/machine_api.py +++ b/src/api/machine_api.py @@ -8,7 +8,9 @@ def __init__(self): self.client = APIClient() @handle_api_errors - def create_machine(self, machine_name: str, machine_type: str) -> Dict[str, str]: + def create_fpga_machine( + self, machine_name: str, machine_type: str + ) -> Dict[str, str]: data = { "machine_name": machine_name, "machine_type": machine_type, @@ -16,6 +18,28 @@ def create_machine(self, machine_name: str, machine_type: str) -> Dict[str, str] response = self.client.post("machine/fpga", data=data) return response + @handle_api_errors + def create_gpu_machine( + self, machine_name: str, machine_type: str + ) -> Dict[str, str]: + data = { + "machine_name": machine_name, + "machine_type": machine_type, + } + response = self.client.post("machine/gpu", data=data) + return response + + @handle_api_errors + def create_cpu_machine( + self, machine_name: str, machine_type: str + ) -> Dict[str, str]: + data = { + "machine_name": machine_name, + "machine_type": machine_type, + } + response = self.client.post("machine/cpu", data=data) + return response + @handle_api_errors def list_user_machines(self): return self.client.get("machines") diff --git a/src/commands/machine.py b/src/commands/machine.py index 3f83d97..4a1a99f 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -9,6 +9,16 @@ def __init__(self): self.endpoint = MachineAPI() def create(self, machine_name, machine_type): + """ + create create a machine with name and type. + + :param machine_name: descriptor of machine + :type machine_name: str + :param machine_type: GPU: [g4dn.xlarge] + FPGA: [f1.2xlarge, f1.4xlarge, f1.16xlarge] + CPU: [t2.micro, m5.xlarge, m5.2xlarge] + :type machine_type: TBD + """ click.echo("\nAttempting to create machine...\n") result = self.endpoint.create_machine(machine_name, machine_type) if result["success"]: From 476f69cd7abf720d2204065b753bcb894e509e86 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Tue, 14 Jan 2025 13:28:55 -0330 Subject: [PATCH 02/12] Adds other GPU/CPU endpoints to machine API interface --- src/api/machine_api.py | 50 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/src/api/machine_api.py b/src/api/machine_api.py index dcdb7f7..ba787c5 100644 --- a/src/api/machine_api.py +++ b/src/api/machine_api.py @@ -7,6 +7,8 @@ class MachineAPI: def __init__(self): self.client = APIClient() + """FPGA REQUESTS""" + @handle_api_errors def create_fpga_machine( self, machine_name: str, machine_type: str @@ -18,6 +20,8 @@ def create_fpga_machine( response = self.client.post("machine/fpga", data=data) return response + """GPU REQUESTS""" + @handle_api_errors def create_gpu_machine( self, machine_name: str, machine_type: str @@ -29,6 +33,30 @@ def create_gpu_machine( response = self.client.post("machine/gpu", data=data) return response + @handle_api_errors + def pull_gpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: + data = { + "machine_name": machine_id, + "machine_type": model_name, + } + response = self.client.post("machine/gpu/pull_model", data=data) + return response + + @handle_api_errors + def delete_gpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: + data = { + "machine_name": machine_id, + "machine_type": model_name, + } + response = self.client.delete("machine/gpu/model", data=data) + return response + + @handle_api_errors + def get_gpu_inference_url(self, machine_id: str) -> Dict[str, str]: + return self.client.get(f"machine/gpu/{machine_id}/inference_url") + + """CPU REQUESTS""" + @handle_api_errors def create_cpu_machine( self, machine_name: str, machine_type: str @@ -40,6 +68,28 @@ def create_cpu_machine( response = self.client.post("machine/cpu", data=data) return response + @handle_api_errors + def pull_cpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: + data = { + "machine_name": machine_id, + "machine_type": model_name, + } + response = self.client.post("machine/cpu/pull_model", data=data) + return response + + @handle_api_errors + def delete_cpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: + data = { + "machine_name": machine_id, + "machine_type": model_name, + } + response = self.client.delete("machine/cpu/model", data=data) + return response + + @handle_api_errors + def get_cpu_inference_url(self, machine_id: str) -> Dict[str, str]: + return self.client.get(f"machine/cpu/{machine_id}/inference_url") + @handle_api_errors def list_user_machines(self): return self.client.get("machines") From 569b58c8cc07a42b794817af0e9035858ccbbfc9 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 10:49:48 -0330 Subject: [PATCH 03/12] Move common functions to base class --- src/utils/formatters/help_formatter_base.py | 50 +++++++++++++------- src/utils/formatters/quack_formatter.py | 26 ---------- src/utils/formatters/subcommand_formatter.py | 13 ----- 3 files changed, 33 insertions(+), 56 deletions(-) diff --git a/src/utils/formatters/help_formatter_base.py b/src/utils/formatters/help_formatter_base.py index 00f6338..5e6f70a 100644 --- a/src/utils/formatters/help_formatter_base.py +++ b/src/utils/formatters/help_formatter_base.py @@ -21,17 +21,12 @@ def format_header(self, console: rConsole) -> None: def format_usage(self, console: rConsole, usage_text: str) -> None: pass - @abstractmethod - def format_options(self, console: rConsole, options: list) -> None: - pass - - @abstractmethod - def format_commands(self, console: rConsole, commands: dict, title: str) -> None: - pass - - @abstractmethod - def render(self, console: rConsole) -> None: - pass + def add_section_header(self, title: str): + title = f"━━━ {title.upper()}" + right_sep = "━" * (self.width - len(title) - 3) + self.main_table.add_row( + rAlign(f"[bold cyan]{title} {right_sep}[/]"), style="on black" + ) def create_main_table(self): self.main_table = rTable( @@ -45,9 +40,30 @@ def create_main_table(self): ) self.main_table.add_column("Content", style="bright_white") - def add_section_header(self, title: str): - title = f"━━━ {title.upper()}" - right_sep = "━" * (self.width - len(title) - 3) - self.main_table.add_row( - rAlign(f"[bold cyan]{title} {right_sep}[/]"), style="on black" - ) + def format_commands(self, commands: dict, title: str) -> None: + self.add_section_header(title) + if title == "COMMANDS": + prefix = "❯" + print_newline = True + else: + prefix = "▶" + print_newline = False + + for name, cmd in sorted(commands.items()): + help = cmd.help.split("\n")[0] + + self.main_table.add_row( + f"{self.item_padding}{prefix} [green]{name:<13}[/] { help or ''}" + ) + if print_newline: + self.main_table.add_row() + + def format_options(self, options: list) -> None: + self.add_section_header("OPTIONS") + for opt, desc in options: + self.main_table.add_row(f"{self.item_padding}[magenta]{opt:<15}[/] {desc}") + self.main_table.add_row() + + def render(self) -> None: + if self.main_table: + self.console.print(self.main_table) diff --git a/src/utils/formatters/quack_formatter.py b/src/utils/formatters/quack_formatter.py index e5ddc3a..1a7e583 100644 --- a/src/utils/formatters/quack_formatter.py +++ b/src/utils/formatters/quack_formatter.py @@ -34,29 +34,3 @@ def format_header(self, title: str, help: str) -> None: def format_usage(self, usage_text: str) -> None: self.add_section_header("USAGE") self.main_table.add_row(f"{self.item_padding}{usage_text}\n") - - def format_options(self, options: list) -> None: - self.add_section_header("OPTIONS") - for opt, desc in options: - self.main_table.add_row(f"{self.item_padding}[magenta]{opt:<15}[/] {desc}") - self.main_table.add_row() - - def format_commands(self, commands: dict, title: str) -> None: - self.add_section_header(title) - if title == "COMMANDS": - prefix = "❯" - print_newline = True - else: - prefix = "▶" - print_newline = False - - for name, cmd in sorted(commands.items()): - self.main_table.add_row( - f"{self.item_padding}{prefix} [green]{name:<13}[/] {cmd.help or ''}" - ) - if print_newline: - self.main_table.add_row() - - def render(self) -> None: - if self.main_table: - self.console.print(self.main_table) diff --git a/src/utils/formatters/subcommand_formatter.py b/src/utils/formatters/subcommand_formatter.py index 1ba2107..4717350 100644 --- a/src/utils/formatters/subcommand_formatter.py +++ b/src/utils/formatters/subcommand_formatter.py @@ -30,16 +30,3 @@ def format_options(self, options: list) -> None: for opt, desc in options: self.main_table.add_row(f"{self.item_padding}[magenta]{opt:<15}[/] {desc}") self.main_table.add_row() - - def format_commands(self, commands: dict, title: str) -> None: - self.add_section_header(title) - prefix = "❯" if title == "COMMANDS" else "▶" - - for name, cmd in sorted(commands.items()): - self.main_table.add_row( - f"{self.item_padding}{prefix} [green]{name:<13}[/] {cmd.help or ''}" - ) - - def render(self) -> None: - if self.main_table: - self.console.print(self.main_table) From 5db015fb625399690b17d1ce246414a15c474076 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 14:17:35 -0330 Subject: [PATCH 04/12] GPU, CPU spinup and inference works --- requirements.txt | 1 + src/api/machine_api.py | 16 +-- src/commands/machine.py | 180 +++++++++++++++++++++++++++++---- src/models/default_models.json | 2 +- 4 files changed, 172 insertions(+), 27 deletions(-) diff --git a/requirements.txt b/requirements.txt index f251454..be8439b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,3 +26,4 @@ zipp==3.20.2 pre_commit==4.0.1 rich==13.9.4 pyfiglet==1.0.2 +click-spinner==0.1.10 diff --git a/src/api/machine_api.py b/src/api/machine_api.py index ba787c5..39bc7c7 100644 --- a/src/api/machine_api.py +++ b/src/api/machine_api.py @@ -36,8 +36,8 @@ def create_gpu_machine( @handle_api_errors def pull_gpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: data = { - "machine_name": machine_id, - "machine_type": model_name, + "machine_id": machine_id, + "model_name": model_name, } response = self.client.post("machine/gpu/pull_model", data=data) return response @@ -45,8 +45,8 @@ def pull_gpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: @handle_api_errors def delete_gpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: data = { - "machine_name": machine_id, - "machine_type": model_name, + "machine_id": machine_id, + "model_name": model_name, } response = self.client.delete("machine/gpu/model", data=data) return response @@ -71,8 +71,8 @@ def create_cpu_machine( @handle_api_errors def pull_cpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: data = { - "machine_name": machine_id, - "machine_type": model_name, + "machine_id": machine_id, + "model_name": model_name, } response = self.client.post("machine/cpu/pull_model", data=data) return response @@ -80,8 +80,8 @@ def pull_cpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: @handle_api_errors def delete_cpu_model(self, machine_id: str, model_name: str) -> Dict[str, str]: data = { - "machine_name": machine_id, - "machine_type": model_name, + "machine_id": machine_id, + "model_name": model_name, } response = self.client.delete("machine/cpu/model", data=data) return response diff --git a/src/commands/machine.py b/src/commands/machine.py index 4a1a99f..b5b2ac6 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -1,4 +1,7 @@ +import json + import click +import click_spinner from src.api.machine_api import MachineAPI from src.utils.groups.subcommand_group import SubCommandGroup @@ -8,7 +11,7 @@ class MachineCommands: def __init__(self): self.endpoint = MachineAPI() - def create(self, machine_name, machine_type): + def create(self, hardware_type, machine_name, machine_type): """ create create a machine with name and type. @@ -19,15 +22,94 @@ def create(self, machine_name, machine_type): CPU: [t2.micro, m5.xlarge, m5.2xlarge] :type machine_type: TBD """ - click.echo("\nAttempting to create machine...\n") - result = self.endpoint.create_machine(machine_name, machine_type) + with click_spinner.spinner(): + match hardware_type: + case "GPU": + result = self.endpoint.create_gpu_machine( + machine_name, machine_type + ) + case "FPGA": + result = self.endpoint.create_fpga_machine( + machine_name, machine_type + ) + case "CPU": + result = self.endpoint.create_cpu_machine( + machine_name, machine_type + ) + case _: + click.echo("Invalid hardware type.") + return -1 + if result["success"]: - click.echo(f"Machine created successfully. Details: {result['data']}") + click.echo("Machine created successfully. Details:") + click.echo(json.dumps(result["data"], indent=2)) else: click.echo( "Failed to create machine. Check machine name and type and try again." ) + def pull_model(self, hardware_type, machine_id, model_name): + click.echo(f"Pulling model {model_name} for machine {machine_id}") + with click_spinner.spinner(): + match hardware_type: + case "GPU": + result = self.endpoint.pull_gpu_model(machine_id, model_name) + case "FPGA": + click.secho("FPGA pull model not implemented.", fg="red") + return + case "CPU": + result = self.endpoint.pull_cpu_model(machine_id, model_name) + case _: + click.echo("Invalid hardware type.") + return -1 + + if result["success"]: + click.echo(f"Model {model_name} pulled successfully.") + else: + click.echo( + "Failed to pull model. Check machine ID and model name and try again." + ) + + def delete_machine_model(self, hardware_type, machine_id, model_name): + click.echo(f"Deleting model {model_name} for machine {machine_id}") + with click_spinner.spinner(): + match hardware_type: + case "GPU": + result = self.endpoint.delete_gpu_model(machine_id, model_name) + case "FPGA": + click.secho("FPGA delete model not implemented.", fg="red") + return + case "CPU": + result = self.endpoint.delete_cpu_model(machine_id, model_name) + case _: + click.echo("Invalid hardware type.") + return -1 + if result["success"]: + click.echo(f"Model {model_name} deleted successfully.") + else: + click.echo( + "Failed to delete model. Check machine ID and model name and try again." + ) + + def get_inference_url(self, hardware_type, machine_id): + with click_spinner.spinner(): + match hardware_type: + case "GPU": + result = self.endpoint.get_gpu_inference_url(machine_id) + case "FPGA": + click.secho("FPGA inference URL not implemented.", fg="red") + return + case "CPU": + result = self.endpoint.get_cpu_inference_url(machine_id) + case _: + click.echo("Invalid hardware type.") + return -1 + + if result["success"]: + click.echo(f"Inference URL: {result['data']['inference_url']}") + else: + click.echo("Failed to get inference URL. Check machine ID and try again.") + def list(self): result = self.endpoint.list_user_machines() if result["success"]: @@ -36,7 +118,7 @@ def list(self): return click.echo("Machines:") for machine in result["data"]: - click.echo(machine) + click.echo(json.dumps(machine, indent=2)) return result else: click.echo("Failed to retrieve list of machines.") @@ -69,7 +151,7 @@ def get_details(self, machine_id): result = self.endpoint.get_machine(machine_id) if result["success"]: click.echo(f"Machine {machine_id} details:") - click.echo(result["data"]) + click.echo(json.dumps(result["data"], indent=2)) return result["data"] else: click.echo("Failed to get machine. Check machine ID and try again.") @@ -84,19 +166,81 @@ def machine(ctx): @machine.command() @click.pass_context +@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) @click.option("--machine-name", "-n", required=True, prompt="Enter machine name") -@click.option( - "--machine-type", - "-t", - required=True, - prompt="Choose machine type", - type=click.Choice(["t2.micro", "f1.2xlarge", "f1.4xlarge", "f1.16xlarge"]), - default="f1.2xlarge", - show_default=True, -) -def create(ctx, machine_name, machine_type): - """Create a machine with name and type.""" - ctx.obj.create(machine_name=machine_name, machine_type=machine_type) +@click.option("--machine-type", "-t", required=True, prompt="Choose machine type") +def create(ctx, hardware_type, machine_name, machine_type): + """Create a machine with name and type. + + Valid machine types: + + - CPU: \t[t2.micro, m5.xlarge, m5.2xlarge] + + - GPU: \t[g4dn.xlarge] + + - FPGA: \t[f1.2xlarge, f1.4xlarge, f1.16xlarge] + """ + valid_machine_types = { + "GPU": ["g4dn.xlarge"], + "FPGA": ["f1.2xlarge", "f1.4xlarge", "f1.16xlarge"], + "CPU": ["t2.micro", "m5.xlarge", "m5.2xlarge"], + } + + if machine_type not in valid_machine_types[hardware_type]: + click.secho( + f"Invalid machine type for {hardware_type}. See `quack machine create --help` for options.", + fg="red", + ) + return + + ctx.obj.create( + hardware_type=hardware_type, + machine_name=machine_name, + machine_type=machine_type, + ) + + +@machine.command() +@click.pass_context +@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") +@click.option("--model-name", "-m", required=True, prompt="Enter model name") +def pull_model(ctx, hardware_type, machine_id, model_name): + """Pull model for machine to machine with machine ID. + + See `quack model list` for available models. + """ + with open("src/models/default_models.json", "r") as f: + models = json.loads(f.read()) + for item in models: + if model_name == item["name"] and item["available"][f"{hardware_type}"]: + ctx.obj.pull_model(hardware_type, machine_id, model_name) + return + else: + click.secho( + "Invalid parameters. See `quack machine pull-model --help` for options.", + fg="red", + ) + return + + +@machine.command() +@click.pass_context +@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") +@click.option("--model-name", "-m", required=True, prompt="Enter model name") +def delete_model(ctx, hardware_type, machine_id, model_name): + """Delete model for machine with machine ID.""" + ctx.obj.delete_machine_model(hardware_type, machine_id, model_name) + + +@machine.command() +@click.pass_context +@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") +def infer_url(ctx, hardware_type, machine_id): + """Get inference URL for machine with machine ID.""" + ctx.obj.get_inference_url(hardware_type, machine_id) @machine.command() diff --git a/src/models/default_models.json b/src/models/default_models.json index 03b07e7..3db1aa7 100644 --- a/src/models/default_models.json +++ b/src/models/default_models.json @@ -18,7 +18,7 @@ "description": "Meta's Llama 2" }, { - "name": "smollm:135b", + "name": "smollm:135m", "available": { "FPGA": false, "GPU": true, From cfcee31949163859d73d585811a4e93e6e1c4cb2 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 17:51:21 -0330 Subject: [PATCH 05/12] Tested FPGA creation and inference --- src/api/machine_api.py | 4 ++++ src/commands/machine.py | 7 +++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/api/machine_api.py b/src/api/machine_api.py index 39bc7c7..a8e51c8 100644 --- a/src/api/machine_api.py +++ b/src/api/machine_api.py @@ -20,6 +20,10 @@ def create_fpga_machine( response = self.client.post("machine/fpga", data=data) return response + @handle_api_errors + def get_fpga_inference_url(self, machine_id: str) -> Dict[str, str]: + return self.client.get(f"machine/fpga/{machine_id}/inference_url") + """GPU REQUESTS""" @handle_api_errors diff --git a/src/commands/machine.py b/src/commands/machine.py index b5b2ac6..0e947fa 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -55,7 +55,7 @@ def pull_model(self, hardware_type, machine_id, model_name): case "GPU": result = self.endpoint.pull_gpu_model(machine_id, model_name) case "FPGA": - click.secho("FPGA pull model not implemented.", fg="red") + click.secho("FPGA pull model not supported yet.", fg="yellow") return case "CPU": result = self.endpoint.pull_cpu_model(machine_id, model_name) @@ -77,7 +77,7 @@ def delete_machine_model(self, hardware_type, machine_id, model_name): case "GPU": result = self.endpoint.delete_gpu_model(machine_id, model_name) case "FPGA": - click.secho("FPGA delete model not implemented.", fg="red") + click.secho("FPGA delete model not supported yet.", fg="yellow") return case "CPU": result = self.endpoint.delete_cpu_model(machine_id, model_name) @@ -97,8 +97,7 @@ def get_inference_url(self, hardware_type, machine_id): case "GPU": result = self.endpoint.get_gpu_inference_url(machine_id) case "FPGA": - click.secho("FPGA inference URL not implemented.", fg="red") - return + result = self.endpoint.get_fpga_inference_url(machine_id) case "CPU": result = self.endpoint.get_cpu_inference_url(machine_id) case _: From 12dcb21bdfb7d5e6dbfd28561467636338d4da95 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 18:36:23 -0330 Subject: [PATCH 06/12] test machine api updated, added pytest-cov to requirements.txt --- requirements.txt | 1 + test/api/test_machine_api.py | 89 ++++++++++++++++++++++++++++++++++-- 2 files changed, 86 insertions(+), 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index be8439b..96bb634 100644 --- a/requirements.txt +++ b/requirements.txt @@ -27,3 +27,4 @@ pre_commit==4.0.1 rich==13.9.4 pyfiglet==1.0.2 click-spinner==0.1.10 +pytest-cov==6.0.0 diff --git a/test/api/test_machine_api.py b/test/api/test_machine_api.py index 299bf1a..c36748d 100644 --- a/test/api/test_machine_api.py +++ b/test/api/test_machine_api.py @@ -10,15 +10,96 @@ def setUp(self): self.api = MachineAPI() @patch.object(APIClient, "post") - def test_create_machine(self, mock_post): + def test_create_fpga_machine(self, mock_post): mock_post.return_value = {"status": "success"} - response = self.api.create_machine("test_machine", "type_a") - self.assertEqual(response, {"success": True, "data": {"status": "success"}}) + response = self.api.create_fpga_machine("test_fpga", "type_a") + self.assertEqual(response["data"], {"status": "success"}) mock_post.assert_called_with( "machine/fpga", - data={"machine_name": "test_machine", "machine_type": "type_a"}, + data={"machine_name": "test_fpga", "machine_type": "type_a"}, ) + @patch.object(APIClient, "get") + def test_get_fpga_inference_url(self, mock_get): + mock_get.return_value = {"inference_url": "http://example.com"} + response = self.api.get_fpga_inference_url("123") + self.assertEqual(response["data"], {"inference_url": "http://example.com"}) + mock_get.assert_called_with("machine/fpga/123/inference_url") + + @patch.object(APIClient, "post") + def test_create_gpu_machine(self, mock_post): + mock_post.return_value = {"status": "success"} + response = self.api.create_gpu_machine("test_gpu", "type_b") + self.assertEqual(response["data"], {"status": "success"}) + mock_post.assert_called_with( + "machine/gpu", + data={"machine_name": "test_gpu", "machine_type": "type_b"}, + ) + + @patch.object(APIClient, "post") + def test_pull_gpu_model(self, mock_post): + mock_post.return_value = {"status": "pulled"} + response = self.api.pull_gpu_model("123", "model_a") + self.assertEqual(response["data"], {"status": "pulled"}) + mock_post.assert_called_with( + "machine/gpu/pull_model", + data={"machine_id": "123", "model_name": "model_a"}, + ) + + @patch.object(APIClient, "delete") + def test_delete_gpu_model(self, mock_delete): + mock_delete.return_value = {"status": "deleted"} + response = self.api.delete_gpu_model("123", "model_a") + self.assertEqual(response["data"], {"status": "deleted"}) + mock_delete.assert_called_with( + "machine/gpu/model", + data={"machine_id": "123", "model_name": "model_a"}, + ) + + @patch.object(APIClient, "get") + def test_get_gpu_inference_url(self, mock_get): + mock_get.return_value = {"inference_url": "http://example.com"} + response = self.api.get_gpu_inference_url("123") + self.assertEqual(response["data"], {"inference_url": "http://example.com"}) + mock_get.assert_called_with("machine/gpu/123/inference_url") + + @patch.object(APIClient, "post") + def test_create_cpu_machine(self, mock_post): + mock_post.return_value = {"status": "success"} + response = self.api.create_cpu_machine("test_cpu", "type_c") + self.assertEqual(response["data"], {"status": "success"}) + mock_post.assert_called_with( + "machine/cpu", + data={"machine_name": "test_cpu", "machine_type": "type_c"}, + ) + + @patch.object(APIClient, "post") + def test_pull_cpu_model(self, mock_post): + mock_post.return_value = {"status": "pulled"} + response = self.api.pull_cpu_model("123", "model_b") + self.assertEqual(response["data"], {"status": "pulled"}) + mock_post.assert_called_with( + "machine/cpu/pull_model", + data={"machine_id": "123", "model_name": "model_b"}, + ) + + @patch.object(APIClient, "delete") + def test_delete_cpu_model(self, mock_delete): + mock_delete.return_value = {"status": "deleted"} + response = self.api.delete_cpu_model("123", "model_b") + self.assertEqual(response["data"], {"status": "deleted"}) + mock_delete.assert_called_with( + "machine/cpu/model", + data={"machine_id": "123", "model_name": "model_b"}, + ) + + @patch.object(APIClient, "get") + def test_get_cpu_inference_url(self, mock_get): + mock_get.return_value = {"inference_url": "http://example.com"} + response = self.api.get_cpu_inference_url("123") + self.assertEqual(response["data"], {"inference_url": "http://example.com"}) + mock_get.assert_called_with("machine/cpu/123/inference_url") + @patch.object(APIClient, "get") def test_list_user_machines(self, mock_get): mock_get.return_value = {"machines": []} From 43f95d6704d9ddcd4346dc9082965e4094fe8fb7 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 19:09:32 -0330 Subject: [PATCH 07/12] update user_auth tests --- src/commands/user_auth.py | 1 + test/commands/test_user_auth.py | 121 +++++++++++++++++--------------- 2 files changed, 67 insertions(+), 55 deletions(-) diff --git a/src/commands/user_auth.py b/src/commands/user_auth.py index f8cb449..ceb73bc 100644 --- a/src/commands/user_auth.py +++ b/src/commands/user_auth.py @@ -1,4 +1,5 @@ import subprocess + import click from src.api.api_client import APIClient diff --git a/test/commands/test_user_auth.py b/test/commands/test_user_auth.py index 9962385..65efb96 100644 --- a/test/commands/test_user_auth.py +++ b/test/commands/test_user_auth.py @@ -10,17 +10,26 @@ def setUp(self): self.runner = CliRunner() # Create and setup mock_auth_api - self.patcher = patch("src.api.auth_api.AuthAPI") - self.mock_auth_api = self.patcher.start() + self.auth_patcher = patch("src.api.auth_api.AuthAPI") + self.mock_auth_api = self.auth_patcher.start() self.mock_instance = self.mock_auth_api.return_value self.endpoint = patch( "src.commands.user_auth.endpoint", self.mock_instance ).start() + # Create and setup mock_user_api + self.user_patcher = patch("src.api.user_api.UserAPI") + self.mock_user_api = self.user_patcher.start() + self.mock_user_instance = self.mock_user_api.return_value + + self.user_endpoint = patch( + "src.commands.user_auth.user_endpoint", self.mock_user_instance + ).start() + def tearDown(self): """Clean up after each test method.""" - self.patcher.stop() + self.auth_patcher.stop() self.endpoint.stop() def test_login_success(self): @@ -80,56 +89,58 @@ def test_logout_failure(self): def test_register_success(self): """Test successful registration.""" + self.mock_user_instance.register.return_value = { + "success": True, + "data": { + "user_name": "test_user", + "email": "test@user.ca", + "user_id": "123", + }, + } + self.mock_instance.login.return_value = { + "success": True, + "data": "Login successful", + } + + with patch("click.echo") as mock_print: + result = self.runner.invoke( + register, + [ + "--username", + "test_user", + "--email", + "test@user.ca", + "--password", + "password", + ], + ) + mock_print.assert_any_call( + "Welcome to Duckington Labs, you've successfully registered as:" + ) + mock_print.assert_any_call(" - Username: test_user") + mock_print.assert_any_call(" - Email: test@user.ca") + mock_print.assert_any_call(" - User ID: 123") + mock_print.assert_any_call("Successfully logged in. Login successful") + self.assertEqual(result.exit_code, 0) - def test_register_success(self): - """Test successful registration.""" - self.mock_instance.register.return_value = { - "success": True, - "data": { - "user_name": "testuser", - "email": "test@test.com", - "user_id": "12345", - }, - } - - with patch("click.echo") as mock_print: - result = self.runner.invoke( - register, - [ - "--username", - "testuser", - "--email", - "test@test.com", - "--password", - "password", - ], - ) - mock_print.assert_any_call( - "Welcome to Duckington Labs, you've successfully registered as:" - ) - mock_print.assert_any_call(" - Username: testuser") - mock_print.assert_any_call(" - Email: test@test.com") - mock_print.assert_any_call(" - User ID: 12345") - self.assertEqual(result.exit_code, 0) - - def test_register_failure(self): - """Test failed registration.""" - self.mock_instance.register.return_value = { - "success": False, - "message": "Registration failed", - } - - with patch("click.echo") as mock_print: - result = self.runner.invoke( - register, - [ - "--username", - "testuser", - "--email", - "test@test.com", - "--password", - "password", - ], - ) - mock_print.assert_called_with("Login failed. Registration failed") - self.assertEqual(result.exit_code, 0) + def test_register_failure(self): + """Test failed registration.""" + self.mock_user_instance.register.return_value = { + "success": False, + "message": "Registration failed", + } + + with patch("click.echo") as mock_print: + result = self.runner.invoke( + register, + [ + "--username", + "test_user", + "--email", + "test@user.ca", + "--password", + "password", + ], + ) + mock_print.assert_called_with("Login failed. Registration failed") + self.assertEqual(result.exit_code, 0) From b19702e9449b0655241b2e4e8a83c86074700922 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 19:39:52 -0330 Subject: [PATCH 08/12] update machine command tests --- src/commands/machine.py | 6 +- test/commands/test_machine.py | 499 ++++++++++++++++++++++++++++++---- 2 files changed, 448 insertions(+), 57 deletions(-) diff --git a/src/commands/machine.py b/src/commands/machine.py index 0e947fa..788b853 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -192,11 +192,7 @@ def create(ctx, hardware_type, machine_name, machine_type): ) return - ctx.obj.create( - hardware_type=hardware_type, - machine_name=machine_name, - machine_type=machine_type, - ) + ctx.obj.create(hardware_type, machine_name, machine_type) @machine.command() diff --git a/test/commands/test_machine.py b/test/commands/test_machine.py index 772acd0..6031b37 100644 --- a/test/commands/test_machine.py +++ b/test/commands/test_machine.py @@ -1,10 +1,13 @@ +import json import unittest from unittest.mock import MagicMock, patch + from click.testing import CliRunner + from src.commands.machine import MachineCommands, machine -class TestMachineCommands(unittest.TestCase): +class TestNewMachineCommands(unittest.TestCase): def setUp(self): """Set up test fixtures before each test method.""" self.runner = CliRunner() @@ -15,7 +18,15 @@ def setUp(self): self.mock_instance = self.mock_machine_api.return_value # Set up default successful responses - self.mock_instance.create_machine.return_value = { + self.mock_instance.create_gpu_machine.return_value = { + "success": True, + "data": { + "machine_id": "test123", + "machine_name": "test-machine", + "machine_type": "g4dn.xlarge", + }, + } + self.mock_instance.create_fpga_machine.return_value = { "success": True, "data": { "machine_id": "test123", @@ -23,21 +34,37 @@ def setUp(self): "machine_type": "f1.2xlarge", }, } - self.mock_instance.list_user_machines.return_value = { + self.mock_instance.create_cpu_machine.return_value = { "success": True, - "data": [], + "data": { + "machine_id": "test123", + "machine_name": "test-machine", + "machine_type": "t2.micro", + }, } - self.mock_instance.start_machine.return_value = { + self.mock_instance.pull_gpu_model.return_value = { "success": True, - "data": {"message": "Machine started successfully"}, + "data": {"message": "Model pulled successfully"}, } - self.mock_instance.stop_machine.return_value = { + self.mock_instance.pull_cpu_model.return_value = { "success": True, - "data": {"message": "Machine stopped successfully"}, + "data": {"message": "Model pulled successfully"}, } - self.mock_instance.terminate_machine.return_value = { + self.mock_instance.delete_gpu_model.return_value = { "success": True, - "data": {"message": "Machine terminated successfully"}, + "data": {"message": "Model deleted successfully"}, + } + self.mock_instance.delete_cpu_model.return_value = { + "success": True, + "data": {"message": "Model deleted successfully"}, + } + self.mock_instance.get_gpu_inference_url.return_value = { + "success": True, + "data": {"inference_url": "http://gpu-inference-url"}, + } + self.mock_instance.get_cpu_inference_url.return_value = { + "success": True, + "data": {"inference_url": "http://cpu-inference-url"}, } # Setup machine_commands @@ -49,103 +76,471 @@ def tearDown(self): """Clean up after each test method.""" self.patcher.stop() - def test_create_machine_success(self): + def test_create_gpu_machine_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.create("GPU", "test-machine", "g4dn.xlarge") + mock_print.assert_any_call("Machine created successfully. Details:") + mock_print.assert_any_call( + json.dumps( + self.mock_instance.create_gpu_machine.return_value["data"], indent=2 + ) + ) + + def test_create_fpga_machine_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.create("FPGA", "test-machine", "f1.2xlarge") + mock_print.assert_any_call("Machine created successfully. Details:") + mock_print.assert_any_call( + json.dumps( + self.mock_instance.create_fpga_machine.return_value["data"], + indent=2, + ) + ) + + def test_create_cpu_machine_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.create("CPU", "test-machine", "t2.micro") + mock_print.assert_any_call("Machine created successfully. Details:") + mock_print.assert_any_call( + json.dumps( + self.mock_instance.create_cpu_machine.return_value["data"], indent=2 + ) + ) + + def test_create_gpu_machine_failure(self): + self.mock_instance.create_gpu_machine.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.create("test-machine", "f1.2xlarge") + self.machine_commands.create("GPU", "test-machine", "g4dn.xlarge") mock_print.assert_called_with( - f"Machine created successfully. Details: {self.mock_instance.create_machine.return_value['data']}" + "Failed to create machine. Check machine name and type and try again." ) - def test_create_machine_failure(self): - self.mock_instance.create_machine.return_value = { - "success": False, - "message": "Failed to create machine", - } + def test_create_fpga_machine_failure(self): + self.mock_instance.create_fpga_machine.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.create("FPGA", "test-machine", "f1.2xlarge") + mock_print.assert_called_with( + "Failed to create machine. Check machine name and type and try again." + ) + def test_create_cpu_machine_failure(self): + self.mock_instance.create_cpu_machine.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.create("test-machine", "f1.2xlarge") + self.machine_commands.create("CPU", "test-machine", "t2.micro") mock_print.assert_called_with( "Failed to create machine. Check machine name and type and try again." ) - def test_list_machines_with_data(self): - mock_machines = [ - {"machine_id": "test123", "machine_name": "test-machine-1"}, - {"machine_id": "test456", "machine_name": "test-machine-2"}, - ] - self.mock_instance.list_user_machines.return_value = { + def test_create_machine_invalid_hardware_type(self): + with patch("click.echo") as mock_print: + self.machine_commands.create("INVALID", "test-machine", "t2.micro") + mock_print.assert_called_with("Invalid hardware type.") + + def test_pull_gpu_model_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.pull_model("GPU", "test123", "test-model") + mock_print.assert_called_with("Model test-model pulled successfully.") + + def test_pull_cpu_model_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.pull_model("CPU", "test123", "test-model") + mock_print.assert_called_with("Model test-model pulled successfully.") + + def test_pull_fpga_model_not_supported(self): + with patch("click.secho") as mock_print: + self.machine_commands.pull_model("FPGA", "test123", "test-model") + mock_print.assert_called_with( + "FPGA pull model not supported yet.", fg="yellow" + ) + + def test_pull_gpu_model_failure(self): + self.mock_instance.pull_gpu_model.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.pull_model("GPU", "test123", "test-model") + mock_print.assert_called_with( + "Failed to pull model. Check machine ID and model name and try again." + ) + + def test_pull_cpu_model_failure(self): + self.mock_instance.pull_cpu_model.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.pull_model("CPU", "test123", "test-model") + mock_print.assert_called_with( + "Failed to pull model. Check machine ID and model name and try again." + ) + + def test_pull_model_invalid_hardware_type(self): + with patch("click.echo") as mock_print: + self.machine_commands.pull_model("INVALID", "test123", "test-model") + mock_print.assert_called_with("Invalid hardware type.") + + def test_delete_gpu_model_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.delete_machine_model("GPU", "test123", "test-model") + mock_print.assert_called_with("Model test-model deleted successfully.") + + def test_delete_cpu_model_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.delete_machine_model("CPU", "test123", "test-model") + mock_print.assert_called_with("Model test-model deleted successfully.") + + def test_delete_fpga_model_not_supported(self): + with patch("click.secho") as mock_print: + self.machine_commands.delete_machine_model("FPGA", "test123", "test-model") + mock_print.assert_called_with( + "FPGA delete model not supported yet.", fg="yellow" + ) + + def test_delete_gpu_model_failure(self): + self.mock_instance.delete_gpu_model.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.delete_machine_model("GPU", "test123", "test-model") + mock_print.assert_called_with( + "Failed to delete model. Check machine ID and model name and try again." + ) + + def test_delete_cpu_model_failure(self): + self.mock_instance.delete_cpu_model.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.delete_machine_model("CPU", "test123", "test-model") + mock_print.assert_called_with( + "Failed to delete model. Check machine ID and model name and try again." + ) + + def test_delete_model_invalid_hardware_type(self): + with patch("click.echo") as mock_print: + self.machine_commands.delete_machine_model( + "INVALID", "test123", "test-model" + ) + mock_print.assert_called_with("Invalid hardware type.") + + def test_get_gpu_inference_url_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("GPU", "test123") + mock_print.assert_called_with("Inference URL: http://gpu-inference-url") + + def test_get_cpu_inference_url_success(self): + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("CPU", "test123") + mock_print.assert_called_with("Inference URL: http://cpu-inference-url") + + def test_get_fpga_inference_url_success(self): + self.mock_instance.get_fpga_inference_url.return_value = { "success": True, - "data": mock_machines, + "data": {"inference_url": "http://fpga-inference-url"}, } + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("FPGA", "test123") + mock_print.assert_called_with("Inference URL: http://fpga-inference-url") + + def test_get_inference_url_invalid_hardware_type(self): + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("INVALID", "test123") + mock_print.assert_called_with("Invalid hardware type.") + def test_get_cpu_inference_url_failure(self): + self.mock_instance.get_cpu_inference_url.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("CPU", "test123") + mock_print.assert_called_with( + "Failed to get inference URL. Check machine ID and try again." + ) + + def test_get_fpga_inference_url_failure(self): + self.mock_instance.get_fpga_inference_url.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("FPGA", "test123") + mock_print.assert_called_with( + "Failed to get inference URL. Check machine ID and try again." + ) + + def test_get_gpu_inference_url_failure(self): + self.mock_instance.get_gpu_inference_url.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.get_inference_url("GPU", "test123") + mock_print.assert_called_with( + "Failed to get inference URL. Check machine ID and try again." + ) + + def test_list_machines_success(self): + self.mock_instance.list_user_machines.return_value = { + "success": True, + "data": [ + { + "machine_id": "test123", + "machine_name": "test-machine", + "machine_type": "g4dn.xlarge", + } + ], + } with patch("click.echo") as mock_print: self.machine_commands.list() - self.assertEqual(mock_print.call_count, 3) # Header + 2 machines + mock_print.assert_any_call("Machines:") + mock_print.assert_any_call( + json.dumps( + self.mock_instance.list_user_machines.return_value["data"][0], + indent=2, + ) + ) - def test_list_machines_empty(self): + def test_list_machines_no_machines(self): + self.mock_instance.list_user_machines.return_value = { + "success": True, + "data": [], + } with patch("click.echo") as mock_print: self.machine_commands.list() mock_print.assert_called_with("No machines to list.") - def test_start_machine_success(self): + def test_list_machines_failure(self): + self.mock_instance.list_user_machines.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.start("test123") - mock_print.assert_called_with("Machine started successfully") + self.machine_commands.list() + mock_print.assert_called_with("Failed to retrieve list of machines.") def test_stop_machine_success(self): + self.mock_instance.stop_machine.return_value = { + "success": True, + "data": {"message": "Machine stopped successfully"}, + } with patch("click.echo") as mock_print: self.machine_commands.stop("test123") mock_print.assert_called_with("Machine stopped successfully") + def test_stop_machine_failure(self): + self.mock_instance.stop_machine.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.stop("test123") + mock_print.assert_called_with( + "Failed to stop machine. Check machine ID and try again." + ) + + def test_start_machine_success(self): + self.mock_instance.start_machine.return_value = { + "success": True, + "data": {"message": "Machine started successfully"}, + } + with patch("click.echo") as mock_print: + self.machine_commands.start("test123") + mock_print.assert_called_with("Machine started successfully") + + def test_start_machine_failure(self): + self.mock_instance.start_machine.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.start("test123") + mock_print.assert_called_with( + "Failed to start machine. Check machine ID and try again." + ) + def test_terminate_machine_success(self): + self.mock_instance.terminate_machine.return_value = { + "success": True, + "data": {"message": "Machine terminated successfully"}, + } with patch("click.echo") as mock_print: self.machine_commands.terminate("test123") mock_print.assert_called_with("Machine terminated successfully") - def test_get_machine_details_success(self): + def test_terminate_machine_failure(self): + self.mock_instance.terminate_machine.return_value = {"success": False} + with patch("click.echo") as mock_print: + self.machine_commands.terminate("test123") + mock_print.assert_called_with( + "Failed to terminate machine. Check machine ID and try again." + ) + + def test_get_details_success(self): self.mock_instance.get_machine.return_value = { "success": True, "data": { "machine_id": "test123", "machine_name": "test-machine", - "machine_type": "f1.2xlarge", + "machine_type": "g4dn.xlarge", }, } + with patch("click.echo") as mock_print: + self.machine_commands.get_details("test123") + mock_print.assert_any_call("Machine test123 details:") + mock_print.assert_any_call( + json.dumps( + self.mock_instance.get_machine.return_value["data"], indent=2 + ) + ) + def test_get_details_failure(self): + self.mock_instance.get_machine.return_value = {"success": False} with patch("click.echo") as mock_print: self.machine_commands.get_details("test123") - self.assertEqual(mock_print.call_count, 2) + mock_print.assert_called_with( + "Failed to get machine. Check machine ID and try again." + ) + + +class TestMachineCommandsInvoke(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + + @patch("src.commands.machine.MachineCommands.create") + def test_create_success_command(self, mock_create): + result = self.runner.invoke( + machine, + [ + "create", + "GPU", + "--machine-name", + "test-machine", + "--machine-type", + "g4dn.xlarge", + ], + ) + self.assertEqual(result.exit_code, 0) + mock_create.assert_called_once_with("GPU", "test-machine", "g4dn.xlarge") + + @patch("src.commands.machine.MachineCommands.create") + def test_create_failure_command(self, mock_create): + mock_create.return_value = -1 + result = self.runner.invoke( + machine, + [ + "create", + "GPU", + "--machine-name", + "test-machine", + "--machine-type", + "t2.micro", + ], + ) + self.assertEqual(result.exit_code, 0) + mock_create.assert_not_called() - # CLI command tests - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_create_command(self, machine_commands): + @patch("src.commands.machine.MachineCommands.pull_model") + def test_pull_model_command(self, mock_pull_model): + with patch( + "builtins.open", + unittest.mock.mock_open( + read_data='[{"name": "test-model", "available": {"GPU": true}}]' + ), + ): + result = self.runner.invoke( + machine, + [ + "pull-model", + "GPU", + "--machine-id", + "test123", + "--model-name", + "test-model", + ], + ) + self.assertEqual(result.exit_code, 0) + mock_pull_model.assert_called_once_with("GPU", "test123", "test-model") + + @patch("src.commands.machine.MachineCommands.pull_model") + def test_pull_model_failure_command(self, mock_pull_model): + with patch( + "builtins.open", + unittest.mock.mock_open( + read_data='[{"name": "test-model", "available": {"GPU": false}}]' + ), + ): + result = self.runner.invoke( + machine, + [ + "pull-model", + "GPU", + "--machine-id", + "test123", + "--model-name", + "test-model", + ], + ) + self.assertEqual(result.exit_code, 0) + mock_pull_model.assert_not_called() + + @patch("src.commands.machine.MachineCommands.delete_machine_model") + def test_delete_model_command(self, mock_delete_model): result = self.runner.invoke( machine, - ["create", "-n", "test-machine", "-t", "f1.2xlarge"], + [ + "delete-model", + "GPU", + "--machine-id", + "test123", + "--model-name", + "test-model", + ], ) self.assertEqual(result.exit_code, 0) + mock_delete_model.assert_called_once_with("GPU", "test123", "test-model") - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_list_command(self, machine_commands): + @patch("src.commands.machine.MachineCommands.get_inference_url") + def test_infer_url_command(self, mock_get_inference_url): + result = self.runner.invoke( + machine, + [ + "infer-url", + "GPU", + "--machine-id", + "test123", + ], + ) + self.assertEqual(result.exit_code, 0) + mock_get_inference_url.assert_called_once_with("GPU", "test123") + + @patch("src.commands.machine.MachineCommands.list") + def test_list_command(self, mock_list): result = self.runner.invoke(machine, ["list"]) self.assertEqual(result.exit_code, 0) + mock_list.assert_called_once() - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_start_command(self, machine_commands): - result = self.runner.invoke(machine, ["start", "--machine-id", "test123"]) + @patch("src.commands.machine.MachineCommands.stop") + def test_stop_command(self, mock_stop): + result = self.runner.invoke( + machine, + [ + "stop", + "--machine-id", + "test123", + ], + ) self.assertEqual(result.exit_code, 0) + mock_stop.assert_called_once_with("test123") - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_stop_command(self, machine_commands): - result = self.runner.invoke(machine, ["stop", "--machine-id", "test123"]) + @patch("src.commands.machine.MachineCommands.start") + def test_start_command(self, mock_start): + result = self.runner.invoke( + machine, + [ + "start", + "--machine-id", + "test123", + ], + ) self.assertEqual(result.exit_code, 0) + mock_start.assert_called_once_with("test123") - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_terminate_command(self, machine_commands): - result = self.runner.invoke(machine, ["terminate", "--machine-id", "test123"]) + @patch("src.commands.machine.MachineCommands.terminate") + def test_terminate_command(self, mock_terminate): + result = self.runner.invoke( + machine, + [ + "terminate", + "--machine-id", + "test123", + ], + ) self.assertEqual(result.exit_code, 0) + mock_terminate.assert_called_once_with("test123") - @patch("src.commands.machine.MachineCommands", autospec=True) - def test_cli_details_command(self, machine_commands): - result = self.runner.invoke(machine, ["details", "--machine-id", "test123"]) + @patch("src.commands.machine.MachineCommands.get_details") + def test_details_command(self, mock_get_details): + result = self.runner.invoke( + machine, + [ + "details", + "--machine-id", + "test123", + ], + ) self.assertEqual(result.exit_code, 0) + mock_get_details.assert_called_once_with("test123") From c94d0d922496813764542dceab141ab1812e9cfa Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 19:53:54 -0330 Subject: [PATCH 09/12] add click_spinner to all commands --- src/commands/key.py | 16 +++++++++++----- src/commands/machine.py | 15 ++++++++++----- src/commands/model_file.py | 26 +++++++++++++++++--------- src/commands/user_auth.py | 10 +++++++--- 4 files changed, 45 insertions(+), 22 deletions(-) diff --git a/src/commands/key.py b/src/commands/key.py index ba549d3..716f94e 100644 --- a/src/commands/key.py +++ b/src/commands/key.py @@ -1,4 +1,5 @@ import click +import click_spinner from src.api.api_client import APIClient from src.api.auth_api import AuthAPI @@ -12,14 +13,16 @@ def __init__(self): self.endpoint = AuthAPI(self.client) def create_api_key(self, validity): - result = self.endpoint.create_api_key(ValidityEnum[validity]) + with click_spinner.spinner(): + result = self.endpoint.create_api_key(ValidityEnum[validity]) if result["success"]: click.echo(f"API key created successfully: {result['data']}") else: click.echo(f"Failed to create API key. {result['message']}") def list_api_keys(self): - result = self.endpoint.list_api_keys() + with click_spinner.spinner(): + result = self.endpoint.list_api_keys() if result["success"]: click.echo("API Keys:") for key in result["data"]: @@ -30,19 +33,22 @@ def list_api_keys(self): click.echo(f"Failed to retrieve API keys. {result['message']}") def delete_api_key(self, token): - result = self.endpoint.delete_api_key(token) + with click_spinner.spinner(): + result = self.endpoint.delete_api_key(token) if result["success"]: click.echo(f"API key deleted successfully. {result['data']}") else: click.echo(f"Failed to delete API key. {result['message']}") def set_api_key(self, api_key): - result = self.endpoint.set_api_key(api_key) + with click_spinner.spinner(): + result = self.endpoint.set_api_key(api_key) if result: click.echo("API key set successfully.") def remove_api_key(self): - result = self.endpoint.clear_api_key() + with click_spinner.spinner(): + result = self.endpoint.clear_api_key() if result: click.echo("API key removed successfully.") else: diff --git a/src/commands/machine.py b/src/commands/machine.py index 788b853..e0ae9c5 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -110,7 +110,8 @@ def get_inference_url(self, hardware_type, machine_id): click.echo("Failed to get inference URL. Check machine ID and try again.") def list(self): - result = self.endpoint.list_user_machines() + with click_spinner.spinner(): + result = self.endpoint.list_user_machines() if result["success"]: if not result["data"]: click.echo("No machines to list.") @@ -124,7 +125,8 @@ def list(self): def stop(self, machine_id): click.echo("\nAttempting to stop machine...\n") - result = self.endpoint.stop_machine(machine_id) + with click_spinner.spinner(): + result = self.endpoint.stop_machine(machine_id) if result["success"]: click.echo(f"{result['data']['message']}") else: @@ -132,7 +134,8 @@ def stop(self, machine_id): def start(self, machine_id): click.echo("\nAttempting to start machine...\n") - result = self.endpoint.start_machine(machine_id) + with click_spinner.spinner(): + result = self.endpoint.start_machine(machine_id) if result["success"]: click.echo(f"{result['data']['message']}") else: @@ -140,14 +143,16 @@ def start(self, machine_id): def terminate(self, machine_id): click.echo("\nAttempting to terminate machine...\n") - result = self.endpoint.terminate_machine(machine_id) + with click_spinner.spinner(): + result = self.endpoint.terminate_machine(machine_id) if result["success"]: click.echo(f"{result['data']['message']}") else: click.echo("Failed to terminate machine. Check machine ID and try again.") def get_details(self, machine_id): - result = self.endpoint.get_machine(machine_id) + with click_spinner.spinner(): + result = self.endpoint.get_machine(machine_id) if result["success"]: click.echo(f"Machine {machine_id} details:") click.echo(json.dumps(result["data"], indent=2)) diff --git a/src/commands/model_file.py b/src/commands/model_file.py index 824ad27..f698c2f 100644 --- a/src/commands/model_file.py +++ b/src/commands/model_file.py @@ -2,6 +2,7 @@ import json import click +import click_spinner from src.api.api_client import APIClient from src.api.model_file_api import ModelFileAPI @@ -20,7 +21,8 @@ def upload( model_id: str | None = None, ): click.echo("\nAttempting to upload model file...\n") - result = self.endpoint.upload_model_file(model_name, model_id, file_path) + with click_spinner.spinner(): + result = self.endpoint.upload_model_file(model_name, model_id, file_path) if result["success"]: upload_date = datetime.strptime( result["data"]["upload_date"], "%Y-%m-%dT%H:%M:%S.%f%z" @@ -45,8 +47,9 @@ def list_default(self, file_path="src/models/default_models.json"): click.echo(f" • {key}: {'Yes' if value else 'No'}") def list(self): + with click_spinner.spinner(): + result = self.endpoint.get_all_models() self.list_default() - result = self.endpoint.get_all_models() if result["success"]: if not result["data"]: click.echo("No models to list.") @@ -70,7 +73,8 @@ def list(self): click.echo(f"Failed to list models: {result['response']['detail']}") def get_model(self, model_id): - result = self.endpoint.get_model(model_id) + with click_spinner.spinner(): + result = self.endpoint.get_model(model_id) if result["success"]: click.echo(f"Name: {result['data']['model_name']}") click.echo(f"ID: {result['data']['model_id']}") @@ -88,7 +92,8 @@ def get_model(self, model_id): click.echo(f"Failed to get model. {result['response']['detail']}") def read_file(self, model_id, file_name): - result = self.endpoint.read_model_file(model_id, file_name) + with click_spinner.spinner(): + result = self.endpoint.read_model_file(model_id, file_name) if result["success"]: click.echo(f"File contents for {file_name}:") click.echo(result["data"]["content"]) @@ -102,9 +107,10 @@ def update( model_id: str | None = None, ): click.echo("\nAttempting to update model file...\n") - result = self.endpoint.update_model_file( - model_name=model_name, model_id=model_id, file_path=file_path - ) + with click_spinner.spinner(): + result = self.endpoint.update_model_file( + model_name=model_name, model_id=model_id, file_path=file_path + ) if result["success"]: upload_date = datetime.strptime( result["data"]["upload_date"], "%Y-%m-%dT%H:%M:%S.%f%z" @@ -119,7 +125,8 @@ def update( def delete_file(self, model_id, file_name): click.echo("\nAttempting to delete model file...\n") - result = self.endpoint.delete_model_file(model_id, file_name) + with click_spinner.spinner(): + result = self.endpoint.delete_model_file(model_id, file_name) if result["success"] and result["data"] is None: click.echo(f"Model file {file_name} deleted successfully") else: @@ -127,7 +134,8 @@ def delete_file(self, model_id, file_name): def delete_model(self, model_id): click.echo("\nAttempting to delete model...\n") - result = self.endpoint.delete_model(model_id) + with click_spinner.spinner(): + result = self.endpoint.delete_model(model_id) if result["success"] and result["data"] is None: click.echo(f"Model {model_id} deleted successfully") else: diff --git a/src/commands/user_auth.py b/src/commands/user_auth.py index ceb73bc..ff5c58e 100644 --- a/src/commands/user_auth.py +++ b/src/commands/user_auth.py @@ -1,6 +1,7 @@ import subprocess import click +import click_spinner from src.api.api_client import APIClient from src.api.auth_api import AuthAPI @@ -15,7 +16,8 @@ @click.option("--password", prompt=True, hide_input=True) def login(email, password): """Login to the application.""" - result = endpoint.login(email, password) + with click_spinner.spinner(): + result = endpoint.login(email, password) if result["success"]: click.echo(f"Successfully logged in. {result['data']}") subprocess.run("quack", shell=True) @@ -26,7 +28,8 @@ def login(email, password): @click.command() def logout(): """Logout from the application.""" - result = endpoint.logout() + with click_spinner.spinner(): + result = endpoint.logout() if result: click.echo("Successfully logged out.") else: @@ -40,7 +43,8 @@ def logout(): @click.option("--password", prompt=True, hide_input=True) def register(ctx, username, email, password): """Register with Duckington Labs.""" - result = user_endpoint.register(username, email, password) + with click_spinner.spinner(): + result = user_endpoint.register(username, email, password) if result["success"]: click.echo("Welcome to Duckington Labs, you've successfully registered as:") click.echo(f" - Username: {result['data']['user_name']}") From 3a4204341687c0894b409f11594ba574cd956527 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Thu, 16 Jan 2025 20:34:23 -0330 Subject: [PATCH 10/12] updated requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 96bb634..10c9163 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ certifi==2024.8.30 charset-normalizer==3.3.2 click==8.1.3 colorama==0.4.6 -coverage==7.2.7 +coverage==7.5 exceptiongroup==1.2.2 idna==3.10 importlib_metadata==8.5.0 From 8f4b7f89ffc53d1f68ea5e3a0a4fe9144e561854 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Sun, 19 Jan 2025 14:50:32 -0330 Subject: [PATCH 11/12] update machine_types to lowercase --- src/commands/machine.py | 44 +++++++++++++++++++--------------- src/models/default_models.json | 24 +++++++++---------- 2 files changed, 37 insertions(+), 31 deletions(-) diff --git a/src/commands/machine.py b/src/commands/machine.py index e0ae9c5..2d92baf 100644 --- a/src/commands/machine.py +++ b/src/commands/machine.py @@ -7,8 +7,14 @@ from src.utils.groups.subcommand_group import SubCommandGroup +machine_types = ["cpu", "gpu", "fpga"] + + class MachineCommands: def __init__(self): + self.cpu = machine_types[0] + self.gpu = machine_types[1] + self.fpga = machine_types[2] self.endpoint = MachineAPI() def create(self, hardware_type, machine_name, machine_type): @@ -24,15 +30,15 @@ def create(self, hardware_type, machine_name, machine_type): """ with click_spinner.spinner(): match hardware_type: - case "GPU": + case self.gpu: result = self.endpoint.create_gpu_machine( machine_name, machine_type ) - case "FPGA": + case self.fpga: result = self.endpoint.create_fpga_machine( machine_name, machine_type ) - case "CPU": + case self.cpu: result = self.endpoint.create_cpu_machine( machine_name, machine_type ) @@ -52,12 +58,12 @@ def pull_model(self, hardware_type, machine_id, model_name): click.echo(f"Pulling model {model_name} for machine {machine_id}") with click_spinner.spinner(): match hardware_type: - case "GPU": + case self.gpu: result = self.endpoint.pull_gpu_model(machine_id, model_name) - case "FPGA": + case self.fpga: click.secho("FPGA pull model not supported yet.", fg="yellow") return - case "CPU": + case self.cpu: result = self.endpoint.pull_cpu_model(machine_id, model_name) case _: click.echo("Invalid hardware type.") @@ -74,12 +80,12 @@ def delete_machine_model(self, hardware_type, machine_id, model_name): click.echo(f"Deleting model {model_name} for machine {machine_id}") with click_spinner.spinner(): match hardware_type: - case "GPU": + case self.gpu: result = self.endpoint.delete_gpu_model(machine_id, model_name) - case "FPGA": + case self.fpga: click.secho("FPGA delete model not supported yet.", fg="yellow") return - case "CPU": + case self.cpu: result = self.endpoint.delete_cpu_model(machine_id, model_name) case _: click.echo("Invalid hardware type.") @@ -94,11 +100,11 @@ def delete_machine_model(self, hardware_type, machine_id, model_name): def get_inference_url(self, hardware_type, machine_id): with click_spinner.spinner(): match hardware_type: - case "GPU": + case self.gpu: result = self.endpoint.get_gpu_inference_url(machine_id) - case "FPGA": + case self.fpga: result = self.endpoint.get_fpga_inference_url(machine_id) - case "CPU": + case self.cpu: result = self.endpoint.get_cpu_inference_url(machine_id) case _: click.echo("Invalid hardware type.") @@ -170,7 +176,7 @@ def machine(ctx): @machine.command() @click.pass_context -@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.argument("hardware-type", type=click.Choice(machine_types)) @click.option("--machine-name", "-n", required=True, prompt="Enter machine name") @click.option("--machine-type", "-t", required=True, prompt="Choose machine type") def create(ctx, hardware_type, machine_name, machine_type): @@ -185,9 +191,9 @@ def create(ctx, hardware_type, machine_name, machine_type): - FPGA: \t[f1.2xlarge, f1.4xlarge, f1.16xlarge] """ valid_machine_types = { - "GPU": ["g4dn.xlarge"], - "FPGA": ["f1.2xlarge", "f1.4xlarge", "f1.16xlarge"], - "CPU": ["t2.micro", "m5.xlarge", "m5.2xlarge"], + "gpu": ["g4dn.xlarge"], + "fpga": ["f1.2xlarge", "f1.4xlarge", "f1.16xlarge"], + "cpu": ["t2.micro", "m5.xlarge", "m5.2xlarge"], } if machine_type not in valid_machine_types[hardware_type]: @@ -202,7 +208,7 @@ def create(ctx, hardware_type, machine_name, machine_type): @machine.command() @click.pass_context -@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.argument("hardware-type", type=click.Choice(machine_types)) @click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") @click.option("--model-name", "-m", required=True, prompt="Enter model name") def pull_model(ctx, hardware_type, machine_id, model_name): @@ -226,7 +232,7 @@ def pull_model(ctx, hardware_type, machine_id, model_name): @machine.command() @click.pass_context -@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.argument("hardware-type", type=click.Choice(machine_types)) @click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") @click.option("--model-name", "-m", required=True, prompt="Enter model name") def delete_model(ctx, hardware_type, machine_id, model_name): @@ -236,7 +242,7 @@ def delete_model(ctx, hardware_type, machine_id, model_name): @machine.command() @click.pass_context -@click.argument("hardware-type", type=click.Choice(["CPU", "GPU", "FPGA"])) +@click.argument("hardware-type", type=click.Choice(machine_types)) @click.option("--machine-id", "-id", required=True, prompt="Enter machine ID") def infer_url(ctx, hardware_type, machine_id): """Get inference URL for machine with machine ID.""" diff --git a/src/models/default_models.json b/src/models/default_models.json index 3db1aa7..4c2d64d 100644 --- a/src/models/default_models.json +++ b/src/models/default_models.json @@ -2,36 +2,36 @@ { "name": "llama3:8b", "available": { - "FPGA": false, - "GPU": true, - "CPU": false + "fpga": false, + "gpu": true, + "cpu": false }, "description": "Meta's Llama 3" }, { "name": "llama2:7b", "available": { - "FPGA": false, - "GPU": true, - "CPU": true + "fpga": false, + "gpu": true, + "cpu": true }, "description": "Meta's Llama 2" }, { "name": "smollm:135m", "available": { - "FPGA": false, - "GPU": true, - "CPU": true + "fpga": false, + "gpu": true, + "cpu": true }, "description": "Small Language Model" }, { "name": "tinyllamas-stories-110m", "available": { - "FPGA": true, - "GPU": true, - "CPU": true + "fpga": true, + "gpu": true, + "cpu": true }, "description": "kparthy's Tiny Llama Stories" } From 45beb61a11ccf051098dfc7520d16d16671c8408 Mon Sep 17 00:00:00 2001 From: K-rolls Date: Sun, 19 Jan 2025 14:57:50 -0330 Subject: [PATCH 12/12] update tests --- test/commands/test_key.py | 2 +- test/commands/test_machine.py | 68 +++++++++++++++++------------------ 2 files changed, 35 insertions(+), 35 deletions(-) diff --git a/test/commands/test_key.py b/test/commands/test_key.py index 25211cb..eec00fb 100644 --- a/test/commands/test_key.py +++ b/test/commands/test_key.py @@ -70,7 +70,7 @@ def test_list_api_keys_success(self): with patch("click.echo") as mock_print: self.key_commands.list_api_keys() - self.assertEqual(mock_print.call_count, 3) # Header + 2 keys + self.assertEqual(mock_print.call_count, 3) def test_list_api_keys_empty(self): self.mock_instance.list_api_keys.return_value = {"success": True, "data": []} diff --git a/test/commands/test_machine.py b/test/commands/test_machine.py index 6031b37..f7f8e33 100644 --- a/test/commands/test_machine.py +++ b/test/commands/test_machine.py @@ -78,7 +78,7 @@ def tearDown(self): def test_create_gpu_machine_success(self): with patch("click.echo") as mock_print: - self.machine_commands.create("GPU", "test-machine", "g4dn.xlarge") + self.machine_commands.create("gpu", "test-machine", "g4dn.xlarge") mock_print.assert_any_call("Machine created successfully. Details:") mock_print.assert_any_call( json.dumps( @@ -88,7 +88,7 @@ def test_create_gpu_machine_success(self): def test_create_fpga_machine_success(self): with patch("click.echo") as mock_print: - self.machine_commands.create("FPGA", "test-machine", "f1.2xlarge") + self.machine_commands.create("fpga", "test-machine", "f1.2xlarge") mock_print.assert_any_call("Machine created successfully. Details:") mock_print.assert_any_call( json.dumps( @@ -99,7 +99,7 @@ def test_create_fpga_machine_success(self): def test_create_cpu_machine_success(self): with patch("click.echo") as mock_print: - self.machine_commands.create("CPU", "test-machine", "t2.micro") + self.machine_commands.create("cpu", "test-machine", "t2.micro") mock_print.assert_any_call("Machine created successfully. Details:") mock_print.assert_any_call( json.dumps( @@ -110,7 +110,7 @@ def test_create_cpu_machine_success(self): def test_create_gpu_machine_failure(self): self.mock_instance.create_gpu_machine.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.create("GPU", "test-machine", "g4dn.xlarge") + self.machine_commands.create("gpu", "test-machine", "g4dn.xlarge") mock_print.assert_called_with( "Failed to create machine. Check machine name and type and try again." ) @@ -118,7 +118,7 @@ def test_create_gpu_machine_failure(self): def test_create_fpga_machine_failure(self): self.mock_instance.create_fpga_machine.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.create("FPGA", "test-machine", "f1.2xlarge") + self.machine_commands.create("fpga", "test-machine", "f1.2xlarge") mock_print.assert_called_with( "Failed to create machine. Check machine name and type and try again." ) @@ -126,7 +126,7 @@ def test_create_fpga_machine_failure(self): def test_create_cpu_machine_failure(self): self.mock_instance.create_cpu_machine.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.create("CPU", "test-machine", "t2.micro") + self.machine_commands.create("cpu", "test-machine", "t2.micro") mock_print.assert_called_with( "Failed to create machine. Check machine name and type and try again." ) @@ -138,17 +138,17 @@ def test_create_machine_invalid_hardware_type(self): def test_pull_gpu_model_success(self): with patch("click.echo") as mock_print: - self.machine_commands.pull_model("GPU", "test123", "test-model") + self.machine_commands.pull_model("gpu", "test123", "test-model") mock_print.assert_called_with("Model test-model pulled successfully.") def test_pull_cpu_model_success(self): with patch("click.echo") as mock_print: - self.machine_commands.pull_model("CPU", "test123", "test-model") + self.machine_commands.pull_model("cpu", "test123", "test-model") mock_print.assert_called_with("Model test-model pulled successfully.") def test_pull_fpga_model_not_supported(self): with patch("click.secho") as mock_print: - self.machine_commands.pull_model("FPGA", "test123", "test-model") + self.machine_commands.pull_model("fpga", "test123", "test-model") mock_print.assert_called_with( "FPGA pull model not supported yet.", fg="yellow" ) @@ -156,7 +156,7 @@ def test_pull_fpga_model_not_supported(self): def test_pull_gpu_model_failure(self): self.mock_instance.pull_gpu_model.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.pull_model("GPU", "test123", "test-model") + self.machine_commands.pull_model("gpu", "test123", "test-model") mock_print.assert_called_with( "Failed to pull model. Check machine ID and model name and try again." ) @@ -164,7 +164,7 @@ def test_pull_gpu_model_failure(self): def test_pull_cpu_model_failure(self): self.mock_instance.pull_cpu_model.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.pull_model("CPU", "test123", "test-model") + self.machine_commands.pull_model("cpu", "test123", "test-model") mock_print.assert_called_with( "Failed to pull model. Check machine ID and model name and try again." ) @@ -176,17 +176,17 @@ def test_pull_model_invalid_hardware_type(self): def test_delete_gpu_model_success(self): with patch("click.echo") as mock_print: - self.machine_commands.delete_machine_model("GPU", "test123", "test-model") + self.machine_commands.delete_machine_model("gpu", "test123", "test-model") mock_print.assert_called_with("Model test-model deleted successfully.") def test_delete_cpu_model_success(self): with patch("click.echo") as mock_print: - self.machine_commands.delete_machine_model("CPU", "test123", "test-model") + self.machine_commands.delete_machine_model("cpu", "test123", "test-model") mock_print.assert_called_with("Model test-model deleted successfully.") def test_delete_fpga_model_not_supported(self): with patch("click.secho") as mock_print: - self.machine_commands.delete_machine_model("FPGA", "test123", "test-model") + self.machine_commands.delete_machine_model("fpga", "test123", "test-model") mock_print.assert_called_with( "FPGA delete model not supported yet.", fg="yellow" ) @@ -194,7 +194,7 @@ def test_delete_fpga_model_not_supported(self): def test_delete_gpu_model_failure(self): self.mock_instance.delete_gpu_model.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.delete_machine_model("GPU", "test123", "test-model") + self.machine_commands.delete_machine_model("gpu", "test123", "test-model") mock_print.assert_called_with( "Failed to delete model. Check machine ID and model name and try again." ) @@ -202,7 +202,7 @@ def test_delete_gpu_model_failure(self): def test_delete_cpu_model_failure(self): self.mock_instance.delete_cpu_model.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.delete_machine_model("CPU", "test123", "test-model") + self.machine_commands.delete_machine_model("cpu", "test123", "test-model") mock_print.assert_called_with( "Failed to delete model. Check machine ID and model name and try again." ) @@ -216,12 +216,12 @@ def test_delete_model_invalid_hardware_type(self): def test_get_gpu_inference_url_success(self): with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("GPU", "test123") + self.machine_commands.get_inference_url("gpu", "test123") mock_print.assert_called_with("Inference URL: http://gpu-inference-url") def test_get_cpu_inference_url_success(self): with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("CPU", "test123") + self.machine_commands.get_inference_url("cpu", "test123") mock_print.assert_called_with("Inference URL: http://cpu-inference-url") def test_get_fpga_inference_url_success(self): @@ -230,7 +230,7 @@ def test_get_fpga_inference_url_success(self): "data": {"inference_url": "http://fpga-inference-url"}, } with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("FPGA", "test123") + self.machine_commands.get_inference_url("fpga", "test123") mock_print.assert_called_with("Inference URL: http://fpga-inference-url") def test_get_inference_url_invalid_hardware_type(self): @@ -241,7 +241,7 @@ def test_get_inference_url_invalid_hardware_type(self): def test_get_cpu_inference_url_failure(self): self.mock_instance.get_cpu_inference_url.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("CPU", "test123") + self.machine_commands.get_inference_url("cpu", "test123") mock_print.assert_called_with( "Failed to get inference URL. Check machine ID and try again." ) @@ -249,7 +249,7 @@ def test_get_cpu_inference_url_failure(self): def test_get_fpga_inference_url_failure(self): self.mock_instance.get_fpga_inference_url.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("FPGA", "test123") + self.machine_commands.get_inference_url("fpga", "test123") mock_print.assert_called_with( "Failed to get inference URL. Check machine ID and try again." ) @@ -257,7 +257,7 @@ def test_get_fpga_inference_url_failure(self): def test_get_gpu_inference_url_failure(self): self.mock_instance.get_gpu_inference_url.return_value = {"success": False} with patch("click.echo") as mock_print: - self.machine_commands.get_inference_url("GPU", "test123") + self.machine_commands.get_inference_url("gpu", "test123") mock_print.assert_called_with( "Failed to get inference URL. Check machine ID and try again." ) @@ -386,7 +386,7 @@ def test_create_success_command(self, mock_create): machine, [ "create", - "GPU", + "gpu", "--machine-name", "test-machine", "--machine-type", @@ -394,7 +394,7 @@ def test_create_success_command(self, mock_create): ], ) self.assertEqual(result.exit_code, 0) - mock_create.assert_called_once_with("GPU", "test-machine", "g4dn.xlarge") + mock_create.assert_called_once_with("gpu", "test-machine", "g4dn.xlarge") @patch("src.commands.machine.MachineCommands.create") def test_create_failure_command(self, mock_create): @@ -403,7 +403,7 @@ def test_create_failure_command(self, mock_create): machine, [ "create", - "GPU", + "gpu", "--machine-name", "test-machine", "--machine-type", @@ -418,14 +418,14 @@ def test_pull_model_command(self, mock_pull_model): with patch( "builtins.open", unittest.mock.mock_open( - read_data='[{"name": "test-model", "available": {"GPU": true}}]' + read_data='[{"name": "test-model", "available": {"gpu": true}}]' ), ): result = self.runner.invoke( machine, [ "pull-model", - "GPU", + "gpu", "--machine-id", "test123", "--model-name", @@ -433,21 +433,21 @@ def test_pull_model_command(self, mock_pull_model): ], ) self.assertEqual(result.exit_code, 0) - mock_pull_model.assert_called_once_with("GPU", "test123", "test-model") + mock_pull_model.assert_called_once_with("gpu", "test123", "test-model") @patch("src.commands.machine.MachineCommands.pull_model") def test_pull_model_failure_command(self, mock_pull_model): with patch( "builtins.open", unittest.mock.mock_open( - read_data='[{"name": "test-model", "available": {"GPU": false}}]' + read_data='[{"name": "test-model", "available": {"gpu": false}}]' ), ): result = self.runner.invoke( machine, [ "pull-model", - "GPU", + "gpu", "--machine-id", "test123", "--model-name", @@ -463,7 +463,7 @@ def test_delete_model_command(self, mock_delete_model): machine, [ "delete-model", - "GPU", + "gpu", "--machine-id", "test123", "--model-name", @@ -471,7 +471,7 @@ def test_delete_model_command(self, mock_delete_model): ], ) self.assertEqual(result.exit_code, 0) - mock_delete_model.assert_called_once_with("GPU", "test123", "test-model") + mock_delete_model.assert_called_once_with("gpu", "test123", "test-model") @patch("src.commands.machine.MachineCommands.get_inference_url") def test_infer_url_command(self, mock_get_inference_url): @@ -479,13 +479,13 @@ def test_infer_url_command(self, mock_get_inference_url): machine, [ "infer-url", - "GPU", + "gpu", "--machine-id", "test123", ], ) self.assertEqual(result.exit_code, 0) - mock_get_inference_url.assert_called_once_with("GPU", "test123") + mock_get_inference_url.assert_called_once_with("gpu", "test123") @patch("src.commands.machine.MachineCommands.list") def test_list_command(self, mock_list):