def start():
while not pathlib.Path("/mnt/workspace/init_script_run.txt").is_file():
time.sleep(0.2)
cache_dir = os.environ["CORTEX_CACHE_DIR"]
api_spec_path = os.environ["CORTEX_API_SPEC"]
job_spec_path = os.environ["CORTEX_JOB_SPEC"]
project_dir = os.environ["CORTEX_PROJECT_DIR"]
model_dir = os.getenv("CORTEX_MODEL_DIR")
tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT", "9000")
tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST", "localhost")
region = os.getenv("AWS_REGION")
has_multiple_servers = os.getenv("CORTEX_MULTIPLE_TF_SERVERS")
if has_multiple_servers:
with LockedFile("/run/used_ports.json", "r+") as f:
used_ports = json.load(f)
for port in used_ports.keys():
if not used_ports[port]:
tf_serving_port = port
used_ports[port] = True
break
f.seek(0)
json.dump(used_ports, f)
f.truncate()
api = get_api(api_spec_path, model_dir, cache_dir)
with open(job_spec_path) as json_file:
job_spec = json.load(json_file)
sqs_client = boto3.client("sqs", region_name=region)
client = api.predictor.initialize_client(
tf_serving_host=tf_serving_host, tf_serving_port=tf_serving_port
)
try:
log.info("loading the predictor from {}".format(api.predictor.path))
metrics_client = MetricsClient(api.statsd)
predictor_impl = api.predictor.initialize_impl(
project_dir=project_dir,
client=client,
metrics_client=metrics_client,
job_spec=job_spec,
)
except UserRuntimeException as err:
err.wrap(f"failed to start job {job_spec['job_id']}")
log.error(str(err), exc_info=True)
sys.exit(1)
except Exception as err:
capture_exception(err)
log.error(f"failed to start job {job_spec['job_id']}", exc_info=True)
sys.exit(1)
# crons only stop if an unhandled exception occurs
def check_if_crons_have_failed():
while True:
for cron in api.predictor.crons:
if not cron.is_alive():
os.kill(os.getpid(), signal.SIGQUIT)
time.sleep(1)
threading.Thread(target=check_if_crons_have_failed, daemon=True).start()
local_cache["api"] = api
local_cache["job_spec"] = job_spec
local_cache["predictor_impl"] = predictor_impl
local_cache["predict_fn_args"] = inspect.getfullargspec(predictor_impl.predict).args
local_cache["sqs_client"] = sqs_client
open("/mnt/workspace/api_readiness.txt", "a").close()
log.info("polling for batches...")
try:
sqs_handler = SQSHandler(
sqs_client=sqs_client,
queue_url=job_spec["sqs_url"],
renewal_period=MESSAGE_RENEWAL_PERIOD,
visibility_timeout=INITIAL_MESSAGE_VISIBILITY,
not_found_sleep_time=MESSAGE_NOT_FOUND_SLEEP,
message_wait_time=SQS_POLL_WAIT_TIME,
dead_letter_queue_url=job_spec.get("sqs_dead_letter_queue"),
stop_if_no_messages=True,
)
sqs_handler.start(
message_fn=handle_batch_message,
message_failure_fn=handle_batch_failure,
on_job_complete_fn=handle_on_job_complete,
)
except UserRuntimeException as err:
err.wrap(f"failed to run job {job_spec['job_id']}")
log.error(str(err), exc_info=True)
sys.exit(1)
except Exception as err:
capture_exception(err)
log.error(f"failed to run job {job_spec['job_id']}", exc_info=True)
sys.exit(1)
def main():
while not pathlib.Path("/mnt/workspace/init_script_run.txt").is_file():
time.sleep(0.2)
model_dir = os.getenv("CORTEX_MODEL_DIR")
cache_dir = os.environ["CORTEX_CACHE_DIR"]
api_spec_path = os.environ["CORTEX_API_SPEC"]
workload_path = os.environ["CORTEX_ASYNC_WORKLOAD_PATH"]
project_dir = os.environ["CORTEX_PROJECT_DIR"]
readiness_file = os.getenv("CORTEX_READINESS_FILE",
"/mnt/workspace/api_readiness.txt")
region = os.getenv("AWS_REGION")
queue_url = os.environ["CORTEX_QUEUE_URL"]
statsd_host = os.getenv("HOST_IP")
statsd_port = os.getenv("CORTEX_STATSD_PORT", "9125")
tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST")
tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT")
storage, api_spec = get_spec(api_spec_path, cache_dir, region)
sqs_client = boto3.client("sqs", region_name=region)
api = AsyncAPI(
api_spec=api_spec,
storage=storage,
storage_path=workload_path,
statsd_host=statsd_host,
statsd_port=int(statsd_port),
model_dir=model_dir,
)
try:
log.info(f"loading the predictor from {api.path}")
metrics_client = MetricsClient(api.statsd)
predictor_impl = api.initialize_impl(
project_dir,
metrics_client,
tf_serving_host=tf_serving_host,
tf_serving_port=tf_serving_port,
)
except UserRuntimeException as err:
err.wrap(f"failed to initialize predictor implementation")
log.error(str(err), exc_info=True)
sys.exit(1)
except Exception as err:
capture_exception(err)
log.error(f"failed to initialize predictor implementation",
exc_info=True)
sys.exit(1)
local_cache["api"] = api
local_cache["predictor_impl"] = predictor_impl
local_cache["sqs_client"] = sqs_client
local_cache["storage_client"] = storage
local_cache["predict_fn_args"] = inspect.getfullargspec(
predictor_impl.predict).args
open(readiness_file, "a").close()
log.info("polling for workloads...")
try:
sqs_handler = SQSHandler(
sqs_client=sqs_client,
queue_url=queue_url,
renewal_period=MESSAGE_RENEWAL_PERIOD,
visibility_timeout=INITIAL_MESSAGE_VISIBILITY,
not_found_sleep_time=MESSAGE_NOT_FOUND_SLEEP,
message_wait_time=SQS_POLL_WAIT_TIME,
)
sqs_handler.start(message_fn=handle_workload,
message_failure_fn=handle_workload_failure)
except UserRuntimeException as err:
log.error(str(err), exc_info=True)
sys.exit(1)
except Exception as err:
capture_exception(err)
log.error(str(err), exc_info=True)
sys.exit(1)