def override_parameters_if_possible(parameters: AppAllocationParameters,
nodes: Optional[int],
cores: Optional[int],
memory_per_node: Optional[str],
walltime: Optional[str],
native_args: List[Tuple[str, str]]):
"""Overrides parameters with values from the command line, if they were
provided.
:param parameters: Parameters to override.
:param nodes: Nodes from command line.
:param cores: Cores from command line.
:param memory_per_node: Memory from command line.
:param walltime: Walltime from command line.
:param native_args: Native args from command line.
"""
if nodes is not None:
parameters.nodes = nodes
if cores is not None:
parameters.cores = cores
if memory_per_node is not None:
parameters.memory_per_node = memory_per_node
if walltime is not None:
parameters.walltime = walltime
if native_args:
parameters.native_args = [[key for (key, _) in native_args],
[value for (_, value) in native_args]]
def test_serialize_deserialize_empty():
value = AppAllocationParameters()
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.APP_ALLOCATION_PARAMETERS',
'nodes': 1,
'cores': 1,
'memory_per_node': '1GiB',
'walltime': '0:10:00',
'native_args': [[], []]
}
deserialized = AppAllocationParameters.deserialize(serialized=serialized)
assert deserialized == value
def get_parameters_for_test() -> AppAllocationParameters:
return AppAllocationParameters(
nodes=10,
cores=20,
memory_per_node='30GiB',
walltime='0:40:00',
native_args=[['--arg1'],
['value1']])
def test_serialize_deserialize():
value = AppAllocationParameters(
nodes=20,
cores=10,
memory_per_node='20 GiB',
walltime='1-2:03:04',
native_args=[['--arg1', '--arg2', '--arg3'], ['value 1', 'None',
'65']])
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.APP_ALLOCATION_PARAMETERS',
'nodes': 20,
'cores': 10,
'memory_per_node': '20 GiB',
'walltime': '1-2:03:04',
'native_args': [['--arg1', '--arg2', '--arg3'],
['value 1', 'None', '65']]
}
deserialized = AppAllocationParameters.deserialize(serialized=serialized)
assert deserialized == value
def test_format_allocation_parameters_no_native_args():
formatted = format_allocation_parameters(
AppAllocationParameters(nodes=20,
cores=10,
memory_per_node='20 GiB',
walltime='1-2:03:04'))
assert formatted == ("Allocation parameters:\n"
" Nodes: 20\n"
" Cores: 10\n"
" Memory per node: 20 GiB\n"
" Walltime: 1-2:03:04\n"
" No native arguments.\n")
def test_override_all():
parameters = get_parameters_for_test()
override_parameters_if_possible(parameters=parameters,
nodes=15,
cores=25,
memory_per_node='35GiB',
walltime='0:45:00',
native_args=[('--arg2', 'value2')])
assert parameters == AppAllocationParameters(
nodes=15,
cores=25,
memory_per_node='35GiB',
walltime='0:45:00',
native_args=[['--arg2'],
['value2']])
def test_override_nodes():
parameters = get_parameters_for_test()
override_parameters_if_possible(parameters=parameters,
nodes=15,
cores=None,
memory_per_node=None,
walltime=None,
native_args=[])
assert parameters == AppAllocationParameters(
nodes=15,
cores=20,
memory_per_node='30GiB',
walltime='0:40:00',
native_args=[['--arg1'],
['value1']])
def test_format_allocation_parameters():
formatted = format_allocation_parameters(
AppAllocationParameters(nodes=20,
cores=10,
memory_per_node='20 GiB',
walltime='1-2:03:04',
native_args=[['--arg1', '--arg2', '--arg3'],
['value 1', 'None', '65']]))
assert formatted == ("Allocation parameters:\n"
" Nodes: 20\n"
" Cores: 10\n"
" Memory per node: 20 GiB\n"
" Walltime: 1-2:03:04\n"
" Native arguments:\n"
" --arg1 -> value 1\n"
" --arg2\n"
" --arg3 -> 65\n")
def test_deserialize_empty_dict():
serialized = {}
deserialized = AppAllocationParameters.deserialize(serialized=serialized)
assert deserialized == AppAllocationParameters()
def test_missing_serialized_keys():
serialized = {'type': 'SerializableTypes.APP_ALLOCATION_PARAMETERS'}
deserialized = AppAllocationParameters.deserialize(serialized=serialized)
assert deserialized == AppAllocationParameters()
def test_invalid_serialized_type():
serialized = {'type': 'SerializableTypes.APP_ALLOCATION_PARAMETERS2'}
with pytest.raises(AssertionError):
AppAllocationParameters.deserialize(serialized=serialized)
def main(cluster_name: str, environment: Optional[str], save_defaults: bool,
reset_defaults: bool, nodes: Optional[int], cores: Optional[int],
memory_per_node: Optional[str], walltime: Optional[str],
native_arg: List[Tuple[str, str]]) -> int:
"""A console script that executes a Jupyter Notebook instance on
an allocated cluster node, and makes it accessible
in the local browser.
CLUSTER_NAME argument is the cluster name to execute the notebook on.
It must already be present in the config file.
"""
ensure_stdin_has_fileno()
log = None
try:
with ExitStack() as stack:
click.echo("Loading environment.")
load_environment(path=environment)
log = get_logger(__name__)
cluster = show_cluster(name=cluster_name)
config = cluster.config
assert isinstance(config, ClusterConfigImpl)
if reset_defaults:
click.echo("Resetting allocation parameters to defaults.")
config.notebook_defaults = {}
parameters = AppAllocationParameters.deserialize(
serialized=config.notebook_defaults)
override_parameters_if_possible(parameters=parameters,
nodes=nodes,
cores=cores,
memory_per_node=memory_per_node,
walltime=walltime,
native_args=native_arg)
if save_defaults:
click.echo("Saving defaults.")
config.notebook_defaults = parameters.serialize()
save_environment(path=environment)
click.echo(format_allocation_parameters(parameters=parameters))
click.echo("Allocating nodes.")
nodes = cluster.allocate_nodes(
nodes=parameters.nodes,
cores=parameters.cores,
memory_per_node=parameters.memory_per_node,
walltime=parameters.walltime,
native_args=convert_native_args_from_command_line_to_dict(
native_args=parameters.native_args))
stack.enter_context(cancel_on_exit(nodes))
nodes.wait()
notebook = nodes[0].deploy_notebook()
stack.enter_context(cancel_local_on_exit(notebook))
click.echo("Pushing the allocation deployment.")
cluster.push_deployment(nodes)
click.echo("Pushing the notebook deployment.")
cluster.push_deployment(notebook)
click.echo(format_deployments_info(cluster_name=cluster_name))
click.echo("Notebook address: ", nl=False)
click.echo(click.style(notebook.address, fg='red'))
notebook.open_in_browser()
sleep_until_allocation_ends(nodes=nodes)
except: # noqa, pylint: disable=broad-except
if log is not None:
log.error("Exception raised.", exc_info=1)
return 1
raise
return 0