<commit_before># -*- coding: utf-8 -*-
"""TODO(e-carlin): Doc

:copyright: Copyright (c) 2019 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function
from pykern import pkconfig
from pykern import pkio
from pykern import pkjson
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdp, pkdlog, pkdexc
from sirepo import job
from sirepo import job_driver
from sirepo import util
import asyncssh
import sirepo.simulation_db
import sirepo.srdb
import time
import tornado.ioloop


cfg = None

_KNOWN_HOSTS = None


class SbatchDriver(job_driver.DriverBase):

    instances = PKDict()

    def __init__(self, req):
        super().__init__(req)
#TODO(robnagler) read a db for an sbatch_user
        self._srdb_root = None
        self.instances[self.uid] = self

    @classmethod
    async def get_instance(cls, req):
        u = req.content.uid
        return cls.instances.pksetdefault(u, lambda: cls(req))[u]

    @classmethod
    def init_class(cls):
        return cls

    async def kill(self):
        if not self.websocket:
            # if there is no websocket then we don't know about the agent
            # so we can't do anything
            return
        # hopefully the agent is nice and listens to the kill
        self.websocket.write_message(PKDict(opName=job.OP_KILL))

    def run_scheduler(self, exclude_self=False):
        for d in self.instances.values():
            if exclude_self and d == self:
                continue
            for o in d.get_ops_with_send_allocation():
                assert o.opId not in d.ops_pending_done
                d.ops_pending_send.remove(o)
                d.ops_pending_done[o.opId] = o
#TODO(robnagler) encapsulation is incorrect. Superclass should make
# decisions about send_ready.
                o.send_ready.set()

    async def send(self, op):
        m = op.msg
        try:
            self._creds = m.pkdel('sbatchCredentials')
            if self._srdb_root is None:
                assert not self.websocket, \
                    'expected no agent if _srdb_root not set msg={}'.format(m)
                if not self._creds or 'username' not in self._creds:
                    self._raise_sbatch_login_srexception('no-creds', m)
                self._srdb_root = cfg.srdb_root.format(
                    sbatch_user=self._creds.username,
                )
            m.userDir = '/'.join(
                (
                    str(self._srdb_root),
                    sirepo.simulation_db.USER_ROOT_DIR,
                    m.uid,
                )
            )
            m.runDir = '/'.join((m.userDir, m.simulationType, m.computeJid))
            if op.opName == job.OP_RUN:
                assert m.sbatchHours
                if cfg.cores:
                    m.sbatchCores = min(m.mpiCores, cfg.cores)
                m.mpiCores = m.sbatchCores
                if op.kind == job.PARALLEL:
                    op.maxRunSecs = 0
            m.shifterImage = cfg.shifter_image
            return await super().send(op)
        finally:
            self.pkdel('_creds')

    def _agent_env(self):
        return super()._agent_env(
            env=PKDict(
                SIREPO_SRDB_ROOT=self._srdb_root,
            )
        )

    async def _do_agent_start(self, msg):
        log_file = 'job_agent.log'
        agent_start_dir = self._srdb_root
        script = f'''#!/bin/bash
{self._agent_start_dev()}
set -e
mkdir -p '{agent_start_dir}'
cd '{self._srdb_root}'
{self._agent_env()}
(/usr/bin/env; setsid {cfg.sirepo_cmd} job_agent start_sbatch) >& {log_file} &
disown
'''

        def write_to_log(stdout, stderr, filename):
            p = pkio.py_path(msg.userDir).join('log')
            if not p.exists():
                pkio.mkdir_parent(p)
            pkjson.dump_pretty(
                PKDict(stdout=stdout, stderr=stderr),
                p.join(f'{int(time.time())}-{filename}.log'),
            )

        async def get_agent_log(connection):
            await tornado.gen.sleep(cfg.agent_log_read_sleep)
            async with connection.create_process(
                    f'/bin/bash -c "/usr/bin/cat {agent_start_dir}/{log_file}"'
            ) as p:
                o, e = await p.communicate()
                write_to_log(o, e, 'remote-job-agent-log')

        try:
            async with asyncssh.connect(
                cfg.host,
                username=self._creds.username,<commit_msg>resolve pr comments<commit_after># -*- coding: utf-8 -*-
"""TODO(e-carlin): Doc

:copyright: Copyright (c) 2019 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function
from pykern import pkconfig
from pykern import pkio
from pykern import pkjson
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdp, pkdlog, pkdexc
from sirepo import job
from sirepo import job_driver
from sirepo import util
import asyncssh
import datetime
import sirepo.simulation_db
import sirepo.srdb
import tornado.ioloop


cfg = None

_KNOWN_HOSTS = None


class SbatchDriver(job_driver.DriverBase):

    instances = PKDict()

    def __init__(self, req):
        super().__init__(req)
#TODO(robnagler) read a db for an sbatch_user
        self._srdb_root = None
        self.instances[self.uid] = self

    @classmethod
    async def get_instance(cls, req):
        u = req.content.uid
        return cls.instances.pksetdefault(u, lambda: cls(req))[u]

    @classmethod
    def init_class(cls):
        return cls

    async def kill(self):
        if not self.websocket:
            # if there is no websocket then we don't know about the agent
            # so we can't do anything
            return
        # hopefully the agent is nice and listens to the kill
        self.websocket.write_message(PKDict(opName=job.OP_KILL))

    def run_scheduler(self, exclude_self=False):
        for d in self.instances.values():
            if exclude_self and d == self:
                continue
            for o in d.get_ops_with_send_allocation():
                assert o.opId not in d.ops_pending_done
                d.ops_pending_send.remove(o)
                d.ops_pending_done[o.opId] = o
#TODO(robnagler) encapsulation is incorrect. Superclass should make
# decisions about send_ready.
                o.send_ready.set()

    async def send(self, op):
        m = op.msg
        try:
            self._creds = m.pkdel('sbatchCredentials')
            if self._srdb_root is None:
                assert not self.websocket, \
                    'expected no agent if _srdb_root not set msg={}'.format(m)
                if not self._creds or 'username' not in self._creds:
                    self._raise_sbatch_login_srexception('no-creds', m)
                self._srdb_root = cfg.srdb_root.format(
                    sbatch_user=self._creds.username,
                )
            m.userDir = '/'.join(
                (
                    str(self._srdb_root),
                    sirepo.simulation_db.USER_ROOT_DIR,
                    m.uid,
                )
            )
            m.runDir = '/'.join((m.userDir, m.simulationType, m.computeJid))
            if op.opName == job.OP_RUN:
                assert m.sbatchHours
                if cfg.cores:
                    m.sbatchCores = min(m.mpiCores, cfg.cores)
                m.mpiCores = m.sbatchCores
                if op.kind == job.PARALLEL:
                    op.maxRunSecs = 0
            m.shifterImage = cfg.shifter_image
            return await super().send(op)
        finally:
            self.pkdel('_creds')

    def _agent_env(self):
        return super()._agent_env(
            env=PKDict(
                SIREPO_SRDB_ROOT=self._srdb_root,
            )
        )

    async def _do_agent_start(self, msg):
        log_file = 'job_agent.log'
        agent_start_dir = self._srdb_root
        script = f'''#!/bin/bash
{self._agent_start_dev()}
set -e
mkdir -p '{agent_start_dir}'
cd '{self._srdb_root}'
{self._agent_env()}
(/usr/bin/env; setsid {cfg.sirepo_cmd} job_agent start_sbatch) >& {log_file} &
disown
'''

        def write_to_log(stdout, stderr, filename):
            p = pkio.py_path(msg.userDir).join('log')
            pkio.mkdir_parent(p)
            pkjson.dump_pretty(
                PKDict(stdout=stdout, stderr=stderr),
                p.join(f'{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}-{filename}.log'),
            )

        async def get_agent_log(connection):
            await tornado.gen.sleep(cfg.agent_log_read_sleep)
            async with connection.create_process(
                    f'/bin/bash -c "/usr/bin/cat {agent_start_dir}/{log_file}"'
            ) as p:
                o, e = await p.communicate()
                write_to_log(o, e, 'remote-job-agent-log')

        try:
            async with asyncssh.connect(
                cfg.host,
                username=self._creds.username,
{'fraction_non_alphanumeric': 0.09901646590783512, 'fraction_numerical': 0.001768151176925627, 'mean_word_length': 1.9052969502407704, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7192070', 'n_tokens_mistral': 2932, 'n_tokens_neox': 2791, 'n_words': 636}