def get_spec( provider: str, spec_path: str, cache_dir: str, region: Optional[str] = None, spec_name: str = "api_spec.json", ) -> Tuple[Union[S3, GCS], dict]: """ Args: provider: "aws" or "gcp". spec_path: Path to API spec (i.e. "s3://cortex-dev-0/apis/iris-classifier/api/69b93378fa5c0218-jy1fjtyihu-9fcc10739e7fc8050cefa8ca27ece1ee/master-spec.json"). cache_dir: Local directory where the API spec gets saved to. region: Region of the bucket. Only required for "S3" provider. spec_name: File name of the spec as it is saved on disk. """ if provider == "aws": bucket, key = S3.deconstruct_s3_path(spec_path) storage = S3(bucket=bucket, region=region) elif provider == "gcp": bucket, key = GCS.deconstruct_gcs_path(spec_path) storage = GCS(bucket=bucket) else: raise ValueError('invalid "provider" argument') local_spec_path = os.path.join(cache_dir, spec_name) if not os.path.isfile(local_spec_path): storage.download_file(key, local_spec_path) return storage, read_json(local_spec_path)
def start(args): download_config = json.loads(base64.urlsafe_b64decode(args.download)) for download_arg in download_config["download_args"]: from_path = download_arg["from"] to_path = download_arg["to"] item_name = download_arg.get("item_name", "") bucket_name, prefix = S3.deconstruct_s3_path(from_path) client = S3(bucket_name) if item_name != "": if download_arg.get("hide_from_log", False): logger.info("downloading {}".format(item_name)) else: logger.info("downloading {} from {}".format( item_name, from_path)) if download_arg.get("to_file", False): client.download_file(prefix, to_path) else: client.download(prefix, to_path) if download_arg.get("unzip", False): if item_name != "" and not download_arg.get( "hide_unzipping_log", False): logger.info("unzipping {}".format(item_name)) if download_arg.get("to_file", False): util.extract_zip(to_path, delete_zip_file=True) else: util.extract_zip(os.path.join(to_path, os.path.basename(from_path)), delete_zip_file=True) if download_config.get("last_log", "") != "": logger.info(download_config["last_log"])
def start(args): download_config = json.loads(base64.urlsafe_b64decode(args.download)) for download_arg in download_config["download_args"]: from_path = download_arg["from"] to_path = download_arg["to"] item_name = download_arg.get("item_name", "") if from_path.startswith("s3://"): bucket_name, prefix = S3.deconstruct_s3_path(from_path) client = S3(bucket_name, client_config={}) elif from_path.startswith("gs://"): bucket_name, prefix = GCS.deconstruct_gcs_path(from_path) client = GCS(bucket_name) else: raise ValueError( '"from" download arg can either have the "s3://" or "gs://" prefixes' ) if item_name != "": if download_arg.get("hide_from_log", False): logger().info("downloading {}".format(item_name)) else: logger().info("downloading {} from {}".format( item_name, from_path)) if download_arg.get("to_file", False): client.download_file(prefix, to_path) else: client.download(prefix, to_path) if download_arg.get("unzip", False): if item_name != "" and not download_arg.get( "hide_unzipping_log", False): logger().info("unzipping {}".format(item_name)) if download_arg.get("to_file", False): util.extract_zip(to_path, delete_zip_file=True) else: util.extract_zip(os.path.join(to_path, os.path.basename(from_path)), delete_zip_file=True) if download_config.get("last_log", "") != "": logger().info(download_config["last_log"])
def get_spec( spec_path: str, cache_dir: str, region: str, spec_name: str = "api_spec.json", ) -> Tuple[S3, dict]: """ Args: spec_path: Path to API spec (i.e. "s3://cortex-dev-0/apis/iris-classifier/api/69b93378fa5c0218-jy1fjtyihu-9fcc10739e7fc8050cefa8ca27ece1ee/master-spec.json"). cache_dir: Local directory where the API spec gets saved to. region: Region of the bucket. spec_name: File name of the spec as it is saved on disk. """ bucket, key = S3.deconstruct_s3_path(spec_path) storage = S3(bucket=bucket, region=region) local_spec_path = os.path.join(cache_dir, spec_name) if not os.path.isfile(local_spec_path): storage.download_file(key, local_spec_path) return storage, read_json(local_spec_path)
def get_job_spec(storage, cache_dir, job_spec_path): local_spec_path = os.path.join(cache_dir, "job_spec.json") _, key = S3.deconstruct_s3_path(job_spec_path) storage.download_file(key, local_spec_path) with open(local_spec_path) as f: return json.load(f)
def find_all_cloud_models( is_dir_used: bool, models_dir: str, predictor_type: PredictorType, cloud_paths: List[str], cloud_model_names: List[str], ) -> Tuple[List[str], Dict[str, List[str]], List[str], List[List[str]], List[List[datetime.datetime]], List[str], List[str], ]: """ Get updated information on all models that are currently present on the cloud upstreams. Information on the available models, versions, last edit times, the subpaths of each model, and so on. Args: is_dir_used: Whether predictor:models:dir is used or not. models_dir: The value of predictor:models:dir in case it's present. Ignored when not required. predictor_type: The predictor type. cloud_paths: The cloud model paths as they are specified in predictor:models:path/predictor:models:paths/predictor:models:dir is used. Ignored when not required. cloud_model_names: The cloud model names as they are specified in predictor:models:paths:name when predictor:models:paths is used or the default name of the model when predictor:models:path is used. Ignored when not required. Returns: The tuple with the following elements: model_names - a list with the names of the models (i.e. bert, gpt-2, etc) and they are unique versions - a dictionary with the keys representing the model names and the values being lists of versions that each model has. For non-versioned model paths ModelVersion.NOT_PROVIDED, the list will be empty. model_paths - a list with the prefix of each model. sub_paths - a list of filepaths lists for each file of each model. timestamps - a list of timestamps lists representing the last edit time of each versioned model. bucket_providers - a list of the bucket providers for each model. Can be "s3" or "gs". bucket_names - a list of the bucket names of each model. """ # validate models stored in cloud (S3 or GS) that were specified with predictor:models:dir field if is_dir_used: if S3.is_valid_s3_path(models_dir): bucket_name, models_path = S3.deconstruct_s3_path(models_dir) client = S3(bucket_name) if GCS.is_valid_gcs_path(models_dir): bucket_name, models_path = GCS.deconstruct_gcs_path(models_dir) client = GCS(bucket_name) sub_paths, timestamps = client.search(models_path) model_paths, ooa_ids = validate_models_dir_paths( sub_paths, predictor_type, models_path) model_names = [ os.path.basename(model_path) for model_path in model_paths ] model_paths = [ model_path for model_path in model_paths if os.path.basename(model_path) in model_names ] model_paths = [ model_path + "/" * (not model_path.endswith("/")) for model_path in model_paths ] if S3.is_valid_s3_path(models_dir): bucket_providers = len(model_paths) * ["s3"] if GCS.is_valid_gcs_path(models_dir): bucket_providers = len(model_paths) * ["gs"] bucket_names = len(model_paths) * [bucket_name] sub_paths = len(model_paths) * [sub_paths] timestamps = len(model_paths) * [timestamps] # validate models stored in cloud (S3 or GS) that were specified with predictor:models:paths field if not is_dir_used: sub_paths = [] ooa_ids = [] model_paths = [] model_names = [] timestamps = [] bucket_providers = [] bucket_names = [] for idx, path in enumerate(cloud_paths): if S3.is_valid_s3_path(path): bucket_name, model_path = S3.deconstruct_s3_path(path) client = S3(bucket_name) elif GCS.is_valid_gcs_path(path): bucket_name, model_path = GCS.deconstruct_gcs_path(path) client = GCS(bucket_name) else: continue sb, model_path_ts = client.search(model_path) try: ooa_ids.append( validate_model_paths(sb, predictor_type, model_path)) except CortexException: continue model_paths.append(model_path) model_names.append(cloud_model_names[idx]) bucket_names.append(bucket_name) sub_paths += [sb] timestamps += [model_path_ts] if S3.is_valid_s3_path(path): bucket_providers.append("s3") if GCS.is_valid_gcs_path(path): bucket_providers.append("gs") # determine the detected versions for each cloud model # if the model was not versioned, then leave the version list empty versions = {} for model_path, model_name, model_ooa_ids, bucket_sub_paths in zip( model_paths, model_names, ooa_ids, sub_paths): if ModelVersion.PROVIDED not in model_ooa_ids: versions[model_name] = [] continue model_sub_paths = [ os.path.relpath(sub_path, model_path) for sub_path in bucket_sub_paths ] model_versions_paths = [ path for path in model_sub_paths if not path.startswith("../") ] model_versions = [ util.get_leftmost_part_of_path(model_version_path) for model_version_path in model_versions_paths ] model_versions = list(set(model_versions)) versions[model_name] = model_versions # pick up the max timestamp for each versioned model aux_timestamps = [] for model_path, model_name, bucket_sub_paths, sub_path_timestamps in zip( model_paths, model_names, sub_paths, timestamps): model_ts = [] if len(versions[model_name]) == 0: masks = list( map( lambda x: x.startswith(model_path + "/" * (model_path[-1] != "/")), bucket_sub_paths, )) model_ts = [max(itertools.compress(sub_path_timestamps, masks))] for version in versions[model_name]: masks = list( map( lambda x: x.startswith( os.path.join(model_path, version) + "/"), bucket_sub_paths, )) model_ts.append(max(itertools.compress(sub_path_timestamps, masks))) aux_timestamps.append(model_ts) timestamps = aux_timestamps # type: List[List[datetime.datetime]] # model_names - a list with the names of the models (i.e. bert, gpt-2, etc) and they are unique # versions - a dictionary with the keys representing the model names and the values being lists of versions that each model has. # For non-versioned model paths ModelVersion.NOT_PROVIDED, the list will be empty # model_paths - a list with the prefix of each model # sub_paths - a list of filepaths lists for each file of each model # timestamps - a list of timestamps lists representing the last edit time of each versioned model # bucket_providers - bucket providers # bucket_names - names of the buckets return model_names, versions, model_paths, sub_paths, timestamps, bucket_providers, bucket_names
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") bucket, key = S3.deconstruct_s3_path(workload_path) storage = S3(bucket=bucket, region=region) with open(api_spec_path) as json_file: api_spec = json.load(json_file) sqs_client = boto3.client("sqs", region_name=region) api = AsyncAPI( api_spec=api_spec, storage=storage, storage_path=key, 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)