-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Expand file tree
/
Copy pathuseLoadComputeEngine.ts
More file actions
121 lines (107 loc) · 3.44 KB
/
Copy pathuseLoadComputeEngine.ts
File metadata and controls
121 lines (107 loc) · 3.44 KB
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
import { useContext } from "react";
import { useQuery } from "react-query";
import RegistryPathContext from "../contexts/RegistryPathContext";
import { useDataMode } from "../contexts/DataModeContext";
import restFetch, { RestApiError } from "./restApiClient";
export interface ComputeEngineConfig {
type: string;
[key: string]: any;
}
export interface ComputeEngineInfo {
engineType: string;
engineClass: string;
config: ComputeEngineConfig;
featureViewCount: number;
}
export interface FeatureViewEngineInfo {
name: string;
type: string;
online: boolean;
lastMaterialized?: string;
hasOverride: boolean;
overrides?: Record<string, any>;
materializationIntervals: Array<{
startTime?: string;
endTime?: string;
start_time?: string;
end_time?: string;
}>;
}
const ENGINE_TYPE_TO_CLASS: Record<string, string> = {
local: "LocalComputeEngine",
"spark.engine": "SparkComputeEngine",
"ray.engine": "RayComputeEngine",
"flink.engine": "FlinkComputeEngine",
"snowflake.engine": "SnowflakeComputeEngine",
lambda: "LambdaComputeEngine",
k8s: "KubernetesComputeEngine",
};
function extractFeatureViewInfos(featureViews: any[]): FeatureViewEngineInfo[] {
return featureViews.map((fv: any) => ({
name: fv.name,
type: fv.type || "Batch",
online: fv.online ?? true,
lastMaterialized: fv.lastMaterialized,
hasOverride: fv.hasOverride ?? false,
overrides: fv.overrides,
materializationIntervals: fv.materializationIntervals || [],
}));
}
export function useLoadComputeEngine(projectName?: string) {
const registryUrl = useContext(RegistryPathContext);
const { fetchOptions } = useDataMode();
const enginePath =
projectName && projectName !== "all"
? `/compute_engines?project=${encodeURIComponent(projectName)}`
: "/compute_engines/all?limit=100";
const engineQuery = useQuery<any, Error>(
["rest", "compute-engine", registryUrl, projectName || "all"],
() => restFetch<any>(registryUrl, enginePath, fetchOptions),
{
enabled: !!registryUrl,
staleTime: 30_000,
retry: (failureCount, error) => {
if (error instanceof RestApiError && error.status === 403) return false;
return failureCount < 3;
},
},
);
const isPermissionDenied =
engineQuery.isError &&
engineQuery.error instanceof RestApiError &&
engineQuery.error.status === 403;
let engineInfo: ComputeEngineInfo | null = null;
let featureViewInfos: FeatureViewEngineInfo[] = [];
if (engineQuery.isSuccess && engineQuery.data) {
const engine = engineQuery.data.engine || engineQuery.data.engines?.[0];
if (engine) {
engineInfo = {
engineType: engine.engineType || "local",
engineClass:
engine.engineClass ||
ENGINE_TYPE_TO_CLASS[engine.engineType] ||
"LocalComputeEngine",
config: engine.config || { type: engine.engineType || "local" },
featureViewCount: engine.featureViewCount || 0,
};
}
const rawFvs = engineQuery.data.featureViews || [];
featureViewInfos = extractFeatureViewInfos(rawFvs);
}
if (!engineInfo) {
engineInfo = {
engineType: "local",
engineClass: "LocalComputeEngine",
config: { type: "local" },
featureViewCount: featureViewInfos.length,
};
}
return {
isLoading: engineQuery.isLoading,
isSuccess: engineQuery.isSuccess,
isError: engineQuery.isError,
isPermissionDenied,
engineInfo,
featureViewInfos,
};
}