Machine Learning application design with Pure Functional ZIO in Scala

MLToyApp.scala

package ziomlscoring

import zio._
import ziomlscoring.SizeMyShirt.moduleExperimentFramework.ExperimentFramework
import ziomlscoring.SizeMyShirt.moduleInferenceLogger.InferenceLogger
import ziomlscoring.SizeMyShirt.moduleTShirtSizer.TShirtSizer
import ziomlscoring.SizeMyShirt.sizeMyShirtMicroService

object SizeMyShirt {
  case class UserId(value: Long) extends AnyVal
  case class UserFeatures(id: UserId, heightM: Double, weightKg: Double, age: Int, sex: Boolean, fitness: Double)
  object UserFeatures {
    val testId: UserId = UserId(123L)
    val test           = UserFeatures(testId, 1.88, 82, 40, true, 82)
  }

  /*
  sealed trait Err extends Throwable
  case object NotImplementedErr extends Err
  case class OtherErr(msg: String) extends Err
  case class NetworkErr(msg: String) extends Err
   */

  sealed trait TShirtSize

  sealed trait StandardTShirtSize extends TShirtSize
  case object XS                  extends StandardTShirtSize
  case object S                   extends StandardTShirtSize
  case object M                   extends StandardTShirtSize
  case object L                   extends StandardTShirtSize
  case object XL                  extends StandardTShirtSize
  case object XXL                 extends StandardTShirtSize

  case class TailoredTShirtSize(trunkCm: Int, waistCm: Int, neckCm: Int, chestCm: Int) extends TShirtSize
  object TailoredTShirtSize {
    val avg: TailoredTShirtSize = TailoredTShirtSize(82, 82, 35, 85)
  }
  case class Request(userId: UserId, requestId: Long)

  val getRequest: Task[Request]                             = IO(Request(UserFeatures.testId, 345L))
  def getUserFeatures(request: Request): Task[UserFeatures] = IO(UserFeatures.test)
  def sendResponse(request: Request, stdSize: StandardTShirtSize, tailoredSize: TailoredTShirtSize): Task[Unit] =
    IO.unit

  object moduleTShirtSizer {
    type TShirtSizer = Has[TShirtSizer.Service]

    object TShirtSizer {

      trait Service {
        def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize]
        def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize]
      }

      val manual: ZLayer[Any, Throwable, Has[Service]] = ZLayer.succeed {
        new TShirtSizer.Service {
          def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize] =
            for { // the "model" is a heuristic that comes from author's mind
              bmi: Double <- IO(user.weightKg / (math.pow(user.heightM, 2.0))).catchAll(_ => IO(24.0))
              sz          = if (bmi < 19) S else if (bmi < 25) M else L
            } yield sz
          def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize] = IO(TailoredTShirtSize.avg)
        }
      }

      val manualWithLog: ZLayer[InferenceLogger, Throwable, Has[Service]] = ZLayer.fromService { logger =>
        new TShirtSizer.Service {
          def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize] =
            for {
              bmi: Double <- IO(user.weightKg / (math.pow(user.heightM, 2.0))).catchAll(_ => IO(24.0))
              sz          = if (bmi < 19) S else if (bmi < 25) M else L
              _: Unit     <- logger.log(user, sz) // JUST ONE NEW FUNCTION COMPOSED
            } yield sz
          def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize] =
            for {
              sz      <- IO(TailoredTShirtSize.avg)
              _: Unit <- logger.log(user, sz)
            } yield sz
        }
      }

      def testService(logger: InferenceLogger.Service): Service =
        new TShirtSizer.Service {
          def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize] =
            for {
              s       <- IO(M)
              _: Unit <- logger.log(user, s)
            } yield s
          def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize] =
            for {
              s       <- IO.succeed(TailoredTShirtSize.avg)
              _: Unit <- logger.log(user, s)
            } yield s
        }

      val test: ZLayer[InferenceLogger, Throwable, Has[TShirtSizer.Service]] =
        ZLayer.fromService { logger: InferenceLogger.Service => testService(logger) }

      def getLogRegModelStd(modelPath: String, l: InferenceLogger.Service): UserFeatures => Task[StandardTShirtSize] =
        ???
      def getLogRegModelTailored(
        modelPath: String,
        l: InferenceLogger.Service
      ): UserFeatures => Task[TailoredTShirtSize] = ???
      def logRegPmmlSizer(modelPath: String): ZLayer[InferenceLogger, Throwable, Has[Service]] =
        ZLayer.fromService { logger =>
          new TShirtSizer.Service {
            private val stdModel: UserFeatures => Task[StandardTShirtSize] = getLogRegModelStd(modelPath, logger)
            private val tailoredModel: UserFeatures => Task[TailoredTShirtSize] =
              getLogRegModelTailored(modelPath, logger)
            def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize] = stdModel(user)
            def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize] =
              IO.succeed(TailoredTShirtSize.avg) // just average, LogReg is making bad model
          }
        }

      def getTensorflowModelStd(
        modelPath: String,
        l: InferenceLogger.Service
      ): UserFeatures => Task[StandardTShirtSize] = ???
      def getTensorflowModelTailored(
        modelPath: String,
        l: InferenceLogger.Service
      ): UserFeatures => Task[TailoredTShirtSize] = ???
      def dnnTensorflowSizer(modelPath: String): ZLayer[InferenceLogger, Throwable, Has[Service]] = ZLayer.fromService {
        logger =>
          new Service {
            private val stdModel: UserFeatures => Task[StandardTShirtSize] = getTensorflowModelStd(modelPath, logger)
            private val tailoredModel: UserFeatures => Task[TailoredTShirtSize] =
              getTensorflowModelTailored(modelPath, logger)
            def computeStandardSize(user: UserFeatures): Task[StandardTShirtSize] = stdModel(user)
            def computeTailoredSize(user: UserFeatures): Task[TailoredTShirtSize] = tailoredModel(user)
          }
      }

      val experimentBasedSizer: ZLayer[ExperimentFramework, Throwable, Has[Service]] =
        ZLayer.fromService[ExperimentFramework.Service, Service] { experimentFramework =>
          new Service {
            override def computeStandardSize(user: UserFeatures): IO[Throwable, StandardTShirtSize] =
              for {
                selectedAModel         <- IO(experimentFramework.selectModel(user.id))
                sz: StandardTShirtSize <- selectedAModel.computeStandardSize(user)
              } yield sz
            override def computeTailoredSize(user: UserFeatures): IO[Throwable, TailoredTShirtSize] =
              for {
                selectedAModel         <- IO(experimentFramework.selectModel(user.id))
                sz: TailoredTShirtSize <- selectedAModel.computeTailoredSize(user)
              } yield sz
          }
        }

    }
  }

  type Model = String
  object moduleExperimentFramework {
    type ExperimentFramework = Has[ExperimentFramework.Service]
    object ExperimentFramework {
      trait Service {
        def selectModel(userId: UserId): TShirtSizer.Service
      }
      val live: ZLayer[InferenceLogger, Throwable, Has[Service]] = ZLayer.fromService { logger =>
        new ExperimentFramework.Service {
          val allModels: Map[UserId, Model]                             = ???
          override def selectModel(userId: UserId): TShirtSizer.Service = ???
        }
      }
      val test: ZLayer[InferenceLogger, Throwable, Has[Service]] =
        ZLayer.fromService { logger: InferenceLogger.Service =>
          new ExperimentFramework.Service {
            override def selectModel(userId: UserId): TShirtSizer.Service = TShirtSizer.testService(logger)
          }
        }
    }
  }

  object moduleInferenceLogger {
    type InferenceLogger = Has[InferenceLogger.Service]

    object InferenceLogger {
      trait Service {
        def log(user: UserFeatures, size: TShirtSize): UIO[Unit]
      }
      val live: Layer[Throwable, Has[Service]] = ZLayer.succeed {
        new InferenceLogger.Service {
          override def log(user: UserFeatures, size: TShirtSize): UIO[Unit] = ??? // maybe StackDriver
        }
      }
      val test: Layer[Throwable, Has[Service]] = ZLayer.succeed {
        new InferenceLogger.Service {
          override def log(user: UserFeatures, size: TShirtSize): UIO[Unit] =
            UIO(println(s"LOGGED $user and size ${size}"))
        }
      }
    }

  }

  import moduleTShirtSizer._

  val sizeMyShirtMicroService: ZIO[TShirtSizer, Throwable, Nothing] =
    ZIO.accessM { env =>
      (for {
        request: Request        <- getRequest
        user: UserFeatures      <- getUserFeatures(request)
        std: StandardTShirtSize <- env.get[TShirtSizer.Service].computeStandardSize(user)
        tld: TailoredTShirtSize <- env.get[TShirtSizer.Service].computeTailoredSize(user)
        _: Unit                 <- sendResponse(request, std, tld)
      } yield ()).forever
    }

}

object SizeMyShirtApp extends zio.App {
  val layers: ZLayer[Any, Throwable, TShirtSizer] =
    (InferenceLogger.test >>> ExperimentFramework.test) >>> TShirtSizer.experimentBasedSizer
  override def run(args: List[String]): URIO[zio.ZEnv, ExitCode] =
    sizeMyShirtMicroService.provideLayer(layers).exitCode
}