simple_databricks_asset

simple_databricks_asset.py

# hooli-ml/databricks_mlops/dagster_pipeline.py
import dagster as dg
from hooli_ml.defs.resources import get_env, MLWorkflowConfig, DatabricksResource
from databricks.sdk.service import jobs

NOTEBOOK_ROOT_PATH = "/Users/[email protected]/.bundle/databricks_mlops/dev/files/"

@dg.asset(kinds={"databricks", "feature_engineering"})
def feature_engineering_pickup(context: dg.AssetExecutionContext, databricks_resource: DatabricksResource):
    run_id = databricks_resource.run_and_stream_notebook_logs_sdk(
        notebook_path=NOTEBOOK_ROOT_PATH + "feature_engineering/notebooks/GenerateAndWriteFeatures",
        parameters={"model_type": config.model_type,
                    "training_data_path": config.input_table_path,
                    "model_name": config.model_name,
                    "pickup_features_table": config.pickup_features_table,
                    "dropoff_features_table": config.dropoff_features_table,
        },
    )
    context.log.info(f"Feature engineering notebook submitted. Run ID: {run_id}")