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 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_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