feat: support model caching (#317)

Because

- the model downloading time is quite long for huge models

This commit

- support caching model into` ~/.cache/instill/models` when enabling the
config `cache.model = true`

Dockerfile

@@ -26,6 +26,7 @@ RUN tar -xf ArtiVC-v${ARTIVC_VERSION}-${TARGETOS}-${TARGETARCH}.tar.gz -C /usr/l
 RUN mkdir /etc/vdp
 RUN mkdir /vdp
 RUN mkdir /model-repository
+RUN mkdir /.cache
 
 FROM --platform=$BUILDPLATFORM ubuntu:${UBUNTU_VERSION}
 
@@ -68,4 +69,5 @@ COPY --from=build --chown=nobody:nogroup /usr/local/bin/avc /usr/local/bin/avc
 
 COPY --from=build --chown=nobody:nogroup /etc/vdp /etc/vdp
 COPY --from=build --chown=nobody:nogroup /vdp /vdp
-COPY --from=build --chown=nobody:nogroup /model-repository /model-repository
+COPY --from=build --chown=nobody:nogroup /model-repository /model-repository
+COPY --from=build --chown=nobody:nogroup /.cache /.cache

Dockerfile.dev

@@ -47,6 +47,7 @@ COPY . .
 
 RUN chown -R nobody:nogroup /go
 RUN chown -R nobody:nogroup /tmp
+RUN chown -R nobody:nogroup /.cache
 
 ENV GOCACHE /go/.cache/go-build
 ENV GOENV /go/.config/go/env

Makefile

@@ -17,6 +17,7 @@ dev:							## Run dev container
 	@docker run -d --rm \
 		-v $(PWD):/${SERVICE_NAME} \
 		-v model-repository:/model-repository \
+		-v ~/.cache/instill:/.cache \
 		-p ${SERVICE_PORT}:${SERVICE_PORT} \
 		--network instill-network \
 		--name ${SERVICE_NAME} \

config/config.go

@@ -68,6 +68,7 @@ type CacheConfig struct {
 	Redis struct {
 		RedisOptions redis.Options `koanf:"redisoptions"`
 	}
+	Model bool `koanf:"model"`
 }
 
 // UsageServerConfig related to usage-server

config/config.yaml

@@ -35,6 +35,7 @@ cache:
   redis:
     redisoptions:
       addr: redis:6379
+  model: true
 usageserver:
   tlsenabled: true
   host: usage.instill.tech

pkg/handler/public_handler.go

@@ -781,17 +781,19 @@ func createGitHubModel(h *PublicHandler, ctx context.Context, req *modelPB.Creat
 			Tag:        tag.Name,
 		}
 		rdid, _ := uuid.NewV4()
-		modelSrcDir := fmt.Sprintf("/tmp/%v", rdid.String())
+		modelSrcDir := fmt.Sprintf("/tmp/%v", rdid.String()) + ""
+		if config.Config.Cache.Model { // cache model into ~/.cache/instill/models
+			modelSrcDir = util.MODEL_CACHE_DIR + "/" + instanceConfig.Repository + instanceConfig.Tag
+		}
 
 		if config.Config.Server.ItMode { // use local model for testing to remove internet connection issue while testing
-			cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir %s; cp -rf assets/model-dummy-cls/* %s", modelSrcDir, modelSrcDir))
+			cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir -p %s > /dev/null; cp -rf assets/model-dummy-cls/* %s", modelSrcDir, modelSrcDir))
 			if err := cmd.Run(); err != nil {
 				util.RemoveModelRepository(config.Config.TritonServer.ModelStore, owner, githubModel.ID, tag.Name)
 				return &modelPB.CreateModelResponse{}, err
 			}
 		} else {
-			cloneLargeFile := false
-			err = util.GitHubClone(cloneLargeFile, modelSrcDir, instanceConfig)
+			err = util.GitHubClone(modelSrcDir, instanceConfig, false)
 			if err != nil {
 				st, err := sterr.CreateErrorResourceInfo(
 					codes.FailedPrecondition,
@@ -816,7 +818,9 @@ func createGitHubModel(h *PublicHandler, ctx context.Context, req *modelPB.Creat
 		}
 
 		readmeFilePath, ensembleFilePath, err := util.UpdateModelPath(modelSrcDir, config.Config.TritonServer.ModelStore, owner, githubModel.ID, &instance)
-		_ = os.RemoveAll(modelSrcDir) // remove uploaded temporary files
+		if !config.Config.Cache.Model {
+			_ = os.RemoveAll(modelSrcDir) // remove uploaded temporary files
+		}
 		if err != nil {
 			st, err := sterr.CreateErrorResourceInfo(
 				codes.FailedPrecondition,
@@ -1008,7 +1012,7 @@ func createArtiVCModel(h *PublicHandler, ctx context.Context, req *modelPB.Creat
 		rdid, _ := uuid.NewV4()
 		modelSrcDir := fmt.Sprintf("/tmp/%v", rdid.String())
 		if config.Config.Server.ItMode { // use local model for testing to remove internet connection issue while testing
-			cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir %s; cp -rf assets/model-dummy-cls/* %s", modelSrcDir, modelSrcDir))
+			cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir -p %s > /dev/null; cp -rf assets/model-dummy-cls/* %s", modelSrcDir, modelSrcDir))
 			if err := cmd.Run(); err != nil {
 				util.RemoveModelRepository(config.Config.TritonServer.ModelStore, owner, artivcModel.ID, tag)
 				return &modelPB.CreateModelResponse{}, err
@@ -1209,7 +1213,7 @@ func createHuggingFaceModel(h *PublicHandler, ctx context.Context, req *modelPB.
 	rdid, _ := uuid.NewV4()
 	configTmpDir := fmt.Sprintf("/tmp/%s", rdid.String())
 	if config.Config.Server.ItMode { // use local model for testing to remove internet connection issue while testing
-		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir %s; cp -rf assets/tiny-vit-random/* %s", configTmpDir, configTmpDir))
+		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("mkdir -p %s > /dev/null; cp -rf assets/tiny-vit-random/* %s", configTmpDir, configTmpDir))
 		if err := cmd.Run(); err != nil {
 			_ = os.RemoveAll(configTmpDir)
 			return &modelPB.CreateModelResponse{}, err

pkg/triton/triton.go

@@ -933,11 +933,5 @@ func (ts *triton) IsTritonServerReady() bool {
 		return false
 	}
 	fmt.Printf("Triton Health - Live: %v\n", serverLiveResponse.Live)
-	if !serverLiveResponse.Live {
-		return false
-	}
-
-	serverReadyResponse := ts.ServerReadyRequest()
-	fmt.Printf("Triton Health - Ready: %v\n", serverReadyResponse.Ready)
-	return serverReadyResponse.Ready
+	return serverLiveResponse.Live
 }

pkg/util/const.go

@@ -88,3 +88,6 @@ const (
 	TEXT_GENERATION_TOP_K      = int64(1)
 	TEXT_GENERATION_SEED       = int64(0)
 )
+
+const MODEL_CACHE_DIR = "/.cache/models"
+const MODEL_CACHE_FILE = "cached_models.json"

pkg/util/file.go

@@ -212,7 +212,7 @@ func UpdateModelPath(modelDir string, dstDir string, owner string, modelID strin
 	var readmeFilePath string
 	files := []FileMeta{}
 	var configFiles []string
-	var fileRe = regexp.MustCompile(`.git|.dvc|.dvcignore`)
+	var fileRe = regexp.MustCompile(`/.git|/.dvc|/.dvcignore`)
 	err := filepath.Walk(modelDir, func(path string, f os.FileInfo, err error) error {
 		if !fileRe.MatchString(path) {
 			files = append(files, FileMeta{
@@ -243,7 +243,6 @@ func UpdateModelPath(modelDir string, dstDir string, owner string, modelID strin
 		oldModelName := subStrs[0]
 		subStrs[0] = fmt.Sprintf("%v#%v#%v#%v", owner, modelID, oldModelName, modelInstance.ID)
 		var filePath = filepath.Join(dstDir, strings.Join(subStrs, "/"))
-
 		if f.fInfo.IsDir() { // create new folder
 			err = os.MkdirAll(filePath, os.ModePerm)
 
@@ -389,3 +388,16 @@ func SaveFile(stream modelPB.ModelPublicService_CreateModelBinaryFileUploadServe
 	}
 	return tmpFile, &uploadedModel, modelDefinitionID, nil
 }
+
+// CreateFolder creates a folder if it does not exist.
+func CreateFolder(dstDir string) error {
+	if _, err := os.Stat(dstDir); os.IsNotExist(err) {
+		err := os.MkdirAll(dstDir, os.ModePerm)
+		if err != nil {
+			return err
+		}
+	} else {
+		return fmt.Errorf("folder already exists")
+	}
+	return nil
+}

pkg/util/util.go

@@ -122,39 +122,113 @@ func getInferModelConfigPath(modelRepository string, tritonModels []datamodel.Tr
 	return modelPath
 }
 
-func GitHubClone(isWithLargeFile bool, dir string, instanceConfig datamodel.GitHubModelInstanceConfiguration) error {
+type CacheModel struct {
+	ModelRepo string `json:"model_repo"`
+	State     string `json:"state"`
+}
+
+// GitHubClone clones a repository from GitHub.
+func GitHubClone(dir string, instanceConfig datamodel.GitHubModelInstanceConfiguration, isWithLargeFile bool) error {
 	urlRepo := instanceConfig.Repository
-	if !strings.HasPrefix(urlRepo, "https://github.com") {
-		urlRepo = "https://github.com/" + urlRepo
-	}
-	if !strings.HasSuffix(urlRepo, ".git") {
-		urlRepo = urlRepo + ".git"
+
+	// Check in the cache first.
+	var cacheModels []CacheModel
+	if config.Config.Cache.Model {
+		_ = CreateFolder(MODEL_CACHE_DIR) // create model cache folder if not exist.
+		if _, err := os.Stat(MODEL_CACHE_DIR + "/" + MODEL_CACHE_FILE); !os.IsNotExist(err) {
+			f, err := os.ReadFile(MODEL_CACHE_DIR + "/" + MODEL_CACHE_FILE)
+			if err != nil {
+				return err
+			}
+			if err := json.Unmarshal([]byte(f), &cacheModels); err != nil {
+				return err
+			}
+			for _, cacheModel := range cacheModels {
+				if cacheModel.ModelRepo != (instanceConfig.Repository + instanceConfig.Tag) {
+					continue
+				}
+				if cacheModel.State == "done" { // everything is cached.
+					return nil
+				} else if cacheModel.State == "without_large_file" && !isWithLargeFile { // the GitHub repo is being cached.
+					return nil
+				}
+			}
+		}
 	}
+	if !isWithLargeFile || isWithLargeFile && !config.Config.Cache.Model {
+		if !strings.HasPrefix(urlRepo, "https://github.com") {
+			urlRepo = "https://github.com/" + urlRepo
+		}
+		if !strings.HasSuffix(urlRepo, ".git") {
+			urlRepo = urlRepo + ".git"
+		}
 
-	extraFlag := ""
-	if !isWithLargeFile {
-		extraFlag = "GIT_LFS_SKIP_SMUDGE=1"
+		extraFlag := ""
+		if !isWithLargeFile {
+			extraFlag = "GIT_LFS_SKIP_SMUDGE=1"
+		}
+		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("%s git clone -b %s %s %s", extraFlag, instanceConfig.Tag, urlRepo, dir))
+		err := cmd.Run()
+		if err != nil {
+			return err
+		}
 	}
 
-	cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("%s git clone -b %s %s %s", extraFlag, instanceConfig.Tag, urlRepo, dir))
-	err := cmd.Run()
-	if err != nil {
-		return err
+	var f *os.File
+	var err error
+	if config.Config.Cache.Model {
+		f, err = os.Create(MODEL_CACHE_DIR + "/" + MODEL_CACHE_FILE)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
 	}
 
 	if isWithLargeFile {
 		dvcPaths := findDVCPaths(dir)
 		for _, dvcPath := range dvcPaths {
-			cmd = exec.Command("/bin/sh", "-c", fmt.Sprintf("cd %s; dvc pull %s", dir, dvcPath))
+			cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cd %s; dvc pull %s", dir, dvcPath))
 			err = cmd.Run()
 			if err != nil {
 				return err
 			}
 		}
+		if config.Config.Cache.Model {
+			for i, cacheModel := range cacheModels {
+				if cacheModel.ModelRepo == (instanceConfig.Repository + instanceConfig.Tag) {
+					cacheModels[i].State = "done"
+					break
+				}
+			}
+			b, err := json.Marshal(cacheModels)
+			if err != nil {
+				return err
+			}
+			if _, err := f.Write(b); err != nil {
+				return err
+			}
+		}
+	} else {
+		if config.Config.Cache.Model {
+			cacheFile := CacheModel{
+				ModelRepo: instanceConfig.Repository + instanceConfig.Tag,
+				State:     "without_large_file",
+			}
+			cacheModels = append(cacheModels, cacheFile)
+			b, err := json.Marshal(cacheModels)
+			if err != nil {
+				return err
+			}
+			if _, err := f.Write(b); err != nil {
+				return err
+			}
+		}
 	}
+
 	return nil
 }
 
+// CopyModelFileToModelRepository copies model files to model repository.
 func CopyModelFileToModelRepository(modelRepository string, dir string, tritonModels []datamodel.TritonModel) error {
 	modelPaths := findModelFiles(dir)
 	for _, modelPath := range modelPaths {

pkg/worker/model.go

@@ -114,26 +114,28 @@ func (w *worker) DeployModelActivity(ctx context.Context, param *ModelInstancePa
 	}
 
 	// downloading model weight when making inference
+	rdid, _ := uuid.NewV4()
+	modelSrcDir := fmt.Sprintf("/tmp/%s", rdid.String())
 	switch modelDef.ID {
 	case "github":
 		if !config.Config.Server.ItMode && !util.HasModelWeightFile(config.Config.TritonServer.ModelStore, tritonModels) {
 			var instanceConfig datamodel.GitHubModelInstanceConfiguration
 			if err := json.Unmarshal(dbModelInstance.Configuration, &instanceConfig); err != nil {
 				return err
 			}
-			rdid, _ := uuid.NewV4()
-			modelSrcDir := fmt.Sprintf("/tmp/%s", rdid.String())
 
-			cloneLargeFile := true
-			if err := util.GitHubClone(cloneLargeFile, modelSrcDir, instanceConfig); err != nil {
+			if config.Config.Cache.Model { // cache model into ~/.cache/instill/models
+				modelSrcDir = util.MODEL_CACHE_DIR + "/" + instanceConfig.Repository + instanceConfig.Tag
+			}
+
+			if err := util.GitHubClone(modelSrcDir, instanceConfig, true); err != nil {
 				_ = os.RemoveAll(modelSrcDir)
 				return err
 			}
 			if err := util.CopyModelFileToModelRepository(config.Config.TritonServer.ModelStore, modelSrcDir, tritonModels); err != nil {
 				_ = os.RemoveAll(modelSrcDir)
 				return err
 			}
-			_ = os.RemoveAll(modelSrcDir)
 		}
 	case "huggingface":
 		if !util.HasModelWeightFile(config.Config.TritonServer.ModelStore, tritonModels) {
@@ -142,14 +144,16 @@ func (w *worker) DeployModelActivity(ctx context.Context, param *ModelInstancePa
 				return err
 			}
 
+			if config.Config.Cache.Model { // cache model into ~/.cache/instill/models
+				modelSrcDir = util.MODEL_CACHE_DIR + "/" + instanceConfig.RepoId
+			}
+
 			var modelConfig datamodel.HuggingFaceModelConfiguration
 			err = json.Unmarshal([]byte(dbModel.Configuration), &modelConfig)
 			if err != nil {
 				return err
 			}
 
-			rdid, _ := uuid.NewV4()
-			modelSrcDir := fmt.Sprintf("/tmp/%s", rdid.String())
 			if config.Config.Server.ItMode { // use local model to remove internet connection issue while integration testing
 				if err = util.HuggingFaceExport(modelSrcDir, datamodel.HuggingFaceModelConfiguration{
 					RepoId: "assets/tiny-vit-random",
@@ -172,7 +176,6 @@ func (w *worker) DeployModelActivity(ctx context.Context, param *ModelInstancePa
 			if err := util.UpdateModelConfig(config.Config.TritonServer.ModelStore, tritonModels); err != nil {
 				return err
 			}
-			_ = os.RemoveAll(modelSrcDir)
 		}
 	case "artivc":
 		if !config.Config.Server.ItMode && !util.HasModelWeightFile(config.Config.TritonServer.ModelStore, tritonModels) {
@@ -187,8 +190,6 @@ func (w *worker) DeployModelActivity(ctx context.Context, param *ModelInstancePa
 				return err
 			}
 
-			rdid, _ := uuid.NewV4()
-			modelSrcDir := fmt.Sprintf("/tmp/%v", rdid.String())
 			err = util.ArtiVCClone(modelSrcDir, modelConfig, instanceConfig, true)
 			if err != nil {
 				_ = os.RemoveAll(modelSrcDir)
@@ -201,6 +202,10 @@ func (w *worker) DeployModelActivity(ctx context.Context, param *ModelInstancePa
 		}
 	}
 
+	if !config.Config.Cache.Model {
+		_ = os.RemoveAll(modelSrcDir)
+	}
+
 	tEnsembleModel, _ := w.repository.GetTritonEnsembleModel(param.ModelInstanceUID)
 	for _, tModel := range tritonModels {
 		if tEnsembleModel.Name != "" && tEnsembleModel.Name == tModel.Name { // load ensemble model last.