Skip to content

Move workdir to top level of config #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andreasjansson merged 2 commits into from
Apr 22, 2021
Merged

Move workdir to top level of config #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d93844e
Move helper scripts and fix workdir
andreasjansson Apr 22, 2021
90e06c4
Move workdir to top level of config
andreasjansson Apr 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion pkg/cli/debug.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"io/ioutil"
"os"
"path"
"runtime"

"github.com/spf13/cobra"

Expand Down Expand Up @@ -66,7 +67,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
if err != nil {
return err
}
generator := &docker.DockerfileGenerator{Config: config, Arch: arch}
generator := &docker.DockerfileGenerator{Config: config, Arch: arch, GOOS: runtime.GOOS, GOARCH: runtime.GOARCH}
out, err := generator.Generate()
if err != nil {
return err
Expand Down
5 changes: 3 additions & 2 deletions pkg/cli/test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ package cli

import (
"fmt"
"path/filepath"
"os"
"path/filepath"
"runtime"

"github.com/spf13/cobra"

Expand Down Expand Up @@ -64,7 +65,7 @@ func Test(cmd *cobra.Command, args []string) error {
if _, ok := archMap[arch]; !ok {
return fmt.Errorf("Architecture %s is not defined for model", arch)
}
generator := &docker.DockerfileGenerator{Config: config, Arch: arch}
generator := &docker.DockerfileGenerator{Config: config, Arch: arch, GOOS: runtime.GOOS, GOARCH: runtime.GOARCH}
dockerfileContents, err := generator.Generate()
if err != nil {
return fmt.Errorf("Failed to generate Dockerfile for %s: %w", arch, err)
Expand Down
30 changes: 23 additions & 7 deletions pkg/docker/generate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@ var cogLibrary []byte
type DockerfileGenerator struct {
Config *model.Config
Arch string

// these are here to make this type testable
GOOS string
GOARCH string
}

func (g *DockerfileGenerator) Generate() (string, error) {
Expand All @@ -52,6 +56,7 @@ func (g *DockerfileGenerator) Generate() (string, error) {
pythonRequirements, err := g.pythonRequirements()
if err != nil {
return "", err

}
pipInstalls, err := g.pipInstalls()
if err != nil {
Expand Down Expand Up @@ -149,13 +154,17 @@ func (g *DockerfileGenerator) installHelperScripts() string {
}

func (g *DockerfileGenerator) serverHelperScript(serverClass string, filename string) string {
scriptPath := "/code/" + filename
scriptPath := "/usr/bin/" + filename
name := g.Config.Model
parts := strings.Split(name, ".py:")
module := parts[0]
class := parts[1]
script := `#!/usr/bin/env python
import sys
import cog
import os
os.chdir("` + g.getWorkdir() + `")
sys.path.append("` + g.getWorkdir() + `")
from ` + module + ` import ` + class + `
cog.` + serverClass + `(` + class + `()).start_server()`
scriptString := strings.ReplaceAll(script, "\n", "\\n")
Expand All @@ -165,14 +174,17 @@ RUN chmod +x ` + scriptPath
}

func (g *DockerfileGenerator) queueWorkerHelperScript() string {
scriptPath := "/code/cog-redis-queue-worker"
scriptPath := "/usr/bin/cog-redis-queue-worker"
name := g.Config.Model
parts := strings.Split(name, ".py:")
module := parts[0]
class := parts[1]
script := `#!/usr/bin/env python
import sys
import cog
import os
os.chdir("` + g.getWorkdir() + `")
sys.path.append("` + g.getWorkdir() + `")
from ` + module + ` import ` + class + `
cog.RedisQueueWorker(` + class + `(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()`
scriptString := strings.ReplaceAll(script, "\n", "\\n")
Expand All @@ -191,7 +203,7 @@ RUN pip install -r /tmp/requirements.txt && rm /tmp/requirements.txt`, reqs), ni
}

func (g *DockerfileGenerator) pipInstalls() (string, error) {
packages, indexURLs, err := g.Config.PythonPackagesForArch(g.Arch)
packages, indexURLs, err := g.Config.PythonPackagesForArch(g.Arch, g.GOOS, g.GOARCH)
if err != nil {
return "", err
}
Expand Down Expand Up @@ -221,15 +233,19 @@ func (g *DockerfileGenerator) copyCode() string {
func (g *DockerfileGenerator) command() string {
// TODO: handle infer scripts in subdirectories
// TODO: check this actually exists
return `CMD /code/cog-http-server`
return `CMD /usr/bin/cog-http-server`
}

func (g *DockerfileGenerator) workdir() string {
return "WORKDIR " + g.getWorkdir()
}

func (g *DockerfileGenerator) getWorkdir() string {
wd := "/code"
if g.Config.Environment.Workdir != "" {
wd += "/" + g.Config.Environment.Workdir
if g.Config.Workdir != "" {
wd += "/" + g.Config.Workdir
}
return "WORKDIR " + wd
return wd
}

func (g *DockerfileGenerator) preInstall() string {
Expand Down
58 changes: 22 additions & 36 deletions pkg/docker/generate_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,15 +64,9 @@ ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu
` + installPython("3.8") + installCog() + `
RUN ### --> Copying code
COPY . /code

RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /code/cog-http-server
RUN chmod +x /code/cog-http-server
RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /code/cog-ai-platform-prediction-server
RUN chmod +x /code/cog-ai-platform-prediction-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()' > /code/cog-redis-queue-worker
RUN chmod +x /code/cog-redis-queue-worker
` + helperScripts() + `
WORKDIR /code
CMD /code/cog-http-server`
CMD /usr/bin/cog-http-server`

expectedGPU := `FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu16.04
ENV DEBIAN_FRONTEND=noninteractive
Expand All @@ -81,20 +75,14 @@ ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu
` + installPython("3.8") + installCog() + `
RUN ### --> Copying code
COPY . /code

RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /code/cog-http-server
RUN chmod +x /code/cog-http-server
RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /code/cog-ai-platform-prediction-server
RUN chmod +x /code/cog-ai-platform-prediction-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()' > /code/cog-redis-queue-worker
RUN chmod +x /code/cog-redis-queue-worker
` + helperScripts() + `
WORKDIR /code
CMD /code/cog-http-server`
CMD /usr/bin/cog-http-server`

gen := DockerfileGenerator{conf, "cpu"}
gen := DockerfileGenerator{Config: conf, Arch: "cpu"}
actualCPU, err := gen.Generate()
require.NoError(t, err)
gen = DockerfileGenerator{conf, "gpu"}
gen = DockerfileGenerator{Config: conf, Arch: "gpu"}
actualGPU, err := gen.Generate()
require.NoError(t, err)

Expand Down Expand Up @@ -131,15 +119,9 @@ RUN pip install -f https://download.pytorch.org/whl/torch_stable.html torch==1
` + installCog() + `
RUN ### --> Copying code
COPY . /code

RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /code/cog-http-server
RUN chmod +x /code/cog-http-server
RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /code/cog-ai-platform-prediction-server
RUN chmod +x /code/cog-ai-platform-prediction-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()' > /code/cog-redis-queue-worker
RUN chmod +x /code/cog-redis-queue-worker
` + helperScripts() + `
WORKDIR /code
CMD /code/cog-http-server`
CMD /usr/bin/cog-http-server`

expectedGPU := `FROM nvidia/cuda:10.2-cudnn8-devel-ubuntu18.04
ENV DEBIAN_FRONTEND=noninteractive
Expand All @@ -155,23 +137,27 @@ RUN pip install torch==1.5.1 pandas==1.2.0.12
` + installCog() + `
RUN ### --> Copying code
COPY . /code

RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /code/cog-http-server
RUN chmod +x /code/cog-http-server
RUN echo '#!/usr/bin/env python\nimport cog\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /code/cog-ai-platform-prediction-server
RUN chmod +x /code/cog-ai-platform-prediction-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()' > /code/cog-redis-queue-worker
RUN chmod +x /code/cog-redis-queue-worker
` + helperScripts() + `
WORKDIR /code
CMD /code/cog-http-server`
CMD /usr/bin/cog-http-server`

gen := DockerfileGenerator{conf, "cpu"}
gen := DockerfileGenerator{Config: conf, Arch: "cpu"}
actualCPU, err := gen.Generate()
require.NoError(t, err)
gen = DockerfileGenerator{conf, "gpu"}
gen = DockerfileGenerator{Config: conf, Arch: "gpu"}
actualGPU, err := gen.Generate()
require.NoError(t, err)

require.Equal(t, expectedCPU, actualCPU)
require.Equal(t, expectedGPU, actualGPU)
}

func helperScripts() string {
return `
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /usr/bin/cog-http-server
RUN chmod +x /usr/bin/cog-http-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /usr/bin/cog-ai-platform-prediction-server
RUN chmod +x /usr/bin/cog-ai-platform-prediction-server
RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4]).start()' > /usr/bin/cog-redis-queue-worker
RUN chmod +x /usr/bin/cog-redis-queue-worker`
}
13 changes: 6 additions & 7 deletions pkg/model/compatibility.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
_ "embed"
"encoding/json"
"fmt"
"runtime"
"sort"
"strings"

Expand Down Expand Up @@ -249,10 +248,10 @@ func tfGPUPackage(ver string, cuda string) (name string, cpuVersion string, err
return "", "", fmt.Errorf("No matching tensorflow GPU package for version %s and CUDA %s", ver, cuda)
}

func torchCPUPackage(ver string) (name string, cpuVersion string, indexURL string, err error) {
func torchCPUPackage(ver string, goos string, goarch string) (name string, cpuVersion string, indexURL string, err error) {
for _, compat := range TorchCompatibilityMatrix {
if compat.TorchVersion() == ver && compat.CUDA == nil {
return "torch", torchStripCPUSuffixForM1(compat.Torch), compat.IndexURL, nil
return "torch", torchStripCPUSuffixForM1(compat.Torch, goos, goarch), compat.IndexURL, nil
}
}

Expand Down Expand Up @@ -296,10 +295,10 @@ func torchGPUPackage(ver string, cuda string) (name string, cpuVersion string, i
return "torch", latest.Torch, latest.IndexURL, nil
}

func torchvisionCPUPackage(ver string) (name string, cpuVersion string, indexURL string, err error) {
func torchvisionCPUPackage(ver string, goos string, goarch string) (name string, cpuVersion string, indexURL string, err error) {
for _, compat := range TorchCompatibilityMatrix {
if compat.TorchvisionVersion() == ver && compat.CUDA == nil {
return "torchvision", torchStripCPUSuffixForM1(compat.Torchvision), compat.IndexURL, nil
return "torchvision", torchStripCPUSuffixForM1(compat.Torchvision, goos, goarch), compat.IndexURL, nil
}
}
return "", "", "", fmt.Errorf("No matching torchvision CPU package for version %s", ver)
Expand Down Expand Up @@ -344,9 +343,9 @@ func torchvisionGPUPackage(ver string, cuda string) (name string, cpuVersion str

// aarch64 packages don't have +cpu suffix: https://download.pytorch.org/whl/torch_stable.html
// TODO(andreas): clean up this hack by actually parsing the torch_stable.html list in the generator
func torchStripCPUSuffixForM1(version string) string {
func torchStripCPUSuffixForM1(version string, goos string, goarch string) string {
// TODO(andreas): clean up this hack
if runtime.GOOS == "darwin" && runtime.GOARCH == "arm64" {
if goos == "darwin" && goarch == "arm64" {
return strings.ReplaceAll(version, "+cpu", "")
}
return version
Expand Down
12 changes: 6 additions & 6 deletions pkg/model/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ type Environment struct {
PythonFindLinks []string `json:"python_find_links" yaml:"python_find_links"`
PythonPackages []string `json:"python_packages" yaml:"python_packages"`
SystemPackages []string `json:"system_packages" yaml:"system_packages"`
Workdir string `json:"workdir" yaml:"workdir"`
PreInstall []string `json:"pre_install" yaml:"pre_install"`
PostInstall []string `json:"post_install" yaml:"post_install"`
Architectures []string `json:"architectures" yaml:"architectures"`
Expand All @@ -39,6 +38,7 @@ type Config struct {
Environment *Environment `json:"environment" yaml:"environment"`
Model string `json:"model" yaml:"model"`
Examples []*Example `json:"examples" yaml:"examples"`
Workdir string `json:"workdir" yaml:"workdir"`
}

func DefaultConfig() *Config {
Expand Down Expand Up @@ -138,11 +138,11 @@ func (c *Config) ValidateAndCompleteConfig() error {
return nil
}

func (c *Config) PythonPackagesForArch(arch string) (packages []string, indexURLs []string, err error) {
func (c *Config) PythonPackagesForArch(arch string, goos string, goarch string) (packages []string, indexURLs []string, err error) {
packages = []string{}
indexURLSet := map[string]bool{}
for _, pkg := range c.Environment.PythonPackages {
archPkg, indexURL, err := c.pythonPackageForArch(pkg, arch)
archPkg, indexURL, err := c.pythonPackageForArch(pkg, arch, goos, goarch)
if err != nil {
return nil, nil, err
}
Expand All @@ -158,7 +158,7 @@ func (c *Config) PythonPackagesForArch(arch string) (packages []string, indexURL
return packages, indexURLs, nil
}

func (c *Config) pythonPackageForArch(pkg string, arch string) (actualPackage string, indexURL string, err error) {
func (c *Config) pythonPackageForArch(pkg string, arch string, goos string, goarch string) (actualPackage string, indexURL string, err error) {
name, version, err := splitPythonPackage(pkg)
if err != nil {
return "", "", err
Expand All @@ -177,7 +177,7 @@ func (c *Config) pythonPackageForArch(pkg string, arch string) (actualPackage st
}
} else if name == "torch" {
if arch == "cpu" {
name, version, indexURL, err = torchCPUPackage(version)
name, version, indexURL, err = torchCPUPackage(version, goos, goarch)
if err != nil {
return "", "", err
}
Expand All @@ -189,7 +189,7 @@ func (c *Config) pythonPackageForArch(pkg string, arch string) (actualPackage st
}
} else if name == "torchvision" {
if arch == "cpu" {
name, version, indexURL, err = torchvisionCPUPackage(version)
name, version, indexURL, err = torchvisionCPUPackage(version, goos, goarch)
if err != nil {
return "", "", err
}
Expand Down
6 changes: 3 additions & 3 deletions pkg/model/config_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ func TestPythonPackagesForArchTorchGPU(t *testing.T) {
require.Equal(t, "10.1", config.Environment.CUDA)
require.Equal(t, "8", config.Environment.CuDNN)

packages, indexURLs, err := config.PythonPackagesForArch("gpu")
packages, indexURLs, err := config.PythonPackagesForArch("gpu", "", "")
require.NoError(t, err)
expectedPackages := []string{
"torch==1.7.1+cu101",
Expand Down Expand Up @@ -116,7 +116,7 @@ func TestPythonPackagesForArchTorchCPU(t *testing.T) {
require.Equal(t, "10.1", config.Environment.CUDA)
require.Equal(t, "8", config.Environment.CuDNN)

packages, indexURLs, err := config.PythonPackagesForArch("gpu")
packages, indexURLs, err := config.PythonPackagesForArch("gpu", "", "")
require.NoError(t, err)
expectedPackages := []string{
"torch==1.7.1+cu101",
Expand Down Expand Up @@ -145,7 +145,7 @@ func TestPythonPackagesForArchTensorflowGPU(t *testing.T) {
require.Equal(t, "10.0", config.Environment.CUDA)
require.Equal(t, "7", config.Environment.CuDNN)

packages, indexURLs, err := config.PythonPackagesForArch("gpu")
packages, indexURLs, err := config.PythonPackagesForArch("gpu", "", "")
require.NoError(t, err)
expectedPackages := []string{
"tensorflow_gpu==1.15.0",
Expand Down
3 changes: 2 additions & 1 deletion pkg/server/build.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net/http"
"os"
"path/filepath"
"runtime"
"time"

"github.com/mholt/archiver/v3"
Expand Down Expand Up @@ -148,7 +149,7 @@ func (s *Server) buildDockerImages(dir string, config *model.Config, name string
for _, arch := range config.Environment.Architectures {
logWriter.WriteStatus("Building %s image", arch)

generator := &docker.DockerfileGenerator{Config: config, Arch: arch}
generator := &docker.DockerfileGenerator{Config: config, Arch: arch, GOOS: runtime.GOOS, GOARCH: runtime.GOARCH}
dockerfileContents, err := generator.Generate()
if err != nil {
return nil, fmt.Errorf("Failed to generate Dockerfile for %s: %w", arch, err)
Expand Down
4 changes: 2 additions & 2 deletions pkg/serving/test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ func TestModel(servingPlatform Platform, imageTag string, config *model.Config,
if example.Output != "" {
if strings.HasPrefix(example.Output, "@") {
outputIsFile = true
expectedOutput, err = os.ReadFile(filepath.Join(dir, example.Output[1:]))
expectedOutput, err = os.ReadFile(filepath.Join(dir, config.Workdir, example.Output[1:]))
if err != nil {
return nil, fmt.Errorf("Failed to read example output file %s: %w", example.Output[1:], err)
}
Expand All @@ -44,7 +44,7 @@ func TestModel(servingPlatform Platform, imageTag string, config *model.Config,
}
}

input := NewExampleWithBaseDir(example.Input, dir)
input := NewExampleWithBaseDir(example.Input, filepath.Join(dir, config.Workdir))

result, err := deployment.RunInference(input, logWriter)
if err != nil {
Expand Down