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.