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