Add SetTag endpoint

.devcontainer/devcontainer.json

@@ -59,8 +59,8 @@
 	// "forwardPorts": [5432],
 
 	// Use 'postCreateCommand' to run commands after the container is created.
-	"postCreateCommand": ".devcontainer/postCreate.sh"
+	"postCreateCommand": ".devcontainer/postCreate.sh",
 
 	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
-	// "remoteUser": "root"
+	"remoteUser": "root"
 }

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_RunId":												"required",
+	"SetTag_Key":													"required",
+	"DeleteTag_RunId":										"required",
+	"DeleteTag_Key":											"required",
 }

magefiles/tests.go

@@ -64,6 +64,33 @@ func (Test) Python() error {
 	return nil
 }
 
+// Run specific Python test against the Go backend.
+func (Test) PythonSpecific(testName string) error {
+	libpath, err := os.MkdirTemp("", "")
+	if err != nil {
+		return err
+	}
+
+	defer os.RemoveAll(libpath)
+	defer cleanUpMemoryFile()
+
+	if err := sh.RunV("python", "-m", "mlflow_go.lib", ".", libpath); err != nil {
+		return nil
+	}
+
+	if err := sh.RunWithV(map[string]string{
+		"MLFLOW_GO_LIBRARY_PATH": libpath,
+	}, "pytest",
+		"--confcutdir=.",
+		".mlflow.repo/tests/tracking/test_rest_tracking.py",
+		"-k", testName,
+	); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // 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,27 @@
+package service
+
+import (
+	"context"
+	"fmt"
+
+	"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) {
+	// Print input
+	fmt.Println(input)
+	if err := ts.Store.SetTag(ctx, input.GetRunId(), 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
+}