feat(lang): implement expose func (#780)

* WIP: add expose ir func Signed-off-by: nullday <aseaday@hotmail.com> * WIP: modify named keyword of expose Signed-off-by: nullday <aseaday@hotmail.com> * Feat: finish expose func Signed-off-by: nullday <aseaday@hotmail.com> * feat: fix the proposal Signed-off-by: nullday <aseaday@hotmail.com> * Feat: move expose func to the module runtime Signed-off-by: nullday <aseaday@hotmail.com> * Refactor: rename ir.Expose to ir.RuntimeExpose Signed-off-by: nullday <aseaday@hotmail.com> * Fix the error of judge port range Signed-off-by: nullday <aseaday@hotmail.com> Co-authored-by: nullday <aseaday@hotmail.com> Co-authored-by: WeiZhang <kweizh@gmail.com> Signed-off-by: nullday <aseaday@hotmail.com> Co-authored-by: WeiZhang <kweizh@gmail.com>
tensorchord · Aug 17, 2022 · e30866f · e30866f
1 parent c49863e
commit e30866f
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 2 deletions.
diff --git a/docs/proposals/20220707-service-primitive.md b/docs/proposals/20220707-service-primitive.md
@@ -19,7 +19,9 @@ In the future, we should also keep the possibility for remote or cloud access.
     - wsl access from the windows
 
 ## API
+
+```python
+expose(envd_port=<LPORT>, host_port=<RPORT>, service=<SERVICE_NAME>)
 ```
-expose(local_port=<LPORT>, host_port=<RPORT>, svc=<SERVICE_NAME>)
-```
+
 It will bind the host port to the local port.
diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
@@ -508,6 +508,26 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 		config.ExposedPorts[natPort] = struct{}{}
 	}
 
+	if g.ExposeConfig != nil && len(g.ExposeConfig.ExposeItems) > 0 {
+		for _, item := range g.ExposeConfig.ExposeItems {
+			var err error
+			if item.HostPort == 0 {
+				item.HostPort, err = netutil.GetFreePort()
+				if err != nil {
+					return "", "", errors.Wrap(err, "failed to get a free port")
+				}
+			}
+			natPort := nat.Port(fmt.Sprintf("%d/tcp", item.EnvdPort))
+			hostConfig.PortBindings[natPort] = []nat.PortBinding{
+				{
+					HostIP:   localhost,
+					HostPort: strconv.Itoa(item.HostPort),
+				},
+			}
+			config.ExposedPorts[natPort] = struct{}{}
+		}
+	}
+
 	if gpuEnabled {
 		logger.Debug("GPU is enabled.")
 		hostConfig.DeviceRequests = deviceRequests(numGPUs)

diff --git a/pkg/lang/frontend/starlark/runtime/const.go b/pkg/lang/frontend/starlark/runtime/const.go
@@ -16,5 +16,6 @@ package runtime
 
 const (
 	ruleCommand = "runtime.command"
+	ruleExpose  = "runtime.expose"
 	ruleDaemon  = "runtime.daemon"
 )
diff --git a/pkg/lang/frontend/starlark/runtime/runtime.go b/pkg/lang/frontend/starlark/runtime/runtime.go
@@ -15,6 +15,7 @@
 package runtime
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/sirupsen/logrus"
@@ -33,6 +34,7 @@ var Module = &starlarkstruct.Module{
 	Members: starlark.StringDict{
 		"command": starlark.NewBuiltin(ruleCommand, ruleFuncCommand),
 		"daemon":  starlark.NewBuiltin(ruleDaemon, ruleFuncDaemon),
+		"expose":  starlark.NewBuiltin(ruleExpose, ruleFuncExpose),
 	},
 }
 
@@ -90,3 +92,30 @@ func ruleFuncDaemon(thread *starlark.Thread, _ *starlark.Builtin,
 	}
 	return starlark.None, nil
 }
+
+func ruleFuncExpose(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var (
+		envdPort    starlark.Int
+		hostPort    = starlark.MakeInt(0)
+		serviceName = starlark.String("")
+	)
+
+	if err := starlark.UnpackArgs(ruleExpose,
+		args, kwargs, "envd_port", &envdPort, "host_port?", &hostPort, "service?", &serviceName); err != nil {
+		return nil, err
+	}
+	envdPortInt, ok := envdPort.Int64()
+	if !ok || envdPortInt < 1 || envdPortInt > 65535 {
+		return nil, errors.New("envd_port must be a positive integer less than 65535")
+	}
+	hostPortInt, ok := hostPort.Int64()
+	if !ok || hostPortInt < 1 || hostPortInt > 65535 {
+		return nil, errors.New("envd_port must be a positive integer less than 65535")
+	}
+	serviceNameStr := serviceName.GoString()
+
+	logger.Debugf("rule `%s` is invoked, envd_port=%d, host_port=%d, service=%s", ruleExpose, envdPortInt, hostPortInt, serviceNameStr)
+	err := ir.RuntimeExpose(int(envdPortInt), int(hostPortInt), serviceNameStr)
+	return starlark.None, err
+}
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -217,3 +217,17 @@ func RuntimeCommands(commands map[string]string) {
 func RuntimeDaemon(commands [][]string) {
 	DefaultGraph.RuntimeDaemon = append(DefaultGraph.RuntimeDaemon, commands...)
 }
+
+func RuntimeExpose(envdPort, hostPort int, serviceName string) error {
+	if DefaultGraph.ExposeConfig == nil {
+		DefaultGraph.ExposeConfig = &ExposeConfig{
+			ExposeItems: make([]ExposeItem, 0),
+		}
+	}
+	DefaultGraph.ExposeConfig.ExposeItems = append(DefaultGraph.ExposeConfig.ExposeItems, ExposeItem{
+		EnvdPort:    envdPort,
+		HostPort:    hostPort,
+		ServiceName: serviceName,
+	})
+	return nil
+}
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -61,6 +61,7 @@ type Graph struct {
 	*GitConfig
 	*CondaConfig
 	*RStudioServerConfig
+	*ExposeConfig
 
 	Writer      compileui.Writer
 	CachePrefix string
@@ -96,6 +97,16 @@ type GitConfig struct {
 	Editor string
 }
 
+type ExposeItem struct {
+	EnvdPort    int
+	HostPort    int
+	ServiceName string
+}
+
+type ExposeConfig struct {
+	ExposeItems []ExposeItem
+}
+
 type JupyterConfig struct {
 	Token string
 	Port  int64