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