WIP: finetuning CLI

truss/cli/cli.py

@@ -16,6 +16,7 @@
 import rich.table
 import rich.traceback
 import rich_click as click
+import yaml
 from InquirerPy import inquirer
 from rich.console import Console
 
@@ -76,6 +77,13 @@
                 "row_styles": ["red"],
             },
         },
+        {
+            "name": "Finetune",
+            "commands": ["finetune"],
+            "table_styles": {  # type: ignore
+                "row_styles": ["orange"],
+            },
+        },
     ]
 }
 
@@ -814,6 +822,7 @@ def init_chain(
 def _load_example_chainlet_code() -> str:
     try:
         from truss_chains import example_chainlet
+
     # if the example is faulty, a validation error would be raised
     except Exception as e:
         raise Exception("Failed to load starter code. Please notify support.") from e
@@ -1326,9 +1335,41 @@ def _get_truss_from_directory(target_directory: Optional[str] = None):
     return truss.load(target_directory)
 
 
+# Finetuning Stuff #########################################################################
+
+
+@truss_cli.command()
+@click.argument("file_path", type=click.Path(exists=True))
+@click.argument("name", type=str, required=True)
+@click.option(
+    "--remote",
+    type=str,
+    required=False,
+    help="Name of the remote in .trussrc to push to",
+)
+def finetune(
+    file_path: Path,
+    name: str,
+    remote: str,
+):
+    """Finetune on Baseten"""
+    print(f"We finetuning now with file: {file_path}, name: {name}")
+
+    if not remote:
+        remote = inquire_remote_name(RemoteFactory.get_available_config_names())
+
+    remote_provider = RemoteFactory.create(remote=remote)
+    with open(file_path, "r") as file:
+        config = yaml.safe_load(file)
+
+    id = remote_provider.finetune(name, config)
+    print(id)
+
+
 truss_cli.add_command(container)
 truss_cli.add_command(image)
 truss_cli.add_command(chains)
+truss_cli.add_command(finetune)
 
 if __name__ == "__main__":
     truss_cli()

truss/remote/baseten/api.py

@@ -84,6 +84,28 @@ def _post_graphql_query(self, query_string: str) -> dict:
             raise ApiError(message, error_code)
         return resp_dict
 
+    def kickoff_finetuning_job(
+        self,
+        name: str,
+        config: str,
+        client_version: str,
+    ):
+        query_string = f"""
+        mutation {{
+            kickoff_finetuning_job(
+                config: "{config}",
+                client_version: "{client_version}"
+                name: "{name}"
+                s3_key: ""
+            ) {{
+                id
+            }}
+        }}
+        """
+
+        resp = self._post_graphql_query(query_string)
+        return resp["data"]["kickoff_finetuning_job"]
+
     def model_s3_upload_credentials(self):
         query_string = """
         {

truss/remote/baseten/remote.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple
 
+import truss
 import yaml
 from requests import ReadTimeout
 from truss.constants import PRODUCTION_ENVIRONMENT_NAME
@@ -115,6 +116,16 @@ def get_chainlets(
             )
         ]
 
+    def finetune(self, name: str, config: dict):
+        encoded_config_str = base64_encoded_json_str(config)
+
+        ft_job_id = self.api.kickoff_finetuning_job(
+            name=name,
+            config=encoded_config_str,
+            client_version=f"truss=={truss.version()}",
+        )
+        return ft_job_id
+
     def push(  # type: ignore
         self,
         truss_handle: TrussHandle,