flyteplugins/go/tasks/pluginmachinery/catalog/async_client.go at 60459dcb59290bfb921bc9ce20b5d15b67ea7c4a · flyteorg/flyteplugins · GitHub

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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package catalog

import (
	"context"

	"github.com/lyft/flytestdlib/storage"

	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"

	"github.com/lyft/flytestdlib/bitarray"

	"github.com/lyft/flytestdlib/errors"

	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io"
)

type ResponseStatus uint8

const (
	ResponseStatusNotReady ResponseStatus = iota
	ResponseStatusReady
)

const (
	ErrResponseNotReady errors.ErrorCode = "RESPONSE_NOT_READY"
	ErrSystemError      errors.ErrorCode = "SYSTEM_ERROR"
)

type UploadRequest struct {
	Key              Key
	ArtifactData     io.OutputFilePaths
	DataStore        *storage.DataStore
	ArtifactMetadata Metadata
}

type ReadyHandler func(ctx context.Context, future Future)

// A generic Future interface to represent async operations results
type Future interface {
	// Gets the response status for the future. If the future represents multiple operations, the status will only be
	// ready if all of them are.
	GetResponseStatus() ResponseStatus

	// Sets a callback handler to be called when the future status changes to ready.
	OnReady(handler ReadyHandler)

	GetResponseError() error
}

// Catalog Upload future to represent async process of uploading catalog artifacts.
type UploadFuture interface {
	Future
}

// Catalog Download Request to represent async operation download request.
type DownloadRequest struct {
	Key       Key
	Target    io.OutputWriter
	DataStore *storage.DataStore
}

// Catalog download future to represent async process of downloading catalog artifacts.
type DownloadFuture interface {
	Future

	// Gets the actual response from the future. This will return an error if the future isn't ready yet.
	GetResponse() (DownloadResponse, error)
}

// Catalog download response.
type DownloadResponse interface {
	// Gets a bit set representing which items from the request were cached.
	GetCachedResults() *bitarray.BitSet

	// Gets the total size of the cached result.
	GetResultsSize() int

	// A convenience method to retrieve the number of cached items.
	GetCachedCount() int
}

// Catalog Download Request to represent async operation download request.
type DownloadArrayRequest struct {
	// Identifier is the same among all subtasks of the array
	Identifier core.Identifier
	// Cache version is the same among all subtasks of the array
	CacheVersion string
	// Interface is the same among all subtasks of the array
	TypedInterface core.TypedInterface

	dataStore *storage.DataStore

	// Base input reader to build subtasks input readers from
	BaseInputReader io.InputReader
	// Base output writer to build subtasks input readers from
	BaseTarget io.OutputWriter
	Indexes    *bitarray.BitSet
	Count      int
}

type UploadArrayRequest struct {
	// Identifier is the same among all subtasks of the array
	Identifier core.Identifier
	// Cache version is the same among all subtasks of the array
	CacheVersion string
	// Interface is the same among all subtasks of the array
	TypedInterface core.TypedInterface
	// ArtifactMetadata is the same among all subtasks of the array
	ArtifactMetadata Metadata

	dataStore *storage.DataStore

	// Base input reader to build subtasks input readers from
	BaseInputReader io.InputReader
	// Base output reader to build subtasks input readers from
	BaseArtifactData io.OutputFilePaths
	Indexes          *bitarray.BitSet
	Count            int
}

// An interface that helps async interaction with catalog service
type AsyncClient interface {
	// Returns if an entry exists for the given task and input. It returns the data as a LiteralMap
	Download(ctx context.Context, request DownloadRequest) (outputFuture DownloadFuture, err error)

	// Adds a new entry to catalog for the given task execution context and the generated output
	Upload(ctx context.Context, request UploadRequest) (putFuture UploadFuture, err error)

	// Returns if an entry exists for the given task and input. It returns the data as a LiteralMap
	DownloadArray(ctx context.Context, request DownloadArrayRequest) (outputFuture DownloadFuture, err error)

	// Adds a new entry to catalog for the given task execution context and the generated output
	UploadArray(ctx context.Context, requests UploadArrayRequest) (putFuture UploadFuture, err error)
}

var _ AsyncClient = AsyncClientImpl{}