Source code for gyoza.worker.execution_pipeline.steps.prepare_inputs

import json
import os

from gyoza.worker import shared_runtime
from gyoza.worker.execution_pipeline.types import ExecutionContext

GYOZA_WORKING_DIRECTORY = os.getenv("GYOZA_WORKING_DIRECTORY", "/tmp")

_INPUT_FILENAME = "input.json"
_OUTPUT_FILENAME = "output.json"




class PrepareInputs:
    """Prepare the inputs for the container execution.

    Creates a unique gyoza directory for the run, saves inputs as a JSON file,
    and sets GYOZA_INPUT_PATH and GYOZA_OUTPUT_PATH environment variables.

    When the run carries an outputs template (fields pre-filled with upload
    URLs by the caller), it is pre-written to the output file so the op
    runtime inside the container can deliver those outputs.
    """

    def __call__(self, ctx: ExecutionContext) -> ExecutionContext:
        """Execute the step.

        Parameters
        ----------
        ctx : ExecutionContext
            Execution context with inputs to convert.

        Returns
        -------
        ExecutionContext
            Context with gyoza_dir, volumes, and env_vars configured.
        """
        gyoza_dir = self._generate_op_dir(ctx.oprun_id)
        self._save_inputs(ctx.inputs, gyoza_dir)
        self._save_output_template(ctx.outputs, gyoza_dir)
        self._setup_volumes(ctx, gyoza_dir)
        self._setup_env_vars(ctx, gyoza_dir)

        # Store gyoza_dir in extras for cleanup later
        ctx.extras["gyoza_dir"] = gyoza_dir

        return ctx

    def _generate_op_dir(self, oprun_id: str) -> str:
        """Generate the operation directory path using the oprun_id.

        Parameters
        ----------
        oprun_id : str
            OpRun identifier used as directory name.

        Returns
        -------
        str
            Path to the gyoza directory.
        """
        return os.path.join(GYOZA_WORKING_DIRECTORY, oprun_id)

    def _save_inputs(self, inputs: dict, gyoza_dir: str) -> None:
        """Save inputs as a JSON file in the gyoza directory.

        Parameters
        ----------
        inputs : dict
            Input parameters to save.
        gyoza_dir : str
            Directory where the JSON file will be saved.
        """
        os.makedirs(gyoza_dir, exist_ok=True)
        os.chmod(gyoza_dir, 0o777)
        self._write_json(inputs, os.path.join(gyoza_dir, _INPUT_FILENAME))

    def _save_output_template(self, outputs: dict, gyoza_dir: str) -> None:
        """Pre-write the outputs template as the output JSON file.

        Skipped when the template is empty, so runs without pre-filled
        outputs behave exactly as before (no output file until the
        container writes one).

        The file is made writable by anybody, for the same reason the
        directory is: the worker writes it and the *container* rewrites it at
        the end of the run, and the two are not the same user — op images run
        as an unprivileged user while the worker is usually root. Without this
        the container fails with ``Permission denied`` on the very last step,
        after the op has finished and its outputs are already uploaded, and the
        run is marked FAILED with everything about it successful.

        Parameters
        ----------
        outputs : dict
            Partial outputs template from the run.
        gyoza_dir : str
            Directory where the JSON file will be saved.
        """
        if not outputs:
            return
        output_path = os.path.join(gyoza_dir, _OUTPUT_FILENAME)
        self._write_json(outputs, output_path)
        # Same reason the directory is 0o777: the container runs under a
        # different uid. Without this the file keeps the worker's umask and
        # ownership, so the container can create files here but gets [Errno 13]
        # writing its results over this one.
        os.chmod(output_path, 0o666)

    @staticmethod
    def _write_json(data: dict, json_path: str) -> None:
        """Write *data* as JSON to *json_path*.

        Parameters
        ----------
        data : dict
            Data to serialise.
        json_path : str
            Destination file path.
        """
        with open(json_path, "w") as f:
            json.dump(data, f)

    def _setup_volumes(self, ctx: ExecutionContext, gyoza_dir: str) -> None:
        """Mount the worker's shared paths and the run's own directory.

        Parameters
        ----------
        ctx : ExecutionContext
            Execution context to update.
        gyoza_dir : str
            Path to mount inside the container.
        """
        ctx.volumes.update(shared_runtime.mounts())

        cache_dir = shared_runtime.cache_dir()
        ctx.volumes[cache_dir] = {"bind": cache_dir, "mode": "rw"}
        ctx.volumes[gyoza_dir] = {"bind": gyoza_dir, "mode": "rw"}

    def _setup_env_vars(self, ctx: ExecutionContext, gyoza_dir: str) -> None:
        """Set the worker's shared variables and the run's own paths.

        Gyoza's own variables are set last, so a shared one cannot shadow them.

        Parameters
        ----------
        ctx : ExecutionContext
            Execution context to update.
        gyoza_dir : str
            Base directory for input/output files.
        """
        ctx.env_vars.update(shared_runtime.env())

        ctx.env_vars[shared_runtime.CACHE_DIR_VAR] = shared_runtime.cache_dir()
        ctx.env_vars["GYOZA_INPUT_PATH"] = os.path.join(gyoza_dir, _INPUT_FILENAME)
        ctx.env_vars["GYOZA_OUTPUT_PATH"] = os.path.join(gyoza_dir, _OUTPUT_FILENAME)



prepare_inputs = PrepareInputs()