This repository has been archived by the owner on Oct 9, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 60
/
future_file_reader.go
110 lines (95 loc) · 4.22 KB
/
future_file_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package task
import (
"context"
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
"github.com/flyteorg/flytestdlib/errors"
"github.com/flyteorg/flytestdlib/logger"
"github.com/flyteorg/flytestdlib/storage"
"github.com/flyteorg/flytepropeller/pkg/apis/flyteworkflow/v1alpha1"
"github.com/flyteorg/flytepropeller/pkg/utils"
)
// TODO this file exists only until we need to support dynamic nodes instead of closure.
// Once closure migration is done, this file should be deleted.
const implicitFutureFileName = "futures.pb"
const implicitCompileWorkflowsName = "futures_compiled.pb"
const implicitCompiledWorkflowClosureName = "dynamic_compiled.pb"
type FutureFileReader struct {
RemoteFileWorkflowStore
loc storage.DataReference
flyteWfCRDCacheLoc storage.DataReference
flyteWfClosureCacheLoc storage.DataReference
store *storage.DataStore
}
func (f FutureFileReader) Exists(ctx context.Context) (bool, error) {
metadata, err := f.store.Head(ctx, f.loc)
// If no futures file produced, then declare success and return.
if err != nil {
logger.Warnf(ctx, "Failed to read futures file. Error: %v", err)
return false, errors.Wrapf(utils.ErrorCodeUser, err, "Failed to do HEAD on futures file.")
}
return metadata.Exists(), nil
}
func (f FutureFileReader) Read(ctx context.Context) (*core.DynamicJobSpec, error) {
djSpec := &core.DynamicJobSpec{}
if err := f.store.ReadProtobuf(ctx, f.loc, djSpec); err != nil {
logger.Warnf(ctx, "Failed to read futures file. Error: %v", err)
return nil, errors.Wrapf(utils.ErrorCodeSystem, err, "Failed to read futures protobuf file.")
}
return djSpec, nil
}
func (f FutureFileReader) CacheExists(ctx context.Context) (bool, error) {
exists, err := f.RemoteFileWorkflowStore.Exists(ctx, f.flyteWfCRDCacheLoc)
if err != nil || !exists {
return exists, err
}
return f.RemoteFileWorkflowStore.Exists(ctx, f.flyteWfClosureCacheLoc)
}
func (f FutureFileReader) Cache(ctx context.Context, wf *v1alpha1.FlyteWorkflow, workflowClosure *core.CompiledWorkflowClosure) error {
err := f.RemoteFileWorkflowStore.PutFlyteWorkflowCRD(ctx, wf, f.flyteWfCRDCacheLoc)
if err != nil {
return err
}
return f.RemoteFileWorkflowStore.PutCompiledFlyteWorkflow(ctx, workflowClosure, f.flyteWfClosureCacheLoc)
}
type CacheContents struct {
WorkflowCRD *v1alpha1.FlyteWorkflow
CompiledWorkflow *core.CompiledWorkflowClosure
}
func (f FutureFileReader) RetrieveCache(ctx context.Context) (CacheContents, error) {
workflowCRD, err := f.RemoteFileWorkflowStore.GetWorkflowCRD(ctx, f.flyteWfCRDCacheLoc)
if err != nil {
return CacheContents{}, err
}
compiledWorkflow, err := f.RemoteFileWorkflowStore.GetCompiledWorkflow(ctx, f.flyteWfClosureCacheLoc)
if err != nil {
return CacheContents{}, err
}
return CacheContents{
WorkflowCRD: workflowCRD,
CompiledWorkflow: compiledWorkflow,
}, nil
}
func NewRemoteFutureFileReader(ctx context.Context, dataDir storage.DataReference, store *storage.DataStore) (FutureFileReader, error) {
loc, err := store.ConstructReference(ctx, dataDir, implicitFutureFileName)
if err != nil {
logger.Warnf(ctx, "Failed to construct data path for futures file. Error: %v", err)
return FutureFileReader{}, errors.Wrapf(utils.ErrorCodeSystem, err, "failed to construct data path")
}
flyteWfCRDCacheLoc, err := store.ConstructReference(ctx, dataDir, implicitCompileWorkflowsName)
if err != nil {
logger.Warnf(ctx, "Failed to construct data path for compile workflows file, error: %s", err)
return FutureFileReader{}, errors.Wrapf(utils.ErrorCodeSystem, err, "failed to construct reference for workflow CRD cache location")
}
flyteWfClosureCacheLoc, err := store.ConstructReference(ctx, dataDir, implicitCompiledWorkflowClosureName)
if err != nil {
logger.Warnf(ctx, "Failed to construct data path for compile workflows file, error: %s", err)
return FutureFileReader{}, errors.Wrapf(utils.ErrorCodeSystem, err, "failed to construct reference for compiled workflow closure cache location")
}
return FutureFileReader{
loc: loc,
flyteWfCRDCacheLoc: flyteWfCRDCacheLoc,
flyteWfClosureCacheLoc: flyteWfClosureCacheLoc,
store: store,
RemoteFileWorkflowStore: NewRemoteWorkflowStore(store),
}, nil
}