From 9229bfa01c5377eae6db79cb263d97f565458618 Mon Sep 17 00:00:00 2001
From: Sam Skjonsberg <sams@allenai.org>
Date: Thu, 21 Apr 2022 14:40:08 -0700
Subject: [PATCH] Add flag for exposing session ports. (#365)

This adds the `--port` flag to the `beaker session create`
command.
---
 cmd/beaker/session.go | 19 +++++++++++++++++++
 go.mod                |  4 ++--
 go.sum                |  4 ++--
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/cmd/beaker/session.go b/cmd/beaker/session.go
index 1aea3db1..913742db 100644
--- a/cmd/beaker/session.go
+++ b/cmd/beaker/session.go
@@ -116,10 +116,17 @@ To pass flags, use "--" e.g. "create -- ls -l"`,
 	var gpus int
 	var memory string
 	var sharedMemory string
+	var ports []string
 	cmd.Flags().Float64Var(&cpus, "cpus", 0, "Minimum CPU cores to reserve, e.g. 7.5")
 	cmd.Flags().IntVar(&gpus, "gpus", 0, "Minimum number of GPUs to reserve")
 	cmd.Flags().StringVar(&memory, "memory", "", "Minimum memory to reserve, e.g. 6.5GiB")
 	cmd.Flags().StringVar(&sharedMemory, "shared-memory", "", "Shared memory (size of /dev/shm), e.g. 1GiB")
+	cmd.Flags().StringSliceVar(
+		&ports,
+		"port",
+		[]string{},
+		"TCP container ports to expose. Each will be assigned a random, ephemeral port on the host.",
+	)
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		rt, err := docker.NewRuntime()
@@ -178,6 +185,17 @@ To pass flags, use "--" e.g. "create -- ls -l"`,
 			})
 		}
 
+		var tcpPorts []api.TCPPort
+		for _, p := range ports {
+			pp, err := strconv.ParseInt(p, 10, 32)
+			if err != nil {
+				return fmt.Errorf("invalid port: %d", pp)
+			}
+			// We know the conversion to int32 is safe because of the bitsize
+			// passed to strconv.ParseInt() above.
+			tcpPorts = append(tcpPorts, int32(pp))
+		}
+
 		session, err := beaker.CreateJob(ctx, api.JobSpec{
 			Session: &api.SessionJobSpec{
 				Workspace: workspace,
@@ -194,6 +212,7 @@ To pass flags, use "--" e.g. "create -- ls -l"`,
 				Datasets:  mounts,
 				Image:     *imageSource,
 				SaveImage: saveImage,
+				TCPPorts:  tcpPorts,
 			},
 		})
 		if err != nil {
diff --git a/go.mod b/go.mod
index b718281a..3f4854db 100644
--- a/go.mod
+++ b/go.mod
@@ -4,7 +4,7 @@ go 1.17
 
 require (
 	github.com/allenai/bytefmt v0.1.2
-	github.com/beaker/client v0.0.0-20220118174312-992c493d18d8
+	github.com/beaker/client v0.0.0-20220421175640-599703769000
 	github.com/beaker/fileheap v0.0.0-20211007204440-1bd3920c4320
 	github.com/beaker/runtime v0.0.0-20211213171103-95151aa06fad
 	github.com/docker/distribution v2.8.1+incompatible
@@ -33,6 +33,7 @@ require (
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-retryablehttp v0.7.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/kr/pretty v0.2.1 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6 // indirect
@@ -54,7 +55,6 @@ require (
 	google.golang.org/genproto v0.0.0-20220208230804-65c12eb4c068 // indirect
 	google.golang.org/grpc v1.44.0 // indirect
 	google.golang.org/protobuf v1.27.1 // indirect
-	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 )
 
 // See https://github.com/advisories/GHSA-c2h3-6mxw-7mvq
diff --git a/go.sum b/go.sum
index 3d80b89b..573c9180 100644
--- a/go.sum
+++ b/go.sum
@@ -103,8 +103,8 @@ github.com/allenai/bytefmt v0.1.2 h1:27970/ph1eYlv9e5i6Jc6SZoIyXy+DVD4hDDiF0dQGI
 github.com/allenai/bytefmt v0.1.2/go.mod h1:3hhcDqHXjwM3JBHx4cX79jG5G3V8zKsjeutbhJ3XjI0=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
-github.com/beaker/client v0.0.0-20220118174312-992c493d18d8 h1:rDcZ19MPKf20umUkBeQV5hw3R8q7UIK/aOIz2V0nvew=
-github.com/beaker/client v0.0.0-20220118174312-992c493d18d8/go.mod h1:gneEdoTEWgYoS32NWMZN6bsZWqlnYSj3NxFawTTJPWI=
+github.com/beaker/client v0.0.0-20220421175640-599703769000 h1:+q57+NHrNvNBscKC5TxKY/rcZRqGzVhj32VdqrtRUIg=
+github.com/beaker/client v0.0.0-20220421175640-599703769000/go.mod h1:3YTEVm+/Yb/iPCh9tq0PzotY/DTdEYoEO6WvU882XJU=
 github.com/beaker/fileheap v0.0.0-20211007204440-1bd3920c4320 h1:z3Ed6zWWi5wz+AUojzbxZglrfNLM8YeeMrmAv7WbsQQ=
 github.com/beaker/fileheap v0.0.0-20211007204440-1bd3920c4320/go.mod h1:B39KYPwS7QX1cXqPUKmC0ypuZy5idMGwPsLrY1N7dWQ=
 github.com/beaker/runtime v0.0.0-20211213171103-95151aa06fad h1:m6VFGPVTfwKAxjRakGf5tc9vy5e5SPVbRNfTpvEmPxY=