Skip to content

Commit bdb6e54

Browse files
ganisbackJames
ganisback
and
James
authored
support path mapping (#266)
Co-authored-by: James <xzgan@opencsg.com>
1 parent 4d68912 commit bdb6e54

File tree

8 files changed

+208
-134
lines changed

8 files changed

+208
-134
lines changed

‎_mocks/opencsg.com/csghub-server/builder/store/database/mock_MirrorStore.go‎

Lines changed: 7 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎api/middleware/repo.go‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -34,29 +34,31 @@ func RepoMapping(repo_type types.RepositoryType) gin.HandlerFunc {
3434
ctx.Next()
3535
return
3636
}
37-
mirror, err := mirrorStore.FindWithMapping(ctx, repo_type, namespace, name, mapping)
37+
repo, err := mirrorStore.FindWithMapping(ctx, repo_type, namespace, name, mapping)
3838
//if found mirror, that means this is a synced source, otherwise it's may a user-upload repo
3939
if err == nil {
40-
repo_id := strings.Split(mirror.Repository.Path, "/")
40+
namespace, name = repo.NamespaceAndName()
4141
//set the real namespace, the name was unchange
42-
slog.Info("namespace changed: ", "namespace", repo_id[0])
43-
ctx.Set("namespace_mapped", repo_id[0])
44-
ctx.Set("name_mapped", repo_id[1])
45-
// for modelscope, the default branch is master, we should mapp it to real branch
46-
if (branch == "main" || branch == "master") && mirror.Repository.DefaultBranch != branch {
47-
ctx.Set("branch_mapped", mirror.Repository.DefaultBranch)
42+
slog.Info("namespace changed: ", "namespace", namespace)
43+
ctx.Set("namespace_mapped", namespace)
44+
ctx.Set("name_mapped", name)
45+
// for modelscope, the default branch is master, we should map it to real branch
46+
if (branch == "main" || branch == "master") && repo.DefaultBranch != branch {
47+
ctx.Set("branch_mapped", repo.DefaultBranch)
4848
}
49-
ctx.Next()
50-
return
5149
}
5250
ctx.Next()
5351
}
5452
}
5553

5654
func GetMapping(ctx *gin.Context) types.Mapping {
57-
rawRp := ctx.Query("mirror")
58-
if rawRp == "" {
59-
return types.AutoMapping
55+
fullPath := ctx.FullPath()
56+
if strings.HasPrefix(fullPath, "/hf/") {
57+
return types.HFMapping
58+
}
59+
if strings.HasPrefix(fullPath, "/ms/") {
60+
return types.ModelScopeMapping
6061
}
61-
return types.Mapping(rawRp)
62+
//csg
63+
return types.CSGHubMapping
6264
}

‎api/router/api.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -147,8 +147,12 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
147147
if err != nil {
148148
return nil, fmt.Errorf("error creating HF dataset handler: %w", err)
149149
}
150-
151-
createHFRoutes(r, hfdsHandler, repoCommonHandler, modelHandler, userHandler)
150+
//create routes for hf
151+
createMappingRoutes(r, "/hf", hfdsHandler, repoCommonHandler, modelHandler, userHandler)
152+
//create routes for ms
153+
createMappingRoutes(r, "/ms", hfdsHandler, repoCommonHandler, modelHandler, userHandler)
154+
//create routes for csg
155+
createMappingRoutes(r, "/csg", hfdsHandler, repoCommonHandler, modelHandler, userHandler)
152156

153157
apiGroup := r.Group("/api/v1")
154158
// TODO:use middleware to handle common response
@@ -758,9 +762,9 @@ func createAccountRoutes(apiGroup *gin.RouterGroup, needAPIKey gin.HandlerFunc,
758762
}
759763
}
760764

761-
func createHFRoutes(r *gin.Engine, hfdsHandler *handler.HFDatasetHandler, repoCommonHandler *handler.RepoHandler, modelHandler *handler.ModelHandler, userHandler *handler.UserHandler) {
765+
func createMappingRoutes(r *gin.Engine, group string, hfdsHandler *handler.HFDatasetHandler, repoCommonHandler *handler.RepoHandler, modelHandler *handler.ModelHandler, userHandler *handler.UserHandler) {
762766
// Huggingface SDK routes
763-
hfGroup := r.Group("/hf")
767+
hfGroup := r.Group(group)
764768
{
765769
hfGroup.GET("/:namespace/:name/resolve/:branch/*file_path", middleware.RepoMapping(types.ModelRepo), repoCommonHandler.SDKDownload)
766770
hfGroup.HEAD("/:namespace/:name/resolve/:branch/*file_path", middleware.RepoMapping(types.ModelRepo), repoCommonHandler.HeadSDKDownload)

‎builder/store/database/mirror.go‎

Lines changed: 18 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ type MirrorStore interface {
2020
FindByRepoID(ctx context.Context, repoID int64) (*Mirror, error)
2121
FindByID(ctx context.Context, ID int64) (*Mirror, error)
2222
FindByRepoPath(ctx context.Context, repoType types.RepositoryType, namespace, name string) (*Mirror, error)
23-
FindWithMapping(ctx context.Context, repoType types.RepositoryType, namespace, name string, mapping types.Mapping) (*Mirror, error)
23+
FindWithMapping(ctx context.Context, repoType types.RepositoryType, namespace, name string, mapping types.Mapping) (*Repository, error)
2424
Create(ctx context.Context, mirror *Mirror) (*Mirror, error)
2525
WithPagination(ctx context.Context) ([]Mirror, error)
2626
WithPaginationWithRepository(ctx context.Context) ([]Mirror, error)
@@ -141,37 +141,25 @@ func (s *mirrorStoreImpl) FindByRepoPath(ctx context.Context, repoType types.Rep
141141
return &mirror, nil
142142
}
143143

144-
func (s *mirrorStoreImpl) FindWithMapping(ctx context.Context, repoType types.RepositoryType, namespace, name string, mapping types.Mapping) (*Mirror, error) {
145-
var mirror Mirror
146-
var err error
147-
if mapping == types.CSGHubMapping {
148-
return s.FindByRepoPath(ctx, repoType, namespace, name)
149-
} else if mapping == types.HFMapping {
150-
err = s.db.Operator.Core.NewSelect().
151-
Model(&mirror).
152-
Relation("Repository").
153-
Where("mirror.source_repo_path=?", fmt.Sprintf("%s/%s", namespace, name)).
154-
Where("repository.repository_type=?", repoType).
155-
Scan(ctx)
144+
func (s *mirrorStoreImpl) FindWithMapping(ctx context.Context, repoType types.RepositoryType, namespace, name string, mapping types.Mapping) (*Repository, error) {
145+
resRepo := new(Repository)
146+
query := s.db.Operator.Core.
147+
NewSelect().
148+
Model(resRepo)
149+
path := fmt.Sprintf("%s/%s", namespace, name)
150+
query.Where("repository_type = ?", repoType)
151+
if mapping == types.HFMapping {
152+
//compatiebility with old data
153+
//TODO: remove path after sdk 0.4.6
154+
query.Where("hf_path = ? or path = ?", path, path)
155+
} else if mapping == types.ModelScopeMapping {
156+
query.Where("ms_path = ?", path, path)
156157
} else {
157-
// auto mapping
158-
//fix some repo id has mirror but it's not public,for example: https://opencsg.com/models/Qwen/Qwen_Qwen2-7B-Instruct
159-
exist, _ := s.IsRepoExist(ctx, repoType, namespace, name)
160-
if exist {
161-
// no need mapping if repo id already exists in reporitory
162-
return nil, fmt.Errorf("repo already exists, no need mapping")
163-
}
164-
err = s.db.Operator.Core.NewSelect().
165-
Model(&mirror).
166-
Relation("Repository").
167-
Where("mirror.source_repo_path=?", fmt.Sprintf("%s/%s", namespace, name)).
168-
Where("repository.repository_type=?", repoType).
169-
Scan(ctx)
158+
// for csg path
159+
query.Where("path = ?", path)
170160
}
171-
if err != nil {
172-
return nil, err
173-
}
174-
return &mirror, nil
161+
err := query.Limit(1).Scan(ctx)
162+
return resRepo, err
175163
}
176164

177165
func (s *mirrorStoreImpl) Create(ctx context.Context, mirror *Mirror) (*Mirror, error) {

‎builder/store/database/mirror_test.go‎

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package database_test
22

33
import (
44
"context"
5-
"strings"
65
"testing"
76
"time"
87

@@ -129,42 +128,34 @@ func TestMirrorStore_FindWithMapping(t *testing.T) {
129128
store := database.NewMirrorStoreWithDB(db)
130129

131130
repos := []*database.Repository{
132-
{Name: "repo1", RepositoryType: types.ModelRepo, Path: "models_ns/repo1"},
133-
{Name: "repo2", RepositoryType: types.DatasetRepo, Path: "datasets_ns/repo2"},
134-
{Name: "repo3", RepositoryType: types.PromptRepo, Path: "prompts_ns/repo3"},
131+
{Name: "repo1", RepositoryType: types.ModelRepo, Path: "ns/repo1", HFPath: "hf/repo1"},
132+
{Name: "repo2", RepositoryType: types.DatasetRepo, Path: "ns/repo2", MSPath: "ms/repo2"},
133+
{Name: "repo3", RepositoryType: types.PromptRepo, Path: "ns/repo3"},
135134
}
136135

137136
for _, repo := range repos {
138137
repo.GitPath = repo.Path
139138
err := db.Core.NewInsert().Model(repo).Scan(ctx, repo)
140139
require.Nil(t, err)
141-
sp := strings.Split(repo.Path, "_")
142-
_, err = store.Create(ctx, &database.Mirror{
143-
RepositoryID: repo.ID,
144-
SourceRepoPath: strings.ReplaceAll(sp[1], "ns/", "nsn/"),
145-
Interval: repo.Name,
146-
})
147-
require.Nil(t, err)
148140
}
149141

150142
mi, err := store.FindWithMapping(ctx, types.ModelRepo, "ns", "repo1", types.CSGHubMapping)
151143
require.Nil(t, err)
152-
require.Equal(t, "repo1", mi.Interval)
153-
154-
_, err = store.FindWithMapping(ctx, types.ModelRepo, "ns", "repo1", types.HFMapping)
155-
require.NotNil(t, err)
144+
require.Equal(t, "repo1", mi.Name)
156145

157-
mi, err = store.FindWithMapping(ctx, types.ModelRepo, "nsn", "repo1", types.HFMapping)
146+
_, err = store.FindWithMapping(ctx, types.ModelRepo, "hf", "repo1", types.HFMapping)
158147
require.Nil(t, err)
159-
require.Equal(t, "repo1", mi.Interval)
160148

161-
mi, err = store.FindWithMapping(ctx, types.DatasetRepo, "nsn", "repo2", types.HFMapping)
149+
_, err = store.FindWithMapping(ctx, types.ModelRepo, "aaa", "repo1", types.HFMapping)
150+
require.NotNil(t, err)
151+
152+
mi, err = store.FindWithMapping(ctx, types.DatasetRepo, "ms", "repo2", types.ModelScopeMapping)
162153
require.Nil(t, err)
163-
require.Equal(t, "repo2", mi.Interval)
154+
require.Equal(t, "repo2", mi.Name)
164155

165-
mi, err = store.FindWithMapping(ctx, types.PromptRepo, "nsn", "repo3", types.AutoMapping)
156+
mi, err = store.FindWithMapping(ctx, types.PromptRepo, "ns", "repo3", types.CSGHubMapping)
166157
require.Nil(t, err)
167-
require.Equal(t, "repo3", mi.Interval)
158+
require.Equal(t, "repo3", mi.Name)
168159
}
169160

170161
func TestMirrorStore_ToSync(t *testing.T) {

0 commit comments

Comments
 (0)