From 38377ad098f07eb0f2889ebcf09dd1b14e2566c6 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 25 Apr 2023 05:49:47 -0700 Subject: [PATCH 1/2] Add support for run commands with mounted secrets Signed-off-by: Mattt Zmuda --- pkg/config/config.go | 65 ++++++++++++++++--- pkg/config/data/config_schema_v1.0.json | 29 +++++++++ pkg/dockerfile/generator.go | 24 +++++-- test-integration/test_integration/test_run.py | 29 +++++++++ 4 files changed, 133 insertions(+), 14 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 8248cc6efc..bcf3b1a451 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -19,16 +19,25 @@ import ( // TODO(andreas): custom cpu/gpu installs // TODO(andreas): suggest valid torchvision versions (e.g. if the user wants to use 0.8.0, suggest 0.8.1) +type RunItem struct { + Command string `json:"command,omitempty" yaml:"command"` + Mounts []struct { + Type string `json:"type,omitempty" yaml:"type"` + ID string `json:"id,omitempty" yaml:"id"` + Target string `json:"target,omitempty" yaml:"target"` + } `json:"mounts,omitempty" yaml:"mounts"` +} + type Build struct { - GPU bool `json:"gpu,omitempty" yaml:"gpu"` - PythonVersion string `json:"python_version,omitempty" yaml:"python_version"` - PythonRequirements string `json:"python_requirements,omitempty" yaml:"python_requirements"` - PythonPackages []string `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility - Run []string `json:"run,omitempty" yaml:"run"` - SystemPackages []string `json:"system_packages,omitempty" yaml:"system_packages"` - PreInstall []string `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility - CUDA string `json:"cuda,omitempty" yaml:"cuda"` - CuDNN string `json:"cudnn,omitempty" yaml:"cudnn"` + GPU bool `json:"gpu,omitempty" yaml:"gpu"` + PythonVersion string `json:"python_version,omitempty" yaml:"python_version"` + PythonRequirements string `json:"python_requirements,omitempty" yaml:"python_requirements"` + PythonPackages []string `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility + Run []RunItem `json:"run,omitempty" yaml:"run"` + SystemPackages []string `json:"system_packages,omitempty" yaml:"system_packages"` + PreInstall []string `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility + CUDA string `json:"cuda,omitempty" yaml:"cuda"` + CuDNN string `json:"cudnn,omitempty" yaml:"cudnn"` pythonRequirementsContent []string } @@ -54,6 +63,44 @@ func DefaultConfig() *Config { } } +func (r *RunItem) UnmarshalYAML(unmarshal func(interface{}) error) error { + var commandOrMap interface{} + if err := unmarshal(&commandOrMap); err != nil { + return err + } + + switch v := commandOrMap.(type) { + case string: + r.Command = v + case map[interface{}]interface{}: + var data []byte + var err error + + if data, err = yaml.Marshal(v); err != nil { + return err + } + + aux := struct { + Command string `yaml:"command"` + Mounts []struct { + Type string `yaml:"type"` + ID string `yaml:"id"` + Target string `yaml:"target"` + } `yaml:"mounts,omitempty"` + }{} + + if err := yaml.Unmarshal(data, &aux); err != nil { + return err + } + + *r = RunItem(aux) + default: + return fmt.Errorf("unexpected type %T for RunItem", v) + } + + return nil +} + func FromYAML(contents []byte) (*Config, error) { config := DefaultConfig() if err := yaml.Unmarshal(contents, config); err != nil { diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index eeb4e44fc4..a32884a38a 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -85,6 +85,35 @@ { "$id": "#/properties/build/properties/run/items/anyOf/0", "type": "string" + }, + { + "$id": "#/properties/build/properties/run/items/anyOf/1", + "type": "object", + "properties": { + "command": { + "type": "string" + }, + "mounts": { + "type": "array", + "items": { + "type": "object", + "properties": { + "type": { + "type": "string", + "enum": ["secret"] + }, + "id": { + "type": "string" + }, + "target": { + "type": "string" + } + }, + "required": ["type", "id", "target"] + } + } + }, + "required": ["command"] } ] } diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index bbb4d5bbf0..8c9f1b1958 100644 --- a/pkg/dockerfile/generator.go +++ b/pkg/dockerfile/generator.go @@ -228,17 +228,31 @@ func (g *Generator) run() (string, error) { runCommands := g.Config.Build.Run // For backwards compatibility - runCommands = append(runCommands, g.Config.Build.PreInstall...) + for _, command := range g.Config.Build.PreInstall { + runCommands = append(runCommands, config.RunItem{Command: command}) + } lines := []string{} for _, run := range runCommands { - run = strings.TrimSpace(run) - if strings.Contains(run, "\n") { + command := strings.TrimSpace(run.Command) + if strings.Contains(command, "\n") { return "", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command. -This is the offending line: %s`, run) +This is the offending line: %s`, command) + } + + if len(run.Mounts) > 0 { + mounts := []string{} + for _, mount := range run.Mounts { + if mount.Type == "secret" { + secretMount := fmt.Sprintf("--mount=type=secret,id=%s,target=%s", mount.ID, mount.Target) + mounts = append(mounts, secretMount) + } + } + lines = append(lines, fmt.Sprintf("RUN %s %s", strings.Join(mounts, " "), command)) + } else { + lines = append(lines, "RUN "+command) } - lines = append(lines, "RUN "+run) } return strings.Join(lines, "\n"), nil } diff --git a/test-integration/test_integration/test_run.py b/test-integration/test_integration/test_run.py index d9948e19de..ce35f805c7 100644 --- a/test-integration/test_integration/test_run.py +++ b/test-integration/test_integration/test_run.py @@ -17,3 +17,32 @@ def test_run(tmpdir_factory): capture_output=True, ) assert b"hello world" in result.stdout + + +def test_run_with_secret(tmpdir_factory): + tmpdir = tmpdir_factory.mktemp("project") + with open(tmpdir / "cog.yaml", "w") as f: + cog_yaml = """ +build: + python_version: "3.8" + run: + - echo hello world + - command: >- + echo shh + mounts: + - type: secret + id: foo + target: secret.txt + """ + f.write(cog_yaml) + with open(tmpdir / "secret.txt", "w") as f: + f.write("🤫") + + result = subprocess.run( + ["cog", "debug"], + cwd=tmpdir, + check=True, + capture_output=True, + ) + assert b"RUN echo hello world" in result.stdout + assert b"RUN --mount=type=secret,id=foo,target=secret.txt echo shh" in result.stdout From 080e0fae4760f52b17649d0678abc1a305d4d4aa Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 25 Apr 2023 13:26:04 -0700 Subject: [PATCH 2/2] Document new run item format Signed-off-by: Mattt Zmuda --- docs/yaml.md | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/docs/yaml.md b/docs/yaml.md index cda3b9f145..b956a9a97b 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -86,6 +86,20 @@ build: Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. +Each command in `run` can be either a string or a dictionary in the following format: + +````yaml +build: + run: + - command: pip install + mounts: + - type: secret + id: pip + target: /etc/pip.conf +``` + +You can use secret mounts to securely pass credentials to setup commands, without baking them into the image. For more information, see [Dockerfile reference](https://docs.docker.com/engine/reference/builder/#run---mounttypesecret). + ### `system_packages` A list of Ubuntu APT packages to install. For example: @@ -95,7 +109,7 @@ build: system_packages: - "ffmpeg" - "libavcodec-dev" -``` +```` ## `image`