Add dependencies to task definition (#1927)

Author: vishalbollu

docs/workloads/task/configuration.md

@@ -7,6 +7,10 @@
   definition:
     path: <string>  # path to a python file with a Task class definition, relative to the Cortex root (required)
     config: <string: value>  # arbitrary dictionary passed to the callable method of the Task class (can be overridden by config passed in job submission) (optional)
+    dependencies: # (optional)
+      pip: <string>  # relative path to requirements.txt (default: requirements.txt)
+      conda: <string>  # relative path to conda-packages.txt (default: conda-packages.txt)
+      shell: <string>  # relative path to a shell script for system package installation (default: dependencies.sh)
     python_path: <string>  # path to the root of your Python folder that will be appended to PYTHONPATH (default: folder containing cortex.yaml)
     image: <string> # docker image to use for the Task (default: quay.io/cortexlabs/python-predictor-cpu:master, quay.io/cortexlabs/python-predictor-gpu:master-cuda10.2-cudnn8, or quay.io/cortexlabs/python-predictor-inf:master based on compute)
     env: <string: string>  # dictionary of environment variables

pkg/operator/operator/k8s.go

@@ -521,6 +521,18 @@ func getTaskEnvVars(api *spec.API, container string) []kcore.EnvVar {
 				Name:  "CORTEX_PROJECT_DIR",
 				Value: path.Join(_emptyDirMountPath, "project"),
 			},
+			kcore.EnvVar{
+				Name:  "CORTEX_DEPENDENCIES_PIP",
+				Value: api.TaskDefinition.Dependencies.Pip,
+			},
+			kcore.EnvVar{
+				Name:  "CORTEX_DEPENDENCIES_CONDA",
+				Value: api.TaskDefinition.Dependencies.Conda,
+			},
+			kcore.EnvVar{
+				Name:  "CORTEX_DEPENDENCIES_SHELL",
+				Value: api.TaskDefinition.Dependencies.Shell,
+			},
 			kcore.EnvVar{
 				Name:  "CORTEX_CACHE_DIR",
 				Value: _specCacheDir,

pkg/types/spec/validations.go

@@ -311,6 +311,7 @@ func taskDefinitionValidation() *cr.StructFieldValidation {
 						AllowEmpty: true,
 					},
 				},
+				dependencyPathValidation(),
 			},
 		},
 	}

pkg/types/userconfig/api.go

@@ -65,12 +65,13 @@ type Predictor struct {
 }
 
 type TaskDefinition struct {
-	Path       string                 `json:"path" yaml:"path"`
-	PythonPath *string                `json:"python_path" yaml:"python_path"`
-	Image      string                 `json:"image" yaml:"image"`
-	LogLevel   LogLevel               `json:"log_level" yaml:"log_level"`
-	Config     map[string]interface{} `json:"config" yaml:"config"`
-	Env        map[string]string      `json:"env" yaml:"env"`
+	Path         string                 `json:"path" yaml:"path"`
+	PythonPath   *string                `json:"python_path" yaml:"python_path"`
+	Image        string                 `json:"image" yaml:"image"`
+	LogLevel     LogLevel               `json:"log_level" yaml:"log_level"`
+	Config       map[string]interface{} `json:"config" yaml:"config"`
+	Env          map[string]string      `json:"env" yaml:"env"`
+	Dependencies *Dependencies          `json:"dependencies" yaml:"dependencies"`
 }
 
 type MultiModels struct {
@@ -373,6 +374,14 @@ func (api *API) UserStr(provider types.ProviderType) string {
 	return sb.String()
 }
 
+func (dependencies Dependencies) UserStr() string {
+	var sb strings.Builder
+	sb.WriteString(fmt.Sprintf("%s: %s\n", PipKey, dependencies.Pip))
+	sb.WriteString(fmt.Sprintf("%s: %s\n", CondaKey, dependencies.Conda))
+	sb.WriteString(fmt.Sprintf("%s: %s\n", ShellKey, dependencies.Shell))
+	return sb.String()
+}
+
 func (trafficSplit *TrafficSplit) UserStr() string {
 	var sb strings.Builder
 	sb.WriteString(fmt.Sprintf("%s: %s\n", NameKey, trafficSplit.Name))
@@ -399,6 +408,8 @@ func (task *TaskDefinition) UserStr() string {
 		d, _ := yaml.Marshal(&task.Env)
 		sb.WriteString(s.Indent(string(d), "  "))
 	}
+	sb.WriteString(fmt.Sprintf("%s:\n", DependenciesKey))
+	sb.WriteString(s.Indent(task.Dependencies.UserStr(), "  "))
 
 	return sb.String()
 }
@@ -450,6 +461,9 @@ func (predictor *Predictor) UserStr() string {
 		d, _ := yaml.Marshal(&predictor.Env)
 		sb.WriteString(s.Indent(string(d), "  "))
 	}
+	sb.WriteString(fmt.Sprintf("%s:\n", DependenciesKey))
+	sb.WriteString(s.Indent(predictor.Dependencies.UserStr(), "  "))
+
 	return sb.String()
 }
 

pkg/types/userconfig/config_key.go

@@ -45,6 +45,12 @@ const (
 	ConfigKey                 = "config"
 	EnvKey                    = "env"
 
+	// Predictor/TaskDefinition.Dependencies
+	DependenciesKey = "dependencies"
+	PipKey          = "pip"
+	ShellKey        = "shell"
+	CondaKey        = "conda"
+
 	// MultiModelReloading
 	MultiModelReloadingKey = "multi_model_reloading"