feat: add config.repo in envd (#1101)

* add

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix lint

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

envd/api/config/__init__.py

@@ -136,3 +136,12 @@ def rstudio_server():
     """
     Enable the RStudio Server (only work for `base(os="ubuntu20.04", language="r")`)
     """
+
+
+def repo(url: str, description: str):
+    """Setup repo related information. Will save to the image labels.
+
+    Args:
+        url (str): repo URL
+        description (str): repo description
+    """

pkg/lang/frontend/starlark/config/config.go

@@ -44,6 +44,7 @@ var Module = &starlarkstruct.Module{
 			ruleJuliaPackageServer, ruleFuncJuliaPackageServer),
 		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
 		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
+		"repo":           starlark.NewBuiltin(ruleRepo, ruleFuncRepo),
 	},
 }
 
@@ -212,3 +213,16 @@ func ruleFuncEntrypoint(thread *starlark.Thread, _ *starlark.Builtin,
 	ir.Entrypoint(argList)
 	return starlark.None, nil
 }
+
+func ruleFuncRepo(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var url, description string
+
+	if err := starlark.UnpackArgs(ruleRepo, args, kwargs, "url", &url, "description?", &description); err != nil {
+		return nil, err
+	}
+
+	logger.Debugf("repo info: url=%s, description=%s", url, description)
+	ir.Repo(url, description)
+	return starlark.None, nil
+}

pkg/lang/frontend/starlark/config/const.go

@@ -24,4 +24,5 @@ const (
 	ruleJuliaPackageServer = "config.julia_pkg_server"
 	ruleRStudioServer      = "config.rstudio_server"
 	ruleEntrypoint         = "config.entrypoint"
+	ruleRepo               = "config.repo"
 )

pkg/lang/ir/compile.go

@@ -181,6 +181,12 @@ func (g Graph) Labels() (map[string]string, error) {
 	}
 	labels[types.ImageLabelPorts] = string(portsData)
 
+	repoInfo, err := json.Marshal(g.Repo)
+	if err != nil {
+		return labels, err
+	}
+	labels[types.ImageLabelRepo] = string(repoInfo)
+
 	return labels, nil
 }
 

pkg/lang/ir/interface.go

@@ -19,6 +19,7 @@ import (
 	"github.com/opencontainers/go-digest"
 
 	"github.com/tensorchord/envd/pkg/editor/vscode"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 func Base(os, language, image string) error {
@@ -229,3 +230,10 @@ func RuntimeEnviron(env map[string]string) {
 func RuntimeInitScript(commands []string) {
 	DefaultGraph.RuntimeInitScript = append(DefaultGraph.RuntimeInitScript, commands)
 }
+
+func Repo(url, description string) {
+	DefaultGraph.Repo = types.RepoInfo{
+		Description: description,
+		URL:         url,
+	}
+}

pkg/lang/ir/types.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/tensorchord/envd/pkg/editor/vscode"
 	"github.com/tensorchord/envd/pkg/progress/compileui"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 // A Graph contains the state,
@@ -61,6 +62,8 @@ type Graph struct {
 	HTTP       []HTTPInfo
 	Entrypoint []string
 
+	Repo types.RepoInfo
+
 	*JupyterConfig
 	*GitConfig
 	*CondaConfig

pkg/types/envd.go

@@ -145,6 +145,11 @@ type Dependency struct {
 	PyPIPackages []string `json:"pypi_packages,omitempty"`
 }
 
+type RepoInfo struct {
+	URL         string `json:"url,omitempty"`
+	Description string `json:"description,omitempty"`
+}
+
 type PortBinding struct {
 	Port     string
 	Protocol string

pkg/types/label.go

@@ -22,6 +22,7 @@ const (
 
 	ImageLabelVendor    = "ai.tensorchord.envd.vendor"
 	ImageLabelGPU       = "ai.tensorchord.envd.gpu"
+	ImageLabelRepo      = "ai.tensorchord.envd.repo"
 	ImageLabelPorts     = "ai.tensorchord.envd.ports"
 	ImageLabelAPT       = "ai.tensorchord.envd.apt.packages"
 	ImageLabelPyPI      = "ai.tensorchord.envd.pypi.commands"