Skip to content

Commit 7f0662b

Browse files
skipiskipi
authored
Zebra machine type translation (#674)
## πŸ“ Description <!-- Describe your changes in detail --> ## βœ… Checklist - [x] I have tested this change - [ ] This change requires documentation update
1 parent 77001c1 commit 7f0662b

File tree

4 files changed

+186
-4
lines changed

4 files changed

+186
-4
lines changed

β€Žzebra/lib/zebra/workers/agent/hosted_agent.exβ€Ž

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,15 @@ defmodule Zebra.Workers.Agent.HostedAgent do
6262
def occupy(job) do
6363
Task.async(fn ->
6464
Watchman.benchmark("zebra.external.chmura.occupy", fn ->
65+
{machine_type, machine_os_image} = translate_machine(job)
66+
6567
request =
6668
OccupyAgentRequest.new(
6769
request_id: job.id,
6870
machine:
6971
InternalApi.Chmura.Agent.Machine.new(
70-
type: job.machine_type,
71-
os_image: job.machine_os_image
72+
type: machine_type,
73+
os_image: machine_os_image
7274
)
7375
)
7476

@@ -122,4 +124,59 @@ defmodule Zebra.Workers.Agent.HostedAgent do
122124
{:ok, channel} = GRPC.Stub.connect(endpoint)
123125
channel
124126
end
127+
128+
@spec translate_machine(job :: Zebra.Models.Job.t()) :: {String.t(), String.t()}
129+
defp translate_machine(job) do
130+
original_type = job.machine_type || ""
131+
original_os_image = job.machine_os_image || ""
132+
133+
cond do
134+
not e1_family?(original_type) ->
135+
{original_type, original_os_image}
136+
137+
not migration_enabled?(job.organization_id) ->
138+
{original_type, original_os_image}
139+
140+
true ->
141+
{new_type, new_os_image} = map_machine_type(original_type, original_os_image)
142+
143+
Watchman.increment(
144+
{"zebra.occupy.translation", [original_type, new_type, job.organization_id]}
145+
)
146+
147+
{new_type, new_os_image}
148+
end
149+
end
150+
151+
@spec map_machine_type(String.t(), String.t()) :: {String.t(), String.t()}
152+
defp map_machine_type(original_type, original_os_image) do
153+
original_type
154+
|> case do
155+
"e1-standard-2" ->
156+
{"f1-standard-2", original_os_image}
157+
158+
"e1-standard-4" ->
159+
{"f1-standard-2", original_os_image}
160+
161+
"e1-standard-8" ->
162+
{"f1-standard-4", original_os_image}
163+
164+
_ ->
165+
{original_type, original_os_image}
166+
end
167+
end
168+
169+
@spec migration_enabled?(nil | String.t()) :: boolean()
170+
defp migration_enabled?(nil), do: false
171+
172+
defp migration_enabled?(org_id) do
173+
FeatureProvider.feature_enabled?("e1_to_f1_migration", param: org_id)
174+
end
175+
176+
@spec e1_family?(String.t()) :: boolean()
177+
defp e1_family?(machine_type) when is_binary(machine_type) do
178+
String.starts_with?(machine_type, "e1-")
179+
end
180+
181+
defp e1_family?(_), do: false
125182
end

β€Žzebra/test/support/stubbed_provider.exβ€Ž

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,22 @@
11
defmodule Support.StubbedProvider do
22
use FeatureProvider.Provider
33

4+
@e1_to_f1_org_id "org-e1-to-f1-enabled"
5+
46
@impl FeatureProvider.Provider
57
def provide_features(org_id \\ nil, _opts \\ []) do
68
{:ok,
79
[
810
feature("max_paralellism_in_org", [:enabled, {:quantity, 500}]),
911
feature("cache_cli_parallel_archive_method", [:hidden]),
1012
feature("some_custom_feature", [:hidden]),
11-
max_job_time_limit_feature(org_id)
13+
max_job_time_limit_feature(org_id),
14+
feature("e1_to_f1_migration", e1_to_f1_traits(org_id))
1215
]}
1316
end
1417

18+
def e1_to_f1_org_id, do: @e1_to_f1_org_id
19+
1520
defp max_job_time_limit_feature("enabled_30") do
1621
feature("max_job_execution_time_limit", [:enabled, {:quantity, 30}])
1722
end
@@ -24,6 +29,9 @@ defmodule Support.StubbedProvider do
2429
feature("max_job_execution_time_limit", [:hidden])
2530
end
2631

32+
defp e1_to_f1_traits(@e1_to_f1_org_id), do: [:enabled]
33+
defp e1_to_f1_traits(_org_id), do: [:hidden]
34+
2735
@impl FeatureProvider.Provider
2836
def provide_machines(_org_id \\ nil, _opts \\ []) do
2937
{:ok,

β€Žzebra/test/zebra/workers/agent/hosted_agent.exsβ€Ž

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,123 @@ defmodule Zebra.Workers.Agent.HostedAgentTest do
5454
assert Agent.occupy(job) == {:error, error.message}
5555
end
5656
end
57+
58+
test "uses original machine values when migration flag is disabled" do
59+
{:ok, job} =
60+
Support.Factories.Job.create(:scheduled, %{
61+
machine_type: "e1-standard-2",
62+
machine_os_image: "ubuntu2004"
63+
})
64+
65+
response = %InternalApi.Chmura.OccupyAgentResponse{
66+
agent: %InternalApi.Chmura.Agent{id: @agent_id}
67+
}
68+
69+
with_mock Stub,
70+
occupy_agent: fn _, request, _ ->
71+
assert request.machine.type == "e1-standard-2"
72+
assert request.machine.os_image == "ubuntu2004"
73+
{:ok, response}
74+
end do
75+
assert Agent.occupy(job) == {:ok, Agent.construct_agent(response)}
76+
end
77+
end
78+
79+
test "remaps e1-standard-2 machines to f1-standard-2 when migration flag is enabled" do
80+
org_id = Support.StubbedProvider.e1_to_f1_org_id()
81+
82+
{:ok, job} =
83+
Support.Factories.Job.create(:scheduled, %{
84+
organization_id: org_id,
85+
machine_type: "e1-standard-2",
86+
machine_os_image: "ubuntu2004"
87+
})
88+
89+
response = %InternalApi.Chmura.OccupyAgentResponse{
90+
agent: %InternalApi.Chmura.Agent{id: @agent_id}
91+
}
92+
93+
with_mock Stub,
94+
occupy_agent: fn _, request, _ ->
95+
assert request.machine.type == "f1-standard-2"
96+
assert request.machine.os_image == "ubuntu2004"
97+
{:ok, response}
98+
end do
99+
assert Agent.occupy(job) == {:ok, Agent.construct_agent(response)}
100+
end
101+
end
102+
103+
test "remaps e1-standard-4 machines to f1-standard-2 when migration flag is enabled" do
104+
org_id = Support.StubbedProvider.e1_to_f1_org_id()
105+
106+
{:ok, job} =
107+
Support.Factories.Job.create(:scheduled, %{
108+
organization_id: org_id,
109+
machine_type: "e1-standard-4",
110+
machine_os_image: "ubuntu2004"
111+
})
112+
113+
response = %InternalApi.Chmura.OccupyAgentResponse{
114+
agent: %InternalApi.Chmura.Agent{id: @agent_id}
115+
}
116+
117+
with_mock Stub,
118+
occupy_agent: fn _, request, _ ->
119+
assert request.machine.type == "f1-standard-2"
120+
assert request.machine.os_image == "ubuntu2004"
121+
{:ok, response}
122+
end do
123+
assert Agent.occupy(job) == {:ok, Agent.construct_agent(response)}
124+
end
125+
end
126+
127+
test "remaps e1-standard-8 machines to f1-standard-4 when migration flag is enabled" do
128+
org_id = Support.StubbedProvider.e1_to_f1_org_id()
129+
130+
{:ok, job} =
131+
Support.Factories.Job.create(:scheduled, %{
132+
organization_id: org_id,
133+
machine_type: "e1-standard-8",
134+
machine_os_image: "ubuntu2004"
135+
})
136+
137+
response = %InternalApi.Chmura.OccupyAgentResponse{
138+
agent: %InternalApi.Chmura.Agent{id: @agent_id}
139+
}
140+
141+
with_mock Stub,
142+
occupy_agent: fn _, request, _ ->
143+
assert request.machine.type == "f1-standard-4"
144+
assert request.machine.os_image == "ubuntu2004"
145+
{:ok, response}
146+
end do
147+
assert Agent.occupy(job) == {:ok, Agent.construct_agent(response)}
148+
end
149+
end
150+
151+
test "does not remap different machine types when migration flag is enabled" do
152+
org_id = Support.StubbedProvider.e1_to_f1_org_id()
153+
154+
{:ok, job} =
155+
Support.Factories.Job.create(:scheduled, %{
156+
organization_id: org_id,
157+
machine_type: "e2-standard-2",
158+
machine_os_image: "ubuntu2004"
159+
})
160+
161+
response = %InternalApi.Chmura.OccupyAgentResponse{
162+
agent: %InternalApi.Chmura.Agent{id: @agent_id}
163+
}
164+
165+
with_mock Stub,
166+
occupy_agent: fn _, request, _ ->
167+
assert request.machine.type == "e2-standard-2"
168+
assert request.machine.os_image == "ubuntu2004"
169+
{:ok, response}
170+
end do
171+
assert Agent.occupy(job) == {:ok, Agent.construct_agent(response)}
172+
end
173+
end
57174
end
58175

59176
describe ".release" do

β€Žzebra/test/zebra/workers/feature_provider_invalidator_worker_test.exsβ€Ž

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ defmodule Zebra.Workers.FeatureProviderInvalidatorWorkerTest do
9999
Worker.features_changed(callback_message)
100100

101101
{:ok, features} = FeatureProvider.list_features()
102-
assert length(features) == 4
102+
assert length(features) == 5
103103
end
104104

105105
test "when the organization feature state changes, organization feature caches are invalidated" do

0 commit comments

Comments
Β (0)