SUBFLOW.yml

id: model-training
namespace: ml.pipeline
description: "Model Training Flow"

inputs:
  - id: reference_data
    type: FILE
  - id: GITHUB_USERNAME
    type: STRING
    defaults: "{{ secret:GITHUB_USERNAME }}"
  - id: GITHUB_TOKEN
    type: STRING
    defaults: "{{ secret:GITHUB_TOKEN }}"


  
tasks:
  - id: train_model
    type: io.kestra.plugin.scripts.python.Script
    containerImage: ghcr.io/kestra-io/pydata:latest
    namespaceFiles:
      enabled: true
      include:
        - train_model.py
    beforeCommands:
      - pip install -r requirements.txt
    inputFiles:
      requirements.txt: |
        pandas
        numpy
        scikit-learn
        joblib
        requests
      reference_data.csv: "{{ inputs.reference_data }}"
    outputFiles:
      - "model.joblib"
      - "metrics.json"
    script: |
      import pandas as pd
      import numpy as np
      from sklearn.ensemble import RandomForestRegressor
      from sklearn.model_selection import train_test_split
      from sklearn.metrics import mean_squared_error, r2_score
      import joblib
      import json
      import os


      model_params = {
          "test_size": 0.2,
          "random_state": 42,
          "n_estimators": 100
      }
   
      data = pd.read_csv("reference_data.csv")

      data['last_update'] = pd.to_datetime(data['last_update'])

      # Drop rows with NA values
      data = data.dropna(subset=['pollutant_min', 'pollutant_max', 'pollutant_avg'])
      print("Data shape after dropping NaN:", data.shape)

      # Feature engineering
      data['hour'] = data['last_update'].dt.hour
      data['day'] = data['last_update'].dt.day
      data['month'] = data['last_update'].dt.month


      feature_columns = [
          'latitude', 'longitude', 'pollutant_min', 'pollutant_max',
          'hour', 'day', 'month'
      ]

      X = data[feature_columns]
      y = data['pollutant_avg']

      X_train, X_test, y_train, y_test = train_test_split(
          X, y,
          test_size=model_params['test_size'],
          random_state=model_params['random_state']
      )

      model = RandomForestRegressor(
          n_estimators=model_params['n_estimators'],
          random_state=model_params['random_state']
      )
      model.fit(X_train, y_train)


      y_pred = model.predict(X_test)
      mse = mean_squared_error(y_test, y_pred)
      r2 = r2_score(y_test, y_pred)

      metrics = {
          "mse": float(mse),
          "r2_score": float(r2),
          "feature_importance": dict(zip(feature_columns, model.feature_importances_.tolist())),
          "data_stats": {
              "total_samples": len(data),
              "training_samples": len(X_train),
              "test_samples": len(X_test)
          }
      }


      joblib.dump(model, 'model.joblib')
      with open('metrics.json', 'w') as f:
          json.dump(metrics, f, indent=4)

  - id: clone_repo
    type: io.kestra.plugin.git.Clone
    url: https://github.com/anuj846k/kestra_demo
    branch: main
    username: "{{ inputs.GITHUB_USERNAME }}"
    password: "{{ inputs.GITHUB_TOKEN }}"


  - id: push_to_git
    type: io.kestra.plugin.git.Push
    username: "{{ inputs.GITHUB_USERNAME }}"
    password: "{{ inputs.GITHUB_TOKEN }}"
    url: https://github.com/anuj846k/kestra_demo
    branch: main
    commitMessage: "Add trained model and metrics"
    inputFiles:
      model.joblib: "{{ outputs.train_model.outputFiles['model.joblib'] }}"
      metrics.json: "{{ outputs.train_model.outputFiles['metrics.json'] }}"
    addFilesPattern:
      - model.joblib
      - metrics.json