From d93844e0beab51e88473a9fdac70ff8d038e016e Mon Sep 17 00:00:00 2001 From: andreasjansson Date: Thu, 22 Apr 2021 15:52:40 -0700 Subject: [PATCH 1/2] Move helper scripts and fix workdir Signed-off-by: andreasjansson --- pkg/docker/generate.go | 23 +++++++++++++++++------ pkg/model/config.go | 2 +- pkg/serving/test.go | 4 ++-- 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/pkg/docker/generate.go b/pkg/docker/generate.go index 1fca101da6..bd73f45dc8 100644 --- a/pkg/docker/generate.go +++ b/pkg/docker/generate.go @@ -149,13 +149,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") @@ -165,7 +169,7 @@ 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] @@ -173,6 +177,9 @@ func (g *DockerfileGenerator) queueWorkerHelperScript() string { 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") @@ -221,15 +228,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 { diff --git a/pkg/model/config.go b/pkg/model/config.go index 849817a9da..6c8d77dae7 100644 --- a/pkg/model/config.go +++ b/pkg/model/config.go @@ -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"` @@ -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 { diff --git a/pkg/serving/test.go b/pkg/serving/test.go index 3ee72e58cd..b36950353f 100644 --- a/pkg/serving/test.go +++ b/pkg/serving/test.go @@ -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.Environment.Workdir, example.Output[1:])) if err != nil { return nil, fmt.Errorf("Failed to read example output file %s: %w", example.Output[1:], err) } @@ -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.Environment.Workdir)) result, err := deployment.RunInference(input, logWriter) if err != nil { From 90e06c4aab02184a0d7a2cd28f11c7edbe5c6a48 Mon Sep 17 00:00:00 2001 From: andreasjansson Date: Thu, 22 Apr 2021 16:15:30 -0700 Subject: [PATCH 2/2] Move workdir to top level of config Plus make tests pass on M1 Signed-off-by: andreasjansson --- pkg/cli/debug.go | 3 +- pkg/cli/test.go | 5 ++-- pkg/docker/generate.go | 7 ++++- pkg/docker/generate_test.go | 58 ++++++++++++++----------------------- pkg/model/compatibility.go | 13 ++++----- pkg/model/config.go | 10 +++---- pkg/model/config_test.go | 6 ++-- pkg/server/build.go | 3 +- pkg/serving/test.go | 4 +-- 9 files changed, 51 insertions(+), 58 deletions(-) diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go index 36fff825ee..dac6b4cdfb 100644 --- a/pkg/cli/debug.go +++ b/pkg/cli/debug.go @@ -5,6 +5,7 @@ import ( "io/ioutil" "os" "path" + "runtime" "github.com/spf13/cobra" @@ -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 diff --git a/pkg/cli/test.go b/pkg/cli/test.go index 03f7d4adf8..e460784868 100644 --- a/pkg/cli/test.go +++ b/pkg/cli/test.go @@ -2,8 +2,9 @@ package cli import ( "fmt" - "path/filepath" "os" + "path/filepath" + "runtime" "github.com/spf13/cobra" @@ -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) diff --git a/pkg/docker/generate.go b/pkg/docker/generate.go index bd73f45dc8..ac893bba7d 100644 --- a/pkg/docker/generate.go +++ b/pkg/docker/generate.go @@ -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) { @@ -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 { @@ -198,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 } diff --git a/pkg/docker/generate_test.go b/pkg/docker/generate_test.go index ddb02c3ba9..5d0792cffc 100644 --- a/pkg/docker/generate_test.go +++ b/pkg/docker/generate_test.go @@ -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 @@ -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) @@ -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 @@ -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` +} diff --git a/pkg/model/compatibility.go b/pkg/model/compatibility.go index c9b1588417..f82f101272 100644 --- a/pkg/model/compatibility.go +++ b/pkg/model/compatibility.go @@ -4,7 +4,6 @@ import ( _ "embed" "encoding/json" "fmt" - "runtime" "sort" "strings" @@ -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 } } @@ -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) @@ -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 diff --git a/pkg/model/config.go b/pkg/model/config.go index 6c8d77dae7..49048a5bb0 100644 --- a/pkg/model/config.go +++ b/pkg/model/config.go @@ -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 } @@ -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 @@ -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 } @@ -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 } diff --git a/pkg/model/config_test.go b/pkg/model/config_test.go index d609de1291..87d4a7012a 100644 --- a/pkg/model/config_test.go +++ b/pkg/model/config_test.go @@ -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", @@ -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", @@ -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", diff --git a/pkg/server/build.go b/pkg/server/build.go index 33c997c4aa..4b3c2ceaa6 100644 --- a/pkg/server/build.go +++ b/pkg/server/build.go @@ -10,6 +10,7 @@ import ( "net/http" "os" "path/filepath" + "runtime" "time" "github.com/mholt/archiver/v3" @@ -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) diff --git a/pkg/serving/test.go b/pkg/serving/test.go index b36950353f..f4c861345d 100644 --- a/pkg/serving/test.go +++ b/pkg/serving/test.go @@ -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, config.Environment.Workdir, 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) } @@ -44,7 +44,7 @@ func TestModel(servingPlatform Platform, imageTag string, config *model.Config, } } - input := NewExampleWithBaseDir(example.Input, filepath.Join(dir, config.Environment.Workdir)) + input := NewExampleWithBaseDir(example.Input, filepath.Join(dir, config.Workdir)) result, err := deployment.RunInference(input, logWriter) if err != nil {