@@ -3,18 +3,19 @@ package flytek8s
33import (
44 "context"
55
6- "github.com/flyteorg/flytestdlib/logger "
6+ "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core "
77
8+ "github.com/flyteorg/flyteplugins/go/tasks/errors"
9+ pluginscore "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
810 "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/template"
9- "k8s.io/apimachinery/pkg/util/validation"
11+ "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
12+
13+ "github.com/flyteorg/flytestdlib/logger"
1014
11- "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
1215 v1 "k8s.io/api/core/v1"
1316 "k8s.io/apimachinery/pkg/api/resource"
1417 "k8s.io/apimachinery/pkg/util/rand"
15-
16- "github.com/flyteorg/flyteplugins/go/tasks/errors"
17- "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
18+ "k8s.io/apimachinery/pkg/util/validation"
1819)
1920
2021const resourceGPU = "gpu"
@@ -193,22 +194,16 @@ func ApplyResourceOverrides(resources, platformResources v1.ResourceRequirements
193194 return resources
194195}
195196
196- // ToK8sContainer transforms a task template target of type core.Container into a bare-bones kubernetes container, which
197- // can be further modified with flyte-specific customizations specified by various static and run-time attributes.
198- func ToK8sContainer (ctx context.Context , taskContainer * core.Container , iFace * core.TypedInterface , parameters template.Parameters ) (* v1.Container , error ) {
199- // Perform preliminary validations
200- if parameters .TaskExecMetadata .GetOverrides () == nil {
201- return nil , errors .Errorf (errors .BadTaskSpecification , "platform/compiler error, overrides not set for task" )
202- }
203- if parameters .TaskExecMetadata .GetOverrides () == nil || parameters .TaskExecMetadata .GetOverrides ().GetResources () == nil {
204- return nil , errors .Errorf (errors .BadTaskSpecification , "resource requirements not found for container task, required!" )
205- }
197+ // BuildRawContainer constructs a Container based on the definition passed by the taskContainer and
198+ // TaskExecutionMetadata.
199+ func BuildRawContainer (ctx context.Context , taskContainer * core.Container , taskExecMetadata pluginscore.TaskExecutionMetadata ) (* v1.Container , error ) {
206200 // Make the container name the same as the pod name, unless it violates K8s naming conventions
207201 // Container names are subject to the DNS-1123 standard
208- containerName := parameters . TaskExecMetadata .GetTaskExecutionID ().GetGeneratedName ()
202+ containerName := taskExecMetadata .GetTaskExecutionID ().GetGeneratedName ()
209203 if errs := validation .IsDNS1123Label (containerName ); len (errs ) > 0 {
210204 containerName = rand .String (4 )
211205 }
206+
212207 container := & v1.Container {
213208 Name : containerName ,
214209 Image : taskContainer .GetImage (),
@@ -217,12 +212,49 @@ func ToK8sContainer(ctx context.Context, taskContainer *core.Container, iFace *c
217212 Env : ToK8sEnvVar (taskContainer .GetEnv ()),
218213 TerminationMessagePolicy : v1 .TerminationMessageFallbackToLogsOnError ,
219214 }
220- if err := AddCoPilotToContainer (ctx , config .GetK8sPluginConfig ().CoPilot , container , iFace , taskContainer .DataConfig ); err != nil {
215+
216+ return container , nil
217+ }
218+
219+ // ToK8sContainer builds a Container based on the definition passed by the TaskExecutionContext. This involves applying
220+ // all Flyte configuration including k8s plugins and resource requests.
221+ func ToK8sContainer (ctx context.Context , tCtx pluginscore.TaskExecutionContext ) (* v1.Container , error ) {
222+ taskTemplate , err := tCtx .TaskReader ().Read (ctx )
223+ if err != nil {
224+ logger .Warnf (ctx , "failed to read task information when trying to construct container, err: %s" , err .Error ())
225+ return nil , err
226+ }
227+
228+ // validate arguments
229+ if taskTemplate .GetContainer () == nil {
230+ return nil , errors .Errorf (errors .BadTaskSpecification , "unable to create container with no definition in TaskTemplate" )
231+ }
232+ if tCtx .TaskExecutionMetadata ().GetOverrides () == nil || tCtx .TaskExecutionMetadata ().GetOverrides ().GetResources () == nil {
233+ return nil , errors .Errorf (errors .BadTaskSpecification , "resource requirements not found for container task, required!" )
234+ }
235+
236+ // build raw container
237+ container , err := BuildRawContainer (ctx , taskTemplate .GetContainer (), tCtx .TaskExecutionMetadata ())
238+ if err != nil {
221239 return nil , err
222240 }
241+
223242 if container .SecurityContext == nil && config .GetK8sPluginConfig ().DefaultSecurityContext != nil {
224243 container .SecurityContext = config .GetK8sPluginConfig ().DefaultSecurityContext .DeepCopy ()
225244 }
245+
246+ // add flyte resource customizations to the container
247+ templateParameters := template.Parameters {
248+ TaskExecMetadata : tCtx .TaskExecutionMetadata (),
249+ Inputs : tCtx .InputReader (),
250+ OutputPath : tCtx .OutputWriter (),
251+ Task : tCtx .TaskReader (),
252+ }
253+
254+ if err := AddFlyteCustomizationsToContainer (ctx , templateParameters , ResourceCustomizationModeAssignResources , container ); err != nil {
255+ return nil , err
256+ }
257+
226258 return container , nil
227259}
228260
0 commit comments