Add SetTag endpoint

CONTRIBUTING.md

@@ -8,7 +8,7 @@ to configure all the development environment just run `mage` target:
 mage configure
 ```
 
-it will configure mlflow and all the Python dependencies required by the project or run each step manually:
+it will configure MLflow and all the Python dependencies required by the project or run each step manually:
 
 ```bash
 # Install our Python package and its dependencies
@@ -17,7 +17,7 @@ pip install -e .
 # Install the dreaded psycho
 pip install psycopg2-binary
 
-# Archive the MLFlow pre-built UI
+# Archive the MLflow pre-built UI
 tar -C /usr/local/python/current/lib/python3.8/site-packages/mlflow -czvf ./ui.tgz ./server/js/build
 
 # Clone the MLflow repo
@@ -30,7 +30,7 @@ tar -C .mlflow.repo/mlflow -xzvf ./ui.tgz
 pip install -e .mlflow.repo
 ```
 
-## Run Go Mlflow server
+## Run Go MLflow server
 
 To start the mlflow-go dev server connecting to postgres just run next `mage` target:
 
@@ -61,10 +61,18 @@ mage test:all
 ```
 
 ```bash
-# Run just MLFlow Python tests
+# Run just MLflow Python tests
 mage test:python
 ```
 
+```bash
+# Run specific MLflow Python tests (matches all tests containing the argument)
+mage test:pythonSpecific <test_file::test_name>
+
+#Example
+mage test:pythonSpecific ".mlflow.repo/tests/tracking/test_rest_tracking.py::test_rename_experiment"
+```
+
 ```bash
 # Run just unit tests 
 mage test:unit

conftest.py

@@ -44,6 +44,11 @@ def pytest_configure(config):
             "tests.store.tracking.test_sqlalchemy_store.test_log_param_max_length_value",
             "tests/override_test_sqlalchemy_store.py",
         ),
+        # This test uses monkeypatch.setenv which does not flow through to the Go side.
+        (
+            "tests.store.tracking.test_sqlalchemy_store.test_set_tag",
+            "tests/override_test_sqlalchemy_store.py",
+        ),
         # This tests calls the store using invalid metric entity that cannot be converted
         # to its proto counterpart.
         # Example: entities.Metric("invalid_metric", None, (int(time.time() * 1000)), 0).to_proto()

docs/porting-a-new-endpoint.md

@@ -219,7 +219,7 @@ An example use case where unit tests proved to be highly beneficial is the `filt
 
 ## Run Tests
 
-Run `mage test:python` and verify that our Go implementation passes the existing tests.
+Run `mage test:python` and verify that our Go implementation passes the existing tests. You can also use `mage test:pythonSpecific <testname>` to run a specific set of tests.
 
 There is one caveat to these tests; occasionally, they may have a Python bias, meaning that the tests pass due to Python's dynamic nature, while our Go tests might fail because they are strongly typed. Another issue may arise if the Python implementation does not consistently return the same error messages. Therefore, it may be necessary to submit a PR to [mlflow](https://github.com/mlflow/mlflow) to adjust the existing tests.
 

magefiles/configure.go

@@ -35,7 +35,7 @@ func Configure() error {
 		return err
 	}
 
-	// Archive the MLFlow pre-built UI
+	// Archive the MLflow pre-built UI
 	if err := tar(
 		"-C", "/usr/local/python/current/lib/python3.8/site-packages/mlflow",
 		"-czvf",

magefiles/generate/endpoints.go

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

magefiles/generate/validations.go

@@ -34,6 +34,8 @@ var validations = map[string]string{
 	"SetExperimentTag_ExperimentId":      "required",
 	"SetExperimentTag_Key":               "required,max=250,validMetricParamOrTagName",
 	"SetExperimentTag_Value":             "max=5000",
+	"SetTag_Key":                         "required,max=1000,validMetricParamOrTagName,pathIsUnique",
+	"SetTag_Value":                       "omitempty,truncate=8000",
 	"LogInputs_RunId":                    "required,runId",
 	"LogInputs_Datasets":                 "required",
 	"DatasetInput_Dataset":               "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(pytestArgs []string) error {
 	libpath, err := os.MkdirTemp("", "")
 	if err != nil {
 		return err
@@ -45,25 +44,42 @@ func (Test) Python() error {
 		return nil
 	}
 
+	args := []string{
+		"--confcutdir=.",
+		"-k", "not [file",
+	}
+	args = append(args, pytestArgs...)
+
 	//  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",
-		// "-vv",
+	}, "pytest", args...,
+	// "-vv",
 	); err != nil {
 		return err
 	}
 
 	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{
+		testName,
+		"-vv",
+	})
+}
+
 // Run the Go unit tests.
 func (Test) Unit() error {
 	return sh.RunV("go", "test", "./pkg/...")

mlflow_go/store/tracking.py

@@ -26,6 +26,7 @@
     RestoreExperiment,
     RestoreRun,
     SearchRuns,
+    SetTag,
     SetTraceTag,
     UpdateExperiment,
     UpdateRun,
@@ -200,6 +201,10 @@ def delete_trace_tag(self, request_id: str, key: str):
         )
         self.service.call_endpoint(get_lib().TrackingServiceDeleteTraceTag, 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 TrackingStore(cls):
     return type(cls.__name__, (_TrackingStore, cls), {})

pkg/tracking/service/tags.go

@@ -16,3 +16,17 @@ func (ts TrackingService) DeleteTag(
 
 	return &protos.DeleteTag_Response{}, nil
 }
+
+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
+}