easy_ml 0.1.4 → 0.2.0.pre.rc1

data/app/models/easy_ml/dataset_history.rb

@@ -0,0 +1,56 @@
+# == Schema Information
+#
+# Table name: easy_ml_dataset_histories
+#
+#  id                      :bigint           not null, primary key
+#  dataset_id              :integer          not null
+#  name                    :string           not null
+#  description             :string
+#  dataset_type            :string
+#  status                  :string
+#  version                 :string
+#  datasource_id           :integer
+#  root_dir                :string
+#  configuration           :json
+#  num_rows                :integer
+#  workflow_status         :string
+#  statistics              :json
+#  preprocessor_statistics :json
+#  schema                  :json
+#  refreshed_at            :datetime
+#  created_at              :datetime         not null
+#  updated_at              :datetime         not null
+#  history_started_at      :datetime         not null
+#  history_ended_at        :datetime
+#  history_user_id         :integer
+#  snapshot_id             :string
+#
+module EasyML
+  class DatasetHistory < ActiveRecord::Base
+    self.table_name = "easy_ml_dataset_histories"
+    include Historiographer::History
+
+    has_many :columns,
+      ->(dataset_history) { where(snapshot_id: dataset_history.snapshot_id) },
+      class_name: "EasyML::ColumnHistory",
+      foreign_key: "dataset_id",
+      primary_key: "dataset_id",
+      extend: EasyML::ColumnList
+
+    def root_dir
+      read_attribute(:root_dir)
+    end
+
+    def fit
+      false
+    end
+
+    def processed?
+      true
+    end
+
+    def should_split?
+      false
+    end
+  end
+end

data/app/models/easy_ml/datasource.rb

@@ -0,0 +1,182 @@
+# == Schema Information
+#
+# Table name: easy_ml_datasources
+#
+#  id              :bigint           not null, primary key
+#  name            :string           not null
+#  datasource_type :string
+#  root_dir        :string
+#  configuration   :json
+#  refreshed_at    :datetime
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#
+module EasyML
+  class Datasource < ActiveRecord::Base
+    self.table_name = "easy_ml_datasources"
+    include Historiographer::Silent
+    historiographer_mode :snapshot_only
+    include EasyML::Concerns::Configurable
+
+    DATASOURCE_OPTIONS = {
+      "s3" => "EasyML::Datasources::S3Datasource",
+      "file" => "EasyML::Datasources::FileDatasource",
+      "polars" => "EasyML::Datasources::PolarsDatasource",
+    }
+    DATASOURCE_TYPES = [
+      {
+        value: "s3",
+        label: "Amazon S3",
+        description: "Connect to data stored in Amazon Simple Storage Service (S3) buckets",
+      },
+      {
+        value: "file",
+        label: "Local Files",
+        description: "Connect to data stored in local files",
+      },
+      {
+        value: "polars",
+        label: "Polars DataFrame",
+        description: "In-memory dataframe storage using Polars",
+      },
+    ].freeze
+    DATASOURCE_NAMES = DATASOURCE_OPTIONS.keys.freeze
+    DATASOURCE_CONSTANTS = DATASOURCE_OPTIONS.values.map(&:constantize)
+
+    validates :name, presence: true
+    validates :datasource_type, presence: true
+    validates :datasource_type, inclusion: { in: DATASOURCE_NAMES }
+    # validate :validate_datasource_exists
+
+    before_save :set_root_dir
+    after_initialize :read_adapter_from_configuration, if: :persisted?
+    after_find :read_adapter_from_configuration
+    before_save :store_adapter_in_configuration
+    after_create :refresh_async
+
+    has_many :events, as: :eventable, class_name: "EasyML::Event", dependent: :destroy
+    attr_accessor :schema, :columns, :num_rows, :is_syncing
+
+    add_configuration_attributes :schema, :columns, :num_rows, :polars_args, :verbose, :is_syncing
+    DATASOURCE_CONSTANTS.flat_map(&:configuration_attributes).each do |attribute|
+      add_configuration_attributes attribute
+    end
+
+    delegate :query, :in_batches, :files, :all_files, :last_updated_at, :data, :needs_refresh?,
+             :should_sync?, :files_to_sync, :s3_access_key_id, :s3_secret_access_key,
+             :download_file, :clean, to: :adapter
+
+    def self.constants
+      {
+        DATASOURCE_TYPES: DATASOURCE_TYPES,
+        s3: EasyML::Datasources::S3Datasource.constants,
+      }
+    end
+
+    def reread(columns = nil)
+      return false unless adapter.respond_to?(:convert_to_parquet)
+
+      adapter.convert_to_parquet(columns)
+    end
+
+    def available_files
+      all_files.select { |f| File.exist?(f) && Pathname.new(f).extname == ".csv" }.map { |f| f.gsub(Regexp.new(Rails.root.to_s), "") }
+    end
+
+    def in_memory?
+      datasource_type == "polars"
+    end
+
+    def root_dir
+      persisted = read_attribute(:root_dir)
+      return persisted if persisted.present? && !persisted.to_s.blank?
+
+      default_root_dir
+    end
+
+    def refresh_async
+      update(is_syncing: true)
+      EasyML::SyncDatasourceJob.perform_later(id)
+    end
+
+    def before_sync
+      update!(is_syncing: true)
+      adapter.before_sync
+      Rails.logger.info("Starting sync for datasource #{id}")
+    end
+
+    def after_sync
+      adapter.after_sync
+      self.schema = data.schema.reduce({}) do |h, (k, v)|
+        h.tap do
+          h[k] = EasyML::Data::PolarsColumn.polars_to_sym(v)
+        end
+      end
+      self.columns = data.columns
+      self.num_rows = data.shape[0]
+      self.is_syncing = false
+      self.refreshed_at = Time.now
+      save
+    end
+
+    def refresh
+      unless adapter.needs_refresh?
+        update!(is_syncing: false)
+        return
+      end
+
+      syncing do
+        adapter.refresh
+      end
+    end
+
+    def refresh!
+      syncing do
+        adapter.refresh!
+      end
+    end
+
+    def syncing
+      before_sync
+      yield.tap do
+        after_sync
+      end
+    end
+
+    private
+
+    def adapter
+      @adapter ||= begin
+          adapter_class = DATASOURCE_OPTIONS[datasource_type]
+          raise "Don't know how to use datasource adapter #{datasource_type}!" unless adapter_class.present?
+
+          adapter_class.constantize.new(self)
+        end
+    end
+
+    def default_root_dir
+      folder = name.gsub(/\s{2,}/, " ").split(" ").join("_").downcase
+      EasyML::Engine.root_dir.join("datasources").join(folder)
+    end
+
+    def set_root_dir
+      write_attribute(:root_dir, default_root_dir) unless read_attribute(:root_dir).present?
+    end
+
+    def read_adapter_from_configuration
+      return unless persisted?
+
+      adapter.read_from_configuration if adapter.respond_to?(:read_from_configuration)
+    end
+
+    def store_adapter_in_configuration
+      adapter.store_in_configuration if adapter.respond_to?(:store_in_configuration)
+    end
+
+    def validate_datasource_exists
+      return if adapter.exists?
+
+      errors.add(:root_dir, adapter.error_not_exists)
+    end
+  end
+end

data/app/models/easy_ml/datasource_history.rb

@@ -0,0 +1,24 @@
+# == Schema Information
+#
+# Table name: easy_ml_datasource_histories
+#
+#  id                 :bigint           not null, primary key
+#  datasource_id      :integer          not null
+#  name               :string           not null
+#  datasource_type    :string
+#  root_dir           :string
+#  configuration      :json
+#  refreshed_at       :datetime
+#  created_at         :datetime         not null
+#  updated_at         :datetime         not null
+#  history_started_at :datetime         not null
+#  history_ended_at   :datetime
+#  history_user_id    :integer
+#  snapshot_id        :string
+#
+module EasyML
+  class DatasourceHistory < ActiveRecord::Base
+    self.table_name = "easy_ml_datasource_histories"
+    include Historiographer::History
+  end
+end

data/app/models/easy_ml/datasources/base_datasource.rb

@@ -0,0 +1,54 @@
+module EasyML
+  module Datasources
+    class BaseDatasource
+      include ActiveModel::Validations
+      include EasyML::Concerns::Configurable
+
+      attr_reader :datasource
+
+      def clean; end
+
+      def before_sync; end
+
+      def after_sync; end
+
+      def initialize(datasource)
+        @datasource = datasource
+      end
+
+      def query(*)
+        raise NotImplementedError
+      end
+
+      def in_batches(*)
+        raise NotImplementedError
+      end
+
+      def files
+        raise NotImplementedError
+      end
+
+      def last_updated_at
+        raise NotImplementedError
+      end
+
+      def data
+        raise NotImplementedError
+      end
+
+      def needs_refresh?
+        false
+      end
+
+      def refresh
+        datasource.syncing do
+          # Default implementation does nothing
+        end
+      end
+
+      def refresh!
+        refresh
+      end
+    end
+  end
+end

data/app/models/easy_ml/datasources/file_datasource.rb

@@ -0,0 +1,58 @@
+module EasyML
+  module Datasources
+    class FileDatasource < BaseDatasource
+      delegate :query, :convert_to_parquet, to: :reader
+
+      def after_sync
+        reader.normalize
+      end
+
+      def in_batches(&block)
+        reader.in_batches(&block)
+      end
+
+      def all_files
+        reader.all_files
+      end
+
+      def files
+        reader.files
+      end
+
+      def last_updated_at
+        files.map { |file| File.mtime(file) }.max
+      end
+
+      def needs_refresh?
+        false
+      end
+
+      def data
+        return @combined_df if @combined_df.present?
+
+        combined_df = nil
+        reader.in_batches do |df|
+          combined_df = combined_df.nil? ? df : combined_df.vstack(df)
+        end
+        @combined_df = combined_df
+      end
+
+      def exists?
+        Dir.glob(File.join(datasource.root_dir, "**/*.{csv,parquet}")).any?
+      end
+
+      def error_not_exists
+        "Expected to find datasource files at #{datasource.root_dir}"
+      end
+
+      private
+
+      def reader
+        @reader ||= EasyML::Data::PolarsReader.new(
+          root_dir: datasource.root_dir,
+          polars_args: (datasource.configuration || {}).dig("polars_args"),
+        )
+      end
+    end
+  end
+end

data/app/models/easy_ml/datasources/polars_datasource.rb

@@ -0,0 +1,89 @@
+module EasyML
+  module Datasources
+    class PolarsDatasource < BaseDatasource
+      validates :df, presence: true
+      add_configuration_attributes :df
+
+      def query(drop_cols: [], filter: nil, limit: nil, select: nil, unique: nil, sort: nil, descending: false)
+        return if df.nil?
+
+        df = self.df.clone
+        df = df.filter(filter) if filter
+        df = df.select(select) if select.present?
+        df = df.unique if unique
+        drop_cols &= df.columns
+        df = df.drop(drop_cols) unless drop_cols.empty?
+        df = df.sort(sort, reverse: descending) if sort
+        df = df.limit(limit) if limit
+        df
+      end
+
+      def in_batches(of: 10_000)
+        total_rows = df.shape[0]
+        (0...total_rows).step(of) do |start|
+          end_index = [start + of, total_rows].min
+          yield df.slice(start, end_index - start)
+        end
+      end
+
+      def all_files
+        []
+      end
+
+      def files
+        []
+      end
+
+      def last_updated_at
+        datasource.updated_at
+      end
+
+      def data
+        df
+      end
+
+      def df
+        datasource.df
+      end
+
+      def exists?
+        df.present?
+      end
+
+      def error_not_exists
+        "Must have a dataframe"
+      end
+
+      def store_in_configuration
+        return unless df
+
+        datasource.configuration = (datasource.configuration || {}).merge(
+          "df" => JSON.parse(df.write_json),
+        )
+      end
+
+      def read_from_configuration
+        return unless datasource.configuration&.key?("df")
+
+        df_data = datasource.configuration["df"]
+        return unless df_data.present? && df_data.key?("columns")
+
+        columns = df_data["columns"].map do |col|
+          dtype = case col["datatype"]
+            when Hash
+              if col["datatype"]["Datetime"]
+                Polars::Datetime.new(col["datatype"]["Datetime"][0].downcase.to_sym).class
+              else
+                Polars::Utf8
+              end
+            else
+              Polars.const_get(col["datatype"])
+            end
+          Polars::Series.new(col["name"], col["values"], dtype: dtype)
+        end
+
+        datasource.df = Polars::DataFrame.new(columns)
+      end
+    end
+  end
+end

data/app/models/easy_ml/datasources/s3_datasource.rb

@@ -0,0 +1,97 @@
+module EasyML
+  module Datasources
+    class S3Datasource < BaseDatasource
+      REGIONS = [
+        { value: "us-east-1", label: "US East (N. Virginia)" },
+        { value: "us-east-2", label: "US East (Ohio)" },
+        { value: "us-west-1", label: "US West (N. California)" },
+        { value: "us-west-2", label: "US West (Oregon)" },
+      ].freeze
+
+      def self.constants
+        { S3_REGIONS: REGIONS }
+      end
+
+      validates :s3_bucket, :s3_access_key_id, :s3_secret_access_key, presence: true
+
+      add_configuration_attributes :s3_bucket, :s3_prefix, :s3_region, :cache_for
+
+      delegate :query, :data, :s3_access_key_id, :s3_secret_access_key, :before_sync, :after_sync, :clean,
+               to: :synced_directory
+
+      def in_batches(&block)
+        synced_directory.in_batches(&block)
+      end
+
+      def all_files
+        synced_directory.all_files
+      end
+
+      def files
+        synced_directory.files
+      end
+
+      def last_updated_at
+        synced_directory.last_updated_at
+      end
+
+      def needs_refresh?
+        synced_directory.should_sync?
+      end
+
+      def refresh
+        synced_directory.sync
+      end
+
+      def refresh!
+        synced_directory.sync!
+      end
+
+      def files_to_sync
+        synced_directory.files_to_sync
+      end
+
+      def download_file(file)
+        synced_directory.download_file(file)
+      end
+
+      def exists?
+        synced_directory.files_to_sync.any?
+      end
+
+      def error_not_exists
+        "No files found at s3://#{File.join(s3_bucket, s3_prefix)}"
+      end
+
+      def s3_bucket
+        datasource_config.dig("s3_bucket") || EasyML::Configuration.s3_bucket
+      end
+
+      def s3_prefix
+        datasource_config.dig("s3_prefix")
+      end
+
+      def cache_for
+        datasource_config.dig("cache_for") || 0
+      end
+
+      private
+
+      def datasource_config
+        @datasource_config ||= datasource.configuration || {}
+      end
+
+      def synced_directory
+        @synced_directory ||= EasyML::Data::SyncedDirectory.new(
+          root_dir: datasource.root_dir,
+          s3_bucket: s3_bucket,
+          s3_prefix: s3_prefix,
+          s3_access_key_id: EasyML::Configuration.s3_access_key_id,
+          s3_secret_access_key: EasyML::Configuration.s3_secret_access_key,
+          polars_args: datasource_config.dig("polars_args") || {},
+          cache_for: cache_for,
+        )
+      end
+    end
+  end
+end

data/app/models/easy_ml/deploy.rb

@@ -0,0 +1,114 @@
+# == Schema Information
+#
+# Table name: easy_ml_deploys
+#
+#  id                :bigint           not null, primary key
+#  model_id          :bigint
+#  model_history_id  :bigint
+#  retraining_run_id :bigint
+#  model_file_id     :bigint
+#  status            :string           not null
+#  trigger           :string           default("manual")
+#  stacktrace        :text
+#  snapshot_id       :string
+#  created_at        :datetime         not null
+#  updated_at        :datetime         not null
+#
+module EasyML
+  class Deploy < ActiveRecord::Base
+    self.table_name = "easy_ml_deploys"
+
+    belongs_to :model, class_name: "EasyML::Model"
+    belongs_to :model_file, class_name: "EasyML::ModelFile", optional: true
+    belongs_to :retraining_run, class_name: "EasyML::RetrainingRun"
+    belongs_to :model_version, class_name: "EasyML::ModelHistory", optional: true, foreign_key: :model_history_id
+
+    validates :status, presence: true
+    after_initialize :set_defaults, if: :new_record?
+    before_save :set_model_file, if: :new_record?
+    validates :status, presence: true, inclusion: { in: %w[pending running success failed] }
+
+    scope :latest, -> { select("DISTINCT ON (model_id) *").order("model_id, id DESC") }
+
+    def unlocked?
+      EasyML::Deploy.where(model_id: model_id).where.not(locked_at: nil).where(status: ["pending", "running"]).empty?
+    end
+
+    def locked?
+      !unlocked?
+    end
+
+    def deploy(async: true)
+      if async
+        EasyML::DeployJob.perform_later(id)
+      else
+        actually_deploy
+      end
+    end
+
+    def actually_deploy
+      lock_deploy do
+        update(status: "running")
+        EasyML::Event.create_event(self, "started")
+
+        if identical_deploy.present?
+          self.model_file = identical_deploy.model_file
+          self.model_version = identical_deploy.model_version
+        else
+          if model_file.present?
+            model.model_file = model_file
+          end
+          model.load_model
+          self.model_version = model.actually_deploy
+        end
+
+        EasyML::Deploy.transaction do
+          update(model_history_id: self.model_version.id, snapshot_id: self.model_version.snapshot_id, status: :success)
+          model.retraining_runs.where(status: :deployed).update_all(status: :success)
+          retraining_run.update(model_history_id: self.model_version.id, snapshot_id: self.model_version.snapshot_id, deploy_id: id, status: :deployed, is_deploying: false)
+        end
+
+        model_version.tap do
+          EasyML::Event.create_event(self, "success")
+        end
+      end
+    end
+
+    alias_method :rollback, :deploy
+
+    def unlock!
+      Support::Lockable.unlock!(lock_key)
+    end
+
+    def lock_deploy
+      with_lock do |client|
+        yield
+      end
+    end
+
+    def identical_deploy
+      EasyML::Deploy.where(retraining_run_id: retraining_run_id).
+        where.not(id: id).where(status: :success).limit(1).first
+    end
+
+    private
+
+    def with_lock
+      EasyML::Support::Lockable.with_lock(lock_key, stale_timeout: 60, resources: 1) do |client|
+        yield client
+      end
+    end
+
+    def lock_key
+      "deploy:#{self.model.name}:#{self.model.id}"
+    end
+
+    def set_defaults
+      self.status ||= :pending
+    end
+
+    def set_model_file
+      self.model_file ||= retraining_run.model_file
+    end
+  end
+end