diff --git a/core/services/addons/hostops/broker_test.go b/core/services/addons/hostops/broker_test.go index 84215608a..ecb16bddb 100644 --- a/core/services/addons/hostops/broker_test.go +++ b/core/services/addons/hostops/broker_test.go @@ -7,10 +7,11 @@ import ( "encoding/json" "net/http" "net/http/httptest" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) func env(addon, op, scope string, exp time.Time) contract.Envelope { @@ -27,22 +28,18 @@ type fixture struct { seen *Call } -func newFixture(t *testing.T, has func(user, feature string) bool) *fixture { - t.Helper() +func newFixture(has func(user, feature string) bool) *fixture { f := &fixture{inflight: NewInFlight(), seen: &Call{}} f.b = New(f.inflight, has, contract.DefaultLimits()) err := f.b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: func(_ context.Context, c Call) (json.RawMessage, error) { *f.seen = c return json.RawMessage(`{"ok":true}`), nil }}) - if err != nil { - t.Fatal(err) - } - if err := f.b.Register(Spec{Op: "finetune.start", Scope: "apply", Feature: "fine_tuning", Handler: func(context.Context, Call) (json.RawMessage, error) { + Expect(err).ToNot(HaveOccurred()) + err = f.b.Register(Spec{Op: "finetune.start", Scope: "apply", Feature: "fine_tuning", Handler: func(context.Context, Call) (json.RawMessage, error) { return json.RawMessage(`{"id":"j1"}`), nil - }}); err != nil { - t.Fatal(err) - } + }}) + Expect(err).ToNot(HaveOccurred()) f.h = f.b.Handler("fine-tuning", "secret", []string{"finetune."}) return f } @@ -60,103 +57,100 @@ func (f *fixture) call(credential, operationID, op, body string) (int, contract. func always(string, string) bool { return true } -func TestAWrongCredentialIsRefused(t *testing.T) { - f := newFixture(t, always) - if code, _ := f.call("nope", "op1", "finetune.get", `{}`); code != http.StatusUnauthorized { - t.Fatalf("got %d", code) +var _ = Describe("Host operation broker", func() { + live := func(f *fixture, scope string) { + end := f.inflight.Begin(env("fine-tuning", "op1", scope, time.Now().Add(time.Minute))) + DeferCleanup(end) } -} -func TestACallWithoutALiveEnvelopeIsRefused(t *testing.T) { - f := newFixture(t, always) - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses a wrong credential", func() { + f := newFixture(always) + code, _ := f.call("nope", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusUnauthorized)) + }) -func TestAnExpiredAuthorityIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(-time.Second))) - defer end() - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses a call without a live envelope", func() { + f := newFixture(always) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestAReadScopeCannotRunAnApplyOperation(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - code, out := f.call("secret", "op1", "finetune.start", `{}`) - if code != http.StatusForbidden || out.Error.Code != "scope" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an expired authority", func() { + f := newFixture(always) + end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(-time.Second))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestTheFeatureOfTheActorIsChecked(t *testing.T) { - f := newFixture(t, func(user, feature string) bool { return feature != "fine_tuning" }) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error.Code != "feature" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("does not let a read scope run an apply operation", func() { + f := newFixture(always) + live(f, "read") + code, out := f.call("secret", "op1", "finetune.start", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("scope")) + }) -func TestTheOperationRunsAsTheActorOfTheEnvelope(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - // An argument that names a user is refused, whatever the schema says. - code, out := f.call("secret", "op1", "finetune.get", `{"userId":"admin"}`) - if code != http.StatusBadRequest || out.Error.Code != "bad_request" { - t.Fatalf("%d %+v", code, out.Error) - } - if code, _ := f.call("secret", "op1", "finetune.get", `{"jobId":"j1"}`); code != http.StatusOK { - t.Fatalf("got %d", code) - } - if f.seen.Actor.UserID != "u1" || f.seen.AddonID != "fine-tuning" { - t.Fatalf("the operation did not run as the envelope's actor: %+v", f.seen) - } -} + It("checks the feature of the actor", func() { + f := newFixture(func(user, feature string) bool { return feature != "fine_tuning" }) + live(f, "read") + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("feature")) + }) -func TestAnOperationOutsideTheAllowedPrefixIsRefused(t *testing.T) { - f := newFixture(t, always) - _ = f.b.Register(Spec{Op: "quantize.get", Scope: "read", Feature: "quantization", Handler: func(context.Context, Call) (json.RawMessage, error) { return nil, nil }}) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "quantize.get", `{}`); code != http.StatusForbidden || out.Error.Code != "unknown_op" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("runs the operation as the actor of the envelope", func() { + f := newFixture(always) + live(f, "read") + // An argument that names a user is refused, whatever the schema says. + code, out := f.call("secret", "op1", "finetune.get", `{"userId":"admin"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("bad_request")) + code, _ = f.call("secret", "op1", "finetune.get", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(f.seen.Actor.UserID).To(Equal("u1"), "the operation did not run as the envelope's actor: %+v", f.seen) + Expect(f.seen.AddonID).To(Equal("fine-tuning")) + }) -func TestAnOversizedBodyIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - big := `{"x":"` + string(bytes.Repeat([]byte("a"), int(contract.DefaultLimits().PayloadBytes)+1)) + `"}` - if code, out := f.call("secret", "op1", "finetune.get", big); code != http.StatusRequestEntityTooLarge || out.Error.Code != "too_large" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an operation outside the allowed prefix", func() { + f := newFixture(always) + _ = f.b.Register(Spec{Op: "quantize.get", Scope: "read", Feature: "quantization", Handler: func(context.Context, Call) (json.RawMessage, error) { return nil, nil }}) + live(f, "read") + code, out := f.call("secret", "op1", "quantize.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("unknown_op")) + }) -func TestAnIdempotentOperationRunsOncePerOperationID(t *testing.T) { - f := newFixture(t, always) - runs := 0 - _ = f.b.Register(Spec{Op: "finetune.once", Scope: "apply", Feature: "fine_tuning", Idempotent: true, Handler: func(context.Context, Call) (json.RawMessage, error) { - runs++ - return json.RawMessage(`{"id":"j9"}`), nil - }}) - end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) - defer end() - for i := 0; i < 2; i++ { - if code, _ := f.call("secret", "op1", "finetune.once", `{}`); code != http.StatusOK { - t.Fatalf("call %d: %d", i, code) + It("refuses an oversized body", func() { + f := newFixture(always) + live(f, "read") + big := `{"x":"` + string(bytes.Repeat([]byte("a"), int(contract.DefaultLimits().PayloadBytes)+1)) + `"}` + code, out := f.call("secret", "op1", "finetune.get", big) + Expect(code).To(Equal(http.StatusRequestEntityTooLarge)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("too_large")) + }) + + It("runs an idempotent operation once per operation ID", func() { + f := newFixture(always) + runs := 0 + _ = f.b.Register(Spec{Op: "finetune.once", Scope: "apply", Feature: "fine_tuning", Idempotent: true, Handler: func(context.Context, Call) (json.RawMessage, error) { + runs++ + return json.RawMessage(`{"id":"j9"}`), nil + }}) + live(f, "apply") + for i := 0; i < 2; i++ { + code, _ := f.call("secret", "op1", "finetune.once", `{}`) + Expect(code).To(Equal(http.StatusOK), "call %d", i) } - } - if runs != 1 { - t.Fatalf("a repeated operationId must not run the operation twice, ran %d", runs) - } -} + Expect(runs).To(Equal(1), "a repeated operationId must not run the operation twice") + }) +}) diff --git a/core/services/addons/hostops/engine_datasets_test.go b/core/services/addons/hostops/engine_datasets_test.go index b8a9f470a..123ac1b61 100644 --- a/core/services/addons/hostops/engine_datasets_test.go +++ b/core/services/addons/hostops/engine_datasets_test.go @@ -3,12 +3,12 @@ package hostops import ( "encoding/json" - "errors" "fmt" "strings" - "testing" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeDatasets struct { @@ -37,160 +37,123 @@ func (f *fakeDatasets) Resolve(ref string) (string, bool) { return p, ok } -func datasetArgs(t *testing.T, name, content string) string { - t.Helper() +func datasetArgs(name, content string) string { b, err := json.Marshal(map[string]string{"name": name, "content": content}) - if err != nil { - t.Fatal(err) - } + Expect(err).ToNot(HaveOccurred()) return string(b) } -func datasetBroker(t *testing.T, ds DatasetStore) (*Broker, Spec) { - t.Helper() +func datasetBroker(ds DatasetStore) (*Broker, Spec) { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineDatasets(b, ds); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineDatasets(b, ds)).To(Succeed()) return b, b.specs["finetune.dataset"] } -func TestDatasetStoresTheContentAndReturnsARefNotAPath(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - raw, err := call(t, s, "named-user", "u1", datasetArgs(t, "train.jsonl", "{\"a\":1}\n")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if out.DatasetRef == "" || strings.Contains(out.DatasetRef, "/") { - t.Fatalf("ref %q", out.DatasetRef) - } - if ds.gotName != "train.jsonl" || string(ds.gotData) != "{\"a\":1}\n" { - t.Fatalf("stored %q %q", ds.gotName, ds.gotData) - } - if strings.Contains(string(raw), "/srv") { - t.Fatalf("a path in the result: %s", raw) - } +type datasetResult struct { + DatasetRef string `json:"datasetRef"` } -func TestDatasetIsIdempotentAndFineTuningOnly(t *testing.T) { - _, s := datasetBroker(t, &fakeDatasets{}) - if !s.Idempotent || s.Scope != "apply" || s.Feature != "fine_tuning" { - t.Fatalf("%+v", s) - } -} +var _ = Describe("Dataset operation", func() { + It("stores the content and returns a ref, not a path", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + raw, err := call(s, "named-user", "u1", datasetArgs("train.jsonl", "{\"a\":1}\n")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + Expect(out.DatasetRef).ToNot(BeEmpty()) + Expect(out.DatasetRef).ToNot(ContainSubstring("/")) + Expect(ds.gotName).To(Equal("train.jsonl")) + Expect(string(ds.gotData)).To(Equal("{\"a\":1}\n")) + Expect(string(raw)).ToNot(ContainSubstring("/srv"), "a path in the result") + }) -func TestDatasetRefusesAnEmptyOrOversizedContent(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - limit := 900 * 1024 - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, "a.txt", strings.Repeat("x", limit))); err != nil { - t.Fatalf("exactly 900 KiB refused: %v", err) - } - for name, content := range map[string]string{ - "empty": "", - "one over": strings.Repeat("x", limit+1), - "far over": strings.Repeat("x", 2*limit), - } { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, "a.txt", content)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", name, err) - } - } - if ds.puts != 1 { - t.Fatalf("a refused dataset was stored: %d puts", ds.puts) - } -} + It("is idempotent and fine-tuning only", func() { + _, s := datasetBroker(&fakeDatasets{}) + Expect(s.Idempotent).To(BeTrue()) + Expect(s.Scope).To(Equal("apply")) + Expect(s.Feature).To(Equal("fine_tuning")) + }) -func TestDatasetRefusesANameWithADirectory(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - for _, name := range []string{"../x.jsonl", "a/b.jsonl", "/etc/x.jsonl", `a\b.jsonl`, ".."} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("refuses an empty or oversized content", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + limit := 900 * 1024 + _, err := call(s, "named-user", "u1", datasetArgs("a.txt", strings.Repeat("x", limit))) + Expect(err).ToNot(HaveOccurred(), "exactly 900 KiB refused") + for name, content := range map[string]string{ + "empty": "", + "one over": strings.Repeat("x", limit+1), + "far over": strings.Repeat("x", 2*limit), + } { + _, err := call(s, "named-user", "u1", datasetArgs("a.txt", content)) + Expect(err).To(MatchError(ErrBadRequest), name) } - } - if ds.puts != 0 { - t.Fatalf("a refused name was stored") - } -} + Expect(ds.puts).To(Equal(1), "a refused dataset was stored") + }) -func TestDatasetRefusesAnExtensionThatIsNotATextDataset(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - for _, name := range []string{"x.exe", "x.sh", "x", "x.jsonl.exe", ".jsonl.bin"} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("refuses a name with a directory", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + for _, name := range []string{"../x.jsonl", "a/b.jsonl", "/etc/x.jsonl", `a\b.jsonl`, ".."} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).To(MatchError(ErrBadRequest), name) } - } - for _, name := range []string{"x.jsonl", "x.json", "x.csv", "x.txt"} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); err != nil { - t.Fatalf("%q refused: %v", name, err) - } - } -} + Expect(ds.puts).To(Equal(0), "a refused name was stored") + }) -func TestDatasetRefusesContentThatIsNotUTF8(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - if _, err := call(t, s, "named-user", "u1", "{\"name\":\"x.csv\",\"content\":\"a,b\\n\xff\xfe,1\\n\"}"); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } - if ds.puts != 0 { - t.Fatalf("non-UTF-8 content was stored") - } -} + It("refuses an extension that is not a text dataset", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + for _, name := range []string{"x.exe", "x.sh", "x", "x.jsonl.exe", ".jsonl.bin"} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).To(MatchError(ErrBadRequest), name) + } + for _, name := range []string{"x.jsonl", "x.json", "x.csv", "x.txt"} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).ToNot(HaveOccurred(), "%q refused", name) + } + }) -func TestResolveOnlyReturnsAReferenceTheHostIssuedToThatUser(t *testing.T) { - ds := &fakeDatasets{} - b, s := datasetBroker(t, ds) - raw, err := call(t, s, "named-user", "userA", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if _, ok := b.resolveDataset("userB", out.DatasetRef); ok { - t.Fatal("user B resolved the reference of user A") - } - if p, ok := b.resolveDataset("userA", out.DatasetRef); !ok || p == "" { - t.Fatalf("the issuing user cannot resolve its own reference: %q %v", p, ok) - } -} + It("refuses content that is not UTF-8", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + _, err := call(s, "named-user", "u1", "{\"name\":\"x.csv\",\"content\":\"a,b\\n\xff\xfe,1\\n\"}") + Expect(err).To(MatchError(ErrBadRequest)) + Expect(ds.puts).To(Equal(0), "non-UTF-8 content was stored") + }) -func TestTheIssuedReferencesAreBoundedTo256(t *testing.T) { - ds := &fakeDatasets{} - b, s := datasetBroker(t, ds) - var first, last string - for i := range 257 { - raw, err := call(t, s, "named-user", "u1", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) + It("resolves only a reference the host issued to that user", func() { + ds := &fakeDatasets{} + b, s := datasetBroker(ds) + raw, err := call(s, "named-user", "userA", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + _, ok := b.resolveDataset("userB", out.DatasetRef) + Expect(ok).To(BeFalse(), "user B resolved the reference of user A") + p, ok := b.resolveDataset("userA", out.DatasetRef) + Expect(ok).To(BeTrue(), "the issuing user cannot resolve its own reference") + Expect(p).ToNot(BeEmpty()) + }) + + It("bounds the issued references to 256", func() { + ds := &fakeDatasets{} + b, s := datasetBroker(ds) + var first, last string + for i := range 257 { + raw, err := call(s, "named-user", "u1", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + if i == 0 { + first = out.DatasetRef + } + last = out.DatasetRef } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if i == 0 { - first = out.DatasetRef - } - last = out.DatasetRef - } - if _, ok := b.resolveDataset("u1", first); ok { - t.Fatal("the oldest reference is still issued after 257 datasets") - } - if _, ok := b.resolveDataset("u1", last); !ok { - t.Fatal("the newest reference is not issued") - } -} + _, ok := b.resolveDataset("u1", first) + Expect(ok).To(BeFalse(), "the oldest reference is still issued after 257 datasets") + _, ok = b.resolveDataset("u1", last) + Expect(ok).To(BeTrue(), "the newest reference is not issued") + }) +}) diff --git a/core/services/addons/hostops/engine_jobs_test.go b/core/services/addons/hostops/engine_jobs_test.go index 1710f467d..59ea1d199 100644 --- a/core/services/addons/hostops/engine_jobs_test.go +++ b/core/services/addons/hostops/engine_jobs_test.go @@ -5,15 +5,14 @@ import ( "bytes" "context" "encoding/json" - "errors" "net/http" "net/http/httptest" - "strings" - "testing" "time" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeFT struct { @@ -78,185 +77,29 @@ type backends struct{} func (backends) FineTune() []string { return []string{"trl"} } func (backends) Quantization() []string { return []string{"llama-cpp-quantization"} } -func call(t *testing.T, spec Spec, mode, user, args string) (json.RawMessage, error) { - t.Helper() +func call(spec Spec, mode, user, args string) (json.RawMessage, error) { return spec.Handler(context.Background(), Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: user, Mode: mode}, Arguments: json.RawMessage(args)}) } -func register(t *testing.T, ft FineTuner, q Quantizer) map[string]Spec { - t.Helper() +func register(ft FineTuner, q Quantizer) map[string]Spec { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, q, backends{}); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineJobs(b, ft, q, backends{})).To(Succeed()) return b.specs } -func specs(t *testing.T, ft FineTuner) map[string]Spec { - t.Helper() - return register(t, ft, nil) -} - -func TestStartPassesTheSameUserKeyAsTheRoute(t *testing.T) { - ft := &fakeFT{} - s := specs(t, ft) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data"}`); err != nil || ft.startUser != "u1" { - t.Fatalf("named user: %v %q", err, ft.startUser) - } - // Authentication off: the route passes "" so the job shows in the UI list. - if _, err := call(t, s["finetune.start"], "local-owner", "local-owner", `{"model":"org/base","datasetSource":"org/data"}`); err != nil || ft.startUser != "" { - t.Fatalf("local owner: %v %q", err, ft.startUser) - } -} - -func TestStartRefusesALocalPathAsModelOrDataset(t *testing.T) { - s := specs(t, &fakeFT{}) - for _, args := range []string{ - `{"model":"/etc/passwd","datasetSource":"org/data"}`, - `{"model":"../x","datasetSource":"org/data"}`, - `{"model":"org/base","datasetSource":"/var/data.jsonl"}`, - `{"model":"org\\base","datasetSource":"org/data"}`, - } { - if _, err := call(t, s["finetune.start"], "named-user", "u1", args); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", args, err) - } - } -} - -func TestStartRefusesABackendThatIsNotAFineTuningBackend(t *testing.T) { - s := specs(t, &fakeFT{}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","backend":"llama-cpp"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestStartRefusesAFieldThatIsNotInTheToolSchema(t *testing.T) { - s := specs(t, &fakeFT{}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","extraOptions":{"x":"y"}}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestGetOfAnotherUsersJobIsNotFound(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "other"}}} - s := specs(t, ft) - if _, err := call(t, s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} - -func TestListReturnsOnlyTheCallersJobs(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}}} - s := specs(t, ft) - out, err := call(t, s["finetune.list"], "named-user", "u1", `{}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Jobs []struct { - ID string `json:"id"` - } `json:"jobs"` - } - _ = json.Unmarshal(out, &got) - if len(got.Jobs) != 1 || got.Jobs[0].ID != "a" { - t.Fatalf("a caller must see only its own jobs: %s", out) - } -} - -func TestAJobResultNeverCarriesAFilesystemPath(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1", OutputDir: "/data/fine-tune/j1", ExportStatus: "completed", ExportModelName: "my-model", Message: "saved to /data/fine-tune/j1/out"}}} - s := specs(t, ft) - out, _ := call(t, s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) - if strings.Contains(string(out), "/data/") || strings.Contains(string(out), "/models/") { - t.Fatalf("a result must not leak a server path: %s", out) - } - if !strings.Contains(string(out), `"exportedModelName":"my-model"`) { - t.Fatalf("the registered model name is part of the result: %s", out) - } -} - -func TestQuantizeStartPassesTheSameUserKeyAsTheRoute(t *testing.T) { - q := &fakeQ{} - s := register(t, nil, q) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","quantizationType":"q8_0"}`); err != nil || q.startUser != "u1" { - t.Fatalf("named user: %v %q", err, q.startUser) - } - if _, err := call(t, s["quantize.start"], "local-owner", "local-owner", `{"model":"org/base"}`); err != nil || q.startUser != "" { - t.Fatalf("local owner: %v %q", err, q.startUser) - } - if got := q.started[0].QuantizationType; got != "q8_0" { - t.Fatalf("quantization type not passed: %q", got) - } -} - -func TestQuantizeStartRefusesALocalPathAsModel(t *testing.T) { - s := register(t, nil, &fakeQ{}) - for _, args := range []string{`{"model":"/models/x.gguf"}`, `{"model":"./x"}`, `{"model":"a\\b"}`} { - if _, err := call(t, s["quantize.start"], "named-user", "u1", args); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", args, err) - } - } -} - -func TestQuantizeStartRefusesABackendThatIsNotAQuantizationBackend(t *testing.T) { - s := register(t, nil, &fakeQ{}) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","backend":"trl"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeStartRefusesAFieldThatIsNotInTheToolSchema(t *testing.T) { - s := register(t, nil, &fakeQ{}) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","extraOptions":{"x":"y"}}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeGetOfAnotherUsersJobIsNotFound(t *testing.T) { - q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "other", OutputFile: "/data/q1/m.gguf"}}} - s := register(t, nil, q) - if _, err := call(t, s["quantize.get"], "named-user", "u1", `{"jobId":"q1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeListReturnsOnlyTheCallersJobsWithoutAPath(t *testing.T) { - q := &fakeQ{jobs: map[string]*schema.QuantizationJob{ - "a": {ID: "a", UserID: "u1", OutputFile: "/data/a/m.gguf", OutputDir: "/data/a"}, - "b": {ID: "b", UserID: "u2"}, - }} - s := register(t, nil, q) - out, err := call(t, s["quantize.list"], "named-user", "u1", `{}`) - if err != nil { - t.Fatal(err) - } - if strings.Contains(string(out), "/data/") || strings.Contains(string(out), `"id":"b"`) || !strings.Contains(string(out), `"id":"a"`) { - t.Fatalf("a caller sees only its own jobs and no path: %s", out) - } -} - -func TestNilServicesRegisterNoOperation(t *testing.T) { - b := New(NewInFlight(), always, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, nil, nil, backends{}); err != nil { - t.Fatal(err) - } - if ops := b.Ops(); len(ops) != 0 { - t.Fatalf("no service, no operation: %v", ops) - } +func specs(ft FineTuner) map[string]Spec { + return register(ft, nil) } // serveEngine runs one engine operation through the broker's HTTP channel, as // the add-on process would call it, for a named user holding the features // that has allows. The jobs the fakes hold belong to u1 so that a call with // valid arguments answers 200 for an actor who holds the feature. -func serveEngine(t *testing.T, op, args string, has func(user, feature string) bool) (int, contract.HostResponse) { - t.Helper() +func serveEngine(op, args string, has func(user, feature string) bool) (int, contract.HostResponse) { ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1"}}} q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "u1"}}} b := New(NewInFlight(), has, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, q, backends{}); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineJobs(b, ft, q, backends{})).To(Succeed()) end := b.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Hour))) defer end() h := b.Handler("fine-tuning", "secret", []string{"finetune.", "quantize."}) @@ -270,154 +113,261 @@ func serveEngine(t *testing.T, op, args string, has func(user, feature string) b return rec.Code, out } -// TestEachEngineOperationNeedsItsFeature checks every engine operation against -// its own feature: refused without it, refused with only the other family's -// feature, and answered when the actor holds the operation's feature. -func TestEachEngineOperationNeedsItsFeature(t *testing.T) { - const fineTune, quant = "fine_tuning", "quantization" - cases := []struct { - op, args, feature, other string - }{ - {"finetune.start", `{"model":"org/base","datasetSource":"org/data"}`, fineTune, quant}, - {"finetune.get", `{"jobId":"j1"}`, fineTune, quant}, - {"finetune.list", `{}`, fineTune, quant}, - {"finetune.checkpoints", `{"jobId":"j1"}`, fineTune, quant}, - {"quantize.start", `{"model":"org/base"}`, quant, fineTune}, - {"quantize.get", `{"jobId":"q1"}`, quant, fineTune}, - {"quantize.list", `{}`, quant, fineTune}, - } - for _, c := range cases { - t.Run(c.op, func(t *testing.T) { - none := func(string, string) bool { return false } - if code, out := serveEngine(t, c.op, c.args, none); code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("without any feature: %d %+v", code, out.Error) - } - otherOnly := func(_, f string) bool { return f == c.other } - if code, out := serveEngine(t, c.op, c.args, otherOnly); code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("with only %s: %d %+v", c.other, code, out.Error) - } - ownOnly := func(_, f string) bool { return f == c.feature } - if code, out := serveEngine(t, c.op, c.args, ownOnly); code != http.StatusOK { - t.Fatalf("with %s: %d %+v", c.feature, code, out.Error) +var _ = Describe("Fine-tuning and quantization job operations", func() { + Describe("finetune", func() { + It("start passes the same user key as the route", func() { + ft := &fakeFT{} + s := specs(ft) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.startUser).To(Equal("u1"), "named user") + // Authentication off: the route passes "" so the job shows in the UI list. + _, err = call(s["finetune.start"], "local-owner", "local-owner", `{"model":"org/base","datasetSource":"org/data"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.startUser).To(BeEmpty(), "local owner") + }) + + It("start refuses a local path as model or dataset", func() { + s := specs(&fakeFT{}) + for _, args := range []string{ + `{"model":"/etc/passwd","datasetSource":"org/data"}`, + `{"model":"../x","datasetSource":"org/data"}`, + `{"model":"org/base","datasetSource":"/var/data.jsonl"}`, + `{"model":"org\\base","datasetSource":"org/data"}`, + } { + _, err := call(s["finetune.start"], "named-user", "u1", args) + Expect(err).To(MatchError(ErrBadRequest), args) } }) - } -} -// TestDatasetSourceMustBeAGalleryOrHuggingFaceID refuses a dataset reference -// that is not an id, names the field, and never reaches the service. -func TestDatasetSourceMustBeAGalleryOrHuggingFaceID(t *testing.T) { - bad := []string{ - "https://example.com/data.jsonl", // URL - "hf://org/data", // scheme - "file:///etc/passwd", // scheme and absolute path - "/var/data.jsonl", // absolute path - "../data", // parent segment - "org/../data", // parent segment inside an id - "org/..", // parent segment at the end - "org data", // space - "org/data/extra", // more than one separator - `C:\data`, // backslash - "", // empty - } - for _, src := range bad { - ft := &fakeFT{} - s := specs(t, ft) - args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) - _, err := call(t, s["finetune.start"], "named-user", "u1", string(args)) - if !errors.Is(err, ErrBadRequest) || !strings.Contains(err.Error(), "datasetSource") { - t.Fatalf("%q: want a bad request naming datasetSource, got %v", src, err) - } - if len(ft.started) != 0 { - t.Fatalf("%q reached the service", src) - } - } -} + It("start refuses a backend that is not a fine-tuning backend", func() { + s := specs(&fakeFT{}) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","backend":"llama-cpp"}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) -// TestDatasetSourceRefusalIsA400NamingTheField checks the answer the add-on -// process receives through the channel. -func TestDatasetSourceRefusalIsA400NamingTheField(t *testing.T) { - code, out := serveEngine(t, "finetune.start", `{"model":"org/base","datasetSource":"https://example.com/d.jsonl"}`, func(string, string) bool { return true }) - if code != http.StatusBadRequest || out.Error == nil || out.Error.Code != contract.HostBadRequest || !strings.Contains(out.Error.Message, "datasetSource") { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("start refuses a field that is not in the tool schema", func() { + s := specs(&fakeFT{}) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","extraOptions":{"x":"y"}}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) -// TestDatasetSourceAcceptsHuggingFaceAndGalleryIDs keeps the ids the UI uses. -func TestDatasetSourceAcceptsHuggingFaceAndGalleryIDs(t *testing.T) { - for _, src := range []string{"org/data", "HuggingFaceH4/ultrachat_200k", "my-dataset", "org/data-v1.0"} { - ft := &fakeFT{} - s := specs(t, ft) - args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", string(args)); err != nil { - t.Fatalf("%q: %v", src, err) - } - } -} + It("get of another user's job is not found", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "other"}}} + s := specs(ft) + _, err := call(s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) -func TestStartAcceptsADatasetRefAsDatasetSource(t *testing.T) { - ft := &fakeFT{} - ds := &fakeDatasets{} - b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, nil, backends{}); err != nil { - t.Fatal(err) - } - if err := RegisterEngineDatasets(b, ds); err != nil { - t.Fatal(err) - } - raw, err := call(t, b.specs["finetune.dataset"], "named-user", "u1", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - start := b.specs["finetune.start"] - if _, err := call(t, start, "named-user", "u1", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`); err != nil { - t.Fatal(err) - } - if got := ft.started[len(ft.started)-1].DatasetSource; got != ds.paths[out.DatasetRef] { - t.Fatalf("the service got %q, want the stored path", got) - } - if _, err := call(t, start, "named-user", "u2", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("another user's reference: want ErrBadRequest, got %v", err) - } - if _, err := call(t, start, "named-user", "u1", `{"model":"org/base","datasetSource":"ffffffff-nope.jsonl"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("an unknown reference: want ErrBadRequest, got %v", err) - } -} + It("list returns only the caller's jobs", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}}} + s := specs(ft) + out, err := call(s["finetune.list"], "named-user", "u1", `{}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Jobs []struct { + ID string `json:"id"` + } `json:"jobs"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Jobs).To(HaveLen(1), "a caller must see only its own jobs: %s", out) + Expect(got.Jobs[0].ID).To(Equal("a")) + }) -func TestASecondUsersListDoesNotShowTheFirstUsersJobs(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}, "c": {ID: "c", UserID: "u1"}}} - s := specs(t, ft) - ids := func(user string) []string { - out, err := call(t, s["finetune.list"], "named-user", user, `{}`) - if err != nil { - t.Fatal(err) + It("a job result never carries a filesystem path", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1", OutputDir: "/data/fine-tune/j1", ExportStatus: "completed", ExportModelName: "my-model", Message: "saved to /data/fine-tune/j1/out"}}} + s := specs(ft) + out, _ := call(s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(string(out)).ToNot(ContainSubstring("/data/"), "a result must not leak a server path") + Expect(string(out)).ToNot(ContainSubstring("/models/"), "a result must not leak a server path") + Expect(string(out)).To(ContainSubstring(`"exportedModelName":"my-model"`), "the registered model name is part of the result") + }) + + It("a second user's list does not show the first user's jobs", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}, "c": {ID: "c", UserID: "u1"}}} + s := specs(ft) + ids := func(user string) []string { + out, err := call(s["finetune.list"], "named-user", user, `{}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Jobs []struct { + ID string `json:"id"` + } `json:"jobs"` + } + _ = json.Unmarshal(out, &got) + var r []string + for _, j := range got.Jobs { + r = append(r, j.ID) + } + return r + } + Expect(ids("u2")).To(Equal([]string{"b"}), "user 2 must see only its own job") + Expect(ids("u1")).To(HaveLen(2), "user 1 must see its two jobs") + // A user with no job sees none, not the others'. + Expect(ids("u3")).To(BeEmpty(), "user 3 must see nothing") + }) + }) + + Describe("quantize", func() { + It("start passes the same user key as the route", func() { + q := &fakeQ{} + s := register(nil, q) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","quantizationType":"q8_0"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(q.startUser).To(Equal("u1"), "named user") + _, err = call(s["quantize.start"], "local-owner", "local-owner", `{"model":"org/base"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(q.startUser).To(BeEmpty(), "local owner") + Expect(q.started[0].QuantizationType).To(Equal("q8_0"), "quantization type not passed") + }) + + It("start refuses a local path as model", func() { + s := register(nil, &fakeQ{}) + for _, args := range []string{`{"model":"/models/x.gguf"}`, `{"model":"./x"}`, `{"model":"a\\b"}`} { + _, err := call(s["quantize.start"], "named-user", "u1", args) + Expect(err).To(MatchError(ErrBadRequest), args) + } + }) + + It("start refuses a backend that is not a quantization backend", func() { + s := register(nil, &fakeQ{}) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","backend":"trl"}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) + + It("start refuses a field that is not in the tool schema", func() { + s := register(nil, &fakeQ{}) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","extraOptions":{"x":"y"}}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) + + It("get of another user's job is not found", func() { + q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "other", OutputFile: "/data/q1/m.gguf"}}} + s := register(nil, q) + _, err := call(s["quantize.get"], "named-user", "u1", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) + + It("list returns only the caller's jobs without a path", func() { + q := &fakeQ{jobs: map[string]*schema.QuantizationJob{ + "a": {ID: "a", UserID: "u1", OutputFile: "/data/a/m.gguf", OutputDir: "/data/a"}, + "b": {ID: "b", UserID: "u2"}, + }} + s := register(nil, q) + out, err := call(s["quantize.list"], "named-user", "u1", `{}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).ToNot(ContainSubstring("/data/"), "a caller sees no path") + Expect(string(out)).ToNot(ContainSubstring(`"id":"b"`), "a caller sees only its own jobs") + Expect(string(out)).To(ContainSubstring(`"id":"a"`), "a caller sees its own jobs") + }) + }) + + It("registers no operation for nil services", func() { + b := New(NewInFlight(), always, contract.DefaultLimits()) + Expect(RegisterEngineJobs(b, nil, nil, backends{})).To(Succeed()) + Expect(b.Ops()).To(BeEmpty(), "no service, no operation") + }) + + // Every engine operation is checked against its own feature: refused + // without it, refused with only the other family's feature, and answered + // when the actor holds the operation's feature. + Describe("feature gate", func() { + const fineTune, quant = "fine_tuning", "quantization" + for _, c := range []struct { + op, args, feature, other string + }{ + {"finetune.start", `{"model":"org/base","datasetSource":"org/data"}`, fineTune, quant}, + {"finetune.get", `{"jobId":"j1"}`, fineTune, quant}, + {"finetune.list", `{}`, fineTune, quant}, + {"finetune.checkpoints", `{"jobId":"j1"}`, fineTune, quant}, + {"quantize.start", `{"model":"org/base"}`, quant, fineTune}, + {"quantize.get", `{"jobId":"q1"}`, quant, fineTune}, + {"quantize.list", `{}`, quant, fineTune}, + } { + It(c.op+" needs its feature", func() { + none := func(string, string) bool { return false } + code, out := serveEngine(c.op, c.args, none) + Expect(code).To(Equal(http.StatusForbidden), "without any feature") + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature), "without any feature") + otherOnly := func(_, f string) bool { return f == c.other } + code, out = serveEngine(c.op, c.args, otherOnly) + Expect(code).To(Equal(http.StatusForbidden), "with only %s", c.other) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature), "with only %s", c.other) + ownOnly := func(_, f string) bool { return f == c.feature } + code, out = serveEngine(c.op, c.args, ownOnly) + Expect(code).To(Equal(http.StatusOK), "with %s: %+v", c.feature, out.Error) + }) } - var got struct { - Jobs []struct { - ID string `json:"id"` - } `json:"jobs"` - } - _ = json.Unmarshal(out, &got) - var r []string - for _, j := range got.Jobs { - r = append(r, j.ID) - } - return r - } - if got := ids("u2"); len(got) != 1 || got[0] != "b" { - t.Fatalf("user 2 must see only its own job, got %v", got) - } - if got := ids("u1"); len(got) != 2 { - t.Fatalf("user 1 must see its two jobs, got %v", got) - } - // A user with no job sees none, not the others'. - if got := ids("u3"); len(got) != 0 { - t.Fatalf("user 3 must see nothing, got %v", got) - } -} + }) + + Describe("dataset source", func() { + // A dataset reference that is not an id is refused, the field is + // named, and the service is never reached. + It("must be a gallery or Hugging Face ID", func() { + bad := []string{ + "https://example.com/data.jsonl", // URL + "hf://org/data", // scheme + "file:///etc/passwd", // scheme and absolute path + "/var/data.jsonl", // absolute path + "../data", // parent segment + "org/../data", // parent segment inside an id + "org/..", // parent segment at the end + "org data", // space + "org/data/extra", // more than one separator + `C:\data`, // backslash + "", // empty + } + for _, src := range bad { + ft := &fakeFT{} + s := specs(ft) + args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) + _, err := call(s["finetune.start"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", src) + Expect(err.Error()).To(ContainSubstring("datasetSource"), "%q", src) + Expect(ft.started).To(BeEmpty(), "%q reached the service", src) + } + }) + + // The answer the add-on process receives through the channel. + It("refusal is a 400 naming the field", func() { + code, out := serveEngine("finetune.start", `{"model":"org/base","datasetSource":"https://example.com/d.jsonl"}`, func(string, string) bool { return true }) + Expect(code).To(Equal(http.StatusBadRequest)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostBadRequest)) + Expect(out.Error.Message).To(ContainSubstring("datasetSource")) + }) + + // Keeps the ids the UI uses. + It("accepts Hugging Face and gallery IDs", func() { + for _, src := range []string{"org/data", "HuggingFaceH4/ultrachat_200k", "my-dataset", "org/data-v1.0"} { + ft := &fakeFT{} + s := specs(ft) + args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) + _, err := call(s["finetune.start"], "named-user", "u1", string(args)) + Expect(err).ToNot(HaveOccurred(), "%q", src) + } + }) + + It("accepts a dataset ref as dataset source", func() { + ft := &fakeFT{} + ds := &fakeDatasets{} + b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) + Expect(RegisterEngineJobs(b, ft, nil, backends{})).To(Succeed()) + Expect(RegisterEngineDatasets(b, ds)).To(Succeed()) + raw, err := call(b.specs["finetune.dataset"], "named-user", "u1", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + start := b.specs["finetune.start"] + _, err = call(start, "named-user", "u1", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.started[len(ft.started)-1].DatasetSource).To(Equal(ds.paths[out.DatasetRef]), "the service must get the stored path") + _, err = call(start, "named-user", "u2", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`) + Expect(err).To(MatchError(ErrBadRequest), "another user's reference") + _, err = call(start, "named-user", "u1", `{"model":"org/base","datasetSource":"ffffffff-nope.jsonl"}`) + Expect(err).To(MatchError(ErrBadRequest), "an unknown reference") + }) + }) +}) diff --git a/core/services/addons/hostops/engine_models_test.go b/core/services/addons/hostops/engine_models_test.go index da24dc382..60cde1093 100644 --- a/core/services/addons/hostops/engine_models_test.go +++ b/core/services/addons/hostops/engine_models_test.go @@ -5,16 +5,16 @@ import ( "bytes" "context" "encoding/json" - "errors" "fmt" "net/http" "net/http/httptest" "strings" - "testing" "time" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeExporter struct { @@ -68,19 +68,15 @@ func newExporter(path string) *fakeExporter { return &fakeExporter{owner: map[string]string{"j1": "u1"}, cps: oneCheckpoint, path: path} } -func modelSpecs(t *testing.T, ex Exporter, im Importer) map[string]Spec { - t.Helper() +func modelSpecs(ex Exporter, im Importer) map[string]Spec { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineModels(b, ex, im); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineModels(b, ex, im)).To(Succeed()) return b.specs } // serveAs runs one operation through the channel handler as the given actor, // with a live envelope for it, and returns the status and the answer. -func serveAs(t *testing.T, b *Broker, op, mode, user, args string) (int, contract.HostResponse) { - t.Helper() +func serveAs(b *Broker, op, mode, user, args string) (int, contract.HostResponse) { end := b.inflight.Begin(contract.Envelope{ Request: contract.Request{AddonID: "fine-tuning", OperationID: "op1", Generation: 1, Action: "start", Arguments: json.RawMessage(`{}`)}, Authority: contract.Authority{Actor: contract.Actor{UserID: user, CredentialID: "c1", Mode: mode}, Scope: "apply", ExpiresAt: time.Now().Add(time.Minute)}, @@ -99,214 +95,175 @@ func serveAs(t *testing.T, b *Broker, op, mode, user, args string) (int, contrac // adminBroker builds a broker where every user holds the feature, and only the // user "admin" is an administrator. -func adminBroker(t *testing.T, ex Exporter, im Importer) *Broker { - t.Helper() +func adminBroker(ex Exporter, im Importer) *Broker { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) b.SetAdminCheck(func(actor contract.Actor) bool { return actor.UserID == "admin" }) - if err := RegisterEngineModels(b, ex, im); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineModels(b, ex, im)).To(Succeed()) return b } -func TestExportRegistersAModelAndReturnsOnlyItsName(t *testing.T) { - ex := newExporter("/models/my-model") - s := modelSpecs(t, ex, nil) - out, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1","name":"my-model","exportFormat":"lora"}`) - if err != nil { - t.Fatal(err) - } - if string(out) != `{"model":"my-model"}` { - t.Fatalf("got %s", out) - } - if strings.Contains(string(out), "/") { - t.Fatalf("result carries a path: %s", out) - } - if ex.user != "u1" || ex.req.Name != "my-model" || ex.req.ExportFormat != "lora" { - t.Fatalf("service got user=%q req=%+v", ex.user, ex.req) - } -} +var _ = Describe("Model export and import operations", func() { + It("export registers a model and returns only its name", func() { + ex := newExporter("/models/my-model") + s := modelSpecs(ex, nil) + out, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1","name":"my-model","exportFormat":"lora"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).To(Equal(`{"model":"my-model"}`)) + Expect(string(out)).ToNot(ContainSubstring("/"), "result carries a path") + Expect(ex.user).To(Equal("u1")) + Expect(ex.req.Name).To(Equal("my-model")) + Expect(ex.req.ExportFormat).To(Equal("lora")) + }) -func TestExportUsesTheHighestStepCheckpoint(t *testing.T) { - ex := newExporter("/models/my-model") - ex.cps = []Checkpoint{{Step: 100, Path: "/out/ckpt-100"}, {Step: 300, Path: "/out/ckpt-300"}, {Step: 200, Path: "/out/ckpt-200"}} - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); err != nil { - t.Fatal(err) - } - if ex.req.CheckpointPath != "/out/ckpt-300" { - t.Fatalf("export got checkpoint %q, want the highest step", ex.req.CheckpointPath) - } -} + It("export uses the highest step checkpoint", func() { + ex := newExporter("/models/my-model") + ex.cps = []Checkpoint{{Step: 100, Path: "/out/ckpt-100"}, {Step: 300, Path: "/out/ckpt-300"}, {Step: 200, Path: "/out/ckpt-200"}} + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ex.req.CheckpointPath).To(Equal("/out/ckpt-300"), "want the highest step") + }) -func TestExportOfAJobWithNoCheckpointIsAConflictNamingIt(t *testing.T) { - ex := newExporter("/models/my-model") - ex.cps = nil - s := modelSpecs(t, ex, nil) - _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) - if !errors.Is(err, ErrConflict) || !strings.Contains(err.Error(), "checkpoint") { - t.Fatalf("got %v, want a conflict naming the checkpoint", err) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a job with no checkpoint", ex.calls) - } -} + It("export of a job with no checkpoint is a conflict naming it", func() { + ex := newExporter("/models/my-model") + ex.cps = nil + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrConflict)) + Expect(err.Error()).To(ContainSubstring("checkpoint")) + Expect(ex.calls).To(Equal(0), "service called for a job with no checkpoint") + }) -func TestExportRefusesABadName(t *testing.T) { - ex := newExporter("/models/x") - s := modelSpecs(t, ex, nil) - long := strings.Repeat("a", 65) - for _, name := range []string{"../x", "a/b", " ", long} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "name": name}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("export refuses a bad name", func() { + ex := newExporter("/models/x") + s := modelSpecs(ex, nil) + long := strings.Repeat("a", 65) + for _, name := range []string{"../x", "a/b", " ", long} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "name": name}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", name) } - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a refused name", ex.calls) - } -} + Expect(ex.calls).To(Equal(0), "service called for a refused name") + }) -func TestExportRefusesAnExportFormatThatIsNotListed(t *testing.T) { - ex := newExporter("/models/x") - s := modelSpecs(t, ex, nil) - for _, f := range []string{"adapter", "merged_8bit", "GGUF", "gguf "} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", f, err) + It("export refuses an export format that is not listed", func() { + ex := newExporter("/models/x") + s := modelSpecs(ex, nil) + for _, f := range []string{"adapter", "merged_8bit", "GGUF", "gguf "} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", f) } - } - for _, f := range []string{"lora", "merged_16bit", "merged_4bit", "gguf"} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); err != nil { - t.Fatalf("%q: %v", f, err) + for _, f := range []string{"lora", "merged_16bit", "merged_4bit", "gguf"} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).ToNot(HaveOccurred(), "%q", f) } - } -} + }) -func TestExportOfAnotherUsersJobIsNotFound(t *testing.T) { - ex := newExporter("/models/x") - ex.owner = map[string]string{"j1": "u2"} - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} + It("export of another user's job is not found", func() { + ex := newExporter("/models/x") + ex.owner = map[string]string{"j1": "u2"} + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) -func TestExportRunsOncePerOperationID(t *testing.T) { - ex := newExporter("/models/my-model") - b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineModels(b, ex, nil); err != nil { - t.Fatal(err) - } - spec := b.specs["finetune.export"] - c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} - first, err := b.run(context.Background(), spec, c) - if err != nil { - t.Fatal(err) - } - second, err := b.run(context.Background(), spec, c) - if err != nil { - t.Fatal(err) - } - if ex.calls != 1 || string(first) != string(second) { - t.Fatalf("calls=%d first=%s second=%s", ex.calls, first, second) - } -} + It("export runs once per operation ID", func() { + ex := newExporter("/models/my-model") + b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) + Expect(RegisterEngineModels(b, ex, nil)).To(Succeed()) + spec := b.specs["finetune.export"] + c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} + first, err := b.run(context.Background(), spec, c) + Expect(err).ToNot(HaveOccurred()) + second, err := b.run(context.Background(), spec, c) + Expect(err).ToNot(HaveOccurred()) + Expect(ex.calls).To(Equal(1)) + Expect(string(first)).To(Equal(string(second))) + }) -func TestExportRefusesAReturnedPathThatIsNotAModelName(t *testing.T) { - ex := newExporter("/models/bad name") - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); err == nil || errors.Is(err, ErrBadRequest) { - t.Fatalf("want an outage error, got %v", err) - } -} + It("export refuses a returned path that is not a model name", func() { + ex := newExporter("/models/bad name") + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(HaveOccurred(), "want an outage error") + Expect(err).ToNot(MatchError(ErrBadRequest), "want an outage error") + }) -func TestImportRegistersTheQuantizedModelAndReturnsOnlyItsName(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - s := modelSpecs(t, nil, im) - out, err := call(t, s["quantize.import"], "named-user", "u1", `{"jobId":"q1","name":"q-model"}`) - if err != nil { - t.Fatal(err) - } - if string(out) != `{"model":"q-model"}` { - t.Fatalf("got %s", out) - } -} + It("import registers the quantized model and returns only its name", func() { + im := &fakeImporter{path: "/models/q-model"} + s := modelSpecs(nil, im) + out, err := call(s["quantize.import"], "named-user", "u1", `{"jobId":"q1","name":"q-model"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).To(Equal(`{"model":"q-model"}`)) + }) -func TestImportOfAJobThatIsNotCompletedIsAConflict(t *testing.T) { - im := &fakeImporter{err: fmt.Errorf("%w: job q1 is not completed", ErrConflict)} - s := modelSpecs(t, nil, im) - if _, err := call(t, s["quantize.import"], "named-user", "u1", `{"jobId":"q1"}`); !errors.Is(err, ErrConflict) { - t.Fatalf("got %v", err) - } -} + It("import of a job that is not completed is a conflict", func() { + im := &fakeImporter{err: fmt.Errorf("%w: job q1 is not completed", ErrConflict)} + s := modelSpecs(nil, im) + _, err := call(s["quantize.import"], "named-user", "u1", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrConflict)) + }) -func TestModelOperationsAreApplyAndIdempotentAndAdminOnly(t *testing.T) { - s := modelSpecs(t, newExporter("/models/x"), &fakeImporter{}) - for _, op := range []string{"finetune.export", "quantize.import"} { - if s[op].Scope != "apply" || !s[op].Idempotent || !s[op].Admin { - t.Fatalf("%s: scope=%q idempotent=%v admin=%v", op, s[op].Scope, s[op].Idempotent, s[op].Admin) + It("model operations are apply, idempotent and admin only", func() { + s := modelSpecs(newExporter("/models/x"), &fakeImporter{}) + for _, op := range []string{"finetune.export", "quantize.import"} { + Expect(s[op].Scope).To(Equal("apply"), op) + Expect(s[op].Idempotent).To(BeTrue(), op) + Expect(s[op].Admin).To(BeTrue(), op) } - } - if s["finetune.export"].Feature != featureFineTuning || s["quantize.import"].Feature != featureQuantization { - t.Fatalf("features: %q %q", s["finetune.export"].Feature, s["quantize.import"].Feature) - } -} + Expect(s["finetune.export"].Feature).To(Equal(featureFineTuning)) + Expect(s["quantize.import"].Feature).To(Equal(featureQuantization)) + }) -func TestExportIsRefusedForAFeatureHoldingNonAdmin(t *testing.T) { - ex := newExporter("/models/my-model") - b := adminBroker(t, ex, nil) - code, out := serveAs(t, b, "finetune.export", "named-user", "u1", `{"jobId":"j1"}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("%d %+v", code, out.Error) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a non-admin", ex.calls) - } -} + It("export is refused for a feature-holding non-admin", func() { + ex := newExporter("/models/my-model") + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "named-user", "u1", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(ex.calls).To(Equal(0), "service called for a non-admin") + }) -func TestExportIsAllowedForAnAdmin(t *testing.T) { - ex := newExporter("/models/my-model") - ex.owner = map[string]string{"j1": "admin"} - b := adminBroker(t, ex, nil) - code, out := serveAs(t, b, "finetune.export", "named-user", "admin", `{"jobId":"j1"}`) - if code != http.StatusOK || out.Error != nil { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("export is allowed for an admin", func() { + ex := newExporter("/models/my-model") + ex.owner = map[string]string{"j1": "admin"} + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "named-user", "admin", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(out.Error).To(BeNil()) + }) -func TestImportIsRefusedForAFeatureHoldingNonAdmin(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - b := adminBroker(t, nil, im) - code, out := serveAs(t, b, "quantize.import", "named-user", "u1", `{"jobId":"q1"}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("%d %+v", code, out.Error) - } - if im.calls != 0 { - t.Fatalf("service called %d times for a non-admin", im.calls) - } -} + It("import is refused for a feature-holding non-admin", func() { + im := &fakeImporter{path: "/models/q-model"} + b := adminBroker(nil, im) + code, out := serveAs(b, "quantize.import", "named-user", "u1", `{"jobId":"q1"}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(im.calls).To(Equal(0), "service called for a non-admin") + }) -func TestImportIsAllowedForAnAdmin(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - b := adminBroker(t, nil, im) - code, out := serveAs(t, b, "quantize.import", "named-user", "admin", `{"jobId":"q1"}`) - if code != http.StatusOK || out.Error != nil { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("import is allowed for an admin", func() { + im := &fakeImporter{path: "/models/q-model"} + b := adminBroker(nil, im) + code, out := serveAs(b, "quantize.import", "named-user", "admin", `{"jobId":"q1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(out.Error).To(BeNil()) + }) -// The broker does not shortcut the local owner: the admin check decides in -// every mode, so a local-owner actor the check refuses is refused. -func TestTheAdminCheckDecidesInEveryMode(t *testing.T) { - ex := newExporter("/models/my-model") - ex.owner = map[string]string{"j1": "u1"} - b := adminBroker(t, ex, nil) - if code, out := serveAs(t, b, "finetune.export", "local-owner", "u1", `{"jobId":"j1"}`); code != http.StatusForbidden || out.Error.Code != contract.HostFeature { - t.Fatalf("local-owner actor the check refuses: %d %+v", code, out.Error) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a refused actor", ex.calls) - } -} + // The broker does not shortcut the local owner: the admin check decides in + // every mode, so a local-owner actor the check refuses is refused. + It("lets the admin check decide in every mode", func() { + ex := newExporter("/models/my-model") + ex.owner = map[string]string{"j1": "u1"} + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "local-owner", "u1", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusForbidden), "local-owner actor the check refuses") + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(ex.calls).To(Equal(0), "service called for a refused actor") + }) +}) diff --git a/core/services/addons/hostops/engine_progress_test.go b/core/services/addons/hostops/engine_progress_test.go index 542174cdd..914d984fb 100644 --- a/core/services/addons/hostops/engine_progress_test.go +++ b/core/services/addons/hostops/engine_progress_test.go @@ -5,14 +5,14 @@ import ( "bytes" "context" "encoding/json" - "errors" "net/http" "net/http/httptest" "sync" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) // fakeSource stands in for a job service. Events are fed by the test through @@ -76,160 +76,125 @@ func (f *fakeSource) stopCount() int { return len(f.stops) } -func registerProgress(t *testing.T, src ProgressSource) (*Broker, map[string]Spec) { - t.Helper() +func registerProgress(src ProgressSource) (*Broker, map[string]Spec) { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) p := NewProgresser(src) - if err := RegisterEngineProgress(b, p, p); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineProgress(b, p, p)).To(Succeed()) return b, b.specs } -func TestProgressReturnsTheNextEventWithinTheWait(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "training", Percent: 40} - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Percent float32 `json:"percent"` - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if got.Percent != 40 || got.Done { - t.Fatalf("want percent 40 and not done: %s", out) - } -} +var _ = Describe("Progress and stop operations", func() { + var ( + src *fakeSource + b *Broker + s map[string]Spec + ) -func TestProgressWithNoEventReturnsTheSnapshotWhenTheWaitEnds(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - start := time.Now() - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":1}`) - elapsed := time.Since(start) - if err != nil { - t.Fatal(err) - } - if elapsed < 900*time.Millisecond || elapsed >= 3*time.Second { - t.Fatalf("the call should wait about one second, took %s", elapsed) - } - var got struct { - Status string `json:"status"` - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if got.Status != "training" || got.Done { - t.Fatalf("want the current status of the job: %s", out) - } -} + BeforeEach(func() { + src = newFakeSource() + b, s = registerProgress(src) + }) -func TestProgressWaitIsCapped(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "training"} - if _, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":600}`); err != nil { - t.Fatal(err) - } - src.mu.Lock() - left := time.Until(src.deadline) - src.mu.Unlock() - if left <= 24*time.Second || left > 25*time.Second+500*time.Millisecond { - t.Fatalf("the stream should run under a wait of 25 s, got about %s", left) - } -} - -func TestProgressReportsDoneForAFinalStatus(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "completed"} - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if !got.Done { - t.Fatalf("a completed job is done: %s", out) - } -} - -func TestProgressOfAnotherUsersJobIsNotFound(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["finetune.progress"], "named-user", "u2", `{"jobId":"j1","waitSeconds":1}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } - if _, err := call(t, s["quantize.progress"], "named-user", "u2", `{"jobId":"q1","waitSeconds":1}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("quantize: got %v", err) - } -} - -func TestStopIsIdempotentPerOperation(t *testing.T) { - src := newFakeSource() - b, s := registerProgress(t, src) - c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} - for range 2 { - if _, err := b.run(context.Background(), s["finetune.stop"], c); err != nil { - t.Fatal(err) + It("returns the next event within the wait", func() { + src.events <- Progress{JobID: "j1", Status: "training", Percent: 40} + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Percent float32 `json:"percent"` + Done bool `json:"done"` } - } - if n := src.stopCount(); n != 1 { - t.Fatalf("two calls with one operation ID must stop once, stopped %d times", n) - } -} + _ = json.Unmarshal(out, &got) + Expect(got.Percent).To(BeEquivalentTo(40), "want percent 40: %s", out) + Expect(got.Done).To(BeFalse(), "want not done: %s", out) + }) -func TestStopKeepsTheCheckpointOnlyWhenAsked(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["finetune.stop"], "named-user", "u1", `{"jobId":"j1","saveCheckpoint":true}`); err != nil { - t.Fatal(err) - } - if _, err := call(t, s["finetune.stop"], "named-user", "u1", `{"jobId":"j1"}`); err != nil { - t.Fatal(err) - } - if len(src.stops) != 2 || !src.stops[0] || src.stops[1] { - t.Fatalf("saveCheckpoint must reach StopJob only when true: %v", src.stops) - } -} + It("returns the snapshot when the wait ends with no event", func() { + start := time.Now() + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":1}`) + elapsed := time.Since(start) + Expect(err).ToNot(HaveOccurred()) + Expect(elapsed).To(BeNumerically(">=", 900*time.Millisecond), "the call should wait about one second") + Expect(elapsed).To(BeNumerically("<", 3*time.Second), "the call should wait about one second") + var got struct { + Status string `json:"status"` + Done bool `json:"done"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Status).To(Equal("training"), "want the current status of the job: %s", out) + Expect(got.Done).To(BeFalse(), "want the current status of the job: %s", out) + }) -func TestStopOfAnotherUsersJobIsNotFound(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["quantize.stop"], "named-user", "u2", `{"jobId":"q1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } - if src.stopCount() != 0 { - t.Fatal("another user's job must not be stopped") - } -} + It("caps the wait", func() { + src.events <- Progress{JobID: "j1", Status: "training"} + _, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":600}`) + Expect(err).ToNot(HaveOccurred()) + src.mu.Lock() + left := time.Until(src.deadline) + src.mu.Unlock() + Expect(left).To(BeNumerically(">", 24*time.Second), "the stream should run under a wait of 25 s") + Expect(left).To(BeNumerically("<=", 25*time.Second+500*time.Millisecond), "the stream should run under a wait of 25 s") + }) -func TestACancelledContextEndsTheWait(t *testing.T) { - src := newFakeSource() - b, _ := registerProgress(t, src) - in := b.inflight - end := in.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - h := b.Handler("fine-tuning", "secret", []string{"finetune."}) + It("reports done for a final status", func() { + src.events <- Progress{JobID: "j1", Status: "completed"} + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Done bool `json:"done"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Done).To(BeTrue(), "a completed job is done: %s", out) + }) - body, _ := json.Marshal(contract.HostRequest{OperationID: "op1", Op: "finetune.progress", Arguments: json.RawMessage(`{"jobId":"j1","waitSeconds":25}`)}) - ctx, cancel := context.WithCancel(context.Background()) - cancel() - req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body)).WithContext(ctx) - req.Header.Set("X-Addon-Host-Credential", "secret") - rec := httptest.NewRecorder() - start := time.Now() - h.ServeHTTP(rec, req) - if d := time.Since(start); d > time.Second { - t.Fatalf("a cancelled call must end at once, took %s", d) - } - var out contract.HostResponse - _ = json.Unmarshal(rec.Body.Bytes(), &out) - if rec.Code != http.StatusServiceUnavailable || out.Error == nil || out.Error.Code != contract.HostUnavailable { - t.Fatalf("want 503 unavailable, got %d %+v", rec.Code, out.Error) - } -} + It("does not find the progress of another user's job", func() { + _, err := call(s["finetune.progress"], "named-user", "u2", `{"jobId":"j1","waitSeconds":1}`) + Expect(err).To(MatchError(ErrNotFound)) + _, err = call(s["quantize.progress"], "named-user", "u2", `{"jobId":"q1","waitSeconds":1}`) + Expect(err).To(MatchError(ErrNotFound), "quantize") + }) + + It("stops once per operation", func() { + c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} + for range 2 { + _, err := b.run(context.Background(), s["finetune.stop"], c) + Expect(err).ToNot(HaveOccurred()) + } + Expect(src.stopCount()).To(Equal(1), "two calls with one operation ID must stop once") + }) + + It("keeps the checkpoint only when asked", func() { + _, err := call(s["finetune.stop"], "named-user", "u1", `{"jobId":"j1","saveCheckpoint":true}`) + Expect(err).ToNot(HaveOccurred()) + _, err = call(s["finetune.stop"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(src.stops).To(Equal([]bool{true, false}), "saveCheckpoint must reach StopJob only when true") + }) + + It("does not find the stop of another user's job", func() { + _, err := call(s["quantize.stop"], "named-user", "u2", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrNotFound)) + Expect(src.stopCount()).To(Equal(0), "another user's job must not be stopped") + }) + + It("ends the wait when the context is cancelled", func() { + in := b.inflight + end := in.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + h := b.Handler("fine-tuning", "secret", []string{"finetune."}) + + body, _ := json.Marshal(contract.HostRequest{OperationID: "op1", Op: "finetune.progress", Arguments: json.RawMessage(`{"jobId":"j1","waitSeconds":25}`)}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body)).WithContext(ctx) + req.Header.Set("X-Addon-Host-Credential", "secret") + rec := httptest.NewRecorder() + start := time.Now() + h.ServeHTTP(rec, req) + Expect(time.Since(start)).To(BeNumerically("<=", time.Second), "a cancelled call must end at once") + var out contract.HostResponse + _ = json.Unmarshal(rec.Body.Bytes(), &out) + Expect(rec.Code).To(Equal(http.StatusServiceUnavailable)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostUnavailable)) + }) +}) diff --git a/core/services/addons/hostops/registry_test.go b/core/services/addons/hostops/registry_test.go index eea1612a0..609e2a34a 100644 --- a/core/services/addons/hostops/registry_test.go +++ b/core/services/addons/hostops/registry_test.go @@ -5,81 +5,80 @@ import ( "context" "encoding/json" "net/http" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) func noop(context.Context, Call) (json.RawMessage, error) { return nil, nil } -func TestRegisterRefusesWhatIsMalformed(t *testing.T) { - b := New(NewInFlight(), nil, contract.DefaultLimits()) - if err := b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}); err != nil { - t.Fatal(err) - } - for name, s := range map[string]Spec{ - "duplicate": {Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}, - "bad name": {Op: "Finetune", Scope: "read", Feature: "fine_tuning", Handler: noop}, - "bad scope": {Op: "finetune.x", Scope: "preview", Feature: "fine_tuning", Handler: noop}, - "no feature": {Op: "finetune.y", Scope: "read", Handler: noop}, - "no handler": {Op: "finetune.z", Scope: "read", Feature: "fine_tuning"}, - } { - if b.Register(s) == nil { - t.Errorf("%s: accepted", name) +var _ = Describe("Operation registry", func() { + It("refuses what is malformed", func() { + b := New(NewInFlight(), nil, contract.DefaultLimits()) + Expect(b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop})).To(Succeed()) + for name, s := range map[string]Spec{ + "duplicate": {Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}, + "bad name": {Op: "Finetune", Scope: "read", Feature: "fine_tuning", Handler: noop}, + "bad scope": {Op: "finetune.x", Scope: "preview", Feature: "fine_tuning", Handler: noop}, + "no feature": {Op: "finetune.y", Scope: "read", Handler: noop}, + "no handler": {Op: "finetune.z", Scope: "read", Feature: "fine_tuning"}, + } { + Expect(b.Register(s)).ToNot(Succeed(), "%s: accepted", name) } - } -} + }) -func TestAnUnregisteredOperationUnderAnAllowedPrefixIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.nothing", `{}`); code != http.StatusForbidden || out.Error.Code != "unknown_op" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an unregistered operation under an allowed prefix", func() { + f := newFixture(always) + end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.nothing", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("unknown_op")) + }) -func TestABrokerWithoutAFeatureCheckRefusesEverything(t *testing.T) { - f := newFixture(t, nil) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusForbidden || out.Error.Code != "feature" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses everything for a broker without a feature check", func() { + f := newFixture(nil) + end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("feature")) + }) -func TestTheLocalOwnerNeedsNoFeatureCheck(t *testing.T) { - f := newFixture(t, func(string, string) bool { return false }) - e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) - e.Authority.Actor.Mode = "local-owner" - end := f.inflight.Begin(e) - defer end() - if code, _ := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusOK { - t.Fatalf("%d", code) - } -} + It("needs no feature check for the local owner", func() { + f := newFixture(func(string, string) bool { return false }) + e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) + e.Authority.Actor.Mode = "local-owner" + end := f.inflight.Begin(e) + DeferCleanup(end) + code, _ := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusOK)) + }) -func TestAnotherAddonCannotUseTheOperationOfALiveRequest(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("quantization", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusForbidden || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("does not let another add-on use the operation of a live request", func() { + f := newFixture(always) + end := f.inflight.Begin(env("quantization", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestTheEnvelopeIsGoneWhenTheRequestEnds(t *testing.T) { - in := NewInFlight() - e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) - end1, end2 := in.Begin(e), in.Begin(e) - end1() - end1() - if _, ok := in.Lookup("fine-tuning", "op1"); !ok { - t.Fatal("the second request still runs") - } - end2() - if _, ok := in.Lookup("fine-tuning", "op1"); ok { - t.Fatal("the envelope stayed") - } -} + It("drops the envelope when the request ends", func() { + in := NewInFlight() + e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) + end1, end2 := in.Begin(e), in.Begin(e) + end1() + end1() + _, ok := in.Lookup("fine-tuning", "op1") + Expect(ok).To(BeTrue(), "the second request still runs") + end2() + _, ok = in.Lookup("fine-tuning", "op1") + Expect(ok).To(BeFalse(), "the envelope stayed") + }) +}) diff --git a/core/services/addons/hostops/suite_test.go b/core/services/addons/hostops/suite_test.go new file mode 100644 index 000000000..70cc5fcff --- /dev/null +++ b/core/services/addons/hostops/suite_test.go @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +package hostops + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestHostops(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Add-on host operations suite") +}