Skip to main content

​
BohriumToolkit

class BohriumToolkit(BaseToolkit):
A class representing a toolkit for interacting with Bohrium services. Parameters:
  • timeout (Optional[float], optional): The timeout for BohriumToolkit. (default: :obj:None)
  • api_key (Optional[str], optional): The API key for Bohrium client. (default: :obj:None)
  • project_id (Optional[int], optional): The project ID for Bohrium client. (default: :obj:None)
  • yaml_path (Optional[str], optional): The path to the YAML file containing the job parameters. (default: :obj:None)

​
init

def __init__(
    self,
    timeout: Optional[float] = None,
    api_key: Optional[str] = None,
    project_id: Optional[int] = None,
    yaml_path: Optional[str] = None,
    _test_mode: bool = False
):

​
_custom_insert

def _custom_insert(self, data):
refactor insert method, ensure return jobId information

​
submit_job

def submit_job(
    self,
    job_name: str = 'bohr-job',
    machine_type: str = 'c2_m4_cpu',
    cmd: str = 'mpirun -n 2 lmp_mpi -i in.shear',
    image_address: str = 'registry.dp.tech/dptech/lammps:29Sep2021'
):
Submit a job to Bohrium. Parameters:
  • job_name (str): The name of the job. It will be updated when yaml file is provided. The yaml file might be set when initialize BohriumToolkit. (default: :obj:bohr-job)
  • machine_type (str): The type of machine to use. It will be updated when yaml file is provided. The yaml file might be set when initialize BohriumToolkit. (default: :obj:c2_m4_cpu)
  • cmd (str): The command to run. It will be updated when yaml file is provided. The yaml file might be set when initialize (default: :obj:mpirun -n 2 lmp_mpi -i in.shear)
  • image_address (str): The address of the image to use. It will be updated when yaml file is provided. The yaml file might be set when initialize BohriumToolkit. (default: :obj:registry.dp.tech/dptech/lammps:29Sep2021)
Returns: Dict[str, Any]: The result of the job submission.

​
get_job_details

def get_job_details(self, job_id: int):
Get details for a specific job. Parameters:
  • job_id (int): The ID of the job.
Returns: Dict[str, Any]: The job details.

​
terminate_job

def terminate_job(self, job_id: int):
Terminate a running job. Parameters:
  • job_id (int): The ID of the job to terminate.
Returns: Dict[str, Any]: The result of the termination request.

​
kill_job

def kill_job(self, job_id: int):
Kill a running job. Parameters:
  • job_id (int): The ID of the job to kill.
Returns: Dict[str, Any]: The result of the kill request.

​
get_job_logs

def get_job_logs(
    self,
    job_id: int,
    log_file: str = 'STDOUTERR',
    page: int = -1,
    page_size: int = 8192
):
Get logs for a specific job. Parameters:
  • job_id (int): The ID of the job.
  • log_file (str, optional): The log file to get. (default: :obj:STDOUTERR)
  • page (int, optional): The page number. (default: :obj:-1)
  • page_size (int, optional): The page size. (default: :obj:8192)
Returns: str: The log contents.

​
create_job_group

def create_job_group(self, project_id: int, job_group_name: str):
Create a job group. Parameters:
  • project_id (int): The ID of the project.
  • job_group_name (str): The name of the job group.
Returns: Dict[str, Any]: The result of the job group creation.

​
download_job_results

def download_job_results(self, job_id: int, save_path: str):
Download the results of a job. Parameters:
  • job_id (int): The ID of the job.
  • save_path (str): The path to save the results to.
Returns: Dict[str, Any]: The result of the download request.

​
get_tools

def get_tools(self):
Returns: List[FunctionTool]: A list of FunctionTool objects representing the functions in the toolkit.