Add SetTag endpoint

magefiles/generate/endpoints.go

@@ -26,10 +26,10 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			"logMetric",
 			// "logParam",
 			// "setExperimentTag",
-			// "setTag",
+			"setTag",
 			// "setTraceTag",
 			// "deleteTraceTag",
-			// "deleteTag",
+			"deleteTag",
 			"searchRuns",
 			// "listArtifacts",
 			// "getMetricHistory",

magefiles/generate/validations.go

@@ -25,4 +25,8 @@ var validations = map[string]string{
 	"LogMetric_Key":                      "required",
 	"LogMetric_Value":                    "required",
 	"LogMetric_Timestamp":                "required",
+	"SetTag_Key":                         "required,max=1000,validMetricParamOrTagName,pathIsUnique",
+	"SetTag_Value":                       "required,max=8000",
+	"DeleteTag_RunId":                    "required",
+	"DeleteTag_Key":                      "required",
 }

magefiles/tests.go

@@ -28,8 +28,7 @@ func cleanUpMemoryFile() error {
 	return nil
 }
 
-// Run mlflow Python tests against the Go backend.
-func (Test) Python() error {
+func runPythonTests(testFiles []string, testName string) error {
 	libpath, err := os.MkdirTemp("", "")
 	if err != nil {
 		return err
@@ -45,17 +44,22 @@ func (Test) Python() error {
 		return nil
 	}
 
+	args := []string{
+		"--confcutdir=.",
+	}
+	args = append(args, testFiles...)
+
+	// Add testName filter if provided
+	if testName != "" {
+		args = append(args, "-k", testName)
+	} else {
+		args = append(args, "-k", "not [file")
+	}
+
 	//  Run the tests (currently just the server ones)
 	if err := sh.RunWithV(map[string]string{
 		"MLFLOW_GO_LIBRARY_PATH": libpath,
-	}, "pytest",
-		"--confcutdir=.",
-		".mlflow.repo/tests/tracking/test_rest_tracking.py",
-		".mlflow.repo/tests/tracking/test_model_registry.py",
-		".mlflow.repo/tests/store/tracking/test_sqlalchemy_store.py",
-		".mlflow.repo/tests/store/model_registry/test_sqlalchemy_store.py",
-		"-k",
-		"not [file",
+	}, "pytest", args...,
 		// "-vv",
 	); err != nil {
 		return err
@@ -64,6 +68,23 @@ func (Test) Python() error {
 	return nil
 }
 
+// Run mlflow Python tests against the Go backend.
+func (Test) Python() error {
+	return runPythonTests([]string{
+		".mlflow.repo/tests/tracking/test_rest_tracking.py",
+		".mlflow.repo/tests/tracking/test_model_registry.py",
+		".mlflow.repo/tests/store/tracking/test_sqlalchemy_store.py",
+		".mlflow.repo/tests/store/model_registry/test_sqlalchemy_store.py",
+	}, "")
+}
+
+// Run specific Python test against the Go backend.
+func (Test) PythonSpecific(testName string) error {
+	return runPythonTests([]string{
+		".mlflow.repo/tests/tracking/test_rest_tracking.py",
+	}, testName)
+}
+
 // Run the Go unit tests.
 func (Test) Unit() error {
 	return sh.RunV("go", "test", "./pkg/...")

mlflow_go/store/tracking.py

@@ -14,6 +14,7 @@
     CreateRun,
     DeleteExperiment,
     DeleteRun,
+    DeleteTag,
     GetExperiment,
     GetExperimentByName,
     GetRun,
@@ -22,6 +23,7 @@
     RestoreExperiment,
     RestoreRun,
     SearchRuns,
+    SetTag,
     UpdateExperiment,
     UpdateRun,
 )
@@ -165,6 +167,13 @@ def log_metric(self, run_id, metric):
         )
         self.service.call_endpoint(get_lib().TrackingServiceLogMetric, request)
 
+    def set_tag(self, run_id, tag):
+        request = SetTag(run_id=run_id, key=tag.key, value=tag.value)
+        self.service.call_endpoint(get_lib().TrackingServiceSetTag, request)
+
+    def delete_tag(self, run_id, key):
+        request = DeleteTag(run_id=run_id, key=key)
+        self.service.call_endpoint(get_lib().TrackingServiceDeleteTag, request)
 
 def TrackingStore(cls):
     return type(cls.__name__, (_TrackingStore, cls), {})

pkg/tracking/service/tags.go

@@ -0,0 +1,30 @@
+package service
+
+import (
+	"context"
+
+	"github.com/mlflow/mlflow-go/pkg/contract"
+	"github.com/mlflow/mlflow-go/pkg/protos"
+)
+
+func (ts TrackingService) SetTag(ctx context.Context, input *protos.SetTag) (*protos.SetTag_Response, *contract.Error) {
+	runID := input.GetRunId()
+
+	if runID == "" {
+		runID = input.GetRunUuid()
+	}
+
+	if err := ts.Store.SetTag(ctx, runID, input.GetKey(), input.GetValue()); err != nil {
+		return nil, err
+	}
+
+	return &protos.SetTag_Response{}, nil
+}
+
+func (ts TrackingService) DeleteTag(ctx context.Context, input *protos.DeleteTag) (*protos.DeleteTag_Response, *contract.Error) {
+  if err := ts.Store.DeleteTag(ctx, input.GetRunId(), input.GetKey()); err != nil {
+		return nil, err
+	}
+
+	return &protos.DeleteTag_Response{}, nil
+}