logInputs endpoint

.github/workflows/test.yml

@@ -52,7 +52,7 @@ jobs:
       - name: Initialize MLflow repo
         run: mage repo:init
       - name: Install dependencies
-        run: pip install pytest==8.1.1 psycopg2-binary -e .mlflow.repo
+        run: pip install pytest==8.1.1 psycopg2-binary==2.9.9 -e .mlflow.repo
       - name: Run integration tests
         run: mage test:python
         # Temporary workaround for failing tests

magefiles/generate/endpoints.go

@@ -37,7 +37,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			// "getMetricHistoryBulkInterval",
 			"logBatch",
 			// "logModel",
-			// "logInputs",
+			"logInputs",
 			// "startTrace",
 			// "endTrace",
 			// "getTraceInfo",

magefiles/generate/validations.go

@@ -34,4 +34,15 @@ var validations = map[string]string{
 	"SetExperimentTag_ExperimentId":      "required",
 	"SetExperimentTag_Key":               "required,max=250,validMetricParamOrTagName",
 	"SetExperimentTag_Value":             "max=5000",
+	"LogInputs_RunId":                    "required,runId",
+	"LogInputs_Datasets":                 "required",
+	"DatasetInput_Dataset":               "required",
+	"Dataset_Name":                       "required,max=500",
+	"Dataset_Digest":                     "required,max=36",
+	"Dataset_SourceType":                 "required",
+	"Dataset_Source":                     "required,max=65535",
+	"Dataset_Profile":                    "max:16777215",
+	"Dataset_Schema":                     "max:1048575",
+	"InputTag_Key":                       "required,max=255",
+	"InputTag_Value":                     "required,max=500",
 }

pkg/entities/dataset.go

@@ -33,3 +33,14 @@ func (d *Dataset) ToProto() *protos.Dataset {
 		Profile:    profile,
 	}
 }
+
+func NewDatasetFromProto(proto *protos.Dataset) *Dataset {
+	return &Dataset{
+		Name:       proto.GetName(),
+		Digest:     proto.GetDigest(),
+		SourceType: proto.GetSourceType(),
+		Source:     proto.GetSource(),
+		Schema:     proto.GetSchema(),
+		Profile:    proto.GetProfile(),
+	}
+}

pkg/entities/dataset_input.go

@@ -18,3 +18,15 @@ func (ds DatasetInput) ToProto() *protos.DatasetInput {
 		Dataset: ds.Dataset.ToProto(),
 	}
 }
+
+func NewDatasetInputFromProto(proto *protos.DatasetInput) *DatasetInput {
+	tags := make([]*InputTag, 0, len(proto.GetTags()))
+	for _, t := range proto.GetTags() {
+		tags = append(tags, NewInputTagFromProto(t))
+	}
+
+	return &DatasetInput{
+		Dataset: NewDatasetFromProto(proto.GetDataset()),
+		Tags:    tags,
+	}
+}

pkg/entities/input_tag.go

@@ -13,3 +13,10 @@ func (i InputTag) ToProto() *protos.InputTag {
 		Value: &i.Value,
 	}
 }
+
+func NewInputTagFromProto(proto *protos.InputTag) *InputTag {
+	return &InputTag{
+		Key:   proto.GetKey(),
+		Value: proto.GetValue(),
+	}
+}

pkg/tracking/service/inputs.go

@@ -0,0 +1,28 @@
+package service
+
+import (
+	"context"
+
+	"github.com/mlflow/mlflow-go/pkg/contract"
+	"github.com/mlflow/mlflow-go/pkg/entities"
+	"github.com/mlflow/mlflow-go/pkg/protos"
+)
+
+func (ts TrackingService) LogInputs(
+	ctx context.Context, input *protos.LogInputs,
+) (*protos.LogInputs_Response, *contract.Error) {
+	if len(input.GetDatasets()) == 0 {
+		return &protos.LogInputs_Response{}, nil
+	}
+
+	datasets := make([]*entities.DatasetInput, 0, len(input.GetDatasets()))
+	for _, d := range input.GetDatasets() {
+		datasets = append(datasets, entities.NewDatasetInputFromProto(d))
+	}
+
+	if err := ts.Store.LogInputs(ctx, input.GetRunId(), datasets); err != nil {
+		return nil, err
+	}
+
+	return &protos.LogInputs_Response{}, nil
+}