ModelManager

Helper class to manage trained and registered models.

Parameters

Name	Type	Description	Optional
`connector`	`SnowflakeConnector`	The connector object used for sending requests to the GNN engine.	No

Returns

An instance of the ModelManager class.

Example

from relationalai_gnns import ModelManager
model_manager = ModelManager(connector =connector)

Methods

`list_models()`

Lists all models that have been trained and registered.

Returns:
A dictionary containing:

Example:

from relationalai_gnns import ModelManager

model_manager = ModelManager()
model_manager.list_models()

`delete_model()`

Permanently deletes a trained model from the RelationalAI application stage.

Parameters:

Name	Type	Description	Optional
`model_run_id`	str	Unique identifier of the trained model run.	No
`experiment_config`	`ExperimentConfig`	An instance of `ExperimentConfig` which defines the database and schema of the experiment under which that model was trained. Not needed when `delete_experiment_run` is `False`.	Yes
`delete_experiment_run`	bool	If `True`, the Snowflake experiment tracking run is also deleted. Set to `False` to keep the run so that historical metrics, hyperparameters, etc. remain accessible. Defaults to `True`.	Yes

Example:

from relationalai_gnns import ExperimentConfig, ModelManager


model_manager = ModelManager()
experiment_config = ExperimentConfig(database="experiment_database",
                                     schema="experiment_schema")
model_manager.delete_model(
    model_run_id="abc123",
    experiment_config = experiment_config
)

`register_model()`

Registers a trained model in the Snowflake Model Registry.
The model is stored under the specified name and version in the database and schema where the tracker has permissions.
Registration is only allowed if a valid model_run_id exists.
Duplicate version names for a model are not allowed.

⚠️ Note: Ensure the native app has the necessary permissions to the specified database and schema.

-- Grant access to resources needed to register a model in Snowflake.

GRANT USAGE ON DATABASE <DATABASE> TO APPLICATION RELATIONALAI;
GRANT USAGE ON SCHEMA <DATABASE>.<SCHEMA> TO APPLICATION RELATIONALAI;
GRANT CREATE MODEL ON SCHEMA <DATABASE>.<SCHEMA> TO APPLICATION RELATIONALAI;

Parameters:

Name	Type	Description	Optional
`model_run_id`	str	Unique identifier of the trained model run.	No
`model_name`	str	Name to assign to the registered model.	No
`version_name`	str	Version name for the model.	No
`database_name`	str	Database where the model will be registered.	No
`schema_name`	str	Schema in the database where the model will be registered.	No
`comment`	str	Optional comment for model registration.	Yes

Example:

from relationalai_gnns import ModelManager

model_manager = ModelManager()
model_manager.register_model(
    model_run_id=train_job.model_run_id,
    model_name="new_model",
    version_name="v3",
    database_name="database_name",
    schema_name="schema_name",
    comment="Production-ready model",
)

`delete_registered_model()`

Permanently deletes a registered model from the RelationalAI model registry. The corresponding trained model is not deleted from the RelationalAI application stage. delete_model() should be used for that instead.

Parameters:

Name	Type	Description	Optional
`model_name`	str	Name to assign to the registered model.	No
`version_name`	str	Version name for the model.	No
`experiment_config`	`ExperimentConfig`	An instance of `ExperimentConfig` which defines the database and schema where the model was registered.	No

Example:

from relationalai_gnns import ExperimentConfig, ModelManager

model_manager = ModelManager()
experiment_config = ExperimentConfig(database="model_registry_database",
                                     schema="model_registry_schema")
model_manager.delete_registered_model(
    model_name="new_model",
    version_name="v3",
    experiment_config = experiment_config
)