def test_environment_provider():
"""Creating and resetting the environment."""
cluster1 = ClusterConfigImpl(host='host1',
port=22,
user='******',
auth=AuthMethod.ASK)
config = ClientConfig(clusters={'cluster1': cluster1})
environment = EnvironmentImpl(config=config)
EnvironmentProvider._state = None # pylint: disable=protected-access
environment_provider = EnvironmentProvider(initial_environment=environment)
assert EnvironmentProvider._state == environment_provider.__dict__ # noqa, pylint: disable=protected-access,line-too-long
assert environment_provider.environment is environment
cluster2 = ClusterConfigImpl(host='host2',
port=22,
user='******',
auth=AuthMethod.ASK)
config2 = ClientConfig(clusters={'cluster2': cluster2})
environment2 = EnvironmentImpl(config=config2)
environment_provider2 = EnvironmentProvider(
initial_environment=environment2)
assert environment_provider2.environment is environment
environment_provider2.environment = environment2
assert environment_provider2.environment is environment2
assert environment_provider.environment is environment2
EnvironmentProvider._state = None # pylint: disable=protected-access
def show_cluster(name: str) -> Cluster:
"""Returns the cluster with this name.
:param name: Cluster name.
"""
environment = EnvironmentProvider().environment
return environment.clusters[name]
def test_environment_missing_and_defaults():
user = USER_25
with ExitStack() as stack:
stack.enter_context(clear_environment(user))
EnvironmentProvider()._environment = None # noqa, pylint: disable=protected-access,line-too-long
assert show_clusters() == {}
assert not os.path.isfile(os.environ['IDACT_CONFIG_PATH'])
with pytest.raises(ValueError):
load_environment()
add_cluster(name=TEST_CLUSTER, user=user, host='localhost')
config = show_cluster(TEST_CLUSTER).config
set_log_level(logging.DEBUG)
check_config_is_default(config=config, user=user)
try:
save_environment()
with open(os.environ['IDACT_CONFIG_PATH'], 'r') as test_file:
contents = test_file.read().splitlines()
pprint(contents)
assert contents == get_default_config_contents(user=user)
finally:
os.remove(os.environ['IDACT_CONFIG_PATH'])
def remove_cluster(name: str):
"""Removes a cluster with this name
:param name: Name of the cluster to remove
"""
environment = EnvironmentProvider().environment
environment.remove_cluster(name=name)
def reset_environment(user: str, auth: AuthMethod = AuthMethod.ASK):
"""Clears the environment and adds the testing cluster.
:param user: User to connect to the cluster as,
and whose home dir should be cleaned.
:param auth: Authentication method to use.
"""
# pylint: disable=protected-access
saved_state = EnvironmentProvider._state
EnvironmentProvider._state = None
os.environ['IDACT_KEY_LOCATION'] = get_test_key_location(user=user)
os.environ['IDACT_CONFIG_PATH'] = get_test_environment_file(user=user)
cluster = ClusterConfigImpl(host=get_testing_host(),
port=get_testing_port(),
user=user,
auth=auth,
retries=get_default_retries_heavy_load())
cluster.retries[Retry.PORT_INFO] = set_retry(count=0)
EnvironmentProvider(initial_environment=EnvironmentImpl(
config=ClientConfig(clusters={TEST_CLUSTER: cluster})))
set_log_level(DEBUG)
try:
yield
finally:
EnvironmentProvider._state = saved_state
clear_home(user=user)
def push_environment(cluster: Cluster, path: Optional[str] = None):
"""Merges the environment on the cluster with the current environment.
:param cluster: Cluster to push the environment to.
:param path: Path to remote environment file.
Default: Remote IDACT_CONFIG_PATH environment variable,
or ~/.idact.conf
"""
log = get_logger(__name__)
with stage_info(log, "Pushing the environment to cluster."):
try:
remote_environment = deserialize_environment_from_cluster(
cluster=cluster,
path=path)
except RuntimeError:
log.info("Remote environment is missing, current environment will"
" be copied to cluster.")
log.debug("Exception", exc_info=1)
remote_environment = EnvironmentImpl()
local_environment = EnvironmentProvider().environment
merged_environment = merge_environments(local=remote_environment,
remote=local_environment)
serialize_environment_to_cluster(environment=merged_environment,
cluster=cluster,
path=path)
def load_environment(path: Optional[str] = None):
"""Loads the environment from file.
:param path: Path to environment file.
Default: IDACT_CONFIG_PATH environment variable,
or ~/.idact.conf
"""
environment = deserialize_environment_from_file(path=path)
EnvironmentProvider().environment = environment
def save_environment(path: Optional[str] = None):
"""Saves the environment to file.
:param path: Path to environment file.
Default: IDACT_CONFIG_PATH environment variable,
or ~/.idact.conf
"""
environment = EnvironmentProvider().environment
serialize_environment_to_file(environment=environment,
path=path)
def pull_environment(cluster: Cluster,
path: Optional[str] = None):
"""Merges the current environment with the environment on cluster.
:param cluster: Cluster to pull the environment from.
:param path: Path to remote environment file.
Default: Remote IDACT_CONFIG_PATH environment variable,
or ~/.idact.conf
"""
log = get_logger(__name__)
with stage_info(log, "Pulling the environment from cluster."):
remote_environment = deserialize_environment_from_cluster(
cluster=cluster,
path=path)
local_environment = EnvironmentProvider().environment
merged_environment = merge_environments(local=local_environment,
remote=remote_environment)
EnvironmentProvider().environment = merged_environment
def set_log_level(level: int):
"""Sets log level for idact loggers.
If the log level is lower or equal to DEBUG, Fabric logs are also
shown, otherwise they are hidden.
:param level: May be one of CRITICAL, FATAL, ERROR, WARNING,
WARN, INFO, DEBUG, NOTSET. Default: INFO
"""
environment = EnvironmentProvider().environment
environment.set_log_level(level=level)
LoggerProvider().log_level = level
def clear_environment(user: str):
"""Clears the environment, but does not add any clusters.
:param user: User, whose home dir should be cleaned.
"""
# pylint: disable=protected-access
saved_state = EnvironmentProvider._state
EnvironmentProvider._state = None
EnvironmentProvider(initial_environment=EnvironmentImpl())
set_log_level(logging.DEBUG)
os.environ['IDACT_CONFIG_PATH'] = get_test_environment_file(user=user)
try:
yield
finally:
EnvironmentProvider._state = saved_state
clear_home(user=user)
def show_clusters() -> Dict[str, Cluster]:
"""Returns a dictionary of all defined clusters by name."""
environment = EnvironmentProvider().environment
return {key: value for key, value in environment.clusters.items()}
def add_cluster(name: str,
user: str,
host: str,
port: int = 22,
auth: Optional[AuthMethod] = None,
key: Union[None, str, KeyType] = None,
install_key: bool = True,
disable_sshd: bool = False,
setup_actions: Optional[SetupActionsConfig] = None,
scratch: Optional[str] = None,
retries: Optional[Dict[Retry, RetryConfig]] = None,
use_jupyter_lab: bool = True) -> Cluster:
"""Adds a new cluster.
:param name:
Cluster name to identify it by.
:param user:
Remote cluster user to log in and run commands as.
:param host:
Cluster access node hostname.
:param port:
SSH port to access the cluster.
Default: 22.
:param auth:
Authentication method.
Default: :attr:`.AuthMethod.ASK` (password-based).
:param key:
Private key path (if applicable), or key type to generate.
Default: None
:param install_key:
True, if the public key should be installed on the cluster before
use (if applicable).
Default: True
:param disable_sshd:
Should be set to True, if the cluster allows ssh connection
to compute nodes out of the box.
Default: False
:param setup_actions:
Commands to run before deployment.
Default: None
:param setup_actions:
Commands to run before deployment.
Default: None
:param scratch:
Absolute path to a high-performance filesystem for temporary
computation data, or an environment variable that contains it.
Default: $HOME
:param retries:
Retry config by action name.
Defaults: see :func:`.get_default_retries`.
:param use_jupyter_lab:
Use Jupyter Lab instead of Jupyter Notebook.
Default: True.
"""
log = get_logger(__name__)
environment = EnvironmentProvider().environment
if auth is None:
log.info("No auth method specified, defaulting to password-based.")
auth = AuthMethod.ASK
if auth is AuthMethod.PUBLIC_KEY:
if isinstance(key, KeyType):
log.info("Generating public-private key pair.")
key = generate_key(host=host, key_type=key)
elif isinstance(key, str):
key = os.path.expanduser(key)
else:
raise ValueError("Invalid key argument for public key"
" authentication.")
config = ClusterConfigImpl(host=host,
port=port,
user=user,
auth=auth,
key=key,
install_key=install_key,
disable_sshd=disable_sshd,
setup_actions=setup_actions,
scratch=scratch,
retries=retries,
use_jupyter_lab=use_jupyter_lab)
return environment.add_cluster(name=name, config=config)