# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing,

# software distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import os

import shutil

import time

from subprocess import PIPE, STDOUT, CompletedProcess, run

from typing import List

import requests

from sparsezoo import Model

def delete_file(filename: str):

if os.path.exists(filename):

os.remove(filename)

def is_server_up(url: str) -> bool:

try:

res = requests.get(url)

return res.status_code is not None

except Exception:

return False

def wait_for_server(url: str, retries: int, interval: int = 1) -> bool:

for _ in range(retries):

is_up = is_server_up(url)

if is_up:

print(f"server up and ready in ~{_} seconds...")

return True

time.sleep(interval)

return False

def run_command(command: List[str]) -> CompletedProcess:

"""

Run given command with custom config and return the completed process.

:param command: command to be executed (formatted as `subprocess.run` expects)

:return: completed process as received from `subprocess.run`

"""

return run(command, stdout=PIPE, stderr=STDOUT, check=False, encoding="utf-8")

def predownload_stub(stub: str, copy_framework_files: bool = False) -> Model:

"""

Download a model based on SparseZoo stub and return the Model object. If

`copy_framework_files` is True (default: False), also copy model’s config.json and

tokenizer.json files from the framework subfolder (e.g. `pytorch`) up into the model

root folder.

:return: SparseZoo Model object of downloaded model

"""

model = Model(stub)

model_path = model.path

if copy_framework_files:

# required for `deepsparse.transformers.run_inference` on local model files

config_path = model.deployment.default.get_file("config.json").path

tokenizer_config_path = model.deployment.default.get_file("tokenizer.json").path

shutil.copy(config_path, model_path)

shutil.copy(tokenizer_config_path, model_path)

return model

def find_free_port():

import socket

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

s.bind(("0.0.0.0", 0))

portnum = s.getsockname()[1]

s.close()

return portnum