async def test_prompt_cluster_default(make_client: Callable[..., Client]) -> None:
clusters = {
"first": Cluster(
registry_url=URL("https://registry-dev.neu.ro"),
storage_url=URL("https://storage-dev.neu.ro"),
users_url=URL("https://users-dev.neu.ro"),
monitoring_url=URL("https://monitoring-dev.neu.ro"),
presets={"cpu-small": Preset(cpu=1, memory_mb=1024)},
name="first",
),
"second": Cluster(
registry_url=URL("https://registry2-dev.neu.ro"),
storage_url=URL("https://storage2-dev.neu.ro"),
users_url=URL("https://users2-dev.neu.ro"),
monitoring_url=URL("https://monitoring2-dev.neu.ro"),
presets={"cpu-small": Preset(cpu=2, memory_mb=1024)},
name="second",
),
}
client = make_client("https://neu.ro", clusters=clusters)
session = mock.Mock()
loop = asyncio.get_event_loop()
fut = loop.create_future()
fut.set_result("")
session.prompt_async.return_value = fut
ret = await prompt_cluster(client, session=session)
assert ret == "first"
def nmrc_path(tmp_path: Path, token: str, auth_config: _AuthConfig) -> Path:
nmrc_path = tmp_path / "conftest.nmrc"
cluster_config = Cluster(
registry_url=URL("https://registry-dev.neu.ro"),
storage_url=URL("https://storage-dev.neu.ro"),
users_url=URL("https://users-dev.neu.ro"),
monitoring_url=URL("https://monitoring-dev.neu.ro"),
presets={
"gpu-small":
Preset(cpu=7,
memory_mb=30 * 1024,
gpu=1,
gpu_model="nvidia-tesla-k80"),
"gpu-large":
Preset(cpu=7,
memory_mb=60 * 1024,
gpu=1,
gpu_model="nvidia-tesla-v100"),
"cpu-small":
Preset(cpu=7, memory_mb=2 * 1024),
"cpu-large":
Preset(cpu=7, memory_mb=14 * 1024),
},
name="default",
)
config = _ConfigData(
auth_config=auth_config,
auth_token=_AuthToken.create_non_expiring(token),
url=URL("https://dev.neu.ro/api/v1"),
version=neuromation.__version__,
cluster_name="default",
clusters={cluster_config.name: cluster_config},
)
Factory(nmrc_path)._save(config)
return nmrc_path
def cluster_config() -> Cluster:
return Cluster(
registry_url=URL("https://registry-dev.neu.ro"),
storage_url=URL("https://storage-dev.neu.ro"),
users_url=URL("https://users-dev.neu.ro"),
monitoring_url=URL("https://monitoring-dev.neu.ro"),
presets={
"gpu-small":
Preset(cpu=7,
memory_mb=30 * 1024,
gpu=1,
gpu_model="nvidia-tesla-k80"),
"gpu-large":
Preset(cpu=7,
memory_mb=60 * 1024,
gpu=1,
gpu_model="nvidia-tesla-v100"),
"cpu-small":
Preset(cpu=7, memory_mb=2 * 1024),
"cpu-large":
Preset(cpu=7, memory_mb=14 * 1024),
"cpu-large-p":
Preset(cpu=7, memory_mb=14 * 1024, is_preemptible=True),
},
name="default",
)
def go(url_str: str,
*,
registry_url: str = "https://registry-dev.neu.ro",
trace_id: str = "bd7a977555f6b982",
clusters: Optional[Dict[str, Cluster]] = None,
token_url: Optional[URL] = None) -> Client:
url = URL(url_str)
if clusters is None:
cluster_config = Cluster(
registry_url=URL(registry_url),
monitoring_url=(url / "jobs"),
storage_url=(url / "storage"),
users_url=url,
presets={
"gpu-small":
Preset(cpu=7,
memory_mb=30 * 1024,
gpu=1,
gpu_model="nvidia-tesla-k80"),
"gpu-large":
Preset(cpu=7,
memory_mb=60 * 1024,
gpu=1,
gpu_model="nvidia-tesla-v100"),
"cpu-small":
Preset(cpu=7, memory_mb=2 * 1024),
"cpu-large":
Preset(cpu=7, memory_mb=14 * 1024),
},
name="default",
)
clusters = {cluster_config.name: cluster_config}
if token_url is not None:
real_auth_config = replace(auth_config, token_url=token_url)
else:
real_auth_config = auth_config
config = _ConfigData(
auth_config=real_auth_config,
auth_token=_AuthToken.create_non_expiring(token),
url=URL(url),
version=neuromation.__version__,
cluster_name=next(iter(clusters)),
clusters=clusters,
)
config_dir = tmp_path / ".neuro"
_save(config, config_dir)
session = aiohttp.ClientSession(trace_configs=[_make_trace_config()])
return Client._create(session, config_dir, trace_id)
async def test_clusters(
aiohttp_server: _TestServerFactory, make_client: _MakeClient
) -> None:
app = web.Application()
srv = await aiohttp_server(app)
async with make_client(srv.make_url("/")) as client:
assert client.config.clusters == {
"default": Cluster(
name="default",
registry_url=URL("https://registry-dev.neu.ro"),
storage_url=srv.make_url("/storage"),
users_url=srv.make_url("/"),
monitoring_url=srv.make_url("/jobs"),
presets=mock.ANY,
)
}
async def test_get_server_config_with_token(
aiohttp_client: _TestClientFactory) -> None:
registry_url = "https://registry.dev.neuromation.io"
storage_url = "https://storage.dev.neuromation.io"
users_url = "https://dev.neuromation.io/users"
monitoring_url = "https://dev.neuromation.io/monitoring"
auth_url = "https://dev-neuromation.auth0.com/authorize"
token_url = "https://dev-neuromation.auth0.com/oauth/token"
logout_url = "https://dev-neuromation.auth0.com/v2/logout"
client_id = "this_is_client_id"
audience = "https://platform.dev.neuromation.io"
headless_callback_url = "https://dev.neu.ro/oauth/show-code"
success_redirect_url = "https://platform.neuromation.io"
JSON = {
"name":
"default",
"registry_url":
registry_url,
"storage_url":
storage_url,
"users_url":
users_url,
"monitoring_url":
monitoring_url,
"resource_presets": [
{
"name": "gpu-small",
"cpu": 7,
"memory_mb": 30 * 1024,
"gpu": 1,
"gpu_model": "nvidia-tesla-k80",
},
{
"name": "gpu-large",
"cpu": 7,
"memory_mb": 60 * 1024,
"gpu": 1,
"gpu_model": "nvidia-tesla-v100",
},
{
"name": "cpu-small",
"cpu": 2,
"memory_mb": 2 * 1024
},
{
"name": "cpu-large",
"cpu": 3,
"memory_mb": 14 * 1024
},
],
"clusters": [{
"name":
"default",
"registry_url":
registry_url,
"storage_url":
storage_url,
"users_url":
users_url,
"monitoring_url":
monitoring_url,
"secrets_url":
"https://secrets-dev.neu.ro",
"resource_presets": [
{
"name": "gpu-small",
"cpu": 7,
"memory_mb": 30 * 1024,
"gpu": 1,
"gpu_model": "nvidia-tesla-k80",
},
{
"name": "gpu-large",
"cpu": 7,
"memory_mb": 60 * 1024,
"gpu": 1,
"gpu_model": "nvidia-tesla-v100",
},
{
"name": "cpu-small",
"cpu": 2,
"memory_mb": 2 * 1024
},
{
"name": "cpu-large",
"cpu": 3,
"memory_mb": 14 * 1024
},
],
}],
"auth_url":
auth_url,
"token_url":
token_url,
"logout_url":
logout_url,
"client_id":
client_id,
"audience":
audience,
"callback_urls": [
"http://127.0.0.1:54540",
"http://127.0.0.1:54541",
"http://127.0.0.1:54542",
],
"headless_callback_url":
headless_callback_url,
"success_redirect_url":
success_redirect_url,
}
async def handler(request: web.Request) -> web.Response:
assert request.headers["Authorization"] == "Bearer bananatoken"
return web.json_response(JSON)
app = web.Application()
app.router.add_get("/config", handler)
client = await aiohttp_client(app)
config = await get_server_config(client.session,
client.make_url("/"),
token="bananatoken")
cluster_config = Cluster(
registry_url=URL(registry_url),
storage_url=URL(storage_url),
users_url=URL(users_url),
monitoring_url=URL(monitoring_url),
secrets_url=URL("https://secrets-dev.neu.ro"),
presets={
"gpu-small":
Preset(cpu=7,
memory_mb=30 * 1024,
gpu=1,
gpu_model="nvidia-tesla-k80"),
"gpu-large":
Preset(cpu=7,
memory_mb=60 * 1024,
gpu=1,
gpu_model="nvidia-tesla-v100"),
"cpu-small":
Preset(cpu=2, memory_mb=2 * 1024),
"cpu-large":
Preset(cpu=3, memory_mb=14 * 1024),
},
name="default",
)
assert config == _ServerConfig(
auth_config=_AuthConfig(
auth_url=URL(auth_url),
token_url=URL(token_url),
logout_url=URL(logout_url),
client_id=client_id,
audience=audience,
headless_callback_url=URL(headless_callback_url),
success_redirect_url=URL(success_redirect_url),
),
clusters={"default": cluster_config},
)
async def test_fetch(
aiohttp_server: _TestServerFactory, make_client: _MakeClient
) -> None:
registry_url = "https://registry2-dev.neu.ro"
storage_url = "https://storage2-dev.neu.ro"
users_url = "https://users2-dev.neu.ro"
monitoring_url = "https://jobs2-dev.neu.ro"
auth_url = "https://dev-neuromation.auth0.com/authorize"
token_url = "https://dev-neuromation.auth0.com/oauth/token"
client_id = "this_is_client_id"
audience = "https://platform.dev.neuromation.io"
headless_callback_url = "https://dev.neu.ro/oauth/show-code"
success_redirect_url = "https://platform.neuromation.io"
JSON = {
"auth_url": auth_url,
"token_url": token_url,
"client_id": client_id,
"audience": audience,
"headless_callback_url": headless_callback_url,
"success_redirect_url": success_redirect_url,
"clusters": [
{
"name": "default",
"registry_url": registry_url,
"storage_url": storage_url,
"users_url": users_url,
"monitoring_url": monitoring_url,
"resource_presets": [
{"name": "cpu-small", "cpu": 2, "memory_mb": 2 * 1024}
],
}
],
}
async def handler(request: web.Request) -> web.Response:
return web.json_response(JSON)
app = web.Application()
app.add_routes([web.get("/config", handler)])
srv = await aiohttp_server(app)
async with make_client(srv.make_url("/")) as client:
await client.config.fetch()
assert client.config.clusters == {
"default": Cluster(
name="default",
registry_url=URL("https://registry2-dev.neu.ro"),
storage_url=URL("https://storage2-dev.neu.ro"),
users_url=URL("https://users2-dev.neu.ro"),
monitoring_url=URL("https://jobs2-dev.neu.ro"),
presets={
"cpu-small": Preset(
cpu=2,
memory_mb=2048,
is_preemptible=False,
gpu=None,
gpu_model=None,
tpu_type=None,
tpu_software_version=None,
)
},
)
}