Add ability to reference other ModelKits in a Kitfile #260

amisevsk · 2024-04-23T22:20:22Z

Description

This PR makes a number of changes:

We're adding a field to the Kitfile, under model:

model:
  name: <model-name>
  path: <path>
  parts: # New
    - name: readme
      path: my-readme.md
      description: My Readme file
    - <another part>

The parts field is currently a list of structs that look similar to models. It can be used to e.g. store a license/readme alongside binary model data.

Kitfiles now allow you to refer to other modelkits in the model's path field:
```
model:
  name: <model-name>
  path: ghcr.io/jozu-ai/llama-2:7b-text-q4_0
  parts:
    - name: additional-file
      path: ./my-file.bin
```
I've kept the field path even though it could be a model reference now for compatibility.

Testing

For testing purposes, I've packed and pushed two images: docker.io/amisevsk/artifact-test:base and docker.io/amisevsk/artifact-test:sub. The sub image references the base image in its Kitfile (see kit info --remote docker.io/amisevsk/artifact-test:sub)

I intend to add tests for packing+unpacking modelkits with references as above, hopefully soon.

Linked issues

Closes #85

amisevsk · 2024-04-24T21:54:13Z

I've rebased on main and added a few tests for the ModelKit references. For reviewers, commits c2f2652..b1145df are unchanged except for one squashed-in commit that fixes a bug I found while testing: amisevsk@4541148

Add field 'parts' to the model part of a modelkit, allowing for specifying additional components in addition to a "base" model.

Add ability to pack modelkits where the model references a different modelkit in its path

Enable kitops to unpack modelkits where the model references a different modelkit in its path

* Make every command use RunE instead of Run so that we can return an error * Instead of calling os.Exit(1) on failure, return a generic error to signal failure and call os.Exit(1) in the root command's Execute()

pkg/artifact/kit-file.go

gorkem · 2024-04-25T14:56:52Z

pkg/cmd/pull/pull.go

@@ -18,6 +22,67 @@ import (
 	"oras.land/oras-go/v2/registry/remote"
 )

+func runPull(ctx context.Context, opts *pullOptions) (ocispec.Descriptor, error) {


The implementation is good. Is there room for us to improve it so that the actual pulls are async and we can run several pull operations in parallel?

In the future, yes, shouldn't be a problem. For now I think doing it sequentially is better since it leads to the output being sequential ("pulling referenced image X... pulling image Y"). We're bottlenecked by network bandwidth anyways, so I don't know that doing it in parallel would run faster.

I do not think we can saturate the bandwidth that easily, but we can take the iterative approach.

cmd/root.go

The Cobra library has really annoying behaviour around printing the usage/error message (e.g. "expected 2 args, received 1") where the usage/error message is printed both in cases of incorrect usage _and_ when a RunE command fails (returning an error). This is undesirable as a) we can't control the format of the error message, and b) printing usage hides the error under a wall of irrelevant text (the user used the CLI correctly, there was just an error while running -- the fix isn't to correct CLI usage). Setting cmd.SilenceUsage and cmd.SilenceErrors will mean running `kit asdf` will print nothing. Since we need commands to return an error for testing purposes (calling os.Exit(1) exits the test suite), we have to work around this by setting cmd.SilenceUsage and cmd.SilenceErrors within each RunE function rather than on the command as it's being created.

amisevsk · 2024-04-25T15:31:17Z

Cobra is giving me a lot of trouble with printing messages (hence the ugly workaround in 1db15f8). It seems like the only two "easy" options in cobra are

Always print error/usage (default). This means that if e.g. we fail to contact a remote registry, we'll still print the usage message, effectively hiding the actual error

❯ kit pack -t test:test ./scratch/test-not-exists
Configuration already exists in storage: sha256:4bad644891ff9a2c4f95e4a97202e71ed3762874f1bde0853a27f7ab50873019
Failed to pack model kit: model path not-exist.txt does not exist  <-- this is the relevant message
Error: failed to run

Usage:
  kit pack [flags] DIRECTORY
<dozens of lines of help text>...

This is useless for us since a command failing with error means the user used the CLI correctly and we hit an error somewhere else. Telling them how to use the CLI is annoying.

Never print error/usage (SilenceUsage: true, SilenceErrors: true). This means that you can run kit asdf and it will output nothing. Obviously not good.

Since we have to use RunE functions (calling os.Exit(1) in tests will just exit the test suite), our options are

Set SilenceUsage and SilenceErrors inside the command after it starts running, or

Wrap commands so that they don't use RunE:

Run: func(cmd *cobra.Command, args []string) {
  if err := runCommand(opts)(cmd, args); err != nil {
    os.Exit(1)
  }
}

I've chosen the first option here.

gorkem · 2024-04-25T17:48:15Z

pkg/artifact/kit-file.md

+  - `parts`: List of related files for the model (e.g. LoRA weights)
+    - `name`: Identifier for the part
+    - `path`: Location of the file or a directory relative to the context
+    - `description`: Description of the part


Should reflect the change on the kit-file.go

gorkem · 2024-04-25T17:53:46Z

pkg/cmd/info/cmd.go

+	return func(cmd *cobra.Command, args []string) error {
+		// Don't want to output usage and placeholder error if this command fails. We can't set this higher up
+		// since it will suppress usage and error messages for actual incorrect usage (e.g. wrong number of args)
+		cmd.SilenceUsage = true


Would it work if we add these on the root after or on addSubCommands call?

Nope, that still suppresses usage in all cases (no message/error if you use a non-existent command or the wrong number of args).

strange. I would assume looping through all the added commands in here would not be different than adding them on each file.

As far as I understand, it's an execution time switch -- if parsing the command hits an error before it hits a RunE, the error and usage will be printed. Once it hits the RunE and calls the runCommand(opts) function, we disable printing usage/errors and so it won't print anything on errors after that point. It would maybe work in the root command's persistentPreRunE but I'm kind of tired of grappling with such a weird decision.

Turns out that does work: #268

amisevsk requested review from gorkem and nida-hasan April 23, 2024 22:20

amisevsk force-pushed the remote-modelkit-refs branch 2 times, most recently from 5b4d4c3 to 56dfcd2 Compare April 23, 2024 22:36

nida-hasan approved these changes Apr 23, 2024

View reviewed changes

amisevsk force-pushed the remote-modelkit-refs branch from 56dfcd2 to a23e882 Compare April 24, 2024 21:50

amisevsk force-pushed the remote-modelkit-refs branch from a23e882 to 7821984 Compare April 24, 2024 21:55

amisevsk added 9 commits April 25, 2024 10:54

Add parts to modelkit model

d7c1d6a

Add field 'parts' to the model part of a modelkit, allowing for specifying additional components in addition to a "base" model.

Pack and unpack modelkit parts in addition to model

3311002

Add tests for model parts

49356ea

Add support for packing remote model references

3be2d18

Add ability to pack modelkits where the model references a different modelkit in its path

Add ability to unpack modelkits with remote references

cc565fc

Enable kitops to unpack modelkits where the model references a different modelkit in its path

Allow pulling modelkits with remote references.

97fd9bf

Refactor functional testing to allow different test structs

651b781

Refactor all commands to not use os.Exit() to enable testing

d3780c4

* Make every command use RunE instead of Run so that we can return an error * Instead of calling os.Exit(1) on failure, return a generic error to signal failure and call os.Exit(1) in the root command's Execute()

Add tests for packing and unpacking with modelkit references

5478cb7

gorkem requested changes Apr 25, 2024

View reviewed changes

amisevsk force-pushed the remote-modelkit-refs branch from 7821984 to 1db15f8 Compare April 25, 2024 15:20

gorkem reviewed Apr 25, 2024

View reviewed changes

gorkem approved these changes Apr 25, 2024

View reviewed changes

Rename model part "description" field to "type"

ea2d6f5

amisevsk force-pushed the remote-modelkit-refs branch from 7e20d0e to ea2d6f5 Compare April 25, 2024 21:16

amisevsk merged commit 4d73b84 into jozu-ai:main Apr 25, 2024
3 checks passed

amisevsk deleted the remote-modelkit-refs branch April 25, 2024 21:24

amisevsk mentioned this pull request Apr 25, 2024

Set SilenceErrors and SilenceUsage on root command PersistentPreRun #268

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ability to reference other ModelKits in a Kitfile #260

Add ability to reference other ModelKits in a Kitfile #260

amisevsk commented Apr 23, 2024 •

edited

amisevsk commented Apr 24, 2024

gorkem Apr 25, 2024

amisevsk Apr 25, 2024

gorkem Apr 25, 2024

amisevsk commented Apr 25, 2024

gorkem Apr 25, 2024

gorkem Apr 25, 2024

amisevsk Apr 25, 2024

gorkem Apr 25, 2024

amisevsk Apr 25, 2024

amisevsk Apr 25, 2024

Add ability to reference other ModelKits in a Kitfile #260

Add ability to reference other ModelKits in a Kitfile #260

Conversation

amisevsk commented Apr 23, 2024 • edited

Description

Testing

Linked issues

amisevsk commented Apr 24, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amisevsk commented Apr 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amisevsk commented Apr 23, 2024 •

edited