def test_invalid_serialized_type():
config = get_config_for_test()
serialized = {'type': 'SerializableTypes.NODE_IMPL2'}
with pytest.raises(AssertionError):
NodeImpl.deserialize(config=config, serialized=serialized)
def test_missing_serialized_keys():
config = get_config_for_test()
serialized = {'type': 'SerializableTypes.NODE_IMPL'}
with pytest.raises(RuntimeError):
NodeImpl.deserialize(config=config, serialized=serialized)
def get_data_for_test():
config = ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=AuthMethod.ASK)
access_node = NodeImpl(config=config)
nodes = [NodeImpl(config=config), NodeImpl(config=config)]
return access_node, nodes
def test_serialize_deserialize():
config, access_node = get_data_for_test()
nodes = [NodeImpl(config=config), NodeImpl(config=config)]
uuid = '1111'
value = NodesImpl(nodes=nodes,
allocation=SlurmAllocation(
job_id=1,
access_node=access_node,
nodes=nodes,
entry_point_script_path='a',
parameters=AllocationParameters()),
uuid=uuid)
serialized = value.serialize()
assert serialized == {
'type':
'SerializableTypes.NODES_IMPL',
'nodes': [{
'type': 'SerializableTypes.NODE_IMPL',
'host': None,
'port': None,
'cores': None,
'memory': None,
'allocated_until': None
}, {
'type': 'SerializableTypes.NODE_IMPL',
'host': None,
'port': None,
'cores': None,
'memory': None,
'allocated_until': None
}],
'allocation': {
'type': 'SerializableTypes.SLURM_ALLOCATION',
'job_id': 1,
'entry_point_script_path': 'a',
'parameters': {
'type': 'SerializableTypes.ALLOCATION_PARAMETERS',
'nodes': None,
'cores': None,
'memory_per_node': None,
'walltime': None,
'native_args': {}
},
'done_waiting': False
}
}
deserialized = NodesImpl.deserialize(config=config,
access_node=access_node,
serialized=serialized,
uuid=uuid)
assert deserialized == value
def get_access_node(config: ClusterConfig) -> NodeInternal:
"""Returns the cluster access node, identified
by :attr:`.ClusterConfig.host`.
An access node is expected to be available without allocation.
:param config: Client cluster config.
"""
node = NodeImpl(config=config)
node.make_allocated(host=config.host,
port=config.port,
cores=None,
memory=None,
allocated_until=None)
return node
def test_serialize_deserialize():
config, uuid = get_data_for_test()
value = JupyterDeploymentImpl(deployment=GenericDeployment(
node=NodeImpl(config=config), pid=111, runtime_dir='/dir'),
tunnel=FakeTunnel(here=2222, there=1111),
token='abcdefg',
uuid=uuid)
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.JUPYTER_DEPLOYMENT_IMPL',
'deployment': {
'type': 'SerializableTypes.GENERIC_DEPLOYMENT',
'node': {
'type': 'SerializableTypes.NODE_IMPL',
'host': None,
'port': None,
'cores': None,
'memory': None,
'allocated_until': None
},
'pid': 111,
'runtime_dir': '/dir'
},
'tunnel_there': 1111,
'tunnel_here': 2222,
'token': 'abcdefg'
}
with use_fake_tunnel():
deserialized = deserialize_jupyter_deployment_impl(
config=config, uuid=uuid, serialized=serialized)
assert deserialized == value
def test_serialize_deserialize_not_allocated():
config = get_config_for_test()
value = NodeImpl(config=config)
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.NODE_IMPL',
'host': None,
'port': None,
'cores': None,
'memory': None,
'allocated_until': None
}
deserialized = NodeImpl.deserialize(config=config, serialized=serialized)
assert deserialized == value
def get_deployment_for_test(uuid: str, job_id: int) -> NodesImpl:
config = ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=AuthMethod.ASK)
access_node = NodeImpl(config=config)
nodes = [NodeImpl(config=config),
NodeImpl(config=config)]
deployment = NodesImpl(nodes=nodes,
allocation=SlurmAllocation(
job_id=job_id,
access_node=access_node,
nodes=nodes,
entry_point_script_path='a',
parameters=AllocationParameters()),
uuid=uuid)
return deployment
def deserialize_generic_deployment(config: ClusterConfig,
serialized: dict) -> GenericDeployment:
"""Counterpart to :meth:`.GenericDeployment.serialize`."""
try:
assert serialized['type'] == str(
SerializableTypes.GENERIC_DEPLOYMENT)
return GenericDeployment(
node=NodeImpl.deserialize(config=config,
serialized=serialized['node']),
pid=serialized['pid'],
runtime_dir=serialized['runtime_dir'])
except KeyError as e:
raise RuntimeError("Unable to deserialize.") from e
def get_data_for_test():
config = ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=AuthMethod.ASK)
uuid = 'uuid1'
deployment = GenericDeployment(node=NodeImpl(config=config),
pid=111,
runtime_dir='/dir1')
serialized_deployment = deployment.serialize()
return config, uuid, deployment, serialized_deployment
def deserialize(config: ClusterConfig, access_node: NodeInternal,
uuid: str, serialized: dict) -> 'NodesImpl':
try:
assert serialized['type'] == str(SerializableTypes.NODES_IMPL)
nodes = [
NodeImpl.deserialize(config=config, serialized=node)
for node in serialized['nodes']
]
allocation = SlurmAllocation.deserialize(
access_node=access_node,
nodes=nodes,
serialized=serialized['allocation'])
return NodesImpl(nodes=nodes, allocation=allocation, uuid=uuid)
except KeyError as e:
raise RuntimeError("Unable to deserialize.") from e
def allocate_slurm_nodes(parameters: AllocationParameters,
config: ClusterConfig) -> Nodes:
"""Tries to allocate nodes using Slurm.
:param parameters: Allocation parameters.
:param config: Config for the cluster to allocate nodes on.
"""
args = SbatchArguments(params=parameters)
log = get_logger(__name__)
with stage_debug(log, "Executing sbatch on access node."):
access_node = get_access_node(config=config)
job_id, entry_point_script_path = run_sbatch(args=args,
node=access_node)
def run_squeue_task() -> SqueueResult:
job_squeue = run_squeue(node=access_node)
return job_squeue[job_id]
try:
with stage_debug(log, "Obtaining info about job %d using squeue.",
job_id):
job = retry_with_config(run_squeue_task,
name=Retry.SQUEUE_AFTER_SBATCH,
config=config)
except Exception as e: # noqa, pylint: disable=broad-except
run_scancel(job_id=job_id, node=access_node)
raise RuntimeError("Unable to obtain job info"
" after allocation.") from e
node_count = job.node_count
nodes = [NodeImpl(config=config) for _ in range(node_count)]
allocation = SlurmAllocation(
job_id=job_id,
access_node=access_node,
nodes=nodes,
entry_point_script_path=entry_point_script_path,
parameters=parameters)
return NodesImpl(nodes=nodes,
allocation=allocation)
def test_serialize_deserialize():
config = get_config_for_test()
value = GenericDeployment(node=NodeImpl(config=config),
pid=1,
runtime_dir='/dir')
serialized = value.serialize()
assert serialized == {'type': 'SerializableTypes.GENERIC_DEPLOYMENT',
'node': {'type': 'SerializableTypes.NODE_IMPL',
'host': None,
'port': None,
'cores': None,
'memory': None,
'allocated_until': None},
'pid': 1,
'runtime_dir': '/dir'}
deserialized = deserialize_generic_deployment(config=config,
serialized=serialized)
assert deserialized == value
def test_serialize_deserialize():
config = get_config_for_test()
value = NodeImpl(config=config)
value.make_allocated(host='node1',
port=12323,
cores=30,
memory=bitmath.GiB(60),
allocated_until=(datetime.datetime(
2018, 11, 12, 13,
14).replace(tzinfo=dateutil.tz.tzutc())))
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.NODE_IMPL',
'host': 'node1',
'port': 12323,
'cores': 30,
'memory': '60.0 GiB',
'allocated_until': '2018-11-12T13:14:00+00:00'
}
deserialized = NodeImpl.deserialize(config=config, serialized=serialized)
assert deserialized == value