def main(step_run_ref_filepath, pipeline_zip): # Extract any zip files to a temporary directory and add that temporary directory # to the site path so the contained files can be imported. # # We can't rely on pip or other packaging tools because the zipped files might not # even be Python packages. with tempfile.TemporaryDirectory() as tmp: print('Extracting {}'.format(pipeline_zip)) with zipfile.ZipFile(pipeline_zip) as zf: zf.extractall(tmp) site.addsitedir(tmp) print('Loading step run ref') # We can use regular local filesystem APIs to access DBFS inside the Databricks runtime. with open(step_run_ref_filepath, 'rb') as handle: step_run_ref = pickle.load(handle) print('Step run ref:') print(step_run_ref) print('Setting up storage credentials') setup_storage(step_run_ref) print('Running pipeline') events = list(run_step_from_ref(step_run_ref)) print('Saving events to DBFS') events_filepath = os.path.dirname( step_run_ref_filepath) + '/' + PICKLED_EVENTS_FILE_NAME with open(events_filepath, 'wb') as handle: pickle.dump(serialize_value(events), handle)
def main(step_run_ref_bucket, s3_dir_key): session = boto3.client("s3") file_manager = S3FileManager(session, step_run_ref_bucket, "") file_handle = S3FileHandle(step_run_ref_bucket, s3_dir_key) step_run_ref_data = file_manager.read_data(file_handle) step_run_ref = pickle.loads(step_run_ref_data) events_bucket = step_run_ref_bucket events_s3_key = os.path.dirname( s3_dir_key) + "/" + PICKLED_EVENTS_FILE_NAME def put_events(events): file_obj = io.BytesIO(pickle.dumps(events)) session.put_object(Body=file_obj, Bucket=events_bucket, Key=events_s3_key) # Set up a thread to handle writing events back to the plan process, so execution doesn't get # blocked on remote communication events_queue = Queue() event_writing_thread = Thread( target=event_writing_loop, kwargs=dict(events_queue=events_queue, put_events_fn=put_events), ) event_writing_thread.start() with DagsterInstance.ephemeral() as instance: try: for event in run_step_from_ref(step_run_ref, instance): events_queue.put(event) finally: events_queue.put(DONE) event_writing_thread.join()
def main(step_run_ref_filepath, pipeline_zip): # Extract any zip files to a temporary directory and add that temporary directory # to the site path so the contained files can be imported. # # We can't rely on pip or other packaging tools because the zipped files might not # even be Python packages. with tempfile.TemporaryDirectory() as tmp: with zipfile.ZipFile(pipeline_zip) as zf: zf.extractall(tmp) site.addsitedir(tmp) # We can use regular local filesystem APIs to access DBFS inside the Databricks runtime. with open(step_run_ref_filepath, "rb") as handle: step_run_ref = pickle.load(handle) setup_storage(step_run_ref) with DagsterInstance.ephemeral() as instance: events = list(run_step_from_ref(step_run_ref, instance)) events_filepath = os.path.dirname( step_run_ref_filepath) + "/" + PICKLED_EVENTS_FILE_NAME with open(events_filepath, "wb") as handle: pickle.dump(serialize_value(events), handle)
def main(step_run_ref_path): file_manager = LocalFileManager(".") file_handle = LocalFileHandle(step_run_ref_path) step_run_ref = pickle.loads(file_manager.read_data(file_handle)) events = list(run_step_from_ref(step_run_ref)) events_out_path = os.path.join(os.path.dirname(step_run_ref_path), PICKLED_EVENTS_FILE_NAME) with open(events_out_path, "wb") as events_file: pickle.dump(events, events_file)
def main(step_run_ref_path: str) -> None: file_manager = LocalFileManager(".") file_handle = LocalFileHandle(step_run_ref_path) step_run_ref = pickle.loads(file_manager.read_data(file_handle)) with DagsterInstance.ephemeral() as instance: events = list(run_step_from_ref(step_run_ref, instance)) events_out_path = os.path.join(os.path.dirname(step_run_ref_path), PICKLED_EVENTS_FILE_NAME) with open(events_out_path, "wb") as events_file: pickle.dump(events, events_file)
def main(step_run_ref_bucket, s3_dir_key): session = boto3.client('s3') file_manager = S3FileManager(session, step_run_ref_bucket, '') file_handle = S3FileHandle(step_run_ref_bucket, s3_dir_key) step_run_ref_data = file_manager.read_data(file_handle) step_run_ref = pickle.loads(step_run_ref_data) events = list(run_step_from_ref(step_run_ref)) file_obj = io.BytesIO(pickle.dumps(events)) events_key = os.path.dirname(s3_dir_key) + '/' + PICKLED_EVENTS_FILE_NAME session.put_object(Body=file_obj, Bucket=step_run_ref_bucket, Key=events_key)
def main(step_run_ref_path: str) -> None: file_manager = LocalFileManager(".") file_handle = LocalFileHandle(step_run_ref_path) step_run_ref = pickle.loads(file_manager.read_data(file_handle)) all_events: List[EventLogEntry] = [] try: instance = external_instance_from_step_run_ref( step_run_ref, event_listener_fn=all_events.append) # consume entire step iterator list(run_step_from_ref(step_run_ref, instance)) finally: events_out_path = os.path.join(os.path.dirname(step_run_ref_path), PICKLED_EVENTS_FILE_NAME) with open(events_out_path, "wb") as events_file: pickle.dump(serialize_value(all_events), events_file)
def main( step_run_ref_filepath, setup_filepath, dagster_job_zip, ): # Extract any zip files to a temporary directory and add that temporary directory # to the site path so the contained files can be imported. # # We can't rely on pip or other packaging tools because the zipped files might not # even be Python packages. with tempfile.TemporaryDirectory() as tmp: with zipfile.ZipFile(dagster_job_zip) as zf: zf.extractall(tmp) site.addsitedir(tmp) # We can use regular local filesystem APIs to access DBFS inside the Databricks runtime. with open(setup_filepath, "rb") as handle: databricks_config = pickle.load(handle) # sc and dbutils are globally defined in the Databricks runtime. databricks_config.setup(dbutils, sc) # noqa pylint: disable=undefined-variable with open(step_run_ref_filepath, "rb") as handle: step_run_ref = pickle.load(handle) print("Running dagster job") # noqa pylint: disable=print-call events_filepath = os.path.dirname( step_run_ref_filepath) + "/" + PICKLED_EVENTS_FILE_NAME def put_events(events): with open(events_filepath, "wb") as handle: pickle.dump(serialize_value(events), handle) # Set up a thread to handle writing events back to the plan process, so execution doesn't get # blocked on remote communication events_queue = Queue() event_writing_thread = Thread( target=event_writing_loop, kwargs=dict(events_queue=events_queue, put_events_fn=put_events), ) event_writing_thread.start() with StringIO() as stderr, StringIO() as stdout, redirect_stderr( stderr), redirect_stdout(stdout): try: instance = external_instance_from_step_run_ref( step_run_ref, event_listener_fn=events_queue.put) # consume iterator list(run_step_from_ref(step_run_ref, instance)) except Exception as e: # ensure that exceptiosn make their way into stdout traceback.print_exc() raise e finally: events_queue.put(DONE) event_writing_thread.join() # write final stdout and stderr with open( os.path.dirname(step_run_ref_filepath) + "/stderr", "wb") as handle: stderr_str = stderr.getvalue() sys.stderr.write(stderr_str) handle.write(stderr_str.encode()) with open( os.path.dirname(step_run_ref_filepath) + "/stdout", "wb") as handle: stdout_str = stdout.getvalue() sys.stdout.write(stdout_str) handle.write(stdout_str.encode())