def infra_teardown(pytestconfig, core_url, serving_url): client = Client(core_url=core_url, serving_url=serving_url) client.set_project(PROJECT_NAME) marker = pytestconfig.getoption("-m") yield marker if marker == 'dataflow_runner': ingest_jobs = client.list_ingest_jobs() ingest_jobs = [client.list_ingest_jobs(job.id)[0].external_id for job in ingest_jobs if job.status == IngestionJobStatus.RUNNING] cwd = os.getcwd() with open(f"{cwd}/ingesting_jobs.txt", "w+") as output: for job in ingest_jobs: output.write('%s\n' % job) else: print('Cleaning up not required')
def ingest_job_restart(job_id: str): """ Restart job for id. Waits for the job to fully restart. """ # find ingestion job for id feast_client = Client() jobs = feast_client.list_ingest_jobs(job_id=job_id) if len(jobs) < 1: print(f"Ingestion Job with id {job_id} could not be found") sys.exit(1) job = jobs[0] feast_client.restart_ingest_job(job)
def ingest_job_describe(job_id: str): """ Describe the ingestion job with the given id. """ # find ingestion job for id feast_client = Client() jobs = feast_client.list_ingest_jobs(job_id=job_id) if len(jobs) < 1: print(f"Ingestion Job with id {job_id} could not be found") sys.exit(1) job = jobs[0] # pretty render ingestion job as yaml print( yaml.dump(yaml.safe_load(str(job)), default_flow_style=False, sort_keys=False))
def ingest_job_stop(wait: bool, timeout: int, job_id: str): """ Stop ingestion job for id. """ # find ingestion job for id feast_client = Client() jobs = feast_client.list_ingest_jobs(job_id=job_id) if len(jobs) < 1: print(f"Ingestion Job with id {job_id} could not be found") sys.exit(1) job = jobs[0] feast_client.stop_ingest_job(job) # wait for ingestion job to stop if wait: job.wait(IngestionJobStatus.ABORTED, timeout=timeout)
def ingest_job_list(job_id, feature_set_ref, store_name): """ List ingestion jobs """ # parse feature set reference if feature_set_ref is not None: feature_set_ref = FeatureSetRef.from_str(feature_set_ref) # pull & render ingestion jobs as a table feast_client = Client() table = [] for ingest_job in feast_client.list_ingest_jobs( job_id=job_id, feature_set_ref=feature_set_ref, store_name=store_name): table.append( [ingest_job.id, IngestionJobStatus.Name(ingest_job.status)]) from tabulate import tabulate print(tabulate(table, headers=["ID", "STATUS"], tablefmt="plain"))