Repository: G-Research/fasttrackml
Pull request #802 (2024-01-11 21:09:45)

Fix default experiment parameters for new namespaces.
What's inside the box:
- properly initialise default experiment for newly created namespaces. before fix we missing such fields: `CreationTime`, `LastUpdateTime`, `ArtifactLocation`.


Patch:
diff --git a/pkg/api/admin/service/namespace/service.go b/pkg/api/admin/service/namespace/service.go
index 4956636..a10cd88 100644
--- a/pkg/api/admin/service/namespace/service.go
+++ b/pkg/api/admin/service/namespace/service.go
@@ -2,25 +2,37 @@ package namespace
 
 import (
 	"context"
+	"database/sql"
+	"fmt"
+	"net/url"
+	"time"
 
 	"github.com/rotisserie/eris"
 
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/api"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/common"
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/config"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/repositories"
 )
 
 // Service provides service layer to work with `namespace` business logic.
 type Service struct {
-	namespaceRepository repositories.NamespaceRepositoryProvider
+	config               *config.ServiceConfig
+	namespaceRepository  repositories.NamespaceRepositoryProvider
+	experimentRepository repositories.ExperimentRepositoryProvider
 }
 
 // NewService creates new Service instance.
 func NewService(
+	config *config.ServiceConfig,
 	namespaceRepository repositories.NamespaceRepositoryProvider,
+	experimentRepository repositories.ExperimentRepositoryProvider,
 ) *Service {
 	return &Service{
-		namespaceRepository: namespaceRepository,
+		config:               config,
+		namespaceRepository:  namespaceRepository,
+		experimentRepository: experimentRepository,
 	}
 }
 
@@ -47,25 +59,49 @@ func (s Service) CreateNamespace(ctx context.Context, code, description string)
 	if err := ValidateNamespace(code); err != nil {
 		return nil, eris.Wrap(err, "error validating namespace")
 	}
-	exp := &models.Experiment{
-		Name:           "Default",
-		LifecycleStage: models.LifecycleStageActive,
-	}
+
 	namespace := &models.Namespace{
 		Code:                code,
 		Description:         description,
-		Experiments:         []models.Experiment{*exp},
 		DefaultExperimentID: common.GetPointer(models.DefaultExperimentID),
 	}
 	if err := s.namespaceRepository.Create(ctx, namespace); err != nil {
 		return nil, eris.Wrap(err, "error creating namespace")
 	}
+
+	timestamp := time.Now().UTC().UnixMilli()
+	experiment := models.Experiment{
+		Name:           models.DefaultExperimentName,
+		NamespaceID:    namespace.ID,
+		CreationTime:   sql.NullInt64{Int64: timestamp, Valid: true},
+		LifecycleStage: models.LifecycleStageActive,
+		LastUpdateTime: sql.NullInt64{Int64: timestamp, Valid: true},
+	}
+
+	if err := s.experimentRepository.Create(ctx, &experiment); err != nil {
+		return nil, eris.Wrap(err, "error creating experiment")
+	}
+
 	// update Namespace with correct DefaultExperimentID now that it is known
-	namespace.DefaultExperimentID = namespace.Experiments[0].ID
-	err := s.namespaceRepository.Update(ctx, namespace)
-	if err != nil {
+	namespace.DefaultExperimentID = experiment.ID
+	if err := s.namespaceRepository.Update(ctx, namespace); err != nil {
 		return nil, eris.Wrap(err, "error setting namespace default experiment id during create")
 	}
+
+	// setup ArtifactLocation for default experiment.
+	path, err := url.JoinPath(s.config.DefaultArtifactRoot, fmt.Sprintf("%d", *experiment.ID))
+	if err != nil {
+		return nil, api.NewInternalError(
+			"error creating artifact_location for experiment'%s': %s", experiment.Name, err,
+		)
+	}
+	experiment.ArtifactLocation = path
+	if err := s.experimentRepository.Update(ctx, &experiment); err != nil {
+		return nil, api.NewInternalError(
+			"error updating artifact_location for experiment '%s': %s", experiment.Name, err,
+		)
+	}
+
 	return namespace, nil
 }
 
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 0751ccd..86a83fb 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -251,21 +251,33 @@ func createApp(
 	// init `admin` api routes.
 	adminAPI.NewRouter(
 		adminAPIController.NewController(
-			namespace.NewService(namespaceRepository),
+			namespace.NewService(
+				config,
+				namespaceRepository,
+				mlflowRepositories.NewExperimentRepository(db.GormDB()),
+			),
 		),
 	).Init(app)
 
 	// init `admin` UI routes.
 	adminUI.NewRouter(
 		adminUIController.NewController(
-			namespace.NewService(namespaceRepository),
+			namespace.NewService(
+				config,
+				namespaceRepository,
+				mlflowRepositories.NewExperimentRepository(db.GormDB()),
+			),
 		),
 	).Init(app)
 
 	// init `chooser` ui routes.
 	chooser.NewRouter(
 		chooserController.NewController(
-			namespace.NewService(namespaceRepository),
+			namespace.NewService(
+				config,
+				namespaceRepository,
+				mlflowRepositories.NewExperimentRepository(db.GormDB()),
+			),
 		),
 	).AddRoutes(app)
 


Test patch:
diff --git a/pkg/api/admin/service/namespace/service_test.go b/pkg/api/admin/service/namespace/service_test.go
index da3f646..d75403d 100644
--- a/pkg/api/admin/service/namespace/service_test.go
+++ b/pkg/api/admin/service/namespace/service_test.go
@@ -3,12 +3,15 @@ package namespace
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/common"
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/config"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/repositories"
 )
@@ -35,8 +38,36 @@ func TestService_CreateNamespace_Ok(t *testing.T) {
 		}),
 	).Return(nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+	experimentRepository.On(
+		"Create",
+		context.TODO(),
+		mock.MatchedBy(func(experiment *models.Experiment) bool {
+			assert.Equal(t, models.DefaultExperimentName, experiment.Name)
+			assert.Equal(t, models.LifecycleStageActive, experiment.LifecycleStage)
+			assert.NotNil(t, experiment.CreationTime)
+			assert.NotNil(t, experiment.LastUpdateTime)
+			experiment.ID = common.GetPointer(int32(1))
+			return true
+		}),
+	).Return(nil)
+	experimentRepository.On(
+		"Update",
+		context.TODO(),
+		mock.MatchedBy(func(experiment *models.Experiment) bool {
+			assert.Equal(
+				t,
+				fmt.Sprintf("default_artifact_root/%d", *experiment.ID),
+				experiment.ArtifactLocation,
+			)
+			return true
+		}),
+	).Return(nil)
+
 	// call service under testing.
-	service := NewService(&namespaceRepository)
+	service := NewService(&config.ServiceConfig{
+		DefaultArtifactRoot: "default_artifact_root",
+	}, &namespaceRepository, &experimentRepository)
 	_, err := service.CreateNamespace(context.TODO(), "code", "description")
 
 	// compare results.
@@ -55,8 +86,15 @@ func TestService_CreateNamespace_Error(t *testing.T) {
 		"Update", context.TODO(), mock.Anything, mock.Anything,
 	).Return(nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+	experimentRepository.On(
+		"Update",
+		context.TODO(),
+		mock.Anything,
+	).Return(nil)
+
 	// call service under testing.
-	service := NewService(&namespaceRepository)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	_, err = service.CreateNamespace(context.TODO(), "code", "description")
 
 	// compare results.
@@ -77,10 +115,10 @@ func TestService_GetNamespace_Ok(t *testing.T) {
 		"GetByID", context.TODO(), uint(0),
 	).Return(&ns, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	namespace, err := service.GetNamespace(context.TODO(), uint(0))
 
 	// compare results.
@@ -95,10 +133,10 @@ func TestService_GetNamespace_Error(t *testing.T) {
 		"GetByID", context.TODO(), uint(0),
 	).Return(nil, errors.New("something is wrong"))
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	namespace, err := service.GetNamespace(context.TODO(), uint(0))
 
 	// compare results.
@@ -121,10 +159,10 @@ func TestService_ListNamespace_Ok(t *testing.T) {
 		"List", context.TODO(),
 	).Return(testNamespaces, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	namespaces, err := service.ListNamespaces(context.TODO())
 
 	// compare results.
@@ -139,10 +177,10 @@ func TestService_ListNamespaces_Error(t *testing.T) {
 		"List", context.TODO(),
 	).Return(nil, errors.New("error listing namespaces"))
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	namespaces, err := service.ListNamespaces(context.TODO())
 
 	// compare results.
@@ -164,10 +202,10 @@ func TestService_DeleteNamespace_Ok(t *testing.T) {
 		"GetByID", context.TODO(), uint(0),
 	).Return(&ns, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	err := service.DeleteNamespace(context.TODO(), uint(0))
 
 	// compare results.
@@ -181,10 +219,10 @@ func TestService_DeleteNamespace_Error(t *testing.T) {
 		"GetByID", context.TODO(), uint(0),
 	).Return(nil, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	err := service.DeleteNamespace(context.TODO(), uint(0))
 
 	// compare results.
@@ -202,10 +240,10 @@ func TestService_DeleteDefaultNamespace_Error(t *testing.T) {
 		Code: models.DefaultNamespaceCode,
 	}, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(
-		&namespaceRepository,
-	)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	err := service.DeleteNamespace(context.TODO(), uint(0))
 
 	// compare results.
@@ -232,8 +270,10 @@ func TestService_UpdateNamespace_Ok(t *testing.T) {
 		"GetByID", context.TODO(), uint(1),
 	).Return(&ns, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(&namespaceRepository)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	_, err := service.UpdateNamespace(context.TODO(), uint(1), "code", "description")
 
 	// compare results.
@@ -247,8 +287,10 @@ func TestService_UpdateNamespace_Error(t *testing.T) {
 		"GetByID", context.TODO(), uint(1),
 	).Return(nil, nil)
 
+	experimentRepository := repositories.MockExperimentRepositoryProvider{}
+
 	// call service under testing.
-	service := NewService(&namespaceRepository)
+	service := NewService(&config.ServiceConfig{}, &namespaceRepository, &experimentRepository)
 	_, err := service.UpdateNamespace(context.TODO(), uint(1), "code", "description")
 
 	// compare results.