"""

Enum for different options.

"""

NUM_SAMPLES = 5

BASE = "base"

OUTPUT_CSV = "result.csv"

INPUT_CSV = "input.csv"

OUTPUT1_CSV = "output1.csv"

OUTPUT2_CSV = "output2.csv"

OUTPUT_FILES = [

"file.csv",

"file.h5",

"file.json",

]

"""

Enum for different check types.

"""

MODEL_CONSISTENCY = "Check Consistency of Model Output"

IMAGE_SIZE = "Image Size Mb"

PREDEFINED_EXAMPLE = "Check Predefined Example Input"

ENV_SIZE = "Environment Size Mb"

DIR_SIZE = "Directory Size Mb"

# messages

SIZE_CACL_SUCCESS = "Size Successfully Calculated"

SIZE_CACL_FAILED = "Size Calculation Failed"

INCONSISTENCY = "Inconsistent Output Detected"

CONSISTENCY = "Model Output Was Consistent"

"""

Enum for different table types.

"""

MODEL_INFORMATION_CHECKS = "Model Metadata Checks"

MODEL_FILE_CHECKS = "Model File Checks"

MODEL_DIRECTORY_SIZES = "Model Directory Sizes"

MODEL_ENV_SIZES = "Model Environment Sizes"

RUNNER_CHECKUP_STATUS = "Runner Checkup Status"

FINAL_RUN_SUMMARY = "Test Run Summary"

DEPENDECY_CHECK = "Dependency Check"

COMPUTATIONAL_PERFORMANCE = "Computational Performance Summary"

SHALLOW_CHECK_SUMMARY = "Validation and Size Check Summary"

CONSISTENCY_BASH = "Consistency Summary Between Ersilia and Bash Execution Outputs"

MODEL_OUTPUT = "Model Output Content Validation Summary"

"""

Configuration for a table.

"""

title: str

TableType.MODEL_INFORMATION_CHECKS: TableConfig(

title="Model Information Checks", headers=["Check", "Status"]

),

TableType.MODEL_FILE_CHECKS: TableConfig(

title="Model File Checks", headers=["Check", "Status"]

),

TableType.MODEL_DIRECTORY_SIZES: TableConfig(

title="Model Directory Sizes", headers=["Check", "Size"]

),

TableType.RUNNER_CHECKUP_STATUS: TableConfig(

title="Runner Checkup Status",

headers=["Runner", "Status"],

),

TableType.FINAL_RUN_SUMMARY: TableConfig(

title="Test Run Summary", headers=["Check", "Status"]

),

TableType.DEPENDECY_CHECK: TableConfig(

title="Dependency Check", headers=["Check", "Status"]

),

TableType.COMPUTATIONAL_PERFORMANCE: TableConfig(

title="Computational Performance Summary", headers=["Check", "Status"]

),

TableType.SHALLOW_CHECK_SUMMARY: TableConfig(

title="Validation and Size Check Results",

headers=["Check", "Details", "Status"],

),

TableType.MODEL_OUTPUT: TableConfig(

title="Model Output Content Validation Summary",

headers=["Check", "Detail", "Status"],

),

TableType.CONSISTENCY_BASH: TableConfig(

title="Consistency Summary Between Ersilia and Bash Execution Outputs",

headers=["Check", "Result", "Status"],

),

TableType.INSPECT_SUMMARY: TableConfig(

title="Inspect Summary", headers=["Check", "Status"]

),

"""

Enum for status configurations.

"""

PASSED = ("PASSED", "green", "✔")

FAILED = ("FAILED", "red", "✘")

WARNING = ("WARNING", "yellow", "⚠")

SUCCESS = ("SUCCESS", "green", "★")

NA = ("N/A", "dim", "~")

def __init__(self, label, color, icon):

self.label = label

self.color = color

self.icon = icon

def __str__(self):

"""

Execuetd a strategy for checking inspect commands.

Parameters

----------

check_function : callable

The function to check.

success_key : str

The key for success.

details_key : str

The key for details.

"""

def __init__(self, check_function, success_key, details_key):

self.check_function = check_function

self.success_key = success_key

self.details_key = details_key

def execute(self):

"""

Execute the check strategy.

Returns

-------

dict

The results of the check.

"""

if self.check_function is None:

return {}

result = self.check_function()

if result is None:

return {}

return {

self.success_key: result.success,

self.details_key: result.details,

"""

Service for inspecting models and their configurations.

Parameters

----------

dir : str, optional

Directory where the model is located.

model : str, optional

Model identifier.

remote : bool, optional

Flag indicating whether the model is remote.

config_json : str, optional

Path to the configuration JSON file.

credentials_json : str, optional

Path to the credentials JSON file.

Examples

--------

.. code-block:: python

inspector = InspectService(

dir="/path/to/model", model="model_id"

)

results = inspector.run()

"""

def __init__(

self,

dir: str,

model: str,

remote: bool = False,

config_json: str = None,

credentials_json: str = None,

):

super().__init__(config_json, credentials_json)

self.dir = dir

self.model = model

self.remote = remote

self.resolver = TemplateResolver(model_id=model, repo_path=self.dir)

def run(self, check_keys: list = None) -> dict:

"""

Run the inspection checks on the specified model.

Parameters

----------

check_keys : list, optional

A list of check keys to execute. If None, all checks will be executed.

Returns

-------

dict

A dictionary containing the results of the inspection checks.

Raises

------

ValueError

If the model is not specified.

KeyError

If any of the specified keys do not exist.

"""

def _transform_key(value):

if value is True:

return str(STATUS_CONFIGS.PASSED)

elif value is False:

return str(STATUS_CONFIGS.FAILED)

return value

inspector = ModelInspector(self.model, self.dir)

checks = self._get_checks(inspector)

output = {}

if check_keys:

for key in check_keys:

try:

strategy = checks.get(key)

if strategy.check_function:

output.update(strategy.execute())

except KeyError:

raise KeyError(f"Check '{key}' does not exist.")

else:

for key, strategy in checks.all():

if strategy.check_function:

output.update(strategy.execute())

output = {

" ".join(word.capitalize() for word in k.split("_")): _transform_key(v)

for k, v in output.items()

}

output = [(key, value) for key, value in output.items()]

return output

def _get_checks(self, inspector: ModelInspector) -> dict:

def create_check(check_fn, key, details):

return lambda: CheckStrategy(check_fn, key, details)

dependency_check = (

"Dockerfile" if self.resolver.is_bentoml() else "Install_YAML"

)

checks = {

"is_github_url_available": create_check(

inspector.check_repo_exists if self.remote else lambda: None,

"is_github_url_available",

"is_github_url_available_details",

),

"complete_metadata": create_check(

inspector.check_complete_metadata if self.remote else lambda: None,

"complete_metadata",

"complete_metadata_details",

),

"complete_folder_structure": create_check(

inspector.check_complete_folder_structure,

"complete_folder_structure",

"complete_folder_structure_details",

),

"docker_check": create_check(

inspector.check_dependencies_are_valid,

f"{dependency_check}_check",

"check_details",

),

"computational_performance_tracking": create_check(

inspector.check_computational_performance,

"computational_performance_tracking",

"computational_performance_tracking_details",

),

"extra_files_check": create_check(

inspector.check_no_extra_files if self.remote else lambda: None,

"extra_files_check",

"extra_files_check_details",

),

}

class LazyChecks:

def __init__(self, checks):

self._checks = checks

self._loaded = {}

def get(self, key):

if key not in self._loaded:

if key not in self._checks:

raise KeyError(f"Check '{key}' does not exist.")

self._loaded[key] = self._checks[key]()

return self._loaded[key]

def all(self):

for key in self._checks.keys():

yield key, self.get(key)

"""

Service for setting up the environment and fetching the model repository.

Parameters

----------

model_id : str

Identifier of the model.

dir : str

Directory where the model repository will be cloned.

from_github : bool

Flag indicating whether to fetch the repository from GitHub.

from_s3 : bool

Flag indicating whether to fetch the repository from S3.

logger : Any

Logger for logging messages.

"""

BASE_URL = "https://github.com/ersilia-os/"

def __init__(

self,

model_id: str,

dir: str,

from_github: bool,

from_s3: bool,

logger: Any,

):

self.model_id = model_id

self.dir = dir

self.logger = logger

self.from_github = from_github

self.from_s3 = from_s3

self.mc = ModelCard()

self.metadata = self.mc.get(model_id)

self.s3 = self.metadata.get("card", {}).get("S3") or self.metadata.get(

"metadata", {}

).get("S3")

self.repo_url = f"{self.BASE_URL}{self.model_id}"

self.overwrite = self._handle_overwrite()

self.github_down = GitHubDownloader(overwrite=self.overwrite)

self.s3_down = S3Downloader()

self.conda = SimpleConda()

def _handle_overwrite(self) -> bool:

if os.path.exists(self.dir):

self.logger.info(f"Directory {self.dir} already exists.")

return yes_no_input(

f"Directory {self.dir} already exists. Do you want to overwrite it? [Y/n]",

default_answer="n",

)

return False

def _download_s3(self):

if not self.overwrite and os.path.exists(self.dir):

self.logger.info("Skipping S3 download as user chose not to overwrite.")

return

tmp_file = os.path.join(make_temp_dir("ersilia-"), f"{self.model_id}.zip")

self.logger.info(f"Downloading model from S3 to temporary file: {tmp_file}")

self.s3_down.download_from_s3(

bucket_url=S3_BUCKET_URL_ZIP,

file_name=f"{self.model_id}.zip",

destination=tmp_file,

)

self.logger.info(f"Extracting model to: {self.dir}")

with zipfile.ZipFile(tmp_file, "r") as zip_ref:

zip_ref.extractall(EOS_TMP)

def _download_github(self):

try:

if not os.path.exists(EOS_TMP):

self.logger.info(f"Path does not exist. Creating: {EOS_TMP}")

os.makedirs(EOS_TMP, exist_ok=True)

except OSError as e:

self.logger.error(f"Failed to create directory {EOS_TMP}: {e}")

self.logger.info(f"Cloning repository from GitHub to: {EOS_TMP}")

self.github_down.clone(

org=GITHUB_ORG,

repo=self.model_id,

destination=self.dir,

)

def get_model(self):

if self.from_s3:

self._download_s3()

if self.from_github:

self._download_github()

@staticmethod

def run_command(

command: str, logger, capture_output: bool = False, shell: bool = True

) -> str:

"""

Run a shell command.

Parameters

----------

command : str

The command to run.

logger : logging.Logger

Logger for logging messages.

capture_output : bool, optional

Flag indicating whether to capture the command output.

shell : bool, optional

Flag indicating whether to run the command in the shell.

Returns

-------

str

The output of the command.

Raises

------

subprocess.CalledProcessError

If the command returns a non-zero exit code.

"""

try:

if capture_output:

result = subprocess.run(

command,

stdout=subprocess.PIPE,

stderr=subprocess.PIPE,

text=True,

check=True,

shell=shell,

)

return result.stdout

else:

process = subprocess.Popen(

command,

stdout=subprocess.PIPE,

stderr=subprocess.PIPE,

text=True,

shell=shell,

)

stdout_lines, stderr_lines = [], []

for line in iter(process.stdout.readline, ""):

if line.strip():

stdout_lines.append(line.strip())

logger.info(line.strip())

for line in iter(process.stderr.readline, ""):

if line.strip():

stderr_lines.append(line.strip())

logger.info(line.strip())

process.wait()

if process.returncode != 0:

raise subprocess.CalledProcessError(

returncode=process.returncode,

cmd=command,

output="\n".join(stdout_lines),

stderr="\n".join(stderr_lines),

)

return process.stdout

except subprocess.CalledProcessError as e:

logger.error(f"Error executing command: {e}")

if e.output:

logger.debug(f"Output: {e.output.strip()}")

if e.stderr:

logger.error(f"Error: {e.stderr.strip()}")

sys.exit(1)

except Exception as e:

logger.debug(f"Unexpected error: {e}")

sys.exit(1)

@staticmethod

def get_conda_env_location(model_id: str, logger) -> str:

"""

Get the location of the Conda environment for the model.

Parameters

----------

model_id : str

Identifier of the model.

logger : logging.Logger

Logger for logging messages.

Returns

-------

str

The location of the Conda environment.

Raises

------

subprocess.CalledProcessError

If the command to list Conda environments returns a non-zero exit code.

"""

try:

result = SetupService.run_command(

"conda env list", logger=logger, capture_output=True

)

for line in result.splitlines():

if line.startswith("#") or not line.strip():

continue

parts = line.split()

if parts[0] == model_id:

return parts[-1]

except subprocess.CalledProcessError as e:

print(f"Error running conda command: {e.stderr}")

except Exception as e:

print(f"Unexpected error: {e}")

"""

Service for handling input/output operations related to model testing.

Parameters

----------

logger : logging.Logger

Logger for logging messages.

model_id : str

Identifier of the model.

dir : str

Directory where the model repository is located.

Examples

--------

.. code-block:: python

ios = IOService(

logger=logger,

model_id="model_id",

dir="/path/to/dir",

)

ios.read_information()

"""

RUN_FILE = f"model/framework/{RUN_FILE}"

BENTOML_FILES = [

DOCKERFILE_FILE,

METADATA_JSON_FILE,

RUN_FILE,

"src/service.py",

"pack.py",

"README.md",

"LICENSE",

]

ERSILIAPACK_FILES = [

INSTALL_YAML_FILE,

METADATA_YAML_FILE,

PREDEFINED_EXAMPLE_FILES[0],

PREDEFINED_EXAMPLE_FILES[1],

RUN_FILE,

"README.md",

"LICENSE",

]

def __init__(self, logger, model_id: str, dir: str):

self.logger = logger

self.model_id = model_id

self.dir = dir

self.model_size = 0

self.console = Console()

self.check_results = []

self.simple_docker = SimpleDocker()

self.resolver = TemplateResolver(model_id=model_id, repo_path=self.dir)

def _run_check(

self, check_function, data, check_name: str, additional_info=None

) -> bool:

try:

if additional_info is not None:

check_function(additional_info)

else:

check_function(data)

self.check_results.append((check_name, str(STATUS_CONFIGS.PASSED)))

return True

except Exception as e:

self.logger.error(f"Check '{check_name}' failed: {e}")

self.check_results.append((check_name, str(STATUS_CONFIGS.FAILED)))

return False

def _get_metadata(self):

path = METADATA_JSON_FILE if self.resolver.is_bentoml() else METADATA_YAML_FILE

path = os.path.join(self.dir, path)

with open(path, "r") as file:

if path.endswith(".json"):

data = json.load(file)

elif path.endswith((".yml", ".yaml")):

data = yaml.safe_load(file)

else:

raise ValueError(f"Unsupported file format: {path}")

return data

def collect_and_save_json(self, results, output_file):

"""

Helper function to collect JSON results and save them to a file.

"""

aggregated_json = {}

for result in results:

aggregated_json.update(result)

with open(output_file, "w") as f:

json.dump(aggregated_json, f, indent=4)

def _create_json_data(self, rows, key):

def sanitize_name(name):

return re.sub(r"[ \-./]", "_", str(name).lower())

def parse_status(status):

if isinstance(status, str):

status = re.sub(

r"[-+]?\d*\.\d+|\d+", lambda m: str(float(m.group())), status

)

if "[green]✔" in status:

return True

elif "[red]✘" in status:

return False

else:

return re.sub(r"\[.*?\]", "", status).strip()

return status

def parse_performance(status):

return {

f"pred_{match[0]}": float(match[1])

for match in re.findall(

r"(\d+) predictions executed in (\d+\.\d{2}) seconds. \n", status

)

}

key = re.sub(r" ", "_", key.lower())

json_data = {}

for row in rows:

check_name = sanitize_name(row[0])

check_status = row[-1]

if check_name == "computational_performance_tracking_details":

json_data[check_name] = parse_performance(check_status)

else:

json_data[check_name] = parse_status(check_status)

return {key: json_data}

def _generate_table(self, title, headers, rows, large_table=False, merge=False):

f_col_width = 30 if large_table else 30

l_col_width = 50 if large_table else 10

d_col_width = 30 if not large_table else 20

table = Table(

title=Text(title, style="bold light_green"),

border_style="light_green",

show_lines=True,

)

table.add_column(headers[0], justify="left", width=f_col_width, style="bold")

for header in headers[1:-1]:

table.add_column(header, justify="center", width=d_col_width, style="bold")

table.add_column(headers[-1], justify="right", width=l_col_width, style="bold")

prev_value = None

for row in rows:

first_col = str(row[0])

if merge and first_col == prev_value:

first_col = ""

else:

prev_value = first_col

styled_row = [

Text(first_col, style="bold"),

*[str(cell) for cell in row[1:-1]],

row[-1],

]

table.add_row(*styled_row)

json_data = self._create_json_data(rows, title)

self.console.print(table)

return json_data

@staticmethod

def get_model_type(model_id: str, repo_path: str) -> str:

"""

Get the type of the model based on the repository contents.

Parameters

----------

model_id : str

Identifier of the model.

repo_path : str

Path to the model repository.

Returns

-------

str

The type of the model (e.g., PACK_METHOD_BENTOML, PACK_METHOD_FASTAPI).

"""

resolver = TemplateResolver(model_id=model_id, repo_path=repo_path)

if resolver.is_bentoml():

return PACK_METHOD_BENTOML

elif resolver.is_fastapi():

return PACK_METHOD_FASTAPI

else:

return None

def get_file_requirements(self) -> List[str]:

"""

Get the list of required files for the model.

Returns

-------

List[str]

List of required files.

Raises

------

ValueError

If the model type is unsupported.

"""

type = IOService.get_model_type(model_id=self.model_id, repo_path=self.dir)

if type == PACK_METHOD_BENTOML:

return self.BENTOML_FILES

elif type == PACK_METHOD_FASTAPI:

return self.ERSILIAPACK_FILES

else:

raise ValueError(f"Unsupported model type: {type}")

def read_information(self) -> dict:

"""

Read the information file for the model.

Returns

-------

dict

The contents of the information file.

Raises

------

FileNotFoundError

If the information file does not exist.

"""

file = os.path.join(EOS_TMP, self.model_id, METADATA_JSON_FILE)

if not os.path.exists(file):

raise FileNotFoundError(

f"Information file does not exist for model {self.model_id}"

)

with open(file, "r") as f:

return json.load(f)

def get_conda_env_size(self):

"""

Get the size of the Conda environment for the model.

Returns

-------

int

The size of the Conda environment in megabytes.

Raises

------

Exception

If there is an error calculating the size.

"""

try:

loc = SetupService.get_conda_env_location(self.model_id, self.logger)

return self.calculate_directory_size(loc)

except Exception as e:

self.logger.error(

f"Error calculating size of Conda environment '{self.model_id}': {e}"

)

return 0

def calculate_directory_size(self, path: str) -> int:

"""

Calculate the size of a directory.

Parameters

----------

path : str

The path to the directory.

Returns

-------

int

The size of the directory.

"""

try:

size_output = SetupService.run_command(

["du", "-sm", path],

logger=self.logger,

capture_output=True,

shell=False,

)

size = float(size_output.split()[0])

return size

except Exception as e:

self.logger.error(f"Error calculating directory size for {path}: {e}")

return 0

def calculate_image_size(self, tag="latest"):

"""

Calculate the size of a Docker image.

Parameters

----------

tag : str, optional

The tag of the Docker image (default is 'latest').

Returns

-------

str

The size of the Docker image.

"""

image_name = f"{DOCKERHUB_ORG}/{self.model_id}:{tag}"

client = docker.from_env()

try:

image = client.images.get(image_name)

size_bytes = image.attrs["Size"]

size_mb = size_bytes / (1024**2)

return f"{size_mb:.2f} MB", Checks.SIZE_CACL_SUCCESS.value

except docker.errors.ImageNotFound:

return f"Image '{image_name}' not found.", Checks.SIZE_CACL_FAILED.value

except Exception as e:

return f"An error occurred: {e}", Checks.SIZE_CACL_FAILED.value

@throw_ersilia_exception()

def get_directories_sizes(self) -> str:

"""

Get the sizes of the model directory.

Returns

-------

str

A string of containing size of the model directory

"""

dir_size = self.calculate_directory_size(self.dir)

dir_size = f"{dir_size:.2f}"

return dir_size

@throw_ersilia_exception()

def get_env_sizes(self) -> str:

"""

Get the sizes of the Conda environment directory.

Returns

-------

str

A string of containing size of the model environment

"""

env_size = self.get_conda_env_size()

env_size = f"{env_size:.2f}"

return env_size

def _extract_size(self, data, key="validation_and_size_check_results"):

sizes, keys = {}, ("environment_size_mb", "image_size_mb")

validation_results = next((item.get(key) for item in data if key in item), None)

if validation_results:

if keys[0] in validation_results:

env_size = validation_results[keys[0]]

self.logger.info(f"Environment Size: {env_size}")

sizes["Environment Size"] = float(env_size)

if keys[1] in validation_results:

img_size = validation_results[keys[1]]

self.logger.info(f"Image Size: {img_size}")

sizes["Image Size"] = float(img_size)

return sizes

def _extract_execution_times(self, data, key="computational_performance_summary"):

self.logger.info("Performance Extraction is started")

performance = {

"Computational Performance 1": None,

"Computational Performance 10": None,

"Computational Performance 100": None,

}

summary = next((item.get(key) for item in data if key in item), None)

if summary:

preds = summary.get("computational_performance_tracking_details")

performance["Computational Performance 1"] = float(preds.get("pred_1"))

performance["Computational Performance 10"] = float(preds.get("pred_10"))

performance["Computational Performance 100"] = float(preds.get("pred_100"))

return performance

def update_metadata(self, json_data):

"""

Processes JSON/YAML metadata to extract size and performance info and then updates them.

Parameters

----------

json_data : dict

Report data from the command output.

Returns

-------

dict

Updated metadata containing computed performance and size information.

"""

sizes = self._extract_size(json_data)

exec_times = self._extract_execution_times(json_data)

metadata = self._get_metadata()

metadata.update(sizes)

metadata.update(exec_times)

self._save_file(

metadata,

)