dcp-client/libexec/sandbox/pyodide-worktime.js

Core libraries for accessing DCP network

/**
 *  @file pyodide-worktime.js
 *  This file sets up the hooks for the Pyodide worktime. It registers the processSlice
 *  initSandbox hooks, which run slices and setup the sandbox respectively.
 *
 *  @author Wes Garland <wes@distributive.network>
 *  @author Severn Lortie <severn@distributive.network>
 *  @author Will Pringle <will@distributive.network>
 *  @author Hamada Gasmallah <hamada@distributive.network>
 *
 *  @date February, 2024
 *  @copyright Copyright (c) 2018-2024, Distributive, Ltd.  All Rights Reserved.
 */
'use-strict';
/* global self, bravojs, addEventListener, postMessage */
// @ts-nocheck

self.wrapScriptLoading({ scriptName: 'pyodide-worktime'}, function pyodideWorktime$$fn(protectedStorage, ring1PostMessage, wrapPostMessage)
{
  let wrappedPythonSliceHandler;
  /**
   * Registers the worktime callbacks that allow the worktime controller (worktimes.js) to setup our environment and run the worktime
   */
  function registerWorktime()
  {
    protectedStorage.worktimes.registerWorktime({
      name: 'pyodide',
      version: '0.28.0',
      processSlice,
      initSandbox,
    });
  }

  /**
   * Function which generates a "map-basic"-like workFunction
   * out of a Pyodide Worktime job (Python code, files, env variables).
   *
   * It takes any "images" passed in the workFunction "arguments" and
   * writes them to the in memory filesystem provided by Emscripten.
   * It adds any environment variables specified in the workFunction
   * "arguments" to the pseudo-"process" for use.
   * It globally imports a dcp module with function "set_slice_handler"
   * which takes a python function as input. The python function passed
   * to that slice handler is invoked by the function which this
   * function creates.
   *
   * @param {Object} job The job being assigned to the sandbox
   */
  async function initSandbox(job)
  {
    var pythonSliceHandler;

    const pyodide = await pyodideInit();
    const sys = pyodide.pyimport('sys');

    const findImports = pyodide.runPython('import pyodide; pyodide.code.find_imports');
    const findPythonModuleLoader = pyodide.runPython('import importlib; importlib.util.find_spec');

    const parsedArguments = parsePyodideArguments(job.arguments);

    // write images to file and set environment variables
    const prepPyodide = pyodide.runPython(`
import tarfile, io
import os, sys

def prepPyodide(args):
  for image in args['images']:
    image = bytes(image)
    imageFile = io.BytesIO(image)
    tar = tarfile.open(mode='r', fileobj=imageFile)

    # Don't overwrite directories which corrupts Pyodide's in memory filesystem
    def safe_extract(tar, path="/"):
      for member in tar.getmembers():
          if member.isdir():
              dir_path = os.path.join(path, member.name)
              if not os.path.exists(dir_path):
                  os.makedirs(dir_path)
          else:
              tar.extract(member, path)
    safe_extract(tar)

  for item, value in args['environmentVariables'].items():
    os.environ[item] = value

  sys.argv.extend(args['sysArgv'])

  return

prepPyodide`);

    prepPyodide(pyodide.toPy(parsedArguments));

    // register the dcp Python module 
    if (!sys.modules.get('dcp'))
    {
      const create_proxy = pyodide.runPython('import pyodide;pyodide.ffi.create_proxy');

      pyodide.registerJsModule('dcp', {
        set_slice_handler: function pyodide$$dcp$$setSliceHandler(func) {
          pythonSliceHandler = create_proxy(func);
        },
        progress,
      });
    }
    pyodide.runPython( 'import dcp' );

    // attempt to import packages from the package manager (if they're not already loaded)
    const workFunctionPythonImports = findImports(job.workFunction).toJs();
    const packageManagerImports = workFunctionPythonImports.filter(x=>!findPythonModuleLoader(x));
    if (packageManagerImports.length > 0)
    {
      await fetchAndInitPyodidePackages(packageManagerImports);
      await pyodide.loadPackage(packageManagerImports);
    }

    wrappedPythonSliceHandler = workFunctionWrapper;

    /**
     * Evaluates the Python WorkFunction string and then executes the slice
     * handler Python function. If no call to `dcp.set_slice_handler` is passed
     * or a non function is passed to it.
     * This function specifically only takes one parameter since Pyodide Slice
     * Handlers only accept one parameter.
     */
    async function workFunctionWrapper(datum)
    {
      const pyodide = await pyodideInit(); // returns the same promise when called multiple times

      // load and execute the Python Workfunction, this populates the pythonSliceHandler variable
      await pyodide.runPython(job.workFunction);

      // failure to specify a slice handler is considered an error
      if (!pythonSliceHandler)
        throw new Error('ENOSLICEHANDLER: Must specify the slice handler using `dcp.set_slice_handler(fn)`');

      // setting the slice handler to a non function or lambda is not supported
      else if (typeof pythonSliceHandler !== 'function')
        throw new Error('ENOSLICEHANDLER: Slice Handler must be a function');

      const sliceHandlerResult = await pythonSliceHandler(pyodide.toPy(datum));

      // if it is a PyProxy, convert its value to JavaScript
      if (sliceHandlerResult?.toJs)
        return sliceHandlerResult.toJs();

      return sliceHandlerResult;
    }

    /*
     * Refer to the "The Pyodide Worktime"."Work Function (JS)"."Arguments"."Commands"
     * part of the DCP Worktimes spec.
     */
    function parsePyodideArguments(args)
    {
      var index = 1;
      const numArgs = args[0];
      const images = [];
      const environmentVariables = {};
      const sysArgv = args.slice(numArgs);

      while (index < numArgs)
      {
        switch (args[index])
        {
          case 'gzImage':
            const image = args[index+1];
            images.push(image);
            index+=2;
            break;
          case 'env':
            const env = args[index+1].split(/=(.*)/s);
            index+=2;
            environmentVariables[env[0]] = env[1];
            break;
          default:
            throw new Error(`Invalid argument ${args[index]}`);
        }
      }

      return { sysArgv, images, environmentVariables };
    }
  }

  /**
   * The processSlice hook function which runs the work function for input slice data
   * @param {*} data The slice data
   * @returns {*} The slice result
   */
  async function processSlice(data)
  {
    const result = await wrappedPythonSliceHandler(data);
    return result;
  }

  registerWorktime();
}); /* end of fn */