def test_serialize_deserialize():
value = AllocationParameters(nodes=20,
cores=10,
memory_per_node=bitmath.GiB(20),
walltime=Walltime(days=1,
hours=12,
minutes=5,
seconds=32),
native_args={
'--arg1': 'value 1',
'--arg2': None,
'--arg3': '65'
})
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.ALLOCATION_PARAMETERS',
'nodes': 20,
'cores': 10,
'memory_per_node': '20.0 GiB',
'walltime': '1-12:05:32',
'native_args': {
'--arg1': 'value 1',
'--arg2': None,
'--arg3': '65'
}
}
deserialized = AllocationParameters.deserialize(serialized=serialized)
assert deserialized == value
def test_sbatch_arguments_unsupported_provided():
params = AllocationParameters(nodes=1,
cores=2,
walltime=Walltime(minutes=10),
memory_per_node=GiB(1))
params.all['Provided Unsupported Param'] = 12
with pytest.raises(ValueError):
SbatchArguments(params=params)
def test_serialize_deserialize_empty():
value = AllocationParameters()
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.ALLOCATION_PARAMETERS',
'nodes': None,
'cores': None,
'memory_per_node': None,
'walltime': None,
'native_args': {}
}
deserialized = AllocationParameters.deserialize(serialized=serialized)
assert deserialized == value
def allocate_nodes(self,
nodes: int = 1,
cores: int = 1,
memory_per_node: Union[str, bitmath.Byte] = None,
walltime: Union[str, Walltime] = None,
native_args: Optional[
Dict[str, Optional[
str]]] = None) -> Nodes: # noqa, pylint: disable=bad-whitespace,line-too-long
if memory_per_node is None:
memory_per_node = bitmath.GiB(1)
elif isinstance(memory_per_node, str):
memory_per_node = bitmath.parse_string(memory_per_node)
if walltime is None:
walltime = Walltime(minutes=10)
elif isinstance(walltime, str):
walltime = Walltime.from_string(walltime)
parameters = AllocationParameters(nodes=nodes,
cores=cores,
memory_per_node=memory_per_node,
walltime=walltime,
native_args=native_args)
return allocate_slurm_nodes(parameters=parameters,
config=self._config)
def test_serialize_deserialize():
access_node, nodes = get_data_for_test()
value = SlurmAllocation(job_id=111,
access_node=access_node,
nodes=nodes,
entry_point_script_path='/a/b/c',
parameters=AllocationParameters())
serialized = value.serialize()
assert serialized == {
'type': 'SerializableTypes.SLURM_ALLOCATION',
'job_id': 111,
'entry_point_script_path': '/a/b/c',
'parameters': {
'type': 'SerializableTypes.ALLOCATION_PARAMETERS',
'nodes': None,
'cores': None,
'memory_per_node': None,
'walltime': None,
'native_args': {}
},
'done_waiting': False
}
deserialized = SlurmAllocation.deserialize(access_node=access_node,
nodes=nodes,
serialized=serialized)
assert deserialized == value
def test_allocation_parameters_create_empty():
"""Tests default construction of allocation parameters."""
params = AllocationParameters()
assert params.all == {
'nodes': None,
'cores': None,
'memory_per_node': None,
'walltime': None
}
assert params.native_args == {}
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 test_sbatch_arguments_create():
params = AllocationParameters(nodes=1,
cores=2,
memory_per_node=GiB(1),
walltime=Walltime(minutes=10))
sbatch_args = SbatchArguments(params=params)
assert sbatch_args.native_args == {}
assert sbatch_args.args == {'--nodes': '1',
'--cpus-per-task': '2',
'--mem': '1048576K',
'--time': '0-00:10:00'}
def deserialize(access_node: NodeInternal,
nodes: List[NodeImpl],
serialized: dict) -> 'SlurmAllocation':
try:
assert serialized['type'] == str(
SerializableTypes.SLURM_ALLOCATION)
return SlurmAllocation(
job_id=serialized['job_id'],
access_node=access_node,
nodes=nodes,
entry_point_script_path=serialized['entry_point_script_path'],
parameters=AllocationParameters.deserialize(
serialized['parameters']),
done_waiting=serialized['done_waiting'])
except KeyError as e:
raise RuntimeError("Unable to serialize.") from e
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 test_allocation_parameters_create():
"""Tests construction of allocation parameters."""
params = AllocationParameters(nodes=1,
cores=2,
memory_per_node=GiB(1),
walltime=Walltime(minutes=10),
native_args={
'--abc': None,
'--def': '80'
})
assert params.all == {
'nodes': 1,
'cores': 2,
'memory_per_node': GiB(1),
'walltime': Walltime(minutes=10)
}
assert params.nodes == 1
assert params.cores == 2
assert params.memory_per_node == GiB(1)
assert params.walltime == Walltime(minutes=10)
assert params.native_args == {'--abc': None, '--def': '80'}
def test_format_sbatch_allocation_request():
params = AllocationParameters(nodes=1,
cores=2,
memory_per_node=GiB(1),
walltime=Walltime(minutes=10),
native_args={'--arg1': 'def; rm -rf /abc &&',
'--arg2': None,
'--arg3': 'a b c',
'--arg4 ||': '3',
'arg5': '3#',
'--mem': '8G'})
args = SbatchArguments(params=params)
formatted = format_sbatch_allocation_request(
args=args,
entry_point_script='/home/user/script')
expected = ("sbatch"
" --arg1 'def; rm -rf /abc &&'"
" --arg2"
" --arg3 'a b c'"
" '--arg4 ||' 3"
" --mem 8G"
" arg5 '3#'"
" --cpus-per-task 2"
" --mem 1048576K"
" --nodes 1"
" --time 0-00:10:00"
" --tasks-per-node=1"
" --parsable"
" --output=/dev/null"
" --wrap='export IDACT_ALLOCATION_ID=$SLURM_JOB_ID"
" ; srun /home/user/script'")
print()
print(formatted)
print(expected)
assert formatted == expected
def test_missing_serialized_keys():
serialized = {'type': 'SerializableTypes.ALLOCATION_PARAMETERS'}
with pytest.raises(RuntimeError):
AllocationParameters.deserialize(serialized=serialized)
def test_invalid_serialized_type():
serialized = {'type': 'SerializableTypes.ALLOCATION_PARAMETERS2'}
with pytest.raises(AssertionError):
AllocationParameters.deserialize(serialized=serialized)
def test_sbatch_arguments_missing_required():
with pytest.raises(ValueError):
SbatchArguments(params=AllocationParameters(nodes=1,
cores=2,
memory_per_node=GiB(1)))