<commit_before>﻿namespace MBrace.Azure.Runtime.Primitives

open System
open System.IO
open System.Text.RegularExpressions
open System.Runtime.Serialization

open Nessos.Vagabond
open Nessos.Vagabond.AssemblyProtocols

open MBrace.Runtime
open MBrace.Runtime.Vagabond
open MBrace.Azure
open MBrace.Azure.Runtime
open MBrace.Azure.Runtime.Utilities

#nowarn "1571"

[<AutoOpen>]
module private Common =

    /// blob prefix for stored assemblies
    let prefix = "vagabond"

    /// Gets a unique blob filename for provided assembly
    let filename (id : AssemblyId) = Vagabond.GetFileName id

    let assemblyName id = filename id + id.Extension
    let symbolsName id = filename id + ".pdb"
    let metadataName id = filename id + ".vgb"
    let dataFile id (dd : DataDependencyInfo) = sprintf "%s-%d-%d.dat" (filename id) dd.Id dd.Generation

/// Assembly to blob store uploader implementation
type private BlobAssemblyUploader(config : ConfigurationId, logger : ICloudLogger) =

    let sizeOfFile (path:string) = FileInfo(path).Length |> getHumanReadableByteSize

    let getAssemblyLoadInfo (id : AssemblyId) = async {
        let! assemblyExists = Blob.Exists(config, prefix, assemblyName id)
        if not assemblyExists then return NotLoaded id
        else
            let cell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName id)
            let! metadata = cell.GetValue()
            return Loaded(id, false, metadata)
    }

    /// upload assembly to blob store
    let uploadAssembly (va : VagabondAssembly) = async {
        let assemblyName = assemblyName va.Id
        let! assemblyExists = Blob.Exists(config, prefix, assemblyName)

        // 1. Upload assembly image.
        if not assemblyExists then
            /// print upload sizes for given assembly
            let uploadSizes = 
                seq {
                    yield sprintf "DLL %s" (sizeOfFile va.Image)
                    match va.Symbols with
                    | Some s -> yield sprintf "PDB %s" (sizeOfFile s)
                    | None -> ()
                } |> String.concat ", "

            logger.Logf "Uploading '%s' [%s]" va.FullName uploadSizes
            let! _ = Blob.UploadFromFile(config, prefix, assemblyName, va.Image)
            return ()

        // 2. Upload symbols if applicable.
        match va.Symbols with
        | None -> ()
        | Some symbolsPath ->
            let symbolsName = symbolsName va.Id
            let! symbolsExist = Blob.Exists(config, prefix, symbolsName)
            if not symbolsExist then
                let! _ = Blob.UploadFromFile(config, prefix, symbolsName, symbolsPath)
                return ()

        // 3. Upload metadata
        // check current metadata in store
        let metadataName = metadataName va.Id
        let metadataCell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName)
        let! currentMetadata = metadataCell.TryGetValue()

        // detect if metadata in blob store is stale
        let isRequiredUpdate =
            match currentMetadata with
            | None -> true
            | Some md ->
                // require a data dependency whose store generation is older than local
                (md.DataDependencies, va.Metadata.DataDependencies)
                ||> Array.exists2 (fun store local -> local.Generation > store.Generation)

        if not isRequiredUpdate then return Loaded(va.Id, false, va.Metadata) else

        // upload data dependencies
        let files = va.PersistedDataDependencies |> Map.ofArray
        let dataFiles = 
            va.Metadata.DataDependencies 
            |> Seq.filter (fun dd -> match dd.Data with Persisted _ -> true | _ -> false)
            |> Seq.map (fun dd -> dd, files.[dd.Id])
            |> Seq.toArray

        let uploadDataFile (dd : DataDependencyInfo, localPath : string) = async {
            let blobPath = dataFile va.Id dd
            let! dataExists = Blob.Exists(config, prefix, blobPath)
            if not dataExists then
                logger.Logf "Uploading data dependency '%s' [%s]" dd.Name (sizeOfFile localPath)
                let! _ = Blob.UploadFromFile(config, prefix, blobPath, localPath)
                ()
        }

        do! dataFiles |> Seq.map uploadDataFile |> Async.Parallel |> Async.Ignore

        // upload metadata record; TODO: use table store?
        let! _ = Blob<VagabondMetadata>.Create(config, prefix, metadataName, fun () -> va.Metadata)
        return Loaded(va.Id, false, va.Metadata)
    }

    interface IRemoteAssemblyReceiver with
        member x.GetLoadedAssemblyInfo(dependencies: AssemblyId list): Async<AssemblyLoadInfo list> = async {
            let! loadInfo = dependencies |> Seq.map getAssemblyLoadInfo |> Async.Parallel
            return Array.toList loadInfo
        }
        
        member x.PushAssemblies(assemblies: VagabondAssembly list): Async<AssemblyLoadInfo list> =  async {
            let! loadInfo = assemblies |> Seq.map uploadAssembly |> Async.Parallel
            return Array.toList loadInfo
        }


/// Blob store assembly downloader implementation
type private BlobAssemblyDownloader(config : ConfigurationId, logger : ICloudLogger) =
    
    interface IAssemblyImporter with
        member x.GetImageReader(id: AssemblyId): Async<Stream> = async {
            logger.Logf "Downloading '%s'" id.FullName
            let blob = Blob.FromPath(config, prefix, assemblyName id)
            return! blob.OpenRead()
        }
        
        member x.TryGetSymbolReader(id: AssemblyId): Async<Stream option> = async {
            let! exists = Blob.Exists(config, prefix, symbolsName id)
            if exists then
                let blob = Blob.FromPath(config, prefix, symbolsName id)
                let! stream = blob.OpenRead()
                return Some stream
            else
                return None
        }
        
        member x.ReadMetadata(id: AssemblyId): Async<VagabondMetadata> = async {
            let cell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName id)
            return! cell.GetValue()
        }

        member x.GetPersistedDataDependencyReader(id: AssemblyId, dd : DataDependencyInfo): Async<Stream> = async {
            logger.Logf "Downloading data dependency '%s'." dd.Name
            let blob = Blob.FromPath(config, prefix, dataFile id dd)
            return! blob.OpenRead()
        }

/// Assembly manager instance
type BlobAssemblyManager private (config : ConfigurationId, logger : ICloudLogger, includeUnmanagedDependencies : bool) = 
    let uploader = new BlobAssemblyUploader(config, logger)
    let downloader = new BlobAssemblyDownloader(config, logger)

    /// Upload provided dependencies to store
    member __.UploadDependencies(ids : seq<AssemblyId>) = async { 
        logger.Logf "Uploading dependencies"
        let! errors = VagabondRegistry.Instance.SubmitDependencies(uploader, ids)
        if errors.Length > 0 then
            let errors = errors |> Seq.map (fun dd -> dd.Name) |> String.concat ", "
            logger.Logf "Failed to upload bindings: %s" errors
    }

    /// Download provided dependencies from store
    member __.DownloadDependencies(ids : seq<AssemblyId>) = async {
        return! VagabondRegistry.Instance.ImportAssemblies(downloader, ids)
    }

    /// Load local assemblies to current AppDomain
    member __.LoadAssemblies(assemblies : VagabondAssembly list) =
        VagabondRegistry.Instance.LoadVagabondAssemblies(assemblies)

    /// Compute dependencies for provided object graph
    member __.ComputeDependencies(graph : 'T) : AssemblyId list =
        let managedDependencies = VagabondRegistry.Instance.ComputeObjectDependencies(graph, permitCompilation = true) 
        [
            yield! managedDependencies |> VagabondRegistry.Instance.GetVagabondAssemblies
            if includeUnmanagedDependencies then 
                yield! VagabondRegistry.Instance.NativeDependencies
        ] |> List.map (fun va -> va.Id)

    /// Creates a new AssemblyManager instance for given store configuration
    static member Create(config : ConfigurationId, logger : ICloudLogger, ?includeUnmanagedAssemblies : bool) : BlobAssemblyManager = 
        new BlobAssemblyManager(config, logger, defaultArg includeUnmanagedAssemblies true)

    /// <summary>
    ///     Registers a native assembly dependency to client state.
    /// </summary>
    /// <param name="assemblyPath">Path to native assembly.</param>
    static member RegisterNativeDependency(assemblyPath : string) : VagabondAssembly =
        VagabondRegistry.Instance.RegisterNativeDependency assemblyPath

    /// Gets all native dependencies registered in current instance
    static member NativeDependencies =
        VagabondRegistry.Instance.NativeDependencies |> List.map (fun m -> m.Image)<commit_msg>minor improvements to assemly manager<commit_after>﻿namespace MBrace.Azure.Runtime.Primitives

open System
open System.IO
open System.Text.RegularExpressions
open System.Runtime.Serialization

open Nessos.Vagabond
open Nessos.Vagabond.AssemblyProtocols

open MBrace.Runtime
open MBrace.Runtime.Vagabond
open MBrace.Azure
open MBrace.Azure.Runtime
open MBrace.Azure.Runtime.Utilities

#nowarn "1571"

[<AutoOpen>]
module private Common =

    /// blob prefix for stored assemblies
    let prefix = "vagabond"

    /// Gets a unique blob filename for provided assembly
    let filename (id : AssemblyId) = Vagabond.GetFileName id

    let assemblyName id = filename id + id.Extension
    let symbolsName id = filename id + ".pdb"
    let metadataName id = filename id + ".vgb"
    let dataFile id (dd : DataDependencyInfo) = sprintf "%s-%d-%d.dat" (filename id) dd.Id dd.Generation

    type VagabondMetadata with
        member vm.ContainsDataDependencies = vm.DataDependencies.Length > 0

/// Assembly to blob store uploader implementation
type private BlobAssemblyUploader(config : ConfigurationId, logger : ICloudLogger) =

    let sizeOfFile (path:string) = FileInfo(path).Length |> getHumanReadableByteSize

    let getAssemblyLoadInfo (id : AssemblyId) = async {
        let! assemblyExists = Blob.Exists(config, prefix, assemblyName id)
        if not assemblyExists then return NotLoaded id
        else
            let cell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName id)
            let! metadata = cell.GetValue()
            return Loaded(id, false, metadata)
    }

    /// upload assembly to blob store
    let uploadAssembly (va : VagabondAssembly) = async {
        let assemblyName = assemblyName va.Id
        let! assemblyExists = Blob.Exists(config, prefix, assemblyName)

        // 1. Upload assembly image.
        if not assemblyExists then
            /// print upload sizes for given assembly
            let uploadSizes = 
                seq {
                    yield sprintf "IMG %s" (sizeOfFile va.Image)
                    match va.Symbols with
                    | Some s -> yield sprintf "PDB %s" (sizeOfFile s)
                    | None -> ()
                } |> String.concat ", "

            logger.Logf "Uploading '%s' [%s]" va.FullName uploadSizes
            let! _ = Blob.UploadFromFile(config, prefix, assemblyName, va.Image)
            return ()

            // 2. Upload symbols if applicable.
            match va.Symbols with
            | None -> ()
            | Some symbolsPath ->
                let symbolsName = symbolsName va.Id
                let! symbolsExist = Blob.Exists(config, prefix, symbolsName)
                if not symbolsExist then
                    let! _ = Blob.UploadFromFile(config, prefix, symbolsName, symbolsPath)
                    return ()

        // 3. Upload metadata
        // check current metadata in store
        let metadataName = metadataName va.Id
        let metadataCell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName)
        let! currentMetadata = metadataCell.TryGetValue()

        // detect if metadata in blob store is stale
        let isRequiredUpdate =
            match currentMetadata with
            | None -> true
            | Some md ->
                // require a data dependency whose store generation is older than local
                (md.DataDependencies, va.Metadata.DataDependencies)
                ||> Array.exists2 (fun store local -> local.Generation > store.Generation)

        if not isRequiredUpdate then return Loaded(va.Id, false, va.Metadata) else

        // upload data dependencies
        let files = va.PersistedDataDependencies |> Map.ofArray
        let dataFiles = 
            va.Metadata.DataDependencies 
            |> Seq.filter (fun dd -> match dd.Data with Persisted _ -> true | _ -> false)
            |> Seq.map (fun dd -> dd, files.[dd.Id])
            |> Seq.toArray

        let uploadDataFile (dd : DataDependencyInfo, localPath : string) = async {
            let blobPath = dataFile va.Id dd
            let! dataExists = Blob.Exists(config, prefix, blobPath)
            if not dataExists then
                logger.Logf "Uploading data dependency '%s' [%s]" dd.Name (sizeOfFile localPath)
                let! _ = Blob.UploadFromFile(config, prefix, blobPath, localPath)
                ()
        }

        // only print metadata message if updating data dependencies
        if assemblyExists then logger.Logf "Updating metadata for '%s'" va.FullName

        do! dataFiles |> Seq.map uploadDataFile |> Async.Parallel |> Async.Ignore

        // upload metadata record; TODO: use table store?
        let! _ = Blob<VagabondMetadata>.Create(config, prefix, metadataName, fun () -> va.Metadata)
        return Loaded(va.Id, false, va.Metadata)
    }

    interface IRemoteAssemblyReceiver with
        member x.GetLoadedAssemblyInfo(dependencies: AssemblyId list): Async<AssemblyLoadInfo list> = async {
            let! loadInfo = dependencies |> Seq.map getAssemblyLoadInfo |> Async.Parallel
            return Array.toList loadInfo
        }
        
        member x.PushAssemblies(assemblies: VagabondAssembly list): Async<AssemblyLoadInfo list> =  async {
            let! loadInfo = assemblies |> Seq.map uploadAssembly |> Async.Parallel
            return Array.toList loadInfo
        }


/// Blob store assembly downloader implementation
type private BlobAssemblyDownloader(config : ConfigurationId, logger : ICloudLogger) =
    
    interface IAssemblyImporter with
        member x.GetImageReader(id: AssemblyId): Async<Stream> = async {
            logger.Logf "Downloading '%s'" id.FullName
            let blob = Blob.FromPath(config, prefix, assemblyName id)
            return! blob.OpenRead()
        }
        
        member x.TryGetSymbolReader(id: AssemblyId): Async<Stream option> = async {
            let! exists = Blob.Exists(config, prefix, symbolsName id)
            if exists then
                let blob = Blob.FromPath(config, prefix, symbolsName id)
                let! stream = blob.OpenRead()
                return Some stream
            else
                return None
        }
        
        member x.ReadMetadata(id: AssemblyId): Async<VagabondMetadata> = async {
            let cell = Blob<VagabondMetadata>.FromPath(config, prefix, metadataName id)
            return! cell.GetValue()
        }

        member x.GetPersistedDataDependencyReader(id: AssemblyId, dd : DataDependencyInfo): Async<Stream> = async {
            logger.Logf "Downloading data dependency '%s'." dd.Name
            let blob = Blob.FromPath(config, prefix, dataFile id dd)
            return! blob.OpenRead()
        }

/// Assembly manager instance
type BlobAssemblyManager private (config : ConfigurationId, logger : ICloudLogger, includeUnmanagedDependencies : bool) = 
    let uploader = new BlobAssemblyUploader(config, logger)
    let downloader = new BlobAssemblyDownloader(config, logger)

    /// Upload provided dependencies to store
    member __.UploadDependencies(ids : seq<AssemblyId>) = async { 
        logger.Logf "Uploading dependencies"
        let! errors = VagabondRegistry.Instance.SubmitDependencies(uploader, ids)
        if errors.Length > 0 then
            let errors = errors |> Seq.map (fun dd -> dd.Name) |> String.concat ", "
            logger.Logf "Failed to upload bindings: %s" errors
    }

    /// Download provided dependencies from store
    member __.DownloadDependencies(ids : seq<AssemblyId>) = async {
        return! VagabondRegistry.Instance.ImportAssemblies(downloader, ids)
    }

    /// Load local assemblies to current AppDomain
    member __.LoadAssemblies(assemblies : VagabondAssembly list) =
        VagabondRegistry.Instance.LoadVagabondAssemblies(assemblies)

    /// Compute dependencies for provided object graph
    member __.ComputeDependencies(graph : 'T) : AssemblyId list =
        let managedDependencies = VagabondRegistry.Instance.ComputeObjectDependencies(graph, permitCompilation = true) 
        [
            yield! managedDependencies |> VagabondRegistry.Instance.GetVagabondAssemblies
            if includeUnmanagedDependencies then 
                yield! VagabondRegistry.Instance.NativeDependencies
        ] |> List.map (fun va -> va.Id)

    /// Creates a new AssemblyManager instance for given store configuration
    static member Create(config : ConfigurationId, logger : ICloudLogger, ?includeUnmanagedAssemblies : bool) : BlobAssemblyManager = 
        new BlobAssemblyManager(config, logger, defaultArg includeUnmanagedAssemblies true)

    /// <summary>
    ///     Registers a native assembly dependency to client state.
    /// </summary>
    /// <param name="assemblyPath">Path to native assembly.</param>
    static member RegisterNativeDependency(assemblyPath : string) : VagabondAssembly =
        VagabondRegistry.Instance.RegisterNativeDependency assemblyPath

    /// Gets all native dependencies registered in current instance
    static member NativeDependencies =
        VagabondRegistry.Instance.NativeDependencies |> List.map (fun m -> m.Image)
{'fraction_non_alphanumeric': 0.08337033747779751, 'fraction_numerical': 0.0010546181172291296, 'mean_word_length': 2.4815458937198067, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 110, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5710068', 'n_tokens_mistral': 4821, 'n_tokens_neox': 4619, 'n_words': 1379}