def local_client(
aws_access_key_id: str,
aws_secret_access_key: str,
aws_region: str,
) -> Client:
"""
Initialize a client to deploy and manage APIs locally.
The specified AWS credentials will be used by the CLI to download models
from S3 and authenticate to ECR, and will be set in your Predictor.
Args:
aws_access_key_id: AWS access key ID.
aws_secret_access_key: AWS secret access key.
aws_region: AWS region.
Returns:
Cortex client that can be used to deploy and manage APIs locally.
"""
args = [
"env",
"configure",
"--provider",
"local",
"--aws-region",
aws_region,
"--aws-access-key-id",
aws_access_key_id,
"--aws-secret-access-key",
aws_secret_access_key,
]
run_cli(args, hide_output=True)
return Client("local")
def new_client(
name: str,
operator_endpoint: str,
) -> Client:
"""
Create a new environment to connect to an existing Cortex Cluster, and initialize a client to deploy and manage APIs on that cluster.
Args:
name: Name of the environment to create.
operator_endpoint: The endpoint for the operator of your Cortex Cluster. You can get this endpoint by running the CLI command `cortex cluster info` for an AWS provider or `cortex cluster-gcp info` for a GCP provider.
Returns:
Cortex client that can be used to deploy and manage APIs on a Cortex Cluster.
"""
cli_args = [
"env",
"configure",
name,
"--operator-endpoint",
operator_endpoint,
]
run_cli(cli_args, hide_output=True)
return client(name)
def env_delete(name: str):
"""
Delete an environment configured on this machine.
Args:
name: Name of the environment to delete.
"""
run_cli(["env", "delete", name], hide_output=True)
def stream_api_logs(
self,
api_name: str,
):
"""
Stream the logs of an API.
Args:
api_name: Name of the API.
"""
args = ["logs", api_name, "--env", self.env_name, "-y"]
run_cli(args)
def refresh(self, api_name: str, force: bool = False):
"""
Restart all of the replicas for a Realtime API without downtime.
Args:
api_name: Name of the API to refresh.
force: Override an already in-progress API update.
"""
args = ["refresh", api_name, "--env", self.env_name, "-o", "json"]
if force:
args.append("--force")
run_cli(args, hide_output=True)
def stream_job_logs(
self,
api_name: str,
job_id: str,
):
"""
Stream the logs of a Job.
Args:
api_name: Name of the Batch API.
job_id: Job ID.
"""
args = ["logs", api_name, job_id, "--env", self.env_name, "-y"]
run_cli(args)
def stop_job(self, api_name: str, job_id: str, keep_cache: bool = False):
"""
Stop a running job.
Args:
api_name: Name of the Batch/Task API.
job_id: ID of the Job to stop.
"""
args = [
"delete",
api_name,
job_id,
"--env",
self.env_name,
"-o",
"json",
]
run_cli(args)
def deploy_from_file(
self,
config_file: str,
force: bool = False,
wait: bool = False,
) -> Dict:
"""
Deploy or update APIs specified in a configuration file.
Args:
config_file: Local path to a yaml file defining Cortex API(s). See https://docs.cortex.dev/v/master/ for schema.
force: Override any in-progress api updates.
wait: Block until the API is ready.
Returns:
Deployment status, API specification, and endpoint for each API.
"""
args = [
"deploy",
config_file,
"--env",
self.env_name,
"-o",
"json",
"-y",
]
if force:
args.append("--force")
output = run_cli(args, hide_output=True)
deploy_results = json.loads(output.strip())
deploy_result = deploy_results[0]
if not wait:
return deploy_result
api_name = deploy_result["api"]["spec"]["name"]
if (
deploy_result["api"]["spec"]["kind"] != "RealtimeAPI"
and deploy_result["api"]["spec"]["kind"] != "AsyncAPI"
):
return deploy_result
while True:
time.sleep(5)
api = self.get_api(api_name)
if api["status"]["status_code"] != "status_updating":
break
return api
def list_apis(self) -> list:
"""
List all APIs in the environment.
Returns:
List of APIs, including information such as the API specification, endpoint, status, and metrics (if applicable).
"""
args = ["get", "-o", "json", "--env", self.env_name]
output = run_cli(args, hide_output=True)
return json.loads(output.strip())
def cluster_client(
name: str,
provider: str,
operator_endpoint: str,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
) -> Client:
"""
Create a new environment to connect to an existing Cortex Cluster, and initialize a client to deploy and manage APIs on that cluster.
Args:
name: Name of the environment to create.
provider: The provider of your Cortex cluster. Can be "aws" or "gcp".
operator_endpoint: The endpoint for the operator of your Cortex Cluster. You can get this endpoint by running the CLI command `cortex cluster info` for an AWS provider or `cortex cluster-gcp info` for a GCP provider.
aws_access_key_id: AWS access key ID. Required when `provider` is set to "aws".
aws_secret_access_key: AWS secret access key. Required when `provider` is set to "aws".
Returns:
Cortex client that can be used to deploy and manage APIs on a Cortex Cluster.
"""
cli_args = [
"env",
"configure",
name,
"--provider",
provider,
"--operator-endpoint",
operator_endpoint,
]
if provider == "aws":
cli_args += [
"--aws-access-key-id",
aws_access_key_id,
"--aws-secret-access-key",
aws_secret_access_key,
]
run_cli(cli_args, hide_output=True)
return Client(name)
def delete_api(self, api_name: str, keep_cache: bool = False):
"""
Delete an API.
Args:
api_name: Name of the API to delete.
keep_cache: Whether to retain the cached data for this API.
"""
args = [
"delete",
api_name,
"--env",
self.env_name,
"--force",
"-o",
"json",
]
if keep_cache:
args.append("--keep-cache")
run_cli(args, hide_output=True)
def get_api(self, api_name: str) -> dict:
"""
Get information about an API.
Args:
api_name: Name of the API.
Returns:
Information about the API, including the API specification, endpoint, status, and metrics (if applicable).
"""
output = run_cli(["get", api_name, "--env", self.env_name, "-o", "json"], hide_output=True)
apis = json.loads(output.strip())
return apis[0]
def cluster_client(
name: str,
operator_endpoint: str,
aws_access_key_id: str,
aws_secret_access_key: str,
) -> Client:
"""
Create a new environment to connect to an existing Cortex Cluster, and initialize a client to deploy and manage APIs on that cluster.
Args:
name: Name of the environment to create.
operator_endpoint: The endpoint for the operator of your Cortex Cluster. You can get this endpoint by running the CLI command `cortex cluster info`.
aws_access_key_id: AWS access key ID.
aws_secret_access_key: AWS secret access key.
Returns:
Cortex client that can be used to deploy and manage APIs on a Cortex Cluster.
"""
run_cli(
[
"env",
"configure",
name,
"--provider",
"aws",
"--operator-endpoint",
operator_endpoint,
"--aws-access-key-id",
aws_access_key_id,
"--aws-secret-access-key",
aws_secret_access_key,
],
hide_output=True,
)
return Client(name)
def get_job(self, api_name: str, job_id: str) -> dict:
"""
Get information about a submitted job.
Args:
api_name: Name of the Batch/Task API.
job_id: Job ID.
Returns:
Information about the job, including the job status, worker status, and job progress.
"""
args = ["get", api_name, job_id, "--env", self.env_name, "-o", "json"]
output = run_cli(args, hide_output=True)
return json.loads(output.strip())
def patch(self, api_spec: dict, force: bool = False) -> dict:
"""
Update the api specification for an API that has already been deployed.
Args:
api_spec: The new api specification to apply
force: Override an already in-progress API update.
"""
cortex_yaml_file = cli_config_dir() / "deployments" / f"cortex-{str(uuid.uuid4())}.yaml"
with util.open_temporarily(cortex_yaml_file, "w") as f:
yaml.dump([api_spec], f)
args = ["patch", cortex_yaml_file, "--env", self.env_name, "-o", "json"]
if force:
args.append("--force")
output = run_cli(args, hide_output=True)
return json.loads(output.strip())
def deploy(
self,
config_file: str,
force: bool = False,
wait: bool = False,
) -> list:
"""
Deploy or update APIs specified in the config_file.
Args:
config_file: Local path to a yaml file defining Cortex APIs.
force: Override any in-progress api updates.
wait: Streams logs until the APIs are ready.
Returns:
Deployment status, API specification, and endpoint for each API.
"""
args = [
"deploy",
config_file,
"--env",
self.env,
"-o",
"mixed",
]
if force:
args.append("--force")
output = run_cli(args, mixed_output=True)
deploy_results = json.loads(output.strip())
if not wait:
return deploy_results
def stream_to_stdout(process):
for c in iter(lambda: process.stdout.read(1), ""):
sys.stdout.write(c)
for deploy_result in deploy_results:
api_name = deploy_result["api"]["spec"]["name"]
kind = deploy_result["api"]["spec"]["kind"]
if kind != "RealtimeAPI":
continue
env = os.environ.copy()
env["CORTEX_CLI_INVOKER"] = "python"
process = subprocess.Popen(
[get_cli_path(), "logs", "--env", self.env, api_name],
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
encoding="utf8",
env=env,
)
streamer = threading.Thread(target=stream_to_stdout,
args=[process])
streamer.start()
while process.poll() is None:
api = self.get_api(api_name)
if api["status"]["status_code"] != "status_updating":
if api["status"]["status_code"] == "status_live":
time.sleep(2)
process.terminate()
break
time.sleep(2)
return deploy_results
def env_list() -> list:
"""
List all environments configured on this machine.
"""
output = run_cli(["env", "list", "--output", "json"], hide_output=True)
return json.loads(output.strip())
def _deploy(
self,
config_file: str,
force: bool = False,
wait: bool = False,
) -> list:
"""
Deploy or update APIs specified in the config_file.
Args:
config_file: Local path to a yaml file defining Cortex APIs.
force: Override any in-progress api updates.
wait: Streams logs until the APIs are ready.
Returns:
Deployment status, API specification, and endpoint for each API.
"""
args = [
"deploy",
config_file,
"--env",
self.env_name,
"-o",
"json",
"-y",
]
if force:
args.append("--force")
output = run_cli(args, hide_output=True)
deploy_results = json.loads(output.strip())
deploy_result = deploy_results[0]
if not wait:
return deploy_result
# logging immediately will show previous versions of the replica terminating;
# wait a few seconds for the new replicas to start initializing
time.sleep(5)
def stream_to_stdout(process):
for c in iter(lambda: process.stdout.read(1), ""):
sys.stdout.write(c)
sys.stdout.flush()
api_name = deploy_result["api"]["spec"]["name"]
if deploy_result["api"]["spec"]["kind"] != "RealtimeAPI":
return deploy_result
env = os.environ.copy()
env["CORTEX_CLI_INVOKER"] = "python"
process = subprocess.Popen(
[get_cli_path(), "logs", "--env", self.env_name, api_name, "-y"],
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
encoding="utf8",
errors=
"replace", # replace non-utf8 characters with `?` instead of failing
env=env,
)
streamer = threading.Thread(target=stream_to_stdout, args=[process])
streamer.start()
while process.poll() is None:
api = self.get_api(api_name)
if api["status"]["status_code"] != "status_updating":
time.sleep(
5) # accommodate latency in log streaming from the cluster
process.terminate()
break
time.sleep(5)
streamer.join(timeout=10)
return api
