feat: stable code with tests

Signed-off-by: Arsen Gumin <[email protected]>

Author: aagumin

pkg/modelprovider/mlflow/downloader.go

@@ -43,32 +43,38 @@ func NewMlFlowRegistry(mlflowClient *client.DatabricksClient) (MlFlowClient, err
 		return MlFlowClient{registry: registry}, nil
 	}
 
-	//TODO Support more auth methods?
+	// TODO Support more auth methods?
 	cfg := config.Config{
-		//Credentials: config.BasicCredentials{},
+		// Credentials: config.BasicCredentials{},
 	}
 	mlClient, err := client.New(&cfg)
-
 	if err != nil {
 		return MlFlowClient{}, err
 	}
 	registry = ml.NewModelRegistry(mlClient)
 	return MlFlowClient{registry: registry}, nil
 }
 
-func (mlfr *MlFlowClient) PullModelByName(ctx context.Context, modelName string, modelVersion string, destSrc string) error {
-
+func (mlfr *MlFlowClient) PullModelByName(
+	ctx context.Context,
+	modelName string,
+	modelVersion string,
+	destSrc string,
+) (string, error) {
 	if mlfr == nil || mlfr.registry == nil {
-		return errors.New("Mlflow client is not initialized: registry is nil")
+		return "", errors.New("mlflow client is not initialized: registry is nil")
 	}
 
-	versions, err := mlfr.registry.GetLatestVersionsAll(ctx, ml.GetLatestVersionsRequest{Name: modelName})
+	versions, err := mlfr.registry.GetLatestVersionsAll(
+		ctx,
+		ml.GetLatestVersionsRequest{Name: modelName},
+	)
 	if err != nil {
-		return errors.Join(errors.New(fmt.Sprintf("failed to get versions for model: %s", modelName)), err)
+		return "", errors.Join(fmt.Errorf("failed to get versions for model: %s", modelName), err)
 	}
 
 	if len(versions) == 0 {
-		return errors.New(fmt.Sprintf("model %s has versions: %v", modelName, versions))
+		return "", fmt.Errorf("model %s has versions: %v", modelName, versions)
 	}
 
 	var rawVersion []string
@@ -78,9 +84,12 @@ func (mlfr *MlFlowClient) PullModelByName(ctx context.Context, modelName string,
 	contains := slices.Contains(rawVersion, modelVersion)
 	if !contains {
 		msg := fmt.Sprintf(
-			"model %s version %s not found, available version %v", modelName, modelVersion, rawVersion,
+			"model %s version %s not found, available version %v",
+			modelName,
+			modelVersion,
+			rawVersion,
 		)
-		return errors.New(msg)
+		return "", errors.New(msg)
 	}
 
 	fmt.Printf("Found versions: '%v' for model '%s'\n", rawVersion, modelName)
@@ -92,40 +101,43 @@ func (mlfr *MlFlowClient) PullModelByName(ctx context.Context, modelName string,
 		Name:    modelName,
 		Version: modelVersion,
 	})
-
 	if err != nil {
-		return err
+		return "", err
 	}
 	fmt.Printf("Try pull model from uri %s", uri.ArtifactUri)
 	parsed, err := url.Parse(uri.ArtifactUri)
 	if parsed == nil {
-		return errors.New("failed to parse artifact uri")
+		return "", errors.New("failed to parse artifact uri")
 	}
 
 	switch parsed.Scheme {
 	case "s3":
 		s3storage := storageProvider[parsed.Scheme]
+		destSrc = filepath.Join(destSrc, modelName)
 		err = s3storage.DownloadModel(ctx, uri.ArtifactUri+"/", destSrc) // it's dir
 		if err != nil {
-			return err
+			return "", err
 		}
 	default:
 		msg := fmt.Sprintf("Unsupported artifact storage type: %s", parsed.Scheme)
 		err = errors.New(msg)
-		return err
+		return "", err
 	}
 
 	fmt.Printf("✅ Model downloaded")
 
-	return nil
+	return destSrc, nil
 }
 
 type S3StorageBackend struct {
 	addressing string
 }
 
-func (s3back *S3StorageBackend) DownloadModel(ctx context.Context, path string, destPath string) error {
-
+func (s3back *S3StorageBackend) DownloadModel(
+	ctx context.Context,
+	path string,
+	destPath string,
+) error {
 	parsed, err := url.Parse(path)
 	if err != nil {
 		return err
@@ -134,12 +146,10 @@ func (s3back *S3StorageBackend) DownloadModel(ctx context.Context, path string,
 	bucketName := parsed.Host
 	s3FolderPrefix := parsed.Path[1:]
 	fmt.Printf("Parsed s3 bucket %s, path %s from path", bucketName, s3FolderPrefix)
-	if destPath == "" {
-		destPath = "./downloads/"
-	}
+
 	cfg, err := awsconfig.LoadDefaultConfig(ctx)
 	if err != nil {
-		wrap := errors.New(fmt.Sprintf("Error loading AWS config, try change envs or profile: %v\n", err))
+		wrap := fmt.Errorf("Error loading AWS config, try change envs or profile: %v\n", err)
 		return errors.Join(wrap, err)
 	}
 
@@ -179,9 +189,13 @@ func (s3back *S3StorageBackend) DownloadModel(ctx context.Context, path string,
 			localFilePath := filepath.Join(destPath, relativePath)
 
 			// Create local directories if they don't exist
-			err = os.MkdirAll(filepath.Dir(localFilePath), 0755)
+			err = os.MkdirAll(filepath.Dir(localFilePath), 0o755)
 			if err != nil {
-				log.Printf("Error creating local directory %s: %v\n", filepath.Dir(localFilePath), err)
+				log.Printf(
+					"Error creating local directory %s: %v\n",
+					filepath.Dir(localFilePath),
+					err,
+				)
 				continue
 			}
 

pkg/modelprovider/mlflow/downloader_test.go

@@ -0,0 +1,116 @@
+package mlflow
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/client"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+)
+
+func TestMlFlowClient_PullModelByName(t *testing.T) {
+	type fields struct {
+		registry *ml.ModelRegistryAPI
+	}
+	type args struct {
+		ctx          context.Context
+		modelName    string
+		modelVersion string
+		destSrc      string
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name:    "nil receiver returns error",
+			fields:  fields{registry: nil},
+			args:    args{ctx: context.Background(), modelName: "model", modelVersion: "1", destSrc: "/tmp"},
+			want:    "",
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mlfr := &MlFlowClient{
+				registry: tt.fields.registry,
+			}
+			got, err := mlfr.PullModelByName(tt.args.ctx, tt.args.modelName, tt.args.modelVersion, tt.args.destSrc)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("PullModelByName() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("PullModelByName() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestNewMlFlowRegistry(t *testing.T) {
+	type args struct {
+		mlflowClient *client.DatabricksClient
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    MlFlowClient
+		wantErr bool
+	}{
+		{
+			name:    "non-nil client returns registry",
+			args:    args{mlflowClient: &client.DatabricksClient{}},
+			want:    MlFlowClient{},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := NewMlFlowRegistry(tt.args.mlflowClient)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewMlFlowRegistry() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !tt.wantErr && got.registry == nil {
+				t.Errorf("NewMlFlowRegistry() registry is nil")
+			}
+		})
+	}
+}
+
+func TestS3StorageBackend_DownloadModel(t *testing.T) {
+	type fields struct {
+		addressing string
+	}
+	type args struct {
+		ctx      context.Context
+		path     string
+		destPath string
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		args    args
+		wantErr bool
+	}{
+		{
+			name:    "invalid url returns error",
+			fields:  fields{addressing: ""},
+			args:    args{ctx: context.Background(), path: "http://[::1", destPath: "/tmp"},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s3back := &S3StorageBackend{
+				addressing: tt.fields.addressing,
+			}
+			if err := s3back.DownloadModel(tt.args.ctx, tt.args.path, tt.args.destPath); (err != nil) != tt.wantErr {
+				t.Errorf("DownloadModel() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}

pkg/modelprovider/mlflow/provider.go

@@ -20,15 +20,14 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/url"
 	"os"
-	"os/exec"
-	"path/filepath"
 	"strings"
 )
 
 // MlflowProvider implements the modelprovider.Provider interface for Mlflow
 type MlflowProvider struct {
-	mlfClient MlFlowClient
+	mflClient MlFlowClient
 }
 
 // New creates a new ModelScope provider instance
@@ -47,16 +46,25 @@ func (p *MlflowProvider) Name() string {
 func (p *MlflowProvider) SupportsURL(url string) bool {
 	url = strings.TrimSpace(url)
 	// TODO Mlflow API equals with Databricks Model Registry, support later
-	possibleUrls := []string{"models", "runs"}
+	possibleUrls := []string{"models"}
 
 	return hasAnyPrefix(url, possibleUrls)
 }
 
 // DownloadModel downloads a model from ModelScope using the modelscope CLI
 func (p *MlflowProvider) DownloadModel(ctx context.Context, modelURL, destDir string) (string, error) {
 	model, version, err := parseModelURL(modelURL)
-	xclient := NewMlFlowRegistry(nil)
-	xclient.PullModelByName(modelURL)
+	if err != nil {
+		return "", err
+	}
+	registryClient, err := NewMlFlowRegistry(nil)
+	if err != nil {
+		return "", err
+	}
+	downloadPath, err := registryClient.PullModelByName(ctx, model, version, destDir)
+	if err != nil {
+		return "", err
+	}
 	return downloadPath, nil
 }
 
@@ -75,6 +83,7 @@ func hasAnyPrefix(s string, subs []string) bool {
 }
 
 func checkMlflowAuth() error {
+
 	var err error
 
 	host := os.Getenv("DATABRICKS_HOST")
@@ -88,7 +97,10 @@ func checkMlflowAuth() error {
 		fmt.Println("Please set DATABRICKS_HOST environment variable.")
 		fmt.Println("Please set DATABRICKS_USERNAME environment variable.")
 		fmt.Println("Please set DATABRICKS_PASSWORD environment variable.")
-	} else if mlfhost != "" && mlfuser != "" && mlfpass != "" {
+	} else {
+		return nil
+	}
+	if mlfhost != "" && mlfuser != "" && mlfpass != "" {
 		err = os.Setenv("DATABRICKS_HOST", mlfhost)
 		if err != nil {
 			return err
@@ -108,12 +120,33 @@ func checkMlflowAuth() error {
 	return err
 }
 
-func parseModelURL(modelURL string) (name, version, error) {
-	var err error
-	var name string
-	var version string
+func parseModelURL(modelURL string) (string, string, error) {
+	if modelURL == "" {
+		return "", "", errors.New("modelUrl value missing.")
+	}
+
+	if strings.HasPrefix(modelURL, "models:") {
+		parse, err := url.Parse(modelURL)
+		if err != nil {
+			return "", "", err
+		}
+
+		if parse == nil {
+			return "", "", errors.New("model url is nil")
+		}
+
+		return parse.Hostname(), strings.TrimLeft(parse.Path, "/"), nil
+
+	} else if strings.Contains(modelURL, "/") {
 
-	name = ""
-	version = ""
-	return
+		split := strings.Split(modelURL, "/")
+
+		if len(split) != 2 {
+			return "", "", errors.New("model url is invalid, valid mask name/version")
+		}
+		return split[0], split[1], nil
+
+	} else {
+		return modelURL, "", nil
+	}
 }