def steal_config_maybe(dst_path: pathlib.Path) -> None:
if NEURO_STEAL_CONFIG in os.environ:
src = pathlib.Path(os.environ[NEURO_STEAL_CONFIG])
if not src.exists():
return
dst = Factory(dst_path).path
dst.mkdir(mode=0o700)
for f in src.iterdir():
target = dst / f.name
shutil.copy(f, target)
# forbid access to other users
os.chmod(target, 0o600)
f.unlink()
src.rmdir()
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 nmrc_path(tmp_path: Path, token: str, auth_config: _AuthConfig) -> Path:
nmrc_path = tmp_path / "conftest.nmrc"
cluster_config = _ClusterConfig.create(
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"),
resource_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),
},
)
config = _Config(
auth_config=auth_config,
auth_token=_AuthToken.create_non_expiring(token),
cluster_config=cluster_config,
pypi=_PyPIVersion.create_uninitialized(),
url=URL("https://dev.neu.ro/api/v1"),
cookie_session=_CookieSession.create_uninitialized(),
version=neuromation.__version__,
)
Factory(nmrc_path)._save(config)
return nmrc_path
def convert(self, value: str, param: Optional[click.Parameter],
ctx: Optional[click.Context]) -> RemoteImage:
assert ctx is not None
root = cast(Root, ctx.obj)
config = Factory(root.config_path)._read()
image_parser = _ImageNameParser(config.auth_token.username,
config.cluster_config.registry_url)
return image_parser.parse_as_neuro_image(value, allow_tag=False)
async def test_normal_login(self, tmp_home: Path,
mock_for_login: _TestServer) -> None:
await Factory().login_with_token(token="tokenstr",
url=mock_for_login.make_url("/"))
nmrc_path = tmp_home / ".nmrc"
assert Path(nmrc_path).exists(), "Config file not written after login "
saved_config = Factory(nmrc_path)._read()
assert saved_config.auth_config.is_initialized()
assert saved_config.cluster_config.is_initialized()
async def _get_storage_cookie(nmrc_path: Optional[Path]) -> None:
async with api_get(timeout=CLIENT_TIMEOUT, path=nmrc_path) as client:
await client.storage.ls(URL("storage:/"))
cookie = client._get_session_cookie()
if cookie is not None:
new_config = dataclasses.replace(
client._config,
cookie_session=_CookieSession(cookie=cookie.value,
timestamp=int(time())),
)
Factory(nmrc_path)._save(new_config)
def factory(self) -> Factory:
if self._factory is None:
trace_configs: Optional[List[aiohttp.TraceConfig]]
if self.trace:
trace_configs = [self._create_trace_config()]
else:
trace_configs = None
self._factory = Factory(
path=self.config_path,
trace_configs=trace_configs,
trace_id=gen_trace_id(),
)
return self._factory
def convert(self, value: str, param: Optional[click.Parameter],
ctx: Optional[click.Context]) -> LocalImage:
assert ctx is not None
root = cast(Root, ctx.obj)
config = Factory(root.config_path)._read()
image_parser = _ImageNameParser(config.auth_token.username,
config.cluster_config.registry_url)
if image_parser.is_in_neuro_registry(value):
raise click.BadParameter("remote image cannot be used as local",
ctx, param, self.name)
else:
parsed_image = image_parser.parse_as_local_image(value)
return parsed_image
def _create_config(nmrc_path: Path, token: str, auth_config: _AuthConfig,
cluster_config: Cluster) -> str:
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=cluster_config.name,
clusters={cluster_config.name: cluster_config},
)
Factory(nmrc_path)._save(config)
assert nmrc_path.exists()
return token
def _create_config(
nmrc_path: Path,
token: str,
auth_config: _AuthConfig,
cluster_config: _ClusterConfig,
) -> str:
config = _Config(
auth_config=auth_config,
auth_token=_AuthToken.create_non_expiring(token),
cluster_config=cluster_config,
pypi=_PyPIVersion.create_uninitialized(),
url=URL("https://dev.neu.ro/api/v1"),
cookie_session=_CookieSession.create_uninitialized(),
version=neuromation.__version__,
)
Factory(nmrc_path)._save(config)
assert nmrc_path.exists()
return token
async def test_normal_login(self, tmp_home: Path,
mock_for_login: _TestServer) -> None:
async def get_auth_code_cb(url: URL) -> str:
assert url.with_query(None) == mock_for_login.make_url(
"/authorize")
assert dict(url.query) == dict(
response_type="code",
code_challenge=mock.ANY,
code_challenge_method="S256",
client_id="banana",
redirect_uri=str(mock_for_login.make_url("/oauth/show-code")),
scope="offline_access",
audience="https://test.dev.neuromation.io",
)
return "test_code"
await Factory().login_headless(get_auth_code_cb,
url=mock_for_login.make_url("/"))
nmrc_path = tmp_home / ".nmrc"
assert Path(nmrc_path).exists(), "Config file not written after login "
saved_config = Factory(nmrc_path)._read()
assert saved_config.auth_config.is_initialized()
assert saved_config.cluster_config.is_initialized()
def registry_url(self) -> URL:
config = Factory(path=self._nmrc_path)._read()
return config.cluster_config.registry_url
def token(self) -> str:
config = Factory(path=self._nmrc_path)._read()
return config.auth_token.token
def username(self) -> str:
config = Factory(path=self._nmrc_path)._read()
return config.auth_token.username
def config(nmrc_path: Path) -> _Config:
return Factory(path=nmrc_path)._read()
async def init_client(self) -> None:
self._factory = Factory(path=self.config_path)
client = await self._factory.get(timeout=self.timeout)
self._client = client