Esempio n. 1
0
    def __new__(cls,
                cmd,
                *args,
                iterable=False,
                read=False,
                bench_record=None,
                **kwargs):
        if "stepout" in kwargs:
            raise KeyError("Argument stepout is not allowed in shell command.")
        cmd = format(cmd, *args, stepout=2, **kwargs)
        context = inspect.currentframe().f_back.f_locals

        stdout = sp.PIPE if iterable or read else STDOUT

        close_fds = sys.platform != "win32"

        jobid = context.get("jobid")
        if not context.get("is_shell"):
            logger.shellcmd(cmd)

        env_prefix = ""
        conda_env = context.get("conda_env", None)
        singularity_img = context.get("singularity_img", None)
        shadow_dir = context.get("shadow_dir", None)

        cmd = "{} {} {}".format(cls._process_prefix, cmd.strip(),
                                cls._process_suffix).strip()

        conda = None
        if conda_env:
            cmd = Conda(singularity_img).shellcmd(conda_env, cmd)

        if singularity_img:
            args = context.get("singularity_args", "")
            cmd = singularity.shellcmd(
                singularity_img,
                cmd,
                args,
                shell_executable=cls._process_args["executable"],
                container_workdir=shadow_dir,
            )
            logger.info(
                "Activating singularity image {}".format(singularity_img))
        if conda_env:
            logger.info("Activating conda environment: {}".format(conda_env))

        proc = sp.Popen(cmd,
                        bufsize=-1,
                        shell=True,
                        stdout=stdout,
                        universal_newlines=iterable or None,
                        close_fds=close_fds,
                        **cls._process_args)

        if jobid is not None:
            with cls._lock:
                cls._processes[jobid] = proc

        ret = None
        if iterable:
            return cls.iter_stdout(proc, cmd)
        if read:
            ret = proc.stdout.read()
        if bench_record is not None:
            from snakemake.benchmark import benchmarked

            with benchmarked(proc.pid, bench_record):
                retcode = proc.wait()
        else:
            retcode = proc.wait()

        if jobid is not None:
            with cls._lock:
                del cls._processes[jobid]

        if retcode:
            raise sp.CalledProcessError(retcode, cmd)
        return ret
Esempio n. 2
0
    def __new__(
        cls, cmd, *args, iterable=False, read=False, bench_record=None, **kwargs
    ):
        if "stepout" in kwargs:
            raise KeyError("Argument stepout is not allowed in shell command.")
        cmd = format(cmd, *args, stepout=2, **kwargs)
        context = inspect.currentframe().f_back.f_locals
        print(context)

        stdout = sp.PIPE if iterable or read else STDOUT

        close_fds = sys.platform != "win32"

        jobid = context.get("jobid")
        if not context.get("is_shell"):
            logger.shellcmd(cmd)

        env_prefix = ""
        conda_env = context.get("conda_env", None)
        singularity_img = context.get("singularity_img", None)
        shadow_dir = context.get("shadow_dir", None)

        cmd = "{} {} {}".format(
            format(cls._process_prefix, *args, stepout=2, **kwargs),
            cmd.strip(),
            format(cls._process_suffix, *args, stepout=2, **kwargs),
        ).strip()
        print(cmd)

        conda = None
        if conda_env:
            cmd = Conda(singularity_img).shellcmd(conda_env, cmd)

        if singularity_img:
            args = context.get("singularity_args", "")
            cmd = singularity.shellcmd(
                singularity_img,
                cmd,
                args,
                shell_executable=cls._process_args["executable"],
                container_workdir=shadow_dir,
            )
            logger.info("Activating singularity image {}".format(singularity_img))
        if conda_env:
            logger.info("Activating conda environment: {}".format(conda_env))

        proc = sp.Popen(
            cmd,
            bufsize=-1,
            shell=True,
            stdout=stdout,
            universal_newlines=iterable or None,
            close_fds=close_fds,
            **cls._process_args
        )

        if jobid is not None:
            with cls._lock:
                cls._processes[jobid] = proc

        ret = None
        if iterable:
            return cls.iter_stdout(proc, cmd)
        if read:
            ret = proc.stdout.read()
        if bench_record is not None:
            from benchmark import benchmarked

            gpu = None
            try:
                gpu = [int(x) for x in context.get("params").devices.split(",")]
                print("[snakeshell] Attempting to benchmark GPU process with GPUtil on devices: %s" % gpu)
            except Exception as e:
                pass

            rt_bench_path = None
            try:
                rt_bench_path = context.get("output").rtbench
                print("[snakeshell] Attempting to benchmark in real-time to: %s" % rt_bench_path)
            except Exception as e:
                pass

            with benchmarked(proc.pid, bench_record, gpus=gpu, rt_path=rt_bench_path, interval=15):
                retcode = proc.wait()
        else:
            retcode = proc.wait()

        if jobid is not None:
            with cls._lock:
                del cls._processes[jobid]

        if retcode:
            raise sp.CalledProcessError(retcode, cmd)
        return ret
Esempio n. 3
0
 def _get_cmd(self, cmd):
     if self.singularity_img:
         return singularity.shellcmd(self.singularity_img, cmd)
     return cmd
Esempio n. 4
0
    def create(self, dryrun=False):
        """ Create the conda enviroment."""
        from snakemake.shell import shell

        # Read env file and create hash.
        env_file = self.file
        tmp_file = None

        url_scheme, *_ = urlparse(env_file)
        if (url_scheme and not url_scheme == 'file') or \
            (not url_scheme and env_file.startswith("git+file:/")):
            with tempfile.NamedTemporaryFile(delete=False, suffix=".yaml") as tmp:
                tmp.write(self.content)
                env_file = tmp.name
                tmp_file = tmp.name

        env_hash = self.hash
        env_path = self.path

        # Check for broken environment
        if os.path.exists(os.path.join(env_path,"env_setup_start")) and not os.path.exists(os.path.join(env_path,"env_setup_done")):
            if dryrun:
                logger.info("Incomplete Conda environment {} will be recreated.".format(utils.simplify_path(self.file)))
            else:
                logger.info("Removing incomplete Conda environment {}...".format(utils.simplify_path(self.file)))
                shutil.rmtree(env_path, ignore_errors=True)

        # Create environment if not already present.
        if not os.path.exists(env_path):
            if dryrun:
                logger.info("Conda environment {} will be created.".format(utils.simplify_path(self.file)))
                return env_path
            conda = Conda(self._singularity_img)
            logger.info("Creating conda environment {}...".format(
                        utils.simplify_path(self.file)))
            # Check if env archive exists. Use that if present.
            env_archive = self.archive_file
            try:
                # Touch "start" flag file
                os.makedirs(env_path, exist_ok=True)
                with open(os.path.join(env_path,"env_setup_start"), "a") as f:
                    pass

                if os.path.exists(env_archive):
                    logger.info("Using archived local conda packages.")
                    pkg_list = os.path.join(env_archive, "packages.txt")
                    if os.path.exists(pkg_list):
                        # read pacakges in correct order
                        # this is for newer env archives where the package list
                        # was stored
                        packages = [os.path.join(env_archive, pkg.rstrip())
                                    for pkg in open(pkg_list)]
                    else:
                        # guess order
                        packages = glob(os.path.join(env_archive, "*.tar.bz2"))

                    # install packages manually from env archive
                    cmd = " ".join(
                        ["conda", "create", "--copy", "--prefix '{}'".format(env_path)] +
                        packages)
                    if self._singularity_img:
                        cmd = singularity.shellcmd(self._singularity_img.path, cmd)
                    out = shell.check_output(cmd, stderr=subprocess.STDOUT)

                else:
                    # Copy env file to env_path (because they can be on
                    # different volumes and singularity should only mount one).
                    # In addition, this allows to immediately see what an
                    # environment in .snakemake/conda contains.
                    target_env_file = env_path + ".yaml"
                    shutil.copy(env_file, target_env_file)

                    logger.info("Downloading remote packages.")
                    cmd = " ".join(["conda", "env", "create",
                                                "--file '{}'".format(target_env_file),
                                                "--prefix '{}'".format(env_path)])
                    if self._singularity_img:
                        cmd = singularity.shellcmd(self._singularity_img.path, cmd)
                    out = shell.check_output(cmd, stderr=subprocess.STDOUT)
                # Touch "done" flag file
                with open(os.path.join(env_path,"env_setup_done"), "a") as f:
                    pass

                logger.debug(out.decode())
                logger.info("Environment for {} created (location: {})".format(
                            os.path.relpath(env_file), os.path.relpath(env_path)))
            except subprocess.CalledProcessError as e:
                # remove potential partially installed environment
                shutil.rmtree(env_path, ignore_errors=True)
                raise CreateCondaEnvironmentException(
                    "Could not create conda environment from {}:\n".format(env_file) +
                    e.output.decode())

        if tmp_file:
            # temporary file was created
            os.remove(tmp_file)

        return env_path
Esempio n. 5
0
 def get_cmd(cmd, singularity_img=None):
     if singularity_img:
         return singularity.shellcmd(self.singularity_img.path, cmd)
     return cmd
Esempio n. 6
0
    def create(self, dryrun=False):
        """ Create the conda enviroment."""
        if self._singularity_img:
            check_conda(self._singularity_img)

        # Read env file and create hash.
        env_file = self.file
        tmp_file = None

        url_scheme, *_ = urlparse(env_file)
        if url_scheme and not url_scheme == 'file':
            with tempfile.NamedTemporaryFile(delete=False,
                                             suffix=".yaml") as tmp:
                tmp.write(self.content)
                env_file = tmp.name
                tmp_file = tmp.name

        env_hash = self.hash
        env_path = self.path
        # Create environment if not already present.
        if not os.path.exists(env_path):
            if dryrun:
                logger.info("Conda environment {} will be created.".format(
                    utils.simplify_path(self.file)))
                return env_path
            logger.info("Creating conda environment {}...".format(
                utils.simplify_path(self.file)))
            # Check if env archive exists. Use that if present.
            env_archive = self.archive_file
            try:
                if os.path.exists(env_archive):
                    logger.info("Using archived local conda packages.")
                    # install packages manually from env archive
                    cmd = " ".join(
                        ["conda", "create", "--copy", "--prefix", env_path] +
                        glob(os.path.join(env_archive, "*.tar.bz2")))
                    if self._singularity_img:
                        cmd = singularity.shellcmd(self._singularity_img.path,
                                                   cmd)
                    out = subprocess.check_output(cmd,
                                                  shell=True,
                                                  stderr=subprocess.STDOUT)

                else:
                    # Copy env file to env_path (because they can be on
                    # different volumes and singularity should only mount one).
                    # In addition, this allows to immediately see what an
                    # environment in .snakemake/conda contains.
                    target_env_file = env_path + ".yaml"
                    shutil.copy(env_file, target_env_file)

                    logger.info("Downloading remote packages.")
                    cmd = " ".join([
                        "conda", "env", "create", "--file", target_env_file,
                        "--prefix", env_path
                    ])
                    if self._singularity_img:
                        cmd = singularity.shellcmd(self._singularity_img.path,
                                                   cmd)
                    out = subprocess.check_output(cmd,
                                                  shell=True,
                                                  stderr=subprocess.STDOUT)
                logger.debug(out.decode())
                logger.info("Environment for {} created (location: {})".format(
                    os.path.relpath(env_file), os.path.relpath(env_path)))
            except subprocess.CalledProcessError as e:
                # remove potential partially installed environment
                shutil.rmtree(env_path, ignore_errors=True)
                raise CreateCondaEnvironmentException(
                    "Could not create conda environment from {}:\n".format(
                        env_file) + e.output.decode())

        if tmp_file:
            # temporary file was created
            os.remove(tmp_file)

        return env_path