Skip to content

Commit ab49223

Browse files
ptiurinptiurin
authored
feat: Internal args for engine (#135)
1 parent b614ca0 commit ab49223

File tree

4 files changed

+141
-16
lines changed

4 files changed

+141
-16
lines changed

‎src/firebolt/model/engine.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ def default(
6969
engine_type: EngineType = EngineType.GENERAL_PURPOSE,
7070
auto_stop_delay_duration: str = "1200s",
7171
warm_up: WarmupMethod = WarmupMethod.PRELOAD_INDEXES,
72+
minimum_logging_level: str = "ENGINE_SETTINGS_LOGGING_LEVEL_INFO",
7273
) -> EngineSettings:
7374
if engine_type == EngineType.GENERAL_PURPOSE:
7475
preset = engine_type.GENERAL_PURPOSE.api_settings_preset_name # type: ignore # noqa: E501
@@ -80,7 +81,7 @@ def default(
8081
return cls(
8182
preset=preset,
8283
auto_stop_delay_duration=auto_stop_delay_duration,
83-
minimum_logging_level="ENGINE_SETTINGS_LOGGING_LEVEL_INFO",
84+
minimum_logging_level=minimum_logging_level,
8485
is_read_only=is_read_only,
8586
warm_up=warm_up.api_name,
8687
)
@@ -459,3 +460,11 @@ def _send_engine_request(self, url: str) -> Engine:
459460
return Engine.parse_obj_with_service(
460461
obj=response.json()["engine"], engine_service=self._service
461462
)
463+
464+
465+
class _EngineCreateRequest(FireboltBaseModel):
466+
"""Helper model for sending Engine create requests."""
467+
468+
account_id: str
469+
engine: Engine
470+
engine_revision: Optional[EngineRevision]

‎src/firebolt/service/engine.py‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from logging import getLogger
2-
from typing import List, Optional, Union
2+
from typing import Any, Dict, List, Optional, Union
33

44
from firebolt.common.exception import FireboltError
55
from firebolt.common.urls import (
@@ -9,8 +9,7 @@
99
ENGINES_BY_IDS_URL,
1010
)
1111
from firebolt.common.util import prune_dict
12-
from firebolt.model import FireboltBaseModel
13-
from firebolt.model.engine import Engine, EngineSettings
12+
from firebolt.model.engine import Engine, EngineSettings, _EngineCreateRequest
1413
from firebolt.model.engine_revision import (
1514
EngineRevision,
1615
EngineRevisionSpecification,
@@ -116,6 +115,8 @@ def create(
116115
auto_stop: int = 20,
117116
warmup: Union[str, WarmupMethod] = WarmupMethod.PRELOAD_INDEXES,
118117
description: str = "",
118+
engine_settings_kwargs: Dict[str, Any] = {},
119+
revision_spec_kwargs: Dict[str, Any] = {},
119120
) -> Engine:
120121
"""
121122
Create a new Engine.
@@ -165,6 +166,7 @@ def create(
165166
engine_type=engine_type,
166167
auto_stop_delay_duration=f"{auto_stop * 60}s",
167168
warm_up=warmup,
169+
**engine_settings_kwargs,
168170
),
169171
)
170172

@@ -186,11 +188,8 @@ def create(
186188
specification=EngineRevisionSpecification(
187189
db_compute_instances_type_key=instance_type_key,
188190
db_compute_instances_count=scale,
189-
db_compute_instances_use_spot=False,
190-
db_version="",
191191
proxy_instances_type_key=instance_type_key,
192-
proxy_instances_count=1,
193-
proxy_version="",
192+
**revision_spec_kwargs,
194193
)
195194
)
196195

@@ -210,13 +209,6 @@ def _send_create_engine(
210209
The newly created engine.
211210
"""
212211

213-
class _EngineCreateRequest(FireboltBaseModel):
214-
"""Helper model for sending Engine create requests."""
215-
216-
account_id: str
217-
engine: Engine
218-
engine_revision: Optional[EngineRevision]
219-
220212
response = self.client.post(
221213
url=ACCOUNT_ENGINES_URL.format(account_id=self.account_id),
222214
headers={"Content-type": "application/json"},

‎tests/unit/service/conftest.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@
2222
from firebolt.model.binding import Binding, BindingKey
2323
from firebolt.model.database import Database, DatabaseKey
2424
from firebolt.model.engine import Engine, EngineKey, EngineSettings
25+
from firebolt.model.engine_revision import (
26+
EngineRevision,
27+
EngineRevisionSpecification,
28+
)
2529
from firebolt.model.instance_type import InstanceType, InstanceTypeKey
2630
from firebolt.model.region import Region
2731
from tests.unit.util import list_to_paginated_response
@@ -32,6 +36,11 @@ def engine_name() -> str:
3236
return "my_engine"
3337

3438

39+
@pytest.fixture
40+
def engine_scale() -> int:
41+
return 2
42+
43+
3544
@pytest.fixture
3645
def engine_settings() -> EngineSettings:
3746
return EngineSettings.default()
@@ -48,6 +57,22 @@ def mock_engine(engine_name, region_1, engine_settings, account_id, settings) ->
4857
)
4958

5059

60+
@pytest.fixture
61+
def mock_engine_revision_spec(
62+
instance_type_2, engine_scale
63+
) -> EngineRevisionSpecification:
64+
return EngineRevisionSpecification(
65+
db_compute_instances_type_key=instance_type_2.key,
66+
db_compute_instances_count=engine_scale,
67+
proxy_instances_type_key=instance_type_2.key,
68+
)
69+
70+
71+
@pytest.fixture
72+
def mock_engine_revision(mock_engine_revision_spec) -> EngineRevision:
73+
return EngineRevision(specification=mock_engine_revision_spec)
74+
75+
5176
@pytest.fixture
5277
def instance_type_1(provider, region_1) -> InstanceType:
5378
return InstanceType(

‎tests/unit/service/test_engine.py‎

Lines changed: 100 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
from typing import Callable, List
22

3+
from pydantic import ValidationError
34
from pytest import raises
45
from pytest_httpx import HTTPXMock
56

67
from firebolt.common import Settings
78
from firebolt.common.exception import FireboltError, NoAttachedDatabaseError
8-
from firebolt.model.engine import Engine
9+
from firebolt.model.engine import Engine, _EngineCreateRequest
10+
from firebolt.model.engine_revision import EngineRevision
911
from firebolt.model.instance_type import InstanceType
1012
from firebolt.model.region import Region
1113
from firebolt.service.manager import ResourceManager
@@ -47,6 +49,103 @@ def test_engine_create(
4749
assert engine.name == engine_name
4850

4951

52+
def test_engine_create_with_kwargs(
53+
httpx_mock: HTTPXMock,
54+
auth_callback: Callable,
55+
auth_url: str,
56+
provider_callback: Callable,
57+
provider_url: str,
58+
instance_type_region_1_callback: Callable,
59+
instance_type_region_1_url: str,
60+
region_callback: Callable,
61+
region_url: str,
62+
settings: Settings,
63+
mock_engine: Engine,
64+
engine_name: str,
65+
account_id_callback: Callable,
66+
account_id_url: str,
67+
engine_callback: Callable,
68+
engine_url: str,
69+
account_id: str,
70+
mock_engine_revision: EngineRevision,
71+
):
72+
httpx_mock.add_callback(auth_callback, url=auth_url)
73+
httpx_mock.add_callback(provider_callback, url=provider_url)
74+
httpx_mock.add_callback(
75+
instance_type_region_1_callback, url=instance_type_region_1_url
76+
)
77+
httpx_mock.add_callback(account_id_callback, url=account_id_url)
78+
httpx_mock.add_callback(auth_callback, url=auth_url)
79+
httpx_mock.add_callback(region_callback, url=region_url)
80+
# Setting to manager.engines.create defaults
81+
mock_engine.key = None
82+
mock_engine.description = ""
83+
mock_engine.endpoint = None
84+
# Testing kwargs
85+
mock_engine.settings.minimum_logging_level = "ENGINE_SETTINGS_LOGGING_LEVEL_DEBUG"
86+
mock_engine_revision.specification.proxy_version = "0.2.3"
87+
engine_content = _EngineCreateRequest(
88+
account_id=account_id, engine=mock_engine, engine_revision=mock_engine_revision
89+
)
90+
httpx_mock.add_callback(
91+
engine_callback,
92+
url=engine_url,
93+
method="POST",
94+
match_content=engine_content.json(by_alias=True).encode("ascii"),
95+
)
96+
97+
manager = ResourceManager(settings=settings)
98+
engine_settings_kwargs = {
99+
"minimum_logging_level": "ENGINE_SETTINGS_LOGGING_LEVEL_DEBUG"
100+
}
101+
revision_spec_kwargs = {"proxy_version": "0.2.3"}
102+
engine = manager.engines.create(
103+
name=engine_name,
104+
engine_settings_kwargs=engine_settings_kwargs,
105+
revision_spec_kwargs=revision_spec_kwargs,
106+
)
107+
108+
assert engine.name == engine_name
109+
110+
111+
def test_engine_create_with_kwargs_fail(
112+
httpx_mock: HTTPXMock,
113+
auth_callback: Callable,
114+
auth_url: str,
115+
provider_callback: Callable,
116+
provider_url: str,
117+
instance_type_region_1_callback: Callable,
118+
instance_type_region_1_url: str,
119+
region_callback: Callable,
120+
region_url: str,
121+
settings: Settings,
122+
engine_name: str,
123+
account_id_callback: Callable,
124+
account_id_url: str,
125+
):
126+
httpx_mock.add_callback(auth_callback, url=auth_url)
127+
httpx_mock.add_callback(provider_callback, url=provider_url)
128+
httpx_mock.add_callback(
129+
instance_type_region_1_callback, url=instance_type_region_1_url
130+
)
131+
httpx_mock.add_callback(account_id_callback, url=account_id_url)
132+
httpx_mock.add_callback(auth_callback, url=auth_url)
133+
httpx_mock.add_callback(region_callback, url=region_url)
134+
135+
manager = ResourceManager(settings=settings)
136+
revision_spec_kwargs = {"incorrect_kwarg": "val"}
137+
with raises(ValidationError):
138+
manager.engines.create(
139+
name=engine_name, revision_spec_kwargs=revision_spec_kwargs
140+
)
141+
142+
engine_settings_kwargs = {"incorrect_kwarg": "val"}
143+
with raises(TypeError):
144+
manager.engines.create(
145+
name=engine_name, engine_settings_kwargs=engine_settings_kwargs
146+
)
147+
148+
50149
def test_engine_create_no_available_types(
51150
httpx_mock: HTTPXMock,
52151
auth_callback: Callable,

0 commit comments

Comments
 (0)