<|start_filename|>pkg/image/build.go<|end_filename|>
package image

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/replicate/cog/pkg/config"
	"github.com/replicate/cog/pkg/docker"
	"github.com/replicate/cog/pkg/dockerfile"
	"github.com/replicate/cog/pkg/global"
	"github.com/replicate/cog/pkg/util/console"
)

// Build a Cog model from a config
//
// This is separated out from docker.Build(), so that can be as close as possible to the behavior of 'docker build'.
func Build(cfg *config.Config, dir, imageName string, progressOutput string) error {
	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)

	generator, err := dockerfile.NewGenerator(cfg, dir)
	if err != nil {
		return fmt.Errorf("Error creating Dockerfile generator: %w", err)
	}
	defer func() {
		if err := generator.Cleanup(); err != nil {
			console.Warnf("Error cleaning up Dockerfile generator: %s", err)
		}
	}()

	dockerfileContents, err := generator.Generate()
	if err != nil {
		return fmt.Errorf("Failed to generate Dockerfile: %w", err)
	}

	if err := docker.Build(dir, dockerfileContents, imageName, progressOutput); err != nil {
		return fmt.Errorf("Failed to build Docker image: %w", err)
	}

	console.Info("Adding labels to image...")
	signature, err := GetTypeSignature(imageName, cfg.Build.GPU)
	if err != nil {
		return fmt.Errorf("Failed to get type signature: %w", err)
	}
	signatureJSON, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("Failed to convert type signature to JSON: %w", err)
	}
	configJSON, err := json.Marshal(cfg)
	if err != nil {
		return fmt.Errorf("Failed to convert config to JSON: %w", err)
	}
	// We used to set the cog_version and config labels in Dockerfile, because we didn't require running the
	// built image to get those. But, the escaping of JSON inside a label inside a Dockerfile was gnarly, and
	// doesn't seem to be a problem here, so do it here instead.
	labels := map[string]string{
		global.LabelNamespace + "cog_version":    global.Version,
		global.LabelNamespace + "config":         string(bytes.TrimSpace(configJSON)),
		global.LabelNamespace + "type_signature": string(signatureJSON),
	}
	if err := docker.BuildAddLabelsToImage(imageName, labels); err != nil {
		return fmt.Errorf("Failed to add labels to image: %w", err)
	}
	return nil
}

func BuildBase(cfg *config.Config, dir string, progressOutput string) (string, error) {
	// TODO: better image management so we don't eat up disk space
	// https://github.com/replicate/cog/issues/80
	imageName := config.BaseDockerImageName(dir)

	console.Info("Building Docker image from environment in cog.yaml...")
	generator, err := dockerfile.NewGenerator(cfg, dir)
	if err != nil {
		return "", fmt.Errorf("Error creating Dockerfile generator: %w", err)
	}
	defer func() {
		if err := generator.Cleanup(); err != nil {
			console.Warnf("Error cleaning up Dockerfile generator: %s", err)
		}
	}()
	dockerfileContents, err := generator.GenerateBase()
	if err != nil {
		return "", fmt.Errorf("Failed to generate Dockerfile: %w", err)
	}
	if err := docker.Build(dir, dockerfileContents, imageName, progressOutput); err != nil {
		return "", fmt.Errorf("Failed to build Docker image: %w", err)
	}
	return imageName, nil
}


<|start_filename|>pkg/predict/input.go<|end_filename|>
package predict

import (
	"path/filepath"
	"strings"

	"github.com/mitchellh/go-homedir"
	"github.com/replicate/cog/pkg/util/console"
)

type Input struct {
	String *string
	File   *string
}

type Inputs map[string]Input

func NewInputs(keyVals map[string]string) Inputs {
	input := Inputs{}
	for key, val := range keyVals {
		val := val
		if strings.HasPrefix(val, "@") {
			val = val[1:]
			expandedVal, err := homedir.Expand(val)
			if err != nil {
				// FIXME: handle this better?
				console.Warnf("Error expanding homedir: %s", err)
			} else {
				val = expandedVal
			}

			input[key] = Input{File: &val}
		} else {
			input[key] = Input{String: &val}
		}
	}
	return input
}

func NewInputsWithBaseDir(keyVals map[string]string, baseDir string) Inputs {
	input := Inputs{}
	for key, val := range keyVals {
		val := val
		if strings.HasPrefix(val, "@") {
			val = filepath.Join(baseDir, val[1:])
			input[key] = Input{File: &val}
		} else {
			input[key] = Input{String: &val}
		}
	}
	return input
}


<|start_filename|>pkg/predict/predictor.go<|end_filename|>
package predict

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/replicate/cog/pkg/docker"
	"github.com/replicate/cog/pkg/global"
	"github.com/replicate/cog/pkg/util/console"
	"github.com/replicate/cog/pkg/util/shell"
)

type Predictor struct {
	runOptions docker.RunOptions

	// Running state
	containerID string
	port        int
}

func NewPredictor(runOptions docker.RunOptions) Predictor {
	return Predictor{runOptions: runOptions}
}

func (p *Predictor) Start(logsWriter io.Writer) error {
	var err error
	p.port, err = shell.NextFreePort(5000 + rand.Intn(1000))
	if err != nil {
		return err
	}

	containerPort := 5000

	p.runOptions.Ports = append(p.runOptions.Ports, docker.Port{HostPort: p.port, ContainerPort: containerPort})

	p.containerID, err = docker.RunDaemon(p.runOptions)
	if err != nil {
		return fmt.Errorf("Failed to start container: %w", err)
	}
	go func() {
		if err := docker.ContainerLogsFollow(p.containerID, logsWriter); err != nil {
			console.Warnf("Error getting container logs: %s", err)
		}
	}()

	return p.waitForContainerReady()
}

func (p *Predictor) waitForContainerReady() error {
	url := fmt.Sprintf("http://localhost:%d/ping", p.port)

	start := time.Now()
	for {
		now := time.Now()
		if now.Sub(start) > global.StartupTimeout {
			return fmt.Errorf("Timed out")
		}

		time.Sleep(100 * time.Millisecond)

		cont, err := docker.ContainerInspect(p.containerID)
		if err != nil {
			return fmt.Errorf("Failed to get container status: %w", err)
		}
		if cont.State != nil && (cont.State.Status == "exited" || cont.State.Status == "dead") {
			return fmt.Errorf("Container exited unexpectedly")
		}

		resp, err := http.Get(url)
		if err != nil {
			continue
		}
		if resp.StatusCode != http.StatusOK {
			continue
		}
		return nil
	}
}

func (p *Predictor) Stop() error {
	return docker.Stop(p.containerID)
}

func (p *Predictor) Predict(inputs Inputs) (*Output, error) {
	bodyBuffer := new(bytes.Buffer)

	mwriter := multipart.NewWriter(bodyBuffer)
	for key, val := range inputs {
		if val.File != nil {
			w, err := mwriter.CreateFormFile(key, filepath.Base(*val.File))
			if err != nil {
				return nil, err
			}
			file, err := os.Open(*val.File)
			if err != nil {
				return nil, err
			}
			if _, err := io.Copy(w, file); err != nil {
				return nil, err
			}
			if err := file.Close(); err != nil {
				return nil, err
			}
		} else {
			w, err := mwriter.CreateFormField(key)
			if err != nil {
				return nil, err
			}
			if _, err = w.Write([]byte(*val.String)); err != nil {
				return nil, err
			}
		}
	}
	if err := mwriter.Close(); err != nil {
		return nil, fmt.Errorf("Failed to close form mime writer: %w", err)
	}

	url := fmt.Sprintf("http://localhost:%d/predict", p.port)
	req, err := http.NewRequest(http.MethodPost, url, bodyBuffer)
	if err != nil {
		return nil, fmt.Errorf("Failed to create HTTP request to %s: %w", url, err)
	}
	req.Header.Set("Content-Type", mwriter.FormDataContentType())
	req.Close = true

	httpClient := &http.Client{}
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Failed to POST HTTP request to %s: %w", url, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusBadRequest {
		body := struct {
			Message string `json:"message"`
		}{}
		if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
			return nil, fmt.Errorf("/predict call return status 400, and the response body failed to decode: %w", err)
		}
		if body.Message == "" {
			return nil, fmt.Errorf("Bad request")
		}
		return nil, fmt.Errorf("Bad request: %s", body.Message)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("/predict call returned status %d", resp.StatusCode)
	}

	contentType := resp.Header.Get("Content-Type")
	mimeType := strings.Split(contentType, ";")[0]

	buf := new(bytes.Buffer)
	if _, err := io.Copy(buf, resp.Body); err != nil {
		return nil, fmt.Errorf("Failed to read response: %w", err)
	}

	setupTime := -1.0
	runTime := -1.0
	setupTimeStr := resp.Header.Get("X-Setup-Time")
	if setupTimeStr != "" {
		setupTime, err = strconv.ParseFloat(setupTimeStr, 64)
		if err != nil {
			console.Errorf("Failed to parse setup time '%s' as float: %s", setupTimeStr, err)
		}
	}
	runTimeStr := resp.Header.Get("X-Run-Time")
	if runTimeStr != "" {
		runTime, err = strconv.ParseFloat(runTimeStr, 64)
		if err != nil {
			console.Errorf("Failed to parse run time '%s' as float: %s", runTimeStr, err)
		}
	}

	output := &Output{
		Values: map[string]OutputValue{
			// TODO(andreas): support multiple outputs?
			"output": {
				Buffer:   buf,
				MimeType: mimeType,
			},
		},
		SetupTime: setupTime,
		RunTime:   runTime,
	}
	return output, nil
}

func (p *Predictor) Help() (*HelpResponse, error) {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://localhost:%d/help", p.port), nil)
	if err != nil {
		return nil, fmt.Errorf("Failed to create GET request: %w", err)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Failed to GET /help: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("/help call returned status %d", resp.StatusCode)
	}

	help := new(HelpResponse)
	if err := json.NewDecoder(resp.Body).Decode(help); err != nil {
		return nil, fmt.Errorf("Failed to parse /help body: %w", err)
	}

	return help, nil
}


<|start_filename|>pkg/predict/output.go<|end_filename|>
package predict

import "io"

type OutputValue struct {
	Buffer   io.Reader
	MimeType string
}

type Output struct {
	Values    map[string]OutputValue
	SetupTime float64
	RunTime   float64
}


<|start_filename|>pkg/image/type_signature.go<|end_filename|>
package image

import (
	"bytes"
	"encoding/json"

	"github.com/replicate/cog/pkg/docker"
	"github.com/replicate/cog/pkg/util/console"
)

type InputType string

const (
	InputTypeString InputType = "str"
	InputTypeInt    InputType = "int"
	InputTypeFloat  InputType = "float"
	InputTypeBool   InputType = "bool"
	InputTypePath   InputType = "Path"
)

type Input struct {
	Name    string    `json:"name"`
	Type    InputType `json:"type,omitempty"`
	Default *string   `json:"default,omitempty"`
	Min     *string   `json:"min,omitempty"`
	Max     *string   `json:"max,omitempty"`
	Options *[]string `json:"options,omitempty"`
	Help    *string   `json:"help,omitempty"`
}

type TypeSignature struct {
	Inputs []Input `json:"inputs,omitempty"`
}

func GetTypeSignature(imageName string, enableGPU bool) (*TypeSignature, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	// FIXME(bfirsh): we could detect this by reading the config label on the image
	gpus := ""
	if enableGPU {
		gpus = "all"
	}

	err := docker.RunWithIO(docker.RunOptions{
		Image: imageName,
		Args: []string{
			"python", "-m", "cog.command.type_signature",
		},
		GPUs: gpus,
	}, nil, &stdout, &stderr)

	if enableGPU && err == docker.ErrMissingDeviceDriver {
		console.Debug(stdout.String())
		console.Debug(stderr.String())
		console.Debug("Missing device driver, re-trying without GPU")
		return GetTypeSignature(imageName, false)
	}

	if err != nil {
		console.Info(stdout.String())
		console.Info(stderr.String())
		return nil, err
	}
	var signature *TypeSignature
	if err := json.Unmarshal(stdout.Bytes(), &signature); err != nil {
		// Exit code was 0, but JSON was not returned.
		// This is verbose, but print so anything that gets printed in Python bubbles up here.
		console.Info(stdout.String())
		console.Info(stderr.String())
		return nil, err
	}
	return signature, nil
}
