Move /mlflow/traces/{request_id}/tags endpoint. (#11)

Signed-off-by: Software Developer <7852635+dsuhinin@users.noreply.github.com>

magefiles/generate/discovery/discovery.go

@@ -35,7 +35,7 @@ var routeParameterRegex = regexp.MustCompile(`<[^>]+:([^>]+)>`)
 // Get the safe path to use in Fiber registration.
 func (e Endpoint) GetFiberPath() string {
 	// e.Path cannot be trusted, it could be something like /mlflow-artifacts/artifacts/<path:artifact_path>
-	// Which would need to converted to /mlflow-artifacts/artifacts/:path
+	// which would need to be converted to /mlflow-artifacts/artifacts/:path
 	path := routeParameterRegex.ReplaceAllStringFunc(e.Path, func(s string) string {
 		parts := strings.Split(s, ":")
 

magefiles/generate/endpoints.go

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

magefiles/generate/validations.go

@@ -27,6 +27,8 @@ var validations = map[string]string{
 	"LogMetric_Key":                      "required",
 	"LogMetric_Value":                    "required",
 	"LogMetric_Timestamp":                "required",
+	"SetTraceTag_Key":                    "required,max=250,validMetricParamOrTagName,pathIsUnique",
+	"SetTraceTag_Value":                  "omitempty,truncate=8000",
 	"DeleteTag_RunId":                    "required",
 	"DeleteTag_Key":                      "required",
 	"SetExperimentTag_ExperimentId":      "required",

mlflow_go/store/tracking.py

@@ -7,6 +7,7 @@
     RunInfo,
     ViewType,
 )
+from mlflow.environment_variables import MLFLOW_TRUNCATE_LONG_VALUES
 from mlflow.exceptions import MlflowException
 from mlflow.protos import databricks_pb2
 from mlflow.protos.service_pb2 import (
@@ -25,6 +26,7 @@
     RestoreExperiment,
     RestoreRun,
     SearchRuns,
+    SetTraceTag,
     UpdateExperiment,
     UpdateRun,
 )
@@ -47,9 +49,12 @@ def __init__(self, *args, **kwargs):
         )
         config = json.dumps(
             {
-                "default_artifact_root": resolve_uri_if_local(default_artifact_root),
-                "tracking_store_uri": store_uri,
                 "log_level": logging.getLevelName(_logger.getEffectiveLevel()),
+                "python_tests_env": {
+                    "MLFLOW_TRUNCATE_LONG_VALUES": MLFLOW_TRUNCATE_LONG_VALUES.get()
+                },
+                "tracking_store_uri": store_uri,
+                "default_artifact_root": resolve_uri_if_local(default_artifact_root),
             }
         ).encode("utf-8")
         self.service = _ServiceProxy(get_lib().CreateTrackingService(config, len(config)))
@@ -176,6 +181,14 @@ def log_param(self, run_id, param):
         )
         self.service.call_endpoint(get_lib().TrackingServiceLogParam, request)
 
+    def set_trace_tag(self, request_id: str, key: str, value: str):
+        request = SetTraceTag(
+            key=key,
+            value=value,
+            request_id=request_id,
+        )
+        self.service.call_endpoint(get_lib().TrackingServiceSetTraceTag, request)
+
     def delete_tag(self, run_id, key):
         request = DeleteTag(run_id=run_id, key=key)
         self.service.call_endpoint(get_lib().TrackingServiceDeleteTag, request)

pkg/config/config.go

@@ -39,17 +39,18 @@ func (d *Duration) UnmarshalJSON(b []byte) error {
 }
 
 type Config struct {
-	Address               string   `json:"address"`
-	DefaultArtifactRoot   string   `json:"default_artifact_root"`
-	LogLevel              string   `json:"log_level"`
-	ModelRegistryStoreURI string   `json:"model_registry_store_uri"`
-	PythonAddress         string   `json:"python_address"`
-	PythonCommand         []string `json:"python_command"`
-	PythonEnv             []string `json:"python_env"`
-	ShutdownTimeout       Duration `json:"shutdown_timeout"`
-	StaticFolder          string   `json:"static_folder"`
-	TrackingStoreURI      string   `json:"tracking_store_uri"`
-	Version               string   `json:"version"`
+	Address               string                 `json:"address"`
+	DefaultArtifactRoot   string                 `json:"default_artifact_root"`
+	LogLevel              string                 `json:"log_level"`
+	ModelRegistryStoreURI string                 `json:"model_registry_store_uri"`
+	PythonEnv             []string               `json:"python_env"`
+	PythonAddress         string                 `json:"python_address"`
+	PythonCommand         []string               `json:"python_command"`
+	PythonTestsENV        map[string]interface{} `json:"python_tests_env"`
+	ShutdownTimeout       Duration               `json:"shutdown_timeout"`
+	StaticFolder          string                 `json:"static_folder"`
+	TrackingStoreURI      string                 `json:"tracking_store_uri"`
+	Version               string                 `json:"version"`
 }
 
 func NewConfigFromBytes(cfgBytes []byte) (*Config, error) {

pkg/entities/trace_tag.go

@@ -1,7 +1,17 @@
 package entities
 
+import "github.com/mlflow/mlflow-go/pkg/protos"
+
 type TraceTag struct {
 	Key       string
 	Value     string
 	RequestID string
 }
+
+func NewTraceTagFromProto(proto *protos.SetTraceTag) *TraceTag {
+	return &TraceTag{
+		Key:       proto.GetKey(),
+		Value:     proto.GetValue(),
+		RequestID: proto.GetRequestId(),
+	}
+}

pkg/lib/instance_map.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"sync"
 
 	"github.com/sirupsen/logrus"
@@ -48,6 +50,14 @@ func (s *instanceMap[T]) Create(
 		return -1
 	}
 
+	for key, value := range cfg.PythonTestsENV {
+		if err := os.Setenv(key, fmt.Sprintf("%v", value)); err != nil {
+			logrus.Error("Failed to set env: ", err)
+
+			return -1
+		}
+	}
+
 	logger := utils.NewLoggerFromConfig(cfg)
 
 	logger.Debugf("Loaded config: %#v", cfg)

pkg/lib/tracking.go

@@ -13,7 +13,9 @@ var trackingServices = newInstanceMap[*service.TrackingService]()
 //export CreateTrackingService
 func CreateTrackingService(configData unsafe.Pointer, configSize C.int) int64 {
 	//nolint:nlreturn
-	return trackingServices.Create(service.NewTrackingService, C.GoBytes(configData, configSize))
+	return trackingServices.Create(
+		service.NewTrackingService, C.GoBytes(configData, configSize),
+	)
 }
 
 //export DestroyTrackingService

pkg/tracking/service/trace.go

@@ -8,6 +8,18 @@ import (
 	"github.com/mlflow/mlflow-go/pkg/protos"
 )
 
+func (ts TrackingService) SetTraceTag(
+	ctx context.Context, input *protos.SetTraceTag,
+) (*protos.SetTraceTag_Response, *contract.Error) {
+	if err := ts.Store.SetTraceTag(
+		ctx, input.GetRequestId(), input.GetKey(), input.GetValue(),
+	); err != nil {
+		return nil, contract.NewErrorWith(protos.ErrorCode_INTERNAL_ERROR, "failed to create trace_tag", err)
+	}
+
+	return &protos.SetTraceTag_Response{}, nil
+}
+
 func (ts TrackingService) DeleteTraceTag(
 	ctx context.Context, input *protos.DeleteTraceTag,
 ) (*protos.DeleteTraceTag_Response, *contract.Error) {

pkg/tracking/store/sql/models/trace_tag.go

@@ -16,3 +16,11 @@ func (t TraceTag) ToEntity() *entities.TraceTag {
 		RequestID: t.RequestID,
 	}
 }
+
+func NewTraceTagFromEntity(entity *entities.TraceTag) TraceTag {
+	return TraceTag{
+		Key:       entity.Key,
+		Value:     entity.Value,
+		RequestID: entity.RequestID,
+	}
+}

pkg/tracking/store/sql/trace.go

@@ -6,13 +6,31 @@ import (
 	"fmt"
 
 	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
 	"github.com/mlflow/mlflow-go/pkg/entities"
 	"github.com/mlflow/mlflow-go/pkg/protos"
 	"github.com/mlflow/mlflow-go/pkg/tracking/store/sql/models"
 )
 
+func (s TrackingSQLStore) SetTraceTag(
+	ctx context.Context, requestID, key, value string,
+) error {
+	if err := s.db.WithContext(ctx).Clauses(clause.OnConflict{
+		Columns:   []clause.Column{{Name: "key"}, {Name: "request_id"}},
+		DoUpdates: clause.AssignmentColumns([]string{"value"}),
+	}).Create(models.TraceTag{
+		Key:       key,
+		Value:     value,
+		RequestID: requestID,
+	}).Error; err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (s TrackingSQLStore) GetTraceTag(
 	ctx context.Context, requestID, key string,
 ) (*entities.TraceTag, *contract.Error) {

pkg/tracking/store/store.go

@@ -40,6 +40,7 @@ type (
 		DeleteTag(ctx context.Context, runID, key string) *contract.Error
 	}
 	TraceTrackingStore interface {
+		SetTraceTag(ctx context.Context, requestID, key, value string) error
 		GetTraceTag(ctx context.Context, requestID, key string) (*entities.TraceTag, *contract.Error)
 		DeleteTraceTag(ctx context.Context, tag *entities.TraceTag) *contract.Error
 	}

pkg/validation/validation.go

@@ -117,7 +117,6 @@ func truncateFn(fieldLevel validator.FieldLevel) bool {
 
 	truncateLongValues, shouldTruncate := os.LookupEnv("MLFLOW_TRUNCATE_LONG_VALUES")
 	shouldTruncate = shouldTruncate && truncateLongValues == "true"
-
 	field := fieldLevel.Field()
 
 	if field.Kind() == reflect.String {