Stress Tests

A stress test evaluates a machine learning model’s resilience with regards to a specific vulnerability.

Creating a new Stress Test from the SDK

This procedure presumes a model is already loaded into RIME and that all commands are issued in a Python environment.

1. (When a suitable project does not already exist) Create a project using the following SDK command.

   project = rime_client.create_project(
       name="project", description="$DESC", model_task='MODEL_TASK_BINARY_CLASSIFICATION'
   )
   

   You can also access an existing project using its project ID:

   project = rime_client.get_project(MY_PROJECT_ID)
   

   Replace $DESC with a description of the project or modeling TASK_NATURAL_LANGUAGE_INFERENCE

   • MODEL_TASK_FILL_MASK

2. Register a reference dataset.

   This example registers a dataset from a file.

   reference_id = project.register_dataset(
       name=DATASET_NAME,
       data_params={
           "connection_info": {"data_file": {"path": FILE_PATH}},
           "data_params": {"label_col": LABEL_COL},
       },
       integration_id=INTEGRATION_ID,
   )
   

   The SDK returns the ID of the reference dataset.

3. Register an evaluation dataset.

   evaluation_id = project.register_dataset(
       name=DATASET_NAME,
       data_params={
           "connection_info": {"data_file": {"path": FILE_PATH}},
           "data_params": {"label_col": LABEL_COL},
       },
       integration_id=INTEGRATION_ID,
   )
   

   The SDK returns the ID of the evaluation dataset.

4. Register a model.

   This example registers a model from a model.py file. The Specify a Model section has details on how to format model.py files.

   model_id = project.register_model_from_path(
           name=MODEL_NAME,
           remote_path=S3_PATH_TO_MODEL_PY
   )
   

   This example registers a Hugging Face model.

          model_id = project.register_model(
              name=MODEL_NAME,
              model_config={
                  "hugging_face": {
                      "model_uri": URI,
                      "kwargs": {
                          "tokenizer_uri": TOKENIZER_URI,
                          "class_map": MAP,
                          "ignore_class_names": True,
                      },
                  }
              },
          )
   

   The SDK returns the model ID.

5. Create a test run configuration that uses the registered reference and evaluation datasets and the model, specifying a name and model task in addition to the IDs of the datasets and model registered earlier in this procedure.

   config = {
       "data_info": {"ref_dataset_id": reference_id, "eval_dataset_id": evaluation_id},
       "model_id": model_id,
       "run_name": "My Stress Test Run",
   }
   

   You can also specify predictions in the test run configuration.

6. Issue the following command to start the Stress Test, specifying the configuration created in the previous step and the unique ID of the project that contains the Stress Test.

   job = client.start_stress_test(
       test_run_config=config, project_id=project.project_id
   )