def get_jobs(pypi_fetcher_dir, bucket, processed, amount=1000): pypi_dict = LazyBucketDict(f"{pypi_fetcher_dir}/pypi", restrict_to_bucket=bucket) jobs = [] names = list(pypi_dict.by_bucket(bucket).keys()) total_nr = 0 for pkg_name in names: for ver, release_types in pypi_dict[pkg_name].items(): if 'sdist' not in release_types: continue if (pkg_name, ver) in processed: continue total_nr += 1 release = release_types['sdist'] if len(jobs) <= amount: jobs.append( PackageJob( bucket, pkg_name, ver, f"https://files.pythonhosted.org/packages/source/{pkg_name[0]}/{pkg_name}/{release[1]}", release[0], 0, )) shuffle(jobs) for i, job in enumerate(jobs): job.idx = i print( f"Bucket {bucket}: Planning execution of {len(jobs)} jobs out of {total_nr} total jobs for this bucket" ) return jobs
def crawl_pkgs_meta(packages, target_dir, workers): pkgs_dict = LazyBucketDict(target_dir) args_list = [(name, pkgs_dict) for name in packages] if workers > 1: utils.parallel(save_pkg_meta, zip(*args_list), workers=workers) else: [save_pkg_meta(*args) for args in args_list] pkgs_dict.save()
def names_in_buckets(): in_buckets = {} for name in all_packages(): bucket = LazyBucketDict.bucket(name.replace('_', '-').lower()) if bucket not in in_buckets: in_buckets[bucket] = [] in_buckets[bucket].append(name) return in_buckets
def get_names_per_bucket() -> Dict[str, Set[str]]: result = {} hexdigits = "0123456789abcdef" for a in hexdigits: for b in hexdigits: result[a + b] = set() keys = [p.name for p in P.select(P.name).distinct()] for key in keys: result[LazyBucketDict.bucket(key)].add(key) return result
def main(): dump_dir = sys.argv[1] for bucket_key, key_set in get_names_per_bucket().items(): pkgs_dict = LazyBucketDict(f"{dump_dir}", restrict_to_bucket=bucket_key) pkgs = P.select( P.id, P.name, P.version, P.py_ver, P.install_requires, P.setup_requires, P.extras_require, P.tests_require, P.python_requires, ).where(P.error.is_null(), P.name.in_(key_set)) print(f'dumping bucket {bucket_key}') for pkg in sorted(pkgs, key=lambda pkg: (pkg.name, pkg.version, pkg.py_ver)): py_ver = ''.join(filter(lambda c: c.isdigit(), pkg.py_ver)) insert(py_ver, pkg.name, pkg.version, pkg_to_dict(pkg), pkgs_dict) compress(pkgs_dict) pkgs_dict.save()
def purge(pypi_index, pkgs_dict: LazyBucketDict, bucket, py_vers): # purge all versions which are not on pypi anymore for name, vers in pkgs_dict.by_bucket(bucket).copy().items(): if name not in pypi_index: print( f"deleting package {name} from DB because it has been removed from pypi" ) del pkgs_dict[name] continue for ver in tuple(vers.keys()): if ver not in pypi_index[name]: print( f"deleting package {name} version {ver} from DB because it has been removed from pypi" ) del pkgs_dict[name][ver] # purge old python versions for name, vers in pkgs_dict.by_bucket(bucket).copy().items(): for ver, pyvers in vers.copy().items(): for pyver in tuple(pyvers.keys()): if pyver not in py_vers: print( f"deleting package {name} version {ver} for python {pyver}" f" from DB because we dropped support for this python version" ) del pkgs_dict[name][ver][pyver] if len(pkgs_dict[name][ver]) == 0: print( f"deleting package {name} version {ver} from DB" f" because it is not compatible with any of our supported python versions" ) del pkgs_dict[name][ver] if len(pkgs_dict[name]) == 0: print( f"deleting package {name} from DB" f" because it has no releases left which are compatible with any of our supported python versions" ) del pkgs_dict[name]
def get_jobs(bucket, pypi_dict: LazyBucketDict, dump_dict: LazyBucketDict): names = list(pypi_dict.by_bucket(bucket).keys()) jobs = [] for pkg_name in names: for ver, release_types in pypi_dict[pkg_name].items(): if 'wheels' not in release_types: continue for filename, data in release_types['wheels'].items(): pyver = data[1] if is_done(dump_dict, pkg_name, ver, pyver, filename): continue url = construct_url(pkg_name, pyver, filename) jobs.append( dict(name=pkg_name, ver=ver, filename=filename, pyver=pyver, url=url, bucket=bucket)) shuffle(jobs) return [Job(**j, nr=idx) for idx, j in enumerate(jobs)]
def main(): workers = int(os.environ.get('WORKERS', "1")) pypi_fetcher_dir = os.environ.get('pypi_fetcher', '/tmp/pypi_fetcher') ensure_pypi_fetcher(pypi_fetcher_dir) init_db() build_base(store=os.environ.get('STORE', None)) P = Package with Measure('Get processed pkgs from DB'): processed = set((p.name, p.version) for p in P.select(P.name, P.version).distinct()) print(f"DB contains {len(processed)} pkgs at this time") for bucket in LazyBucketDict.bucket_keys(): with Measure("getting jobs"): jobs = get_jobs(pypi_fetcher_dir, bucket, processed, amount=1000) if not jobs: continue with Measure('batch'): if workers > 1: pool_results = utils.parallel(extract_requirements, (jobs, ), workers=workers, use_processes=False) else: pool_results = [extract_requirements(args) for args in jobs] results = [] for i, res in enumerate(pool_results): if isinstance(res, Exception): print(f"Problem with {jobs[i].name}:{jobs[i].version}") if isinstance(res, sp.CalledProcessError): print(res.stderr) traceback.print_exception(res, res, res.__traceback__) else: for r in res: results.append(r) sleep(1) with db.atomic(): with Measure('bulk insert'): Package.bulk_create([Package(**r) for r in results]) if os.environ.get('CLEANUP', None): cleanup()
def main(): dump_dir = sys.argv[1] workers = int(os.environ.get('WORKERS', "1")) pypi_fetcher_dir = os.environ.get('pypi_fetcher') print(f'Index directory: {pypi_fetcher_dir}') assert isdir(pypi_fetcher_dir) for bucket in LazyBucketDict.bucket_keys(): pypi_dict = LazyBucketDict(f"{pypi_fetcher_dir}/pypi") dump_dict = LazyBucketDict(dump_dir, restrict_to_bucket=bucket) print(f"Prune bucket {bucket}") prune_entries(bucket, pypi_dict, dump_dict) pypi_dict.save() dump_dict.save() print(f"Calculating jobs for bucket {bucket}") jobs = list(get_jobs(bucket, pypi_dict, dump_dict)) if not jobs: continue print(f"Starting batch with {len(jobs)} jobs") func = mine_wheel_metadata_full_download if workers > 1: def f(job): return exec_or_return_exc(func, job) result = parallel(f, (jobs, ), workers=workers) else: result = [exec_or_return_exc(func, job) for job in jobs] for r in result: if isinstance(r, Exception): continue name = r.job.name ver = r.job.ver pyver = r.job.pyver fn = r.job.filename if name not in dump_dict: dump_dict[name] = {} if pyver not in dump_dict[name]: dump_dict[name][pyver] = {} if ver not in dump_dict[name][pyver]: dump_dict[name][pyver][ver] = {} dump_dict[name][pyver][ver][fn] = {} for key in ('requires_dist', 'provides_extras', 'requires_external', 'requires_python'): val = getattr(r, key) if val: dump_dict[name][pyver][ver][fn][key] = val compress(dump_dict) dump_dict.save()
def main(): # settings related to performance/parallelization amount_buckets = int(os.environ.get('AMOUNT_BUCKETS', "256")) limit_names = set( filter(lambda n: bool(n), os.environ.get('LIMIT_NAMES', "").split(','))) max_minutes = int(os.environ.get('MAX_MINUTES', "0")) bucket_jobs = int(os.environ.get('BUCKET_JOBS', "0")) start_bucket = int(os.environ.get('BUCKET_START', "0")) workers = int(os.environ.get('WORKERS', multiprocessing.cpu_count() * 2)) # general settings dump_dir = os.environ.get('DUMP_DIR', "./sdist") extractor_src = os.environ.get("EXTRACTOR_SRC") if not extractor_src: raise Exception( "Set env variable 'EXTRACTOR_SRC to {mach-nix}/lib/extractor'") min_free_gb = int(os.environ.get('MIN_FREE_GB', "0")) py_vers_short = os.environ.get('PYTHON_VERSIONS', "27,36,37,38,39,310").strip().split(',') pypi_fetcher_dir = os.environ.get('PYPI_FETCHER', '/tmp/pypi_fetcher') store = os.environ.get('STORE', None) deadline_total = time() + max_minutes * 60 if max_minutes else None # cache build time deps, otherwise first job will be slow with Measure("ensure build time deps"): build_base(extractor_src, py_vers_short, store=store) garbage_collected = False for idx, bucket in enumerate(LazyBucketDict.bucket_keys()): # calculate per bucket deadline if MAX_MINUTES is used if deadline_total: amount = min(amount_buckets, 256 - start_bucket) deadline = time() + (deadline_total - time()) / amount else: deadline = None if idx < start_bucket or idx >= start_bucket + amount_buckets: continue pkgs_dict = LazyBucketDict(dump_dir, restrict_to_bucket=bucket) pypi_index = LazyBucketDict(f"{pypi_fetcher_dir}/pypi", restrict_to_bucket=bucket) # load error data error_dict = LazyBucketDict(dump_dir + "-errors", restrict_to_bucket=bucket) decompress(error_dict.by_bucket(bucket)) with Measure('Get processed pkgs'): print( f"DB contains {len(list(pkgs_dict.keys()))} pkgs at this time for bucket {bucket}" ) with Measure("decompressing data"): decompress(pkgs_dict.by_bucket(bucket)) # purge data for old python versions and packages which got deleted from pypi with Measure("purging packages"): purge(pypi_index, pkgs_dict, bucket, py_vers_short) with Measure("getting jobs"): jobs = get_jobs(pypi_index, error_dict, pkgs_dict, bucket, py_vers_short, limit_num=bucket_jobs, limit_names=limit_names) if not jobs: continue compute_drvs(jobs, extractor_src, store=store) # ensure that all the build time dependencies are cached before starting, # otherwise jobs might time out if garbage_collected: with Measure("ensure build time deps"): build_base(extractor_src, py_vers_short, store=store) with Measure('executing jobs'): if workers > 1: pool_results = utils.parallel(extract_requirements, (jobs, (deadline, ) * len(jobs), (len(jobs), ) * len(jobs), (store, ) * len(jobs)), workers=workers, use_processes=False) else: pool_results = [ extract_requirements(args, deadline, store) for args in jobs ] # filter out exceptions results = [] for i, res in enumerate(pool_results): if not isinstance(res, Exception): for r in res: results.append(r) # insert new data for pkg in sorted(results, key=lambda pkg: (pkg.name, pkg.version, sort_key_pyver(pkg.py_ver))): py_ver = ''.join(filter(lambda c: c.isdigit(), pkg.py_ver)) if pkg.error: target = error_dict else: target = pkgs_dict insert(py_ver, pkg.name, pkg.version, pkg_to_dict(pkg), target, error=pkg.error) # compress and save with Measure("compressing data"): compress(pkgs_dict.by_bucket(bucket)) compress(error_dict.by_bucket(bucket)) print("finished compressing data") with Measure("saving data"): pkgs_dict.save() error_dict.save() # collect garbage if free space < MIN_FREE_GB if shutil.disk_usage(store or "/nix/store").free / (1000** 3) < min_free_gb: with Measure("collecting nix store garbage"): sp.run( f"nix-collect-garbage {f'--store {store}' if store else ''}", capture_output=True, shell=True) garbage_collected = True # stop execution if deadline occurred if deadline_total and time() > deadline_total: print( f"Deadline occurred. Stopping execution. Last Bucket was {bucket}" ) break
def decompress(pkgs_dict: LazyBucketDict): for name, vers in pkgs_dict.items(): decompress_dict(vers) for ver, pyvers in vers.items(): decompress_dict(pyvers)