<commit_before>"""
Unit tests for the router class. Please don't add any test that will involve
controller or the backend worker, use mock if necessary.
"""
import asyncio
from collections import defaultdict
import os

import pytest
from ray.serve.context import TaskContext

import ray
from ray.serve.config import BackendConfig
from ray.serve.controller import TrafficPolicy
from ray.serve.router import Query, ReplicaSet, RequestMetadata, Router
from ray.serve.utils import get_random_letters
from ray.test_utils import SignalActor

pytestmark = pytest.mark.asyncio

@pytest.fixture
def ray_instance():
    os.environ["SERVE_LOG_DEBUG"] = "1"  # Turns on debug log for tests
    ray.init(num_cpus=16)
    yield
    ray.shutdown()

def mock_task_runner():
    @ray.remote(num_cpus=0)
    class TaskRunnerMock:
        def __init__(self):
            self.query = None
            self.queries = []

        async def handle_request(self, request):
            if isinstance(request, bytes):
                request = Query.ray_deserialize(request)
            self.query = request
            self.queries.append(request)
            return "DONE"

        def get_recent_call(self):
            return self.query

        def get_all_calls(self):
            return self.queries

        def clear_calls(self):
            self.queries = []

        def ready(self):
            pass

    return TaskRunnerMock.remote()

@pytest.fixture
def task_runner_mock_actor():
    yield mock_task_runner()

@pytest.fixture
def mock_controller():
    @ray.remote(num_cpus=0)
    class MockControllerActor:
        def __init__(self):
            from ray.serve.long_poll import LongPollerHost
            self.host = LongPollerHost()
            self.backend_replicas = defaultdict(list)
            self.backend_configs = dict()
            self.clear()

        def clear(self):
            self.host.notify_changed("worker_handles", {})
            self.host.notify_changed("traffic_policies", {})
            self.host.notify_changed("backend_configs", {})

        async def listen_for_change(self, snapshot_ids):
            return await self.host.listen_for_change(snapshot_ids)

        def set_traffic(self, endpoint, traffic_policy):
            self.host.notify_changed("traffic_policies",
                                     {endpoint: traffic_policy})

        def add_new_replica(self,
                            backend_tag,
                            runner_actor,
                            backend_config=BackendConfig()):
            self.backend_replicas[backend_tag].append(runner_actor)
            self.backend_configs[backend_tag] = backend_config

            self.host.notify_changed(
                "worker_handles",
                self.backend_replicas,
            )
            self.host.notify_changed("backend_configs", self.backend_configs)

    yield MockControllerActor.remote()

async def test_simple_endpoint_backend_pair(ray_instance, mock_controller,
                                            task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    # Propogate configs
    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-single-prod": 1.0
        }))
    await mock_controller.add_new_replica.remote("backend-single-prod",
                                                 task_runner_mock_actor)

    # Make sure we get the request result back
    ref = await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 1)
    result = await ref
    assert result == "DONE"

    # Make sure it's the right request
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 1
    assert got_work.kwargs == {}

async def test_changing_backend(ray_instance, mock_controller,
                                task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-alter": 1
        }))
    await mock_controller.add_new_replica.remote("backend-alter",
                                                 task_runner_mock_actor)

    await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 1)
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 1

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-alter-2": 1
        }))
    await mock_controller.add_new_replica.remote("backend-alter-2",
                                                 task_runner_mock_actor)
    await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 2)
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 2

async def test_split_traffic_random(ray_instance, mock_controller,
                                    task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-split": 0.5,
            "backend-split-2": 0.5
        }))
    runner_1, runner_2 = [mock_task_runner() for _ in range(2)]
    await mock_controller.add_new_replica.remote("backend-split", runner_1)
    await mock_controller.add_new_replica.remote("backend-split-2", runner_2)

    # assume 50% split, the probability of all 20 requests goes to a
    # single queue is 0.5^20 ~ 1-6
    for _ in range(20):
        await q.assign_request.remote(
            RequestMetadata(get_random_letters(10), "svc", None), 1)

    got_work = [
        await runner.get_recent_call.remote()
        for runner in (runner_1, runner_2)
    ]
    assert [g.args[0] for g in got_work] == [1, 1]

async def test_shard_key(ray_instance, mock_controller,
                         task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    num_backends = 5
    traffic_dict = {}
    runners = [mock_task_runner() for _ in range(num_backends)]
    for i, runner in enumerate(runners):
        backend_name = "backend-split-" + str(i)
        traffic_dict[backend_name] = 1.0 / num_backends
        await mock_controller.add_new_replica.remote(backend_name, runner)
    await mock_controller.set_traffic.remote("svc",
                                             TrafficPolicy(traffic_dict))

    # Generate random shard keys and send one request for each.
    shard_keys = [get_random_letters() for _ in range(100)]
    for shard_key in shard_keys:
        await q.assign_request.remote(
            RequestMetadata(
                get_random_letters(10), "svc", None, shard_key=shard_key),
            shard_key)

    # Log the shard keys that were assigned to each backend.
    runner_shard_keys = defaultdict(set)
    for i, runner in enumerate(runners):
        calls = await runner.get_all_calls.remote()
        for call in calls:
            runner_shard_keys[i].add(call.args[0])
        await runner.clear_calls.remote()

    # Send queries with the same shard keys a second time.
    for shard_key in shard_keys:
        await q.assign_request.remote(
            RequestMetadata(
                get_random_letters(10), "svc", None, shard_key=shard_key),
            shard_key)

    # Check that the requests were all mapped to the same backends.
    for i, runner in enumerate(runners):
        calls = await runner.get_all_calls.remote()
        for call in calls:
            assert call.args[0] in runner_shard_keys[i]

async def test_replica_set(ray_instance):
    signal = SignalActor.remote()

    @ray.remote(num_cpus=0)
    class MockWorker:
        _num_queries = 0

        async def handle_request(self, request):
            self._num_queries += 1
            await signal.wait.remote()
            return "DONE"

        async def num_queries(self):
            return self._num_queries

    # We will test a scenario with two replicas in the replica set.
    rs = ReplicaSet()
    workers = [MockWorker.remote() for _ in range(2)]
    rs.set_max_concurrent_queries(1)
    rs.update_worker_replicas(workers)

    # Send two queries. They should go through the router but blocked by signal
    # actors.
    query = Query([], {}, TaskContext.Python,
                  RequestMetadata("request-id", "endpoint",
                                  TaskContext.Python))
    first_ref = await rs.assign_replica(query)
    second_ref = await rs.assign_replica(query)

    # These should be blocked by signal actor.
    with pytest.raises(ray.exceptions.GetTimeoutError):
        ray.get([first_ref, second_ref], timeout=1)

    # Each replica should have exactly one inflight query. Let make sure the
    # queries arrived there.
    for worker in workers:
        while await worker.num_queries.remote() != 1:
            await asyncio.sleep(1)

    # Let's try to send another query.
    third_ref_pending_task = asyncio.get_event_loop().create_task(
        rs.assign_replica(query))
    # We should fail to assign a replica, so this coroutine should still be
    # pending after some time.
    await asyncio.sleep(0.2)
    assert not third_ref_pending_task.done()

    # Let's unblock the two workers
    await signal.send.remote()
    assert await first_ref == "DONE"
    assert await second_ref == "DONE"

    # The third request should be unblocked and sent to first worker.
    # This meas we should be able to get the object ref.
    third_ref = await third_ref_pending_task

    # Now we got the object ref, let's get it result.
    await signal.send.remote()
    assert await third_ref == "DONE"

    # Finally, make sure that one of the replica processed the third query.
    num_queries_set = {(await worker.num_queries.remote())
                       for worker in workers}
    assert num_queries_set == {2, 1}

if __name__ == "__main__":
    import sys
    sys.exit(pytest.main(["-v", "-s", __file__]))
<commit_msg>Deflake test_router (#12175)<commit_after>"""
Unit tests for the router class. Please don't add any test that will involve
controller or the backend worker, use mock if necessary.
"""
import asyncio
from collections import defaultdict
import os

import pytest
from ray.serve.context import TaskContext

import ray
from ray.serve.config import BackendConfig
from ray.serve.controller import TrafficPolicy
from ray.serve.router import Query, ReplicaSet, RequestMetadata, Router
from ray.serve.utils import get_random_letters
from ray.test_utils import SignalActor

pytestmark = pytest.mark.asyncio

@pytest.fixture
def ray_instance():
    os.environ["SERVE_LOG_DEBUG"] = "1"  # Turns on debug log for tests
    ray.init(num_cpus=16)
    yield
    ray.shutdown()

def mock_task_runner():
    @ray.remote(num_cpus=0)
    class TaskRunnerMock:
        def __init__(self):
            self.query = None
            self.queries = []

        async def handle_request(self, request):
            if isinstance(request, bytes):
                request = Query.ray_deserialize(request)
            self.query = request
            self.queries.append(request)
            return "DONE"

        def get_recent_call(self):
            return self.query

        def get_all_calls(self):
            return self.queries

        def clear_calls(self):
            self.queries = []

        def ready(self):
            pass

    return TaskRunnerMock.remote()

@pytest.fixture
def task_runner_mock_actor():
    yield mock_task_runner()

@pytest.fixture
def mock_controller():
    @ray.remote(num_cpus=0)
    class MockControllerActor:
        def __init__(self):
            from ray.serve.long_poll import LongPollerHost
            self.host = LongPollerHost()
            self.backend_replicas = defaultdict(list)
            self.backend_configs = dict()
            self.clear()

        def clear(self):
            self.host.notify_changed("worker_handles", {})
            self.host.notify_changed("traffic_policies", {})
            self.host.notify_changed("backend_configs", {})

        async def listen_for_change(self, snapshot_ids):
            return await self.host.listen_for_change(snapshot_ids)

        def set_traffic(self, endpoint, traffic_policy):
            self.host.notify_changed("traffic_policies",
                                     {endpoint: traffic_policy})

        def add_new_replica(self,
                            backend_tag,
                            runner_actor,
                            backend_config=BackendConfig()):
            self.backend_replicas[backend_tag].append(runner_actor)
            self.backend_configs[backend_tag] = backend_config

            self.host.notify_changed(
                "worker_handles",
                self.backend_replicas,
            )
            self.host.notify_changed("backend_configs", self.backend_configs)

    yield MockControllerActor.remote()

async def test_simple_endpoint_backend_pair(ray_instance, mock_controller,
                                            task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    # Propogate configs
    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-single-prod": 1.0
        }))
    await mock_controller.add_new_replica.remote("backend-single-prod",
                                                 task_runner_mock_actor)

    # Make sure we get the request result back
    ref = await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 1)
    result = await ref
    assert result == "DONE"

    # Make sure it's the right request
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 1
    assert got_work.kwargs == {}

async def test_changing_backend(ray_instance, mock_controller,
                                task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-alter": 1
        }))
    await mock_controller.add_new_replica.remote("backend-alter",
                                                 task_runner_mock_actor)

    await (await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 1))
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 1

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-alter-2": 1
        }))
    await mock_controller.add_new_replica.remote("backend-alter-2",
                                                 task_runner_mock_actor)
    await (await q.assign_request.remote(
        RequestMetadata(get_random_letters(10), "svc", None), 2))
    got_work = await task_runner_mock_actor.get_recent_call.remote()
    assert got_work.args[0] == 2

async def test_split_traffic_random(ray_instance, mock_controller,
                                    task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    await mock_controller.set_traffic.remote(
        "svc", TrafficPolicy({
            "backend-split": 0.5,
            "backend-split-2": 0.5
        }))
    runner_1, runner_2 = [mock_task_runner() for _ in range(2)]
    await mock_controller.add_new_replica.remote("backend-split", runner_1)
    await mock_controller.add_new_replica.remote("backend-split-2", runner_2)

    # assume 50% split, the probability of all 20 requests goes to a
    # single queue is 0.5^20 ~ 1-6
    object_refs = []
    for _ in range(20):
        ref = await q.assign_request.remote(
            RequestMetadata(get_random_letters(10), "svc", None), 1)
        object_refs.append(ref)
    ray.get(object_refs)

    got_work = [
        await runner.get_recent_call.remote()
        for runner in (runner_1, runner_2)
    ]
    assert [g.args[0] for g in got_work] == [1, 1]

async def test_shard_key(ray_instance, mock_controller,
                         task_runner_mock_actor):
    q = ray.remote(Router).remote(mock_controller)
    await q.setup_in_async_loop.remote()

    num_backends = 5
    traffic_dict = {}
    runners = [mock_task_runner() for _ in range(num_backends)]
    for i, runner in enumerate(runners):
        backend_name = "backend-split-" + str(i)
        traffic_dict[backend_name] = 1.0 / num_backends
        await mock_controller.add_new_replica.remote(backend_name, runner)
    await mock_controller.set_traffic.remote("svc",
                                             TrafficPolicy(traffic_dict))

    # Generate random shard keys and send one request for each.
    shard_keys = [get_random_letters() for _ in range(100)]
    for shard_key in shard_keys:
        await (await q.assign_request.remote(
            RequestMetadata(
                get_random_letters(10), "svc", None, shard_key=shard_key),
            shard_key))

    # Log the shard keys that were assigned to each backend.
    runner_shard_keys = defaultdict(set)
    for i, runner in enumerate(runners):
        calls = await runner.get_all_calls.remote()
        for call in calls:
            runner_shard_keys[i].add(call.args[0])
        await runner.clear_calls.remote()

    # Send queries with the same shard keys a second time.
    for shard_key in shard_keys:
        await (await q.assign_request.remote(
            RequestMetadata(
                get_random_letters(10), "svc", None, shard_key=shard_key),
            shard_key))

    # Check that the requests were all mapped to the same backends.
    for i, runner in enumerate(runners):
        calls = await runner.get_all_calls.remote()
        for call in calls:
            assert call.args[0] in runner_shard_keys[i]

async def test_replica_set(ray_instance):
    signal = SignalActor.remote()

    @ray.remote(num_cpus=0)
    class MockWorker:
        _num_queries = 0

        async def handle_request(self, request):
            self._num_queries += 1
            await signal.wait.remote()
            return "DONE"

        async def num_queries(self):
            return self._num_queries

    # We will test a scenario with two replicas in the replica set.
    rs = ReplicaSet()
    workers = [MockWorker.remote() for _ in range(2)]
    rs.set_max_concurrent_queries(1)
    rs.update_worker_replicas(workers)

    # Send two queries. They should go through the router but blocked by signal
    # actors.
    query = Query([], {}, TaskContext.Python,
                  RequestMetadata("request-id", "endpoint",
                                  TaskContext.Python))
    first_ref = await rs.assign_replica(query)
    second_ref = await rs.assign_replica(query)

    # These should be blocked by signal actor.
    with pytest.raises(ray.exceptions.GetTimeoutError):
        ray.get([first_ref, second_ref], timeout=1)

    # Each replica should have exactly one inflight query. Let make sure the
    # queries arrived there.
    for worker in workers:
        while await worker.num_queries.remote() != 1:
            await asyncio.sleep(1)

    # Let's try to send another query.
    third_ref_pending_task = asyncio.get_event_loop().create_task(
        rs.assign_replica(query))
    # We should fail to assign a replica, so this coroutine should still be
    # pending after some time.
    await asyncio.sleep(0.2)
    assert not third_ref_pending_task.done()

    # Let's unblock the two workers
    await signal.send.remote()
    assert await first_ref == "DONE"
    assert await second_ref == "DONE"

    # The third request should be unblocked and sent to first worker.
    # This meas we should be able to get the object ref.
    third_ref = await third_ref_pending_task

    # Now we got the object ref, let's get it result.
    await signal.send.remote()
    assert await third_ref == "DONE"

    # Finally, make sure that one of the replica processed the third query.
    num_queries_set = {(await worker.num_queries.remote())
                       for worker in workers}
    assert num_queries_set == {2, 1}

if __name__ == "__main__":
    import sys
    sys.exit(pytest.main(["-v", "-s", __file__]))
<|endoftext|>


<commit_before># Copyright 2020 The SQLFlow Authors. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License
""" XGBoost Local Training.
This module launches a XGBoost training task on host.
"""
import types

import xgboost as xgb

def init_xgb_booster(load_pretrained_model, filename="my_model"):
    """
    Initialize XGBoost Booster from local saved model or just return None
    """
    if load_pretrained_model:
        bst = xgb.Booster()
        bst.load_model(filename)
        return bst
    return None

def train(train_dataset,
          train_params,
          model_params,
          val_dataset=None,
          load_xgb_pretrained_model=False):
    """ XGBoost local training API

    Args:
        train_dataset: Generator
            training dataset with XGBoost DMatrix generator
        train_params: dict
            training parameters, passed into `xgboost.train` API
        model_params: dict
            model parameters, the `model_param` arguments of `xgboost.train`
        val_dataset: Generator
            validation datasets generator with XGBoost DMatrix generator
        load_pretrained_model: bool
            load pre-trained model or not
    Returns:
        evaluation result
    """

    bst = init_xgb_booster(load_xgb_pretrained_model)
    eval_result = dict()
    for per_batch_dmatrix in train_dataset:
        watchlist = [(per_batch_dmatrix, "train")]
        if val_dataset:
            # the `xgboost.train` API accepts the XGBoost DMatrix object
            # as the training or validation dataste, so we should convert
            # the generator to DMatrix.
            if isinstance(val_dataset, types.GeneratorType):
                val_dataset = list(val_dataset)[0]
            watchlist.append((val_dataset, "validate"))

        bst = xgb.train(model_params,
                        per_batch_dmatrix,
                        evals=watchlist,
                        evals_result=eval_result,
                        xgb_model=bst,
                        **train_params)
        print("Evaluation result: %s" % eval_result)

    return eval_result
<commit_msg>Polish XGBoost local train code (#2715)<commit_after># Copyright 2020 The SQLFlow Authors. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License
""" XGBoost Local Training.
This module launches a XGBoost training task on host.
"""
import types

import xgboost as xgb

def init_xgb_booster(load_pretrained_model, filename="my_model"):
    """
    Initialize XGBoost Booster from local saved model or just return None
    """
    if load_pretrained_model:
        bst = xgb.Booster()
        bst.load_model(filename)
        return bst
    return None

def train(train_dataset,
          train_params,
          model_params,
          val_dataset=None,
          load_pretrained_model=False):
    """ XGBoost local training API

    Args:
        train_dataset: Generator
            training dataset with XGBoost DMatrix generator
        train_params: dict
            training parameters, passed into `xgboost.train` API
        model_params: dict
            model parameters, the `model_param` arguments of `xgboost.train`
        val_dataset: Generator
            validation datasets generator with XGBoost DMatrix generator
        load_pretrained_model: bool
            load pre-trained model or not
    Returns:
        evaluation result
    """

    bst = init_xgb_booster(load_pretrained_model)
    eval_result = dict()
    watchlist = [None]
    if val_dataset:
        # The `xgboost.train` API only accepts the XGBoost DMatrix
        # object as the training or validation dataset, so we should
        # convert the generator to DMatrix.
        if isinstance(val_dataset, types.GeneratorType):
            val_dataset = list(val_dataset)[0]
        watchlist.append((val_dataset, "validate"))

    for per_batch_dmatrix in train_dataset:
        watchlist[0] = (per_batch_dmatrix, "train")
        bst = xgb.train(model_params,
                        per_batch_dmatrix,
                        evals=watchlist,
                        evals_result=eval_result,
                        xgb_model=bst,
                        **train_params)
        print("Evaluation result: %s" % eval_result)

    return eval_result
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#!/usr/bin/python
##-------------------------------------------------------------------
## @copyright 2015 DennyZhang.com
## File : tcp_port_scan.py
## Author : DennyZhang.com <denny@dennyzhang.com>
## Description :
## --
## Created : <2016-01-15>
## Updated: Time-stamp: <2016-08-12 10:23:00>
##-------------------------------------------------------------------
import argparse
import subprocess
import os, sys

################################################################################
# TODO: move to common library
def strip_remove_comments(string):
    # remove empty lines and comments (# ...) from string
    l = []
    for line in string.split("\n"):
        line = line.strip()
        if line.startswith("#") or line == "":
            continue
        l.append(line)
    return "\n".join(l)

def strip_remove_emptylines(string):
    l = []
    for line in string.split("\n"):
        line = line.strip()
        if line == "":
            continue
        l.append(line)
    return "\n".join(l)

def string_remove_patterns(string, opt_list):
    l = []
    # remove entries from string
    for line in string.split("\n"):
        should_remove = False
        for item in opt_list:
            if item in line:
                should_remove = True
        if should_remove is False:
            l.append(line)
    return "\n".join(l)

def string_remove_extra_whitespace(string):
    return ' '.join(string.split(' '))

def load_paralist_from_file(fname):
    # TODO: defensive coding when file doesn't exist
    l = []
    with open(fname,'r') as f:
        for line in f:
            line = line.strip()
            if not line.startswith("#") and line != "":
                l.append(line)
    return l

# TODO: add common logging
################################################################################
nmap_command = "sudo nmap -sS -PN %s" # ("-p T:XXX,XXX 192.168.0.16")
output_prefix = "============="
open_port_dict = {}
insecure_port_dict = {}

def nmap_check(server_ip, ports):
    if ports == "":
        nmap_opts = server_ip
    else:
        nmap_opts = "-p %s %s" % (ports, server_ip)
    command = nmap_command % (nmap_opts)
    print output_prefix, "Run: %s" % (command)
    nmap_output =subprocess.check_output(command, shell=True)
    print nmap_output
    return nmap_output

def get_portlist_by_nmap_output(nmap_output, server_ip):
    opt_list = ["Starting Nmap ", "Nmap scan report for ", "Host is ", \
                "Not shown: ", " STATE ", " closed ports", " closed unknown"\
                " filtered ", "Nmap done: ", " scanned ports on "]
    output = string_remove_extra_whitespace(output)
    output = string_remove_patterns(nmap_output, opt_list)
    output = strip_remove_emptylines(output)
    return output.split("\n")

def audit_open_ports(port_list, white_list, server_ip):
    insecure_port_list = []
    for item in port_list:
        port = item.split("/")[0]

        if "*:%s" % (port) not in white_list and \
           "%s:%s" % (server_ip, port) not in white_list:
            insecure_port_list.append(item)

    return insecure_port_list

def tcp_port_scan(server_list, white_list, extra_port_list):
    # TODO: change to multi-threading
    for server_ip in server_list:
        nmap_output = nmap_check(server_ip, "")
        nmap_port_list = get_portlist_by_nmap_output(nmap_output, server_ip)
        open_port_dict[server_ip] = nmap_port_list

    # TODO: change to multi-threading
    extra_ports = ",".join(extra_port_list)
    if extra_ports != "":
        print "%s Run extra checks for given ports: %s" % (output_prefix, extra_ports)
        for server_ip in server_list:
            nmap_output = nmap_check(server_ip, "T:%s" % (extra_ports))
            nmap_port_list = get_portlist_by_nmap_output(nmap_output, server_ip)
            open_port_dict[server_ip] = sorted(list(set(open_port_dict[server_ip]) \
                                                    | set(nmap_port_list)))

    # reduce result: whether any insecure ports open
    detected_insecure_ports = False
    for server_ip in server_list:
        ports = audit_open_ports(open_port_dict[server_ip], white_list, server_ip)
        if len(ports) != 0:
            detected_insecure_ports = True
            insecure_port_dict[server_ip] = ports

    if detected_insecure_ports is True:
        print output_prefix, "Error: Detected insecure TCP ports open"
        for server_ip in insecure_port_dict.keys():
            print "\nServer: %s" % (server_ip)
            print "\n".join(insecure_port_dict[server_ip])
        sys.exit(1)
    else:
        print "OK: No insecure TCP ports open"

################################################################################
if __name__=='__main__':
    # Sample:
    # python ./tcp_port_scan.py --server_list_file /tmp/server_list --white_list_file /tmp/white_list --extra_port_list_file /tmp/extra_port_list
    parser = argparse.ArgumentParser()
    parser.add_argument('--server_list_file', required=True,
                        help="ip list to scan", type=str)
    parser.add_argument('--white_list_file', required=True,
                        help="safe ports to allow open", type=str)
    parser.add_argument('--extra_port_list_file', required=True,
                        help="customized tcp ports to scan", type=str)
    args = parser.parse_args()

    white_list = load_paralist_from_file(args.white_list_file)
    server_list = load_paralist_from_file(args.server_list_file)
    extra_port_list = load_paralist_from_file(args.extra_port_list_file)

    tcp_port_scan(server_list, white_list, extra_port_list)
## File : tcp_port_scan.py ends
<commit_msg>fix typo<commit_after># -*- coding: utf-8 -*-
#!/usr/bin/python
##-------------------------------------------------------------------
## @copyright 2015 DennyZhang.com
## File : tcp_port_scan.py
## Author : DennyZhang.com <denny@dennyzhang.com>
## Description :
## --
## Created : <2016-01-15>
## Updated: Time-stamp: <2016-08-12 10:24:04>
##-------------------------------------------------------------------
import argparse
import subprocess
import os, sys

################################################################################
# TODO: move to common library
def strip_remove_comments(string):
    # remove empty lines and comments (# ...) from string
    l = []
    for line in string.split("\n"):
        line = line.strip()
        if line.startswith("#") or line == "":
            continue
        l.append(line)
    return "\n".join(l)

def strip_remove_emptylines(string):
    l = []
    for line in string.split("\n"):
        line = line.strip()
        if line == "":
            continue
        l.append(line)
    return "\n".join(l)

def string_remove_patterns(string, opt_list):
    l = []
    # remove entries from string
    for line in string.split("\n"):
        should_remove = False
        for item in opt_list:
            if item in line:
                should_remove = True
        if should_remove is False:
            l.append(line)
    return "\n".join(l)

def string_remove_extra_whitespace(string):
    return ' '.join(string.split(' '))

def load_paralist_from_file(fname):
    # TODO: defensive coding when file doesn't exist
    l = []
    with open(fname,'r') as f:
        for line in f:
            line = line.strip()
            if not line.startswith("#") and line != "":
                l.append(line)
    return l

# TODO: add common logging
################################################################################
nmap_command = "sudo nmap -sS -PN %s" # ("-p T:XXX,XXX 192.168.0.16")
output_prefix = "============="
open_port_dict = {}
insecure_port_dict = {}

def nmap_check(server_ip, ports):
    if ports == "":
        nmap_opts = server_ip
    else:
        nmap_opts = "-p %s %s" % (ports, server_ip)
    command = nmap_command % (nmap_opts)
    print output_prefix, "Run: %s" % (command)
    nmap_output =subprocess.check_output(command, shell=True)
    print nmap_output
    return nmap_output

def get_portlist_by_nmap_output(nmap_output, server_ip):
    opt_list = ["Starting Nmap ", "Nmap scan report for ", "Host is ", \
                "Not shown: ", " STATE ", " closed ports", " closed unknown"\
                " filtered ", "Nmap done: ", " scanned ports on "]
    output = string_remove_extra_whitespace(nmap_output)
    output = string_remove_patterns(output, opt_list)
    output = strip_remove_emptylines(output)
    return output.split("\n")

def audit_open_ports(port_list, white_list, server_ip):
    insecure_port_list = []
    for item in port_list:
        port = item.split("/")[0]

        if "*:%s" % (port) not in white_list and \
           "%s:%s" % (server_ip, port) not in white_list:
            insecure_port_list.append(item)

    return insecure_port_list

def tcp_port_scan(server_list, white_list, extra_port_list):
    # TODO: change to multi-threading
    for server_ip in server_list:
        nmap_output = nmap_check(server_ip, "")
        nmap_port_list = get_portlist_by_nmap_output(nmap_output, server_ip)
        open_port_dict[server_ip] = nmap_port_list

    # TODO: change to multi-threading
    extra_ports = ",".join(extra_port_list)
    if extra_ports != "":
        print "%s Run extra checks for given ports: %s" % (output_prefix, extra_ports)
        for server_ip in server_list:
            nmap_output = nmap_check(server_ip, "T:%s" % (extra_ports))
            nmap_port_list = get_portlist_by_nmap_output(nmap_output, server_ip)
            open_port_dict[server_ip] = sorted(list(set(open_port_dict[server_ip]) \
                                                    | set(nmap_port_list)))

    # reduce result: whether any insecure ports open
    detected_insecure_ports = False
    for server_ip in server_list:
        ports = audit_open_ports(open_port_dict[server_ip], white_list, server_ip)
        if len(ports) != 0:
            detected_insecure_ports = True
            insecure_port_dict[server_ip] = ports

    if detected_insecure_ports is True:
        print output_prefix, "Error: Detected insecure TCP ports open"
        for server_ip in insecure_port_dict.keys():
            print "\nServer: %s" % (server_ip)
            print "\n".join(insecure_port_dict[server_ip])
        sys.exit(1)
    else:
        print "OK: No insecure TCP ports open"

################################################################################
if __name__=='__main__':
    # Sample:
    # python ./tcp_port_scan.py --server_list_file /tmp/server_list --white_list_file /tmp/white_list --extra_port_list_file /tmp/extra_port_list
    parser = argparse.ArgumentParser()
    parser.add_argument('--server_list_file', required=True,
                        help="ip list to scan", type=str)
    parser.add_argument('--white_list_file', required=True,
                        help="safe ports to allow open", type=str)
    parser.add_argument('--extra_port_list_file', required=True,
                        help="customized tcp ports to scan", type=str)
    args = parser.parse_args()

    white_list = load_paralist_from_file(args.white_list_file)
    server_list = load_paralist_from_file(args.server_list_file)
    extra_port_list = load_paralist_from_file(args.extra_port_list_file)

    tcp_port_scan(server_list, white_list, extra_port_list)
## File : tcp_port_scan.py ends
<|endoftext|>


<commit_before># Copyright 2015 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Simple, end-to-end, LeNet-5-like convolutional MNIST model example.

This should achieve a test error of 0.8%. Please keep this model as simple and
linear as possible, it is meant as a tutorial for simple convolutional models.
Run with --self_test on the command line to exectute a short self-test.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import gzip
import os
import sys
import time

import tensorflow.python.platform

import numpy
from six.moves import urllib
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'
WORK_DIRECTORY = 'data'
IMAGE_SIZE = 28
NUM_CHANNELS = 1
PIXEL_DEPTH = 255
NUM_LABELS = 10
VALIDATION_SIZE = 5000  # Size of the validation set.
SEED = 66478  # Set to None for random seed.
BATCH_SIZE = 64
NUM_EPOCHS = 10

tf.app.flags.DEFINE_boolean("self_test", False, "True if running a self test.")
FLAGS = tf.app.flags.FLAGS

def maybe_download(filename):
  """Download the data from Yann's website, unless it's already here."""
  if not os.path.exists(WORK_DIRECTORY):
    os.mkdir(WORK_DIRECTORY)
  filepath = os.path.join(WORK_DIRECTORY, filename)
  if not os.path.exists(filepath):
    filepath, _ = urllib.request.urlretrieve(SOURCE_URL + filename, filepath)
    statinfo = os.stat(filepath)
    print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')
  return filepath

def extract_data(filename, num_images):
  """Extract the images into a 4D tensor [image index, y, x, channels].

  Values are rescaled from [0, 255] down to [-0.5, 0.5].
  """
  print('Extracting', filename)
  with gzip.open(filename) as bytestream:
    bytestream.read(16)
    buf = bytestream.read(IMAGE_SIZE * IMAGE_SIZE * num_images)
    data = numpy.frombuffer(buf, dtype=numpy.uint8).astype(numpy.float32)
    data = (data - (PIXEL_DEPTH / 2.0)) / PIXEL_DEPTH
    data = data.reshape(num_images, IMAGE_SIZE, IMAGE_SIZE, 1)
    return data

def extract_labels(filename, num_images):
  """Extract the labels into a 1-hot matrix [image index, label index]."""
  print('Extracting', filename)
  with gzip.open(filename) as bytestream:
    bytestream.read(8)
    buf = bytestream.read(1 * num_images)
    labels = numpy.frombuffer(buf, dtype=numpy.uint8)
  # Convert to dense 1-hot representation.
  return (numpy.arange(NUM_LABELS) == labels[:, None]).astype(numpy.float32)

def fake_data(num_images):
  """Generate a fake dataset that matches the dimensions of MNIST."""
  data = numpy.ndarray(
      shape=(num_images, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS),
      dtype=numpy.float32)
  labels = numpy.zeros(shape=(num_images, NUM_LABELS), dtype=numpy.float32)
  for image in xrange(num_images):
    label = image % 2
    data[image, :, :, 0] = label - 0.5
    labels[image, label] = 1.0
  return data, labels

def error_rate(predictions, labels):
  """Return the error rate based on dense predictions and 1-hot labels."""
  return 100.0 - (
      100.0 *
      numpy.sum(numpy.argmax(predictions, 1) == numpy.argmax(labels, 1)) /
      predictions.shape[0])

def main(argv=None):  # pylint: disable=unused-argument
  if FLAGS.self_test:
    print('Running self-test.')
    train_data, train_labels = fake_data(256)
    validation_data, validation_labels = fake_data(16)
    test_data, test_labels = fake_data(256)
    num_epochs = 1
  else:
    # Get the data.
    train_data_filename = maybe_download('train-images-idx3-ubyte.gz')
    train_labels_filename = maybe_download('train-labels-idx1-ubyte.gz')
    test_data_filename = maybe_download('t10k-images-idx3-ubyte.gz')
    test_labels_filename = maybe_download('t10k-labels-idx1-ubyte.gz')

    # Extract it into numpy arrays.
    train_data = extract_data(train_data_filename, 60000)
    train_labels = extract_labels(train_labels_filename, 60000)
    test_data = extract_data(test_data_filename, 10000)
    test_labels = extract_labels(test_labels_filename, 10000)

    # Generate a validation set.
    validation_data = train_data[:VALIDATION_SIZE, :, :, :]
    validation_labels = train_labels[:VALIDATION_SIZE]
    train_data = train_data[VALIDATION_SIZE:, :, :, :]
    train_labels = train_labels[VALIDATION_SIZE:]
    num_epochs = NUM_EPOCHS
  train_size = train_labels.shape[0]

  # This is where training samples and labels are fed to the graph.
  # These placeholder nodes will be fed a batch of training data at each
  # training step using the {feed_dict} argument to the Run() call below.
  train_data_node = tf.placeholder(
      tf.float32,
      shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))
  train_labels_node = tf.placeholder(tf.float32,
                                     shape=(BATCH_SIZE, NUM_LABELS))
  # For the validation and test data, we'll just hold the entire dataset in
  # one constant node.
  validation_data_node = tf.constant(validation_data)
  test_data_node = tf.constant(test_data)

  # The variables below hold all the trainable weights. They are passed an
  # initial value which will be assigned when when we call:
  # {tf.initialize_all_variables().run()}
  conv1_weights = tf.Variable(
      tf.truncated_normal([5, 5, NUM_CHANNELS, 32],  # 5x5 filter, depth 32.
                          stddev=0.1,
                          seed=SEED))
  conv1_biases = tf.Variable(tf.zeros([32]))
  conv2_weights = tf.Variable(
      tf.truncated_normal([5, 5, 32, 64],
                          stddev=0.1,
                          seed=SEED))
  conv2_biases = tf.Variable(tf.constant(0.1, shape=[64]))
  fc1_weights = tf.Variable(  # fully connected, depth 512.
      tf.truncated_normal(
          [IMAGE_SIZE // 4 * IMAGE_SIZE // 4 * 64, 512],
          stddev=0.1,
          seed=SEED))
  fc1_biases = tf.Variable(tf.constant(0.1, shape=[512]))
  fc2_weights = tf.Variable(
      tf.truncated_normal([512, NUM_LABELS],
                          stddev=0.1,
                          seed=SEED))
  fc2_biases = tf.Variable(tf.constant(0.1, shape=[NUM_LABELS]))

  # We will replicate the model structure for the training subgraph, as well
  # as the evaluation subgraphs, while sharing the trainable parameters.
  def model(data, train=False):
    """The Model definition."""
    # 2D convolution, with 'SAME' padding (i.e. the output feature map has
    # the same size as the input). Note that {strides} is a 4D array whose
    # shape matches the data layout: [image index, y, x, depth].
    conv = tf.nn.conv2d(data,
                        conv1_weights,
                        strides=[1, 1, 1, 1],
                        padding='SAME')
    # Bias and rectified linear non-linearity.
    relu = tf.nn.relu(tf.nn.bias_add(conv, conv1_biases))
    # Max pooling. The kernel size spec {ksize} also follows the layout of
    # the data. Here we have a pooling window of 2, and a stride of 2.
    pool = tf.nn.max_pool(relu,
                          ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1],
                          padding='SAME')
    conv = tf.nn.conv2d(pool,
                        conv2_weights,
                        strides=[1, 1, 1, 1],
                        padding='SAME')
    relu = tf.nn.relu(tf.nn.bias_add(conv, conv2_biases))
    pool = tf.nn.max_pool(relu,
                          ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1],
                          padding='SAME')
    # Reshape the feature map cuboid into a 2D matrix to feed it to the
    # fully connected layers.
    pool_shape = pool.get_shape().as_list()
    reshape = tf.reshape(
        pool,
        [pool_shape[0], pool_shape[1] * pool_shape[2] * pool_shape[3]])
    # Fully connected layer. Note that the '+' operation automatically
    # broadcasts the biases.
    hidden = tf.nn.relu(tf.matmul(reshape, fc1_weights) + fc1_biases)
    # Add a 50% dropout during training only. Dropout also scales
    # activations such that no rescaling is needed at evaluation time.
    if train:
      hidden = tf.nn.dropout(hidden, 0.5, seed=SEED)
    return tf.matmul(hidden, fc2_weights) + fc2_biases

  # Training computation: logits + cross-entropy loss.
  logits = model(train_data_node, True)
  loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
      logits, train_labels_node))

  # L2 regularization for the fully connected parameters.
  regularizers = (tf.nn.l2_loss(fc1_weights) + tf.nn.l2_loss(fc1_biases) +
                  tf.nn.l2_loss(fc2_weights) + tf.nn.l2_loss(fc2_biases))
  # Add the regularization term to the loss.
  loss += 5e-4 * regularizers

  # Optimizer: set up a variable that's incremented once per batch and
  # controls the learning rate decay.
  batch = tf.Variable(0)
  # Decay once per epoch, using an exponential schedule starting at 0.01.
  learning_rate = tf.train.exponential_decay(
      0.01,                # Base learning rate.
      batch * BATCH_SIZE,  # Current index into the dataset.
      train_size,          # Decay step.
      0.95,                # Decay rate.
      staircase=True)
  # Use simple momentum for the optimization.
  optimizer = tf.train.MomentumOptimizer(learning_rate,
                                         0.9).minimize(loss,
                                                       global_step=batch)

  # Predictions for the minibatch, validation set and test set.
  train_prediction = tf.nn.softmax(logits)
  # We'll compute them only once in a while by calling their {eval()} method.
  validation_prediction = tf.nn.softmax(model(validation_data_node))
  test_prediction = tf.nn.softmax(model(test_data_node))

  # Create a local session to run this computation.
  start_time = time.time()
  log_each_x_steps = 100
  with tf.Session() as s:
    # Run all the initializers to prepare the trainable parameters.
    tf.initialize_all_variables().run()
    print('Initialized!')
    # Loop through training steps.
    for step in xrange(num_epochs * train_size // BATCH_SIZE):
      # Compute the offset of the current minibatch in the data.
      # Note that we could use better randomization across epochs.
      offset = (step * BATCH_SIZE) % (train_size - BATCH_SIZE)
      batch_data = train_data[offset:(offset + BATCH_SIZE), :, :, :]
      batch_labels = train_labels[offset:(offset + BATCH_SIZE)]
      # This dictionary maps the batch data (as a numpy array) to the
      # node in the graph is should be fed to.
      feed_dict = {train_data_node: batch_data,
                   train_labels_node: batch_labels}
      # Run the graph and fetch some of the nodes.
      _, l, lr, predictions = s.run(
          [optimizer, loss, learning_rate, train_prediction],
          feed_dict=feed_dict)
      if step % log_each_x_steps == 0:
        elapsed_time = time.time() - start_time
        start_time = time.time()
        print('Step %d, %.1f ms'%(step, 1000 * elapsed_time / log_each_x_steps))
        print('Epoch %.2f' % (float(step) * BATCH_SIZE / train_size))
        print('Minibatch loss: %.3f, learning rate: %.6f' % (l, lr))
        print('Minibatch error: %.1f%%' % error_rate(predictions, batch_labels))
        print('Validation error: %.1f%%' %
              error_rate(validation_prediction.eval(), validation_labels))
        sys.stdout.flush()
    # Finally print the result!
    test_error = error_rate(test_prediction.eval(), test_labels)
    print('Test error: %.1f%%' % test_error)
    if FLAGS.self_test:
      print('test_error', test_error)
      assert test_error == 0.0, 'expected 0.0 test_error, got %.2f' % (
          test_error,)

if __name__ == '__main__':
  tf.app.run()
<commit_msg>Make convolutional.py easier to run on GPUs with less RAM. A few cosmetic fixes. Test error was already 0.7%, not 0.8%. https://github.com/tensorflow/tensorflow/issues/609 Change: 111624237<commit_after># Copyright 2015 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Simple, end-to-end, LeNet-5-like convolutional MNIST model example.

This should achieve a test error of 0.7%. Please keep this model as simple and
linear as possible, it is meant as a tutorial for simple convolutional models.
Run with --self_test on the command line to exectute a short self-test.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import gzip
import os
import sys
import time

import tensorflow.python.platform

import numpy
from six.moves import urllib
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'
WORK_DIRECTORY = 'data'
IMAGE_SIZE = 28
NUM_CHANNELS = 1
PIXEL_DEPTH = 255
NUM_LABELS = 10
VALIDATION_SIZE = 5000  # Size of the validation set.
SEED = 66478  # Set to None for random seed.
BATCH_SIZE = 64
NUM_EPOCHS = 10
EVAL_BATCH_SIZE = 64
EVAL_FREQUENCY = 100  # Number of steps between evaluations.

tf.app.flags.DEFINE_boolean("self_test", False, "True if running a self test.")
FLAGS = tf.app.flags.FLAGS

def maybe_download(filename):
  """Download the data from Yann's website, unless it's already here."""
  if not os.path.exists(WORK_DIRECTORY):
    os.mkdir(WORK_DIRECTORY)
  filepath = os.path.join(WORK_DIRECTORY, filename)
  if not os.path.exists(filepath):
    filepath, _ = urllib.request.urlretrieve(SOURCE_URL + filename, filepath)
    statinfo = os.stat(filepath)
    print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')
  return filepath

def extract_data(filename, num_images):
  """Extract the images into a 4D tensor [image index, y, x, channels].

  Values are rescaled from [0, 255] down to [-0.5, 0.5].
  """
  print('Extracting', filename)
  with gzip.open(filename) as bytestream:
    bytestream.read(16)
    buf = bytestream.read(IMAGE_SIZE * IMAGE_SIZE * num_images)
    data = numpy.frombuffer(buf, dtype=numpy.uint8).astype(numpy.float32)
    data = (data - (PIXEL_DEPTH / 2.0)) / PIXEL_DEPTH
    data = data.reshape(num_images, IMAGE_SIZE, IMAGE_SIZE, 1)
    return data

def extract_labels(filename, num_images):
  """Extract the labels into a 1-hot matrix [image index, label index]."""
  print('Extracting', filename)
  with gzip.open(filename) as bytestream:
    bytestream.read(8)
    buf = bytestream.read(1 * num_images)
    labels = numpy.frombuffer(buf, dtype=numpy.uint8)
  # Convert to dense 1-hot representation.
  return (numpy.arange(NUM_LABELS) == labels[:, None]).astype(numpy.float32)

def fake_data(num_images):
  """Generate a fake dataset that matches the dimensions of MNIST."""
  data = numpy.ndarray(
      shape=(num_images, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS),
      dtype=numpy.float32)
  labels = numpy.zeros(shape=(num_images, NUM_LABELS), dtype=numpy.float32)
  for image in xrange(num_images):
    label = image % 2
    data[image, :, :, 0] = label - 0.5
    labels[image, label] = 1.0
  return data, labels

def error_rate(predictions, labels):
  """Return the error rate based on dense predictions and 1-hot labels."""
  return 100.0 - (
      100.0 *
      numpy.sum(numpy.argmax(predictions, 1) == numpy.argmax(labels, 1)) /
      predictions.shape[0])

def main(argv=None):  # pylint: disable=unused-argument
  if FLAGS.self_test:
    print('Running self-test.')
    train_data, train_labels = fake_data(256)
    validation_data, validation_labels = fake_data(EVAL_BATCH_SIZE)
    test_data, test_labels = fake_data(EVAL_BATCH_SIZE)
    num_epochs = 1
  else:
    # Get the data.
    train_data_filename = maybe_download('train-images-idx3-ubyte.gz')
    train_labels_filename = maybe_download('train-labels-idx1-ubyte.gz')
    test_data_filename = maybe_download('t10k-images-idx3-ubyte.gz')
    test_labels_filename = maybe_download('t10k-labels-idx1-ubyte.gz')

    # Extract it into numpy arrays.
    train_data = extract_data(train_data_filename, 60000)
    train_labels = extract_labels(train_labels_filename, 60000)
    test_data = extract_data(test_data_filename, 10000)
    test_labels = extract_labels(test_labels_filename, 10000)

    # Generate a validation set.
    validation_data = train_data[:VALIDATION_SIZE, ...]
    validation_labels = train_labels[:VALIDATION_SIZE]
    train_data = train_data[VALIDATION_SIZE:, ...]
    train_labels = train_labels[VALIDATION_SIZE:]
    num_epochs = NUM_EPOCHS
  train_size = train_labels.shape[0]

  # This is where training samples and labels are fed to the graph.
  # These placeholder nodes will be fed a batch of training data at each
  # training step using the {feed_dict} argument to the Run() call below.
  train_data_node = tf.placeholder(
      tf.float32,
      shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))
  train_labels_node = tf.placeholder(tf.float32,
                                     shape=(BATCH_SIZE, NUM_LABELS))
  eval_data = tf.placeholder(
      tf.float32,
      shape=(EVAL_BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))

  # The variables below hold all the trainable weights. They are passed an
  # initial value which will be assigned when when we call:
  # {tf.initialize_all_variables().run()}
  conv1_weights = tf.Variable(
      tf.truncated_normal([5, 5, NUM_CHANNELS, 32],  # 5x5 filter, depth 32.
                          stddev=0.1,
                          seed=SEED))
  conv1_biases = tf.Variable(tf.zeros([32]))
  conv2_weights = tf.Variable(
      tf.truncated_normal([5, 5, 32, 64],
                          stddev=0.1,
                          seed=SEED))
  conv2_biases = tf.Variable(tf.constant(0.1, shape=[64]))
  fc1_weights = tf.Variable(  # fully connected, depth 512.
      tf.truncated_normal(
          [IMAGE_SIZE // 4 * IMAGE_SIZE // 4 * 64, 512],
          stddev=0.1,
          seed=SEED))
  fc1_biases = tf.Variable(tf.constant(0.1, shape=[512]))
  fc2_weights = tf.Variable(
      tf.truncated_normal([512, NUM_LABELS],
                          stddev=0.1,
                          seed=SEED))
  fc2_biases = tf.Variable(tf.constant(0.1, shape=[NUM_LABELS]))

  # We will replicate the model structure for the training subgraph, as well
  # as the evaluation subgraphs, while sharing the trainable parameters.
  def model(data, train=False):
    """The Model definition."""
    # 2D convolution, with 'SAME' padding (i.e. the output feature map has
    # the same size as the input). Note that {strides} is a 4D array whose
    # shape matches the data layout: [image index, y, x, depth].
    conv = tf.nn.conv2d(data,
                        conv1_weights,
                        strides=[1, 1, 1, 1],
                        padding='SAME')
    # Bias and rectified linear non-linearity.
    relu = tf.nn.relu(tf.nn.bias_add(conv, conv1_biases))
    # Max pooling. The kernel size spec {ksize} also follows the layout of
    # the data. Here we have a pooling window of 2, and a stride of 2.
    pool = tf.nn.max_pool(relu,
                          ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1],
                          padding='SAME')
    conv = tf.nn.conv2d(pool,
                        conv2_weights,
                        strides=[1, 1, 1, 1],
                        padding='SAME')
    relu = tf.nn.relu(tf.nn.bias_add(conv, conv2_biases))
    pool = tf.nn.max_pool(relu,
                          ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1],
                          padding='SAME')
    # Reshape the feature map cuboid into a 2D matrix to feed it to the
    # fully connected layers.
    pool_shape = pool.get_shape().as_list()
    reshape = tf.reshape(
        pool,
        [pool_shape[0], pool_shape[1] * pool_shape[2] * pool_shape[3]])
    # Fully connected layer. Note that the '+' operation automatically
    # broadcasts the biases.
    hidden = tf.nn.relu(tf.matmul(reshape, fc1_weights) + fc1_biases)
    # Add a 50% dropout during training only. Dropout also scales
    # activations such that no rescaling is needed at evaluation time.
    if train:
      hidden = tf.nn.dropout(hidden, 0.5, seed=SEED)
    return tf.matmul(hidden, fc2_weights) + fc2_biases

  # Training computation: logits + cross-entropy loss.
  logits = model(train_data_node, True)
  loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
      logits, train_labels_node))

  # L2 regularization for the fully connected parameters.
  regularizers = (tf.nn.l2_loss(fc1_weights) + tf.nn.l2_loss(fc1_biases) +
                  tf.nn.l2_loss(fc2_weights) + tf.nn.l2_loss(fc2_biases))
  # Add the regularization term to the loss.
  loss += 5e-4 * regularizers

  # Optimizer: set up a variable that's incremented once per batch and
  # controls the learning rate decay.
  batch = tf.Variable(0)
  # Decay once per epoch, using an exponential schedule starting at 0.01.
  learning_rate = tf.train.exponential_decay(
      0.01,                # Base learning rate.
      batch * BATCH_SIZE,  # Current index into the dataset.
      train_size,          # Decay step.
      0.95,                # Decay rate.
      staircase=True)
  # Use simple momentum for the optimization.
  optimizer = tf.train.MomentumOptimizer(learning_rate,
                                         0.9).minimize(loss,
                                                       global_step=batch)

  # Predictions for the current training minibatch.
  train_prediction = tf.nn.softmax(logits)

  # Predictions for the test and validation, which we'll compute less often.
  eval_prediction = tf.nn.softmax(model(eval_data))

  # Small utility function to evaluate a dataset by feeding batches of data to
  # {eval_data} and pulling the results from {eval_predictions}.
  # Saves memory and enables this to run on smaller GPUs.
  def eval_in_batches(data, sess):
    """Get all predictions for a dataset by running it in small batches."""
    size = data.shape[0]
    if size < EVAL_BATCH_SIZE:
      raise ValueError("batch size for evals larger than dataset: %d" % size)
    predictions = numpy.ndarray(shape=(size, NUM_LABELS), dtype=numpy.float32)
    for begin in xrange(0, size, EVAL_BATCH_SIZE):
      end = begin + EVAL_BATCH_SIZE
      if end <= size:
        predictions[begin:end, :] = sess.run(
            eval_prediction,
            feed_dict={eval_data: data[begin:end, ...]})
      else:
        batch_predictions = sess.run(
            eval_prediction,
            feed_dict={eval_data: data[-EVAL_BATCH_SIZE:, ...]})
        predictions[begin:, :] = batch_predictions[begin - size:, :]
    return predictions

  # Create a local session to run the training.
  start_time = time.time()
  with tf.Session() as sess:
    # Run all the initializers to prepare the trainable parameters.
    tf.initialize_all_variables().run()
    print('Initialized!')
    # Loop through training steps.
    for step in xrange(int(num_epochs * train_size) // BATCH_SIZE):
      # Compute the offset of the current minibatch in the data.
      # Note that we could use better randomization across epochs.
      offset = (step * BATCH_SIZE) % (train_size - BATCH_SIZE)
      batch_data = train_data[offset:(offset + BATCH_SIZE), ...]
      batch_labels = train_labels[offset:(offset + BATCH_SIZE)]
      # This dictionary maps the batch data (as a numpy array) to the
      # node in the graph is should be fed to.
      feed_dict = {train_data_node: batch_data,
                   train_labels_node: batch_labels}
      # Run the graph and fetch some of the nodes.
      _, l, lr, predictions = sess.run(
          [optimizer, loss, learning_rate, train_prediction],
          feed_dict=feed_dict)
      if step % EVAL_FREQUENCY == 0:
        elapsed_time = time.time() - start_time
        start_time = time.time()
        print('Step %d (epoch %.2f), %.1f ms' %
              (step, float(step) * BATCH_SIZE / train_size,
               1000 * elapsed_time / EVAL_FREQUENCY))
        print('Minibatch loss: %.3f, learning rate: %.6f' % (l, lr))
        print('Minibatch error: %.1f%%' % error_rate(predictions, batch_labels))
        print('Validation error: %.1f%%' % error_rate(
            eval_in_batches(validation_data, sess), validation_labels))
        sys.stdout.flush()
    # Finally print the result!
    test_error = error_rate(eval_in_batches(test_data, sess), test_labels)
    print('Test error: %.1f%%' % test_error)
    if FLAGS.self_test:
      print('test_error', test_error)
      assert test_error == 0.0, 'expected 0.0 test_error, got %.2f' % (
          test_error,)

if __name__ == '__main__':
  tf.app.run()
<|endoftext|>


<commit_before># Copyright (c) 2016 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

from PyQt5.QtCore import pyqtSignal, pyqtProperty, pyqtSlot, QObject, QVariant #For communicating data and events to Qt.

import UM.Application #To get the global container stack to find the current machine.
import UM.Logger
import UM.Settings.ContainerRegistry #Finding containers by ID.

##  Manages all existing extruder stacks.
#
#   This keeps a list of extruder stacks for each machine.
class ExtruderManager(QObject):
    ##  Signal to notify other components when the list of extruders changes.
    extrudersChanged = pyqtSignal(QVariant)

    ##  Notify when the user switches the currently active extruder.
    activeExtruderChanged = pyqtSignal()

    ##  Registers listeners and such to listen to changes to the extruders.
    def __init__(self, parent = None):
        super().__init__(parent)
        self._extruder_trains = { } #Per machine, a dictionary of extruder container stack IDs.
        self._active_extruder_index = 0
        UM.Application.getInstance().globalContainerStackChanged.connect(self._addCurrentMachineExtruders)

    ##  Gets the unique identifier of the currently active extruder stack.
    #
    #   The currently active extruder stack is the stack that is currently being
    #   edited.
    #
    #   \return The unique ID of the currently active extruder stack.
    @pyqtProperty(str, notify = activeExtruderChanged)
    def activeExtruderStackId(self):
        if not UM.Application.getInstance().getGlobalContainerStack():
            return None #No active machine, so no active extruder.
        try:
            return self._extruder_trains[UM.Application.getInstance().getGlobalContainerStack().getBottom().getId()][str(self._active_extruder_index)]
        except KeyError: #Extruder index could be -1 if the global tab is selected, or the entry doesn't exist if the machine definition is wrong.
            return None

    ##  The instance of the singleton pattern.
    #
    #   It's None if the extruder manager hasn't been created yet.
    __instance = None

    ##  Gets an instance of the extruder manager, or creates one if no instance
    #   exists yet.
    #
    #   This is an implementation of singleton. If an extruder manager already
    #   exists, it is re-used.
    #
    #   \return The extruder manager.
    @classmethod
    def getInstance(cls):
        if not cls.__instance:
            cls.__instance = ExtruderManager()
        return cls.__instance

    @pyqtSlot(int)
    def setActiveExtruderIndex(self, index):
        self._active_extruder_index = index
        self.activeExtruderChanged.emit()

    ##  Adds all extruders of a specific machine definition to the extruder
    #   manager.
    #
    #   \param machine_definition The machine to add the extruders for.
    def addMachineExtruders(self, machine_definition):
        machine_id = machine_definition.getId()
        if machine_id not in self._extruder_trains:
            self._extruder_trains[machine_id] = { }

        container_registry = UM.Settings.ContainerRegistry.getInstance()
        if not container_registry: #Then we shouldn't have any machine definition either. In any case, there are no extruder trains then so bye bye.
            return

        #Add the extruder trains that don't exist yet.
        for extruder_definition in container_registry.findDefinitionContainers(machine = machine_definition.getId()):
            position = extruder_definition.getMetaDataEntry("position", None)
            if not position:
                UM.Logger.log("w", "Extruder definition %s specifies no position metadata entry.", extruder_definition.getId())
            if not container_registry.findContainerStacks(machine = machine_id, position = position): #Doesn't exist yet.
                self.createExtruderTrain(extruder_definition, machine_definition, position)

        #Gets the extruder trains that we just created as well as any that still existed.
        extruder_trains = container_registry.findContainerStacks(type = "extruder_train", machine = machine_definition.getId())
        for extruder_train in extruder_trains:
            self._extruder_trains[machine_id][extruder_train.getMetaDataEntry("position")] = extruder_train.getId()
        if extruder_trains:
            self.extrudersChanged.emit(machine_definition)

    ##  Creates a container stack for an extruder train.
    #
    #   The container stack has an extruder definition at the bottom, which is
    #   linked to a machine definition. Then it has a nozzle profile, a material
    #   profile, a quality profile and a user profile, in that order.
    #
    #   The resulting container stack is added to the registry.
    #
    #   \param extruder_definition The extruder to create the extruder train
    #   for.
    #   \param machine_definition The machine that the extruder train belongs
    #   to.
    #   \param position The position of this extruder train in the extruder
    #   slots of the machine.
    def createExtruderTrain(self, extruder_definition, machine_definition, position):
        #Cache some things.
        container_registry = UM.Settings.ContainerRegistry.getInstance()
        machine_id = machine_definition.getId()

        #Create a container stack for this extruder.
        extruder_stack_id = container_registry.uniqueName(extruder_definition.getId())
        container_stack = UM.Settings.ContainerStack(extruder_stack_id)
        container_stack.setName(extruder_definition.getName()) #Take over the display name to display the stack with.
        container_stack.addMetaDataEntry("type", "extruder_train")
        container_stack.addMetaDataEntry("machine", machine_definition.getId())
        container_stack.addMetaDataEntry("position", position)
        container_stack.addContainer(extruder_definition)

        #Find the nozzle to use for this extruder.
        nozzle = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_nozzles", default = "False") == "True":
            #First add any nozzle. Later, overwrite with preference if the preference is valid.
            nozzles = container_registry.findInstanceContainers(machine = machine_id, type = "nozzle")
            if len(nozzles) >= 1:
                nozzle = nozzles[0]
            preferred_nozzle_id = machine_definition.getMetaDataEntry("preferred_nozzle")
            if preferred_nozzle_id:
                preferred_nozzles = container_registry.findInstanceContainers(id = preferred_nozzle_id, type = "nozzle")
                if len(preferred_nozzles) >= 1:
                    nozzle = preferred_nozzles[0]
                else:
                    UM.Logger.log("w", "The preferred nozzle \"%s\" of machine %s doesn't exist or is not a nozzle profile.", preferred_nozzle_id, machine_id)
                    #And leave it at the default nozzle.
        container_stack.addContainer(nozzle)

        #Find a material to use for this nozzle.
        material = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_materials", default = "False") == "True":
            #First add any material. Later, overwrite with preference if the preference is valid.
            if machine_definition.getMetaDataEntry("has_nozzle_materials", default = "False") == "True":
                materials = container_registry.findInstanceContainers(type = "material", machine = machine_id, nozzle = nozzle.getId())
            else:
                materials = container_registry.findInstanceContainers(type = "material", machine = machine_id)
            if len(materials) >= 1:
                material = materials[0]
            preferred_material_id = machine_definition.getMetaDataEntry("preferred_material")
            if preferred_material_id:
                preferred_materials = container_registry.findInstanceContainers(id = preferred_material_id, type = "material")
                if len(preferred_materials) >= 1:
                    material = preferred_materials[0]
                else:
                    UM.Logger.log("w", "The preferred material \"%s\" of machine %s doesn't exist or is not a material profile.", preferred_material_id, machine_id)
                    #And leave it at the default material.
        container_stack.addContainer(material)

        #Find a quality to use for this extruder.
        quality = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_machine_quality"):
            #First add any quality. Later, overwrite with preference if the preference is valid.
            qualities = container_registry.findInstanceContainers(type = "quality")
            if len(qualities) >= 1:
                quality = qualities[0]
            preferred_quality_id = machine_definition.getMetaDataEntry("preferred_quality")
            if preferred_quality_id:
                preferred_quality = container_registry.findInstanceContainers(id = preferred_quality_id.lower(), type = "quality")
                if len(preferred_quality) >= 1:
                    quality = preferred_quality[0]
                else:
                    UM.Logger.log("w", "The preferred quality \"%s\" of machine %s doesn't exist or is not a quality profile.", preferred_quality_id, machine_id)
                    #And leave it at the default quality.
        container_stack.addContainer(quality)

        user_profile = container_registry.findInstanceContainers(id = extruder_stack_id + "_current_settings")
        if user_profile: #There was already a user profile, loaded from settings.
            user_profile = user_profile[0]
        else:
            user_profile = UM.Settings.InstanceContainer(extruder_stack_id + "_current_settings") #Add an empty user profile.
            user_profile.addMetaDataEntry("type", "user")
            user_profile.setDefinition(machine_definition)
            container_registry.addContainer(user_profile)
        container_stack.addContainer(user_profile)

        container_stack.setNextStack(UM.Application.getInstance().getGlobalContainerStack())

        container_registry.addContainer(container_stack)

    ##  Generates extruders for a specific machine.
    def getMachineExtruders(self, machine_definition):
        container_registry = UM.Settings.ContainerRegistry.getInstance()
        machine_id = machine_definition.getId()
        if not machine_id in self._extruder_trains:
            UM.Logger.log("w", "Tried to get the extruder trains for machine %s, which doesn't exist.", machine_id)
            return
        for _,extruder_train_id in self._extruder_trains[machine_id].items():
            extruder_train = container_registry.findContainerStacks(id = extruder_train_id)
            if extruder_train:
                yield extruder_train[0]
            else:
                UM.Logger.log("w", "Machine %s refers to an extruder train with ID %s, which doesn't exist.", machine_id, extruder_train_id)

    ##  Adds the extruders of the currently active machine.
    def _addCurrentMachineExtruders(self):
        global_stack = UM.Application.getInstance().getGlobalContainerStack()
        if global_stack and global_stack.getBottom():
            self.addMachineExtruders(global_stack.getBottom())<commit_msg>Store the actual extruder stack in ExtruderManager instead of just the ID<commit_after># Copyright (c) 2016 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

from PyQt5.QtCore import pyqtSignal, pyqtProperty, pyqtSlot, QObject, QVariant #For communicating data and events to Qt.

import UM.Application #To get the global container stack to find the current machine.
import UM.Logger
import UM.Settings.ContainerRegistry #Finding containers by ID.

##  Manages all existing extruder stacks.
#
#   This keeps a list of extruder stacks for each machine.
class ExtruderManager(QObject):
    ##  Signal to notify other components when the list of extruders changes.
    extrudersChanged = pyqtSignal(QVariant)

    ##  Notify when the user switches the currently active extruder.
    activeExtruderChanged = pyqtSignal()

    ##  Registers listeners and such to listen to changes to the extruders.
    def __init__(self, parent = None):
        super().__init__(parent)
        self._extruder_trains = { } #Per machine, a dictionary of extruder container stack IDs.
        self._active_extruder_index = 0
        UM.Application.getInstance().globalContainerStackChanged.connect(self._addCurrentMachineExtruders)

    ##  Gets the unique identifier of the currently active extruder stack.
    #
    #   The currently active extruder stack is the stack that is currently being
    #   edited.
    #
    #   \return The unique ID of the currently active extruder stack.
    @pyqtProperty(str, notify = activeExtruderChanged)
    def activeExtruderStackId(self):
        if not UM.Application.getInstance().getGlobalContainerStack():
            return None #No active machine, so no active extruder.
        try:
            return self._extruder_trains[UM.Application.getInstance().getGlobalContainerStack().getBottom().getId()][str(self._active_extruder_index)].getId()
        except KeyError: #Extruder index could be -1 if the global tab is selected, or the entry doesn't exist if the machine definition is wrong.
            return None

    ##  The instance of the singleton pattern.
    #
    #   It's None if the extruder manager hasn't been created yet.
    __instance = None

    ##  Gets an instance of the extruder manager, or creates one if no instance
    #   exists yet.
    #
    #   This is an implementation of singleton. If an extruder manager already
    #   exists, it is re-used.
    #
    #   \return The extruder manager.
    @classmethod
    def getInstance(cls):
        if not cls.__instance:
            cls.__instance = ExtruderManager()
        return cls.__instance

    @pyqtSlot(int)
    def setActiveExtruderIndex(self, index):
        self._active_extruder_index = index
        self.activeExtruderChanged.emit()

    ##  Adds all extruders of a specific machine definition to the extruder
    #   manager.
    #
    #   \param machine_definition The machine to add the extruders for.
    def addMachineExtruders(self, machine_definition):
        machine_id = machine_definition.getId()
        if machine_id not in self._extruder_trains:
            self._extruder_trains[machine_id] = { }

        container_registry = UM.Settings.ContainerRegistry.getInstance()
        if not container_registry: #Then we shouldn't have any machine definition either. In any case, there are no extruder trains then so bye bye.
            return

        #Add the extruder trains that don't exist yet.
        for extruder_definition in container_registry.findDefinitionContainers(machine = machine_definition.getId()):
            position = extruder_definition.getMetaDataEntry("position", None)
            if not position:
                UM.Logger.log("w", "Extruder definition %s specifies no position metadata entry.", extruder_definition.getId())
            if not container_registry.findContainerStacks(machine = machine_id, position = position): #Doesn't exist yet.
                self.createExtruderTrain(extruder_definition, machine_definition, position)

        #Gets the extruder trains that we just created as well as any that still existed.
        extruder_trains = container_registry.findContainerStacks(type = "extruder_train", machine = machine_definition.getId())
        for extruder_train in extruder_trains:
            self._extruder_trains[machine_id][extruder_train.getMetaDataEntry("position")] = extruder_train
        if extruder_trains:
            self.extrudersChanged.emit(machine_definition)

    ##  Creates a container stack for an extruder train.
    #
    #   The container stack has an extruder definition at the bottom, which is
    #   linked to a machine definition. Then it has a nozzle profile, a material
    #   profile, a quality profile and a user profile, in that order.
    #
    #   The resulting container stack is added to the registry.
    #
    #   \param extruder_definition The extruder to create the extruder train
    #   for.
    #   \param machine_definition The machine that the extruder train belongs
    #   to.
    #   \param position The position of this extruder train in the extruder
    #   slots of the machine.
    def createExtruderTrain(self, extruder_definition, machine_definition, position):
        #Cache some things.
        container_registry = UM.Settings.ContainerRegistry.getInstance()
        machine_id = machine_definition.getId()

        #Create a container stack for this extruder.
        extruder_stack_id = container_registry.uniqueName(extruder_definition.getId())
        container_stack = UM.Settings.ContainerStack(extruder_stack_id)
        container_stack.setName(extruder_definition.getName()) #Take over the display name to display the stack with.
        container_stack.addMetaDataEntry("type", "extruder_train")
        container_stack.addMetaDataEntry("machine", machine_definition.getId())
        container_stack.addMetaDataEntry("position", position)
        container_stack.addContainer(extruder_definition)

        #Find the nozzle to use for this extruder.
        nozzle = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_nozzles", default = "False") == "True":
            #First add any nozzle. Later, overwrite with preference if the preference is valid.
            nozzles = container_registry.findInstanceContainers(machine = machine_id, type = "nozzle")
            if len(nozzles) >= 1:
                nozzle = nozzles[0]
            preferred_nozzle_id = machine_definition.getMetaDataEntry("preferred_nozzle")
            if preferred_nozzle_id:
                preferred_nozzles = container_registry.findInstanceContainers(id = preferred_nozzle_id, type = "nozzle")
                if len(preferred_nozzles) >= 1:
                    nozzle = preferred_nozzles[0]
                else:
                    UM.Logger.log("w", "The preferred nozzle \"%s\" of machine %s doesn't exist or is not a nozzle profile.", preferred_nozzle_id, machine_id)
                    #And leave it at the default nozzle.
        container_stack.addContainer(nozzle)

        #Find a material to use for this nozzle.
        material = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_materials", default = "False") == "True":
            #First add any material. Later, overwrite with preference if the preference is valid.
            if machine_definition.getMetaDataEntry("has_nozzle_materials", default = "False") == "True":
                materials = container_registry.findInstanceContainers(type = "material", machine = machine_id, nozzle = nozzle.getId())
            else:
                materials = container_registry.findInstanceContainers(type = "material", machine = machine_id)
            if len(materials) >= 1:
                material = materials[0]
            preferred_material_id = machine_definition.getMetaDataEntry("preferred_material")
            if preferred_material_id:
                preferred_materials = container_registry.findInstanceContainers(id = preferred_material_id, type = "material")
                if len(preferred_materials) >= 1:
                    material = preferred_materials[0]
                else:
                    UM.Logger.log("w", "The preferred material \"%s\" of machine %s doesn't exist or is not a material profile.", preferred_material_id, machine_id)
                    #And leave it at the default material.
        container_stack.addContainer(material)

        #Find a quality to use for this extruder.
        quality = container_registry.getEmptyInstanceContainer()
        if machine_definition.getMetaDataEntry("has_machine_quality"):
            #First add any quality. Later, overwrite with preference if the preference is valid.
            qualities = container_registry.findInstanceContainers(type = "quality")
            if len(qualities) >= 1:
                quality = qualities[0]
            preferred_quality_id = machine_definition.getMetaDataEntry("preferred_quality")
            if preferred_quality_id:
                preferred_quality = container_registry.findInstanceContainers(id = preferred_quality_id.lower(), type = "quality")
                if len(preferred_quality) >= 1:
                    quality = preferred_quality[0]
                else:
                    UM.Logger.log("w", "The preferred quality \"%s\" of machine %s doesn't exist or is not a quality profile.", preferred_quality_id, machine_id)
                    #And leave it at the default quality.
        container_stack.addContainer(quality)

        user_profile = container_registry.findInstanceContainers(id = extruder_stack_id + "_current_settings")
        if user_profile: #There was already a user profile, loaded from settings.
            user_profile = user_profile[0]
        else:
            user_profile = UM.Settings.InstanceContainer(extruder_stack_id + "_current_settings") #Add an empty user profile.
            user_profile.addMetaDataEntry("type", "user")
            user_profile.setDefinition(machine_definition)
            container_registry.addContainer(user_profile)
        container_stack.addContainer(user_profile)

        container_stack.setNextStack(UM.Application.getInstance().getGlobalContainerStack())

        container_registry.addContainer(container_stack)

    ##  Generates extruders for a specific machine.
    def getMachineExtruders(self, machine_definition):
        container_registry = UM.Settings.ContainerRegistry.getInstance()
        machine_id = machine_definition.getId()
        if not machine_id in self._extruder_trains:
            UM.Logger.log("w", "Tried to get the extruder trains for machine %s, which doesn't exist.", machine_id)
            return
        for name in self._extruder_trains[machine_id]:
            yield self._extruder_trains[machine_id][name]

    ##  Adds the extruders of the currently active machine.
    def _addCurrentMachineExtruders(self):
        global_stack = UM.Application.getInstance().getGlobalContainerStack()
        if global_stack and global_stack.getBottom():
            self.addMachineExtruders(global_stack.getBottom())
<|endoftext|>


<commit_before><commit_msg>Add substantial filtering to NED fetch to only retrieve US files.<commit_after><|endoftext|>