dmlc--dgl
Rhett Ying
cf19254a19
* [Dist] enable to partition many chunks into less partitions via pipeline * refine * add meta file for num_parts, add more tests, refine docstring * remove args.num_parts * create pydantic class for partition metadata * refine * rename json file
74 行
2.2 KiB
Python
74 行
2.2 KiB
Python
from typing import Optional
|
|
import pydantic as dt
|
|
import json
|
|
from dgl import DGLError
|
|
|
|
class PartitionMeta(dt.BaseModel):
|
|
""" Metadata that describes the partition assignment results.
|
|
|
|
Regardless of the choice of partitioning algorithm, a metadata JSON file
|
|
will be created in the output directory which includes the meta information
|
|
of the partition algorithm.
|
|
|
|
To generate a metadata JSON:
|
|
|
|
>>> part_meta = PartitionMeta(version='1.0.0', num_parts=4, algo_name='random')
|
|
>>> with open('metadata.json', 'w') as f:
|
|
... json.dump(part_meta.dict(), f)
|
|
|
|
To read a metadata JSON:
|
|
|
|
>>> with open('metadata.json') as f:
|
|
... part_meta = PartitionMeta(**(json.load(f)))
|
|
|
|
"""
|
|
# version of metadata JSON.
|
|
version: Optional[str] = '1.0.0'
|
|
# number of partitions.
|
|
num_parts: int
|
|
# name of partition algorithm.
|
|
algo_name: str
|
|
|
|
def dump_partition_meta(part_meta, meta_file):
|
|
""" Dump partition metadata into json file.
|
|
|
|
Parameters
|
|
----------
|
|
part_meta : PartitionMeta
|
|
The partition metadata.
|
|
meta_file : str
|
|
The target file to save data.
|
|
"""
|
|
with open(meta_file, 'w') as f:
|
|
json.dump(part_meta.dict(), f, sort_keys=True, indent=4)
|
|
|
|
def load_partition_meta(meta_file):
|
|
""" Load partition metadata and do sanity check.
|
|
|
|
Parameters
|
|
----------
|
|
meta_file : str
|
|
The path of the partition metadata file.
|
|
|
|
Returns
|
|
-------
|
|
PartitionMeta
|
|
The partition metadata.
|
|
"""
|
|
with open(meta_file) as f:
|
|
try:
|
|
part_meta = PartitionMeta(**(json.load(f)))
|
|
except dt.ValidationError as e:
|
|
raise DGLError(
|
|
f"Invalid partition metadata JSON. Error details: {e.json()}")
|
|
if part_meta.version != '1.0.0':
|
|
raise DGLError(
|
|
f"Invalid version[{part_meta.version}]. Supported versions: '1.0.0'")
|
|
if part_meta.num_parts <= 0:
|
|
raise DGLError(
|
|
f"num_parts[{part_meta.num_parts}] should be greater than 0.")
|
|
if part_meta.algo_name not in ['random', 'metis']:
|
|
raise DGLError(
|
|
f"algo_name[{part_meta.num_parts}] is not supported.")
|
|
return part_meta
|